homelab-codex-ws/services/kb-query/tests/test_search.py
oskar 8af8c20fc6 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).
2026-07-22 19:01:21 +02:00

237 lines
9.8 KiB
Python

"""Unit tests for /search's core logic (app.search.run_search) -- no real DB, no real Ollama.
Same mocking style as packages/kb-retrieval/tests/test_retrieval.py, extended with an
`envelope` table fixture for the join app/db.py adds on top of kb_retrieval."""
from __future__ import annotations
import pathlib
import sys
sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[1]))
from app.search import run_search # noqa: E402
class _FakeConn:
"""summaries: [(envelope_id, dist), ...] -- cascade stage-1 pre-filter. chunks_by_envelope:
envelope_id -> [(chunk_index, text, dist), ...]. envelopes: envelope_id -> {"source": ...,
"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,
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:
(envelope_ids, _model) = params
return [
{"envelope_id": eid, "summary": s["summary"], "tags": s["tags"]}
for eid, s in self._summary_texts.items()
if eid in envelope_ids
]
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 = [
{"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
_, 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]
if "FROM envelope" in query:
(envelope_ids,) = params
return [
{"id": eid, "source": self._envelopes[eid]["source"], "entities": self._envelopes[eid]["entities"]}
for eid in envelope_ids
if eid in self._envelopes
]
raise AssertionError(f"unexpected query: {query}")
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 TestRunSearchHappyPath:
async def test_cascade_hit_joins_envelope_and_shapes_paperless_link(self):
conn = _FakeConn(
summaries=[("paperless:119", 0.1)],
chunks_by_envelope={"paperless:119": [(2, "hit text", 0.34)]},
envelopes={"paperless:119": {"source": "paperless", "entities": []}},
)
session = _FakeSession()
result = await run_search(
conn, session, "http://fake-ollama", "polisa PZU", "cascade", "bge-m3", "claude-haiku-4-5"
)
assert result["query"] == "polisa PZU"
assert result["mode"] == "cascade"
assert result["sol_status"] == "up"
assert len(result["results"]) == 1
hit = result["results"][0]
assert hit["envelope_id"] == "paperless:119"
assert hit["source"] == "paperless"
assert hit["dist"] == 0.34
assert hit["chunk_index"] == 2
assert hit["text"] == "hit text"
assert hit["link"] == "https://paper.kapala.org/documents/119/details"
async def test_flat_mode_skips_cascade_stage1(self):
conn = _FakeConn(
chunks_by_envelope={"paperless:1": [(0, "a", 0.2)]},
envelopes={"paperless:1": {"source": "paperless", "entities": []}},
)
session = _FakeSession()
result = await run_search(
conn, session, "http://fake-ollama", "q", "flat", "bge-m3", "claude-haiku-4-5"
)
assert result["mode"] == "flat"
assert len(result["results"]) == 1
async def test_summary_attached_when_document_summary_row_exists(self):
conn = _FakeConn(
summaries=[("paperless:119", 0.1)],
chunks_by_envelope={"paperless:119": [(2, "hit text", 0.34)]},
envelopes={"paperless:119": {"source": "paperless", "entities": []}},
summary_texts={"paperless:119": {"summary": "Polisa OC 2024", "tags": ["ubezpieczenia"]}},
)
session = _FakeSession()
result = await run_search(
conn, session, "http://fake-ollama", "polisa PZU", "cascade", "bge-m3", "claude-haiku-4-5"
)
hit = result["results"][0]
assert hit["summary"] == "Polisa OC 2024"
assert hit["summary_tags"] == ["ubezpieczenia"]
async def test_summary_defaults_to_none_when_no_document_summary_row(self):
conn = _FakeConn(
summaries=[("paperless:119", 0.1)],
chunks_by_envelope={"paperless:119": [(2, "hit text", 0.34)]},
envelopes={"paperless:119": {"source": "paperless", "entities": []}},
)
session = _FakeSession()
result = await run_search(
conn, session, "http://fake-ollama", "polisa PZU", "cascade", "bge-m3", "claude-haiku-4-5"
)
hit = result["results"][0]
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)],
chunks_by_envelope={"<msgid@example.com>": [(0, "body text", 0.4)]},
envelopes={
"<msgid@example.com>": {
"source": "gmail",
"entities": [
{"type": "headers", "from": {"name": "A", "address": "a@b.com"}, "subject": "s", "date_raw": "d"}
],
}
},
)
session = _FakeSession()
result = await run_search(
conn, session, "http://fake-ollama", "q", "cascade", "bge-m3", "claude-haiku-4-5"
)
hit = result["results"][0]
assert hit["source"] == "gmail"
assert hit["subject"] == "s"
assert hit["link"] is None
assert hit["mail_ui_url"] is None
class TestRunSearchNoGoodResults:
async def test_results_above_no_answer_threshold_are_still_returned_unfiltered(self):
# Plan §7: the 0.55 "no answer" colour threshold is a frontend concern -- the API
# must not silently drop/hide a poor match, only report its true dist so the caller
# (UI or eval harness) can apply that policy itself.
conn = _FakeConn(
summaries=[("paperless:1", 0.6)],
chunks_by_envelope={"paperless:1": [(0, "unrelated text", 0.62)]},
envelopes={"paperless:1": {"source": "paperless", "entities": []}},
)
session = _FakeSession()
result = await run_search(
conn, session, "http://fake-ollama", "unrelated query", "cascade", "bge-m3", "claude-haiku-4-5"
)
assert len(result["results"]) == 1
assert result["results"][0]["dist"] == 0.62
async def test_no_summaries_yields_empty_results_not_an_error(self):
conn = _FakeConn(summaries=[], chunks_by_envelope={}, envelopes={})
session = _FakeSession()
result = await run_search(
conn, session, "http://fake-ollama", "nothing matches", "cascade", "bge-m3", "claude-haiku-4-5"
)
assert result["results"] == []