feat(kb-retrieval,kb-query): add hybrid retrieval mode (faza mailowa Krok 3)
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).
This commit is contained in:
oskar
parent
ad0ef408f9
commit
a640cf1455
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"])
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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": []},
|
||||
"<msgid@example.com>": {
|
||||
"source": "gmail",
|
||||
"entities": [{"type": "headers", "from": None, "subject": "s", "date_raw": "d"}],
|
||||
},
|
||||
},
|
||||
mail_chunks_by_source={"gmail": [("<msgid@example.com>", 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 == ["<msgid@example.com>", "paperless:1"]
|
||||
|
||||
async def test_gmail_hit_carries_header_metadata_not_a_link(self):
|
||||
conn = _FakeConn(
|
||||
summaries=[("<msgid@example.com>", 0.1)],
|
||||
|
|
|
|||
Loading…
Reference in a new issue