arc53--docsgpt/application/retriever/dispatcher.py

"""Per-source retrieval dispatcher (D4).

Groups a per-source list by its ``config.retrieval.retriever`` key, builds one
retriever instance per group, and merges their results under a single shared
token budget so no group can starve another.

Parity guarantee: when every source is ``classic``/``default`` (the case for
every existing source) all sources flow into ONE ``ClassicRAG`` instance built
exactly as today, so the output — including token-budget behaviour — is
byte-identical to the pre-dispatch path.
"""

from __future__ import annotations

import logging
from typing import Any, Callable, Dict, List, Optional

from application.core.settings import settings
from application.retriever.base import BaseRetriever
from application.retriever.retriever_creator import RetrieverCreator
from application.retriever.stages.prescreen import (
    build_prescreen_stages,
    max_candidate_k,
)
from application.storage.db.source_config import RetrievalConfig
from application.utils import num_tokens_from_string

logger = logging.getLogger(__name__)

# Retriever keys that share ClassicRAG's single-instance / shared-budget model.
# Grouping all of them into one instance is what preserves byte-identical parity
# for the all-classic case.
_CLASSIC_KEYS = frozenset({"classic", "default"})

# Retriever keys that subclass ClassicRAG and accept ``defer_rephrase``, so the
# eager ctor rephrase can be skipped and computed lazily per-source. Unknown
# keys are excluded: a future non-ClassicRAG retriever may not accept the kwarg.
_DEFERRABLE_KEYS = _CLASSIC_KEYS | {"hybrid"}

# Fields that, when left at their defaults, mean the source did not opt into any
# per-source retrieval override — so it can flow through the global ClassicRAG
# path unchanged (byte-identical parity).
_DEFAULT_RETRIEVAL = RetrievalConfig()

# A post-retrieval stage: takes the candidate docs a group produced (plus the
# resolved query and context) and returns a possibly-filtered/reordered list.
# Left a no-op seam for the later prescreen/rerank stages (F1) to bolt onto.
Stage = Callable[[List[Dict[str, Any]], Dict[str, Any]], List[Dict[str, Any]]]


