homelab-codex-ws/packages/kb-retrieval/tests/test_retrieval.py

"""Unit tests for the retrieval cascade -- no DB, no real HTTP/Ollama."""
from __future__ import annotations

import pytest

from kb_retrieval.retrieval import (
    DEFAULT_K,
    DEFAULT_N,
    cascade_query,
    cascade_retrieve,
    flat_query,
    flat_retrieve,
    hybrid_query,
    hybrid_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.

    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):
        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 "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: 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 (kb/phases/kb-m5-faza3.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 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)]})
        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"

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