arc53--docsgpt/tests/test_dispatcher.py

"""Tests for the per-source retrieval Dispatcher (B1a) and the kill-switch."""

from unittest.mock import Mock, patch

import pytest

from application.retriever.dispatcher import Dispatcher, build_dispatcher
from application.storage.db.source_config import RetrievalConfig


@pytest.fixture
def _patch_llm_creator(mock_llm, monkeypatch):
    monkeypatch.setattr(
        "application.retriever.classic_rag.LLMCreator.create_llm",
        Mock(return_value=mock_llm),
    )
    return mock_llm


def _make_doc(page_content, title="t", source="s"):
    doc = Mock()
    doc.page_content = page_content
    doc.metadata = {"title": title, "source": source}
    return doc


@pytest.mark.unit
class TestDispatcherGrouping:
    def test_no_sources_single_classic_group(self, _patch_llm_creator):
        d = Dispatcher(source={"question": "q", "active_docs": ["a", "b"]})
        groups = d._groups
        assert len(groups) == 1
        assert groups[0]["retriever"] == "classic"
        assert groups[0]["doc_ids"] == ["a", "b"]
        assert groups[0]["retrievals"] == {}

    def test_all_classic_collapse_to_one_group(self, _patch_llm_creator):
        sources = [
            {"id": "a", "retrieval": RetrievalConfig()},
            {"id": "b", "retrieval": RetrievalConfig()},
        ]
        d = Dispatcher(source={"question": "q", "active_docs": ["a", "b"]}, sources=sources)
        assert len(d._groups) == 1
        assert d._groups[0]["retriever"] == "classic"
        # All-default sources record no override → byte-identical global path.
        assert d._groups[0]["retrievals"] == {}

    def test_default_and_alias_share_one_group(self, _patch_llm_creator):
        sources = [
            {"id": "a", "retrieval": RetrievalConfig(retriever="classic")},
            {"id": "b", "retrieval": RetrievalConfig(retriever="default")},
        ]
        d = Dispatcher(source={"question": "q", "active_docs": ["a", "b"]}, sources=sources)
        assert len(d._groups) == 1
        assert d._groups[0]["doc_ids"] == ["a", "b"]

    def test_non_classic_gets_own_group(self, _patch_llm_creator):
        sources = [
            {"id": "a", "retrieval": RetrievalConfig(retriever="classic")},
            {"id": "b", "retrieval": RetrievalConfig(retriever="graphrag")},
        ]
        d = Dispatcher(source={"question": "q", "active_docs": ["a", "b"]}, sources=sources)
        keys = sorted(g["retriever"] for g in d._groups)
        assert keys == ["classic", "graphrag"]

    def test_override_recorded_only_when_non_default(self, _patch_llm_creator):
        sources = [
            {"id": "a", "retrieval": RetrievalConfig(chunks=7)},
            {"id": "b", "retrieval": RetrievalConfig()},
        ]
        d = Dispatcher(source={"question": "q", "active_docs": ["a", "b"]}, sources=sources)
        retrievals = d._groups[0]["retrievals"]
        assert "a" in retrievals
        assert "b" not in retrievals


@pytest.mark.unit
class TestDispatcherSharedBudget:
    def test_single_group_full_budget(self, _patch_llm_creator):
        d = Dispatcher(source={"question": "q", "active_docs": ["a"]}, doc_token_limit=5000)
        assert d._budget_for_group(1, 0) == 5000

    def test_multi_group_budget_split_never_exceeds_total(self, _patch_llm_creator):
        d = Dispatcher(source={"question": "q", "active_docs": ["a"]}, doc_token_limit=1000)
        budgets = [d._budget_for_group(3, i) for i in range(3)]
        assert sum(budgets) == 1000
        # Remainder goes to the first groups.
        assert budgets == [334, 333, 333]


