"""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": []}, "": { "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)], chunks_by_envelope={"": [(0, "body text", 0.4)]}, envelopes={ "": { "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"] == []