From a640cf1455297bdf1edb3b0c3bde8c4560bc5b20 Mon Sep 17 00:00:00 2001 From: oskar Date: Wed, 22 Jul 2026 19:01:21 +0200 Subject: [PATCH] feat(kb-retrieval,kb-query): add hybrid retrieval mode (faza mailowa Krok 3) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mail (gmail) envelopes never get a document_summary (Decyzja 6 -- a mail "summary" would usually be longer than the mail itself), so they're invisible to the cascade's stage-1 pre-filter. hybrid_retrieve runs the existing cascade for summarized sources (paperless) and, in parallel, a direct chunk scan restricted to summaryless_sources (gmail), merging both by dist -- same embedder/cosine space, so the merge is a plain sort, no re-normalization. hybrid_query mirrors cascade_query (one shared query embed). kb-query: mode pattern extended to ^(cascade|flat|hybrid)$, /search routes "hybrid" to hybrid_query. Default mode stays "cascade" until the quality gate (plan §8) PASSes on the full mail corpus -- flipping the default, and deploying this to the running kb-query container, are separate follow-ups for the operator; this task only adds the code path + tests (docs/kb/modules/05-faza-mailowa-plan.md, §6, Krok 3). --- .../src/kb_retrieval/retrieval.py | 90 +++++++++++++++++-- packages/kb-retrieval/tests/test_retrieval.py | 79 +++++++++++++++- services/kb-query/app/main.py | 6 +- services/kb-query/app/search.py | 7 +- services/kb-query/tests/test_search.py | 36 +++++++- 5 files changed, 208 insertions(+), 10 deletions(-) diff --git a/packages/kb-retrieval/src/kb_retrieval/retrieval.py b/packages/kb-retrieval/src/kb_retrieval/retrieval.py index ee61f33..443686f 100644 --- a/packages/kb-retrieval/src/kb_retrieval/retrieval.py +++ b/packages/kb-retrieval/src/kb_retrieval/retrieval.py @@ -3,7 +3,7 @@ Moved 1:1 into `packages/kb-retrieval` in phase 4 (docs/kb/modules/05-faza4-plan decision 1) so both `documents-ingest` (venv job) and `kb-query` (Docker service) share one tested module instead of the service image needing to pull in all of `jobs/documents-ingest`. -Two retrieval paths over the same corpus, sharing one query embedding (bge-m3, via Ollama): +Three 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) @@ -15,11 +15,19 @@ Two retrieval paths over the same corpus, sharing one query embedding (bge-m3, v -- 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. +- `hybrid_retrieve` -- module 5, faza mailowa, plan Krok 3 (docs/kb/modules/05-faza-mailowa-plan.md, + §6, decision 6): mail (gmail) envelopes never get a `document_summary` (streszczenie maila + would usually be longer than the mail itself -- decision 6's rejected-summaries reasoning), so + they are invisible to the cascade's stage 1 pre-filter. Hybrid runs the cascade for sources that + DO have summaries (paperless) and, in parallel, a direct flat-style HNSW scan restricted to + `summaryless_sources` (gmail), then merges both result sets by `dist` (same embedder, same + cosine space -- merge is a plain sort, no re-normalization needed). This is the only path that + can surface mail content until/unless mail gets its own summaries later (decision 6, deferred). -`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. +`cascade_query` / `flat_query` / `hybrid_query` are the intended clean entry points for kb-query: +query_text -> chunk hits with `dist` (cosine distance) and `source` ("flat" | "cascade" | +"hybrid"). `envelope`, `document_chunk`, and `document_summary` are read-only here -- this module +only ever `SELECT`s. """ from __future__ import annotations @@ -32,6 +40,7 @@ DEFAULT_SUMMARY_MODEL = "claude-haiku-4-5" # plan §2 D3 resolution 2026-07-17: DEFAULT_EMBED_MODEL = "bge-m3" DEFAULT_N = 10 # plan §6.1 start value DEFAULT_K = 5 # plan §6.1 start value +DEFAULT_SUMMARYLESS_SOURCES = ("gmail",) # faza mailowa plan §6 Krok 3: sources with no summaries async def flat_retrieve(conn: asyncpg.Connection, query_vector: str, k: int = DEFAULT_K) -> list[dict]: @@ -100,6 +109,51 @@ async def cascade_retrieve( return {"stage1_summaries": stage1_summaries, "chunks": chunks} +async def hybrid_retrieve( + conn: asyncpg.Connection, + query_vector: str, + summary_model: str = DEFAULT_SUMMARY_MODEL, + summaryless_sources: tuple[str, ...] = DEFAULT_SUMMARYLESS_SOURCES, + n: int = DEFAULT_N, + k: int = DEFAULT_K, +) -> dict: + """Cascade (for summarized sources) merged with a direct chunk scan over + `summaryless_sources` (plan §6 decision 6) -- mail envelopes never reach the cascade's + stage 1, so this is the only path today that can surface them. + + Both branches share the same query embedding/cosine space, so the merge is a plain sort by + `dist`, then truncate to top-k -- no re-normalization needed (plan §6 decision 6 rationale). + Each returned chunk's `source` is overwritten to `"hybrid"` regardless of which branch it + came from, matching `flat`/`cascade`'s convention that `source` names the retrieval mode, + not an internal sub-path. + """ + cascade = await cascade_retrieve(conn, query_vector, summary_model, n=n, k=k) + + mail_rows = await conn.fetch( + "SELECT c.envelope_id, c.chunk_index, c.text, c.embedding <=> $1::vector AS dist " + "FROM document_chunk c JOIN envelope e ON e.id = c.envelope_id " + "WHERE e.source = ANY($2::text[]) AND c.excluded_reason IS NULL " + "AND c.embedding IS NOT NULL ORDER BY c.embedding <=> $1::vector LIMIT $3", + query_vector, list(summaryless_sources), k, + ) + mail_chunks = [ + { + "envelope_id": r["envelope_id"], + "chunk_index": r["chunk_index"], + "text": r["text"], + "dist": r["dist"], + "source": "hybrid", + } + for r in mail_rows + ] + + merged = sorted( + [{**c, "source": "hybrid"} for c in cascade["chunks"]] + mail_chunks, + key=lambda c: c["dist"], + )[:k] + return {"stage1_summaries": cascade["stage1_summaries"], "chunks": merged} + + async def flat_query( conn: asyncpg.Connection, session: aiohttp.ClientSession, @@ -136,3 +190,29 @@ async def cascade_query( result["n"] = n result["k"] = k return result + + +async def hybrid_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, + summaryless_sources: tuple[str, ...] = DEFAULT_SUMMARYLESS_SOURCES, + n: int = DEFAULT_N, + k: int = DEFAULT_K, +) -> dict: + """query_text -> hybrid chunk hits (cascade + direct mail scan, merged by dist). + + Same single query embedding as `flat_query`/`cascade_query` -- hybrid's only added cost + over cascade is one extra SQL query (the mail branch), never an extra Ollama call. + """ + embedding, _elapsed = await embed_chunk(session, ollama_url, embed_model, query_text) + result = await hybrid_retrieve( + conn, _vector_literal(embedding), summary_model, summaryless_sources, n=n, k=k + ) + result["query"] = query_text + result["n"] = n + result["k"] = k + return result diff --git a/packages/kb-retrieval/tests/test_retrieval.py b/packages/kb-retrieval/tests/test_retrieval.py index ae21819..904c2d8 100644 --- a/packages/kb-retrieval/tests/test_retrieval.py +++ b/packages/kb-retrieval/tests/test_retrieval.py @@ -10,6 +10,8 @@ from kb_retrieval.retrieval import ( cascade_retrieve, flat_query, flat_retrieve, + hybrid_query, + hybrid_retrieve, ) @@ -17,11 +19,16 @@ 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.""" + can assert the cascade excludes chunks the flat path would have surfaced. - def __init__(self, summaries=None, chunks_by_envelope=None): + mail_chunks_by_source: source -> [(envelope_id, chunk_index, text, dist), ...] -- served by + hybrid_retrieve's direct summaryless-source scan (JOIN envelope ... source = ANY), routed + separately from stage 2's envelope_id = ANY query by checking for "JOIN envelope" first.""" + + def __init__(self, summaries=None, chunks_by_envelope=None, mail_chunks_by_source=None): self._summaries = list(summaries or []) self._chunks_by_envelope = chunks_by_envelope or {} + self._mail_chunks_by_source = mail_chunks_by_source or {} self.queries: list[tuple] = [] async def fetch(self, query, *params): @@ -29,6 +36,15 @@ class _FakeConn: if "FROM document_summary" in query: _, _model, limit = params return [{"envelope_id": eid, "dist": dist} for eid, dist in self._summaries[:limit]] + if "JOIN envelope" in query: # hybrid's direct summaryless-source chunk scan + _, sources, limit = params + rows = [ + {"envelope_id": eid, "chunk_index": idx, "text": text, "dist": dist} + for source in sources + for eid, idx, text, dist in self._mail_chunks_by_source.get(source, []) + ] + rows.sort(key=lambda r: r["dist"]) + return rows[:limit] if "FROM document_chunk" in query and "= ANY" in query: _, envelope_ids, limit = params rows = [ @@ -148,6 +164,49 @@ class TestCascadeRetrieve: assert len(conn.queries) == 1 # stage 2 never ran -- nothing to narrow into +class TestHybridRetrieve: + async def test_merges_cascade_and_mail_branches_by_dist(self): + conn = _FakeConn( + summaries=[("paperless:1", 0.1)], + chunks_by_envelope={"paperless:1": [(0, "a", 0.3)]}, + mail_chunks_by_source={"gmail": [("msgid1@x", 0, "b", 0.2)]}, + ) + result = await hybrid_retrieve(conn, "[0.1]", "claude-haiku-4-5", ("gmail",), n=10, k=5) + assert [c["envelope_id"] for c in result["chunks"]] == ["msgid1@x", "paperless:1"] + assert all(c["source"] == "hybrid" for c in result["chunks"]) + + async def test_empty_cascade_branch_still_returns_mail_chunks(self): + # no summaries at all (e.g. corpus not summarized yet) -- cascade short-circuits to []. + conn = _FakeConn( + summaries=[], + mail_chunks_by_source={"gmail": [("msgid1@x", 0, "b", 0.2)]}, + ) + result = await hybrid_retrieve(conn, "[0.1]", "claude-haiku-4-5", ("gmail",), n=10, k=5) + assert [c["envelope_id"] for c in result["chunks"]] == ["msgid1@x"] + + async def test_empty_mail_branch_still_returns_cascade_chunks(self): + conn = _FakeConn( + summaries=[("paperless:1", 0.1)], + chunks_by_envelope={"paperless:1": [(0, "a", 0.3)]}, + mail_chunks_by_source={}, + ) + result = await hybrid_retrieve(conn, "[0.1]", "claude-haiku-4-5", ("gmail",), n=10, k=5) + assert [c["envelope_id"] for c in result["chunks"]] == ["paperless:1"] + + async def test_truncates_merged_results_to_k(self): + conn = _FakeConn( + summaries=[("paperless:1", 0.1)], + chunks_by_envelope={"paperless:1": [(i, f"c{i}", i / 10) for i in range(5)]}, + mail_chunks_by_source={ + "gmail": [(f"msg{i}@x", 0, f"m{i}", i / 10 + 0.05) for i in range(5)] + }, + ) + result = await hybrid_retrieve(conn, "[0.1]", "claude-haiku-4-5", ("gmail",), n=10, k=3) + assert len(result["chunks"]) == 3 + dists = [c["dist"] for c in result["chunks"]] + assert dists == sorted(dists) + + 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)]}) @@ -172,3 +231,19 @@ class TestQueryEntryPoints: # 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" + + async def test_hybrid_query_embeds_once_and_merges(self): + conn = _FakeConn( + summaries=[("paperless:1", 0.1)], + chunks_by_envelope={"paperless:1": [(0, "a", 0.3)]}, + mail_chunks_by_source={"gmail": [("msgid1@x", 0, "b", 0.2)]}, + ) + session = _FakeSession() + result = await hybrid_query( + conn, session, "http://fake-ollama", "od kogo ta wiadomosc", + summary_model="claude-haiku-4-5", summaryless_sources=("gmail",), n=10, k=5, + ) + assert result["query"] == "od kogo ta wiadomosc" + assert len(session.post_calls) == 1 + assert [c["envelope_id"] for c in result["chunks"]] == ["msgid1@x", "paperless:1"] + assert all(c["source"] == "hybrid" for c in result["chunks"]) diff --git a/services/kb-query/app/main.py b/services/kb-query/app/main.py index 4520ea3..1bf7d3d 100644 --- a/services/kb-query/app/main.py +++ b/services/kb-query/app/main.py @@ -11,6 +11,10 @@ rather than a bare 500. `GET /` (Krok 4, plan §7) serves the search UI from this same FastAPI process -- one image, one container (plan §2 decision 4): a Jinja2 shell + a static vanilla-JS file, no node build step. `/` and `/static/*` need no DB/Ollama, so they stay reachable even while `/search` is 503ing. + +`mode=hybrid` (faza mailowa, plan Krok 3, docs/kb/modules/05-faza-mailowa-plan.md §6) is +available explicitly starting here, but the default stays `cascade` until the quality gate +(plan §8) PASSes on the full mail corpus -- flipping the default is a separate, later change. """ from __future__ import annotations @@ -77,7 +81,7 @@ async def healthz() -> dict: @app.get("/search") async def search( q: str = Query(..., min_length=1), - mode: str = Query("cascade", pattern="^(cascade|flat)$"), + mode: str = Query("cascade", pattern="^(cascade|flat|hybrid)$"), ) -> dict: try: async with app.state.pool.acquire() as conn: diff --git a/services/kb-query/app/search.py b/services/kb-query/app/search.py index 19fdd96..be1fee6 100644 --- a/services/kb-query/app/search.py +++ b/services/kb-query/app/search.py @@ -16,7 +16,7 @@ from __future__ import annotations import aiohttp import asyncpg -from kb_retrieval.retrieval import cascade_query, flat_query +from kb_retrieval.retrieval import cascade_query, flat_query, hybrid_query from app.db import fetch_envelopes, fetch_summaries from app.links import build_result @@ -33,6 +33,11 @@ async def run_search( ) -> dict: if mode == "flat": retrieval = await flat_query(conn, session, ollama_url, query_text, embed_model=embed_model) + elif mode == "hybrid": + retrieval = await hybrid_query( + conn, session, ollama_url, query_text, + summary_model=summary_model, embed_model=embed_model, + ) else: retrieval = await cascade_query( conn, session, ollama_url, query_text, diff --git a/services/kb-query/tests/test_search.py b/services/kb-query/tests/test_search.py index 570deb9..cdcb97a 100644 --- a/services/kb-query/tests/test_search.py +++ b/services/kb-query/tests/test_search.py @@ -17,11 +17,15 @@ class _FakeConn: "entities": [...]}. summary_texts: envelope_id -> {"summary": ..., "tags": [...]} -- the document_summary row fetched for the result header (app/db.py fetch_summaries).""" - def __init__(self, summaries=None, chunks_by_envelope=None, envelopes=None, summary_texts=None): + def __init__( + self, summaries=None, chunks_by_envelope=None, envelopes=None, summary_texts=None, + mail_chunks_by_source=None, + ): self._summaries = list(summaries or []) self._chunks_by_envelope = chunks_by_envelope or {} self._envelopes = envelopes or {} self._summary_texts = summary_texts or {} + self._mail_chunks_by_source = mail_chunks_by_source or {} async def fetch(self, query, *params): if "FROM document_summary" in query and "= ANY" in query: @@ -34,6 +38,15 @@ class _FakeConn: if "FROM document_summary" in query: _, _model, limit = params return [{"envelope_id": eid, "dist": dist} for eid, dist in self._summaries[:limit]] + if "JOIN envelope" in query: # hybrid's direct summaryless-source chunk scan + _, sources, limit = params + rows = [ + {"envelope_id": eid, "chunk_index": idx, "text": text, "dist": dist} + for source in sources + for eid, idx, text, dist in self._mail_chunks_by_source.get(source, []) + ] + rows.sort(key=lambda r: r["dist"]) + return rows[:limit] if "FROM document_chunk" in query and "= ANY" in query: _, envelope_ids, limit = params rows = [ @@ -152,6 +165,27 @@ class TestRunSearchHappyPath: assert hit["summary"] is None assert hit["summary_tags"] == [] + async def test_hybrid_mode_merges_cascade_and_mail_branches(self): + conn = _FakeConn( + summaries=[("paperless:1", 0.3)], + chunks_by_envelope={"paperless:1": [(0, "doc text", 0.3)]}, + envelopes={ + "paperless:1": {"source": "paperless", "entities": []}, + "": { + "source": "gmail", + "entities": [{"type": "headers", "from": None, "subject": "s", "date_raw": "d"}], + }, + }, + mail_chunks_by_source={"gmail": [("", 0, "mail text", 0.2)]}, + ) + session = _FakeSession() + result = await run_search( + conn, session, "http://fake-ollama", "q", "hybrid", "bge-m3", "claude-haiku-4-5" + ) + assert result["mode"] == "hybrid" + envelope_ids = [r["envelope_id"] for r in result["results"]] + assert envelope_ids == ["", "paperless:1"] + async def test_gmail_hit_carries_header_metadata_not_a_link(self): conn = _FakeConn( summaries=[("", 0.1)],