@pytest.mark.unit
class TestDispatcherParity:
    """All-classic sources through the Dispatcher == one ClassicRAG today."""

    @patch("application.retriever.classic_rag.VectorCreator")
    @patch("application.retriever.classic_rag.num_tokens_from_string", return_value=10)
    def test_single_group_matches_classic_rag(self, _tok, mock_vc, _patch_llm_creator):
        from application.retriever.classic_rag import ClassicRAG

        docsearch = Mock()
        docsearch.search.return_value = [_make_doc("content one"), _make_doc("content two")]
        mock_vc.create_vectorstore.return_value = docsearch

        common = dict(
            source={"question": "q", "active_docs": ["a", "b"]},
            chat_history=[],
            chunks=2,
            doc_token_limit=50000,
            model_id="m",
            decoded_token={"sub": "u"},
        )

        baseline = ClassicRAG(**common).search("query")
        dispatched = Dispatcher(
            sources=[
                {"id": "a", "retrieval": RetrievalConfig()},
                {"id": "b", "retrieval": RetrievalConfig()},
            ],
            **common,
        ).search("query")

        assert dispatched == baseline

    @patch("application.retriever.classic_rag.VectorCreator")
    @patch("application.retriever.classic_rag.num_tokens_from_string", return_value=10)
    def test_no_sources_matches_classic_rag(self, _tok, mock_vc, _patch_llm_creator):
        from application.retriever.classic_rag import ClassicRAG

        docsearch = Mock()
        docsearch.search.return_value = [_make_doc("a")]
        mock_vc.create_vectorstore.return_value = docsearch

        common = dict(
            source={"question": "q", "active_docs": ["a"]},
            chat_history=[],
            chunks=2,
            doc_token_limit=50000,
            decoded_token={"sub": "u"},
        )
        baseline = ClassicRAG(**common).search("query")
        dispatched = Dispatcher(**common).search("query")
        assert dispatched == baseline


@pytest.mark.unit
class TestDispatcherStageSeam:
    @patch("application.retriever.classic_rag.VectorCreator")
    @patch("application.retriever.classic_rag.num_tokens_from_string", return_value=10)
    def test_stage_applied_to_candidates(self, _tok, mock_vc, _patch_llm_creator):
        docsearch = Mock()
        docsearch.search.return_value = [_make_doc("keep"), _make_doc("drop")]
        mock_vc.create_vectorstore.return_value = docsearch

        def drop_stage(docs, context):
            return [d for d in docs if d["text"] == "keep"]

        d = Dispatcher(
            source={"question": "q", "active_docs": ["a"]},
            stages=[drop_stage],
        )
        out = d.search("query")
        assert [doc["text"] for doc in out] == ["keep"]

    @patch("application.retriever.classic_rag.VectorCreator")
    @patch("application.retriever.classic_rag.num_tokens_from_string", return_value=10)
    def test_default_stages_passthrough(self, _tok, mock_vc, _patch_llm_creator):
        docsearch = Mock()
        docsearch.search.return_value = [_make_doc("a")]
        mock_vc.create_vectorstore.return_value = docsearch
        d = Dispatcher(source={"question": "q", "active_docs": ["a"]})
        assert len(d.search("query")) == 1


@pytest.mark.unit
class TestDispatcherLenientRead:
    """D7: a garbage/legacy per-source config still retrieves via classic."""

    def test_coerce_garbage_falls_back_to_default(self):
        assert Dispatcher._coerce_retrieval("not-a-dict") == RetrievalConfig()
        assert Dispatcher._coerce_retrieval(None) == RetrievalConfig()
        # An invalid dict that fails validation also falls back.
        assert Dispatcher._coerce_retrieval({"chunks": "abc"}) == RetrievalConfig()

    @patch("application.retriever.classic_rag.VectorCreator")
    @patch("application.retriever.classic_rag.num_tokens_from_string", return_value=10)
    def test_garbage_config_retrieves_via_classic(self, _tok, mock_vc, _patch_llm_creator):
        docsearch = Mock()
        docsearch.search.return_value = [_make_doc("ok")]
        mock_vc.create_vectorstore.return_value = docsearch

        d = Dispatcher(
            source={"question": "q", "active_docs": ["a"]},
            sources=[{"id": "a", "retrieval": {"bogus": True, "chunks": "x"}}],
        )
        out = d.search("query")
        assert [doc["text"] for doc in out] == ["ok"]
        # Falls back to the global classic path (no override recorded).
        assert d._groups[0]["retriever"] == "classic"
        assert d._groups[0]["retrievals"] == {}


