harness_zbt/backend/tests/test_agent_knowledge_citations.py

import json
from datetime import UTC, datetime
from typing import Any

from langchain_core.language_models.chat_models import BaseChatModel
from langchain_core.messages import AIMessage
from langchain_core.outputs import ChatGeneration, ChatResult
from pydantic import Field

from zbt.core.config import Settings
from zbt.domains.agent.tools import build_agent_tool_registry
from zbt.domains.catalog.repository import InMemoryCatalogRepository
from zbt.domains.catalog.service import ProductCatalogService
from zbt.domains.enrollment.repository import InMemoryEnrollmentRepository
from zbt.domains.enrollment.service import EnrollmentService
from zbt.domains.identity.models import H5User
from zbt.domains.knowledge.index import InMemoryKnowledgeIndex
from zbt.domains.knowledge.repository import InMemoryKnowledgeRepository
from zbt.domains.knowledge.service import KnowledgeService
from zbt.harness.kernel import (
    AgentInvocation,
    AgentKernel,
    _infer_cited_chunk_ids,
)

NOW = datetime(2026, 7, 28, 10, tzinfo=UTC)


class CitationAwareTestModel(BaseChatModel):
    cited_chunk_ids: list[str] = Field(default_factory=list)
    invocation_count: int = 0
    structured_tool_name: str | None = None

    @property
    def _llm_type(self) -> str:
        return "citation-aware-test-model"

    def bind_tools(
        self,
        tools: Any,
        *,
        tool_choice: str | None = None,
        **kwargs: Any,
    ) -> Any:
        del tool_choice, kwargs
        names = [_tool_name(tool) for tool in tools]
        self.structured_tool_name = next(
            (name for name in names if name == "AgentStructuredOutput"),
            None,
        )
        return self

    def _generate(
        self,
        messages: Any,
        stop: list[str] | None = None,
        run_manager: Any = None,
        **kwargs: Any,
    ) -> ChatResult:
        del messages, stop, run_manager, kwargs
        self.invocation_count += 1
        if self.invocation_count == 1:
            message = AIMessage(
                content="",
                tool_calls=[
                    {
                        "name": "search_insurance_knowledge",
                        "args": {
                            "query": "安心医疗险是否报销海外整形手术，比例是多少？",
                            "limit": 5,
                        },
                        "id": "knowledge-search-1",
                    }
                ],
            )
        elif self.structured_tool_name is not None:
            message = AIMessage(
                content="",
                tool_calls=[
                    {
                        "name": self.structured_tool_name,
                        "args": {
                            "message": (
                                "当前知识库未明确海外整形手术是否属于保障责任，"
                                "也未明确具体报销比例，无法确定。"
                            ),
                            "selected_product_codes": [],
                            "suggested_action": None,
                            "cited_knowledge_chunk_ids": self.cited_chunk_ids,
                        },
                        "id": "structured-answer-1",
                    }
                ],
            )
        else:
            marker = json.dumps(self.cited_chunk_ids)
            message = AIMessage(
                content=(
                    "当前知识库未明确海外整形手术是否属于保障责任，"
                    "也未明确具体报销比例，无法确定。\n\n"
                    f"<!-- ZBT_CITATIONS: {marker} -->"
                )
            )
        return ChatResult(generations=[ChatGeneration(message=message)])


class CitationAwareTestGateway:
    def __init__(self, model: CitationAwareTestModel) -> None:
        self._model = model

    def chat_model(self, persona: str) -> BaseChatModel:
        assert persona == "customer"
        return self._model


def test_unknown_answer_does_not_expose_retrieval_candidates_as_citations() -> None:
    kernel, user, _, _ = _kernel(cited_chunk_ids=[])

    reply = kernel.reply(
        AgentInvocation(
            persona="customer",
            message="安心医疗险是否可以报销海外整形手术？具体报销比例是多少？",
            h5_user=user,
        )
    )

    assert reply.invoked_tools == ("search_insurance_knowledge",)
    assert "无法确定" in reply.text
    assert "ZBT_CITATIONS" not in reply.text
    assert not any(card["type"] == "knowledge_sources" for card in reply.cards)


