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), ...] -- 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"] == []