src/corpus_registry.py

"""Registry of the legal corpora the assistant can search.

Each corpus is one Chroma collection inside the shared persist directory and
gets its own `LegalRetriever`, because the retrieval defaults differ per document
type: the Rahmenvertrag has a main contract container that explicit § lookups
should fall back to, a statute does not.

Retrievers are built lazily. A cold start therefore still only pays for the
collections a request actually touches, and a corpus that is missing from the
persist directory degrades to a clear error on first use instead of breaking
application startup.
"""

from __future__ import annotations

import logging
import os
from dataclasses import dataclass, field
from threading import Lock
from typing import Any, Dict, List, Optional, Sequence

from retriever import LegalRetriever

logger = logging.getLogger(__name__)


def _env_flag(name: str, default: bool) -> bool:
    raw = os.getenv(name)
    if raw is None:
        return default
    return raw.strip().lower() in {"1", "true", "yes", "ja", "on"}


@dataclass(frozen=True)
class CorpusSpec:
    """Static description of one searchable corpus."""

    corpus_id: str
    collection: str
    # Container the retriever falls back to for explicit § lookups.  Empty for a
    # statute, whose §§ are spread across Kapitel containers.
    default_container_id: str = ""
    label: str = ""
    # Lower-case keywords that name this corpus in a question.  Used by routing.
    aliases: tuple[str, ...] = ()

    def with_metadata(self, metadata: Dict[str, Any]) -> "CorpusSpec":
        """Fill in the label from the collection metadata written at ingest time."""
        if self.label:
            return self
        label = str(metadata.get("doc_title") or metadata.get("doc_id") or self.corpus_id)
        return CorpusSpec(
            corpus_id=self.corpus_id,
            collection=self.collection,
            default_container_id=self.default_container_id,
            label=label,
            aliases=self.aliases,
        )


# Aliases for the corpora this assistant ships with.  Unknown corpora still work;
# they are simply routed to by name only.
CORPUS_ALIASES: Dict[str, tuple[str, ...]] = {
    "rv129": ("rahmenvertrag", "rahmenvertrages", "rahmenvertrag nach § 129", "rv129", "rv 129"),
    "sgb5": ("sgb v", "sgb 5", "sgb5", "sozialgesetzbuch", "fünftes buch", "fuenftes buch"),
    # Die Korpus-Id ist bewusst "amabrv", nicht "abrechnung": `mentioned_corpora`
    # nimmt die Id selbst als Alias, und "Abrechnung" ist ein Alltagswort, das in
    # Rahmenvertrag und SGB V laufend vorkommt — als Alias würde es jede zweite
    # Frage auf dieses eine Korpus verengen, statt sie an alle drei zu geben.
    # Hier stehen deshalb nur Formulierungen, die das Dokument wirklich benennen.
    "amabrv": (
        "arzneimittelabrechnungsvereinbarung",
        "abrechnungsvereinbarung",
        "vereinbarung nach § 300",
        "vereinbarung nach §  300",
        "§ 300 abs. 3",
        "§ 300 absatz 3",
    ),
    # Dieselbe Regel wie bei amabrv: Aliase müssen das Dokument benennen, nicht
    # sein Thema.  "arzneimittel" oder "richtlinie" allein stehen deshalb nicht
    # hier — beide kommen in Rahmenvertrag und SGB V laufend vor und würden die
    # Frage auf dieses eine Korpus verengen.  Indiziert ist der Richtlinientext;
    # die Anlagen (Wirkstofflisten) sind bewusst nicht enthalten, worauf
    # `corpus_boundary` in der Antwort hinweist.
    "amrl": (
        "arzneimittel-richtlinie",
        "arzneimittelrichtlinie",
        "arzneimittel richtlinie",
        "am-rl",
        "am rl",
        "richtlinie nach § 92",
        "richtlinie nach §  92",
        "§ 92 abs. 1 satz 2 nr. 6",
        "§ 92 absatz 1 satz 2 nummer 6",
    ),
}


def specs_from_env() -> List[CorpusSpec]:
    """Read the corpus list from the environment.

    `CHROMA_COLLECTIONS` is a comma-separated list; when it is absent the single
    `CHROMA_COLLECTION` is used, so an existing deployment keeps working without
    any configuration change.

    `DEFAULT_CONTAINER_ID` applies to the first corpus only — it is the
    Rahmenvertrag's "Vertrag" and would wrongly restrict a statute.  Per-corpus
    overrides use `CORPUS_CONTAINER__<corpus_id>`.
    """
    raw = os.getenv("CHROMA_COLLECTIONS", "")
    collections = [c.strip() for c in raw.split(",") if c.strip()]
    if not collections:
        collections = [os.getenv("CHROMA_COLLECTION", "rv129").strip()]

    primary_container = os.getenv("DEFAULT_CONTAINER_ID", "Vertrag")

    specs: List[CorpusSpec] = []
    for index, collection in enumerate(collections):
        corpus_id = collection
        container = os.getenv(
            f"CORPUS_CONTAINER__{corpus_id}",
            primary_container if index == 0 else "",
        )
        specs.append(
            CorpusSpec(
                corpus_id=corpus_id,
                collection=collection,
                default_container_id=container.strip(),
                aliases=CORPUS_ALIASES.get(corpus_id.lower(), ()),
            )
        )
    return specs


