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

"""Retrieval quality gate -- module 5, phase 3, plan §6.2 (kb/phases/kb-m5-faza3.md),
extended in faza mailowa Krok 5 (kb/phases/kb-m5-faza-mailowa.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`.

`--transport {direct,http}` (module 5 phase 4 plan §2 decision 6 / §9, added alongside the
fallback task): `direct` (default) is everything above, unchanged. `http` instead calls
`GET {base_url}/search?q=...&mode=flat|cascade|hybrid` on a live `kb-query` and reshapes its
JSON `results` into the same `{"chunks": [...]}` shape the direct-mode functions return, so
`summarize_query_result`/`evaluate_gate` run identically either way. The gate criterion (plan
§9): `dist` for `http` must be **identical** to `direct` against the same live SOLARIA -- same
DB, same retrieval code, HTTP is only a wrapper, so any difference is a serialization/handler
bug, never expected numerical drift. `http` mode cannot sweep `N` (kb-query serves one
server-side default per request, plan §4) -- it reports only at `--gate-n`, and needs no `--dsn`
(kb-query already owns the DB connection; `envelope.source` for criterion 4 comes straight from
each result's `source` field instead of a separate DB lookup).

Usage:
    python retrieval_eval.py --dsn postgresql://kb:<pw>@piha:5433/kb \\
        --ollama-url http://solaria:11434 --n-sweep 5,10,20
    python retrieval_eval.py --transport http --base-url http://192.168.31.5:8230 \\
        --gate-n 10
"""
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}


async def call_search_http(
    session: aiohttp.ClientSession, base_url: str, query_text: str, mode: str
) -> list[dict]:
    """One `GET {base_url}/search?q=...&mode=...` call -> its `results` list. Each result
    already carries `envelope_id`/`dist`/`source` -- exactly the fields `top1_dist`/`hit_at_3`/
    `mail_hit_at_3` need, no DB lookup required on this side."""
    async with session.get(
        f"{base_url}/search", params={"q": query_text, "mode": mode}
    ) as resp:
        resp.raise_for_status()
        data = await resp.json()
    return data["results"]


async def run_query_all_tracks_http(
    session: aiohttp.ClientSession, base_url: str, query: dict, gate_n: int
) -> dict:
    """HTTP-transport equivalent of `run_query_all_tracks` -- three `/search` calls (one per
    mode) instead of embedding+querying locally. `stage1_summaries` isn't part of the HTTP
    response shape (plan §4) so it's reported empty; nothing in `evaluate_gate` reads it."""
    flat_chunks = await call_search_http(session, base_url, query["text"], "flat")
    cascade_chunks = await call_search_http(session, base_url, query["text"], "cascade")
    hybrid_chunks = await call_search_http(session, base_url, query["text"], "hybrid")
    return {
        "query": query,
        "flat": {"chunks": flat_chunks},
        "cascades": {gate_n: {"chunks": cascade_chunks, "stage1_summaries": []}},
        "hybrid": {"chunks": hybrid_chunks},
    }


def envelope_sources_from_results(*chunk_lists: list[dict]) -> dict[str, str]:
    """http transport has no DB to `fetch_envelope_sources` from -- each `/search` result
    already carries its envelope's `source`, so build the same envelope_id -> source mapping
    straight from the response bodies already fetched for this query."""
    return {c["envelope_id"]: c["source"] for chunks in chunk_lists for c in chunks}


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.
        # Since 2026-08-06 one mail query (M5, the requalified N2 control) *does* carry an
        # expected_envelope, verified by hand against the live index. It stays documentation
        # only: this branch deliberately keeps grading every mail_hit by source-match, so the
        # id is never silently half-enforced. Enforcing it would be a change here (identity
        # match when expected_envelope is not None), not in queries.yaml.
        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 (kb/phases/kb-m5-faza3.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_http(args: argparse.Namespace) -> tuple[list[dict], list[dict], list[int]]:
    """`--transport http` path -- no DB connection, three `/search` calls per query. Returns
    only at `--gate-n` (see module docstring: kb-query serves one server-side N per request)."""
    queries = load_queries(Path(args.queries))
    mail_queries = load_mail_queries(Path(args.queries))
    n_values = [args.gate_n]

    async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=60)) as session:
        results = []
        envelope_sources: dict[str, str] = {}
        for query in queries:
            result = await run_query_all_tracks_http(session, args.base_url, query, args.gate_n)
            envelope_sources.update(envelope_sources_from_results(
                result["flat"]["chunks"], result["cascades"][args.gate_n]["chunks"], result["hybrid"]["chunks"],
            ))
            results.append(summarize_query_result(result))

        mail_results_raw = []
        for query in mail_queries:
            result = await run_query_all_tracks_http(session, args.base_url, query, args.gate_n)
            envelope_sources.update(envelope_sources_from_results(result["hybrid"]["chunks"]))
            mail_results_raw.append(result)
        mail_results = [summarize_query_result(r, envelope_sources) for r in mail_results_raw]

    return results, mail_results, n_values


async def main_async_direct(args: argparse.Namespace) -> tuple[list[dict], list[dict], list[int]]:
    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()

    return results, mail_results, n_values


async def main_async(args: argparse.Namespace) -> dict:
    if args.transport == "http":
        if args.n_sweep != "5,10,20":  # the argparse default -- operator didn't ask for a sweep
            print(
                "note: --transport http ignores --n-sweep (kb-query serves a single "
                f"server-side default N per request); reporting only --gate-n={args.gate_n}",
                file=sys.stderr,
            )
        results, mail_results, n_values = await main_async_http(args)
    else:
        results, mail_results, n_values = await main_async_direct(args)

    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("--transport", choices=["direct", "http"], default="direct",
                         help="direct = query DB+Ollama locally (default); http = call a live kb-query's /search")
    parser.add_argument("--base-url", default=None, help="kb-query base URL, required for --transport http (e.g. http://192.168.31.5:8230)")
    parser.add_argument("--dsn", default=os.environ.get("KB_DSN"), help="asyncpg DSN for kb-postgres (or KB_DSN env var); required for --transport direct")
    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; ignored by --transport http)")
    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 args.transport == "http":
        if not args.base_url:
            print("error: --transport http requires --base-url", file=sys.stderr)
            sys.exit(1)
    elif 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()