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

from __future__ import annotations

import re

from app.core.knowledge.access import KnowledgeAccessContext
from app.core.knowledge.retrieval.contracts import KnowledgeEvidence

_SOURCE_PATTERN = re.compile(
    r"\[DATAOPS_SOURCE\s+([A-Za-z][A-Za-z0-9_]*:[^:\]\s]+:[0-9]+)\]"
)


class LightRAGRetriever:
    def __init__(self, client, *, canonical_loader) -> None:
        self._client = client
        self._canonical_loader = canonical_loader

    def retrieve(
        self, query: str, context: KnowledgeAccessContext, limit: int
    ) -> tuple[KnowledgeEvidence, ...]:
        raw_context = self._client.query_context(query, mode="mix", limit=limit)
        external_ids = tuple(dict.fromkeys(_SOURCE_PATTERN.findall(raw_context)))
        if not external_ids:
            return ()
        candidates = self._canonical_loader(external_ids, context, limit)
        return tuple(
            item
            for item in candidates
            if context.permits_domain(item.business_domain_uid)
            and item.freshness_status != "stale"
        )