@dataclass
class CorpusRegistry:
    """Lazily built `LegalRetriever` per corpus."""

    specs: List[CorpusSpec]
    persist_dir: str
    model_name: str = "auto"
    enable_reranker: bool = True
    reranker_model: str = "cross-encoder/mmarco-mMiniLMv2-L12-H384-v1"
    reranker_candidates: int = 20

    _retrievers: Dict[str, LegalRetriever] = field(default_factory=dict, init=False, repr=False)
    _lock: Lock = field(default_factory=Lock, init=False, repr=False)

    @classmethod
    def from_env(cls, *, persist_dir: str, **overrides: Any) -> "CorpusRegistry":
        return cls(specs=specs_from_env(), persist_dir=persist_dir, **overrides)

    # ------------------------------------------------------------------
    # Access
    # ------------------------------------------------------------------

    @property
    def corpus_ids(self) -> List[str]:
        return [spec.corpus_id for spec in self.specs]

    @property
    def primary_id(self) -> str:
        return self.specs[0].corpus_id if self.specs else ""

    def spec(self, corpus_id: str) -> CorpusSpec:
        for spec in self.specs:
            if spec.corpus_id == corpus_id:
                return spec
        raise KeyError(f"unknown corpus: {corpus_id!r} (known: {self.corpus_ids})")

    def retriever(self, corpus_id: str) -> LegalRetriever:
        existing = self._retrievers.get(corpus_id)
        if existing is not None:
            return existing

        with self._lock:
            if corpus_id in self._retrievers:
                return self._retrievers[corpus_id]

            spec = self.spec(corpus_id)
            instance = self._build(spec)
            self._retrievers[corpus_id] = instance

            # Adopt the document title the ingestion wrote into the collection
            # metadata, so labels never drift from the indexed corpus.
            metadata = dict(getattr(instance.col, "metadata", None) or {})
            enriched = spec.with_metadata(metadata)
            if enriched is not spec:
                self.specs = [enriched if s.corpus_id == corpus_id else s for s in self.specs]

            return instance

    def _build(self, spec: CorpusSpec) -> LegalRetriever:
        kwargs: Dict[str, Any] = {
            "persist_dir": self.persist_dir,
            "collection": spec.collection,
            "model_name": self.model_name,
            "default_container_id": spec.default_container_id,
            "enable_reranker": self.enable_reranker,
            "reranker_model": self.reranker_model,
            "reranker_candidates": self.reranker_candidates,
        }
        try:
            return LegalRetriever(**kwargs)
        except TypeError:
            # Older retriever signature.
            return LegalRetriever(
                persist_dir=self.persist_dir,
                collection=spec.collection,
                model_name=self.model_name,
            )

    def available(self) -> List[str]:
        """Corpus ids whose collection can actually be opened."""
        out: List[str] = []
        for corpus_id in self.corpus_ids:
            try:
                self.retriever(corpus_id)
                out.append(corpus_id)
            except Exception as exc:  # noqa: BLE001 - a missing corpus must not break the rest.
                logger.warning("corpus unavailable: %s (%s)", corpus_id, exc)
        return out

    def diagnostics(self) -> Dict[str, Any]:
        corpora: List[Dict[str, Any]] = []
        for spec in self.specs:
            entry: Dict[str, Any] = {
                "corpus_id": spec.corpus_id,
                "collection": spec.collection,
                "default_container_id": spec.default_container_id,
                "label": spec.label,
            }
            try:
                instance = self.retriever(spec.corpus_id)
                entry["count"] = instance.col.count()
                entry["ok"] = True
                metadata = dict(getattr(instance.col, "metadata", None) or {})
                entry["embedding_model"] = metadata.get("embedding_model")
                entry["doc_title"] = metadata.get("doc_title")
            except Exception as exc:  # noqa: BLE001
                entry["ok"] = False
                entry["error"] = f"{type(exc).__name__}: {exc}"
            corpora.append(entry)
        return {"persist_dir": self.persist_dir, "corpora": corpora}

    def total_count(self) -> int:
        total = 0
        for corpus_id in self.corpus_ids:
            try:
                total += int(self.retriever(corpus_id).col.count())
            except Exception:  # noqa: BLE001
                continue
        return total

    def assert_consistent_embedding(self) -> Optional[str]:
        """Return an error message if the corpora were not indexed alike.

        Fusing results across corpora only makes sense when their scores are
        comparable, which requires one embedding model and one distance metric.
        Mixing them would silently rank one corpus above the other.
        """
        models: Dict[str, set] = {"embedding_model": set(), "embedding_dim": set()}
        for corpus_id in self.corpus_ids:
            try:
                metadata = dict(getattr(self.retriever(corpus_id).col, "metadata", None) or {})
            except Exception:  # noqa: BLE001
                continue
            for key in models:
                value = metadata.get(key)
                if value is not None:
                    models[key].add(value)

        mismatched = {key: sorted(map(str, values)) for key, values in models.items() if len(values) > 1}
        if mismatched:
            return f"corpora were indexed with different embeddings: {mismatched}"
        return None


def build_registry(
    *,
    persist_dir: str,
    model_name: str,
    enable_reranker: bool,
    reranker_model: str,
    reranker_candidates: int,
    specs: Optional[Sequence[CorpusSpec]] = None,
) -> CorpusRegistry:
    return CorpusRegistry(
        specs=list(specs) if specs is not None else specs_from_env(),
        persist_dir=persist_dir,
        model_name=model_name,
        enable_reranker=enable_reranker,
        reranker_model=reranker_model,
        reranker_candidates=reranker_candidates,
    )