DataOps-platform/app/core/knowledge/retrieval/pipeline.py

from __future__ import annotations

from collections.abc import Sequence
from typing import Protocol

from app.core.knowledge.access import KnowledgeAccessContext
from app.core.knowledge.retrieval.contracts import KnowledgeEvidence, SearchResult
from app.core.knowledge.retrieval.fusion import reciprocal_rank_fusion
from app.core.knowledge.retrieval.router import route_query


class Retriever(Protocol):
    def retrieve(
        self,
        query: str,
        context: KnowledgeAccessContext,
        limit: int,
    ) -> Sequence[KnowledgeEvidence]: ...


class KnowledgeRetrievalPipeline:
    def __init__(
        self,
        *,
        lexical: Retriever,
        vector: Retriever,
        graph: Retriever | None = None,
        lightrag: Retriever | None = None,
    ):
        self._retrievers = {
            "lexical": lexical,
            "vector": vector,
            "graph": graph,
            "lightrag": lightrag,
        }

    def search(
        self,
        query: str,
        *,
        context: KnowledgeAccessContext,
        mode: str = "auto",
        limit: int = 20,
    ) -> SearchResult:
        resolved_mode = route_query(query) if mode == "auto" else mode
        names = ["lexical", "vector"]
        if resolved_mode == "relationship" and self._retrievers["graph"] is not None:
            names.append("graph")
        if resolved_mode == "global" and self._retrievers["lightrag"] is not None:
            names.append("lightrag")
        ranked: dict[str, Sequence[KnowledgeEvidence]] = {}
        degraded: list[str] = []
        for name in names:
            retriever = self._retrievers[name]
            if retriever is None:
                continue
            try:
                ranked[name] = retriever.retrieve(query, context, limit)
            except Exception:
                degraded.append(name)
        fused = reciprocal_rank_fusion(ranked, limit=limit)
        authorized = tuple(
            evidence
            for evidence in fused
            if context.permits_domain(evidence.business_domain_uid)
            and evidence.freshness_status != "stale"
        )
        return SearchResult(
            evidence=authorized,
            mode=resolved_mode,
            degraded_components=tuple(degraded),
        )