71eb264448
Cherry-pick 1:1 z porzucanego brancha task/kb-f4-fallback (3d4ee38) wg decyzji z raportu dedup (docs/kb/modules/05-fallback-dedup-raport.md): plan §2 D6/§9 wymaga trybu HTTP-equivalence, a master go nie miał —e7625cdnie tknął tego pliku, patch aplikuje się czysto i woła wyłącznie GET /search (pola envelope_id/dist/source zgodne z odpowiedzią mastera). Live-PASS 2026-07-27 na kodzie brancha; smoke na masterze: CLI + pełny przebieg http przeciwko stubowi /search (raport i werdykt bramki generują się poprawnie). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
521 lines
25 KiB
Python
521 lines
25 KiB
Python
"""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`.
|
|
|
|
`--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.
|
|
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_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()
|