"""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"