harness_zbt/backend/tests/test_knowledge_tool.py

import json
from datetime import UTC, datetime

import pytest

from zbt.core.errors import AppError
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.policy import HarnessPolicyEngine
from zbt.harness.tooling import ToolExecutionContext

NOW = datetime(2026, 7, 27, 4, 0, tzinfo=UTC)


def test_customer_knowledge_tool_returns_published_evidence_with_source() -> None:
    catalog = ProductCatalogService(InMemoryCatalogRepository(), lambda: NOW)
    enrollment = EnrollmentService(
        InMemoryEnrollmentRepository(),
        catalog,
        lambda: NOW,
        "x" * 32,
    )
    knowledge = KnowledgeService(
        InMemoryKnowledgeRepository(),
        lambda: NOW,
        InMemoryKnowledgeIndex(),
    )
    document = knowledge.create_document(
        title="安心医疗险保险条款",
        document_type="INSURANCE_TERMS",
        source_name="安心医疗险条款.pdf",
        product_code="MED-BASIC",
        content="疾病医疗等待期为30天，意外伤害不受等待期限制。",
        created_by="admin-01",
    )
    knowledge.index_document(document["document_id"])
    knowledge.test_search(document["document_id"], "医疗险等待期是多少天？")
    knowledge.publish_document(document["document_id"])
    registry = build_agent_tool_registry(catalog, enrollment, knowledge=knowledge)
    context = ToolExecutionContext(
        persona="customer",
        h5_user=H5User(
            id="01H5USER000000000000000001",
            mobile="18800000001",
            mobile_masked="188****0001",
            display_name="体验用户",
            status="ACTIVE",
            created_at=NOW,
        ),
    )
    tool = registry.build(
        names=("search_insurance_knowledge",),
        context=context,
        policy_engine=HarnessPolicyEngine(),
    )[0]

    result = json.loads(tool.invoke({"query": "医疗险等待期是多少天？", "limit": 3}))

    assert result["data"]["total"] == 1
    assert result["data"]["items"][0]["title"] == "安心医疗险保险条款"
    assert result["data"]["items"][0]["source_name"] == "安心医疗险条款.pdf"
    assert result["data"]["items"][0]["document_version"] == 1
    assert result["data"]["answer_constraints"]["direct_evidence_only"] is True
    assert result["data"]["answer_constraints"]["no_scenario_expansion"] is True
    assert result["data"]["answer_constraints"]["no_numeric_inference"] is True
    assert "30天" in result["summary"]
    assert result["blocks"][0]["type"] == "knowledge_sources"
    assert result["blocks"][0]["items"][0]["source_name"] == "安心医疗险条款.pdf"
    assert [execution.name for execution in context.executions] == [
        "search_insurance_knowledge"
    ]


def test_agentic_knowledge_tool_executes_rewrites_with_metadata_filters() -> None:
    catalog = ProductCatalogService(InMemoryCatalogRepository(), lambda: NOW)
    enrollment = EnrollmentService(
        InMemoryEnrollmentRepository(),
        catalog,
        lambda: NOW,
        "x" * 32,
    )
    knowledge = KnowledgeService(
        InMemoryKnowledgeRepository(),
        lambda: NOW,
        InMemoryKnowledgeIndex(),
    )
    for title, product_code, content in (
        (
            "安心医疗险保险条款",
            "MED-BASIC",
            "疾病医疗等待期为30天，意外伤害不受等待期限制。",
        ),
        (
            "家庭意外险常见问题",
            "ACC-FAMILY",
            "意外医疗不设置疾病等待期，具体责任以意外险条款为准。",
        ),
    ):
        document = knowledge.create_document(
            title=title,
            document_type="INSURANCE_TERMS",
            source_name=f"{title}.md",
            product_code=product_code,
            content=content,
            created_by="admin-01",
        )
        knowledge.index_document(document["document_id"])
        knowledge.test_search(document["document_id"], content[:20])
        knowledge.publish_document(document["document_id"])
    registry = build_agent_tool_registry(catalog, enrollment, knowledge=knowledge)
    context = ToolExecutionContext(
        persona="customer",
        h5_user=H5User(
            id="01H5USER000000000000000001",
            mobile="18800000001",
            mobile_masked="188****0001",
            display_name="体验用户",
            status="ACTIVE",
            created_at=NOW,
        ),
    )
    tool = registry.build(
        names=("search_insurance_knowledge",),
        context=context,
        policy_engine=HarnessPolicyEngine(),
    )[0]

    result = json.loads(
        tool.invoke(
            {
                "query": "安心医疗险多久后可以报销疾病医疗？",
                "alternate_queries": [
                    "安心医疗险疾病等待期",
                    "MED-BASIC等待期天数",
                ],
                "product_code": "MED-BASIC",
                "document_types": ["INSURANCE_TERMS"],
                "limit": 3,
            }
        )
    )

    plan = result["data"]["agentic_retrieval"]
    assert plan["executed_queries"] == [
        "安心医疗险多久后可以报销疾病医疗？",
        "安心医疗险疾病等待期",
        "MED-BASIC等待期天数",
    ]
    assert plan["filters"] == {
        "product_code": "MED-BASIC",
        "document_types": ["INSURANCE_TERMS"],
    }
    assert plan["attempt"] == 1
    assert result["data"]["items"]
    assert {
        item["product_code"] for item in result["data"]["items"]
    } == {"MED-BASIC"}


def test_agentic_knowledge_tool_limits_iterative_retrieval_to_three_rounds() -> None:
    catalog = ProductCatalogService(InMemoryCatalogRepository(), lambda: NOW)
    enrollment = EnrollmentService(
        InMemoryEnrollmentRepository(),
        catalog,
        lambda: NOW,
        "x" * 32,
    )
    knowledge = KnowledgeService(
        InMemoryKnowledgeRepository(),
        lambda: NOW,
        InMemoryKnowledgeIndex(),
    )
    context = ToolExecutionContext(
        persona="customer",
        h5_user=H5User(
            id="01H5USER000000000000000001",
            mobile="18800000001",
            mobile_masked="188****0001",
            display_name="体验用户",
            status="ACTIVE",
            created_at=NOW,
        ),
    )
    tool = build_agent_tool_registry(
        catalog,
        enrollment,
        knowledge=knowledge,
    ).build(
        names=("search_insurance_knowledge",),
        context=context,
        policy_engine=HarnessPolicyEngine(),
    )[0]

    for index in range(3):
        result = json.loads(tool.invoke({"query": f"第{index + 1}轮检索问题"}))
        assert result["data"]["agentic_retrieval"]["attempt"] == index + 1

    with pytest.raises(AppError) as error:
        tool.invoke({"query": "第四轮检索问题"})

    assert error.value.code == "AGENTIC_RETRIEVAL_LIMIT_REACHED"