src/federated_retriever.py

"""Search several legal corpora behind one retriever interface.

The design constraint that shaped this module: `LegalAnswerOrchestrator` and
`AnswerComposer` already contain the citation-verification logic this assistant
depends on, and they talk to exactly one object with a `query()` method.  So the
federation preserves that method's signature and returns the same hit
dictionaries.  Nothing upstream has to know that there is more than one corpus.

Score fusion is a plain merge by rank_score, which is only defensible because
every corpus is indexed with the same embedding model and distance metric, and
the optional cross-encoder scores the query against candidates from all corpora
with the same model.  `CorpusRegistry.assert_consistent_embedding()` guards that
precondition.

Fusion by score alone turned out not to be enough, and the counterexample is
worth keeping: a question about Sonderkennzeichen never reached the
Arzneimittelabrechnungsvereinbarung, although the rule is in it and the corpus
was indexed.  With 106 chunks against the SGB V's 6.670, its best hit simply lost
the similarity contest — a structural defeat that has nothing to do with
relevance.  Two mechanisms answer that, and both are supplied from outside so
this module stays a search layer and does not become a legal register:

* a **per-corpus quota** that survives the global cut, so a small corpus cannot
  be crowded out entirely, and
* **mandatory fetches** (`anchors`), addresses of norms that are retrieved by
  metadata regardless of how they embed.

A third mechanism follows the same shape and closes a different gap.  Anchors
know the chain for a *question type*; they are curated and therefore finite.
But a legal text names its own chain: § 10 of the Rahmenvertrag says the
selection is to be made "nach Maßgabe der §§ 11 bis 14", and § 11 selects "aus
dem Auswahlbereich nach § 9".  Those references were read by nobody — the
retriever parses § references out of the *question* only.  `verweise` supplies
the second pass: the references contained in the hits that were already found,
resolved to addresses and fetched with the same `get_units`.  The provider is
injected for the same reason as the anchors: which name means which corpus is a
statement about the corpora, not about searching.
"""

from __future__ import annotations

import inspect
import logging
from collections import OrderedDict
from typing import Any, Callable, Dict, List, Optional, Sequence

from corpus_registry import CorpusRegistry

logger = logging.getLogger(__name__)

# A router decides which corpora a question should reach.
Router = Callable[[str, Sequence[str]], List[str]]

# An anchor provider names norms that must be in the context whatever the
# ranking says.  It returns plain dicts — corpus_id, container_id, section_id,
# subsection, max_chunks — so that the curated register stays a stranger here.
Anchors = Callable[[str], Sequence[Dict[str, Any]]]

# A reference provider reads the hits that were found and returns the addresses
# they point at — same plain dicts as the anchors, so the fetch below does not
# have to know which of the two produced a target.
Verweise = Callable[[Sequence[Dict[str, Any]], Sequence[str]], Sequence[Dict[str, Any]]]

# Referenced norms are fetched, not ranked, so they need a score — and the
# first value chosen was wrong in a way worth recording.
#
# 0.97 was the reasoning: below the anchors' 1.0, above a merely similar
# passage.  Measured on 19.08.2026 over five chain cases it cost ten points of
# chain coverage (85% -> 75%).  The budget is fixed, so a referenced norm that
# ranks high does not *add* to the context, it *replaces* something — and what
# it replaced was § 6 Rahmenvertrag, itself a member of the chain being
# measured.  A norm the text merely points at is not worth more than a norm the
# ranking actually found.
#
# 0.50 puts them in the tail: they fill slots that would otherwise go to the
# weakest semantic hits, and they yield to everything above.
VERWEIS_SCORE = 0.50


def route_to_all(question: str, corpus_ids: Sequence[str]) -> List[str]:
    """Default policy: ask every corpus and let fusion sort it out."""
    return list(corpus_ids)


def _accepted_kwargs(func: Any, kwargs: Dict[str, Any]) -> Dict[str, Any]:
    """Drop kwargs the callable does not declare, keeping the rest intact."""
    try:
        params = inspect.signature(func).parameters
    except (TypeError, ValueError):  # builtins / C-implemented callables
        return dict(kwargs)
    if any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values()):
        return dict(kwargs)
    return {key: value for key, value in kwargs.items() if key in params}