@pytest.mark.unit
class TestDispatcherPrescreen:
    """F1: prescreen bumps candidate_k, trims to max_keep, off == today."""

    @patch("application.retriever.classic_rag.VectorCreator")
    @patch("application.retriever.classic_rag.num_tokens_from_string", return_value=10)
    def test_candidate_k_fetched_and_trimmed(self, _tok, mock_vc, _patch_llm_creator):
        docsearch = Mock()
        # Return 40 candidate docs; prescreen should trim to max_keep=3.
        docsearch.search.return_value = [
            _make_doc(f"c{i}") for i in range(40)
        ]
        mock_vc.create_vectorstore.return_value = docsearch

        prescreen_llm = Mock()
        # Keep the first index of each batch.
        prescreen_llm.gen = Mock(return_value='{"keep": [0]}')
        prescreen_llm.model_id = "m"

        with patch(
            "application.retriever.stages.prescreen.LLMCreator.create_llm",
            return_value=prescreen_llm,
        ):
            d = Dispatcher(
                source={"question": "q", "active_docs": ["a"]},
                doc_token_limit=500000,
                sources=[
                    {
                        "id": "a",
                        "retrieval": {
                            "chunks": 2,
                            "prescreen": {
                                "candidate_k": 40,
                                "batch_size": 10,
                                "max_keep": 3,
                            },
                        },
                    }
                ],
            )
            out = d.search("query")

        # Base retriever asked for >= candidate_k candidates.
        _, kwargs = docsearch.search.call_args
        assert kwargs["k"] >= 40
        # Prescreen ran (4 batches of 10) and trimmed to max_keep.
        assert prescreen_llm.gen.call_count == 4
        assert len(out) == 3

    @patch("application.retriever.stages.prescreen.build_prescreen_stages")
    @patch("application.retriever.classic_rag.VectorCreator")
    @patch("application.retriever.classic_rag.num_tokens_from_string", return_value=10)
    def test_prescreen_none_no_extra_llm_calls(
        self, _tok, mock_vc, mock_build, _patch_llm_creator
    ):
        docsearch = Mock()
        docsearch.search.return_value = [_make_doc("one"), _make_doc("two")]
        mock_vc.create_vectorstore.return_value = docsearch
        # The dispatcher imports the symbol; patch where it's looked up.
        import application.retriever.dispatcher as disp

        with patch.object(disp, "build_prescreen_stages", mock_build):
            mock_build.return_value = []
            d = Dispatcher(
                source={"question": "q", "active_docs": ["a"]},
                sources=[{"id": "a", "retrieval": RetrievalConfig()}],
            )
            out = d.search("query")

        # A default (non-prescreen) source records no override, so the prescreen
        # stage builder is invoked with an empty retrievals map and yields no
        # stages — i.e. zero extra screening calls — and output is unchanged.
        for call in mock_build.call_args_list:
            assert call.args[0] == {}
        assert [doc["text"] for doc in out] == ["one", "two"]

    def test_prescreen_only_source_records_override(self, _patch_llm_creator):
        d = Dispatcher(
            source={"question": "q", "active_docs": ["a"]},
            sources=[
                {
                    "id": "a",
                    "retrieval": {
                        "chunks": 2,
                        "prescreen": {"candidate_k": 20, "max_keep": 5},
                    },
                }
            ],
        )
        # A source that opts into prescreen only must still record an override
        # so the stage actually fires.
        assert "a" in d._groups[0]["retrievals"]


@pytest.mark.unit
class TestKillSwitch:
    def test_disabled_falls_back_to_legacy(self, monkeypatch):
        monkeypatch.setattr(
            "application.retriever.dispatcher.settings.PER_SOURCE_RETRIEVAL_ENABLED",
            False,
        )
        sentinel = object()
        result = build_dispatcher(
            lambda: sentinel,
            source={"question": "q", "active_docs": ["a"]},
            sources=[{"id": "a", "retrieval": RetrievalConfig(chunks=9)}],
        )
        assert result is sentinel

    def test_enabled_returns_dispatcher(self, monkeypatch, _patch_llm_creator):
        monkeypatch.setattr(
            "application.retriever.dispatcher.settings.PER_SOURCE_RETRIEVAL_ENABLED",
            True,
        )
        result = build_dispatcher(
            lambda: object(),
            source={"question": "q", "active_docs": ["a"]},
            sources=[],
        )
        assert isinstance(result, Dispatcher)