Deploy OpenAI classification and protected terminology streaming

Dockerfile

@@ -0,0 +1,9 @@
+FROM python:3.12-slim
+WORKDIR /app
+COPY backend/requirements.txt /app/requirements.txt
+RUN pip install --no-cache-dir -r requirements.txt
+COPY src /app/src
+COPY backend /app/backend
+ENV PORT=7860
+EXPOSE 7860
+CMD ["sh", "-c", "exec uvicorn backend.app:app --host 0.0.0.0 --port ${PORT}"]

README.md

@@ -1,10 +1,9 @@
 ---
-title: Semantic Harmonization Api
-emoji: 🚀
-colorFrom: green
-colorTo: green
+title: Clinical Codelist API
 sdk: docker
-pinned: false
+app_port: 7860
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Token-protected classification and GPT review through OpenAI, plus streaming access to a private terminology dataset.
+
+Configure OPENAI_API_KEY, PIPELINE_API_TOKEN and HF_TOKEN as Space secrets. Configure PIPELINE_HF_DATASET and PIPELINE_ALLOWED_ORIGINS as variables. No terminology vectors are stored in the Space.

backend/app.py

@@ -0,0 +1,180 @@
+"""Small classification API for the static browser app. No embedding dependencies.
+
+Run: python -m uvicorn backend.app:app --host 127.0.0.1 --port 8000
+"""
+from __future__ import annotations
+
+from dataclasses import asdict, fields
+import hmac
+import json
+import logging
+import math
+import os
+import re
+from typing import Annotated, Literal
+
+import httpx
+from dotenv import load_dotenv
+from fastapi import Depends, FastAPI, Header, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import StreamingResponse
+from starlette.background import BackgroundTask
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+
+from src.data.loaders import EventDefinitionForm
+from src.llm.classify import llm_classify
+from src.llm.client import LLMCallError, default_model_for
+
+load_dotenv()
+logger = logging.getLogger(__name__)
+app = FastAPI(title="Semantic Harmonization classification API", docs_url=None, redoc_url=None)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=[origin.strip() for origin in os.getenv(
+        "PIPELINE_ALLOWED_ORIGINS", "http://127.0.0.1:5173,http://localhost:5173,https://ruurdkuiper.github.io"
+    ).split(",") if origin.strip()],
+    allow_methods=["GET", "POST"], allow_headers=["Authorization", "Content-Type"],
+)
+
+
+def authorize(authorization: Annotated[str | None, Header()] = None):
+    token = os.getenv("PIPELINE_API_TOKEN", "")
+    if not token:
+        raise HTTPException(503, "The classification service needs an access token configured by its administrator.")
+    supplied = (authorization or "").removeprefix("Bearer ")
+    if not hmac.compare_digest(supplied.encode(), token.encode()):
+        raise HTTPException(401, "Enter a valid classification service access token.")
+
+
+class Candidate(BaseModel):
+    model_config = ConfigDict(extra="forbid", str_strip_whitespace=True)
+    code: str = Field(min_length=1, max_length=128)
+    vocabulary: Literal["ICD10CM", "ICD9CM", "ICPC", "MDR", "RCD2", "SNOMEDCT_US"]
+    description: str = Field(min_length=1, max_length=4096)
+
+
+class Batch(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+    edf: dict
+    candidates: list[Candidate] = Field(min_length=1, max_length=20)
+    use_possible_category: bool = False
+
+    @model_validator(mode="after")
+    def validate_edf(self):
+        allowed = {field.name: field for field in fields(EventDefinitionForm) if field.name != "name"}
+        if not set(self.edf) <= set(allowed):
+            raise ValueError("Unknown EDF field")
+        if len(json.dumps(self.edf)) > 65_536:
+            raise ValueError("EDF is too large")
+        for key, value in self.edf.items():
+            # Accept cleared YAML fields, matching the browser's normalization.
+            is_list = key in {"inclusion_criteria", "exclusion_criteria", "narrow_decision_rules", "possible_decision_rules",
+                              "exclude_decision_rules", "synonyms", "candidate_anchor_codes", "ambiguities", "references"}
+            if value is None:
+                self.edf[key] = [] if is_list else ""
+            elif is_list and (not isinstance(value, list) or any(not isinstance(item, str) for item in value)):
+                raise ValueError(f"{key} must be a list of text values")
+            elif not is_list and not isinstance(value, str):
+                raise ValueError(f"{key} must be text")
+        if not self.edf.get("preferred_name", "").strip():
+            raise ValueError("preferred_name is required")
+        if len({(row.code, row.vocabulary) for row in self.candidates}) != len(self.candidates):
+            raise ValueError("Duplicate candidates in batch")
+        return self
+
+
+@app.get("/health")
+def health():
+    return {"status": "ok", "retrieval": "browser-minilm"}
+
+
+_ASSET_PATH = re.compile(
+    r"(?:manifest\.json|(?:ICD10CM|ICD9CM|ICPC|MDR|RCD2|SNOMEDCT_US)-[0-9a-f]{20}/"
+    r"(?:lexical\.json\.gz|[0-9]{7}\.(?:json\.gz|f32)))\Z"
+)
+
+
+@app.get("/api/assets/{asset_path:path}", dependencies=[Depends(authorize)])
+def terminology_asset(asset_path: str):
+    """Stream private HF dataset assets to authorized browsers without storing them."""
+    if not _ASSET_PATH.fullmatch(asset_path):
+        raise HTTPException(404, "Unknown terminology asset.")
+    repo = os.getenv("PIPELINE_HF_DATASET", "")
+    revision = os.getenv("PIPELINE_HF_REVISION", "main")
+    if not re.fullmatch(r"[A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+", repo) or not re.fullmatch(r"[A-Za-z0-9_-]+", revision):
+        raise HTTPException(503, "The Hugging Face terminology dataset is not configured.")
+    if not os.getenv("HF_TOKEN"):
+        raise HTTPException(503, "Private terminology downloads need a read-only HF_TOKEN configured in the Space secrets.")
+    headers = {"Accept-Encoding": "identity"}
+    if os.getenv("HF_TOKEN"):
+        headers["Authorization"] = "Bearer " + os.environ["HF_TOKEN"]
+    # HTTPX strips Authorization on cross-origin CDN redirects. Neither the HF
+    # token nor provider keys are forwarded to the browser.
+    client = httpx.Client(timeout=60, follow_redirects=True)
+    try:
+        response = client.send(client.build_request("GET",
+            f"https://huggingface.co/datasets/{repo}/resolve/{revision}/{asset_path}", headers=headers), stream=True)
+    except httpx.HTTPError:
+        client.close()
+        raise HTTPException(502, "The terminology download could not be reached. Please retry.") from None
+    if response.status_code != 200:
+        status = response.status_code
+        response.close()
+        client.close()
+        raise HTTPException(404 if status == 404 else 502,
+                            "The terminology asset is unavailable. The administrator should check the dataset access settings.")
+
+    def close():
+        response.close()
+        client.close()
+
+    def chunks():
+        try:
+            yield from response.iter_raw(chunk_size=65536)
+        finally:
+            close()
+
+    # Private terminology data must not be cached by shared HTTP proxies.
+    # The authorized browser maintains its own content-addressed IndexedDB cache.
+    return StreamingResponse(chunks(), media_type="application/json" if asset_path == "manifest.json" else "application/octet-stream",
+                             headers={"Cache-Control": "private, no-store", "X-Content-Type-Options": "nosniff"},
+                             background=BackgroundTask(close))
+
+
+def classify_batch(batch: Batch, review: bool):
+    provider = "openai" if review else os.getenv("PIPELINE_CLASSIFIER_PROVIDER", "openai")
+    if provider not in {"jev", "openai", "anthropic", "google"}:
+        raise HTTPException(503, "The classification provider is not configured correctly.")
+    model = os.getenv("PIPELINE_REVIEW_MODEL" if review else "PIPELINE_CLASSIFIER_MODEL") or default_model_for(provider)
+    try:
+        records = llm_classify(
+            batch.candidates, EventDefinitionForm(name="browser-edf", **batch.edf),
+            provider=provider, model=model, max_retries=3, batch_size=10 if review else 20,
+            use_possible_category=batch.use_possible_category, reasoning_effort="medium", adaptive_stopping=None,
+        )
+    except (LLMCallError, ValueError, RuntimeError):
+        logger.exception("Classification request failed")
+        # Provider responses may contain sensitive diagnostic text. Keep it server-side.
+        raise HTTPException(502, "The model request failed. Retry the batch or ask the service administrator to check its configuration.") from None
+    expected = {(row.code, row.vocabulary) for row in batch.candidates}
+    labels = {"Narrow", "Exclude", "Possible"} if batch.use_possible_category else {"Narrow", "Exclude"}
+    keys = [(row.code, row.vocabulary) for row in records]
+    if len(keys) != len(expected) or set(keys) != expected or any(
+        row.label not in labels or not math.isfinite(row.confidence) or not 0 <= row.confidence <= 1
+        or row.explanation == "No classification returned by LLM; routed for manual review."
+        for row in records
+    ):
+        raise HTTPException(502, "The model returned an incomplete or invalid batch. Retry to continue.")
+    return {"decisions": [asdict(row) for row in records], "provider": provider, "model": model}
+
+
+@app.post("/api/classify", dependencies=[Depends(authorize)])
+def classify(batch: Batch):
+    return classify_batch(batch, review=False)
+
+
+@app.post("/api/review", dependencies=[Depends(authorize)])
+def review(batch: Batch):
+    if len(batch.candidates) > 10:
+        raise HTTPException(422, "Review accepts at most ten candidates per batch.")
+    return classify_batch(batch, review=True)

backend/requirements.txt

@@ -0,0 +1,9 @@
+fastapi>=0.115,<1
+uvicorn>=0.30,<1
+openai>=1.50,<3
+anthropic>=0.40,<1
+google-genai>=1,<3
+pandas>=2,<4
+pyyaml>=6,<7
+python-dotenv>=1,<2
+httpx>=0.27,<1

src/data/code_systems.py

@@ -0,0 +1,253 @@
+"""Loaders for full clinical coding-system reference lists (ICD10, SNOMED CT,
+MedDRA, ...), used as the retrieval universe for concept retrieval.
+
+Unlike the phenotype-specific AESI review exports (`load_aesi_dataset`), these
+represent the *complete* set of codes in a vocabulary, independent of any
+phenotype. Since they can be large (tens of thousands of rows) and slow to
+load/parse from their native format, the processed corpus for each vocabulary
+is cached to disk as Parquet so the source file only needs to be parsed once.
+"""
+
+from __future__ import annotations
+
+import csv
+from pathlib import Path
+
+import pandas as pd
+
+
+DEFAULT_CODE_SYSTEM_PATHS = {
+    "ICD10CM": "data/codes/csv/ICD10CM@2026-codes.csv",
+    "ICD9CM": "data/codes/csv/ICD9CM@2014-codes.csv",
+    "ICPC": "data/codes/csv/ICPC@1993-codes.csv",
+    "MDR": "data/codes/csv/MDR@28.0-codes.csv",
+    "RCD2": "data/codes/csv/RCD2@20200401-codes.csv",
+    "SNOMEDCT_US": "data/codes/csv/SNOMEDCT_US@2025_09_01-codes definitive.csv",
+}
+
+
+def _load_codelist_csv(
+    path: str | Path,
+    *,
+    vocabulary: str,
+    delimiter: str = ",",
+) -> pd.DataFrame:
+    """Load a code list CSV into the canonical retrieval schema.
+
+    The source files in ``data/codes`` are plain text CSV exports, but the
+    description field may contain commas wrapped in quotes. We therefore rely
+    on the CSV parser rather than manual splitting so those commas stay inside
+    the description instead of being treated as extra columns.
+    """
+    path = Path(path)
+    if not path.exists():
+        raise FileNotFoundError(f"Full {vocabulary} code list not found: {path}")
+
+    raw = pd.read_csv(
+        path,
+        dtype=str,
+        sep=delimiter,
+        engine="python",
+        quotechar='"',
+        escapechar="\\",
+        on_bad_lines="skip",
+    )
+
+    if raw.empty:
+        raise ValueError(f"No rows loaded from {path}")
+
+    # Normalize the first two columns into code/description regardless of the
+    # exact header names used in the export.
+    columns = list(raw.columns)
+    if len(columns) < 2:
+        raise ValueError(f"Unexpected {vocabulary} CSV schema in {path}: columns={columns}")
+
+    code_col = columns[0]
+    desc_col = columns[1]
+    df = raw[[code_col, desc_col]].rename(columns={code_col: "code", desc_col: "description"})
+    df = df.dropna(subset=["code", "description"]).copy()
+    df["code"] = df["code"].astype(str).str.strip().str.strip('"')
+    df["description"] = df["description"].astype(str).str.strip().str.strip('"')
+    df["vocabulary"] = vocabulary
+    df = df[df["code"] != ""]
+    df = df[df["description"] != ""]
+    df = df.drop_duplicates(subset=["code", "vocabulary"], keep="first").reset_index(drop=True)
+    return df
+
+
+def _load_two_column_text(
+    path: str | Path,
+    *,
+    vocabulary: str,
+    delimiter: str,
+) -> pd.DataFrame:
+    """Load a very loose delimited export by splitting only on the first separator.
+
+    Some vendor exports include trailing separators or extra empty columns. This
+    parser keeps the first field as the code and the remainder of the line as the
+    description so embedded punctuation in the description is preserved.
+    """
+    path = Path(path)
+    if not path.exists():
+        raise FileNotFoundError(f"Full {vocabulary} code list not found: {path}")
+
+    rows: list[tuple[str, str]] = []
+    for line in path.read_text().splitlines():
+        line = line.strip()
+        if not line or line.lower().startswith("code"):
+            continue
+        if delimiter not in line:
+            continue
+        code, description = line.split(delimiter, 1)
+        code = code.strip().strip('"')
+        description = description.strip().strip(";").strip(delimiter).strip().strip('"')
+        if code and description:
+            rows.append((code, description))
+
+    if not rows:
+        raise ValueError(f"No rows loaded from {path}")
+
+    df = pd.DataFrame(rows, columns=["code", "description"])
+    df["vocabulary"] = vocabulary
+    df = df.drop_duplicates(subset=["code", "vocabulary"], keep="first").reset_index(drop=True)
+    return df
+
+
+def load_icd10cm_full(csv_path: str | Path = DEFAULT_CODE_SYSTEM_PATHS["ICD10CM"]) -> pd.DataFrame:
+    return _load_codelist_csv(csv_path, vocabulary="ICD10CM", delimiter=",")
+
+
+def load_icd9cm_full(csv_path: str | Path = DEFAULT_CODE_SYSTEM_PATHS["ICD9CM"]) -> pd.DataFrame:
+    return _load_codelist_csv(csv_path, vocabulary="ICD9CM", delimiter=",")
+
+
+def load_icpc_full(csv_path: str | Path = DEFAULT_CODE_SYSTEM_PATHS["ICPC"]) -> pd.DataFrame:
+    return _load_codelist_csv(csv_path, vocabulary="ICPC", delimiter=",")
+
+
+def load_mdr_full(csv_path: str | Path = DEFAULT_CODE_SYSTEM_PATHS["MDR"]) -> pd.DataFrame:
+    return _load_codelist_csv(csv_path, vocabulary="MDR", delimiter=",")
+
+
+def load_rcd2_full(csv_path: str | Path = DEFAULT_CODE_SYSTEM_PATHS["RCD2"]) -> pd.DataFrame:
+    return _load_two_column_text(csv_path, vocabulary="RCD2", delimiter=",")
+
+
+def load_snomedct_full(csv_path: str | Path = DEFAULT_CODE_SYSTEM_PATHS["SNOMEDCT_US"]) -> pd.DataFrame:
+    path = Path(csv_path)
+    if not path.exists():
+        raise FileNotFoundError(f"Full SNOMEDCT_US code list not found: {path}")
+
+    # Excel can wrap each original comma-delimited record in a semicolon-
+    # delimited sheet. Decode the outer quoting before the original CSV quoting.
+    # Terms containing semicolons may occupy more than one outer cell, so join
+    # those cells back together rather than truncating the term at the first one.
+    with path.open(encoding="utf-8-sig", newline="") as handle:
+        reader = csv.reader(handle, delimiter=";")
+        header = next(reader, [])
+        inner_header = next(csv.reader([header[0]]), []) if header else []
+        wrapped = (
+            len(header) > 1
+            and len(inner_header) == 2
+            and inner_header[0].strip().lower() == "code"
+            and inner_header[1].strip().lower() in {"term", "description"}
+            and not any(cell.strip() for cell in header[1:])
+        )
+        if wrapped:
+            rows = []
+            for row in reader:
+                if not row or not any(cell.strip() for cell in row):
+                    continue
+                record = ";".join(row).rstrip(";")
+                try:
+                    fields = next(csv.reader([record], strict=True))
+                except csv.Error as exc:
+                    raise ValueError(
+                        f"Malformed nested SNOMED CSV at source line {reader.line_num}"
+                    ) from exc
+                if len(fields) != 2:
+                    raise ValueError(
+                        f"Expected code and term at SNOMED source line {reader.line_num}"
+                    )
+                code, description = (value.strip() for value in fields)
+                if code and description:
+                    rows.append((code, description))
+            if not rows:
+                raise ValueError(f"No rows loaded from {path}")
+            frame = pd.DataFrame(rows, columns=["code", "description"])
+            frame["vocabulary"] = "SNOMEDCT_US"
+            return frame.drop_duplicates(subset=["code", "vocabulary"], keep="first").reset_index(drop=True)
+
+    return _load_two_column_text(csv_path, vocabulary="SNOMEDCT_US", delimiter=",")
+
+
+CODE_SYSTEM_LOADERS = {
+    "ICD10CM": load_icd10cm_full,
+    "ICD9CM": load_icd9cm_full,
+    "ICPC": load_icpc_full,
+    "MDR": load_mdr_full,
+    "RCD2": load_rcd2_full,
+    "SNOMEDCT_US": load_snomedct_full,
+}
+
+
+def load_code_system_corpus(
+    vocabularies: list[str],
+    source_paths: dict[str, str] | None = None,
+    cache_dir: str | Path = "data/codes/parquet/snomed_definitive_20261030",
+) -> pd.DataFrame:
+    """Load (and cache) the full reference code corpus for the given vocabularies.
+
+    Each vocabulary is parsed from its source file at most once; the processed
+    result (``code``, ``description``, ``vocabulary``) is cached to
+    ``<cache_dir>/codes_<vocabulary>.parquet`` and reused on subsequent calls,
+    so the (potentially slow) source loader only runs once per vocabulary.
+
+    Parameters
+    ----------
+    vocabularies : list[str]
+        Coding systems to include, e.g. ``["ICD10CM"]``.
+    source_paths : dict[str, str], optional
+        Override source file path per vocabulary (defaults to each loader's
+        built-in default path).
+    cache_dir : str or Path
+        Directory to store/read the cached per-vocabulary Parquet files.
+
+    Returns
+    -------
+    pd.DataFrame
+        Concatenated corpus across all requested vocabularies with columns
+        ``code``, ``description``, ``vocabulary``.
+    """
+    if not vocabularies:
+        raise ValueError("At least one vocabulary must be specified.")
+
+    source_paths = source_paths or {}
+    cache_dir = Path(cache_dir)
+    cache_dir.mkdir(parents=True, exist_ok=True)
+
+    frames: list[pd.DataFrame] = []
+    for vocab in vocabularies:
+        if vocab not in CODE_SYSTEM_LOADERS:
+            raise ValueError(
+                f"No loader registered for coding system '{vocab}'. "
+                f"Available: {sorted(CODE_SYSTEM_LOADERS)}"
+            )
+
+        cache_path = cache_dir / f"codes_{vocab}.parquet"
+        source_path = Path(source_paths.get(vocab, DEFAULT_CODE_SYSTEM_PATHS[vocab]))
+        cache_is_fresh = (
+            cache_path.exists()
+            and source_path.exists()
+            and cache_path.stat().st_mtime_ns >= source_path.stat().st_mtime_ns
+        )
+        if cache_is_fresh:
+            frames.append(pd.read_parquet(cache_path))
+            continue
+
+        loader = CODE_SYSTEM_LOADERS[vocab]
+        df = loader(source_path)
+        df.to_parquet(cache_path, index=False)
+        frames.append(df)
+
+    return pd.concat(frames, ignore_index=True) if len(frames) > 1 else frames[0]

src/data/loaders.py

@@ -0,0 +1,181 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import pandas as pd
+import yaml
+
+
+@dataclass
+class EventDefinitionForm:
+    """Structured phenotype definition (EDF) providing semantic context for
+    LLM-assisted concept retrieval and code classification."""
+
+    name: str
+    preferred_name: str
+    definition: str = ""
+    inclusion_criteria: list[str] = field(default_factory=list)
+    narrow_definition: str = ""
+    exclusion_criteria: list[str] = field(default_factory=list)
+    narrow_decision_rules: list[str] = field(default_factory=list)
+    possible_decision_rules: list[str] = field(default_factory=list)
+    exclude_decision_rules: list[str] = field(default_factory=list)
+    synonyms: list[str] = field(default_factory=list)
+    anatomical_location: str = ""
+    population_restrictions: str = ""
+    temporal_characteristics: str = ""
+    etiology: str = ""
+    morphology: str = ""
+    diagnostic_certainty: str = ""
+    administrative_context: str = ""
+    candidate_anchor_codes: list[str] = field(default_factory=list)
+    ambiguities: list[str] = field(default_factory=list)
+    references: list[str] = field(default_factory=list)
+
+    def to_prompt_context(self, include_possible: bool = True) -> str:
+        """Render the EDF as structured text suitable for retrieval or an LLM prompt.
+
+        Possible rules are never discarded. In binary mode they define cases
+        that are insufficient for Narrow and therefore map to Exclude.
+        """
+        lines = [f"Preferred name: {self.preferred_name}"]
+        if self.definition:
+            lines.append(f"Definition: {self.definition}")
+        if self.narrow_definition:
+            lines.append(f"Narrow definition: {self.narrow_definition}")
+        if self.inclusion_criteria:
+            lines.append("Inclusion criteria:")
+            lines += [f"  - {c}" for c in self.inclusion_criteria]
+        if self.exclusion_criteria:
+            lines.append("Exclusion criteria:")
+            lines += [f"  - {c}" for c in self.exclusion_criteria]
+        if self.narrow_decision_rules:
+            lines.append("Criteria for Narrow classification:")
+            lines += [f"  - {c}" for c in self.narrow_decision_rules]
+        if self.possible_decision_rules:
+            if include_possible:
+                lines.append("Criteria for Possible classification (when enabled):")
+            else:
+                lines.append(
+                    "Criteria that are insufficient for Narrow classification "
+                    "(classify as Exclude because Possible is disabled):"
+                )
+            lines += [f"  - {c}" for c in self.possible_decision_rules]
+        if self.exclude_decision_rules:
+            lines.append("Criteria for Exclude classification:")
+            lines += [f"  - {c}" for c in self.exclude_decision_rules]
+        if self.synonyms:
+            lines.append(f"Synonyms: {', '.join(self.synonyms)}")
+        if self.anatomical_location:
+            lines.append(f"Anatomical location: {self.anatomical_location}")
+        if self.population_restrictions:
+            lines.append(f"Population restrictions: {self.population_restrictions}")
+        if self.temporal_characteristics:
+            lines.append(f"Temporal characteristics: {self.temporal_characteristics}")
+        if self.etiology:
+            lines.append(f"Etiology: {self.etiology}")
+        if self.morphology:
+            lines.append(f"Morphology/pathology: {self.morphology}")
+        if self.diagnostic_certainty:
+            lines.append(f"Diagnostic certainty: {self.diagnostic_certainty}")
+        if self.administrative_context:
+            lines.append(f"Administrative/historical context: {self.administrative_context}")
+        if self.candidate_anchor_codes:
+            lines.append("Candidate anchor codes:")
+            lines += [f"  - {c}" for c in self.candidate_anchor_codes]
+        if self.ambiguities:
+            lines.append("Ambiguities:")
+            lines += [f"  - {c}" for c in self.ambiguities]
+        if self.references:
+            lines.append("References:")
+            lines += [f"  - {c}" for c in self.references]
+        return "\n".join(lines)
+
+    def to_lexical_query(self) -> str:
+        """Render the phenotype name and its synonyms for lexical matching."""
+        terms = [self.preferred_name.strip(), *(term.strip() for term in self.synonyms)]
+        return " ".join(dict.fromkeys(term for term in terms if term))
+
+    def to_embedding_query(self) -> str:
+        """Render the EDF subset used for semantic embedding matching."""
+        lines = [f"Preferred name: {self.preferred_name}"]
+        if self.narrow_definition:
+            lines.append(f"Narrow definition: {self.narrow_definition}")
+        if self.synonyms:
+            lines.append(f"Synonyms: {', '.join(self.synonyms)}")
+        if self.narrow_decision_rules:
+            lines.append("Narrow decision rules:")
+            lines += [f"  - {c}" for c in self.narrow_decision_rules]
+        if self.candidate_anchor_codes:
+            lines.append("Candidate anchor codes:")
+            lines += [f"  - {c}" for c in self.candidate_anchor_codes]
+        if self.morphology:
+            lines.append(f"Morphology/pathology: {self.morphology}")
+        if self.anatomical_location:
+            lines.append(f"Anatomical location: {self.anatomical_location}")
+        return "\n".join(lines)
+
+
+def load_edf(name: str, edf_dir: str | Path = "data/edfs/yamls") -> EventDefinitionForm:
+    """Load an Event Definition Form by phenotype name from a YAML file."""
+    path = Path(edf_dir) / f"{name}.yaml"
+    if not path.exists():
+        raise FileNotFoundError(f"EDF not found for phenotype '{name}': {path}")
+    raw = yaml.safe_load(path.read_text()) or {}
+    return EventDefinitionForm(name=name, **raw)
+
+
+def load_aesi_dataset(csv_path: str | Path) -> tuple[pd.DataFrame, pd.DataFrame]:
+    """Load a raw AESI (Adverse Event of Special Interest) review export and
+    split it into a candidate code corpus and expert-curated gold labels.
+
+    These CSV exports (e.g. ``C_MYOCARD_AESI_filtered.csv``) contain the full
+    set of candidate codes considered during expert phenotype curation, with
+    one row per (code, vocabulary) pair and a ``tags`` column holding the
+    reviewer's final classification (``narrow``/``possible``/``exclude``).
+
+    Parameters
+    ----------
+    csv_path : str or Path
+        Path to the raw AESI CSV export.
+
+    Returns
+    -------
+    tuple[pd.DataFrame, pd.DataFrame]
+        ``(codes, gold_labels)`` where ``codes`` has columns
+        ``code``, ``description``, ``vocabulary`` and ``gold_labels`` has
+        columns ``code``, ``code_name``, ``vocabulary``, ``label``
+        (Narrow/Possible/Exclude).
+    """
+    path = Path(csv_path)
+    if not path.exists():
+        raise FileNotFoundError(f"AESI dataset not found: {path}")
+
+    raw = pd.read_csv(path, dtype=str)
+    required = {"code", "code_name", "coding_system", "tags"}
+    missing = required - set(raw.columns)
+    if missing:
+        raise ValueError(f"AESI dataset missing required columns: {missing}")
+
+    df = raw.dropna(subset=["tags"]).copy()
+    df["label"] = df["tags"].str.strip().str.title()
+
+    valid_labels = {"Narrow", "Possible", "Exclude"}
+    invalid = set(df["label"]) - valid_labels
+    if invalid:
+        raise ValueError(f"AESI dataset contains invalid tag values: {invalid}")
+
+    codes = (
+        df[["code", "code_name", "coding_system"]]
+        .rename(columns={"code_name": "description", "coding_system": "vocabulary"})
+        .drop_duplicates(subset=["code", "vocabulary"], keep="first")
+        .reset_index(drop=True)
+    )
+    gold_labels = (
+        df[["code", "code_name", "coding_system", "label"]]
+        .rename(columns={"coding_system": "vocabulary"})
+        .drop_duplicates(subset=["code", "vocabulary"], keep="first")
+        .reset_index(drop=True)
+    )
+    return codes, gold_labels

src/data/preprocessing.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+import re
+import unicodedata
+
+import pandas as pd
+
+_WHITESPACE_RE = re.compile(r"\s+")
+_PUNCT_RE = re.compile(r"[^\w\s]")
+_CODE_RANGE_RE = re.compile(
+    r"^[A-Z]*\d+[A-Z0-9.]*\s*[-\u2010-\u2015\u2212]\s*[A-Z]*\d+[A-Z0-9.]*$",
+    re.IGNORECASE,
+)
+
+
+def normalize_text(text: str) -> str:
+    """Lowercase, strip accents/punctuation, and collapse whitespace for
+    lexical matching and deduplication."""
+    if not isinstance(text, str):
+        return ""
+    text = unicodedata.normalize("NFKD", text)
+    text = "".join(ch for ch in text if not unicodedata.combining(ch))
+    text = text.lower()
+    text = _PUNCT_RE.sub(" ", text)
+    text = _WHITESPACE_RE.sub(" ", text).strip()
+    return text
+
+
+def standardize_code(code: str, vocabulary: str) -> str:
+    """Standardize a code's surface form per vocabulary conventions
+    (uppercase, no internal whitespace; ICD-style codes keep their dot)."""
+    if not isinstance(code, str):
+        return ""
+    code = code.strip().upper()
+    code = re.sub(r"\s+", "", code)
+    vocabulary = (vocabulary or "").strip().upper()
+    if vocabulary in {"ICD10", "ICD10CM"} and "." not in code and len(code) > 3:
+        code = f"{code[:3]}.{code[3:]}"
+    return code
+
+
+def is_code_range(code: object) -> bool:
+    """Return whether *code* is a dictionary grouping range, not a billable code.
+
+    Examples include ``N17-N19``, ``N17–N19``, and ``001-009.99``. Both
+    endpoints must contain a digit, which avoids treating leading-hyphen ICPC
+    process codes such as ``-30`` or ordinary hyphenated text as a range.
+    """
+    return isinstance(code, str) and _CODE_RANGE_RE.fullmatch(code.strip()) is not None
+
+
+def filter_code_ranges(df: pd.DataFrame) -> pd.DataFrame:
+    """Remove non-codable dictionary range rows from a terminology corpus."""
+    if df.empty or "code" not in df.columns:
+        return df.copy()
+    range_mask = df["code"].map(is_code_range)
+    return df.loc[~range_mask].reset_index(drop=True)
+
+
+def deduplicate_codes(df: pd.DataFrame) -> pd.DataFrame:
+    """Remove duplicate (code, vocabulary) pairs, keeping the first occurrence."""
+    if df.empty:
+        return df
+    working = df.copy()
+    working["_code_key"] = working.apply(
+        lambda r: standardize_code(r["code"], r.get("vocabulary", "")), axis=1
+    )
+    working = working.drop_duplicates(subset=["_code_key", "vocabulary"], keep="first")
+    return working.drop(columns=["_code_key"]).reset_index(drop=True)
+
+
+def preprocess_corpus(df: pd.DataFrame) -> pd.DataFrame:
+    """Filter ranges, standardize codes/descriptions, and drop duplicates."""
+    df = filter_code_ranges(df)
+    df = deduplicate_codes(df)
+    df = df.copy()
+    df["code"] = [standardize_code(c, v) for c, v in zip(df["code"], df["vocabulary"])]
+    df["normalized_description"] = df["description"].map(normalize_text)
+    return df

src/evaluation/__init__.py

@@ -0,0 +1 @@
+"""Evaluation metrics against gold-standard phenotype codelists."""

src/evaluation/metrics.py

@@ -0,0 +1,510 @@
+"""Evaluation metrics for concept retrieval and code classification against
+expert-curated gold-standard phenotype codelists."""
+
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass
+import json
+
+import pandas as pd
+
+from src.data.preprocessing import standardize_code
+
+LABELS = ("Narrow", "Possible", "Exclude")
+
+
+@dataclass(slots=True)
+class RetrievalMetrics:
+    true_positives: int
+    false_positives: int
+    false_negatives: int
+    sensitivity: float
+    precision: float
+    f1: float
+
+
+@dataclass(slots=True)
+class ClassificationMetrics:
+    accuracy: float
+    cohens_kappa: float
+    true_negatives: int
+    sensitivity: float
+    precision: float
+    f1: float
+    macro_sensitivity: float
+    macro_precision: float
+    macro_f1: float
+    per_label: dict[str, dict[str, float]]
+
+
+@dataclass(slots=True)
+class NarrowLossBreakdown:
+    """Where available gold-Narrow codes were lost end to end."""
+
+    gold_narrow_available: int
+    correctly_classified: int
+    misclassified: int
+    not_classified: int
+    total_missed: int
+    end_to_end_recall: float
+
+
+@dataclass(slots=True)
+class GPTReviewEvaluation:
+    """Gold-standard impact of the final low-confidence GPT review."""
+
+    candidates_reviewed: int
+    labels_changed: int
+    corrected_changes: int
+    harmful_changes: int
+    unchanged: int
+    correct_before: int
+    correct_after: int
+    net_correct_change: int
+    reviewed_accuracy_before: float
+    reviewed_accuracy_after: float
+
+
+def _standardized_key(code: str, vocabulary: str) -> tuple[str, str]:
+    return standardize_code(code, vocabulary), (vocabulary or "").strip().upper()
+
+
+def filter_records_by_vocabulary(records: list[dict], vocabularies: list[str]) -> list[dict]:
+    """Keep only records whose ``vocabulary`` is in *vocabularies* (case-insensitive)."""
+    vocab_set = {v.strip().upper() for v in vocabularies}
+    return [r for r in records if (r.get("vocabulary") or "").strip().upper() in vocab_set]
+
+
+def filter_gold_by_vocabulary(gold_labels: pd.DataFrame, vocabularies: list[str]) -> pd.DataFrame:
+    """Keep only gold-label rows whose ``vocabulary`` is in *vocabularies* (case-insensitive)."""
+    vocab_set = {v.strip().upper() for v in vocabularies}
+    mask = gold_labels["vocabulary"].str.strip().str.upper().isin(vocab_set)
+    return gold_labels[mask].reset_index(drop=True)
+
+
+def filter_gold_by_available_codes(
+    gold_labels: pd.DataFrame,
+    available_codes: pd.DataFrame,
+) -> pd.DataFrame:
+    """Restrict gold labels to codes present in the loaded terminology corpus.
+
+    This prevents codes from older or otherwise unavailable terminology
+    versions from being counted as false negatives.
+    """
+    required = {"code", "vocabulary"}
+    missing = required - set(available_codes.columns)
+    if missing:
+        raise ValueError(f"Available code corpus missing required columns: {missing}")
+
+    available_keys = {
+        _standardized_key(row["code"], row["vocabulary"])
+        for _, row in available_codes.iterrows()
+    }
+    mask = gold_labels.apply(
+        lambda row: _standardized_key(row["code"], row["vocabulary"]) in available_keys,
+        axis=1,
+    )
+    return gold_labels[mask].reset_index(drop=True)
+
+
+def narrow_loss_breakdown(
+    classified: list[dict],
+    gold_labels: pd.DataFrame,
+) -> NarrowLossBreakdown:
+    """Separate gold-Narrow losses before and during LLM classification.
+
+    ``not_classified`` covers gold-Narrow codes that never reached the model
+    (for example because adaptive stopping or a hard top-k limit ended the
+    scan). ``misclassified`` covers codes that did reach the model but were
+    assigned a label other than Narrow.
+    """
+    gold_narrow_keys = {
+        _standardized_key(row["code"], row["vocabulary"])
+        for _, row in gold_labels.iterrows()
+        if str(row["label"]).strip().title() == "Narrow"
+    }
+    predictions_by_key = {
+        _standardized_key(item["code"], item["vocabulary"]):
+        str(item.get("label", "")).strip().title()
+        for item in classified
+    }
+
+    correctly_classified = sum(
+        predictions_by_key.get(key) == "Narrow" for key in gold_narrow_keys
+    )
+    misclassified = sum(
+        key in predictions_by_key and predictions_by_key[key] != "Narrow"
+        for key in gold_narrow_keys
+    )
+    not_classified = len(gold_narrow_keys - predictions_by_key.keys())
+    total_missed = misclassified + not_classified
+    recall = correctly_classified / len(gold_narrow_keys) if gold_narrow_keys else 0.0
+
+    return NarrowLossBreakdown(
+        gold_narrow_available=len(gold_narrow_keys),
+        correctly_classified=correctly_classified,
+        misclassified=misclassified,
+        not_classified=not_classified,
+        total_missed=total_missed,
+        end_to_end_recall=round(recall, 4),
+    )
+
+
+def evaluate_gpt_review(
+    stage2_classified: list[dict],
+    final_classified: list[dict],
+    gold_labels: pd.DataFrame,
+    *,
+    confidence_threshold: float,
+    max_candidates: int | None = None,
+    include_possible: bool = False,
+) -> GPTReviewEvaluation:
+    """Measure whether GPT's reviewed labels improved stage-2 decisions."""
+    gold_by_key = {
+        _standardized_key(row["code"], row["vocabulary"]): (
+            row["label"]
+            if include_possible or row["label"] != "Possible"
+            else "Exclude"
+        )
+        for _, row in gold_labels.iterrows()
+    }
+    final_by_key = {
+        _standardized_key(item["code"], item["vocabulary"]):
+        str(item.get("label", "")).strip().title()
+        for item in final_classified
+    }
+
+    eligible = [
+        item
+        for item in stage2_classified
+        if float(item.get("confidence", 0.0)) < confidence_threshold
+    ]
+    eligible.sort(key=lambda item: float(item.get("confidence", 0.0)))
+    selected = eligible[:max_candidates] if max_candidates is not None else eligible
+    selected_keys = {
+        _standardized_key(item["code"], item["vocabulary"]) for item in selected
+    }
+
+    reviewed = 0
+    changed = 0
+    corrected = 0
+    harmful = 0
+    correct_before = 0
+    correct_after = 0
+    for original in stage2_classified:
+        key = _standardized_key(original["code"], original["vocabulary"])
+        if key not in selected_keys:
+            continue
+        if key not in final_by_key:
+            continue
+        reviewed += 1
+        original_label = str(original.get("label", "")).strip().title()
+        final_label = final_by_key[key]
+        if not include_possible:
+            original_label = "Exclude" if original_label == "Possible" else original_label
+            final_label = "Exclude" if final_label == "Possible" else final_label
+        gold_label = gold_by_key.get(key, "Exclude")
+        was_correct = original_label == gold_label
+        is_correct = final_label == gold_label
+        correct_before += was_correct
+        correct_after += is_correct
+        if original_label != final_label:
+            changed += 1
+            corrected += is_correct and not was_correct
+            harmful += was_correct and not is_correct
+
+    return GPTReviewEvaluation(
+        candidates_reviewed=reviewed,
+        labels_changed=changed,
+        corrected_changes=corrected,
+        harmful_changes=harmful,
+        unchanged=reviewed - changed,
+        correct_before=correct_before,
+        correct_after=correct_after,
+        net_correct_change=correct_after - correct_before,
+        reviewed_accuracy_before=round(correct_before / reviewed, 4) if reviewed else 0.0,
+        reviewed_accuracy_after=round(correct_after / reviewed, 4) if reviewed else 0.0,
+    )
+
+
+def to_predicted_codelist(classified: list[dict]) -> pd.DataFrame:
+    """Render final classifications in the same schema as the ground-truth AESI
+    export (``coding_system``, ``code``, ``code_name``, ``concept``,
+    ``concept_name``, ``tags``) so predicted and reference codelists are
+    directly comparable/exportable side by side.
+
+    Parameters
+    ----------
+    classified : list[dict]
+        Each dict must contain ``code``, ``vocabulary``, ``description``, and
+        ``label`` keys (as produced by :func:`src.llm.classify.llm_classify`).
+
+    Returns
+    -------
+    pd.DataFrame
+        AESI-compatible columns. When classification annotations are supplied,
+        confidence, explanation, manual-review flag, and review reason are
+        appended. ``concept`` is blank because the corpus has no UMLS mapping.
+    """
+    base_columns = ["coding_system", "code", "code_name", "concept", "concept_name", "tags"]
+    annotation_columns = ["confidence", "explanation", "manual_review", "review_reason"]
+    include_annotations = any(any(column in item for column in annotation_columns) for item in classified)
+    columns = base_columns + annotation_columns if include_annotations else base_columns
+    mapping_columns = ["decision_origin", "mapping_provenance", "prior_label", "prior_confidence"]
+    include_mapping = any(c.get("decision_origin") == "codemapper" for c in classified)
+    if include_mapping:
+        columns += mapping_columns
+    return pd.DataFrame(
+        [
+            {
+                "coding_system": c["vocabulary"],
+                "code": c["code"],
+                "code_name": c.get("description", ""),
+                "concept": "",
+                "concept_name": c.get("description", ""),
+                "tags": (c.get("label") or "").lower(),
+                **(
+                    {
+                        "confidence": c.get("confidence", ""),
+                        "explanation": c.get("explanation", ""),
+                        "manual_review": bool(c.get("manual_review", False) or c.get("decision_origin") == "codemapper"),
+                        "review_reason": c.get("review_reason") or (
+                            "included by CodeMapper mapping; EDF scope not independently assessed"
+                            if c.get("decision_origin") == "codemapper" else ""
+                        ),
+                    }
+                    if include_annotations
+                    else {}
+                ),
+                **({"decision_origin": c.get("decision_origin", "llm"),
+                    "mapping_provenance": json.dumps(c.get("mapping_provenance", [])),
+                    "prior_label": c.get("prior_label"),
+                    "prior_confidence": c.get("prior_confidence")} if include_mapping else {}),
+            }
+            for c in classified
+        ],
+        columns=columns,
+    )
+
+
+def evaluate_retrieval(
+    retrieved_codes: list[dict],
+    gold_labels: pd.DataFrame,
+    narrow_only_as_positive: bool = False,
+) -> RetrievalMetrics:
+    """Evaluate concept retrieval as an information-retrieval task.
+
+    By default, true positives are retrieved codes labeled 'Narrow' or
+    'Possible' in the gold standard. With ``narrow_only_as_positive=True``, only
+    'Narrow' is positive.
+    False positives are retrieved codes labeled 'Exclude' in the gold standard,
+    or codes absent from it entirely.
+    False negatives are positive gold-standard codes that were not retrieved.
+
+    Parameters
+    ----------
+    retrieved_codes : list[dict]
+        Each dict must contain ``code`` and ``vocabulary`` keys.
+    gold_labels : pd.DataFrame
+        Reference standard with ``code``, ``vocabulary``, ``label`` columns.
+
+    Returns
+    -------
+    RetrievalMetrics
+    """
+    gold_by_key = {
+        _standardized_key(row["code"], row["vocabulary"]): row["label"]
+        for _, row in gold_labels.iterrows()
+    }
+    positive_labels = {"Narrow"} if narrow_only_as_positive else {"Narrow", "Possible"}
+    positive_keys = {k for k, label in gold_by_key.items() if label in positive_labels}
+
+    retrieved_keys = {
+        _standardized_key(r["code"], r["vocabulary"]) for r in retrieved_codes
+    }
+
+    tp = len(retrieved_keys & positive_keys)
+    fp = len(retrieved_keys - positive_keys)
+    fn = len(positive_keys - retrieved_keys)
+
+    sensitivity = tp / (tp + fn) if (tp + fn) > 0 else 0.0
+    precision = tp / (tp + fp) if (tp + fp) > 0 else 0.0
+    f1 = (
+        2 * precision * sensitivity / (precision + sensitivity)
+        if (precision + sensitivity) > 0
+        else 0.0
+    )
+
+    return RetrievalMetrics(
+        true_positives=tp,
+        false_positives=fp,
+        false_negatives=fn,
+        sensitivity=round(sensitivity, 4),
+        precision=round(precision, 4),
+        f1=round(f1, 4),
+    )
+
+
+def _cohens_kappa(confusion: dict[tuple[str, str], int], labels: tuple[str, ...], n: int) -> float:
+    if n == 0:
+        return 0.0
+    observed_agreement = sum(confusion.get((lbl, lbl), 0) for lbl in labels) / n
+
+    row_totals = {lbl: sum(confusion.get((lbl, o), 0) for o in labels) for lbl in labels}
+    col_totals = {lbl: sum(confusion.get((o, lbl), 0) for o in labels) for lbl in labels}
+    expected_agreement = sum((row_totals[lbl] * col_totals[lbl]) / n for lbl in labels) / n
+
+    if expected_agreement >= 1.0:
+        return 1.0
+    return (observed_agreement - expected_agreement) / (1 - expected_agreement)
+
+
+def evaluate_classification(
+    classified: list[dict],
+    gold_labels: pd.DataFrame,
+    include_possible: bool = True,
+) -> ClassificationMetrics:
+    """Evaluate code classification (Narrow/Possible/Exclude) against gold labels.
+
+    Only codes present in both the classified set and gold standard are compared
+    (codes absent from either side cannot be scored for classification agreement).
+
+    Parameters
+    ----------
+    classified : list[dict]
+        Each dict must contain ``code``, ``vocabulary``, ``label`` keys.
+    gold_labels : pd.DataFrame
+        Reference standard with ``code``, ``vocabulary``, ``label`` columns.
+
+    Returns
+    -------
+    ClassificationMetrics
+    """
+    labels = LABELS if include_possible else ("Narrow", "Exclude")
+    gold_by_key = {
+        _standardized_key(row["code"], row["vocabulary"]): (
+            row["label"] if include_possible or row["label"] != "Possible" else "Exclude"
+        )
+        for _, row in gold_labels.iterrows()
+    }
+
+    confusion: dict[tuple[str, str], int] = {}
+    correct = 0
+    total = 0
+    per_label_counts = {lbl: {"tp": 0, "tn": 0, "fp": 0, "fn": 0} for lbl in labels}
+
+    for item in classified:
+        key = _standardized_key(item["code"], item["vocabulary"])
+        gold_label = gold_by_key.get(key)
+        pred_label = item["label"]
+        if not include_possible and pred_label == "Possible":
+            pred_label = "Exclude"
+        if gold_label is None:
+            # A retrieved code absent from the curated positive codelist is an
+            # implicit Exclude. Apply this consistently to both correct
+            # Exclude predictions and false-positive Narrow/Possible results.
+            gold_label = "Exclude"
+
+        total += 1
+        confusion[(gold_label, pred_label)] = confusion.get((gold_label, pred_label), 0) + 1
+        if pred_label == gold_label:
+            correct += 1
+        for label in labels:
+            gold_positive = gold_label == label
+            predicted_positive = pred_label == label
+            if gold_positive and predicted_positive:
+                per_label_counts[label]["tp"] += 1
+            elif not gold_positive and not predicted_positive:
+                per_label_counts[label]["tn"] += 1
+            elif predicted_positive:
+                per_label_counts[label]["fp"] += 1
+            else:
+                per_label_counts[label]["fn"] += 1
+
+    accuracy = correct / total if total > 0 else 0.0
+    kappa = _cohens_kappa(confusion, labels, total)
+
+    per_label: dict[str, dict[str, float]] = {}
+    for lbl in labels:
+        tp = per_label_counts[lbl]["tp"]
+        tn = per_label_counts[lbl]["tn"]
+        fp = per_label_counts[lbl]["fp"]
+        fn = per_label_counts[lbl]["fn"]
+        precision = tp / (tp + fp) if (tp + fp) > 0 else 0.0
+        recall = tp / (tp + fn) if (tp + fn) > 0 else 0.0
+        f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0.0
+        per_label[lbl] = {
+            "tp": tp,
+            "tn": tn,
+            "fp": fp,
+            "fn": fn,
+            "precision": round(precision, 4),
+            "recall": round(recall, 4),
+            "f1": round(f1, 4),
+        }
+
+    macro_labels = [
+        label
+        for label, counts in per_label_counts.items()
+        if counts["tp"] + counts["fn"] > 0
+    ] or list(labels)
+
+    return ClassificationMetrics(
+        accuracy=round(accuracy, 4),
+        cohens_kappa=round(kappa, 4),
+        true_negatives=per_label["Narrow"]["tn"],
+        sensitivity=per_label["Narrow"]["recall"],
+        precision=per_label["Narrow"]["precision"],
+        f1=per_label["Narrow"]["f1"],
+        macro_sensitivity=round(sum(per_label[label]["recall"] for label in macro_labels) / len(macro_labels), 4),
+        macro_precision=round(sum(per_label[label]["precision"] for label in macro_labels) / len(macro_labels), 4),
+        macro_f1=round(sum(per_label[label]["f1"] for label in macro_labels) / len(macro_labels), 4),
+        per_label=per_label,
+    )
+
+
+def evaluate(
+    classified: list[dict],
+    gold_labels: pd.DataFrame,
+    narrow_only_as_positive: bool = False,
+    include_possible: bool = True,
+) -> dict:
+    """Compute both retrieval and classification metrics for the full pipeline output.
+
+    Parameters
+    ----------
+    classified : list[dict]
+        Final classified candidates, each with ``code``, ``vocabulary``, ``label`` keys.
+    gold_labels : pd.DataFrame
+        Reference standard with ``code``, ``vocabulary``, ``label`` columns.
+
+    Returns
+    -------
+    dict
+        Dictionary with 'retrieval' and 'classification' metric sub-dictionaries.
+    """
+    retrieval_metrics = evaluate_retrieval(
+        classified,
+        gold_labels,
+        narrow_only_as_positive=narrow_only_as_positive,
+    )
+    classification_metrics = evaluate_classification(
+        classified,
+        gold_labels,
+        include_possible=include_possible,
+    )
+    return {
+        "retrieval": asdict(retrieval_metrics),
+        "classification": {
+            "accuracy": classification_metrics.accuracy,
+            "cohens_kappa": classification_metrics.cohens_kappa,
+            "true_negatives": classification_metrics.true_negatives,
+            "sensitivity": classification_metrics.sensitivity,
+            "precision": classification_metrics.precision,
+            "f1": classification_metrics.f1,
+            "macro_sensitivity": classification_metrics.macro_sensitivity,
+            "macro_precision": classification_metrics.macro_precision,
+            "macro_f1": classification_metrics.macro_f1,
+            "per_label": classification_metrics.per_label,
+        },
+    }

src/llm/__init__.py

@@ -0,0 +1 @@
+"""LLM-based ranking and classification of candidate codes."""

src/llm/classify.py

@@ -0,0 +1,531 @@
+"""LLM classification into Narrow/Exclude with an optional Possible category."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from collections.abc import Callable, Sequence
+import json
+from typing import Protocol
+
+from src.data.loaders import EventDefinitionForm
+from src.llm.client import call_jev_decisions, call_llm_json, resolve_provider
+from src.llm.prompts import build_classify_prompt, classify_system_prompt
+from src.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+VALID_LABELS = {"Narrow", "Possible", "Exclude"}
+
+
+class ClassificationCandidate(Protocol):
+    """Minimal candidate interface consumed by the classifier."""
+
+    code: str
+    description: str
+    vocabulary: str
+
+
+@dataclass(frozen=True, slots=True)
+class AdaptiveStoppingConfig:
+    """Stop classification after Narrow results remain sparse for several batches."""
+
+    sparse_narrow_threshold: int = 1
+    consecutive_sparse_batches: int = 3
+    minimum_batches: int = 3
+
+    def __post_init__(self) -> None:
+        if self.sparse_narrow_threshold < 0:
+            raise ValueError("sparse_narrow_threshold must be at least 0")
+        if self.consecutive_sparse_batches < 1:
+            raise ValueError("consecutive_sparse_batches must be at least 1")
+        if self.minimum_batches < 1:
+            raise ValueError("minimum_batches must be at least 1")
+
+
+@dataclass(frozen=True, slots=True)
+class ClassificationRunInfo:
+    """Summary of how much of the ranked candidate list was classified."""
+
+    candidates_available: int
+    candidates_classified: int
+    batches_completed: int
+    narrow_counts_by_batch: tuple[int, ...]
+    stopped_early: bool
+    stop_reason: str
+
+
+@dataclass(slots=True)
+class ClassifiedCandidate:
+    code: str
+    description: str
+    vocabulary: str
+    label: str
+    confidence: float
+    explanation: str
+    relevance_score: float = 0.0
+    retrieval_score: float = 0.0
+    decision_origin: str = "llm"
+    mapping_provenance: list[dict] = field(default_factory=list)
+    prior_label: str | None = None
+    prior_confidence: float | None = None
+
+
+@dataclass(frozen=True, slots=True)
+class GPTReviewRunInfo:
+    """Summary of final-stage GPT adjudication of uncertain classifications."""
+
+    confidence_threshold: float
+    max_candidates: int | None
+    eligible_candidates: int
+    candidates_reviewed: int
+    labels_changed: int
+    narrow_to_exclude: int
+    exclude_to_narrow: int
+    reasoning_effort: str
+
+
+def select_gpt_review_candidates(
+    classified: Sequence[ClassifiedCandidate],
+    *,
+    confidence_threshold: float,
+    max_candidates: int | None = None,
+) -> list[ClassifiedCandidate]:
+    """Select the lowest-confidence eligible decisions for GPT review."""
+    if not 0.0 <= confidence_threshold <= 1.0:
+        raise ValueError("confidence_threshold must be between 0 and 1")
+    if max_candidates is not None and max_candidates < 1:
+        raise ValueError("max_candidates must be at least 1 when provided")
+
+    eligible = [item for item in classified if item.confidence < confidence_threshold]
+    eligible.sort(key=lambda item: item.confidence)
+    return eligible[:max_candidates] if max_candidates is not None else eligible
+
+
+def gpt_review_low_confidence(
+    classified: Sequence[ClassifiedCandidate],
+    edf: EventDefinitionForm,
+    *,
+    confidence_threshold: float = 0.7,
+    max_candidates: int | None = None,
+    model: str | None = None,
+    max_retries: int = 3,
+    batch_size: int = 10,
+    progress_callback: Callable[[int, int], None] | None = None,
+    use_possible_category: bool = False,
+    reasoning_effort: str = "medium",
+    provider: str = "openai",
+) -> tuple[list[ClassifiedCandidate], GPTReviewRunInfo]:
+    """Use GPT to adjudicate only classifications below the confidence cutoff.
+
+    Results at or above the threshold are preserved exactly. Reviewed results
+    replace the corresponding stage-2 labels while retaining their original
+    list positions and recording the prior label/confidence in the explanation.
+    """
+    eligible_count = sum(item.confidence < confidence_threshold for item in classified)
+    uncertain = select_gpt_review_candidates(
+        classified,
+        confidence_threshold=confidence_threshold,
+        max_candidates=max_candidates,
+    )
+    if not uncertain:
+        return list(classified), GPTReviewRunInfo(
+            confidence_threshold=confidence_threshold,
+            max_candidates=max_candidates,
+            eligible_candidates=eligible_count,
+            candidates_reviewed=0,
+            labels_changed=0,
+            narrow_to_exclude=0,
+            exclude_to_narrow=0,
+            reasoning_effort=reasoning_effort,
+        )
+
+    reviewed = llm_classify(
+        uncertain,
+        edf,
+        provider=provider,
+        model=model,
+        max_retries=max_retries,
+        batch_size=batch_size,
+        progress_callback=progress_callback,
+        use_possible_category=use_possible_category,
+        adaptive_stopping=None,
+        reasoning_effort=reasoning_effort,
+    )
+    reviewed_by_key = {(item.code, item.vocabulary): item for item in reviewed}
+
+    merged: list[ClassifiedCandidate] = []
+    labels_changed = 0
+    narrow_to_exclude = 0
+    exclude_to_narrow = 0
+    for original in classified:
+        replacement = reviewed_by_key.get((original.code, original.vocabulary))
+        if replacement is None:
+            merged.append(original)
+            continue
+
+        if replacement.label != original.label:
+            labels_changed += 1
+            if original.label == "Narrow" and replacement.label == "Exclude":
+                narrow_to_exclude += 1
+            elif original.label == "Exclude" and replacement.label == "Narrow":
+                exclude_to_narrow += 1
+        merged.append(
+            ClassifiedCandidate(
+                code=original.code,
+                description=original.description,
+                vocabulary=original.vocabulary,
+                label=replacement.label,
+                confidence=replacement.confidence,
+                explanation=(
+                    "GPT review of low-confidence stage-2 result "
+                    f"(previous label={original.label}, confidence={original.confidence:.2f}). "
+                    + replacement.explanation
+                ),
+                relevance_score=original.relevance_score,
+                retrieval_score=original.retrieval_score,
+            )
+        )
+
+    return merged, GPTReviewRunInfo(
+        confidence_threshold=confidence_threshold,
+        max_candidates=max_candidates,
+        eligible_candidates=eligible_count,
+        candidates_reviewed=len(uncertain),
+        labels_changed=labels_changed,
+        narrow_to_exclude=narrow_to_exclude,
+        exclude_to_narrow=exclude_to_narrow,
+        reasoning_effort=reasoning_effort,
+    )
+
+
+def llm_classify(
+    candidates: Sequence[ClassificationCandidate],
+    edf: EventDefinitionForm,
+    provider: str = "auto",
+    model: str | None = None,
+    max_retries: int = 3,
+    batch_size: int = 10,
+    progress_callback: Callable[[int, int], None] | None = None,
+    use_possible_category: bool = False,
+    adaptive_stopping: AdaptiveStoppingConfig | None = None,
+    run_info_callback: Callable[[ClassificationRunInfo], None] | None = None,
+    reasoning_effort: str | None = None,
+    thinking_budget: int | None = None,
+) -> list[ClassifiedCandidate]:
+    """Classify candidates as Narrow/Exclude, optionally allowing Possible.
+
+    Parameters
+    ----------
+    candidates : sequence
+        Similarity-ranked hybrid or LLM-ranked candidates.
+    edf : EventDefinitionForm
+        Structured phenotype definition providing semantic context.
+    provider : str
+        "auto", "anthropic", or "openai".
+    model : str, optional
+        Model identifier. Defaults to a provider-specific default.
+    max_retries : int
+        Number of LLM call attempts before failing.
+    batch_size : int
+        Maximum number of candidates to include in a single LLM request.
+    progress_callback : callable, optional
+        Callback invoked after each completed batch as ``callback(completed, total)``.
+    use_possible_category : bool
+        Enable three-way Narrow/Possible/Exclude labeling. Defaults to False.
+    adaptive_stopping : AdaptiveStoppingConfig, optional
+        Stop after Narrow results remain at or below the configured threshold
+        for the configured number of consecutive batches.
+    run_info_callback : callable, optional
+        Invoked once with the completed/stopped run summary.
+    reasoning_effort : str, optional
+        OpenAI reasoning effort, such as ``"none"`` or ``"medium"``.
+
+    Returns
+    -------
+    list[ClassifiedCandidate]
+    """
+    if not candidates:
+        if run_info_callback is not None:
+            run_info_callback(
+                ClassificationRunInfo(0, 0, 0, (), False, "No candidates to classify.")
+            )
+        return []
+
+    if batch_size < 1:
+        raise ValueError("batch_size must be at least 1")
+
+    resolved_provider = resolve_provider(provider)
+    if resolved_provider in {"jev", "openai_decisions", "laya", "jpt"}:
+        from src.llm.decisions import DEFAULT_DECISION_MODELS
+        return _jev_classify(
+            candidates,
+            edf,
+            model=model or DEFAULT_DECISION_MODELS[resolved_provider],
+            max_retries=max_retries,
+            batch_size=batch_size,
+            progress_callback=progress_callback,
+            use_possible_category=use_possible_category,
+            adaptive_stopping=adaptive_stopping,
+            run_info_callback=run_info_callback,
+            provider=resolved_provider,
+        )
+
+    results: list[ClassifiedCandidate] = []
+    narrow_counts: list[int] = []
+    sparse_streak = 0
+
+    total = len(candidates)
+    for start in range(0, total, batch_size):
+        batch = candidates[start : start + batch_size]
+        batch_result_start = len(results)
+        lookup = {(c.code, c.vocabulary): c for c in batch}
+        payload = [
+            {"code": c.code, "vocabulary": c.vocabulary, "description": c.description}
+            for c in batch
+        ]
+        prompt = build_classify_prompt(edf, payload, use_possible_category=use_possible_category)
+        response = call_llm_json(
+            system_prompt=classify_system_prompt(use_possible_category),
+            user_prompt=prompt,
+            provider=resolved_provider,
+            model=model,
+            max_retries=max_retries,
+            reasoning_effort=reasoning_effort,
+            **({'thinking_budget':thinking_budget} if thinking_budget is not None else {}),
+        )
+
+        classifications_raw = response.get("classifications", [])
+        batch_seen: set[tuple[str, str]] = set()
+        for item in classifications_raw:
+            key = (item.get("code", ""), item.get("vocabulary", ""))
+            source = lookup.get(key)
+            if source is None or key in batch_seen:
+                continue
+            fallback_label = "Possible" if use_possible_category else "Exclude"
+            valid_labels = VALID_LABELS if use_possible_category else {"Narrow", "Exclude"}
+            label = item.get("label", fallback_label)
+            if label not in valid_labels:
+                logger.warning(
+                    "LLM returned invalid label %r for code %s; defaulting to %r",
+                    label,
+                    key,
+                    fallback_label,
+                )
+                label = fallback_label
+            batch_seen.add(key)
+            results.append(
+                ClassifiedCandidate(
+                    code=source.code,
+                    description=source.description,
+                    vocabulary=source.vocabulary,
+                    label=label,
+                    confidence=float(item.get("confidence", 0.0)),
+                    explanation=str(item.get("explanation", "")),
+                    relevance_score=float(getattr(source, "relevance_score", 0.0)),
+                    retrieval_score=float(
+                        getattr(source, "retrieval_score", getattr(source, "score", 0.0))
+                    ),
+                )
+            )
+
+        # Keep every candidate in the result. A missing response receives the
+        # least-committal enabled label and zero confidence so review selection
+        # can annotate it without filtering or altering any other result.
+        for c in batch:
+            key = (c.code, c.vocabulary)
+            if key not in batch_seen:
+                results.append(
+                    ClassifiedCandidate(
+                        code=c.code,
+                        description=c.description,
+                        vocabulary=c.vocabulary,
+                        label="Possible" if use_possible_category else "Exclude",
+                        confidence=0.0,
+                        explanation="No classification returned by LLM; routed for manual review.",
+                        relevance_score=float(getattr(c, "relevance_score", 0.0)),
+                        retrieval_score=float(
+                            getattr(c, "retrieval_score", getattr(c, "score", 0.0))
+                        ),
+                    )
+                )
+
+        if progress_callback is not None:
+            progress_callback(min(start + len(batch), total), total)
+
+        narrow_count = sum(
+            item.label == "Narrow" for item in results[batch_result_start:]
+        )
+        narrow_counts.append(narrow_count)
+        sparse_streak = (
+            sparse_streak + 1
+            if adaptive_stopping is not None
+            and narrow_count <= adaptive_stopping.sparse_narrow_threshold
+            else 0
+        )
+        if _adaptive_stop_reached(adaptive_stopping, len(narrow_counts), sparse_streak):
+            break
+
+    _report_run_info(
+        run_info_callback,
+        total=total,
+        results=results,
+        narrow_counts=narrow_counts,
+        adaptive_stopping=adaptive_stopping,
+    )
+    return results
+
+
+def _jev_classify(
+    candidates: Sequence[ClassificationCandidate],
+    edf: EventDefinitionForm,
+    model: str,
+    max_retries: int,
+    batch_size: int,
+    progress_callback: Callable[[int, int], None] | None,
+    use_possible_category: bool,
+    adaptive_stopping: AdaptiveStoppingConfig | None,
+    run_info_callback: Callable[[ClassificationRunInfo], None] | None,
+    provider: str = "jev",
+) -> list[ClassifiedCandidate]:
+    """Classify candidates with Jev typed choice questions."""
+    effective_batch_size = min(batch_size, 20)
+    labels = {
+        "Narrow": "The code is sufficiently specific to the phenotype and meets the Narrow criteria.",
+        "Exclude": "The code is irrelevant, insufficiently specific, historical, or explicitly excluded.",
+    }
+    if use_possible_category:
+        labels = {
+            "Narrow": labels["Narrow"],
+            "Possible": "The code is related but not sufficiently specific in every context.",
+            "Exclude": labels["Exclude"],
+        }
+
+    results: list[ClassifiedCandidate] = []
+    narrow_counts: list[int] = []
+    sparse_streak = 0
+    total = len(candidates)
+    state = "Phenotype definition:\n" + edf.to_prompt_context(
+        include_possible=use_possible_category
+    )
+    for start in range(0, total, effective_batch_size):
+        batch = candidates[start : start + effective_batch_size]
+        batch_result_start = len(results)
+        questions = {
+            f"candidate_{offset}": {
+                "type": "choice",
+                "instructions": (
+                    "Classify this clinical terminology candidate for the phenotype: "
+                    f"code={candidate.code}; vocabulary={candidate.vocabulary}; "
+                    f"description={candidate.description}"
+                ),
+                "criteria": labels,
+            }
+            for offset, candidate in enumerate(batch)
+        }
+        if provider == "jev":
+            decision_call = call_jev_decisions
+        else:
+            from src.llm.decisions import call_decisions
+            decision_call = lambda **kwargs: call_decisions(provider=provider, **kwargs)
+        response = decision_call(
+            state=state,
+            questions=questions,
+            model=model,
+            max_retries=max_retries,
+        )
+        answers = response.get("answers", {})
+        for offset, candidate in enumerate(batch):
+            answer = answers.get(f"candidate_{offset}", {})
+            probabilities = answer.get("probabilities", {})
+            fallback_label = "Possible" if use_possible_category else "Exclude"
+            label = answer.get("choice", fallback_label)
+            if label not in labels:
+                label = fallback_label
+            confidence = answer.get("confidence")
+            if confidence is None:
+                confidence = max((float(value) for value in probabilities.values()), default=0.0)
+            results.append(
+                ClassifiedCandidate(
+                    code=candidate.code,
+                    description=candidate.description,
+                    vocabulary=candidate.vocabulary,
+                    label=label,
+                    confidence=float(confidence),
+                    explanation=f"{provider} choice probabilities: " + json.dumps(probabilities, sort_keys=True)
+                    + "; confidence source: " + answer.get("confidence_source", "provider" if answer.get("confidence") is not None else "max_probability"),
+                    relevance_score=float(getattr(candidate, "relevance_score", 0.0)),
+                    retrieval_score=float(
+                        getattr(
+                            candidate,
+                            "retrieval_score",
+                            getattr(candidate, "score", 0.0),
+                        )
+                    ),
+                )
+            )
+        if progress_callback is not None:
+            progress_callback(min(start + len(batch), total), total)
+        narrow_count = sum(
+            item.label == "Narrow" for item in results[batch_result_start:]
+        )
+        narrow_counts.append(narrow_count)
+        sparse_streak = (
+            sparse_streak + 1
+            if adaptive_stopping is not None
+            and narrow_count <= adaptive_stopping.sparse_narrow_threshold
+            else 0
+        )
+        if _adaptive_stop_reached(adaptive_stopping, len(narrow_counts), sparse_streak):
+            break
+
+    _report_run_info(
+        run_info_callback,
+        total=total,
+        results=results,
+        narrow_counts=narrow_counts,
+        adaptive_stopping=adaptive_stopping,
+    )
+    return results
+
+
+def _adaptive_stop_reached(
+    config: AdaptiveStoppingConfig | None,
+    batches_completed: int,
+    sparse_streak: int,
+) -> bool:
+    return bool(
+        config is not None
+        and batches_completed >= config.minimum_batches
+        and sparse_streak >= config.consecutive_sparse_batches
+    )
+
+
+def _report_run_info(
+    callback: Callable[[ClassificationRunInfo], None] | None,
+    *,
+    total: int,
+    results: list[ClassifiedCandidate],
+    narrow_counts: list[int],
+    adaptive_stopping: AdaptiveStoppingConfig | None,
+) -> None:
+    if callback is None:
+        return
+    stopped_early = len(results) < total
+    if stopped_early and adaptive_stopping is not None:
+        reason = (
+            f"Stopped after {adaptive_stopping.consecutive_sparse_batches} consecutive "
+            f"batches with at most {adaptive_stopping.sparse_narrow_threshold} Narrow "
+            "result(s) per batch."
+        )
+    else:
+        reason = "All available candidates were classified."
+    callback(
+        ClassificationRunInfo(
+            candidates_available=total,
+            candidates_classified=len(results),
+            batches_completed=len(narrow_counts),
+            narrow_counts_by_batch=tuple(narrow_counts),
+            stopped_early=stopped_early,
+            stop_reason=reason,
+        )
+    )

src/llm/client.py

@@ -0,0 +1,376 @@
+"""Shared model clients for generative JSON calls and Jev typed decisions.
+
+Automatic provider selection preserves the original Anthropic-then-OpenAI
+precedence; Google is selected explicitly by experiments that require it.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+import time
+import uuid
+from typing import Any
+from urllib.error import HTTPError, URLError
+from urllib.request import Request, urlopen
+
+from src.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+_JSON_BLOCK_RE = re.compile(r"\{.*\}", re.DOTALL)
+
+DEFAULT_ANTHROPIC_MODEL = "claude-haiku-4-8"
+DEFAULT_OPENAI_MODEL = "gpt-6-luna"
+DEFAULT_GOOGLE_MODEL = "gemini-3-flash-preview"
+DEFAULT_JEV_MODEL = "jev-latest"
+DEFAULT_JEV_API_URL = "https://api.typesafe.ai/v1/systemone"
+
+_PLACEHOLDER_KEY_VALUES = {"", "your-api-key-here"}
+
+
+class LLMCallError(RuntimeError):
+    """Raised when the LLM call fails or returns unparsable output after retries."""
+
+
+def _has_real_key(env_var: str) -> bool:
+    """Return True if *env_var* is set to a non-empty, non-placeholder value.
+
+    Guards against `.env` files copied from `.env.example` where a key is
+    present but still holds the literal placeholder text.
+    """
+    value = os.environ.get(env_var, "").strip()
+    return value.lower() not in _PLACEHOLDER_KEY_VALUES
+
+
+def resolve_provider(provider: str = "auto") -> str:
+    """Resolve the LLM provider to use, based on explicit choice or available API keys.
+
+    Parameters
+    ----------
+    provider : str
+        One of ``"auto"``, ``"anthropic"``, ``"openai"``, ``"google"``, or ``"jev"``.
+
+    Returns
+    -------
+    str
+        The resolved provider name.
+    """
+    provider = (provider or "auto").lower()
+    if provider in {"anthropic", "openai", "google", "jev", "openai_decisions", "laya", "jpt", "mlx"}:
+        return provider
+    if provider != "auto":
+        raise LLMCallError(
+            f"Unknown LLM provider: {provider!r}. "
+            "Use 'auto', 'anthropic', 'openai', 'google', or 'jev'."
+        )
+
+    if _has_real_key("ANTHROPIC_API_KEY"):
+        return "anthropic"
+    if _has_real_key("OPENAI_API_KEY"):
+        return "openai"
+    if _has_real_key("JEV_API_KEY"):
+        return "jev"
+    raise LLMCallError(
+        "No LLM API key found. Set ANTHROPIC_API_KEY and/or OPENAI_API_KEY in your .env file."
+    )
+
+
+def default_model_for(provider: str) -> str:
+    """Return the default model identifier for the given provider."""
+    if provider == "anthropic":
+        return DEFAULT_ANTHROPIC_MODEL
+    if provider == "google":
+        return DEFAULT_GOOGLE_MODEL
+    if provider == "jev":
+        return DEFAULT_JEV_MODEL
+    if provider in {"laya", "jpt", "mlx"}:
+        return {"laya": "multilingual", "jpt": "kirp/jpt-4b",
+                "mlx": "mlx-community/Qwen3.5-9B-8bit"}[provider]
+    return DEFAULT_OPENAI_MODEL
+
+
+def _get_anthropic_client():
+    """Lazily construct the Anthropic client so importing this module does not
+    require an API key to be set (e.g. for unit tests that mock calls)."""
+    import anthropic
+
+    api_key = os.environ.get("ANTHROPIC_API_KEY")
+    if not api_key:
+        raise LLMCallError("ANTHROPIC_API_KEY environment variable is not set.")
+    return anthropic.Anthropic(api_key=api_key)
+
+
+def _get_openai_client():
+    """Lazily construct the OpenAI client so importing this module does not
+    require an API key to be set (e.g. for unit tests that mock calls)."""
+    import openai
+
+    api_key = os.environ.get("OPENAI_API_KEY")
+    if not api_key:
+        raise LLMCallError("OPENAI_API_KEY environment variable is not set.")
+    return openai.OpenAI(api_key=api_key)
+
+
+def _get_google_client():
+    """Lazily construct the Google Gen AI client."""
+    from google import genai
+
+    api_key = os.environ.get("GOOGLE_API_KEY") or os.environ.get("GEMINI_API_KEY")
+    if not api_key:
+        raise LLMCallError("GOOGLE_API_KEY or GEMINI_API_KEY environment variable is not set.")
+    return genai.Client(api_key=api_key)
+
+
+def extract_json(text: str) -> dict[str, Any]:
+    """Extract and parse the first JSON object found in *text*.
+
+    Handles cases where the model wraps JSON in markdown code fences or adds
+    surrounding commentary despite instructions not to.
+    """
+    text = text.strip()
+    try:
+        return json.loads(text)
+    except json.JSONDecodeError:
+        pass
+
+    match = _JSON_BLOCK_RE.search(text)
+    if match:
+        return json.loads(match.group(0))
+
+    raise ValueError(f"Could not extract JSON from LLM response: {text[:200]!r}")
+
+
+def _call_anthropic(system_prompt: str, user_prompt: str, model: str, max_tokens: int) -> str:
+    client = _get_anthropic_client()
+    response = client.messages.create(
+        model=model,
+        max_tokens=max_tokens,
+        system=system_prompt,
+        messages=[{"role": "user", "content": user_prompt}],
+    )
+    return "".join(
+        block.text for block in response.content if getattr(block, "type", None) == "text"
+    )
+
+
+_REASONING_MODEL_PREFIXES = ("gpt-5", "gpt-6", "o1", "o3", "o4")
+
+
+def _is_reasoning_model(model: str) -> bool:
+    """Whether *model* is an OpenAI reasoning model (GPT-5/6 and o-series).
+
+    Reasoning models spend part of their completion token budget on hidden
+    reasoning tokens before producing visible output, and support a
+    ``reasoning_effort`` parameter to control that trade-off. Non-reasoning
+    models (e.g. gpt-4o) don't accept that parameter, so it must only be sent
+    for models in this family.
+    """
+    return model.lower().startswith(_REASONING_MODEL_PREFIXES)
+
+
+def _call_openai(
+    system_prompt: str,
+    user_prompt: str,
+    model: str,
+    max_tokens: int,
+    reasoning_effort: str | None = None,
+) -> str:
+    client = _get_openai_client()
+    kwargs: dict[str, Any] = {}
+    if reasoning_effort and _is_reasoning_model(model):
+        kwargs["reasoning_effort"] = reasoning_effort
+
+    response = client.chat.completions.create(
+        model=model,
+        max_completion_tokens=max_tokens,
+        response_format={"type": "json_object"},
+        messages=[
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": user_prompt},
+        ],
+        **kwargs,
+    )
+    choice = response.choices[0]
+    content = choice.message.content or ""
+    if not content.strip():
+        # Most commonly caused by a reasoning model spending its entire
+        # max_completion_tokens budget on hidden reasoning, leaving nothing
+        # for the visible answer (finish_reason == "length" in that case).
+        raise ValueError(
+            f"Empty response content from model '{model}' "
+            f"(finish_reason={choice.finish_reason!r}, usage={response.usage!r}). "
+            "Consider increasing max_tokens or lowering reasoning_effort."
+        )
+    return content
+
+
+def _call_google(system_prompt: str, user_prompt: str, model: str, max_tokens: int) -> str:
+    from google.genai import types
+
+    client = _get_google_client()
+    response = client.models.generate_content(
+        model=model,
+        contents=user_prompt,
+        config=types.GenerateContentConfig(
+            system_instruction=system_prompt,
+            response_mime_type="application/json",
+            max_output_tokens=max_tokens,
+        ),
+    )
+    text = response.text or ""
+    if not text.strip():
+        raise ValueError(f"Empty response content from Google model '{model}'.")
+    return text
+
+
+def call_jev_decisions(
+    state: str,
+    questions: dict[str, dict[str, Any]],
+    model: str = DEFAULT_JEV_MODEL,
+    max_retries: int = 3,
+) -> dict[str, Any]:
+    """Submit typed questions to Jev and return its structured response.
+
+    Jev is not a generative chat model, so it does not use
+    :func:`call_llm_json`. A request shares one phenotype state across up to 20
+    candidate-classification questions and receives declared labels plus
+    probability distributions.
+    """
+    api_key = (
+        os.environ.get("TYPESAFE_API_KEY", "").strip()
+        or os.environ.get("JEV_API_KEY", "").strip()
+    )
+    if not api_key:
+        raise LLMCallError(
+            "TYPESAFE_API_KEY or JEV_API_KEY environment variable is not set."
+        )
+    if not questions or len(questions) > 20:
+        raise ValueError("Jev requests must contain between 1 and 20 questions.")
+
+    url = os.environ.get("JEV_API_URL", DEFAULT_JEV_API_URL).strip() or DEFAULT_JEV_API_URL
+    payload = json.dumps(
+        {"model": model, "state": state, "questions": questions},
+        separators=(",", ":"),
+    ).encode("utf-8")
+    idempotency_key = str(uuid.uuid5(uuid.NAMESPACE_URL, payload.decode("utf-8")))
+    retryable_statuses = {429, 502, 503, 504, 529}
+    last_error: Exception | None = None
+
+    for attempt in range(1, max_retries + 1):
+        request = Request(
+            url,
+            data=payload,
+            method="POST",
+            headers={
+                "Authorization": f"Bearer {api_key}",
+                "Content-Type": "application/json",
+                "Accept": "application/json",
+                "Idempotency-Key": idempotency_key,
+            },
+        )
+        try:
+            with urlopen(request, timeout=60) as response:
+                result = json.loads(response.read().decode("utf-8"))
+            if not isinstance(result.get("answers"), dict):
+                raise ValueError("Jev response is missing an 'answers' object.")
+            return result
+        except HTTPError as exc:
+            try:
+                error_body = exc.read().decode("utf-8", errors="replace")[:500]
+            except Exception:  # noqa: BLE001 - preserve the original HTTP status
+                error_body = ""
+            last_error = RuntimeError(
+                f"Jev HTTP {exc.code} from {url}"
+                + (f": {error_body}" if error_body else "")
+            )
+            if exc.code not in retryable_statuses or attempt >= max_retries:
+                break
+            retry_after = exc.headers.get("Retry-After") if exc.headers else None
+            delay = min(float(retry_after), 30.0) if retry_after else min(2 ** (attempt - 1), 8)
+            time.sleep(delay)
+        except (URLError, TimeoutError, ValueError, json.JSONDecodeError) as exc:
+            last_error = exc
+            if attempt < max_retries:
+                time.sleep(min(2 ** (attempt - 1), 8))
+
+    raise LLMCallError(
+        f"Jev decision call failed after {max_retries} attempts: {last_error}"
+    ) from last_error
+
+
+def call_llm_json(
+    system_prompt: str,
+    user_prompt: str,
+    provider: str = "auto",
+    model: str | None = None,
+    max_retries: int = 3,
+    max_tokens: int = 8192,
+    reasoning_effort: str | None = None,
+    thinking_budget: int | None = None,
+) -> dict[str, Any]:
+    """Call the configured LLM provider and parse a JSON object from the response,
+    retrying on transient failures or malformed JSON output.
+
+    Parameters
+    ----------
+    system_prompt : str
+        System prompt establishing the LLM's role and output format.
+    user_prompt : str
+        The task-specific user prompt.
+    provider : str
+        ``"auto"``, ``"anthropic"``, ``"openai"``, or ``"google"``. ``"auto"`` selects
+        Anthropic if ``ANTHROPIC_API_KEY`` is set, otherwise OpenAI.
+    model : str, optional
+        Model identifier. Defaults to a sensible default for the resolved provider.
+    max_retries : int
+        Number of attempts before raising ``LLMCallError``.
+    max_tokens : int
+        Maximum completion tokens to generate. For OpenAI reasoning models
+        (gpt-5/o1/o3/o4), this budget is shared with hidden reasoning tokens,
+        so it should be generous enough to leave room for the visible answer.
+    reasoning_effort : str, optional
+        Passed as ``reasoning_effort`` to OpenAI reasoning models only (e.g.
+        ``"low"``/``"medium"``/``"high"``); ignored for Anthropic and for
+        non-reasoning OpenAI models.
+
+    Returns
+    -------
+    dict
+        Parsed JSON response body.
+    """
+    resolved_provider = resolve_provider(provider)
+    if resolved_provider in {"jev", "openai_decisions", "laya", "jpt"}:
+        raise LLMCallError(
+            "Jev uses typed decisions rather than generative JSON. "
+            "Call call_jev_decisions(), or use llm_classify(provider='jev')."
+        )
+    resolved_model = model or default_model_for(resolved_provider)
+
+    last_error: Exception | None = None
+    for attempt in range(1, max_retries + 1):
+        try:
+            if resolved_provider == "anthropic":
+                text = _call_anthropic(system_prompt, user_prompt, resolved_model, max_tokens)
+            elif resolved_provider == "openai":
+                text = _call_openai(
+                    system_prompt, user_prompt, resolved_model, max_tokens, reasoning_effort
+                )
+            elif resolved_provider == "mlx":
+                from src.llm.local import call_mlx_json
+                return call_mlx_json(system_prompt, user_prompt, resolved_model, max_tokens, reasoning_effort,
+                                     thinking_budget if thinking_budget is not None else 1024)
+            else:
+                text = _call_google(system_prompt, user_prompt, resolved_model, max_tokens)
+            return extract_json(text)
+        except Exception as exc:  # noqa: BLE001 - broad by design, retried below
+            last_error = exc
+            logger.warning(
+                "LLM call attempt %d/%d (provider=%s) failed: %s",
+                attempt, max_retries, resolved_provider, exc,
+            )
+
+    raise LLMCallError(
+        f"LLM call failed after {max_retries} attempts (provider={resolved_provider}): {last_error}"
+    ) from last_error

src/llm/decisions.py

@@ -0,0 +1,111 @@
+"""Adapters for typed decision providers using the shared Jev question schema.
+
+Keep provider confidence separate from option probabilities. A fallback to the
+maximum option probability is recorded explicitly; it cannot be below 0.5 for
+a normalized binary choice. No domain calibration is inferred from that value.
+"""
+from __future__ import annotations
+
+from functools import lru_cache
+import json
+import os
+from urllib.request import Request, urlopen
+from typing import Any
+
+from src.llm.client import LLMCallError, _get_openai_client, call_jev_decisions
+
+DECISION_PROVIDERS = frozenset({'jev', 'openai_decisions', 'laya', 'jpt'})
+DEFAULT_DECISION_MODELS = {
+    'jev': 'jev-latest', 'openai_decisions': 'gpt-6-luna',
+    'laya': 'multilingual', 'jpt': 'kirp/jpt-4b',
+}
+
+
+def normalize_answers(response, names):
+    """Validate and normalize typed answers without inventing missing decisions."""
+    raw = response.get('answers')
+    if isinstance(raw, list):
+        keys = [answer.get('name') for answer in raw]
+        if len(set(keys)) != len(keys):
+            raise LLMCallError('Decision provider returned duplicate answer names.')
+        raw = {answer['name']: answer for answer in raw}
+    if not isinstance(raw, dict) or any(name not in raw for name in names):
+        raise LLMCallError('Decision provider omitted one or more requested answers.')
+    answers = {}
+    for name in names:
+        answer = dict(raw[name])
+        if answer.get('type') == 'refusal' or not isinstance(answer.get('choice'), str):
+            raise LLMCallError(f'Decision provider did not answer question {name}.')
+        probabilities = answer.get('probabilities', {})
+        if isinstance(probabilities, list):
+            probabilities = {item['value']: item['probability'] for item in probabilities}
+        probabilities = {key: float(value) for key, value in probabilities.items()}
+        if not probabilities or any(not 0 <= p <= 1 for p in probabilities.values()):
+            raise LLMCallError(f'Invalid choice probabilities for question {name}.')
+        if abs(sum(probabilities.values())-1) > .02 or answer['choice'] not in probabilities:
+            raise LLMCallError(f'Invalid choice distribution for question {name}.')
+        if isinstance(names, dict) and set(probabilities) != set(names[name]['criteria']):
+            raise LLMCallError(f'Unexpected choices for question {name}.')
+        confidence = answer.get('confidence')
+        answer['confidence_source'] = 'provider' if confidence is not None else 'max_probability'
+        confidence = float(confidence) if confidence is not None else max(probabilities.values())
+        if not 0 <= confidence <= 1:
+            raise LLMCallError(f'Invalid confidence for question {name}.')
+        answers[name] = {**answer, 'confidence': confidence, 'probabilities': probabilities}
+    return {**response, 'answers': answers}
+
+
+def call_openai_decisions(state, questions, model='gpt-6-luna', max_retries=3):
+    """POST /v1/decisions, including SDKs predating the typed resource."""
+    client = _get_openai_client().with_options(max_retries=max_retries-1)
+    payload = {'model': model, 'input': state, 'questions': [
+        {'name': name, 'type': 'choice', 'instructions': question['instructions'],
+         'choices': [{'value': label, 'description': description}
+                     for label, description in question['criteria'].items()]}
+        for name, question in questions.items()
+    ]}
+    response = client.post('/decisions', cast_to=dict[str, Any], body=payload)
+    return normalize_answers(response, questions)
+
+
+@lru_cache(maxsize=1)
+def _laya_router():
+    try:
+        from laya import Router
+    except ImportError as exc:
+        raise LLMCallError('Install requirements-decision-models.txt to use Laya.') from exc
+    from src.llm.model_paths import prepared_model_path
+    path=prepared_model_path('convaiinnovations/laya')
+    return Router(models={'multilingual':(path,'multilingual')})
+
+
+def call_decisions(state, questions, provider='jev', model=None, max_retries=3):
+    """Dispatch an existing typed choice request to the selected backend."""
+    if provider not in DECISION_PROVIDERS:
+        raise ValueError(f'Unsupported decision provider: {provider}')
+    model = model or DEFAULT_DECISION_MODELS[provider]
+    if provider == 'jev':
+        return call_jev_decisions(state, questions, model=model, max_retries=max_retries)
+    if provider == 'openai_decisions':
+        return call_openai_decisions(state, questions, model=model, max_retries=max_retries)
+    if provider == 'laya':
+        # Explicit multilingual checkpoint avoids the English router's shorter
+        # input limit. The model/card documents 8192-token inference support.
+        result = _laya_router().predict(state, questions, model=model, max_len=8192)
+        return normalize_answers(result, questions)
+    if not os.environ.get('JPT_API_URL'):
+        from src.llm.jpt_runtime import predict_jpt
+        return normalize_answers(predict_jpt(state, questions, model), questions)
+    url = os.environ.get('JPT_API_URL', 'http://127.0.0.1:8080/v1/systemone')
+    payload = json.dumps({'state': state, 'questions': questions}).encode()
+    headers = {'Content-Type': 'application/json'}
+    token = os.environ.get('JPT_API_KEY')
+    if token:
+        headers['Authorization'] = f'Bearer {token}'
+    request = Request(url, data=payload, headers=headers, method='POST')
+    try:
+        with urlopen(request, timeout=300) as response:
+            result = json.load(response)
+    except Exception as exc:
+        raise LLMCallError('JPT decision endpoint failed. Start llm2jev before running.') from exc
+    return normalize_answers(result, questions)

src/llm/jpt_runtime.py

@@ -0,0 +1,111 @@
+"""JPT-4B through its published llm2jev scoring interface on a local GPU."""
+from functools import lru_cache
+import platform
+
+
+class AppleJPT:
+    """Published llm2jev prompts and probabilities with MLX model execution."""
+    def __init__(self, model_name):
+        from mlx_vlm import load
+        from src.llm.model_paths import prepared_model_path
+        print(f'[JPT] Loading {model_name} with MLX (original BF16 weights)',flush=True)
+        self.model,self.processor=load(prepared_model_path(model_name))
+        self.prefix_tokens=[];self.prefix_cache=None
+
+    def warm(self,prefix,images):
+        if images:raise ValueError('The study uses text-only questions.')
+        import mlx.core as mx
+        # Leave a boundary margin so tokenization cannot merge the final prefix
+        # token with the question. Every scoring request verifies this prefix.
+        tokens=self.processor.tokenizer.encode(prefix,add_special_tokens=False)[:-2]
+        if tokens==self.prefix_tokens:return
+        self.prefix_tokens=tokens;lm=self.model.language_model
+        self.prefix_cache=lm.make_cache()
+        if tokens:
+            lm(mx.array([tokens]),cache=self.prefix_cache,skip_logits=True)
+            mx.eval([c.state for c in self.prefix_cache])
+
+    def score(self,text,images,ids):
+        if images:raise ValueError('The study uses text-only questions.')
+        rows,count=self._score_texts([text],ids)
+        return rows[0],count
+
+    def _score_texts(self,texts,ids):
+        import mlx.core as mx
+        tokens=[self.processor.tokenizer.encode(text,add_special_tokens=False) for text in texts]
+        lm=self.model.language_model
+        cache=lm.make_cache();suffixes=tokens
+        if self.prefix_tokens and all(t[:len(self.prefix_tokens)]==self.prefix_tokens for t in tokens):
+            from mlx.utils import tree_map
+            for fresh,saved in zip(cache,self.prefix_cache):
+                snapshot=tree_map(lambda x:mx.repeat(x,len(texts),axis=0) if isinstance(x,mx.array) else x,
+                                  saved.prefix_cache_snapshot())
+                fresh.prefix_cache_restore(snapshot)
+            suffixes=[t[len(self.prefix_tokens):] for t in tokens]
+        lengths=[len(s) for s in suffixes];width=max(lengths)
+        # Future padding cannot affect earlier positions in this causal model.
+        # Select each row's final real token before projecting to vocabulary.
+        padded=[s+[0]*(width-len(s)) for s in suffixes]
+        out=lm(mx.array(padded),cache=cache,skip_logits=True,return_hidden=True)
+        hidden=out.hidden_states[-1][mx.arange(len(texts)),mx.array(lengths)-1][:,None,:]
+        logits=lm.speculative_logits_from_hidden(hidden)[:,0].astype(mx.float32)
+        values=(logits-mx.logsumexp(logits,axis=-1,keepdims=True))[:,mx.array(ids)].tolist()
+        return values,sum(map(len,tokens))
+
+    def score_many(self,texts,ids):
+        rows=[];count=0
+        for i in range(0,len(texts),4):
+            batch,n=self._score_texts(texts[i:i+4],ids);rows.extend(batch);count+=n
+            print(f'[JPT] {min(i+4,len(texts))}/{len(texts)} questions scored',flush=True)
+        return rows,count
+
+
+@lru_cache(maxsize=1)
+def load_jpt(model_name):
+    from llm2jev.engine import LLM2Jev
+    if platform.system()=='Darwin' and platform.machine()=='arm64':
+        backend=AppleJPT(model_name)
+        return LLM2Jev(backend.processor,backend,temperature=1.036,workers=1)
+    import torch
+    from llm2jev.backends import HF
+
+    class LastTokenHF(HF):
+        def __init__(self, model, device=None, **kwargs):
+            # Loading safetensors directly onto MPS starts concurrent transfer
+            # threads. Load on CPU and transfer serially to avoid a native
+            # allocator/Metal crash in that loading path on this Mac.
+            super().__init__(model, device='cpu' if device=='mps' else device, **kwargs)
+            if device=='mps':self.model.to(device)
+        # Restrict the LM projection to the last position. The upstream HF
+        # reference projects every EDF token, wasting substantial GPU memory.
+        def score(self, text, images, ids):
+            if images:
+                raise ValueError('This experiment uses text-only JPT inputs.')
+            rows, count = self.score_many([text], ids)
+            return rows[0], count
+
+        def score_many(self, texts, ids):
+            self.tok.padding_side = 'left'
+            rows = []; count = 0
+            for start in range(0, len(texts), 4):
+                inputs = self.tok(texts[start:start+4],padding=True,return_tensors='pt',
+                                  add_special_tokens=False)
+                count += int(inputs['attention_mask'].sum())
+                inputs = {k:v.to(self.model.device) for k,v in inputs.items()}
+                with self.lock, torch.inference_mode():
+                    logits = self.model(**inputs,logits_to_keep=1).logits[:,-1].float().log_softmax(-1)
+                rows.extend(logits[:,ids].cpu().tolist())
+            return rows,count
+
+    device = 'mps' if torch.backends.mps.is_available() else 'cuda' if torch.cuda.is_available() else 'cpu'
+    print(f'[JPT] Loading {model_name} on {device}',flush=True)
+    from src.llm.model_paths import prepared_model_path
+    backend = LastTokenHF(prepared_model_path(model_name),device=device,dtype='bfloat16')
+    return LLM2Jev(backend.processor or backend.tok,backend,temperature=1.036,workers=1)
+
+
+def predict_jpt(state, questions, model_name):
+    answers, usage = load_jpt(model_name).run(state, questions)
+    return {'answers':answers,'model':model_name,'usage':usage,
+            'runtime':'llm2jev/mlx' if platform.system()=='Darwin' and platform.machine()=='arm64' else 'llm2jev/hf',
+            'temperature':1.036}

src/llm/local.py

@@ -0,0 +1,147 @@
+"""Local MLX reviewer for Apple silicon; no API calls or model substitution."""
+from functools import lru_cache
+import hashlib
+import json
+import os
+from pathlib import Path
+import re
+from time import perf_counter
+from uuid import uuid4
+
+
+def visible_response(text):
+    """Exclude reasoning, including a thinking block preopened by the template."""
+    text = re.sub(r'<think>.*?</think>', '', text, flags=re.S)
+    if '</think>' in text:
+        text = text.split('</think>',1)[1]
+    if '<think>' in text:
+        text = text.split('<think>',1)[0]
+    return text.strip()
+
+
+def first_response(text):
+    """Read a complete JSON envelope, allowing subsequent prose or repeated output.
+
+    Never repair incomplete JSON or extract a nested classification from a broken
+    envelope. Generation and classification settings remain unchanged.
+    """
+    text=visible_response(text)
+    start=text.find('{')
+    if start<0:raise ValueError('No complete JSON envelope in local response')
+    result,_=json.JSONDecoder().raw_decode(text[start:])
+    if not isinstance(result,dict):raise ValueError('Local response must be a JSON object')
+    if not isinstance(result.get('classifications'),list):
+        raise ValueError('Local response must contain a classifications array')
+    return result
+
+
+def candidate_parts(user_prompt):
+    marker='Classify each of the following candidate codes:\n'
+    if marker not in user_prompt:return None
+    prefix,body=user_prompt.split(marker,1)
+    lines,suffix=body.split('\n\nReturn a JSON object',1)
+    candidates=lines.splitlines()
+    if not candidates or any(not line.startswith('- code: ') for line in candidates):return None
+    return prefix+marker,candidates,'\n\nReturn a JSON object'+suffix
+
+
+def validate_candidates(response,user_prompt):
+    parts=candidate_parts(user_prompt)
+    if parts is None:return
+    expected=[]
+    for line in parts[1]:
+        code,vocabulary,_=line.removeprefix('- code: ').split(' | ',2)
+        expected.append((code,vocabulary.removeprefix('vocabulary: ')))
+    actual=[(r.get('code'),r.get('vocabulary')) for r in response['classifications']]
+    if len(actual)!=len(expected) or set(actual)!=set(expected):
+        raise ValueError('Local answer does not cover every requested code exactly once')
+
+
+@lru_cache(maxsize=1)
+def load_mlx_reviewer(model_name):
+    from mlx_vlm import load
+    from src.llm.model_paths import prepared_model_path
+    return load(prepared_model_path(model_name))
+
+
+def call_mlx_json(system_prompt, user_prompt, model_name, max_tokens=8192, reasoning_effort='medium',
+                  thinking_budget=1024):
+    from mlx_vlm import stream_generate
+    model, processor = load_mlx_reviewer(model_name)
+    messages = [{'role':'system','content':system_prompt},{'role':'user','content':user_prompt}]
+    prompt = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True,
+                                           enable_thinking=reasoning_effort not in (None,'none'))
+    folder=Path(os.environ.get('LOCAL_RESPONSE_DIR',Path(__file__).resolve().parents[2]/'results/local_responses'))
+    folder.mkdir(parents=True,exist_ok=True)
+    prompt_hash=hashlib.sha256(prompt.encode()).hexdigest()
+    failure=folder/'failed'/f'{prompt_hash}.json'
+    replay=json.loads(failure.read_text()) if failure.exists() else None
+    if replay and (replay.get('model'),replay.get('max_tokens'),replay.get('thinking_budget'))!=(model_name,max_tokens,thinking_budget):
+        replay=None
+    if replay:
+        print('[Qwen reviewer] Reusing audited failed output; subdividing its candidate batch.',flush=True)
+    chunks = []; started = perf_counter(); reported = 0; generated_tokens=0
+    finish_reason='model_stop_or_limit'
+    outputs=() if replay else stream_generate(model, processor, prompt=prompt, max_tokens=max_tokens,
+                                   temperature=.6, top_p=.95, top_k=20, seed=42,
+                                   enable_thinking=reasoning_effort not in (None,'none'),
+                                   thinking_budget=thinking_budget)
+    for output in outputs:
+        chunks.append(output.text)
+        generated_tokens=output.generation_tokens
+        if output.generation_tokens >= reported+256:
+            reported = output.generation_tokens
+            print(f'[Qwen reviewer] {reported} tokens generated, {perf_counter()-started:.1f}s',flush=True)
+            current=''.join(chunks)
+            if '</think>' in current or reasoning_effort in (None,'none'):
+                try:
+                    parsed=first_response(current)
+                    validate_candidates(parsed,user_prompt)
+                except json.JSONDecodeError as exc:
+                    # An unescaped control character inside a JSON string cannot
+                    # become valid by generating more tokens. Other incomplete
+                    # prefixes remain eligible to finish normally.
+                    if exc.msg.startswith('Invalid control character'):
+                        finish_reason='invalid_json_syntax'
+                        if hasattr(outputs,'close'):outputs.close()
+                        break
+                except (ValueError,TypeError):pass
+                else:
+                    finish_reason='complete_json'
+                    if hasattr(outputs,'close'):outputs.close()
+                    break
+    text = replay['raw_response'] if replay else ''.join(chunks)
+    path=folder/(uuid4().hex+'.json')
+    audit={'model':model_name,'raw_response':text,'generation_tokens':generated_tokens,
+           'max_tokens':max_tokens,'thinking_budget':thinking_budget,'seconds':perf_counter()-started,
+           'prompt_sha256':prompt_hash,'replayed_failed_response':bool(replay),'finish_reason':finish_reason,
+           'implementation_sha256':hashlib.sha256(Path(__file__).read_bytes()).hexdigest()}
+    try:
+        result=first_response(text)
+        validate_candidates(result,user_prompt)
+        audit['status']='parsed'
+        return result
+    except (ValueError,TypeError) as exc:
+        audit.update(status='parse_error',error=str(exc))
+        failure.parent.mkdir(exist_ok=True)
+        failure.write_text(json.dumps(audit,ensure_ascii=False,indent=2))
+        print(f'[Qwen reviewer] Raw failed response saved: {path}',flush=True)
+        parts=candidate_parts(user_prompt)
+        if parts and len(parts[1])>1:
+            prefix,lines,suffix=parts;middle=len(lines)//2
+            audit['recovery']='subdivide_batch'
+            # Persist the failure before recovery, including if recovery is interrupted.
+            path.write_text(json.dumps(audit,ensure_ascii=False,indent=2))
+            print(f'[Qwen reviewer] Retrying failed batch as {middle} + {len(lines)-middle} candidates.',flush=True)
+            recovered=[]
+            for subset in (lines[:middle],lines[middle:]):
+                answer=call_mlx_json(system_prompt,prefix+'\n'.join(subset)+suffix,model_name,
+                                     max_tokens,reasoning_effort,thinking_budget)
+                recovered.extend(answer['classifications'])
+            result={'classifications':recovered}
+            validate_candidates(result,user_prompt)
+            audit['status']='recovered_by_subdivision'
+            return result
+        raise
+    finally:
+        path.write_text(json.dumps(audit,ensure_ascii=False,indent=2))