class Dispatcher(BaseRetriever):
    """Route per-source retrieval to grouped retrievers under a shared budget."""

    def __init__(
        self,
        source,
        chat_history=None,
        prompt="",
        chunks=2,
        doc_token_limit=50000,
        model_id="docsgpt-local",
        user_api_key=None,
        agent_id=None,
        llm_name=settings.LLM_PROVIDER,
        api_key=settings.API_KEY,
        decoded_token=None,
        model_user_id=None,
        request_id=None,
        sources: Optional[List[Dict[str, Any]]] = None,
        stages: Optional[List[Stage]] = None,
    ):
        """Build the dispatcher.

        Args:
            source: ClassicRAG-style source dict (``{"active_docs": [...]}``)
                plus the original ``question``. Used as the fallback group when
                no per-source ``sources`` list is supplied.
            chunks: Global default top-k, used when a source carries no
                per-source ``chunks`` hint.
            doc_token_limit: Hard cap shared across all groups.
            sources: Per-source list; each entry is ``{"id": str, "retrieval":
                RetrievalConfig | dict | None}``. When empty the dispatcher
                falls back to the single classic group over ``source``.
            stages: Optional post-retrieval stages applied to each group's
                candidates before final budgeting. Default: none (pass-through).
        """
        self._ctor_kwargs = dict(
            chat_history=chat_history,
            prompt=prompt,
            doc_token_limit=doc_token_limit,
            model_id=model_id,
            user_api_key=user_api_key,
            agent_id=agent_id,
            llm_name=llm_name,
            api_key=api_key,
            decoded_token=decoded_token,
            model_user_id=model_user_id,
            request_id=request_id,
        )
        self.source = source or {}
        self.original_question = self.source.get("question", "")
        self.chunks = chunks
        self.doc_token_limit = doc_token_limit
        self.stages: List[Stage] = stages or []
        self._sources = sources or []
        self._groups = self._build_groups()

    def _build_groups(self) -> List[Dict[str, Any]]:
        """Group the per-source list by retriever key.

        All classic/default sources collapse into one group (exact parity);
        each non-classic retriever key gets its own group. With no per-source
        list, a single classic group over ``self.source`` is produced.
        """
        active_docs = self.source.get("active_docs")
        if isinstance(active_docs, str):
            active_docs = [active_docs]
        active_docs = active_docs or []

        if not self._sources:
            return [
                {
                    "retriever": "classic",
                    "doc_ids": list(active_docs),
                    "retrievals": {},
                }
            ]

        grouped: Dict[str, Dict[str, Any]] = {}
        for entry in self._sources:
            doc_id = entry.get("id")
            if not doc_id:
                continue
            retrieval = self._coerce_retrieval(entry.get("retrieval"))
            key = (retrieval.retriever or "classic").lower()
            if key in _CLASSIC_KEYS:
                key = "classic"
            group = grouped.setdefault(
                key, {"retriever": key, "doc_ids": [], "retrievals": {}}
            )
            if doc_id not in group["doc_ids"]:
                group["doc_ids"].append(doc_id)
            # Only record an override when the source opted into a non-default
            # retrieval config; a default classic source stays on the global
            # ClassicRAG path so all-classic output is byte-identical to today.
            if self._is_override(retrieval):
                group["retrievals"][doc_id] = retrieval

        if not grouped:
            return [
                {"retriever": "classic", "doc_ids": list(active_docs), "retrievals": {}}
            ]

        # Defend the parity guarantee: any active_doc the per-source list omitted
        # would otherwise be dropped from every group. Route the strays through
        # the classic group with default retrieval config so they're still
        # retrieved (creating the classic group if no source opted into it).
        grouped_ids = {
            doc_id for group in grouped.values() for doc_id in group["doc_ids"]
        }
        missing = [doc_id for doc_id in active_docs if doc_id not in grouped_ids]
        if missing:
            classic_group = grouped.setdefault(
                "classic", {"retriever": "classic", "doc_ids": [], "retrievals": {}}
            )
            for doc_id in missing:
                if doc_id not in classic_group["doc_ids"]:
                    classic_group["doc_ids"].append(doc_id)
        return list(grouped.values())

    @staticmethod
    def _is_override(retrieval: RetrievalConfig) -> bool:
        """True if the source opted into any read-path retrieval override.

        Compares the read-path knobs ClassicRAG acts on (chunks /
        score_threshold / rephrase_query) plus an opted-in prescreen config; a
        source left at defaults takes the global path so all-classic retrieval
        stays byte-identical with zero extra LLM calls.
        """
        return (
            retrieval.chunks != _DEFAULT_RETRIEVAL.chunks
            or retrieval.score_threshold != _DEFAULT_RETRIEVAL.score_threshold
            or retrieval.rephrase_query != _DEFAULT_RETRIEVAL.rephrase_query
            or retrieval.prescreen is not None
        )

    @staticmethod
    def _coerce_retrieval(raw: Any) -> RetrievalConfig:
        """Coerce a per-source ``retrieval`` value to a ``RetrievalConfig``."""
        if isinstance(raw, RetrievalConfig):
            return raw
        if isinstance(raw, dict):
            try:
                return RetrievalConfig.model_validate(raw)
            except Exception:
                return RetrievalConfig()
        return RetrievalConfig()

    def _budget_for_group(self, n_groups: int, group_idx: int) -> int:
        """Split the shared token budget across groups.

        With one group the full ``doc_token_limit`` is returned, so the single
        ClassicRAG instance reproduces today's budget exactly. With multiple
        groups the budget is divided evenly (remainder to the first groups) so
        the total never exceeds ``doc_token_limit``.
        """
        if n_groups <= 1:
            return self.doc_token_limit
        base = self.doc_token_limit // n_groups
        remainder = self.doc_token_limit % n_groups
        return base + (1 if group_idx < remainder else 0)

    def _build_group_retriever(self, group: Dict[str, Any], budget: int):
        """Build the retriever for ``group`` with its budget and source list."""
        retriever_key = group["retriever"]
        group_source = dict(self.source)
        group_source["active_docs"] = group["doc_ids"]
        kwargs = dict(self._ctor_kwargs)
        kwargs["doc_token_limit"] = budget
        kwargs["source"] = group_source
        # Prescreen fetches a larger candidate set, then a stage trims it; raise
        # the effective top-k to candidate_k so the base retriever fetches
        # enough for the stage to filter down to the final chunks.
        candidate_k = max_candidate_k(group["retrievals"])
        kwargs["chunks"] = max(self.chunks, candidate_k or 0)
        # With per-source configs the rephrase decision is per-source, so defer
        # the eager rephrase to let a rephrase_query=False source skip the call.
        if group["retrievals"] and retriever_key in _DEFERRABLE_KEYS:
            kwargs["defer_rephrase"] = True
        retriever = RetrieverCreator.create_retriever(retriever_key, **kwargs)
        # Hand the per-source retrieval configs to the classic retriever so it
        # can honour per-source chunks/score_threshold/rephrase in its loop.
        if group["retrievals"]:
            setattr(retriever, "per_source_retrieval", group["retrievals"])
        return retriever

    def _group_stages(self, group: Dict[str, Any]) -> List[Stage]:
        """Stages for a group: caller-supplied stages + prescreen stages.

        Prescreen stages are built from the group's per-source prescreen config,
        so a group with no opted-in source adds nothing — the default path stays
        a pure no-op with zero extra LLM calls.
        """
        prescreen = build_prescreen_stages(
            group["retrievals"],
            llm_name=self._ctor_kwargs.get("llm_name"),
            api_key=self._ctor_kwargs.get("api_key"),
            model_id=self._ctor_kwargs.get("model_id"),
            user_api_key=self._ctor_kwargs.get("user_api_key"),
            decoded_token=self._ctor_kwargs.get("decoded_token"),
            agent_id=self._ctor_kwargs.get("agent_id"),
            model_user_id=self._ctor_kwargs.get("model_user_id"),
            request_id=self._ctor_kwargs.get("request_id"),
        )
        return list(self.stages) + prescreen

    def _run_stages(
        self,
        docs: List[Dict[str, Any]],
        context: Dict[str, Any],
        stages: List[Stage],
    ) -> List[Dict[str, Any]]:
        """Apply the given post-retrieval stages in order (no-op when empty)."""
        for stage in stages:
            try:
                docs = stage(docs, context)
            except Exception as exc:
                logger.warning("Retrieval stage failed, skipping: %s", exc)
        return docs

    def search(self, query: str = "") -> List[Dict[str, Any]]:
        """Run every group under the shared budget and merge the results."""
        groups = self._groups
        n_groups = len(groups)

        # Fast path / exact parity: a single group is just the underlying
        # retriever with the full budget — no merge accounting at all.
        if n_groups == 1:
            retriever = self._build_group_retriever(
                groups[0], self.doc_token_limit
            )
            docs = retriever.search(query) if query else retriever.search()
            context = {"query": query, "retriever": groups[0]["retriever"]}
            return self._run_stages(docs, context, self._group_stages(groups[0]))

        merged: List[Dict[str, Any]] = []
        cap = max(int(self.doc_token_limit * 0.9), 100)
        cumulative_tokens = 0
        for idx, group in enumerate(groups):
            budget = self._budget_for_group(n_groups, idx)
            retriever = self._build_group_retriever(group, budget)
            try:
                group_docs = retriever.search(query) if query else retriever.search()
            except Exception as exc:
                # Log the exception type only — a raw exception message can
                # carry the vector-store DSN (with credentials) on connection
                # errors.
                logger.error(
                    "Group '%s' search failed: %s",
                    group["retriever"],
                    type(exc).__name__,
                )
                continue
            context = {"query": query, "retriever": group["retriever"]}
            group_docs = self._run_stages(
                group_docs, context, self._group_stages(group)
            )
            for doc in group_docs:
                if cumulative_tokens >= cap:
                    break
                header = f"{doc.get('filename', '')}\n{doc.get('text', '')}"
                doc_tokens = num_tokens_from_string(header)
                if cumulative_tokens + doc_tokens < cap:
                    merged.append(doc)
                    cumulative_tokens += doc_tokens
            if cumulative_tokens >= cap:
                break
        return merged


def build_dispatcher(create_classic: Callable[[], BaseRetriever], **kwargs):
    """Build a Dispatcher, or fall back to the legacy single retriever.

    Honours ``settings.PER_SOURCE_RETRIEVAL_ENABLED``: when False, returns the
    legacy single ClassicRAG built by ``create_classic`` (the kill-switch).

    Args:
        create_classic: Zero-arg factory returning the legacy single retriever.
        **kwargs: Dispatcher constructor kwargs (including ``sources``).

    Returns:
        A ``Dispatcher`` or the legacy retriever from ``create_classic``.
    """
    if not getattr(settings, "PER_SOURCE_RETRIEVAL_ENABLED", True):
        return create_classic()
    return Dispatcher(**kwargs)