src/backend/models.py

"""Registry of retrieval models the UI can offer.

One place to declare every embedding model, what it can search, and whether this
deployment actually serves it. Adding or swapping a model is an edit to `MODELS`
plus its vectors and weights on disk — no change to the API surface or the frontend.

The registry includes known models that are not deployed. The catalog returns
their availability and reason, while the picker lists models that can serve the
selected search. Unsupported requests return an error.
"""

from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path

from . import paths

CODE_CATEGORIES = ("diagnosis", "medication", "lab", "procedure")
PHENO_CATEGORY = ("phenotype",)


@dataclass(frozen=True)
class ModelSpec:
    id: str
    label: str
    categories: tuple[str, ...]
    # Prebuilt phenotype vectors for this model (None -> no phenotype support).
    pheno_emb_dir: Path | None = None
    # True when the packaged FAISS code indexes were built with this model.
    code_indexes: bool = False
    # False when the vectors and weights are known but this deployment cannot run it.
    served: bool = True
    unavailable_reason: str = ""

    def supports(self, category: str) -> bool:
        return category in self.categories

    def available(self) -> tuple[bool, str]:
        """(is_available, reason_if_not) — declaration plus what is on disk."""
        if not self.served:
            return False, self.unavailable_reason or "Not deployed on this backend."
        if self.pheno_emb_dir is not None:
            config_path = self.pheno_emb_dir / "config.json"
            if not config_path.exists():
                return False, f"Phenotype vectors missing ({self.pheno_emb_dir.name})."
            try:
                config = json.loads(config_path.read_text())
            except (OSError, json.JSONDecodeError) as err:
                return False, f"Invalid phenotype vector config: {err}"
            manifest_id = config.get("model_id")
            if manifest_id is not None and manifest_id != self.id:
                return False, (
                    f"Phenotype vectors identify model '{manifest_id}', not '{self.id}'.")
        return True, ""


# Overridable so a deployment can point at a different vector build without a
# code change (the ENCODE_PARTA_EMB_DIR escape hatch, applied in paths.py).
_FTVA_EMB = paths.PARTA_EMB_DIR

MODELS: tuple[ModelSpec, ...] = (
    ModelSpec(
        id="bge_ft_va",
        label="BGE-FT-VA",
        categories=PHENO_CATEGORY + CODE_CATEGORIES,
        pheno_emb_dir=_FTVA_EMB,
        code_indexes=True,
    ),
)

DEFAULT_MODEL_ID = MODELS[0].id

_BY_ID = {m.id: m for m in MODELS}


class ModelError(Exception):
    """Raised with an HTTP status when a requested model cannot serve a request."""

    def __init__(self, status: int, detail: str):
        super().__init__(detail)
        self.status = status
        self.detail = detail


def resolve(model_id: str | None, category: str) -> ModelSpec:
    """The model to serve `category` with, or ModelError explaining why not."""
    if not model_id:
        model_id = DEFAULT_MODEL_ID
    spec = _BY_ID.get(model_id)
    if spec is None:
        known = ", ".join(_BY_ID)
        raise ModelError(400, f"Unknown model '{model_id}'. Available: {known}")
    if not spec.supports(category):
        raise ModelError(400, f"Model '{spec.label}' does not support {category} search.")
    if category in CODE_CATEGORIES and not spec.code_indexes:
        raise ModelError(503, f"No code indexes were built with '{spec.label}'.")
    if category in PHENO_CATEGORY and spec.pheno_emb_dir is None:
        raise ModelError(503, f"No phenotype vectors were built with '{spec.label}'.")
    ok, reason = spec.available()
    if not ok:
        raise ModelError(503, f"Model '{spec.label}' is not available: {reason}")
    return spec


def catalog() -> dict:
    """What the model picker renders."""
    entries = []
    for spec in MODELS:
        ok, reason = spec.available()
        entries.append({
            "id": spec.id,
            "label": spec.label,
            "categories": list(spec.categories),
            "available": ok,
            "unavailable_reason": reason,
        })
    return {"default": DEFAULT_MODEL_ID, "models": entries}