src/llm/model_paths.py

@@ -0,0 +1,28 @@
+"""Use prepared local snapshots when available, retaining exact model revisions."""
+import json
+import gc
+import sys
+from pathlib import Path
+
+
+def prepared_model_path(repo):
+    root=Path(__file__).resolve().parents[2]
+    receipt=root/'.model-cache/model_revisions.json'
+    if receipt.exists():
+        info=json.loads(receipt.read_text()).get(repo)
+        if info and Path(info['path']).is_dir():return info['path']
+    return repo
+
+
+def release_local_models():
+    """Release cached weights before another local model or phase is loaded."""
+    for module_name,function in (('src.llm.decisions','_laya_router'),
+                                 ('src.llm.jpt_runtime','load_jpt'),
+                                 ('src.llm.local','load_mlx_reviewer')):
+        module=sys.modules.get(module_name)
+        if module is not None:getattr(module,function).cache_clear()
+    gc.collect()
+    torch=sys.modules.get('torch')
+    if torch is not None and torch.backends.mps.is_available():torch.mps.empty_cache()
+    mx=sys.modules.get('mlx.core')
+    if mx is not None:mx.clear_cache()

src/llm/prompts.py

@@ -0,0 +1,79 @@
+"""Reusable prompt templates for LLM-assisted ranking and classification."""
+
+from __future__ import annotations
+
+from src.data.loaders import EventDefinitionForm
+
+RANK_SYSTEM_PROMPT = (
+    "You are a clinical terminology expert supporting semantic harmonization for "
+    "real-world evidence studies. You review candidate medical codes retrieved for "
+    "a phenotype and re-rank them by clinical relevance, using the structured "
+    "phenotype definition (Event Definition Form) provided by the user. "
+    "You must respond only with valid JSON, no additional commentary."
+)
+
+CLASSIFY_SYSTEM_PROMPT = (
+    "You are a clinical terminology expert supporting semantic harmonization for "
+    "real-world evidence studies. You review candidate medical codes for a phenotype "
+    "and classify each one as 'Narrow' (highly specific to the phenotype), "
+    "'Possible' (related but insufficiently specific, or specific in some but not "
+    "all contexts), or 'Exclude' (not relevant, or explicitly an exclusion "
+    "criterion). Base your decision on the structured phenotype definition (Event "
+    "Definition Form) provided by the user. For every code, provide a brief "
+    "explanation followed by the label and a confidence score between 0 and 1. "
+    "You must respond only with valid JSON, no additional commentary."
+)
+
+
+def classify_system_prompt(use_possible_category: bool = False) -> str:
+    """Return classification instructions for binary or three-way labeling."""
+    if use_possible_category:
+        return CLASSIFY_SYSTEM_PROMPT
+    return (
+        "You are a clinical terminology expert supporting semantic harmonization "
+        "for real-world evidence studies. Classify every candidate as 'Narrow' "
+        "when it is sufficiently specific to the phenotype, or 'Exclude' otherwise. "
+        "Do not return a 'Possible' label. Any EDF rule describing a Possible, "
+        "ambiguous, or insufficiently specific case must be treated as Exclude in "
+        "this binary mode. Base decisions on the supplied Event Definition Form. "
+        "For every code, request the output fields in this order: code, vocabulary, "
+        "explanation, label, confidence. Confidence must be from 0 to 1. Respond "
+        "only with valid JSON."
+    )
+
+
+def build_rank_prompt(edf: EventDefinitionForm, candidates: list[dict]) -> str:
+    """Build the user prompt for LLM-based re-ranking of candidate codes."""
+    candidate_lines = "\n".join(
+        f"- code: {c['code']} | vocabulary: {c['vocabulary']} | description: {c['description']}"
+        for c in candidates
+    )
+    return (
+        f"Phenotype definition:\n{edf.to_prompt_context()}\n\n"
+        f"Candidate codes to rank by clinical relevance to this phenotype "
+        f"(most relevant first):\n{candidate_lines}\n\n"
+        "Return a JSON object with a single key 'ranked_codes', containing a list of "
+        "objects with keys 'code', 'vocabulary', and 'relevance_score' (float 0-1, "
+        "higher = more relevant), sorted descending by relevance_score."
+    )
+
+
+def build_classify_prompt(
+    edf: EventDefinitionForm,
+    candidates: list[dict],
+    use_possible_category: bool = False,
+) -> str:
+    """Build the user prompt for LLM-based classification of candidate codes."""
+    candidate_lines = "\n".join(
+        f"- code: {c['code']} | vocabulary: {c['vocabulary']} | description: {c['description']}"
+        for c in candidates
+    )
+    labels = "'Narrow', 'Possible', or 'Exclude'" if use_possible_category else "'Narrow' or 'Exclude'"
+    return (
+        f"Phenotype definition:\n{edf.to_prompt_context(include_possible=use_possible_category)}\n\n"
+        f"Classify each of the following candidate codes:\n{candidate_lines}\n\n"
+        "Return a JSON object with a single key 'classifications', containing a "
+        "list of objects using this requested key order: 'code', 'vocabulary', "
+        "'explanation' (brief string), 'label' "
+        f"(one of {labels}), and 'confidence' (float 0-1)."
+    )

