"""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]