homelab-codex-ws/services/kb-query/tests/test_search.py

"""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), ...]. chunks_by_envelope: envelope_id -> [(chunk_index,
    text, dist), ...]. envelopes: envelope_id -> {"source": ..., "entities": [...]}."""

    def __init__(self, summaries=None, chunks_by_envelope=None, envelopes=None):
        self._summaries = list(summaries or [])
        self._chunks_by_envelope = chunks_by_envelope or {}
        self._envelopes = envelopes or {}

    async def fetch(self, 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
            _, 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_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"] == []