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),
extended in faza mailowa Krok 5 (docs/kb/modules/05-faza-mailowa-plan.md, §8) to add the
`hybrid` track once mail content exists in `document_chunk` (faza mailowa Krok 2/4).

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`'s `queries:` list through flat
(`kb_retrieval.retrieval.flat_query`), cascade (`cascade_query`, swept over N), and hybrid
(`hybrid_query`), and checks the gate criteria:

  1. every query the flat path hits (top-1 dist < 0.45) must still be a hit in cascade AND in
     hybrid (no degradation into the grey zone or a miss) -- faza mailowa's key risk: does
     adding mail chunks to the HNSW index degrade existing paperless retrieval?
  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 `queries:` set;
  3. both negative controls must stay above 0.55 in flat, cascade, AND hybrid;
  4. (faza mailowa, only when `mail_queries` in queries.yaml is non-empty) hit@3 in hybrid for
     the operator-supplied mail queries must be >= 4/5 (or a proportional threshold for fewer
     queries). An empty `mail_queries` list skips this criterion with an explicit note in the
     report, rather than silently passing or failing on absent data.

`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 math
import os
import sys
from pathlib import Path
from typing import Optional

import aiohttp
import asyncpg
import yaml

_REPO_ROOT = Path(__file__).resolve().parent.parent.parent.parent
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
sys.path.insert(0, str(_REPO_ROOT / "packages" / "kb-retrieval" / "src"))

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

HIT_THRESHOLD = 0.45
NO_ANSWER_THRESHOLD = 0.55
MAIL_HIT3_MIN_FRACTION = 0.8  # plan §8: ">= 4/5 (lub 3/3-4/4 przy mniejszej liczbie)"
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 load_mail_queries(path: Path) -> list[dict]:
    """Placeholder section (plan §8, Krok 5) -- empty until the operator supplies 3-5 queries."""
    with open(path, encoding="utf-8") as f:
        data = yaml.safe_load(f)
    return data.get("mail_queries") or []


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}


def top3_envelopes(chunks: list[dict]) -> list[tuple[str, float]]:
    """Top-3 *distinct* envelopes by best (lowest) distance -- shared by `hit_at_3` (identity
    match against `expected_envelope`) and `mail_hit_at_3` (source match, no expected id)."""
    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"]
    return sorted(best_per_envelope.items(), key=lambda kv: kv[1])[:3]


def mail_hit_at_3(chunks: list[dict], envelope_sources: dict[str, str]) -> bool:
    """`mail_queries` (plan §8 Krok 5) carry `expected_envelope: null` -- the operator supplied
    the query text, not a Message-ID, so `hit_at_3`'s envelope-identity match always returns
    None for them (the bug the operator flagged: it prints as '--' and can never count towards
    criterion 4). The intended semantics is "hybrid actually surfaces mail content here": hit
    iff one of the top-3 distinct hybrid envelopes is mail-sourced (`envelope.source` in
    `DEFAULT_SUMMARYLESS_SOURCES`, i.e. gmail today) with dist < HIT_THRESHOLD."""
    return any(
        envelope_sources.get(eid) in DEFAULT_SUMMARYLESS_SOURCES and dist < HIT_THRESHOLD
        for eid, dist in top3_envelopes(chunks)
    )


async def fetch_envelope_sources(conn: asyncpg.Connection, envelope_ids: list[str]) -> dict[str, str]:
    """Read-only lookup of `envelope.source` for the envelope ids seen in a result set --
    a hybrid chunk's `source` field is always overwritten to `"hybrid"` on merge
    (`hybrid_retrieve`'s docstring), so origin (paperless vs. gmail) has to come from `envelope`
    itself, not from the chunk dict."""
    if not envelope_ids:
        return {}
    rows = await conn.fetch(
        "SELECT id, source FROM envelope WHERE id = ANY($1::text[])", envelope_ids
    )
    return {r["id"]: r["source"] for r in rows}


async def run_query_all_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,
        )
    hybrid = await hybrid_query(
        conn, session, ollama_url, query["text"],
        summary_model=summary_model, embed_model=embed_model, k=k,
    )
    return {"query": query, "flat": flat, "cascades": cascades, "hybrid": hybrid}


