diff --git a/docs/kb/modules/05-faza3-plan.md b/docs/kb/modules/05-faza3-plan.md index dfcc00a..b2856cf 100644 --- a/docs/kb/modules/05-faza3-plan.md +++ b/docs/kb/modules/05-faza3-plan.md @@ -144,6 +144,21 @@ Uzasadnienie i kryteria: **rekomendowany start**) lub `qwen3:14b` (~9 GB, mocny PL, kontekst 32k). Pilot rozstrzyga empirycznie, nie przesądzamy w planie. +**ROZSTRZYGNIĘCIE (Oskar, 2026-07-17, z oceny A/B pilota dwutorowego 186/186 dokumentów):** + +- **Tor kompilacyjny = `claude-haiku-4-5`** — 157/157 streszczeń bez ani jednego JSON-faila, + 100% dyscypliny słownika tagów, wyraźnie lepsza hierarchia treści na długich dokumentach + (OWU, umowy wielostronicowe) niż `gemma3:12b` (155/157, słabsza struktura na długich + dokumentach). Kaskada retrieval (§6) filtruje `document_summary` po `model = + 'claude-haiku-4-5'` z tego powodu. +- **`gemma3:12b` zostaje w odwodzie jako tor lokalny** — dane już w bazie (155 streszczeń + + embeddingi), job (`summarize.py`) wspiera oba backendy bez zmian kodu, więc przełączenie + jest jednym flagowym argumentem, nie migracją. +- **Decyzja o modelu dla fazy mailowej (skala: dziesiątki tysięcy streszczeń) odłożona do + reconu fazy mailowej** — tam ekonomia jest inna (koszt API rzędu setek USD zamiast ~1.5 USD) + i wymaga własnej oceny z flagą prywatności (dane finansowo-tożsamościowe w skali masowej) + i szacunkiem kosztu, nie automatycznego przeniesienia wniosku z pilota 186-dokumentowego. + ### Decyzja 4 — Tagi: słownik kontrolowany vs free-form **Rekomendacja: hybryda — startowy słownik kontrolowany w repo + max 3 tagi free-form @@ -469,6 +484,55 @@ jakość. (c) dopiero potem wniosek o architekturze. Kaskada nieprzechodząca bramki **nie zastępuje** płaskiego retrievalu — płaski zostaje domyślny do skutku. +### 6.3 Implementacja i wynik bramki (2026-07-17) — **PASS, kaskada = domyślna ścieżka** + +**Kod**: `jobs/documents-ingest/src/documents_ingest/retrieval.py` — dwie czyste funkcje +wejściowe pod przyszłe kb-query (faza 4), `query_text → chunki z dist i source`: +`flat_query` (baseline) i `cascade_query` (`summary_model='claude-haiku-4-5'` — decyzja 3 +rozstrzygnięta wyżej). Oba dzielą **jedno** wywołanie embeddingu zapytania (bge-m3) — kaskada +dokłada tylko jeden dodatkowy SQL (stage 1), zero dodatkowych wywołań Ollamy. Testy +jednostkowe z mockami: `tests/test_retrieval.py` (stage1→stage2 zawężenie, koperta bez +aktywnych chunków, N większe niż liczba kopert, no-summaries short-circuit) — 166/166 testów +pakietu przechodzi. + +**Eval-set utrwalony**: `jobs/documents-ingest/eval/queries.yaml` (7 zapytań z pilota +07-16, 1:1 z `docs/kb/eval/retrieval-pilot-2026-07-16.md`, ten plik pozostał nietknięty — +`queries.yaml` to jego wersjonowana kopia robocza). Skrypt bramki (read-only, integracyjny, +**nie wchodzi do pytest**): `jobs/documents-ingest/eval/retrieval_eval.py`. + +**Kalibracja N** (sweep na żywej bazie, k=5 stałe): przy N∈{1,2,3} kaskada **degraduje** +zapytania 2–5 (np. zapytanie 4 spada z 0.448 do 0.532 przy N=1 — trafienie znika); **N=5 jest +najmniejszym progiem bezpiecznym** — od N=5 wzwyż wynik identyczny z płaskim dla wszystkich +7 zapytań. Rekomendacja planu (**N=10, k=5**) ma więc 2× margines nad zmierzonym progiem +minimalnym — potwierdzona bez zmian, nie tylko przyjęta z założenia. + +**Wynik bramki przy N=10, k=5** (żywa baza, `claude-haiku-4-5`, 157 streszczeń, +2545 aktywnych chunków z 2683): + +| Kryterium | Wynik | Werdykt | +|---|---|---| +| 1. żaden płaski hit nie degraduje | 0 degradacji na 5 zapytań-trafień | PASS | +| 2. hit@3 kaskada ≥ hit@3 płaski | 5/5 vs 5/5 | PASS | +| 3. negatywne kontrole > 0.55 (oba tory) | płaski 0.621/0.553, kaskada 0.644/0.553 | PASS | + +**Czy kaskada COŚ poprawia?** Nie przy tej skali (186 dok.) — dystanse top-1 kaskady są +**identyczne co do bita** z płaskim dla wszystkich 5 zapytań-trafień przy N≥5, w tym dla +zapytania 3 (szara strefa "FLL scoring PL", 0.4289 w obu torach — pre-filtr po streszczeniu +ani nie pomaga, ani nie szkodzi). To zgodne z przewidywaniem planu (§6.1): przy 157 +skompilowanych kopertach pre-filtr nie ma czego przyspieszać ani czego poprawiać — test +weryfikuje **architekturę** pod skalę mailową (225k kopert), nie jakość na tym korpusie. +Filtr `ocr_junk` z kroku 1 potwierdził się przy okazji: chunk z kodami kreskowymi, który w +pilocie zajmował poz. 3/5 dla zapytania 2, już nie występuje w wynikach (usunięty z indeksu). + +**Koszt na zapytanie**: płaski = 1 embed + 1 SQL; kaskada = 1 embed (dzielony) + 2 SQL +(stage 1 top-N streszczeń, stage 2 top-k chunków w ich obrębie) — jeden dodatkowy SQL, +zero dodatkowych wywołań Ollamy. + +**Decyzja**: bramka PASS → **kaskada (`cascade_query`, N=10, k=5, `summary_model= +'claude-haiku-4-5'`) jest odtąd domyślną ścieżką retrievalu** dla przyszłego kb-query +(faza 4). Płaski (`flat_query`) zostaje w kodzie jako baseline/fallback do porównań i do +architektur bez warstwy streszczeń. + --- ## 7. Krok 5 — cykliczny ingest (adapter + embed jako timer) diff --git a/jobs/documents-ingest/README.md b/jobs/documents-ingest/README.md index 40a1ef9..689fcc8 100644 --- a/jobs/documents-ingest/README.md +++ b/jobs/documents-ingest/README.md @@ -479,3 +479,59 @@ kb-postgres@PIHA: (50 sekwencyjnych wywołań /api/embeddings, ~600-tok prompt) vs 790ms/chunk CPU baseline — ~3.8× szybciej sekwencyjnie; przy pojedynczych requestach dominuje overhead HTTP/tokenizacji, realny skok da dopiero batching (backlog). + +--- + +## Phase 3 step 4 — retrieval cascade (`documents_ingest.retrieval`) + quality gate + +Module 5, phase 3, plan step 4 (`docs/kb/modules/05-faza3-plan.md`, §6). Two retrieval +paths, both `query_text -> chunk hits (dist, source)` — the intended clean API surface for +phase 4's kb-query, not just this eval: + +- `flat_query` — baseline: rank every active `document_chunk` row directly. Formalizes the + phase-2 pilot's ad hoc `/tmp/kbq.sh` query into a tested module. +- `cascade_query` — pre-filter to the top-N `document_summary` envelopes (one `model`, + default `claude-haiku-4-5` — plan §2 decision 3, resolved 2026-07-17) before ranking + `document_chunk` within just those envelopes. Both share **one** query embedding call; + the cascade only adds one extra SQL query (stage 1), never an extra Ollama call. + +`envelope`, `document_chunk`, and `document_summary` are read-only — this module only ever +`SELECT`s. + +### Quality gate + +`eval/queries.yaml` — 7 queries transcribed 1:1 from the phase-2 pilot baseline +(`docs/kb/eval/retrieval-pilot-2026-07-16.md`, left untouched — this is its versioned working +copy) with expected envelope / kind (`hit`, `grey_zone`, `negative_control`, +`negative_control_borderline`) per query. + +`eval/retrieval_eval.py` — read-only integration script against the live DB + live Ollama, +**not collected by pytest** (same reasoning as the plan: an eval gate against live data isn't +a mocked unit test). Runs every query through both tracks across an N sweep and checks the +plan's three gate criteria (no flat hit degrades, hit@3 cascade ≥ flat, negative controls +stay > 0.55). Exits 0 on PASS, 1 on FAIL. + +```bash +pip install -e packages/kb-mail/ -e jobs/documents-ingest/ +python eval/retrieval_eval.py --dsn postgresql://kb:@piha:5433/kb \ + --ollama-url http://solaria:11434 --n-sweep 5,10,20 +``` + +**Result (2026-07-17, live run)**: PASS at N=10, k=5 — see plan §6.3 for the full table, +the N-sweep calibration (N=5 is the measured safety floor; the plan's N=10 default carries a +2× margin), and the cost/improvement analysis. `cascade_query` (N=10, k=5, +`summary_model='claude-haiku-4-5'`) is now the default retrieval path for phase 4's kb-query; +`flat_query` stays as the baseline/fallback. + +### Tests + +```bash +pip install -e packages/kb-mail/ +pip install -e jobs/documents-ingest/ +cd jobs/documents-ingest && pytest +``` + +`tests/test_retrieval.py` — pure unit tests, no DB or real HTTP. Covers: flat ranking across +all envelopes, cascade stage-1-narrows-stage-2, an envelope whose summary exists but has no +active chunks, N larger than the number of summarized envelopes, the no-summaries +short-circuit (stage 2 never queried), and both query entry points embedding exactly once. diff --git a/jobs/documents-ingest/eval/queries.yaml b/jobs/documents-ingest/eval/queries.yaml new file mode 100644 index 0000000..5ef06a4 --- /dev/null +++ b/jobs/documents-ingest/eval/queries.yaml @@ -0,0 +1,73 @@ +# Eval-set for the retrieval quality gate (module 5, phase 3, plan §6.2, +# docs/kb/modules/05-faza3-plan.md). Transcribed 1:1 from the pilot baseline, +# docs/kb/eval/retrieval-pilot-2026-07-16.md (read-only source -- this file is the versioned +# copy the plan asked for, so the set stops living only in a session transcript). +# +# kind: +# hit -- pilot found a correct top-1 match (dist < 0.45) +# grey_zone -- pilot landed in the 0.45-0.55 band (on-topic, imprecise) +# negative_control -- pilot correctly found nothing (dist > 0.55) +# negative_control_borderline -- pilot correctly found nothing, but close to the 0.55 edge +# +# expected_envelope is null for both negative-control kinds: there is no document these +# queries should match, by design. + +queries: + - id: "1" + text: "polisa ubezpieczeniowa PZU warunki odpowiedzialności" + kind: hit + expected_envelope: "paperless:14" + baseline_top1_dist: 0.3418 + note: > + Pilot top-1 = paperless:14 (OWU PZU Auto); ujawnił duplikat paperless:14 ≡ paperless:74 + (flagowany excluded_reason='duplicate' w kroku 1 fazy 3). + + - id: "2" + text: "faktura za usługi telekomunikacyjne kwota do zapłaty" + kind: hit + expected_envelope: "paperless:192" + baseline_top1_dist: 0.3447 + note: > + Pilot top-1 = paperless:192 (faktura P4); poz. 3/5 w pilocie był chunk z OCR-śmieciem + (kody kreskowe) -- powinien wypaść po excluded_reason='ocr_junk' z kroku 1. + + - id: "3" + text: "zasady punktacji FLL Challenge robot game (PL)" + kind: grey_zone + expected_envelope: "paperless:5" + baseline_top1_dist: 0.4154 + note: > + Pilot: szara strefa -- rodzina dokumentów OK, sedno (zasady punktacji) nie trafione + precyzyjnie. Kandydat na poprawę: czy pre-filtr po streszczeniu podnosi trafność? + + - id: "4" + text: "FLL robot game mission scoring points table (EN)" + kind: hit + expected_envelope: "paperless:5" + baseline_top1_dist: 0.4114 + note: > + Cross-lingual (EN zapytanie) trafia lepiej niż PL (zapytanie 3); poz. 2+ w pilocie = + paperless:119 (scoresheet, mojibake) -- mojibake nie jest śmieciem, niesie sygnał. + + - id: "5" + text: "innovation project scoring" + kind: hit + expected_envelope: "paperless:3" + baseline_top1_dist: 0.3869 + note: "Najlepszy wynik pilota: top-5 spójnie z jednego właściwego dokumentu (arkusz ocen IP)." + + - id: "N" + text: "przepis na sernik z rodzynkami" + kind: negative_control + expected_envelope: null + baseline_top1_dist: 0.6210 + note: "Poprawny brak w pilocie; separacja od trafień wyraźna." + + - id: "N2" + text: "piaskownica plastikowa" + kind: negative_control_borderline + expected_envelope: null + baseline_top1_dist: 0.5533 + note: > + Poprawnie na granicy "brak" w pilocie -- semantycznie sąsiednie dokumenty wspólnoty + mieszkaniowej, nie odpowiedź na zapytanie. diff --git a/jobs/documents-ingest/eval/retrieval_eval.py b/jobs/documents-ingest/eval/retrieval_eval.py new file mode 100644 index 0000000..588afb2 --- /dev/null +++ b/jobs/documents-ingest/eval/retrieval_eval.py @@ -0,0 +1,264 @@ +"""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:@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() diff --git a/jobs/documents-ingest/src/documents_ingest/retrieval.py b/jobs/documents-ingest/src/documents_ingest/retrieval.py new file mode 100644 index 0000000..f4320aa --- /dev/null +++ b/jobs/documents-ingest/src/documents_ingest/retrieval.py @@ -0,0 +1,135 @@ +"""Retrieval module -- module 5, phase 3, plan step 4 (docs/kb/modules/05-faza3-plan.md, §6). + +Two retrieval paths over the same corpus, sharing one query embedding (bge-m3, via Ollama): + +- `flat_retrieve`: baseline -- ranks every active `document_chunk` row directly. This + formalizes the pilot's ad hoc `/tmp/kbq.sh` query (docs/kb/eval/retrieval-pilot-2026-07-16.md) + into a tested, versioned module instead of a script living only in a session transcript. +- `cascade_retrieve`: pre-filters to the top-N `document_summary` matches for one configured + `model` (plan §2 decision 3, resolved 2026-07-17 as D3: `claude-haiku-4-5` is the compilation + track; `gemma3:12b` stays in reserve as the local track) before ranking `document_chunk` + within just those envelopes. At the 186-document pilot scale this doesn't speed anything up + -- it is an architecture test for the mail-scale corpus (225k envelopes, plan §1.1) where a + flat chunk scan stops being cheap. `eval/retrieval_eval.py` runs the quality gate (plan §6.2) + that decides whether it becomes the default path. + +`cascade_query` / `flat_query` are the intended clean entry points for phase 4's kb-query: +query_text -> chunk hits with `dist` (cosine distance) and `source` ("flat" | "cascade"). +`envelope`, `document_chunk`, and `document_summary` are read-only here -- this module only +ever `SELECT`s. +""" +from __future__ import annotations + +import aiohttp +import asyncpg + +from documents_ingest.chunk_embed import _vector_literal, embed_chunk + +DEFAULT_SUMMARY_MODEL = "claude-haiku-4-5" # plan §2 D3 resolution 2026-07-17: compilation track +DEFAULT_EMBED_MODEL = "bge-m3" +DEFAULT_N = 10 # plan §6.1 start value +DEFAULT_K = 5 # plan §6.1 start value + + +async def flat_retrieve(conn: asyncpg.Connection, query_vector: str, k: int = DEFAULT_K) -> list[dict]: + """Baseline: rank every active chunk directly against the query embedding, no pre-filter.""" + rows = await conn.fetch( + "SELECT envelope_id, chunk_index, text, embedding <=> $1::vector AS dist " + "FROM document_chunk WHERE excluded_reason IS NULL AND embedding IS NOT NULL " + "ORDER BY embedding <=> $1::vector LIMIT $2", + query_vector, k, + ) + return [ + { + "envelope_id": r["envelope_id"], + "chunk_index": r["chunk_index"], + "text": r["text"], + "dist": r["dist"], + "source": "flat", + } + for r in rows + ] + + +async def cascade_retrieve( + conn: asyncpg.Connection, + query_vector: str, + summary_model: str = DEFAULT_SUMMARY_MODEL, + n: int = DEFAULT_N, + k: int = DEFAULT_K, +) -> dict: + """Stage 1: top-N `document_summary` envelopes for `summary_model`. Stage 2: top-k + `document_chunk` ranked within just those envelopes. + + `n` exceeding the number of summarized envelopes is not an error -- plain SQL `LIMIT` + semantics just return all of them. An empty stage 1 (no summaries for this model, or a + corpus that hasn't been summarized yet) short-circuits before stage 2 runs at all: a + cascade can never rank chunks in envelopes it didn't pre-filter into, so there is nothing + for stage 2 to query. + """ + stage1 = await conn.fetch( + "SELECT envelope_id, embedding <=> $1::vector AS dist FROM document_summary " + "WHERE model = $2 AND embedding IS NOT NULL ORDER BY embedding <=> $1::vector LIMIT $3", + query_vector, summary_model, n, + ) + stage1_summaries = [{"envelope_id": r["envelope_id"], "dist": r["dist"]} for r in stage1] + + envelope_ids = [s["envelope_id"] for s in stage1_summaries] + if not envelope_ids: + return {"stage1_summaries": stage1_summaries, "chunks": []} + + stage2 = await conn.fetch( + "SELECT envelope_id, chunk_index, text, embedding <=> $1::vector AS dist " + "FROM document_chunk WHERE envelope_id = ANY($2::text[]) AND excluded_reason IS NULL " + "AND embedding IS NOT NULL ORDER BY embedding <=> $1::vector LIMIT $3", + query_vector, envelope_ids, k, + ) + chunks = [ + { + "envelope_id": r["envelope_id"], + "chunk_index": r["chunk_index"], + "text": r["text"], + "dist": r["dist"], + "source": "cascade", + } + for r in stage2 + ] + return {"stage1_summaries": stage1_summaries, "chunks": chunks} + + +async def flat_query( + conn: asyncpg.Connection, + session: aiohttp.ClientSession, + ollama_url: str, + query_text: str, + embed_model: str = DEFAULT_EMBED_MODEL, + k: int = DEFAULT_K, +) -> dict: + """query_text -> flat chunk hits. One Ollama embed call, one SQL query.""" + embedding, _elapsed = await embed_chunk(session, ollama_url, embed_model, query_text) + chunks = await flat_retrieve(conn, _vector_literal(embedding), k=k) + return {"query": query_text, "k": k, "chunks": chunks} + + +async def cascade_query( + conn: asyncpg.Connection, + session: aiohttp.ClientSession, + ollama_url: str, + query_text: str, + summary_model: str = DEFAULT_SUMMARY_MODEL, + embed_model: str = DEFAULT_EMBED_MODEL, + n: int = DEFAULT_N, + k: int = DEFAULT_K, +) -> dict: + """query_text -> cascade chunk hits, plus the stage-1 envelope pre-filter (needed for the + gate's diagnosis order, plan §6.2: "N too small" is diagnosed by looking at stage 1). + + Same query embedding as `flat_query`'s single Ollama call -- the cascade's only added + cost over the flat path is one extra SQL query (stage 1), never an extra embed. + """ + embedding, _elapsed = await embed_chunk(session, ollama_url, embed_model, query_text) + result = await cascade_retrieve(conn, _vector_literal(embedding), summary_model, n=n, k=k) + result["query"] = query_text + result["n"] = n + result["k"] = k + return result diff --git a/jobs/documents-ingest/tests/test_retrieval.py b/jobs/documents-ingest/tests/test_retrieval.py new file mode 100644 index 0000000..8cccaaf --- /dev/null +++ b/jobs/documents-ingest/tests/test_retrieval.py @@ -0,0 +1,174 @@ +"""Unit tests for the retrieval cascade -- no DB, no real HTTP/Ollama.""" +from __future__ import annotations + +import pytest + +from documents_ingest.retrieval import ( + DEFAULT_K, + DEFAULT_N, + cascade_query, + cascade_retrieve, + flat_query, + flat_retrieve, +) + + +class _FakeConn: + """summaries: [(envelope_id, dist), ...] already in distance order (mirrors what the real + `ORDER BY embedding <=> $1` would hand back). chunks_by_envelope: envelope_id -> [(chunk_index, + text, dist), ...]. Both stage-1 and the flat path are served from the same fixture so a test + can assert the cascade excludes chunks the flat path would have surfaced.""" + + def __init__(self, summaries=None, chunks_by_envelope=None): + self._summaries = list(summaries or []) + self._chunks_by_envelope = chunks_by_envelope or {} + self.queries: list[tuple] = [] + + async def fetch(self, query, *params): + self.queries.append((query, params)) + if "FROM document_summary" in query: + _, _model, limit = params + return [{"envelope_id": eid, "dist": dist} for eid, dist in self._summaries[:limit]] + if "FROM document_chunk" in query and "= ANY" in query: + _, envelope_ids, limit = params + rows = [ + {"envelope_id": eid, "chunk_index": idx, "text": text, "dist": dist} + for eid in envelope_ids + for idx, text, dist in self._chunks_by_envelope.get(eid, []) + ] + rows.sort(key=lambda r: r["dist"]) + return rows[:limit] + if "FROM document_chunk" in query: # flat path: no envelope pre-filter + _, limit = params + rows = [ + {"envelope_id": eid, "chunk_index": idx, "text": text, "dist": dist} + for eid, chunk_list in self._chunks_by_envelope.items() + for idx, text, dist in chunk_list + ] + rows.sort(key=lambda r: r["dist"]) + return rows[:limit] + raise AssertionError(f"unexpected query: {query}") + + async def close(self): + pass + + +class _FakeEmbedResponse: + def __init__(self, payload): + self._payload = payload + + async def __aenter__(self): + return self + + async def __aexit__(self, *exc): + return False + + def raise_for_status(self): + pass + + async def json(self): + return self._payload + + +class _FakeSession: + def __init__(self): + self.post_calls: list[dict] = [] + + def post(self, url, json): + self.post_calls.append({"url": url, "json": json}) + return _FakeEmbedResponse({"embedding": [0.01] * 1024}) + + +class TestDefaults: + def test_plan_start_values(self): + # plan §6.1: "Start: N=10, k=5" -- a regression guard against silently drifting off + # the value the quality gate (docs/kb/modules/05-faza3-plan.md §6.2) was run against. + assert DEFAULT_N == 10 + assert DEFAULT_K == 5 + + +class TestFlatRetrieve: + async def test_ranks_across_all_envelopes(self): + conn = _FakeConn( + chunks_by_envelope={ + "paperless:1": [(0, "a", 0.5)], + "paperless:2": [(0, "b", 0.1)], + } + ) + result = await flat_retrieve(conn, "[0.1]", k=5) + assert [c["envelope_id"] for c in result] == ["paperless:2", "paperless:1"] + assert all(c["source"] == "flat" for c in result) + + async def test_respects_k(self): + conn = _FakeConn( + chunks_by_envelope={"paperless:1": [(i, f"c{i}", i / 10) for i in range(10)]} + ) + result = await flat_retrieve(conn, "[0.1]", k=3) + assert len(result) == 3 + + +class TestCascadeRetrieve: + async def test_stage1_narrows_stage2(self): + # paperless:3 has the single closest chunk overall, but it never enters stage 1's + # top-N summaries -- the cascade must not surface it, unlike a flat scan would. + conn = _FakeConn( + summaries=[("paperless:1", 0.1), ("paperless:2", 0.2)], + chunks_by_envelope={ + "paperless:1": [(0, "a", 0.3)], + "paperless:2": [(0, "b", 0.4)], + "paperless:3": [(0, "c", 0.01)], + }, + ) + result = await cascade_retrieve(conn, "[0.1]", "claude-haiku-4-5", n=2, k=5) + envelope_ids = {c["envelope_id"] for c in result["chunks"]} + assert envelope_ids == {"paperless:1", "paperless:2"} + assert all(c["source"] == "cascade" for c in result["chunks"]) + + async def test_envelope_without_active_chunks_yields_no_chunks_for_it(self): + # e.g. every chunk in this envelope is excluded_reason-flagged (junk/duplicate) -- + # the summary still exists and surfaces in stage 1, stage 2 just finds nothing there. + conn = _FakeConn(summaries=[("paperless:1", 0.1)], chunks_by_envelope={}) + result = await cascade_retrieve(conn, "[0.1]", "claude-haiku-4-5", n=10, k=5) + assert result["stage1_summaries"] == [{"envelope_id": "paperless:1", "dist": 0.1}] + assert result["chunks"] == [] + + async def test_n_larger_than_available_summaries_returns_all_of_them(self): + conn = _FakeConn( + summaries=[("paperless:1", 0.1), ("paperless:2", 0.2)], + chunks_by_envelope={"paperless:1": [(0, "a", 0.3)], "paperless:2": [(0, "b", 0.4)]}, + ) + result = await cascade_retrieve(conn, "[0.1]", "claude-haiku-4-5", n=1000, k=5) + assert len(result["stage1_summaries"]) == 2 + assert {c["envelope_id"] for c in result["chunks"]} == {"paperless:1", "paperless:2"} + + async def test_no_summaries_short_circuits_before_stage2_query(self): + conn = _FakeConn(summaries=[], chunks_by_envelope={"paperless:1": [(0, "a", 0.1)]}) + result = await cascade_retrieve(conn, "[0.1]", "claude-haiku-4-5", n=10, k=5) + assert result == {"stage1_summaries": [], "chunks": []} + assert len(conn.queries) == 1 # stage 2 never ran -- nothing to narrow into + + +class TestQueryEntryPoints: + async def test_flat_query_embeds_once_and_returns_query_text(self): + conn = _FakeConn(chunks_by_envelope={"paperless:1": [(0, "a", 0.2)]}) + session = _FakeSession() + result = await flat_query(conn, session, "http://fake-ollama", "sernik z rodzynkami", k=5) + assert result["query"] == "sernik z rodzynkami" + assert len(session.post_calls) == 1 + assert result["chunks"][0]["source"] == "flat" + + async def test_cascade_query_embeds_once_shared_across_both_stages(self): + conn = _FakeConn( + summaries=[("paperless:1", 0.1)], + chunks_by_envelope={"paperless:1": [(0, "a", 0.2)]}, + ) + session = _FakeSession() + result = await cascade_query( + conn, session, "http://fake-ollama", "polisa PZU", summary_model="claude-haiku-4-5", n=10, k=5 + ) + assert result["query"] == "polisa PZU" + assert result["n"] == 10 + assert result["k"] == 5 + # one embed call total, reused for both the stage-1 and stage-2 SQL queries + assert len(session.post_calls) == 1 + assert result["chunks"][0]["source"] == "cascade"