DataOps-platform/app/core/knowledge/query_audit.py

from __future__ import annotations

import hashlib
import json
from collections import Counter
from dataclasses import dataclass
from typing import Any
from uuid import uuid4

from sqlalchemy import text

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


@dataclass(frozen=True)
class KnowledgeQueryAuditRecord:
    uid: str
    query_hash: str
    user_id: str | None
    roles: tuple[str, ...]
    business_domain_uids: tuple[str, ...]
    mode: str
    retriever_counts: dict[str, int]
    cited_points: tuple[dict[str, Any], ...]
    degraded_components: tuple[str, ...]
    correlation_id: str
    latency_ms: int


def normalize_audited_query(query: str) -> str:
    return " ".join(str(query).split())


def build_query_audit(
    *,
    query: str,
    context: KnowledgeAccessContext,
    mode: str,
    evidence: tuple[KnowledgeEvidence, ...] | list[KnowledgeEvidence],
    cited_points: tuple[tuple[str, int], ...] = (),
    degraded_components: tuple[str, ...] = (),
    latency_ms: int,
) -> KnowledgeQueryAuditRecord:
    normalized = normalize_audited_query(query)
    retriever_counts: Counter[str] = Counter()
    for item in evidence:
        retriever_counts.update(set(item.retriever.split("+")))
    minimized_points = tuple(
        {
            "point_key": point_key,
            "point_revision": int(point_revision),
        }
        for point_key, point_revision in dict.fromkeys(cited_points)
    )
    return KnowledgeQueryAuditRecord(
        uid=str(uuid4()),
        query_hash=hashlib.sha256(normalized.encode("utf-8")).hexdigest(),
        user_id=context.subject_id or None,
        roles=tuple(sorted(context.roles)),
        business_domain_uids=tuple(sorted(context.business_domain_uids)),
        mode=str(mode),
        retriever_counts=dict(sorted(retriever_counts.items())),
        cited_points=minimized_points,
        degraded_components=tuple(sorted(set(degraded_components))),
        correlation_id=context.correlation_id,
        latency_ms=max(0, int(latency_ms)),
    )


class SqlKnowledgeQueryAuditRepository:
    def __init__(self, session):
        self._session = session

    def record(self, record: KnowledgeQueryAuditRecord) -> None:
        try:
            self._session.execute(
                text(
                    """
                    INSERT INTO public.knowledge_query_audits (
                        id, query_hash, user_id, roles,
                        business_domain_uids, mode, retriever_counts,
                        cited_points, degraded_components, correlation_id,
                        latency_ms
                    ) VALUES (
                        CAST(:id AS uuid), :query_hash,
                        CAST(:user_id AS uuid), CAST(:roles AS jsonb),
                        CAST(:business_domain_uids AS jsonb), :mode,
                        CAST(:retriever_counts AS jsonb),
                        CAST(:cited_points AS jsonb),
                        CAST(:degraded_components AS jsonb),
                        CAST(:correlation_id AS uuid), :latency_ms
                    )
                    """
                ),
                {
                    "id": record.uid,
                    "query_hash": record.query_hash,
                    "user_id": record.user_id,
                    "roles": json.dumps(record.roles),
                    "business_domain_uids": json.dumps(
                        record.business_domain_uids
                    ),
                    "mode": record.mode,
                    "retriever_counts": json.dumps(record.retriever_counts),
                    "cited_points": json.dumps(record.cited_points),
                    "degraded_components": json.dumps(
                        record.degraded_components
                    ),
                    "correlation_id": record.correlation_id,
                    "latency_ms": record.latency_ms,
                },
            )
            self._session.commit()
        except Exception:
            self._session.rollback()
            raise

    def list(self, *, limit: int = 100) -> tuple[dict[str, Any], ...]:
        rows = (
            self._session.execute(
                text(
                    """
                    SELECT id::text, query_hash, user_id::text,
                           roles, business_domain_uids, mode,
                           retriever_counts, cited_points,
                           degraded_components, correlation_id::text,
                           latency_ms, created_at
                    FROM public.knowledge_query_audits
                    ORDER BY created_at DESC, id DESC
                    LIMIT :limit
                    """
                ),
                {"limit": max(1, min(int(limit), 200))},
            )
            .mappings()
            .all()
        )
        return tuple(
            {
                **dict(row),
                "created_at": (
                    row["created_at"].isoformat()
                    if row["created_at"] is not None
                    else None
                ),
            }
            for row in rows
        )