class FederatedRetriever:
    """Drop-in replacement for `LegalRetriever` that spans several corpora."""

    def __init__(
        self,
        registry: CorpusRegistry,
        *,
        router: Optional[Router] = None,
        max_corpora_per_query: int = 4,
        anchors: Optional[Anchors] = None,
        min_hits_per_corpus: int = 2,
        verweise: Optional[Verweise] = None,
        max_verweis_normen: int = 3,
        named_corpora: Optional[Router] = None,
    ) -> None:
        self.registry = registry
        self.router = router or route_to_all
        self.max_corpora_per_query = max_corpora_per_query
        self.anchors = anchors
        self.min_hits_per_corpus = min_hits_per_corpus
        self.verweise = verweise
        self.max_verweis_normen = max_verweis_normen
        self.named_corpora = named_corpora

        warning = registry.assert_consistent_embedding()
        if warning:
            # Not fatal: a single-corpus deployment is unaffected, and refusing to
            # start would take the whole assistant down for a ranking problem.
            logger.warning("federated retrieval may rank badly: %s", warning)

    # ------------------------------------------------------------------
    # Compatibility surface
    # ------------------------------------------------------------------

    @property
    def col(self) -> Any:
        """Chroma collection of the primary corpus.

        Kept so legacy diagnostics that reach for `.col` keep working; corpus-aware
        callers should use `registry.diagnostics()` instead.
        """
        return self.registry.retriever(self.registry.primary_id).col

    @property
    def collection_name(self) -> str:
        return self.registry.spec(self.registry.primary_id).collection

    @property
    def persist_dir(self) -> str:
        return self.registry.persist_dir

    def diagnostics(self, *, sample: int = 3) -> Dict[str, Any]:
        out = self.registry.diagnostics()
        out["federated"] = True
        out["router"] = getattr(self.router, "__name__", type(self.router).__name__)
        out["min_hits_per_corpus"] = self.min_hits_per_corpus
        out["norm_anchors"] = self.anchors is not None
        return out

    def explain_selection(self, question: str) -> Dict[str, Any]:
        """Why these corpora — for /debug/routing and for tests.

        Kept separate from `corpus_router.explain_routing` because the router
        only sees the question's wording.  Whether a corpus is queried on top of
        that is this layer's decision, and a debug endpoint that showed only half
        of it would be misleading in exactly the case worth debugging.
        """
        available = self.registry.available()
        named = [c for c in self.router(question or "", available) if c in available]
        required = [c for c in self.required_corpora(question) if c in available]
        return {
            "available": available,
            "routed": named,
            "required": required,
            "selected": self.select_corpora(question),
            "anchor_targets": self.anchor_targets(question),
        }

    # ------------------------------------------------------------------
    # Retrieval
    # ------------------------------------------------------------------

    def anchor_targets(self, question: str) -> List[Dict[str, Any]]:
        """Norms that have to be in the context, whatever the ranking says."""
        if self.anchors is None:
            return []
        try:
            return [dict(target) for target in self.anchors(question or "") or []]
        except Exception as exc:  # noqa: BLE001 - a register must never break retrieval.
            logger.warning("norm anchors skipped (%s)", exc)
            return []

    def required_corpora(self, question: str) -> List[str]:
        """Corpora an anchored question must reach even without naming them.

        This is the half of the Sonderkennzeichen fix that routing owns: the word
        names no corpus, so `route_question` fans out to all of them and the
        smallest one is then outranked.  Naming it here makes the query
        unconditional; the quota below makes sure its hits survive the cut.
        """
        out: List[str] = []
        for target in self.anchor_targets(question):
            corpus_id = str(target.get("corpus_id") or "")
            if corpus_id and corpus_id not in out:
                out.append(corpus_id)
        return out

    def select_corpora(self, question: str) -> List[str]:
        available = self.registry.available()
        if not available:
            return []
        selected = [c for c in self.router(question or "", available) if c in available]
        if not selected:
            selected = available

        # Required corpora go first so the cap below can never drop one, and the
        # cap yields to them outright rather than silently discarding a norm the
        # register calls decisive.
        required = [c for c in self.required_corpora(question) if c in available]
        ordered = required + [c for c in selected if c not in required]
        return ordered[: max(self.max_corpora_per_query, len(required))]

    def query(self, question: str = "", **kwargs: Any) -> List[Dict[str, Any]]:
        # The orchestrator sometimes passes the question positionally, sometimes
        # as a keyword; accept both like LegalRetriever does.
        question = question or str(kwargs.pop("question", "") or "")
        corpora = self.select_corpora(question)
        if not corpora:
            return []

        max_final_results = kwargs.get("max_final_results")
        top_k = kwargs.get("top_k")

        merged: List[Dict[str, Any]] = []
        for corpus_id in corpora:
            try:
                hits = self._query_one(corpus_id, question, kwargs)
            except Exception as exc:  # noqa: BLE001 - one broken corpus must not kill the answer.
                # Der Typ gehört in die Meldung: ohne ihn steht dort eine
                # Fehlerzeile, die nicht sagt, ob das Korpus fehlte, der
                # Speicher ausging oder die Datenbank gesperrt war — und genau
                # das ist die Frage, wenn diese Zeile einmal erscheint.
                logger.warning(
                    "corpus query failed: %s (%s: %s)", corpus_id, exc.__class__.__name__, exc
                )
                continue
            merged.extend(hits)

        merged.extend(self._fetch_anchors(question, corpora))

        merged = _fuse(merged)

        # Erst nach der Fusion: die Verweise werden aus den *besten* Treffern
        # gelesen, und welche das sind, steht vor der Fusion noch nicht fest.
        nachgeladen = self._fetch_referenced_norms(merged, corpora)
        if nachgeladen:
            merged = _fuse(merged + nachgeladen)

        limit = max_final_results or top_k
        ergebnis = _apply_corpus_quota(merged, limit=limit, min_per_corpus=self.min_hits_per_corpus)
        return self._benanntes_korpus_zuerst(ergebnis, question, corpora)

    def _fetch_anchors(self, question: str, corpora: Sequence[str]) -> List[Dict[str, Any]]:
        """Retrieve the anchored norms by metadata instead of by similarity.

        These hits are scored 1.0, the same as any explicit section lookup the
        retriever performs on its own — the register's claim is that they decide
        the question, so ranking them below a merely similar passage would defeat
        the purpose.  They stay cheap because the address carries the Absatz: the
        parent chunk of one Absatz is usually the whole norm.
        """
        out: List[Dict[str, Any]] = []
        for target in self.anchor_targets(question):
            corpus_id = str(target.get("corpus_id") or "")
            if corpus_id not in corpora:
                continue
            retriever = self.registry.retriever(corpus_id)
            fetch = getattr(retriever, "get_units", None)
            if fetch is None:
                # An older retriever without the targeted lookup. Skipping is the
                # honest outcome: the rest of the answer is unaffected.
                continue
            try:
                hits = fetch(
                    target.get("section_id") or "",
                    container=target.get("container_id") or None,
                    subsection=target.get("subsection") or None,
                    max_chunks=int(target.get("max_chunks") or 2),
                )
            except Exception as exc:  # noqa: BLE001 - see above.
                logger.warning(
                    "norm anchor lookup failed: %s %s (%s: %s)",
                    corpus_id,
                    target,
                    exc.__class__.__name__,
                    exc,
                )
                continue

            if not hits:
                logger.warning(
                    "norm anchor matched nothing: %s %s %s",
                    corpus_id,
                    target.get("container_id"),
                    target.get("section_id"),
                )
            for hit in hits:
                hit["corpus_id"] = corpus_id
                metadata = hit.get("metadata")
                if isinstance(metadata, dict):
                    metadata.setdefault("corpus_id", corpus_id)
                # "section_lookup" keeps the composer's existing preference for
                # explicit lookups working; "norm_anchor" is what monitoring reads.
                hit["retrieval_kinds"] = sorted(
                    {*(hit.get("retrieval_kinds") or []), "norm_anchor", "section_lookup"}
                )
                out.append(hit)
        return out

    def _benanntes_korpus_zuerst(
        self, hits: List[Dict[str, Any]], question: str, corpora: Sequence[str]
    ) -> List[Dict[str, Any]]:
        """Nennt die Frage ein Regelwerk, führt dessen bester Treffer.

        Der Pflichtabruf holt seine Normen mit Score 1.0 und steht damit vorn —
        richtig, solange die Frage kein Regelwerk nennt. Nennt sie eines, ist es
        falsch: „Was regelt **die Arzneimittel-Richtlinie** zur Austauschbarkeit
        von Darreichungsformen?" bekam § 9 des Rahmenvertrags auf Rang 1, weil
        der Anker `darreichungsform_austausch` ihn als tragende Norm führt. Die
        Auskunft ist nicht falsch, die Reihenfolge schon: gefragt war, was in
        der Richtlinie steht.

        Nur die **Reihenfolge** ändert sich, nichts fällt weg — der Anker bleibt
        im Kontext, er führt ihn nur nicht mehr an. Der Score des vorgezogenen
        Treffers wird auf den bisherigen Höchstwert gehoben, damit eine spätere
        Sortierung nach `rank_score` die Entscheidung nicht wieder umdreht.
        """
        if self.named_corpora is None or len(hits) < 2:
            return hits

        try:
            benannt = [c for c in self.named_corpora(question or "", corpora) if c in corpora]
        except Exception as exc:  # noqa: BLE001 - eine Reihung darf nie die Suche kippen.
            logger.warning("named-corpus ordering skipped (%s: %s)", exc.__class__.__name__, exc)
            return hits

        # Ohne Nennung oder wenn die Frage alle abgefragten Korpora nennt, gibt
        # es nichts zu entscheiden.
        if not benannt or len(benannt) >= len(corpora):
            return hits
        if str(hits[0].get("corpus_id") or "") in benannt:
            return hits

        for i, hit in enumerate(hits):
            if str(hit.get("corpus_id") or "") in benannt:
                hoechstwert = max(_rank(h) for h in hits)
                hit["rank_score"] = max(_rank(hit), hoechstwert)
                return [hit, *hits[:i], *hits[i + 1 :]]

        # Das benannte Korpus hat nichts geliefert. Das ist eine Aussage über
        # das Korpus und keine über die Reihenfolge — hier bleibt alles stehen.
        return hits

    def _fetch_referenced_norms(
        self, hits: Sequence[Dict[str, Any]], corpora: Sequence[str]
    ) -> List[Dict[str, Any]]:
        """Die Normen nachladen, auf die die gefundenen Texte selbst verweisen.

        Dieselbe Mechanik wie `_fetch_anchors`, aber die Adressen kommen nicht
        aus einem Register, sondern aus dem Text.  Der Unterschied steht im
        Score: ein kuratierter Anker behauptet, *diese* Norm entscheide die
        Frage; ein Verweis behauptet nur, der gefundene Text nehme sie in Bezug.

        Der Container ist Pflicht, wo das Korpus einen führt.  Ohne ihn holt
        „§ 10" im Rahmenvertrag auch den § 10 der Anlage 11 — die Adresse wäre
        dann mehrdeutig, und zwar still.
        """
        if self.verweise is None or not hits or self.max_verweis_normen <= 0:
            return []

        try:
            ziele = list(self.verweise(hits, corpora) or [])
        except Exception as exc:  # noqa: BLE001 - a parser must never break retrieval.
            logger.warning("norm references skipped (%s: %s)", exc.__class__.__name__, exc)
            return []

        out: List[Dict[str, Any]] = []
        for ziel in ziele[: self.max_verweis_normen]:
            corpus_id = str(ziel.get("corpus_id") or "")
            if corpus_id not in corpora:
                continue
            retriever = self.registry.retriever(corpus_id)
            fetch = getattr(retriever, "get_units", None)
            if fetch is None:
                continue

            container = ziel.get("container_id") or self.registry.spec(corpus_id).default_container_id
            try:
                treffer = fetch(
                    ziel.get("section_id") or "",
                    container=container or None,
                    subsection=ziel.get("subsection") or None,
                    max_chunks=int(ziel.get("max_chunks") or 1),
                )
            except Exception as exc:  # noqa: BLE001 - see above.
                logger.warning(
                    "norm reference lookup failed: %s %s (%s: %s)",
                    corpus_id,
                    ziel.get("section_id"),
                    exc.__class__.__name__,
                    exc,
                )
                continue

            # Ein Verweis, der ins Leere zeigt, ist keine Warnung wert: der Text
            # darf auf eine Vorschrift verweisen, die es im Bestand nicht gibt
            # (aufgehobener Paragraph, andere Fassung). Für den Anker gilt das
            # Gegenteil — dort ist die Adresse eine Behauptung.
            for hit in treffer:
                hit["corpus_id"] = corpus_id
                hit["score"] = VERWEIS_SCORE
                hit["rank_score"] = VERWEIS_SCORE
                metadata = hit.get("metadata")
                if isinstance(metadata, dict):
                    metadata.setdefault("corpus_id", corpus_id)
                # Gesetzt, nicht ergänzt: `get_units` schreibt voreingestellt
                # "norm_anchor" hinein, und das liest die Überwachung als
                # kuratierten Pflichtabruf. Ein Verweis ist etwas anderes.
                #
                # Und bewusst OHNE "section_lookup": der Composer wertet diese
                # Art in `_hit_relevance_key` mit Faktor 50 auf (expliziter
                # Paragraphentreffer). Ein Verweis ist kein expliziter Treffer —
                # niemand hat nach dieser Vorschrift gefragt.
                hit["retrieval_kinds"] = ["verweis"]
                out.append(hit)

        return out

    def _query_one(self, corpus_id: str, question: str, kwargs: Dict[str, Any]) -> List[Dict[str, Any]]:
        retriever = self.registry.retriever(corpus_id)
        spec = self.registry.spec(corpus_id)

        call_kwargs = dict(kwargs)
        call_kwargs.pop("question", None)

        # A container filter is corpus-specific: "Vertrag" does not exist in a
        # statute, and passing it through would return nothing at all.
        where = call_kwargs.get("where")
        if isinstance(where, dict) and "container_id" in where:
            if not spec.default_container_id or where["container_id"] != spec.default_container_id:
                call_kwargs["where"] = None

        hits = self._call(retriever, question, call_kwargs)
        for hit in hits:
            hit["corpus_id"] = corpus_id
            metadata = hit.get("metadata")
            if isinstance(metadata, dict):
                metadata.setdefault("corpus_id", corpus_id)
        return hits

    @staticmethod
    def _call(retriever: Any, question: str, kwargs: Dict[str, Any]) -> List[Dict[str, Any]]:
        # Only the kwargs this retriever actually declares.  The orchestrator
        # probes for optional features (e.g. `enable_definition_lookup`) and
        # relies on a TypeError to retry without them — but that TypeError never
        # reaches it through the federation.  Filtering here keeps the valuable
        # kwargs (explicit_sections, include_neighbors, max_final_results)
        # instead of collapsing to the three-parameter emergency fallback below.
        accepted = _accepted_kwargs(retriever.query, kwargs)
        try:
            return list(retriever.query(question=question, **accepted) or [])
        except TypeError:
            reduced = {k: v for k, v in accepted.items() if k in {"top_k", "where", "fetch_k"}}
            try:
                return list(retriever.query(question=question, **reduced) or [])
            except TypeError:
                return list(retriever.query(question, top_k=kwargs.get("top_k", 8)) or [])

    def verify_negative_result(self, question: str, **kwargs: Any) -> Dict[str, Any]:
        """A negative answer is only safe when no corpus can refute it."""
        corpora = self.select_corpora(question)
        results: List[Dict[str, Any]] = []
        blocking: List[str] = []
        reasons: List[str] = []

        for corpus_id in corpora:
            retriever = self.registry.retriever(corpus_id)
            try:
                check = retriever.verify_negative_result(question, **kwargs)
            except TypeError:
                check = retriever.verify_negative_result(question=question)
            except Exception as exc:  # noqa: BLE001
                logger.warning("negative recheck failed for %s (%s)", corpus_id, exc)
                continue

            for hit in check.get("results") or []:
                hit["corpus_id"] = corpus_id
                results.append(hit)
            if not check.get("safe_to_answer_negative", True):
                blocking.append(corpus_id)
                reasons.append(f"{corpus_id}: {check.get('reason', '')}".strip())

        return {
            "safe_to_answer_negative": not blocking,
            "reason": " | ".join(reasons)
            or "Auch der breite Kontrollabruf hat in keinem Korpus belastbare Treffer gefunden.",
            "results": _fuse(results),
            "strong_result_count": len(blocking),
            "corpora_checked": corpora,
            "corpora_with_findings": blocking,
        }


