homelab-codex-ws/jobs/documents-ingest/eval/retrieval_eval.py

"""Retrieval quality gate -- module 5, phase 3, plan §6.2 (docs/kb/modules/05-faza3-plan.md).

Read-only integration script (NOT collected by pytest -- it hits the live kb-postgres DB and
the live Ollama instance, exactly like the plan asked for a separate eval script rather than a
mocked test). Runs every query in `queries.yaml` through both the flat baseline
(`documents_ingest.retrieval.flat_query`) and the cascade (`cascade_query`), for a sweep of N
values, and checks the plan's three gate criteria:

  1. every query the flat path hits (top-1 dist < 0.45) must still be a hit in the cascade
     (no degradation into the grey zone or a miss);
  2. hit@3 (expected envelope among the top-3 *distinct* envelopes by best distance) for the
     cascade must be >= hit@3 for the flat baseline, across the whole set;
  3. both negative controls must stay above 0.55 in both tracks.

`envelope`, `document_chunk`, and `document_summary` are only ever `SELECT`ed -- this script
writes nothing. Query embeddings go through Ollama on localhost/SOLARIA (bge-m3), same as
`chunk_embed.py`/`summarize.py`.

Usage:
    python retrieval_eval.py --dsn postgresql://kb:<pw>@piha:5433/kb \\
        --ollama-url http://solaria:11434 --n-sweep 5,10,20
"""
from __future__ import annotations

import argparse
import asyncio
import json
import os
import sys
from pathlib import Path
from typing import Optional

import aiohttp
import asyncpg
import yaml

sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))

from documents_ingest.retrieval import (  # noqa: E402
    DEFAULT_EMBED_MODEL,
    DEFAULT_K,
    DEFAULT_N,
    DEFAULT_SUMMARY_MODEL,
    cascade_query,
    flat_query,
)

HIT_THRESHOLD = 0.45
NO_ANSWER_THRESHOLD = 0.55
DEFAULT_QUERIES_PATH = Path(__file__).resolve().parent / "queries.yaml"


def load_queries(path: Path) -> list[dict]:
    with open(path, encoding="utf-8") as f:
        data = yaml.safe_load(f)
    return data["queries"]


def top1_dist(chunks: list[dict]) -> Optional[float]:
    return min((c["dist"] for c in chunks), default=None)


def hit_at_3(chunks: list[dict], expected_envelope: Optional[str]) -> Optional[bool]:
    """Is `expected_envelope` among the top-3 *distinct* envelopes, ordered by each envelope's
    best (lowest) distance? Returns None for negative controls (no expected envelope -- hit@3
    isn't a meaningful concept for them, they're graded on distance alone)."""
    if expected_envelope is None:
        return None
    best_per_envelope: dict[str, float] = {}
    for c in chunks:
        prev = best_per_envelope.get(c["envelope_id"])
        if prev is None or c["dist"] < prev:
            best_per_envelope[c["envelope_id"]] = c["dist"]
    top3 = sorted(best_per_envelope.items(), key=lambda kv: kv[1])[:3]
    return expected_envelope in {eid for eid, _ in top3}


async def run_query_both_tracks(
    conn: asyncpg.Connection,
    session: aiohttp.ClientSession,
    ollama_url: str,
    query: dict,
    summary_model: str,
    embed_model: str,
    n_values: list[int],
    k: int,
) -> dict:
    flat = await flat_query(conn, session, ollama_url, query["text"], embed_model=embed_model, k=k)
    cascades = {}
    for n in n_values:
        cascades[n] = await cascade_query(
            conn, session, ollama_url, query["text"],
            summary_model=summary_model, embed_model=embed_model, n=n, k=k,
        )
    return {"query": query, "flat": flat, "cascades": cascades}