src/llm/rank.py

@@ -0,0 +1,139 @@
+"""LLM-based re-ranking of candidate codes using full EDF semantic context."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from src.data.loaders import EventDefinitionForm
+from src.llm.client import call_llm_json
+from src.llm.prompts import RANK_SYSTEM_PROMPT, build_rank_prompt
+from src.retrieval.hybrid import HybridCandidate
+from src.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+@dataclass(slots=True)
+class RankedCandidate:
+    rank: int
+    code: str
+    description: str
+    vocabulary: str
+    relevance_score: float
+    retrieval_score: float
+
+    def __lt__(self, other: object) -> bool:
+        if not isinstance(other, RankedCandidate):
+            return NotImplemented
+        return self.relevance_score > other.relevance_score
+
+
+def skip_rank(candidates: list[HybridCandidate]) -> list[RankedCandidate]:
+    """Wrap hybrid retrieval candidates into ``RankedCandidate`` objects without
+    calling the LLM, preserving their original hybrid retrieval order.
+
+    The LLM ranking step (:func:`llm_rank`) is currently disabled in the
+    pipeline/UI (an extra LLM call that duplicates work also done by
+    classification) but is kept available for future use. This helper lets
+    downstream code (``llm_classify``) keep consuming ``RankedCandidate``
+    objects uniformly regardless of whether ranking actually ran.
+    """
+    results = [
+        RankedCandidate(
+            rank=c.rank,
+            code=c.code,
+            description=c.description,
+            vocabulary=c.vocabulary,
+            relevance_score=0.0,
+            retrieval_score=c.score,
+        )
+        for c in candidates
+    ]
+    return results
+
+
+def llm_rank(
+    candidates: list[HybridCandidate],
+    edf: EventDefinitionForm,
+    provider: str = "auto",
+    model: str | None = None,
+    max_retries: int = 3,
+) -> list[RankedCandidate]:
+    """Re-rank hybrid retrieval candidates using an LLM given the full EDF context.
+
+    Parameters
+    ----------
+    candidates : list[HybridCandidate]
+        Candidates produced by :func:`src.retrieval.hybrid.hybrid_retrieval`.
+    edf : EventDefinitionForm
+        Structured phenotype definition providing semantic context.
+    provider : str
+        "auto", "anthropic", or "openai".
+    model : str, optional
+        Model identifier. Defaults to a provider-specific default.
+    max_retries : int
+        Number of LLM call attempts before failing.
+
+    Returns
+    -------
+    list[RankedCandidate]
+        Candidates ordered by descending LLM-assigned relevance score.
+    """
+    if not candidates:
+        return []
+
+    lookup = {(c.code, c.vocabulary): c for c in candidates}
+    payload = [
+        {"code": c.code, "vocabulary": c.vocabulary, "description": c.description}
+        for c in candidates
+    ]
+    prompt = build_rank_prompt(edf, payload)
+    response = call_llm_json(
+        system_prompt=RANK_SYSTEM_PROMPT,
+        user_prompt=prompt,
+        provider=provider,
+        model=model,
+        max_retries=max_retries,
+    )
+
+    ranked_raw = response.get("ranked_codes", [])
+    results: list[RankedCandidate] = []
+    seen: set[tuple[str, str]] = set()
+
+    for item in ranked_raw:
+        key = (item.get("code", ""), item.get("vocabulary", ""))
+        source = lookup.get(key)
+        if source is None:
+            continue
+        seen.add(key)
+        results.append(
+            RankedCandidate(
+                rank=0,
+                code=source.code,
+                description=source.description,
+                vocabulary=source.vocabulary,
+                relevance_score=float(item.get("relevance_score", 0.0)),
+                retrieval_score=source.score,
+            )
+        )
+
+    # Append any candidates the LLM omitted, ordered by original retrieval score,
+    # so no candidate is silently dropped from downstream classification.
+    for c in candidates:
+        key = (c.code, c.vocabulary)
+        if key not in seen:
+            results.append(
+                RankedCandidate(
+                    rank=0,
+                    code=c.code,
+                    description=c.description,
+                    vocabulary=c.vocabulary,
+                    relevance_score=0.0,
+                    retrieval_score=c.score,
+                )
+            )
+
+    results.sort(key=lambda r: (-r.relevance_score, -r.retrieval_score))
+    for r, item in enumerate(results):
+        item.rank = r + 1
+    return results

src/mapping/__init__.py

@@ -0,0 +1 @@
+"""Optional terminology mapping stages."""

src/mapping/codemapper.py

@@ -0,0 +1,295 @@
+"""One-round cross-vocabulary CUI expansion, with caching and paced sessions."""
+from __future__ import annotations
+
+from dataclasses import replace
+try:
+    import fcntl
+except ImportError:  # Keep the disabled optional stage importable on Windows.
+    fcntl = None
+import hashlib
+from http.cookiejar import CookieJar
+import json
+import logging
+import os
+from pathlib import Path
+import time
+from urllib.error import HTTPError, URLError
+from urllib.parse import urlencode
+from urllib.request import HTTPCookieProcessor, Request, build_opener
+
+import pandas as pd
+
+from src.data.preprocessing import standardize_code
+from src.llm.classify import ClassifiedCandidate
+from src.utils.config import CodeMapperConfig
+
+
+class CodeMapperError(RuntimeError):
+    """Sanitized failure: never include credentials, cookies, or server bodies."""
+
+
+def normalize_code_versions(response,vocabulary):
+    """Select the pinned release from the current all-version API response."""
+    _,separator,version=vocabulary.partition('@')
+    if not separator:raise CodeMapperError('CodeMapper requires a pinned vocabulary version')
+    result={}
+    for code,value in response.items():
+        if isinstance(value,list):
+            matching=[item for item in value if isinstance(item,dict) and str(item.get('version'))==version]
+            if len(matching)>1:raise CodeMapperError('Ambiguous CodeMapper code version response')
+            if matching:result[code]=matching[0]
+        elif isinstance(value,dict):
+            result[code]=value  # Older cached/API responses were already release-specific.
+        else:raise CodeMapperError('Unexpected CodeMapper code version response')
+    return result
+
+
+def normalize_concept_versions(response,targets):
+    """Project series-level codes to requested releases, never to latest by default."""
+    versions={target.partition('@')[0]:target.partition('@')[2] for target in targets}
+    result=[]
+    for concept in response:
+        sources=[]
+        for source in concept.get('sourceConcepts',[]):
+            if 'codingSystem' in source:
+                if source['codingSystem'] in targets:sources.append(source)
+                continue
+            series=source.get('vocId');version=versions.get(series)
+            if version is None:continue
+            matching=[item for item in source.get('versions',[]) if str(item.get('version'))==version]
+            if matching:
+                sources.append({**source,'codingSystem':series+'@'+version,
+                                'preferredTerm':matching[0].get('term',source.get('preferredTerm'))})
+        result.append({**concept,'sourceConcepts':sources})
+    return result
+
+
+class CodeMapperClient:
+    def __init__(self, config: CodeMapperConfig):
+        self.config = config
+        self.cache_dir = Path(config.cache_dir)
+        self.cache_dir.mkdir(parents=True, exist_ok=True)
+        self.opener = build_opener(HTTPCookieProcessor(CookieJar()))
+        self.logged_in = False
+        self.calls: list[dict] = []
+        self.cache_hits = 0
+
+    def _pace(self):
+        if fcntl is None:
+            raise CodeMapperError("Shared CodeMapper pacing currently requires macOS or Linux")
+        # Share the limiter across client instances/processes using this cache.
+        with (self.cache_dir / "rate.lock").open("a+") as handle:
+            fcntl.flock(handle, fcntl.LOCK_EX)
+            handle.seek(0)
+            previous = float(handle.read() or 0)
+            delay = previous + 60 / self.config.requests_per_minute - time.time()
+            if delay > 0:
+                time.sleep(delay)
+            handle.seek(0)
+            handle.truncate()
+            handle.write(str(time.time()))
+            handle.flush()
+
+    def _request(self, path, method="POST", form=None):
+        fields = [(key, item) for key, value in (form or {}).items()
+                  for item in (value if isinstance(value, list) else [value])]
+        request = Request(self.config.base_url.rstrip("/") + path,
+                          data=urlencode(fields).encode() if form is not None else None,
+                          method=method, headers={"Accept": "application/json"})
+        for attempt in range(self.config.max_retries):
+            self._pace()
+            logging.getLogger(__name__).info('CodeMapper %s %s (attempt %s)',method,path,attempt+1)
+            started = time.perf_counter()
+            try:
+                with self.opener.open(request, timeout=self.config.timeout_seconds) as response:
+                    raw = response.read()
+                    status = response.status
+                self.calls.append({"path": path, "status": status,
+                                   "seconds": round(time.perf_counter() - started, 4),
+                                   "bytes": len(raw)})
+                return json.loads(raw) if raw else None
+            except HTTPError as exc:
+                self.calls.append({"path": path, "status": exc.code,
+                                   "seconds": round(time.perf_counter() - started, 4)})
+                if exc.code not in {429, 500, 502, 503, 504}:
+                    raise CodeMapperError(f"CodeMapper {path}: HTTP {exc.code}") from None
+                try:
+                    retry_after = float(exc.headers.get("Retry-After", 0))
+                except ValueError:
+                    retry_after = 0
+                if attempt + 1 < self.config.max_retries:
+                    time.sleep(max(retry_after, 2 ** attempt))
+            except (URLError, TimeoutError, OSError):
+                if attempt + 1 < self.config.max_retries:
+                    time.sleep(2 ** attempt)
+            except (ValueError, UnicodeDecodeError):
+                raise CodeMapperError(f"CodeMapper {path}: invalid JSON") from None
+        raise CodeMapperError(f"CodeMapper {path}: failed after retries")
+
+    def _login(self):
+        if self.logged_in:
+            return
+        username, password = os.getenv("CODEMAPPER_USERNAME"), os.getenv("CODEMAPPER_PASSWORD")
+        if not username or not password:
+            raise CodeMapperError("Set CODEMAPPER_USERNAME and CODEMAPPER_PASSWORD in .env")
+        response = self._request("/session", form={"username": username, "password": password})
+        if not isinstance(response, dict) or response.get("success") is not True:
+            raise CodeMapperError("CodeMapper login unsuccessful")
+        self.logged_in = True
+
+    def close(self):
+        if self.logged_in:
+            try:
+                self._request("/session", "DELETE")
+            finally:
+                self.logged_in = False
+
+    def _cache_path(self, kind, identity):
+        key = hashlib.sha256(json.dumps([self.config.base_url, kind, identity],
+                                       sort_keys=True).encode()).hexdigest()
+        return self.cache_dir / f"{key}.json"
+
+    def _load(self, kind, identity):
+        path = self._cache_path(kind, identity)
+        if path.exists():
+            self.cache_hits += 1
+            return json.loads(path.read_text())
+        return None
+
+    def _save(self, kind, identity, value):
+        path = self._cache_path(kind, identity)
+        temporary = path.with_suffix(f".{os.getpid()}.tmp")
+        temporary.write_text(json.dumps(value))
+        temporary.replace(path)
+
+    def cuis_for_codes(self, vocabulary, codes):
+        result, missing = {}, []
+        for code in sorted(set(codes)):
+            cached = self._load("code_cuis", [vocabulary, code])
+            if cached is None:
+                missing.append(code)
+            elif cached:
+                result[code] = cached
+        for start in range(0, len(missing), self.config.batch_size):
+            batch = missing[start:start + self.config.batch_size]
+            self._login()
+            response = self._request("/legacy/code-mapper/cuis-for-codes",
+                                     form={"vocId": vocabulary.partition('@')[0], "codes": batch})
+            if not isinstance(response, dict):
+                raise CodeMapperError("Unexpected code-to-CUI response")
+            response=normalize_code_versions(response,vocabulary)
+            for code in batch:
+                value = response.get(code, {})
+                self._save("code_cuis", [vocabulary, code], value)
+                if value:
+                    result[code] = value
+        return result
+
+    def concepts(self, cuis, targets):
+        result, missing = [], []
+        targets = sorted(set(targets))
+        for cui in sorted(set(cuis)):
+            cached = self._load("concept", [targets, cui])
+            if cached is None:
+                missing.append(cui)
+            else:
+                result.extend(cached)
+        for start in range(0, len(missing), self.config.batch_size):
+            batch = missing[start:start + self.config.batch_size]
+            self._login()
+            response = self._request("/legacy/code-mapper/umls-concepts",
+                                     form={"cuis": batch, "vocIds": sorted({target.partition('@')[0] for target in targets})})
+            if not isinstance(response, list):
+                raise CodeMapperError("Unexpected concept response")
+            response=normalize_concept_versions(response,targets)
+            for cui in batch:
+                values = [c for c in response if c.get("cui") == cui]
+                self._save("concept", [targets, cui], values)
+                result.extend(values)
+        return result
+
+
+def expand_narrow_matches(classified: list[ClassifiedCandidate], corpus: pd.DataFrame,
+                          vocabularies: list[str], config: CodeMapperConfig,
+                          *, client=None):
+    """Promote/add direct cross-vocabulary matches from a frozen Narrow seed set.
+
+    Only codes in the loaded/preprocessed corpus are eligible. Same-vocabulary
+    links and descendants are excluded. No new additions become seeds this round.
+    Mapping confidence is 0 (unassessed), not an invented model probability.
+    """
+    seeds = [c for c in classified if c.label == "Narrow" and c.vocabulary in vocabularies]
+    versions = {v: f"{v}@{config.vocabulary_versions[v]}" for v in vocabularies}
+    reverse = {version: voc for voc, version in versions.items()}
+    by_cui = {}
+    client = client or CodeMapperClient(config)
+    started = time.perf_counter()
+    seed_counts, unmapped = {}, []
+    try:
+        for voc in vocabularies:
+            group = [c for c in seeds if c.vocabulary == voc]
+            if not group:
+                continue
+            seed_counts[voc] = len(group)
+            response = client.cuis_for_codes(versions[voc], [c.code for c in group])
+            for seed in group:
+                cuis = response.get(seed.code, {}).get("cuis", [])
+                if not cuis:
+                    unmapped.append({"code": seed.code, "vocabulary": voc})
+                for cui in cuis:
+                    by_cui.setdefault(cui, []).append(seed)
+        concepts = client.concepts(by_cui, list(versions.values())) if by_cui else []
+    finally:
+        client.close()
+    available = {(standardize_code(c, v), v): (c, d)
+                 for c, v, d in corpus[["code", "vocabulary", "description"]].itertuples(index=False, name=None)}
+    links, unavailable = {}, {}
+    for concept in concepts:
+        for source in concept.get("sourceConcepts", []):
+            voc = reverse.get(source["codingSystem"])
+            if voc is None:
+                continue
+            key = (standardize_code(source["id"], voc), voc)
+            provenance = [{"seed_code": s.code, "seed_vocabulary": s.vocabulary,
+                           "seed_label": s.label, "cui": concept["cui"],
+                           "source_version": versions[s.vocabulary],
+                           "target_version": source["codingSystem"]}
+                          for s in by_cui.get(concept["cui"], []) if s.vocabulary != voc]
+            if not provenance:
+                continue
+            target = links if key in available else unavailable
+            target.setdefault(key, []).extend(provenance)
+    merged = list(classified)
+    positions = {(standardize_code(c.code, c.vocabulary), c.vocabulary): i
+                 for i, c in enumerate(merged)}
+    added, promoted, already_narrow = [], [], 0
+    for key, provenance in sorted(links.items()):
+        index = positions.get(key)
+        previous = merged[index] if index is not None else None
+        if previous is not None and previous.label == "Narrow":
+            already_narrow += 1
+            continue
+        code, description = available[key]
+        item = ClassifiedCandidate(code, description, key[1], "Narrow", 0.0,
+                                   "Included by direct cross-vocabulary CodeMapper CUI mapping; "
+                                   "not independently classified for EDF scope.",
+                                   decision_origin="codemapper", mapping_provenance=provenance,
+                                   prior_label=previous.label if previous else None,
+                                   prior_confidence=previous.confidence if previous else None)
+        if previous is None:
+            merged.append(item)
+            added.append({"code": code, "vocabulary": key[1]})
+        else:
+            merged[index] = replace(item, retrieval_score=previous.retrieval_score)
+            promoted.append({"code": code, "vocabulary": key[1], "prior_label": previous.label})
+    info = {"enabled": True, "rounds": 1, "descendants_used": False,
+            "direct_links": [{"code": available[key][0], "vocabulary": key[1],
+                              "provenance": provenance} for key, provenance in sorted(links.items())],
+            "seeds": len(seeds), "seed_counts": seed_counts, "unique_cuis": len(by_cui),
+            "unmapped_seeds": unmapped, "new_codes": added, "promoted_codes": promoted,
+            "already_narrow": already_narrow,
+            "unavailable_codes": [{"code": c, "vocabulary": v, "provenance": p}
+                                  for (c, v), p in sorted(unavailable.items())],
+            "calls": client.calls, "cache_hits": client.cache_hits,
+            "seconds": round(time.perf_counter() - started, 3)}
+    return merged, info

src/retrieval/__init__.py

@@ -0,0 +1 @@
+"""Retrieval components: lexical, embedding, and hybrid candidate generation."""

src/retrieval/embeddings.py

@@ -0,0 +1,548 @@
+"""Embedding-based retrieval using local or API-backed encoders."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+import os
+from pathlib import Path
+import re
+from typing import Protocol, Sequence
+from urllib.error import URLError
+from urllib.request import urlopen
+
+import numpy as np
+import pandas as pd
+
+DEFAULT_LOCAL_EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
+DEFAULT_OPENAI_EMBEDDING_MODEL = "text-embedding-3-large"
+
+
+def _create_local_embedding_model(model_name: str):
+    """Import the optional local stack only when that backend is selected."""
+    from sentence_transformers import SentenceTransformer
+
+    return SentenceTransformer(model_name)
+
+
+class EmbeddingEncoder(Protocol):
+    """Common interface for local and API-backed embedding models."""
+
+    def encode(self, sentences, **kwargs): ...
+
+
+class OpenAIEmbeddingModel:
+    """SentenceTransformer-compatible adapter for OpenAI embeddings."""
+
+    def __init__(
+        self,
+        model_name: str = DEFAULT_OPENAI_EMBEDDING_MODEL,
+        *,
+        dimensions: int | None = 3072,
+        batch_size: int = 512,
+    ) -> None:
+        if dimensions is not None and dimensions < 1:
+            raise ValueError("Embedding dimensions must be at least 1")
+        if batch_size < 1:
+            raise ValueError("Embedding batch size must be at least 1")
+        self.model_name = model_name
+        self.dimensions = dimensions
+        self.batch_size = batch_size
+        self._client = None
+
+    def _get_client(self):
+        if self._client is None:
+            import openai
+
+            self._client = openai.OpenAI()
+        return self._client
+
+    def encode(
+        self,
+        sentences,
+        *,
+        normalize_embeddings: bool = True,
+        show_progress_bar: bool = False,
+        **_kwargs,
+    ) -> np.ndarray:
+        del show_progress_bar
+        single = isinstance(sentences, str)
+        inputs = [sentences] if single else [str(item) for item in sentences]
+        if not inputs:
+            width = self.dimensions or 0
+            return np.empty((0, width), dtype=np.float32)
+
+        # Convert each response to float32 immediately. Keeping hundreds of
+        # thousands of embeddings as Python float lists can require many times
+        # more memory than the final NumPy matrix.
+        chunks: list[np.ndarray] = []
+        client = self._get_client()
+        for start in range(0, len(inputs), self.batch_size):
+            request = {
+                "model": self.model_name,
+                "input": inputs[start : start + self.batch_size],
+                "encoding_format": "float",
+            }
+            if self.dimensions is not None:
+                request["dimensions"] = self.dimensions
+            response = client.embeddings.create(**request)
+            chunks.append(
+                np.asarray(
+                    [
+                        item.embedding
+                        for item in sorted(response.data, key=lambda item: item.index)
+                    ],
+                    dtype=np.float32,
+                )
+            )
+
+        result = np.vstack(chunks)
+        if normalize_embeddings:
+            norms = np.linalg.norm(result, axis=1, keepdims=True)
+            result = result / np.maximum(norms, 1e-12)
+        return result[0] if single else result
+
+
+def create_embedding_model(
+    provider: str,
+    model_name: str,
+    *,
+    dimensions: int | None = None,
+    batch_size: int = 512,
+) -> EmbeddingEncoder:
+    """Construct a local SentenceTransformer or OpenAI embedding adapter."""
+    provider = provider.strip().lower()
+    if provider == "local":
+        return _create_local_embedding_model(model_name)
+    if provider == "openai":
+        return OpenAIEmbeddingModel(
+            model_name,
+            dimensions=dimensions,
+            batch_size=batch_size,
+        )
+    raise ValueError(f"Unsupported embedding provider: {provider!r}")
+
+
+def embedding_cache_key(
+    provider: str,
+    model_name: str,
+    dimensions: int | None = None,
+) -> str:
+    """Return a filesystem-safe cache identity without mixing providers/models."""
+    model_key = re.sub(r"[^A-Za-z0-9_.-]+", "_", model_name).strip("_")
+    if provider.strip().lower() == "local":
+        return model_key
+    dimension_key = f"_{dimensions}" if dimensions is not None else ""
+    return f"{provider.strip().lower()}_{model_key}{dimension_key}"
+
+
+@dataclass(slots=True)
+class EmbeddingCandidate:
+    rank: int
+    code: str
+    description: str
+    vocabulary: str
+    score: float
+
+    def __lt__(self, other: object) -> bool:
+        if not isinstance(other, EmbeddingCandidate):
+            return NotImplemented
+        return self.score > other.score
+
+
+@dataclass
+class EmbeddingIndex:
+    """Wraps embedding model + corpus embeddings for fast cosine-sim search."""
+
+    model: EmbeddingEncoder | None
+    descriptions: list[str]
+    codes: list[str]
+    vocabularies: list[str]
+    embeddings: np.ndarray  # shape (n, dim)
+
+    @classmethod
+    def from_codes(
+        cls,
+        codes: pd.DataFrame,
+        model_name: str = DEFAULT_LOCAL_EMBEDDING_MODEL,
+        model: EmbeddingEncoder | None = None,
+        show_progress_bar: bool = False,
+    ) -> "EmbeddingIndex":
+        df = codes.copy()
+        for col in ("code", "vocabulary"):
+            if col not in df.columns:
+                df[col] = ""
+        embedding_model = model or _create_local_embedding_model(model_name)
+        descs = df["description"].fillna("").astype(str).tolist()
+        embs = embedding_model.encode(
+            descs,
+            show_progress_bar=show_progress_bar,
+            normalize_embeddings=True,
+        )
+        return cls(
+            model=embedding_model,
+            descriptions=descs,
+            codes=df["code"].astype(str).tolist(),
+            vocabularies=df["vocabulary"].astype(str).tolist(),
+            embeddings=np.asarray(embs, dtype=np.float32),
+        )
+
+    def save(self, path: str | Path) -> None:
+        """Persist the corpus embeddings + metadata to disk (npz) so they can
+        be reloaded without re-running the embedding model."""
+        path = Path(path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        np.savez_compressed(
+            path,
+            embeddings=np.asarray(self.embeddings, dtype=np.float32),
+            codes=np.array(self.codes, dtype=object),
+            descriptions=np.array(self.descriptions, dtype=object),
+            vocabularies=np.array(self.vocabularies, dtype=object),
+        )
+
+    @classmethod
+    def load(
+        cls,
+        path: str | Path,
+        model_name: str = DEFAULT_LOCAL_EMBEDDING_MODEL,
+        model: EmbeddingEncoder | None = None,
+        load_model: bool = True,
+    ) -> "EmbeddingIndex":
+        """Load a previously-saved embedding index from disk."""
+        with np.load(path, allow_pickle=True) as data:
+            return cls(
+                model=model or (_create_local_embedding_model(model_name) if load_model else None),
+                descriptions=list(data["descriptions"]),
+                codes=list(data["codes"]),
+                vocabularies=list(data["vocabularies"]),
+                embeddings=np.asarray(data["embeddings"], dtype=np.float32),
+            )
+
+    @classmethod
+    def from_cache_or_build(
+        cls,
+        codes: pd.DataFrame,
+        cache_path: str | Path,
+        model_name: str = DEFAULT_LOCAL_EMBEDDING_MODEL,
+        model: EmbeddingEncoder | None = None,
+        load_model: bool = True,
+        show_progress_bar: bool = False,
+        status_callback: Callable[[str], None] | None = None,
+    ) -> "EmbeddingIndex":
+        """Load a cached embedding index if present, otherwise build it from
+        *codes* and cache the result for subsequent runs.
+
+        This is the key mechanism that avoids re-embedding large reference
+        code systems (e.g. the ~74k-row full ICD-10-CM corpus) on every run.
+        """
+        cache_path = Path(cache_path)
+        if cache_path.exists():
+            cached = cls.load(
+                cache_path,
+                model_name=model_name,
+                model=model,
+                load_model=load_model,
+            )
+            expected_codes = codes.get("code", pd.Series(dtype=str)).astype(str).tolist()
+            expected_vocabularies = (
+                codes.get("vocabulary", pd.Series(dtype=str)).astype(str).tolist()
+            )
+            expected_descriptions = codes["description"].fillna("").astype(str).tolist()
+            if (
+                cached.codes == expected_codes
+                and cached.vocabularies == expected_vocabularies
+                and cached.descriptions == expected_descriptions
+            ):
+                return cached
+            subset_positions = _ordered_subset_positions(
+                cached.codes,
+                cached.vocabularies,
+                expected_codes,
+                expected_vocabularies,
+            )
+            if subset_positions is not None and all(
+                cached.descriptions[position] == description
+                for position, description in zip(subset_positions, expected_descriptions)
+            ):
+                # Corpus preprocessing can intentionally remove non-codable
+                # dictionary rows (for example ICD range headers). Reuse the
+                # matching cached vectors instead of re-embedding every code.
+                subset = cls(
+                    model=cached.model,
+                    descriptions=expected_descriptions,
+                    codes=expected_codes,
+                    vocabularies=expected_vocabularies,
+                    embeddings=np.asarray(cached.embeddings[subset_positions], dtype=np.float32),
+                )
+                subset.save(cache_path)
+                return subset
+        if status_callback is not None:
+            status_callback(
+                f"Creating embeddings for {len(codes):,} descriptions; no matching saved index. "
+                "The first run can take several minutes."
+            )
+        index = cls.from_codes(
+            codes,
+            model_name=model_name,
+            model=model,
+            show_progress_bar=show_progress_bar,
+        )
+        index.save(cache_path)
+        return index
+
+    @classmethod
+    def combine(cls, indexes: Sequence["EmbeddingIndex"]) -> "EmbeddingIndex":
+        """Concatenate multiple embedding indexes into a single searchable index."""
+        if not indexes:
+            raise ValueError("At least one embedding index must be provided.")
+
+        descriptions: list[str] = []
+        codes: list[str] = []
+        vocabularies: list[str] = []
+        embeddings: list[np.ndarray] = []
+
+        for index in indexes:
+            descriptions.extend(index.descriptions)
+            codes.extend(index.codes)
+            vocabularies.extend(index.vocabularies)
+            embeddings.append(np.asarray(index.embeddings))
+
+        return cls(
+            model=indexes[0].model,
+            descriptions=descriptions,
+            codes=codes,
+            vocabularies=vocabularies,
+            embeddings=(embeddings[0] if len(embeddings) == 1 else np.vstack(embeddings)).astype(np.float32, copy=False)
+            if embeddings
+            else np.empty((0, 0), dtype=np.float32),
+        )
+
+    @classmethod
+    def load_for_vocabularies(
+        cls,
+        codes: pd.DataFrame,
+        vocabularies: Sequence[str],
+        cache_dir: str | Path,
+        model_name: str = DEFAULT_LOCAL_EMBEDDING_MODEL,
+        model: EmbeddingEncoder | None = None,
+        provider: str = "local",
+        dimensions: int | None = None,
+        progress_callback: Callable[[int, int], None] | None = None,
+        status_callback: Callable[[str], None] | None = None,
+    ) -> "EmbeddingIndex":
+        """Load per-vocabulary embedding caches and concatenate them.
+
+        Each vocabulary is cached independently by vocabulary, provider,
+        model, and output dimensions.
+        Selected vocabularies are then combined in memory for retrieval.
+        """
+        cache_dir = Path(cache_dir)
+        indexes: list[EmbeddingIndex] = []
+        total = len(vocabularies)
+        if progress_callback is not None:
+            progress_callback(0, total)
+        for position, vocabulary in enumerate(vocabularies, start=1):
+            vocab_df = codes[
+                codes["vocabulary"].astype(str).str.upper() == vocabulary.upper()
+            ].reset_index(drop=True)
+            if vocab_df.empty:
+                continue
+            cache_identity = embedding_cache_key(provider, model_name, dimensions)
+            cache_path = cache_dir / f"embeddings_{vocabulary}_{cache_identity}.npz"
+            if status_callback is not None:
+                status_callback(f"Preparing embeddings for {vocabulary} ({len(vocab_df):,} descriptions)...")
+            # Existing release assets contain the legacy local MiniLM indexes.
+            # OpenAI indexes are private, paid-to-generate local artifacts and
+            # must not be mistaken for downloadable release assets.
+            if provider.strip().lower() == "local":
+                _ensure_embedding_asset(cache_path)
+            if status_callback is not None:
+                status_callback(
+                    f"Loading and checking saved embeddings for {vocabulary}..."
+                    if cache_path.exists()
+                    else f"Creating embeddings for {vocabulary} ({len(vocab_df):,} descriptions). "
+                    "The first run can take several minutes."
+                )
+            indexes.append(
+                cls.from_cache_or_build(
+                    vocab_df,
+                    cache_path=cache_path,
+                    model_name=model_name,
+                    model=model,
+                    load_model=False,
+                    show_progress_bar=progress_callback is not None,
+                    status_callback=(
+                        (lambda message, vocabulary=vocabulary: status_callback(f"{vocabulary}: {message}"))
+                        if status_callback is not None else None
+                    ),
+                )
+            )
+            if progress_callback is not None:
+                progress_callback(position, total)
+        if status_callback is not None:
+            status_callback("Combining terminology embeddings...")
+        matrices = [np.asarray(index.embeddings, dtype=np.float32) for index in indexes]
+        return cls(
+            model=None,
+            descriptions=[d for index in indexes for d in index.descriptions],
+            codes=[c for index in indexes for c in index.codes],
+            vocabularies=[v for index in indexes for v in index.vocabularies],
+            embeddings=(matrices[0] if len(matrices) == 1 else np.vstack(matrices))
+            if indexes
+            else np.empty((0, 0), dtype=np.float32),
+        )
+
+    def retrieve(
+        self,
+        query: str,
+        top_k: int | None = None,
+        model: EmbeddingEncoder | None = None,
+    ) -> list[EmbeddingCandidate]:
+        query_model = model or self.model
+        if query_model is None:
+            raise ValueError("An embedding model is required to embed the query.")
+        q_emb = query_model.encode(query, show_progress_bar=False, normalize_embeddings=True)
+        scores = self.embeddings @ q_emb  # cosine similarity (both normalised)
+        if top_k is not None and len(scores) > top_k:
+            top_indices = np.argpartition(scores, -top_k)[-top_k:]
+            sorted_local = sorted(top_indices, key=lambda i: -scores[i])[:top_k]
+        else:
+            sorted_local = np.argsort(-scores)
+
+        return [
+            EmbeddingCandidate(
+                rank=r + 1,
+                code=self.codes[idx],
+                description=self.descriptions[idx],
+                vocabulary=self.vocabularies[idx],
+                score=float(round(scores[idx], 6)),
+            )
+            for r, idx in enumerate(sorted_local)
+        ]
+
+    def score_all(
+        self,
+        query: str,
+        model: EmbeddingEncoder | None = None,
+    ) -> list[float]:
+        """Return corpus-order cosine scores without ranking or object creation.
+
+        Hybrid retrieval needs every embedding score for global min-max
+        normalization, but it does not need the embedding-only ranking. This
+        path preserves the same six-decimal scores as :meth:`retrieve` while
+        avoiding a full sort and hundreds of thousands of temporary
+        ``EmbeddingCandidate`` objects.
+        """
+        query_model = model or self.model
+        if query_model is None:
+            raise ValueError("An embedding model is required to embed the query.")
+        q_emb = query_model.encode(query, show_progress_bar=False, normalize_embeddings=True)
+        scores = self.embeddings @ q_emb
+        return [float(round(score, 6)) for score in scores]
+
+
+def _ordered_subset_positions(
+    cached_codes: Sequence[str],
+    cached_vocabularies: Sequence[str],
+    expected_codes: Sequence[str],
+    expected_vocabularies: Sequence[str],
+) -> list[int] | None:
+    """Locate an order-preserving expected corpus inside a cached corpus."""
+    if (
+        len(expected_codes) != len(expected_vocabularies)
+        or len(expected_codes) > len(cached_codes)
+    ):
+        return None
+    positions: list[int] = []
+    expected_index = 0
+    for cached_index, (code, vocabulary) in enumerate(
+        zip(cached_codes, cached_vocabularies)
+    ):
+        if expected_index >= len(expected_codes):
+            break
+        if (
+            code == expected_codes[expected_index]
+            and vocabulary == expected_vocabularies[expected_index]
+        ):
+            positions.append(cached_index)
+            expected_index += 1
+    return positions if expected_index == len(expected_codes) else None
+
+
+def _release_asset_url(asset_name: str) -> str | None:
+    owner = os.getenv("GITHUB_REPOSITORY_OWNER")
+    repo = os.getenv("GITHUB_REPOSITORY_NAME")
+    tag = os.getenv("GITHUB_RELEASE_TAG")
+    if not owner or not repo or not tag:
+        return None
+    return f"https://github.com/{owner}/{repo}/releases/download/{tag}/{asset_name}"
+
+
+def _ensure_embedding_asset(cache_path: Path) -> None:
+    """Ensure an embedding asset exists locally, downloading it from a GitHub Release if needed."""
+    if cache_path.exists():
+        return
+
+    asset_name = cache_path.name
+    url = _release_asset_url(asset_name)
+    if url is None:
+        return
+
+    cache_path.parent.mkdir(parents=True, exist_ok=True)
+    try:
+        with urlopen(url) as response, cache_path.open("wb") as fh:
+            fh.write(response.read())
+    except URLError:
+        if cache_path.exists():
+            cache_path.unlink(missing_ok=True)
+        raise
+
+
+def retrieve_embeddings(
+    query: str,
+    codes: pd.DataFrame,
+    model_name: str = DEFAULT_LOCAL_EMBEDDING_MODEL,
+    top_k: int | None = None,
+) -> list[EmbeddingCandidate]:
+    """Build an embedding index on *codes* and retrieve candidates for *query*.
+
+    Parameters
+    ----------
+    query : str
+        Phenotype / EDF description to search.
+    codes : pd.DataFrame
+        Corpus with a ``description`` column (and optionally ``code``/``vocabulary``).
+    model_name : str
+        sentence-transformers model identifier.
+    top_k : int, optional
+        Number of results to return.
+
+    Returns
+    -------
+    list[EmbeddingCandidate]
+    """
+    df = codes.copy()
+    for col in ("code", "vocabulary"):
+        if col not in df.columns:
+            df[col] = ""
+
+    descs = df["description"].fillna("").astype(str).tolist()
+    model = _create_local_embedding_model(model_name)
+    embs = model.encode(descs, show_progress_bar=False, normalize_embeddings=True)
+    q_emb = model.encode(query, show_progress_bar=False, normalize_embeddings=True)
+    scores = embs @ q_emb
+
+    sorted_idx = np.argsort(-scores)
+    if top_k is not None:
+        sorted_idx = sorted_idx[:top_k]
+
+    return [
+        EmbeddingCandidate(
+            rank=r + 1,
+            code=df.iloc[idx]["code"],
+            description=descs[idx],
+            vocabulary=str(df.iloc[idx].get("vocabulary", "")),
+            score=float(round(scores[idx], 6)),
+        )
+        for r, idx in enumerate(sorted_idx)
+    ]

src/retrieval/hybrid.py

@@ -0,0 +1,203 @@
+"""Hybrid retrieval combining lexical (BM25) and embedding (semantic) scores."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import pandas as pd
+
+from src.retrieval.embeddings import EmbeddingEncoder, EmbeddingIndex, retrieve_embeddings
+from src.retrieval.lexical import build_lexical_index, score_lexical
+
+
+@dataclass(slots=True)
+class HybridCandidate:
+    rank: int
+    code: str
+    description: str
+    vocabulary: str
+    lexical_score: float
+    embedding_score: float
+    score: float
+
+    def __lt__(self, other: object) -> bool:
+        if not isinstance(other, HybridCandidate):
+            return NotImplemented
+        return self.score > other.score
+
+
+def _min_max_normalize(values: list[float]) -> list[float]:
+    if not values:
+        return values
+    lo, hi = min(values), max(values)
+    if hi - lo < 1e-12:
+        return [0.0 for _ in values]
+    return [(v - lo) / (hi - lo) for v in values]
+
+
+def hybrid_retrieval(
+    query: str,
+    codes: pd.DataFrame,
+    lexical_weight: float = 0.5,
+    embedding_weight: float = 0.5,
+    embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2",
+    top_k: int | None = 25,
+    per_vocabulary_top_k: bool = True,
+    embedding_index: EmbeddingIndex | None = None,
+    query_model: EmbeddingEncoder | None = None,
+    lexical_query: str | None = None,
+    embedding_query: str | None = None,
+) -> list[HybridCandidate]:
+    """Combine BM25 lexical retrieval with embedding-based semantic retrieval.
+
+    Both signals are min-max normalized to [0, 1] and combined with the given
+    weights before ranking. The combined ranking is truncated to *top_k*.
+
+    Parameters
+    ----------
+    query : str
+        Phenotype / EDF text used as the search query.
+    codes : pd.DataFrame
+        Candidate code corpus with ``code``, ``description``, ``vocabulary`` columns.
+    lexical_weight, embedding_weight : float
+        Relative weighting of each retrieval signal. Need not sum to 1.
+    embedding_model : str
+        sentence-transformers model identifier used when *embedding_index* is None.
+    top_k : int, optional
+        Number of ranked results to return per vocabulary. ``None`` returns the
+        full corpus.
+    per_vocabulary_top_k : bool
+        If ``True`` (default), apply ``top_k`` separately within each selected
+        vocabulary. If ``False``, apply ``top_k`` after combining all selected
+        vocabularies.
+    embedding_index : EmbeddingIndex, optional
+        Pre-built embedding index to reuse across multiple queries.
+    query_model : EmbeddingEncoder, optional
+        Cached local or API-backed encoder used to embed the EDF/query text.
+    lexical_query : str, optional
+        Text used for the BM25 lexical score. If omitted, *query* is used.
+    embedding_query : str, optional
+        Text used for the embedding score. If omitted, *query* is used.
+
+    Returns
+    -------
+    list[HybridCandidate]
+    """
+    if codes.empty:
+        return []
+
+    df = codes.reset_index(drop=True).copy()
+    for col in ("code", "vocabulary"):
+        if col not in df.columns:
+            df[col] = ""
+
+    lexical_query = lexical_query or query
+    embedding_query = embedding_query or query
+    embedding_scores_by_code: dict[tuple[str, str], float] | None = None
+    if embedding_index is not None:
+        # Score the combined corpus exactly once. Previously this full matrix
+        # multiplication, full sort, and candidate-object construction ran
+        # once per vocabulary even though every pass produced identical raw
+        # scores. The corpus is deduplicated on (code, vocabulary), so the map
+        # is identical to the maps produced by the former repeated searches.
+        all_embedding_scores = embedding_index.score_all(embedding_query, model=query_model)
+        embedding_scores_by_code = {
+            (code, vocabulary): score
+            for code, vocabulary, score in zip(
+                embedding_index.codes,
+                embedding_index.vocabularies,
+                all_embedding_scores,
+            )
+        }
+
+    per_vocab_candidates: list[
+        tuple[str, list[tuple[str, str, str, float, float]]]
+    ] = []
+    for vocabulary, vocab_df in df.groupby("vocabulary", sort=False):
+        lexical_index = build_lexical_index(vocab_df)
+        lexical_scores = score_lexical(lexical_query, lexical_index, len(vocab_df))
+
+        if embedding_scores_by_code is None:
+            embedding_results = retrieve_embeddings(embedding_query, vocab_df, model_name=embedding_model, top_k=None)
+            vocabulary_embedding_scores = {
+                (c.code, c.vocabulary): c.score for c in embedding_results
+            }
+        else:
+            vocabulary_embedding_scores = embedding_scores_by_code
+
+        codes_values = vocab_df["code"].tolist()
+        description_values = vocab_df["description"].tolist()
+        vocabulary_values = vocab_df["vocabulary"].tolist()
+        vocab_candidates = [
+            (
+                code,
+                description,
+                candidate_vocabulary,
+                lexical_score,
+                vocabulary_embedding_scores.get((code, candidate_vocabulary), 0.0),
+            )
+            for code, description, candidate_vocabulary, lexical_score in zip(
+                codes_values,
+                description_values,
+                vocabulary_values,
+                lexical_scores,
+            )
+        ]
+        per_vocab_candidates.append((vocabulary, vocab_candidates))
+
+    all_lexical = [
+        lexical
+        for _, vocab_candidates in per_vocab_candidates
+        for _, _, _, lexical, _ in vocab_candidates
+    ]
+    all_embedding = [
+        embedding
+        for _, vocab_candidates in per_vocab_candidates
+        for _, _, _, _, embedding in vocab_candidates
+    ]
+    norm_all_lexical = _min_max_normalize(all_lexical)
+    norm_all_embedding = _min_max_normalize(all_embedding)
+
+    # Keep lightweight tuples until after top-k selection. Constructing a
+    # dataclass (and formerly a Pandas Series) for all 620k rows dominated the
+    # runtime even when only a handful of results were requested.
+    selected: list[tuple[str, str, str, float, float, float]] = []
+    flat_idx = 0
+    for _, vocab_candidates in per_vocab_candidates:
+        scored: list[tuple[str, str, str, float, float, float]] = []
+        for idx, (code, description, vocabulary, _, _) in enumerate(vocab_candidates):
+            norm_lexical = norm_all_lexical[flat_idx + idx]
+            norm_embedding = norm_all_embedding[flat_idx + idx]
+            score = lexical_weight * norm_lexical + embedding_weight * norm_embedding
+            scored.append(
+                (
+                    code,
+                    description,
+                    vocabulary,
+                    round(norm_lexical, 6),
+                    round(norm_embedding, 6),
+                    round(score, 6),
+                )
+            )
+        if per_vocabulary_top_k and top_k is not None:
+            scored.sort(key=lambda candidate: -candidate[5])
+            scored = scored[:top_k]
+        selected.extend(scored)
+        flat_idx += len(vocab_candidates)
+
+    selected.sort(key=lambda candidate: -candidate[5])
+    if not per_vocabulary_top_k and top_k is not None:
+        selected = selected[:top_k]
+    return [
+        HybridCandidate(
+            rank=rank,
+            code=code,
+            description=description,
+            vocabulary=vocabulary,
+            lexical_score=lexical_score,
+            embedding_score=embedding_score,
+            score=score,
+        )
+        for rank, (code, description, vocabulary, lexical_score, embedding_score, score)
+        in enumerate(selected, start=1)
+    ]

src/retrieval/lexical.py

@@ -0,0 +1,141 @@
+"""BM25 lexical retrieval over code descriptions."""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+
+import numpy as np
+import pandas as pd
+from rank_bm25 import BM25Okapi
+
+
+def _tokenize(text: str) -> list[str]:
+    """Simple whitespace + punctuation tokenizer for BM25 indexing."""
+    if not text or not isinstance(text, str):
+        return []
+    text = re.sub(r"[^\w\s]", " ", text.lower())
+    return [tok for tok in text.split() if len(tok) > 1]
+
+
+@dataclass(slots=True)
+class LexicalCandidate:
+    rank: int
+    code: str
+    description: str
+    vocabulary: str
+    score: float
+
+    def __lt__(self, other: object) -> bool:
+        if not isinstance(other, LexicalCandidate):
+            return NotImplemented
+        return self.score > other.score  # higher = more relevant
+
+
+def build_lexical_index(codes: pd.DataFrame) -> BM25Okapi:
+    """Build a BM25 index from the code corpus descriptions.
+
+    Parameters
+    ----------
+    codes : pd.DataFrame
+        Must contain ``description`` column with full description strings.
+
+    Returns
+    -------
+    BM25Okapi
+        The built tokenised index.
+    """
+    doc_tokens: list[list[str]] = []
+    for desc in codes["description"].astype(str):
+        tokens = _tokenize(desc)
+        doc_tokens.append(tokens if tokens else [desc])
+
+    return BM25Okapi(doc_tokens)
+
+
+def score_lexical(query: str, index: BM25Okapi, corpus_size: int) -> list[float]:
+    """Return corpus-order BM25 scores without allocating result objects."""
+    query_tokens = _tokenize(query)
+    if not query_tokens:
+        return [0.0] * corpus_size
+    return [round(float(score), 6) for score in index.get_scores(query_tokens)]
+
+
+def retrieve_lexical(
+    query: str,
+    codes: pd.DataFrame | list[dict],
+    index: BM25Okapi | None = None,
+    top_k: int | None = None,
+) -> list[LexicalCandidate]:
+    """Retrieve candidate codes ranked by BM25 lexical similarity.
+
+    Parameters
+    ----------
+    query : str
+        The phenotype description / EDF summary to search against.
+    codes : pd.DataFrame or list[dict]
+        Code corpus containing ``description`` (mandatory) and optionally
+        ``code``, ``vocabulary`` columns (DataFrame) or dicts with those keys.
+    index : BM25Okapi, optional
+        Pre-built index. If omitted one is built on-the-fly from *codes*.
+    top_k : int, optional
+        Maximum number of results to return.  Defaults to the full corpus size.
+
+    Returns
+    -------
+    list[LexicalCandidate]
+        Ranked candidates sorted descending by BM25 score.
+    """
+    if isinstance(codes, pd.DataFrame):
+        df = codes.copy()
+        keys = ("description", "code", "vocabulary")
+        for k in keys[1:]:
+            if k not in df.columns:
+                df[k] = ""  # type: ignore[operator]
+    else:
+        df = pd.DataFrame(codes)
+        for k in ("code", "vocabulary"):
+            if k not in df.columns:
+                df[k] = ""
+
+    if index is None:
+        index = build_lexical_index(df)
+
+    query_tokens = _tokenize(query)
+    if not query_tokens:
+        # Fallback: rank all codes uniformly (no lexical signal)
+        return [
+            LexicalCandidate(
+                rank=0, code=r["code"], description=r["description"], vocabulary=r["vocabulary"], score=0.0
+            )
+            for i, r in df.iterrows()
+        ]
+
+    scores = index.get_scores(query_tokens).tolist()  # type: ignore[union-attr]
+
+    if top_k is not None:
+        top_indices = np.argpartition(scores, -top_k)[-top_k:]
+        top_scores = [scores[i] for i in sorted(top_indices, key=lambda x: -scores[x])]
+        results: list[LexicalCandidate] = [
+            LexicalCandidate(
+                rank=r + 1,
+                code=df.iloc[idx]["code"],
+                description=df.iloc[idx]["description"],
+                vocabulary=df.iloc[idx].get("vocabulary", ""),
+                score=round(float(s), 6),
+            )
+            for r, (idx, s) in enumerate(sorted(enumerate(top_scores), key=lambda x: -x[1]))
+        ]
+    else:
+        results = [
+            LexicalCandidate(
+                rank=r + 1,
+                code=df.iloc[r]["code"],
+                description=df.iloc[r]["description"],
+                vocabulary=df.iloc[r].get("vocabulary", ""),
+                score=round(float(scores[r]), 6),
+            )
+            for r in range(len(df))
+        ]
+
+    return results

src/uncertainty/__init__.py

@@ -0,0 +1 @@
+"""Uncertainty-based selection of candidates for human review."""

src/uncertainty/selection.py

@@ -0,0 +1,76 @@
+"""Select uncertain classifications requiring targeted human review."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from src.llm.classify import ClassifiedCandidate
+
+
+@dataclass(slots=True)
+class ReviewItem:
+    code: str
+    description: str
+    vocabulary: str
+    label: str
+    confidence: float
+    explanation: str
+    reason: str
+
+
+def select_uncertain(
+    classified: list[ClassifiedCandidate],
+    confidence_threshold: float = 0.7,
+    possible_requires_review: bool = False,
+) -> list[ReviewItem]:
+    """Select classified candidates that require human validation.
+
+    A candidate is flagged for review when:
+    - its LLM confidence is below *confidence_threshold*, or
+    - it was classified as 'Possible' and *possible_requires_review* is True.
+
+    Parameters
+    ----------
+    classified : list[ClassifiedCandidate]
+        Output of :func:`src.llm.classify.llm_classify`.
+    confidence_threshold : float
+        Minimum confidence below which a candidate is routed for review.
+    possible_requires_review : bool
+        Whether 'Possible' classifications are always routed for review,
+        regardless of confidence.
+
+    Returns
+    -------
+    list[ReviewItem]
+        Items requiring human review, each annotated with the reason.
+    """
+    review_items: list[ReviewItem] = []
+    for c in classified:
+        reasons: list[str] = []
+        if c.decision_origin == "codemapper":
+            reasons.append("included by CodeMapper mapping; EDF scope not independently assessed")
+        if c.confidence < confidence_threshold:
+            reasons.append(f"confidence {c.confidence:.2f} below threshold {confidence_threshold:.2f}")
+        if possible_requires_review and c.label == "Possible":
+            reasons.append("classified as 'Possible'")
+
+        if reasons:
+            review_items.append(
+                ReviewItem(
+                    code=c.code,
+                    description=c.description,
+                    vocabulary=c.vocabulary,
+                    label=c.label,
+                    confidence=c.confidence,
+                    explanation=c.explanation,
+                    reason="; ".join(reasons),
+                )
+            )
+    return review_items
+
+
+def review_rate(classified: list[ClassifiedCandidate], review_items: list[ReviewItem]) -> float:
+    """Fraction of classified candidates flagged for human review."""
+    if not classified:
+        return 0.0
+    return len(review_items) / len(classified)

src/utils/config.py

@@ -0,0 +1,135 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import yaml
+
+
+@dataclass
+class PathsConfig:
+    edf_dir: str = "data/edfs/yamls"
+    codes_parquet_dir: str = "data/codes/parquet/snomed_definitive_20261030"
+    embeddings_dir: str = "data/codes/embeddings/snomed_definitive_20261030"
+    results_dir: str = "results"
+    aesi_datasets: dict[str, str] = field(default_factory=dict)
+    code_systems: dict[str, str] = field(
+        default_factory=lambda: {
+            "ICD10CM": "data/codes/csv/ICD10CM@2026-codes.csv",
+            "ICD9CM": "data/codes/csv/ICD9CM@2014-codes.csv",
+            "ICPC": "data/codes/csv/ICPC@1993-codes.csv",
+            "MDR": "data/codes/csv/MDR@28.0-codes.csv",
+            "RCD2": "data/codes/csv/RCD2@20200401-codes.csv",
+            "SNOMEDCT_US": "data/codes/csv/SNOMEDCT_US@2025_09_01-codes definitive.csv",
+        }
+    )
+
+
+@dataclass
+class RetrievalConfig:
+    lexical_weight: float = 0.1
+    embedding_weight: float = 0.9
+    embedding_provider: str = "openai"
+    embedding_model: str = "text-embedding-3-large"
+    embedding_dimensions: int | None = 3072
+    embedding_batch_size: int = 512
+    top_k: int = 50
+    use_top_k_limit: bool = False
+    per_vocabulary_top_k: bool = True
+    vocabularies: list[str] = field(
+        default_factory=lambda: ["ICD10CM", "ICD9CM", "ICPC", "MDR", "RCD2", "SNOMEDCT_US"]
+    )
+
+
+@dataclass
+class LLMConfig:
+    provider: str = "jev"  # Stage 2 defaults to Jev; other providers remain available.
+    anthropic_model: str = "claude-opus-4-8"
+    openai_model: str = "gpt-6-luna"
+    google_model: str = "gemini-3-flash-preview"
+    jev_model: str = "jev-latest"
+    openai_decisions_model: str = "gpt-6-luna"
+    laya_model: str = "multilingual"
+    jpt_model: str = "kirp/jpt-4b"
+    mlx_model: str = "mlx-community/Qwen3.5-9B-8bit"
+    rank_effort: str = "low"
+    classify_effort: str = "medium"
+    max_retries: int = 3
+    use_possible_category: bool = False
+    adaptive_stopping_enabled: bool = True
+    sparse_narrow_threshold: int = 0
+    consecutive_sparse_batches: int = 10
+    minimum_batches: int = 3
+
+
+@dataclass
+class UncertaintyConfig:
+    review_provider: str = "openai"
+    review_model: str | None = None
+    confidence_threshold: float = 0.7
+    gpt_review_max_candidates: int = 200
+    possible_requires_review: bool = False
+    gpt_review_enabled: bool = True
+    gpt_review_batch_size: int = 10
+    gpt_review_reasoning_enabled: bool = True
+
+
+@dataclass
+class EvaluationConfig:
+    narrow_only_as_positive: bool = True
+    compare_vocabularies: list[str] = field(
+        default_factory=lambda: ["ICD10CM", "ICD9CM", "ICPC", "MDR", "RCD2", "SNOMEDCT_US"]
+    )
+
+
+@dataclass
+class CodeMapperConfig:
+    enabled: bool = False
+    base_url: str = "https://app.vac4eu.org/codemapper-testing/api"
+    cache_dir: str = "results/codemapper/cache"
+    batch_size: int = 100
+    requests_per_minute: float = 6
+    timeout_seconds: float = 30
+    max_retries: int = 3
+    vocabulary_versions: dict[str, str] = field(default_factory=lambda: {
+        "ICD10CM": "2026", "ICD9CM": "2014", "ICPC": "1993",
+        "MDR": "28.0", "RCD2": "20200401", "SNOMEDCT_US": "2025_09_01",
+    })
+
+    def __post_init__(self):
+        if self.batch_size < 1 or self.requests_per_minute <= 0:
+            raise ValueError("CodeMapper batch size and request rate must be positive")
+        if self.timeout_seconds <= 0 or self.max_retries < 1:
+            raise ValueError("CodeMapper timeout and retry count must be positive")
+        if not self.base_url.startswith("https://"):
+            raise ValueError("CodeMapper requires an HTTPS base URL")
+
+
+@dataclass
+class PipelineConfig:
+    phenotype: str = "myocarditis"
+    verbose: bool = True
+    paths: PathsConfig = field(default_factory=PathsConfig)
+    retrieval: RetrievalConfig = field(default_factory=RetrievalConfig)
+    llm: LLMConfig = field(default_factory=LLMConfig)
+    uncertainty: UncertaintyConfig = field(default_factory=UncertaintyConfig)
+    evaluation: EvaluationConfig = field(default_factory=EvaluationConfig)
+    codemapper: CodeMapperConfig = field(default_factory=CodeMapperConfig)
+
+
+def load_config(path: str | Path) -> PipelineConfig:
+    """Load a YAML config file into a PipelineConfig, falling back to defaults
+    for any keys not present in the file."""
+    path = Path(path)
+    raw = yaml.safe_load(path.read_text()) or {}
+
+    return PipelineConfig(
+        phenotype=raw.get("phenotype", PipelineConfig.phenotype),
+        verbose=raw.get("verbose", PipelineConfig.verbose),
+        paths=PathsConfig(**raw.get("paths", {})),
+        retrieval=RetrievalConfig(**raw.get("retrieval", {})),
+        llm=LLMConfig(**raw.get("llm", {})),
+        uncertainty=UncertaintyConfig(**raw.get("uncertainty", {})),
+        evaluation=EvaluationConfig(**raw.get("evaluation", {})),
+        codemapper=CodeMapperConfig(**raw.get("codemapper", {})),
+    )

src/utils/logging.py

@@ -0,0 +1,32 @@
+import logging
+import sys
+
+_CONFIGURED = False
+
+# Silence noisy third-party libraries (sentence-transformers/httpx/huggingface_hub
+# log INFO-level HTTP requests and an "unauthenticated requests" warning on every
+# model load) — none of this is actionable for local, cached model usage.
+_QUIET_LOGGERS = ("sentence_transformers", "httpx", "huggingface_hub")
+
+
+def get_logger(name: str) -> logging.Logger:
+    global _CONFIGURED
+    if not _CONFIGURED:
+        logging.basicConfig(
+            level=logging.INFO,
+            format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
+            stream=sys.stdout,
+        )
+        for quiet_logger in _QUIET_LOGGERS:
+            logging.getLogger(quiet_logger).setLevel(logging.ERROR)
+        try:
+            # huggingface_hub emits its "unauthenticated requests" notice through
+            # its own verbosity control rather than a plain logger level, so the
+            # logger.setLevel above alone doesn't always suppress it.
+            from huggingface_hub.utils import logging as hf_logging
+
+            hf_logging.set_verbosity_error()
+        except ImportError:
+            pass
+        _CONFIGURED = True
+    return logging.getLogger(name)