def _rank(hit: Dict[str, Any]) -> float:
    try:
        return float(hit.get("rank_score", hit.get("score", 0.0)) or 0.0)
    except (TypeError, ValueError):
        return 0.0


def _key(hit: Dict[str, Any]) -> tuple:
    metadata = hit.get("metadata") or {}
    text_hash = metadata.get("text_hash")
    if text_hash:
        return ("hash", hit.get("corpus_id"), text_hash)
    return (
        "location",
        hit.get("corpus_id"),
        metadata.get("container_id") or hit.get("container"),
        metadata.get("section_id") or hit.get("section"),
        metadata.get("chunk_index_in_section") or hit.get("chunk_index"),
        (hit.get("text") or "")[:160],
    )


def _fuse(hits: Sequence[Dict[str, Any]]) -> List[Dict[str, Any]]:
    """Merge hits from several corpora into one ranked list."""
    best: Dict[tuple, Dict[str, Any]] = {}
    for hit in hits:
        key = _key(hit)
        current = best.get(key)
        if current is None:
            best[key] = hit
            continue
        # The same chunk can arrive twice — once from the semantic pass, once
        # from a mandatory fetch.  Keeping only the better score would drop the
        # other's retrieval_kind, and downstream ranking reads exactly that.
        kinds = sorted({*(current.get("retrieval_kinds") or []), *(hit.get("retrieval_kinds") or [])})
        winner = hit if _rank(hit) > _rank(current) else current
        winner["retrieval_kinds"] = kinds
        best[key] = winner
    return sorted(best.values(), key=_rank, reverse=True)


