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

from __future__ import annotations

from collections.abc import Mapping, Sequence
from typing import Any, Protocol

from sqlalchemy import text

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


class QueryEmbeddingProvider(Protocol):
    def embed(self, texts: list[str]) -> list[list[float]]: ...


def _vector(value: Sequence[float]) -> str:
    return "[" + ",".join(str(float(item)) for item in value) + "]"


def _as_tuple(value: Any) -> tuple:
    if value is None:
        return ()
    if isinstance(value, (list, tuple)):
        return tuple(value)
    return tuple(value)


def _evidence(row: Mapping[str, Any], retriever: str) -> KnowledgeEvidence:
    updated_at = row["source_updated_at"]
    return KnowledgeEvidence(
        chunk_id=str(row["chunk_id"]),
        content=row["content"],
        score=float(row["score"]),
        retriever=retriever,
        object_uid=str(row["object_uid"]),
        object_type=row["object_type"],
        object_version=int(row["object_version"]),
        business_domain_uid=(
            str(row["business_domain_uid"])
            if row["business_domain_uid"] is not None
            else None
        ),
        point_keys=tuple(str(value) for value in _as_tuple(row["point_keys"])),
        point_revisions=tuple(
            int(value) for value in _as_tuple(row["point_revisions"])
        ),
        generation=int(row["generation"]),
        source_updated_at=updated_at.isoformat() if updated_at else None,
        freshness_status=row["freshness_status"],
        section_path=row["section_path"],
    )


_SELECT = """
    SELECT c.id AS chunk_id, c.content, c.section_path,
           d.object_uid, d.object_type, d.object_version,
           d.business_domain_uid, d.active_generation AS generation,
           d.source_updated_at, c.point_keys,
           COALESCE(points.point_revisions, ARRAY[]::bigint[]) AS point_revisions,
           CASE WHEN pending.source_revision > d.object_version
                THEN 'updating' ELSE 'fresh' END AS freshness_status,
           {score} AS score
    FROM public.governance_chunks c
    JOIN public.governance_documents d ON d.id = c.document_id
    LEFT JOIN LATERAL (
        SELECT array_agg(p.point_revision ORDER BY keys.ordinality) AS point_revisions
        FROM jsonb_array_elements_text(c.point_keys) WITH ORDINALITY AS keys(point_key, ordinality)
        JOIN public.knowledge_points p
          ON p.point_key = keys.point_key AND p.status = 'active'
    ) points ON TRUE
    LEFT JOIN LATERAL (
        SELECT MAX(cs.source_revision) AS source_revision
        FROM public.knowledge_change_sets cs
        WHERE cs.source_uid = d.object_uid
          AND cs.status NOT IN ('canonical_active','complete','failed','rolled_back')
    ) pending ON TRUE
    {embedding_join}
    WHERE d.status = 'active'
      AND d.active_generation IS NOT NULL
      AND (:global_access OR d.business_domain_uid = ANY(CAST(:domains AS uuid[])))
      {predicate}
"""


class SqlLexicalRetriever:
    def __init__(self, session) -> None:
        self._session = session

    def retrieve(
        self, query: str, context: KnowledgeAccessContext, limit: int
    ) -> tuple[KnowledgeEvidence, ...]:
        score = """
            GREATEST(
                CASE WHEN lower(c.lexical_text) = lower(:query) THEN 1.0 ELSE 0.0 END,
                CASE WHEN lower(c.lexical_text) LIKE lower(:query) || '%' THEN 0.9 ELSE 0.0 END,
                similarity(COALESCE(c.lexical_text, ''), :query)
            )
        """
        statement = text(
            _SELECT.format(
                score=score,
                embedding_join="",
                predicate="""
                    AND (
                        lower(c.lexical_text) = lower(:query)
                        OR lower(c.lexical_text) LIKE lower(:query) || '%'
                        OR c.lexical_text % :query
                        OR c.search_vector @@ plainto_tsquery('simple', :query)
                    )
                    ORDER BY score DESC, c.id
                    LIMIT :limit
                """,
            )
        )
        rows = self._session.execute(
            statement,
            {
                "query": query.strip(),
                "global_access": context.global_access,
                "domains": list(context.business_domain_uids),
                "limit": limit,
            },
        ).mappings()
        return tuple(_evidence(row, "lexical") for row in rows)


class SqlVectorRetriever:
    def __init__(self, session, embedder: QueryEmbeddingProvider) -> None:
        self._session = session
        self._embedder = embedder

    def retrieve(
        self, query: str, context: KnowledgeAccessContext, limit: int
    ) -> tuple[KnowledgeEvidence, ...]:
        vectors = self._embedder.embed([query])
        if len(vectors) != 1:
            raise RuntimeError("query embedder returned an unexpected vector count")
        statement = text(
            _SELECT.format(
                score="1 - (e.embedding <=> CAST(:query_embedding AS vector))",
                embedding_join="""
                    JOIN public.knowledge_chunk_embeddings e ON e.chunk_id = c.id
                    JOIN public.knowledge_embedding_profiles profile
                      ON profile.id = e.profile_id AND profile.status = 'active'
                """,
                predicate="""
                    ORDER BY e.embedding <=> CAST(:query_embedding AS vector), c.id
                    LIMIT :limit
                """,
            )
        )
        rows = self._session.execute(
            statement,
            {
                "query_embedding": _vector(vectors[0]),
                "global_access": context.global_access,
                "domains": list(context.business_domain_uids),
                "limit": limit,
            },
        ).mappings()
        return tuple(_evidence(row, "vector") for row in rows)


class UnavailableVectorRetriever:
    def __init__(self, reason: str = "query embedding is not configured") -> None:
        self._reason = reason

    def retrieve(self, _query, _context, _limit):
        raise RuntimeError(self._reason)