def test_answer_only_exposes_selected_candidates_as_citations() -> None:
    kernel, user, chunk_ids, model = _kernel(cited_chunk_ids=[])
    model.cited_chunk_ids = [
        chunk_ids[0],
        chunk_ids[1],
        "not-returned-by-the-knowledge-tool",
        chunk_ids[0],
    ]

    reply = kernel.reply(
        AgentInvocation(
            persona="customer",
            message="非医学必需的美容项目是否属于保障责任？",
            h5_user=user,
        )
    )

    source_cards = [
        card for card in reply.cards if card["type"] == "knowledge_sources"
    ]
    assert len(source_cards) == 1
    assert len(source_cards[0]["items"]) == 1
    source = source_cards[0]["items"][0]
    assert source["excerpts"] == [
        "非医学必需的美容项目不属于本产品保障责任。",
        "具体报销比例以电子保单对应的保障计划为准。",
    ]


def test_exact_quoted_knowledge_can_recover_missing_model_citation_marker() -> None:
    text = (
        "依据《安心医疗险保险条款》（版本 v1）："
        "疾病医疗责任的等待期为30天。"
    )
    executions = [
        {
            "chunk_id": "doc-1:1:1",
            "title": "安心医疗险保险条款",
            "source_name": "安心医疗险保险条款.md",
            "content": "疾病医疗责任的等待期为30天。因意外伤害不受等待期限制。",
        },
        {
            "chunk_id": "doc-1:1:2",
            "title": "安心医疗险保险条款",
            "source_name": "安心医疗险保险条款.md",
            "content": "具体责任以电子保单对应的保障计划为准。",
        },
    ]

    assert _infer_cited_chunk_ids(text, executions) == ["doc-1:1:1"]


def _kernel(
    *,
    cited_chunk_ids: list[str],
) -> tuple[AgentKernel, H5User, list[str], CitationAwareTestModel]:
    knowledge = KnowledgeService(
        InMemoryKnowledgeRepository(),
        lambda: NOW,
        InMemoryKnowledgeIndex(),
    )
    document = knowledge.create_document(
        title="安心医疗险保险条款",
        document_type="INSURANCE_TERMS",
        source_name="安心医疗险保险条款.md",
        product_code="MED-BASIC",
        content=(
            "非医学必需的美容项目不属于本产品保障责任。\n\n"
            "具体报销比例以电子保单对应的保障计划为准。"
        ),
        created_by="admin-01",
    )
    knowledge.index_document(document["document_id"])
    knowledge.test_search(document["document_id"], "美容项目是否保障")
    knowledge.publish_document(document["document_id"])
    chunk_ids = [
        f"{document['document_id']}:1:1",
        f"{document['document_id']}:1:2",
    ]

    catalog = ProductCatalogService(InMemoryCatalogRepository(), lambda: NOW)
    enrollment = EnrollmentService(
        InMemoryEnrollmentRepository(),
        catalog,
        lambda: NOW,
        "x" * 32,
    )
    tools = build_agent_tool_registry(catalog, enrollment, knowledge=knowledge)
    user = H5User(
        id="01H5USER000000000000000001",
        mobile="18800000001",
        mobile_masked="188****0001",
        display_name="测试用户",
        status="ACTIVE",
        created_at=NOW,
    )
    settings = Settings(
        app_env="test",
        jwt_access_secret="a" * 32,
        jwt_refresh_secret="b" * 32,
        field_encryption_key="c" * 32,
        langsmith_tracing=False,
    )
    model = CitationAwareTestModel(cited_chunk_ids=cited_chunk_ids)
    return (
        AgentKernel(
            settings=settings,
            model_gateway=CitationAwareTestGateway(model),
            tools=tools,
        ),
        user,
        chunk_ids,
        model,
    )


def _tool_name(tool: Any) -> str:
    if hasattr(tool, "name"):
        return str(tool.name)
    if isinstance(tool, dict):
        function = tool.get("function", {})
        if isinstance(function, dict):
            return str(function.get("name", ""))
    return ""