RAG/backend/app/tools/evaluate_demo.py

"""Run a reproducible retrieval evaluation on the public synthetic corpus."""

from __future__ import annotations

import asyncio
import hashlib
import json
import sys
import uuid
from dataclasses import asdict
from pathlib import Path
from typing import Any, Protocol

from app.adapters.fake import FakeEmbeddingProvider, FakeReranker
from app.core.config import Settings
from app.core.demo_identity import ACCESS_SCOPE_ID, KNOWLEDGE_BASE_ID
from app.persistence.retrieval import PostgresRetrievalRepository
from app.services.evaluation import (
    RankingMetrics,
    bootstrap_mean_confidence_interval,
    evaluate_ranking,
    freeze_run_config,
)
from app.services.retrieval import RetrievalActor, RetrievalGrant, RetrievalResult, RetrievalService
from app.tools.seed_demo import (
    DEFAULT_SAMPLE_ROOT,
    DemoDocument,
    DemoQuery,
    load_documents,
    load_queries,
)


def _sha256_file(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def _source_id(source_name: str) -> str:
    return Path(source_name).stem


def _mean(values: list[float]) -> float:
    return sum(values) / len(values) if values else 0.0


class DemoRetrievalService(Protocol):
    async def search(
        self,
        *,
        actor: RetrievalActor,
        knowledge_base_id: uuid.UUID,
        query: str,
        vector_top_k: int,
        rerank_top_n: int,
    ) -> RetrievalResult: ...


async def evaluate_demo_queries(
    *,
    service: DemoRetrievalService,
    actor: RetrievalActor,
    documents: list[DemoDocument],
    queries: list[DemoQuery],
    vector_top_k: int = 20,
    rerank_top_n: int = 10,
    metric_cutoff: int = 3,
) -> dict[str, Any]:
    """Evaluate answerable queries with a fully judged synthetic corpus pool."""

    corpus_ids = frozenset(document.source_id for document in documents)
    if len(corpus_ids) != len(documents):
        raise ValueError("synthetic corpus document IDs must be unique")

    cases: list[dict[str, Any]] = []
    scored: list[RankingMetrics] = []
    active_profile_hash: str | None = None
    for query in queries:
        result = await service.search(
            actor=actor,
            knowledge_base_id=KNOWLEDGE_BASE_ID,
            query=query.query,
            vector_top_k=vector_top_k,
            rerank_top_n=rerank_top_n,
        )
        if active_profile_hash is None:
            active_profile_hash = result.profile.profile_hash
        elif active_profile_hash != result.profile.profile_hash:
            raise ValueError("active profile changed during evaluation")

        ranked_ids = [_source_id(hit.source_name) for hit in result.results]
        case: dict[str, Any] = {
            "qid": query.qid,
            "answerable": query.answerable,
            "ranked_document_ids": ranked_ids,
            "retrieval_status": result.status,
            "rerank_status": result.rerank_status,
        }
        if query.answerable:
            relevance = {document_id: 0.0 for document_id in corpus_ids}
            for expected_id in query.expected_doc_ids:
                if expected_id not in corpus_ids:
                    raise ValueError("expected document is outside the corpus manifest")
                relevance[expected_id] = 1.0
            groups = tuple(frozenset({expected_id}) for expected_id in query.expected_doc_ids)
            metrics = evaluate_ranking(
                ranked_ids,
                relevance=relevance,
                judged_document_ids=corpus_ids,
                evidence_groups=groups,
                k=metric_cutoff,
            )
            scored.append(metrics)
            case["metrics"] = asdict(metrics)
        else:
            case["metrics"] = None
        cases.append(case)

    if active_profile_hash is None:
        raise ValueError("evaluation query set is empty")
    hit_values = [metric.hit_at_k for metric in scored]
    hit_ci = bootstrap_mean_confidence_interval(
        hit_values,
        seed=20260713,
        iterations=2_000,
    )
    return {
        "status": "ok",
        "dataset": "synthetic-demo",
        "case_count": len(queries),
        "answerable_case_count": len(scored),
        "active_embedding_profile_hash": active_profile_hash,
        "metrics": {
            f"hit_at_{metric_cutoff}": _mean(hit_values),
            "mrr": _mean([metric.reciprocal_rank for metric in scored]),
            f"ndcg_at_{metric_cutoff}": _mean([metric.ndcg_at_k for metric in scored]),
            f"complete_hit_at_{metric_cutoff}": _mean(
                [metric.complete_hit_at_k for metric in scored]
            ),
            f"evidence_group_recall_at_{metric_cutoff}": _mean(
                [metric.evidence_group_recall_at_k for metric in scored]
            ),
            f"hit_at_{metric_cutoff}_confidence_interval": asdict(hit_ci),
        },
        "cases": cases,
    }


async def async_main() -> int:
    document_path = Path(
        sys.argv[1] if len(sys.argv) > 1 else DEFAULT_SAMPLE_ROOT / "demo_documents.jsonl"
    )
    query_path = Path(
        sys.argv[2] if len(sys.argv) > 2 else DEFAULT_SAMPLE_ROOT / "demo_queries.jsonl"
    )
    try:
        settings = Settings()
        documents = load_documents(document_path)
        queries = load_queries(query_path)
        service = RetrievalService(
            repository=PostgresRetrievalRepository(settings),
            embedding_provider=FakeEmbeddingProvider(settings.embedding_dimension),
            reranker=FakeReranker(),
            synthetic_embedding_provider=FakeEmbeddingProvider(settings.embedding_dimension),
            synthetic_reranker=FakeReranker(),
        )
        actor = RetrievalActor(
            subject="synthetic-evaluation-runner",
            grants=(
                RetrievalGrant(
                    knowledge_base_id=KNOWLEDGE_BASE_ID,
                    access_scope_ids=(ACCESS_SCOPE_ID,),
                ),
            ),
        )
        artifact = await evaluate_demo_queries(
            service=service,
            actor=actor,
            documents=documents,
            queries=queries,
        )
        config, config_hash = freeze_run_config(
            {
                "corpus_sha256": _sha256_file(document_path),
                "query_set_sha256": _sha256_file(query_path),
                "embedding_profile_hash": artifact["active_embedding_profile_hash"],
                "vector_top_k": 20,
                "rerank_top_n": 10,
                "metric_cutoff": 3,
                "bootstrap_seed": 20260713,
            }
        )
        artifact["frozen_config"] = json.loads(config)
        artifact["frozen_config_sha256"] = config_hash
        sys.stdout.write(json.dumps(artifact, ensure_ascii=False, sort_keys=True) + "\n")
        return 0
    except Exception:
        # CLI output remains fixed and does not echo paths, document text, DSNs, or errors.
        sys.stdout.write(
            json.dumps(
                {"status": "failed", "error_kind": "evaluation_failed"},
                sort_keys=True,
            )
            + "\n"
        )
        return 1


def main() -> None:
    raise SystemExit(asyncio.run(async_main()))


if __name__ == "__main__":
    main()