def summarize_query_result(result: dict) -> dict:
    query = result["query"]
    expected = query["expected_envelope"]
    flat_chunks = result["flat"]["chunks"]

    row = {
        "id": query["id"],
        "kind": query["kind"],
        "text": query["text"],
        "expected_envelope": expected,
        "flat_top1_dist": top1_dist(flat_chunks),
        "flat_hit3": hit_at_3(flat_chunks, expected),
        "cascade": {},
    }
    for n, cascade_result in result["cascades"].items():
        chunks = cascade_result["chunks"]
        row["cascade"][n] = {
            "top1_dist": top1_dist(chunks),
            "hit3": hit_at_3(chunks, expected),
            "stage1_count": len(cascade_result["stage1_summaries"]),
            "chunk_count": len(chunks),
        }
    return row


def evaluate_gate(rows: list[dict], gate_n: int) -> dict:
    """Plan §6.2's three criteria, evaluated at one chosen N (the sweep is diagnostic, the
    gate verdict is always for one specific configuration)."""
    reasons = []

    # 1. no flat hit (dist < 0.45) may degrade under the cascade.
    degraded = []
    for row in rows:
        flat_dist = row["flat_top1_dist"]
        if flat_dist is not None and flat_dist < HIT_THRESHOLD:
            cascade_dist = row["cascade"][gate_n]["top1_dist"]
            if cascade_dist is None or cascade_dist >= HIT_THRESHOLD:
                degraded.append(row["id"])
    criterion_1 = not degraded
    if not criterion_1:
        reasons.append(f"criterion 1 FAILED: flat hits degraded under cascade N={gate_n}: {degraded}")

    # 2. hit@3 (cascade) >= hit@3 (flat), over queries that have an expected envelope.
    scored = [row for row in rows if row["expected_envelope"] is not None]
    flat_hit3_count = sum(1 for row in scored if row["flat_hit3"])
    cascade_hit3_count = sum(1 for row in scored if row["cascade"][gate_n]["hit3"])
    criterion_2 = cascade_hit3_count >= flat_hit3_count
    if not criterion_2:
        reasons.append(
            f"criterion 2 FAILED: hit@3 cascade={cascade_hit3_count}/{len(scored)} "
            f"< flat={flat_hit3_count}/{len(scored)} (N={gate_n})"
        )

    # 3. negative controls stay above 0.55 in both tracks.
    control_failures = []
    for row in rows:
        if row["kind"] not in ("negative_control", "negative_control_borderline"):
            continue
        if row["flat_top1_dist"] is not None and row["flat_top1_dist"] <= NO_ANSWER_THRESHOLD:
            control_failures.append(f"{row['id']} flat={row['flat_top1_dist']:.4f}")
        cascade_dist = row["cascade"][gate_n]["top1_dist"]
        if cascade_dist is not None and cascade_dist <= NO_ANSWER_THRESHOLD:
            control_failures.append(f"{row['id']} cascade(N={gate_n})={cascade_dist:.4f}")
    criterion_3 = not control_failures
    if not criterion_3:
        reasons.append(f"criterion 3 FAILED: negative control(s) crossed {NO_ANSWER_THRESHOLD}: {control_failures}")

    passed = criterion_1 and criterion_2 and criterion_3
    return {
        "gate_n": gate_n,
        "passed": passed,
        "criterion_1_no_degradation": criterion_1,
        "criterion_2_hit3": {"cascade": cascade_hit3_count, "flat": flat_hit3_count, "total": len(scored)},
        "criterion_3_negative_controls": criterion_3,
        "reasons": reasons,
    }