def summarize_query_result(result: dict, envelope_sources: Optional[dict[str, str]] = None) -> dict:
    query = result["query"]
    expected = query["expected_envelope"]
    flat_chunks = result["flat"]["chunks"]
    hybrid_chunks = result["hybrid"]["chunks"]

    if query["kind"] == "mail_hit":
        # plan §8 Krok 5: no expected_envelope (operator gave query text, not a Message-ID) --
        # graded on mail_hit_at_3's source-match semantics instead of hit_at_3's identity match.
        flat_hit3 = None  # flat never reaches summaryless (gmail) chunks -- not a meaningful axis
        hybrid_hit3 = mail_hit_at_3(hybrid_chunks, envelope_sources or {})
    else:
        flat_hit3 = hit_at_3(flat_chunks, expected)
        hybrid_hit3 = hit_at_3(hybrid_chunks, expected)

    row = {
        "id": query["id"],
        "kind": query["kind"],
        "text": query["text"],
        "expected_envelope": expected,
        "flat_top1_dist": top1_dist(flat_chunks),
        "flat_hit3": flat_hit3,
        "hybrid_top1_dist": top1_dist(hybrid_chunks),
        "hybrid_hit3": hybrid_hit3,
        # per-query override (queries.yaml `no_answer_threshold`) for criterion 3 -- defaults
        # to the module-wide NO_ANSWER_THRESHOLD; see N2's note (faza mailowa false-positive
        # collision with an unrelated ski-newsletter mail, 2026-07-23).
        "no_answer_threshold": query.get("no_answer_threshold", NO_ANSWER_THRESHOLD),
        "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, mail_rows: Optional[list[dict]] = None) -> dict:
    """Plan §6.2/§8's criteria, evaluated at one chosen N (the sweep is diagnostic, the gate
    verdict is always for one specific configuration). `mail_rows` is None or empty when the
    operator hasn't supplied `mail_queries` yet (plan §8 Krok 5 placeholder) -- criterion 4 is
    then skipped with an explicit note, never silently PASSed or FAILed on absent data."""
    reasons = []

    # 1. no flat hit (dist < 0.45) may degrade under cascade OR hybrid (faza mailowa's core
    # regression risk: does the mail chunk mass added to HNSW push existing hits out?).
    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(f"{row['id']} (cascade)")
            hybrid_dist = row["hybrid_top1_dist"]
            if hybrid_dist is None or hybrid_dist >= HIT_THRESHOLD:
                degraded.append(f"{row['id']} (hybrid)")
    criterion_1 = not degraded
    if not criterion_1:
        reasons.append(f"criterion 1 FAILED: flat hits degraded: {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 their pass bar (module default NO_ANSWER_THRESHOLD,
    # unless queries.yaml gives this query its own `no_answer_threshold` -- see N2) in flat,
    # cascade, AND hybrid.
    control_failures = []
    for row in rows:
        if row["kind"] not in ("negative_control", "negative_control_borderline"):
            continue
        threshold = row["no_answer_threshold"]
        if row["flat_top1_dist"] is not None and row["flat_top1_dist"] <= threshold:
            control_failures.append(f"{row['id']} flat={row['flat_top1_dist']:.4f} (bar {threshold})")
        cascade_dist = row["cascade"][gate_n]["top1_dist"]
        if cascade_dist is not None and cascade_dist <= threshold:
            control_failures.append(f"{row['id']} cascade(N={gate_n})={cascade_dist:.4f} (bar {threshold})")
        hybrid_dist = row["hybrid_top1_dist"]
        if hybrid_dist is not None and hybrid_dist <= threshold:
            control_failures.append(f"{row['id']} hybrid={hybrid_dist:.4f} (bar {threshold})")
    criterion_3 = not control_failures
    if not criterion_3:
        reasons.append(f"criterion 3 FAILED: negative control(s) crossed their pass bar: {control_failures}")

    # 4. (faza mailowa, plan §8) mail queries' hit@3 in hybrid -- skipped entirely when the
    # operator hasn't supplied queries yet (mail_rows empty), never PASSed/FAILed on no data.
    mail_rows = mail_rows or []
    if mail_rows:
        mail_hit3_count = sum(1 for row in mail_rows if row["hybrid_hit3"])
        required = math.ceil(len(mail_rows) * MAIL_HIT3_MIN_FRACTION)
        criterion_4 = mail_hit3_count >= required
        criterion_4_detail = {
            "hit3": mail_hit3_count, "total": len(mail_rows), "required": required, "skipped": False,
        }
        if not criterion_4:
            reasons.append(
                f"criterion 4 FAILED: mail hit@3 (hybrid) = {mail_hit3_count}/{len(mail_rows)}, "
                f"required >= {required}"
            )
    else:
        criterion_4 = True  # does not block PASS -- see note below
        criterion_4_detail = {"hit3": 0, "total": 0, "required": 0, "skipped": True}
        reasons.append(
            "criterion 4 SKIPPED: mail_queries is empty in queries.yaml -- operator hasn't "
            "supplied the 3-5 mail queries yet (plan §8, Krok 5). Gate PASS below reflects "
            "only the paperless regression check, not full plan §8 completion."
        )

    passed = criterion_1 and criterion_2 and criterion_3 and criterion_4
    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,
        "criterion_4_mail_hit3": criterion_4_detail,
        "reasons": reasons,
    }


def print_report(
    rows: list[dict], n_values: list[int], gate_result: dict, mail_rows: Optional[list[dict]] = None
) -> None:
    print("=" * 100)
    print("RETRIEVAL QUALITY GATE -- plan §6.2 (docs/kb/modules/05-faza3-plan.md), "
          "+ hybrid/mail extension (05-faza-mailowa-plan.md §8)")
    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"
    header += "  hybrid d1  hyb@3"
    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}"
        hyb_d1 = f"{row['hybrid_top1_dist']:.4f}" if row["hybrid_top1_dist"] is not None else "--"
        hyb_h3 = "y" if row["hybrid_hit3"] else ("--" if row["hybrid_hit3"] is None else "n")
        line += f"  {hyb_d1:>9} {hyb_h3:>6}"
        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, cascade+hybrid): {'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} [per-query override possible], flat+cascade+hybrid): {'PASS' if gate_result['criterion_3_negative_controls'] else 'FAIL'}")
    c4 = gate_result["criterion_4_mail_hit3"]
    if c4["skipped"]:
        print("  criterion 4 (mail hit@3 in hybrid)       : SKIPPED (mail_queries empty)")
    else:
        print(f"  criterion 4 (mail hit@3 in hybrid)       : {c4['hit3']}/{c4['total']} (required >= {c4['required']}) "
              f"-> {'PASS' if gate_result['criterion_1_no_degradation'] else 'FAIL'}")
    print(f"\n  OVERALL: {'PASS' if gate_result['passed'] else 'FAIL'}")
    for reason in gate_result["reasons"]:
        print(f"    - {reason}")

    if mail_rows:
        print("\nMail queries (hybrid only, plan §8 criterion 4):")
        for row in mail_rows:
            d1 = f"{row['hybrid_top1_dist']:.4f}" if row["hybrid_top1_dist"] is not None else "--"
            h3 = "y" if row["hybrid_hit3"] else ("--" if row["hybrid_hit3"] is None else "n")
            print(f"  {row['id']:<4} {row['text']:<50} d1={d1:>8} hit3={h3}")

    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). "
          "hybrid = 1 embed (shared) + cascade's queries + 1 extra SQL query (mail branch).")


async def main_async(args: argparse.Namespace) -> dict:
    queries = load_queries(Path(args.queries))
    mail_queries = load_mail_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_all_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))

            mail_results_raw = []
            for query in mail_queries:
                result = await run_query_all_tracks(
                    conn, session, args.ollama_url, query,
                    summary_model=args.summary_model, embed_model=args.embed_model,
                    n_values=n_values, k=args.k,
                )
                mail_results_raw.append(result)

            mail_envelope_ids = {
                c["envelope_id"] for r in mail_results_raw for c in r["hybrid"]["chunks"]
            }
            envelope_sources = await fetch_envelope_sources(conn, list(mail_envelope_ids))
            mail_results = [
                summarize_query_result(r, envelope_sources) for r in mail_results_raw
            ]
    finally:
        await conn.close()

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

    output = {
        "rows": results, "mail_rows": mail_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()