def _apply_corpus_quota(
    hits: List[Dict[str, Any]],
    *,
    limit: Optional[int],
    min_per_corpus: int,
    min_ratio: float = 0.25,
) -> List[Dict[str, Any]]:
    """Cut to `limit`, but not before every corpus has had its share.

    Without this the cut is purely global, and a corpus that is an order of
    magnitude smaller than its neighbours never appears at all — not because it
    has nothing to say but because it has fewer chances to say it.

    Three limits keep the guarantee from turning into noise:

    * The reservation is capped at **half** the result list.  A floor against
      structural defeat is not a mandate to fill the context with four documents
      when the question belongs to one — with four corpora and ten results that
      is one slot each and six left to the ranking.
    * Only **whole rounds** are handed out.  The corpus that would lose the last
      slot of an incomplete round is always the weakest one, which is the one
      this exists to protect.
    * A corpus qualifies only if its best hit reaches `min_ratio` of the best hit
      overall.  Measured on the five live questions, that is the difference
      between rescuing the Abrechnungsvereinbarung's 0.850 — genuinely
      comparable, merely outnumbered — and admitting a 0.008 from a corpus that
      simply has nothing to say.  Having fewer chances is a defeat worth undoing;
      being irrelevant is not.
    """
    if not limit or limit <= 0:
        return hits
    if min_per_corpus <= 0 or len(hits) <= limit:
        return hits[:limit]

    # `hits` is ranked, so insertion order is the order of each corpus's best hit.
    by_corpus: "OrderedDict[str, List[Dict[str, Any]]]" = OrderedDict()
    for hit in hits:
        by_corpus.setdefault(str(hit.get("corpus_id") or ""), []).append(hit)
    if len(by_corpus) <= 1:
        return hits[:limit]

    floor = _rank(hits[0]) * min_ratio
    by_corpus = OrderedDict(
        (corpus_id, bucket)
        for corpus_id, bucket in by_corpus.items()
        if _rank(bucket[0]) >= floor
    )
    if len(by_corpus) <= 1:
        return hits[:limit]

    budget = max(1, limit // 2)
    reserved: List[Dict[str, Any]] = []
    chosen: set = set()

    for round_index in range(min_per_corpus):
        runde = [b[round_index] for b in by_corpus.values() if round_index < len(b)]
        if not runde or len(reserved) + len(runde) > budget:
            break
        for hit in runde:
            reserved.append(hit)
            chosen.add(id(hit))

    for hit in hits:
        if len(reserved) >= limit:
            break
        if id(hit) not in chosen:
            reserved.append(hit)
            chosen.add(id(hit))

    return sorted(reserved, key=_rank, reverse=True)[:limit]