DataOps-platform/tests/knowledge/test_retrieval.py

from __future__ import annotations


def _evidence(key: str, score: float = 0.8, domain: str = "domain-a"):
    from app.core.knowledge.retrieval.contracts import KnowledgeEvidence

    return KnowledgeEvidence(
        chunk_id=f"chunk-{key}",
        content=f"content-{key}",
        score=score,
        retriever="test",
        object_uid=f"object-{key}",
        object_type="DataFlow",
        object_version=2,
        business_domain_uid=domain,
        point_keys=(f"DataFlow/object-{key}/purpose",),
        point_revisions=(3,),
        generation=1,
        source_updated_at="2026-07-23T00:00:00+00:00",
    )


def test_router_is_deterministic_for_exact_relationship_and_semantic_queries():
    from app.core.knowledge.retrieval.router import route_query

    assert route_query('字段 "customer_id"') == "exact"
    assert route_query("customer_sync 的上游是什么") == "relationship"
    assert route_query("客户同步的用途是什么") == "semantic"


def test_rrf_fuses_duplicate_stable_chunks_and_keeps_provenance():
    from app.core.knowledge.retrieval.fusion import reciprocal_rank_fusion

    lexical = [_evidence("a", 0.8), _evidence("b", 0.7)]
    vector = [_evidence("b", 0.9), _evidence("c", 0.6)]
    fused = reciprocal_rank_fusion({"lexical": lexical, "vector": vector}, limit=3)

    assert [item.chunk_id for item in fused] == ["chunk-b", "chunk-a", "chunk-c"]
    assert fused[0].retriever == "lexical+vector"


def test_pipeline_reauthorizes_all_candidates_and_degrades_failed_retriever():
    from app.core.knowledge.access import KnowledgeAccessContext
    from app.core.knowledge.retrieval.pipeline import KnowledgeRetrievalPipeline

    class Retriever:
        def __init__(self, values=None, error=None):
            self.values = values or []
            self.error = error

        def retrieve(self, _query, _context, _limit):
            if self.error:
                raise self.error
            return self.values

    context = KnowledgeAccessContext(
        subject_id="user-1",
        roles=frozenset({"viewer"}),
        permissions=frozenset({"governance:read"}),
        business_domain_uids=frozenset({"domain-a"}),
        correlation_id="correlation-1",
    )
    pipeline = KnowledgeRetrievalPipeline(
        lexical=Retriever([_evidence("allowed"), _evidence("denied", 0.9, "domain-b")]),
        vector=Retriever(error=RuntimeError("pgvector unavailable")),
    )

    result = pipeline.search("用途", context=context, mode="semantic")

    assert [item.chunk_id for item in result.evidence] == ["chunk-allowed"]
    assert result.degraded_components == ("vector",)


def test_pipeline_fuses_device_evidence_and_reauthorizes_it_after_retrieval():
    from app.core.knowledge.access import KnowledgeAccessContext
    from app.core.knowledge.retrieval.pipeline import KnowledgeRetrievalPipeline

    class Retriever:
        def __init__(self, values=()):
            self.values = values

        def retrieve(self, _query, _context, _limit):
            return self.values

    context = KnowledgeAccessContext(
        subject_id="user-1",
        roles=frozenset({"viewer"}),
        permissions=frozenset({"governance:read"}),
        business_domain_uids=frozenset({"domain-a"}),
        correlation_id="correlation-1",
    )
    allowed = _evidence("device-a", 0.95, "domain-a")
    denied = _evidence("device-b", 0.99, "domain-b")
    pipeline = KnowledgeRetrievalPipeline(
        lexical=Retriever(),
        vector=Retriever(),
        device=Retriever((denied, allowed)),
    )

    result = pipeline.search("循环泵", context=context, mode="exact")

    assert [item.chunk_id for item in result.evidence] == [
        "chunk-device-a"
    ]