def print_report(rows: list[dict], n_values: list[int], gate_result: dict) -> None:
    print("=" * 100)
    print("RETRIEVAL QUALITY GATE -- plan §6.2 (docs/kb/modules/05-faza3-plan.md)")
    print("=" * 100)
    header = f"{'id':<3} {'kind':<28} {'expected':<16} {'flat d1':>8} {'flat@3':>7}"
    for n in n_values:
        header += f"  N={n:<3} d1  hit3"
    print(header)
    for row in rows:
        flat_d1 = f"{row['flat_top1_dist']:.4f}" if row["flat_top1_dist"] is not None else "--"
        flat_h3 = "y" if row["flat_hit3"] else ("--" if row["flat_hit3"] is None else "n")
        line = f"{row['id']:<3} {row['kind']:<28} {str(row['expected_envelope']):<16} {flat_d1:>8} {flat_h3:>7}"
        for n in n_values:
            c = row["cascade"][n]
            d1 = f"{c['top1_dist']:.4f}" if c["top1_dist"] is not None else "--"
            h3 = "y" if c["hit3"] else ("--" if c["hit3"] is None else "n")
            line += f"  {d1:>8} {h3:>4}"
        print(line)
    print("-" * 100)

    print(f"\nGate verdict at N={gate_result['gate_n']}, k={DEFAULT_K}:")
    print(f"  criterion 1 (no flat hit degrades)      : {'PASS' if gate_result['criterion_1_no_degradation'] else 'FAIL'}")
    h3 = gate_result["criterion_2_hit3"]
    print(f"  criterion 2 (hit@3 cascade >= flat)      : cascade={h3['cascade']}/{h3['total']} flat={h3['flat']}/{h3['total']} "
          f"-> {'PASS' if gate_result['criterion_1_no_degradation'] and h3['cascade'] >= h3['flat'] else 'FAIL'}")
    print(f"  criterion 3 (negative controls > {NO_ANSWER_THRESHOLD})  : {'PASS' if gate_result['criterion_3_negative_controls'] else 'FAIL'}")
    print(f"\n  OVERALL: {'PASS' if gate_result['passed'] else 'FAIL'}")
    for reason in gate_result["reasons"]:
        print(f"    - {reason}")

    print("\nCost per query: flat = 1 embed + 1 SQL query. "
          f"cascade = 1 embed (shared) + 2 SQL queries (stage1 top-N summaries, stage2 top-k chunks). "
          "Cascade never costs an extra Ollama call, only one extra SQL round-trip.")


async def main_async(args: argparse.Namespace) -> dict:
    queries = load_queries(Path(args.queries))
    n_values = [int(n) for n in args.n_sweep.split(",")]
    if args.gate_n not in n_values:
        n_values = sorted(set(n_values) | {args.gate_n})

    conn = await asyncpg.connect(args.dsn)
    try:
        async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=60)) as session:
            results = []
            for query in queries:
                result = await run_query_both_tracks(
                    conn, session, args.ollama_url, query,
                    summary_model=args.summary_model, embed_model=args.embed_model,
                    n_values=n_values, k=args.k,
                )
                results.append(summarize_query_result(result))
    finally:
        await conn.close()

    gate_result = evaluate_gate(results, gate_n=args.gate_n)
    print_report(results, n_values, gate_result)

    output = {"rows": results, "gate": gate_result, "n_sweep": n_values, "k": args.k}
    if args.json_out:
        Path(args.json_out).write_text(json.dumps(output, indent=2, default=str), encoding="utf-8")
        print(f"\nFull JSON written to {args.json_out}")
    return output


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--dsn", default=os.environ.get("KB_DSN"), help="asyncpg DSN for kb-postgres (or KB_DSN env var)")
    parser.add_argument("--ollama-url", default=os.environ.get("OLLAMA_URL", "http://localhost:11434"))
    parser.add_argument("--embed-model", default=DEFAULT_EMBED_MODEL)
    parser.add_argument("--summary-model", default=DEFAULT_SUMMARY_MODEL,
                         help="document_summary.model to pre-filter on (plan §2 D3 resolution)")
    parser.add_argument("--k", type=int, default=DEFAULT_K)
    parser.add_argument("--gate-n", type=int, default=DEFAULT_N, help="N used for the PASS/FAIL verdict")
    parser.add_argument("--n-sweep", default="5,10,20", help="comma-separated N values to report (diagnostic)")
    parser.add_argument("--queries", default=str(DEFAULT_QUERIES_PATH))
    parser.add_argument("--json-out", default=None, help="optional path to dump full results as JSON")
    args = parser.parse_args()

    if not args.dsn:
        print("error: pass --dsn or set KB_DSN", file=sys.stderr)
        sys.exit(1)

    output = asyncio.run(main_async(args))
    sys.exit(0 if output["gate"]["passed"] else 1)


if __name__ == "__main__":
    main()