6krrt/src/local_encoder.py

"""Zero-shot category classification via sentence-embedding + nearest-centroid.

Backs ``classifier.mode: local_encoder``. Exists to structurally rule out the
one classifier failure mode that has cost this project two generations of
local model already (see ``docs/local-models.md``): a generative model
spending its output budget on an unbounded reasoning trace and returning no
parseable JSON. An embedding-scored classification cannot exhibit that
failure — it scores a fixed label set against the input; there is no trace
to run away.

Zero-shot at the core, with one optional, committed exception: a fitted
logistic-regression head (``_TrainableHead`` below) can sit on top of the
same frozen embeddings, trained offline on a synthetic corpus in
agent-traffic shape (``scripts/train_encoder_head.py``). No raw task text is
ever stored — the corpus is generated from the committed eval task set, and
the fitted weights are a small, diffable JSON array. When the coefficients
file is absent, or does not cover the configured model and candidate
categories, the zero-shot nearest-centroid path below runs unchanged.

Originally used HuggingFace's ``zero-shot-classification`` pipeline
(``facebook/bart-large-mnli`` — a cross-encoder NLI model). Replaced with a
sentence-embedding + nearest-centroid architecture for Wave 5.1 of the
token-waste reduction plan (``plans/token-waste-waves.md``). An embedding
model does one forward pass for the input, then cheap cosine similarity
against precomputed category-description embeddings — cost independent of
category count — while the NLI pipeline did one cross-encoder pass *per
category*. This also affords a stronger backbone on the same CPU/RAM budget:
``BAAI/bge-large-en-v1.5`` (the current default) is a general-purpose
embedding model known for strong retrieval/classification performance and
good CPU viability at 1.3 GB.

``transformers``/``torch`` are imported LAZILY, inside ``classify_zero_shot``,
never at module import time — the same rule ``tui.py`` follows for
``textual``: "only imported by tui.py — never by the service dispatch path,
so the router itself has no UI dependency." A deployment that never selects
``local_encoder`` mode needs neither package installed.
"""

from __future__ import annotations

import json
import math
import re
from pathlib import Path
from typing import Any, Final, Optional

import logs

_MISSING_DEPENDENCY_MESSAGE = (
    "classifier.mode is 'local_encoder' but the 'transformers'/'torch' "
    "packages are not installed. Install them with: "
    "pip install -r requirements-encoder.txt"
)

# HF's zero-shot pipeline scored each candidate label against the input using
# a hypothesis template ("This example is {}."), so the label itself had to
# read as natural language for entailment scoring. The embedding-based
# backend reuses the same natural-language descriptions the NLI path used —
# they serve as description embeddings for the centroid comparison. A
# category missing from this map falls back to its raw string rather than
# raising, so a newly-added config category degrades gracefully instead of
# crashing.
_CATEGORY_DESCRIPTIONS: dict[str, str] = {
    "coding_general": "writing new code",
    "coding_refactor": (
        "refactoring or restructuring existing code without changing its behavior"
    ),
    "debugging": "finding and fixing a bug in code",
    "docs_writing": "writing documentation, comments, or explanations",
    "summarization": "summarizing or condensing a long text",
    "file_summarization": "summarizing the contents of a file",
    "diff_checking": "reviewing a code diff or comparing code changes",
    "translation": "translating text from one language to another",
    "reasoning_math": "solving a math problem or a logical reasoning puzzle",
    "tool_use_agentic": (
        "a multi-step task that requires using tools, running commands, "
        "or reading and writing files"
    ),
    "general_chat": "casual conversation or a general question",
}

# Query-side instruction prefix for BAAI/bge-large-en-v1.5 and other BGE
# models. BGE's training uses asymmetric instruction prefixes: the input text
# (the "query") gets prefixed, while category descriptions (the "passages")
# are encoded as-is. This is harmless for non-BGE embedding models — the
# extra tokens are just part of the input context.
_BGE_QUERY_INSTRUCTION = "Represent this sentence for searching relevant passages: "

# Softmax temperature for converting cosine similarities to confidence
# scores. Cosine similarities for this task typically span [0.30, 0.90]
# (unrelated vs closely related). At T=0.10 the softmax sharply amplifies
# the winner: a similarity of 0.80 vs 0.50 runners-up gives the winner ~0.95
# confidence, while a close match (0.75 vs 0.72) produces a more balanced
# spread. This makes the confidence scale meaningful and comparable to the
# old NLI pipeline's [0,1] output — /metrics' classifier-degradation-share
# warning and config.py's confidence_min validator both depend on
# confidence being a genuine probability, not a monotonic score.
_SOFTMAX_TEMPERATURE = 0.10

# ── Task-text isolation (structural agent-session noise) ──────────────
#
# Production task text is raw user_content (dispatcher passes the RAW
# task, not the clamped/framed user_content): fenced code blocks, tool
# call/result stanzas, and system-reminder-style wrappers surround the
# prose instruction. Mean pooling weights every token equally, so when
# noise tokens outnumber the instruction words the pooled embedding —
# and therefore the verdict — describes the noise, not the intent.
# Observed on the real BAAI/bge-large-en-v1.5 (2026-09-19): 8 clean
# one-sentence tasks scored 8/8 correct, but the same instructions
# wrapped in a fenced code block + a "Tool result:" line + a
# <system-reminder> tag classified 2/8 — every non-summarization
# instruction flipped to a file/code-flavored category. Tail-biased
# windowing (the fit below) does NOT fix this: measured 2/8 on long
# noisy inputs too, because the window keeps code and instruction
# mixed. The helper below strips that structural noise BEFORE the fit +
# embed pass; cleaning first also means a noisy task often fits the
# window outright, so the fit sees prose, not a code wall.
#
# Deliberately conservative, pure, stdlib-only (``re``), deterministic,
# and free of false strip targets for ordinary prose: inline code spans
# (``like this``) are KEPT — instructions reference real file names and
# commands, which are signal, not noise.

# Fenced code blocks: strip the content outright. Empirically decided
# placeholder vs pure removal on the real model (2026-09-19, 8 noisy
# pairs x short+long): removal scored 8/8 on both, while '[code]' and
# '[elided]' placeholders scored 7/8 and 6/8 on long pairs — the
# placeholder token itself pulls the pooled embedding toward code-
# flavored categories and costs accuracy, so the winner is removal.
_FENCE_PAIR_RE: Final = re.compile(r"```.*?```", re.DOTALL)
_FENCE_OPEN_TAIL_RE: Final = re.compile(r"```.*\Z", re.DOTALL)

# System-reminder-style wrapper tags: agent harnesses inject session
# metadata inside custom tags; the content is context, not the task.
# Closed-tag spans only — an unclosed tag is left alone rather than
# guessing where it ends (the floor guard still protects the verdict).
_ISOLATE_WRAPPER_TAGS: Final[tuple[str, ...]] = ("system-reminder",)
_WRAPPER_TAG_RE: Final = re.compile(
    r"<({})\b[^>]*>.*?</\s*\1\s*>".format("|".join(_ISOLATE_WRAPPER_TAGS)),
    re.DOTALL | re.IGNORECASE,
)

# Tool-result lines: "Tool result: file read successfully, 12 lines" and
# the call/output variants seen in agent-session transcripts. Line-level
# and conservative: the line must START with the marker (case-
# insensitive), so prose merely mentioning a tool result survives.
_TOOL_LINE_RE: Final = re.compile(
    r"^[ \t]*tool[\s_-]*(?:result|output|call)[\s_-]*:[^\n]*\n?",
    re.IGNORECASE | re.MULTILINE,
)

# Floor guard: if stripping leaves (almost) nothing, the input was
# itself one big code block — e.g. a bare file dump with no instruction.
# Embedding a bare placeholder would score a meaningless stub, so fall
# back to the original text, whose code content is then the only signal
# available. Measured (2026-09-19, real model): code-GROUNDED tasks
# (instruction + pasted code) classify BETTER stripped even though
# stripping keeps well under 20% of the characters, so the guard keys on
# an absolute character floor only — a ratio guard would revert exactly
# the noise-wrapped short inputs the isolation exists to fix.
_ISOLATE_MIN_CHARS: Final[int] = 24


def _isolate_task_text(task: str) -> str:
    """Strip structural agent-session noise from *task* before embedding.

    Removes fenced code block bodies, system-reminder-style wrapper tag
    spans, and tool call/result lines — the shapes that dominate a
    mean-pooled embedding when they outnumber the prose instruction.
    Inline code spans and ordinary prose are kept. Pure string in,
    string out; no I/O, no state.

    If stripping leaves fewer than _ISOLATE_MIN_CHARS — the input was
    essentially only noise (a bare code dump, no instruction) — the
    original text is returned: the code content is then the only signal
    there is, and a stripped stub would be a worse embedding target.

    Idempotent: already-isolated text contains none of the stripped
    shapes, so a second pass returns it unchanged.
    """
    if not task:
        return task
    text = _FENCE_PAIR_RE.sub("", task)
    text = _FENCE_OPEN_TAIL_RE.sub("", text)
    text = _WRAPPER_TAG_RE.sub("", text)
    text = _TOOL_LINE_RE.sub("", text)
    if len(text) < _ISOLATE_MIN_CHARS:
        return task
    return text


# ── Token-budget fitting (pre-tokenization window selection) ──────────
#
# The final ``tokenizer(...)`` call in ``classify_zero_shot`` passes
# ``truncation=True`` with no ``max_length``, so HuggingFace truncates to
# the tokenizer's ``model_max_length`` — 512 tokens for the BERT-family
# BGE models — and, with the default ``truncation_side='right'``, KEEPS
# the head. The head of an agent-session prompt is thousands of chars of
# shared boilerplate while the task-specific content lives in the tail,
# so every long request scored the same truncated prefix and returned the
# same verdict (observed live on 2026-09-18: six distinct 6.4k-6.6k-char
# requests, all classified ``diff_checking`` at confidence 0.192363).
# The helpers below fit the task to the real token window BEFORE that
# final tokenize, choosing ourselves which tokens survive.

# Tail-biased head+tail split: when the task overflows the budget, the
# head slice keeps this share of the content budget and the tail gets the
# rest. A head slice survives at all because the raw task text can carry
# its instruction at either end ("Translate this: <doc>" vs "<doc> —
# translate this") — the same rationale as dispatcher.clamp_for_classifier,
# and it applies here directly because the local_encoder path receives the
# RAW task (dispatcher._classify_via_configured_mode), not the clamped/
# framed user_content the chat modes see. The head share stays small
# because the live failure mode was the opposite collapse: shared head
# boilerplate crowding out the tail.
_FIT_HEAD_FRACTION: Final[float] = 0.25

# Marker spliced between the kept head and tail, mirroring
# dispatcher.clamp_for_classifier's elision convention. Its own token
# cost is measured per tokenizer and reserved from the budget.
_FIT_ELISION_MARKER: Final[str] = "\n\n[... content elided ...]\n\n"

# decode() -> re-tokenize is not guaranteed token-identical for every
# input (unicode normalization, odd graphemes), so the fitted text is
# built a few tokens UNDER budget. If drift ever exceeds the margin, the
# final ``truncation=True`` clips a token or two off the tail — a
# bounded, graceful degradation, never a crash.
_FIT_SAFETY_MARGIN_TOKENS: Final[int] = 8

# model_max_length values some tokenizers report when the true window is
# unknown (HF's VERY_LARGE_INTEGER is int(1e30)). Unusable for budget
# math; falls back to the documented constant below.
_MODEL_MAX_LENGTH_SENTINEL: Final[int] = int(1e30)

# Conservative window used when model_max_length is missing or a
# sentinel: 512 is the BERT-family window the BGE models actually have.
_FALLBACK_MODEL_MAX_LENGTH: Final[int] = 512

# Set once the fallback above has been logged, so a misconfigured
# tokenizer logs one warning instead of one per request.
_budget_fallback_logged = False


def _special_tokens_allowance(tokenizer: Any) -> int:
    """Tokens the final tokenize adds around the text ([CLS]/[SEP] for
    BERT-family tokenizers, including the BGE models)."""
    add_fn = getattr(tokenizer, "num_special_tokens_to_add", None)
    if add_fn is None:
        return 2
    return max(0, int(add_fn()))


def _usable_token_budget(tokenizer: Any) -> int:
    """Tokens available for task content under the model's real window.

    usable = model_max_length - len(BGE instruction prefix tokens)
             - special tokens added by the final tokenize.

    A missing or sentinel ``model_max_length`` falls back to
    _FALLBACK_MODEL_MAX_LENGTH (logged once) rather than raising — an
    unusual tokenizer must degrade conservatively, not kill the request.
    """
    global _budget_fallback_logged
    max_len = getattr(tokenizer, "model_max_length", None)
    if (
        not isinstance(max_len, int)
        or max_len <= 0
        or max_len >= _MODEL_MAX_LENGTH_SENTINEL
    ):
        if not _budget_fallback_logged:
            _budget_fallback_logged = True
            logs.warning(
                "encoder_token_budget_fallback",
                reported_model_max_length=str(max_len),
                fallback=_FALLBACK_MODEL_MAX_LENGTH,
            )
        max_len = _FALLBACK_MODEL_MAX_LENGTH
    prefix_ids = tokenizer(
        _BGE_QUERY_INSTRUCTION, add_special_tokens=False,
    )["input_ids"]
    return max_len - len(list(prefix_ids)) - _special_tokens_allowance(tokenizer)


def _fit_task_to_token_budget(task: str, tokenizer: Any) -> str:
    """Fit *task* into the model's token window, biased toward the tail.

    Tokenizes the bare task (no special tokens, no truncation) and, when
    it overflows the usable budget, keeps a small head slice plus a
    tail-majority window (see _FIT_HEAD_FRACTION for why a head slice
    survives at all), spliced with the elision marker. The window is
    selected on real token ids and decoded back to text — no
    chars-per-token arithmetic anywhere.

    Returns *task* unchanged when it already fits (including the empty
    string). Never mutates tokenizer state — the cached tokenizer is
    shared across requests — and the final ``truncation=True`` in
    classify_zero_shot stays as a last-resort backstop.
    """
    budget = _usable_token_budget(tokenizer)
    if budget <= 0:
        # Pathological window (smaller than prefix + specials): leave the
        # task alone; the final tokenize does whatever it can.
        return task
    ids = list(tokenizer(task, add_special_tokens=False)["input_ids"])
    if len(ids) <= budget:
        return task
    marker_tokens = len(list(
        tokenizer(_FIT_ELISION_MARKER, add_special_tokens=False)["input_ids"]
    ))
    content_budget = budget - marker_tokens - _FIT_SAFETY_MARGIN_TOKENS
    if content_budget <= 0:
        # Degenerate window: no room for the marker or the margin — keep
        # a pure tail, which is where task-specific content lives.
        return tokenizer.decode(ids[-budget:], skip_special_tokens=True)
    head_n = int(content_budget * _FIT_HEAD_FRACTION)
    tail_n = content_budget - head_n
    head_text = tokenizer.decode(ids[:head_n], skip_special_tokens=True)
    tail_text = tokenizer.decode(ids[-tail_n:], skip_special_tokens=True)
    return head_text + _FIT_ELISION_MARKER + tail_text


# Cached loaded model + tokenizer, keyed by (model_id, device) so a config
# change to either picks up a fresh instance rather than silently reusing
# one built for a different model.
_model_cache: dict[tuple[str, str], tuple[Any, Any]] = {}

# Precomputed L2-normalized description embeddings, keyed by
# (model_id, device). Built once when the model is first loaded and cached
# for the lifetime of the process — descriptions don't change per call, only
# the input text needs embedding each time.
_desc_cache: dict[tuple[str, str], dict[str, Any]] = {}

# Per-model pooling strategy ("cls" or "mean") read from the model
# snapshot's 1_Pooling/config.json at load time. Populated by
# _load_model_tokenizer; consumed by _mean_pool.
_model_pooling_strategies: dict[str, str] = {}

# Per-model query prefix read from tokenizer_config.json's "prompts" dict
# at load time. Fallback is _BGE_QUERY_INSTRUCTION when empty.
_model_query_prefixes: dict[str, str] = {}


def _read_pooling_strategy(model_id: str) -> str:
    """Read pooling strategy from model snapshot's ``1_Pooling/config.json``.

    Returns ``"cls"`` when ``pooling_mode_cls_token`` is true, otherwise
    ``"mean"``. Falls back to ``"mean"`` when the file is absent, unreadable,
    or unparseable (e.g. under test when transformers is mocked).
    """
    try:
        from transformers.utils.hub import cached_file

        path = cached_file(model_id, "1_Pooling/config.json")
        if path is None:
            return "mean"
        with open(path) as f:
            config = json.load(f)
        if config.get("pooling_mode_cls_token", False):
            return "cls"
        return "mean"
    except Exception:
        return "mean"


def _read_query_prefix(model_id: str) -> str:
    """Read query prefix from ``tokenizer_config.json``'s ``"prompts"`` dict.

    BGE models ship a ``"prompts"`` key with a query instruction; GTE models
    have an empty or absent prompts dict. Returns the query text when found,
    empty string otherwise.
    """
    try:
        from transformers.utils.hub import cached_file

        path = cached_file(model_id, "tokenizer_config.json")
        if path is None:
            return ""
        with open(path) as f:
            config = json.load(f)
        prompts = config.get("prompts", {})
        if isinstance(prompts, dict):
            query_prompt = prompts.get("query", "")
            if isinstance(query_prompt, dict):
                query_prompt = query_prompt.get("text", "")
            if query_prompt:
                return str(query_prompt)
        return ""
    except Exception:
        return ""


def _pool_embeddings(
    token_embeddings: Any,
    attention_mask: Any,
    strategy: str,
) -> Any:
    """Pool token embeddings according to *strategy*.

    ``"cls"`` returns the [CLS] token embedding (first token of each
    sequence). Otherwise performs attention-mask-weighted mean pooling.
    """
    if strategy == "cls":
        return token_embeddings[:, 0]
    input_mask_expanded = (
        attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
    )
    sum_embeddings = (token_embeddings * input_mask_expanded).sum(dim=1)
    sum_mask = input_mask_expanded.sum(dim=1)
    return sum_embeddings / sum_mask.clamp(min=1e-9)


def _mean_pool(
    token_embeddings: Any, attention_mask: Any, model_id: str
) -> Any:
    """Mean-pool token embeddings weighted by the attention mask.

    Masks out padding tokens so the mean reflects only real content.

    The pooling strategy (CLS vs mean) is read from the model snapshot's
    ``1_Pooling/config.json`` at load time and stored in
    ``_model_pooling_strategies`` keyed by model_id. This function
    dispatches to the appropriate strategy via the explicit *model_id*
    — NOT the last-inserted value in the global dict — so it works
    correctly when more than one model is loaded in a single process.
    """
    strategy = _model_pooling_strategies.get(model_id, "mean")
    return _pool_embeddings(token_embeddings, attention_mask, strategy)


def _l2_normalize(embeddings: Any) -> Any:
    """L2-normalize embeddings along the feature dimension."""
    return embeddings / embeddings.norm(dim=1, keepdim=True).clamp(min=1e-9)


def _build_description_embeddings(
    model: Any,
    tokenizer: Any,
    device_str: str,
    categories: list[str],
    model_id: str,
) -> dict[str, Any]:
    """Precompute L2-normalized embeddings for each category description.

    Each description is encoded through the model, mean-pooled, and
    L2-normalised. Results are cached so only the input text needs to be
    encoded per call.
    """
    import torch

    desc_embeddings: dict[str, Any] = {}
    for cat in categories:
        desc = _CATEGORY_DESCRIPTIONS.get(cat, cat)
        encoded = tokenizer(
            desc,
            padding=True,
            truncation=True,
            return_tensors="pt",
        )
        encoded = {k: v.to(device_str) for k, v in encoded.items()}
        with torch.no_grad():
            outputs = model(**encoded)
        emb = _mean_pool(outputs.last_hidden_state, encoded["attention_mask"], model_id)
        emb = _l2_normalize(emb)
        desc_embeddings[cat] = emb
    return desc_embeddings


def _resolve_device(device: str) -> tuple[str, str]:
    """Convert config-style ``"cpu"``/``"cuda"`` to torch device strings.

    Returns (model_device, tensor_device) — the model is pinned to CPU explicitly
    (the default) or to ``cuda:0``, and tensors follow the same mapping.
    """
    if device == "cpu":
        return "cpu", "cpu"
    # CUDA: model.go("cuda:0"), tensors go to "cuda:0"
    return "cuda:0", "cuda:0"


def _load_model_tokenizer(model_id: str, device: str) -> tuple[Any, Any]:
    """Load and cache (model, tokenizer) for *model_id* on *device*."""
    key = (model_id, device)
    cached = _model_cache.get(key)
    if cached is not None:
        return cached

    try:
        from transformers import AutoModel, AutoTokenizer
    except ImportError as exc:
        raise ImportError(_MISSING_DEPENDENCY_MESSAGE) from exc

    model_device, _tensor_device = _resolve_device(device)
    _model_pooling_strategies[model_id] = _read_pooling_strategy(model_id)
    _model_query_prefixes[model_id] = _read_query_prefix(model_id)
    model = AutoModel.from_pretrained(model_id).to(model_device)
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    # Tokenizer always stays on CPU; encoded tensors are moved to device
    # at inference time.

    _model_cache[key] = (model, tokenizer)
    return model, tokenizer


def _get_or_build_descs(
    model: Any,
    tokenizer: Any,
    device_str: str,
    model_id: str,
    device: str,
    categories: list[str],
) -> dict[str, Any]:
    """Return cached description embeddings or build + cache them."""
    key = (model_id, device)
    cached = _desc_cache.get(key)
    if cached is not None:
        # Check if all requested categories are already cached
        missing = [c for c in categories if c not in cached]
        if not missing:
            return cached
        # Extend the cache with any new categories
        new_embs = _build_description_embeddings(model, tokenizer, device_str, missing, model_id)
        cached.update(new_embs)
        return cached

    built = _build_description_embeddings(model, tokenizer, device_str, categories, model_id)
    _desc_cache[key] = built
    return built


def ensure_available(model_id: str, device: str) -> None:
    """Force the lazy import (and model load) now, not on the first request.

    Called at config-load time when ``classifier.mode == "local_encoder"``,
    so a missing dependency or a bad model id fails loudly at startup with
    the actionable message above, rather than as an opaque traceback on the
    first live classification.
    """
    _load_model_tokenizer(model_id, device)


def _embed_task(
    task: str,
    model: Any,
    tokenizer: Any,
    model_id: str,
    tensor_device: str,
) -> Any:
    """One pooled, L2-normalized embedding of *task* through the frozen pipeline.

    THE single embedding pipeline every consumer of this module shares —
    ``classify_zero_shot``'s nearest-centroid path, the fitted
    ``_TrainableHead`` at serving time, and the offline training script all
    call this — so a head is always scored on exactly the preprocessing it
    was trained on. Order matters and is load-bearing:

      1. ``_isolate_task_text`` strips structural agent-session noise
         (fenced code blocks, tool call/result lines, system-reminder-style
         tags) — mean/CLS pooling weights every token equally, so unwrapped
         noise outnumbers the instruction and the embedding scores the
         noise, not the intent.
      2. ``_fit_task_to_token_budget`` fits what survives to the model's
         real token window, tail-biased; the tokenizer's own silent
         truncation would otherwise keep only the shared head of a long
         agent-session prompt and discard the task-specific tail.
      3. The per-model query prefix (``_model_query_prefixes`` read from the
         snapshot, falling back to ``_BGE_QUERY_INSTRUCTION``) is prepended.
      4. One forward pass, pool per the model snapshot's pooling strategy
         (CLS for the BGE models), L2-normalize.
    """
    import torch

    isolated_task = _isolate_task_text(task)
    fitted_task = _fit_task_to_token_budget(isolated_task, tokenizer)
    prefixed_task = (
        _model_query_prefixes.get(model_id)
        or _BGE_QUERY_INSTRUCTION
    ) + fitted_task
    encoded = tokenizer(
        prefixed_task,
        padding=True,
        truncation=True,
        return_tensors="pt",
    )
    encoded = {k: v.to(tensor_device) for k, v in encoded.items()}

    with torch.no_grad():
        outputs = model(**encoded)

    task_emb = _mean_pool(outputs.last_hidden_state, encoded["attention_mask"], model_id)
    return _l2_normalize(task_emb)


# ── Trainable head: a fitted LR over the frozen embeddings ────────────
#
# Nearest-centroid IS a linear classifier whose weights are hand-written
# description strings (spec §D of plans/local-encoder-backbone-and-
# accuracy-measurement.md); the head below fits those weights instead. The
# fit happens offline (``scripts/train_encoder_head.py``) on a synthetic
# corpus generated from the committed eval task set — no raw task text is
# ever captured, so ``docs/operations.md``'s never-store rule stands. The
# artifact is a small, diffable JSON of ``coef``/``intercept`` plus per-class
# Platt (sigmoid) calibration parameters; it is committed next to the corpus
# it was trained on. At serving time it is pure json + math — no sklearn, no
# numpy — so the optional-dependency story is untouched.
#
# Activation is deliberately narrow: the head runs only when a coefficients
# file was trained for exactly THIS model id and covers every requested
# candidate category. Anything else — file absent, different backbone, a
# candidate set the fit never saw (tests use stub categories like ["a"]) —
# falls back to the zero-shot nearest-centroid path, unchanged.

_HEAD_COEFFICIENTS_PATH: Final = "evals/synthetic/encoder-head-coefficients.json"

# Resolved once per process: either a parsed head or a definitive None, so
# the coefficients file is read at most once and a missing file costs one
# stat() per process lifetime, not one per request.
_TRAINABLE_HEAD: Optional["_TrainableHead"] = None
_TRAINABLE_HEAD_RESOLVED: bool = False

# Logged once per unexpected model id, so a backbone swap that silently
# disables the fitted head is visible in the logs without spamming requests.
_head_mismatch_logged = False


class _TrainableHead:
    """Fitted multinomial logistic regression over frozen encoder embeddings.

    Trained by ``scripts/train_encoder_head.py`` on the committed synthetic
    corpus (``evals/synthetic/encoder-training.jsonl``): the same frozen
    embedding pipeline as the nearest-centroid path (see ``_embed_task``),
    then ``LogisticRegression(solver="lbfgs")`` — convex, so a seeded run is
    reproducible — with per-class Platt (sigmoid) calibration fitted by
    ``CalibratedClassifierCV(method="sigmoid")``. The serialized artifact is
    ``evals/synthetic/encoder-head-coefficients.json``:

        {
          "model_id": "...",       # the backbone the embeddings came from
          "embedding_dim": 1024,
          "classes": [...],
          "coef": [[...], ...],    # n_classes x embedding_dim
          "intercept": [...],
          "calibration": "platt",
          "calibration_a": [...],  # per-class sigmoid slope, aligned with classes
          "calibration_b": [...]   # per-class sigmoid intercept
        }

    ``predict_proba`` returns probabilities that mean P(correct) — the
    calibration is what makes the router's confidence threshold read as a
    real abstention signal rather than a monotonic score (spec item C, free
    with item D).
    """

    def __init__(self, payload: dict[str, Any], path: str) -> None:
        self.path = path
        self.model_id = str(payload["model_id"])
        self.embedding_dim = int(payload["embedding_dim"])
        self.classes: list[str] = [str(c) for c in payload["classes"]]
        self.class_set = set(self.classes)
        self.coef: list[list[float]] = payload["coef"]
        self.intercept: list[float] = payload["intercept"]
        self.calibration_a: list[float] = payload["calibration_a"]
        self.calibration_b: list[float] = payload["calibration_b"]
        # Shape-check eagerly: a truncated or hand-edited artifact must
        # degrade to the centroid fallback with one warning, never misroute.
        if len(self.coef) != len(self.classes):
            raise ValueError(
                f"coef has {len(self.coef)} rows for {len(self.classes)} classes"
            )
        for row in self.coef:
            if len(row) != self.embedding_dim:
                raise ValueError(
                    f"coef row has {len(row)} features, expected {self.embedding_dim}"
                )
        for name, values in (
            ("intercept", self.intercept),
            ("calibration_a", self.calibration_a),
            ("calibration_b", self.calibration_b),
        ):
            if len(list(payload[name])) != len(self.classes):
                raise ValueError(
                    f"{name} has {len(payload[name])} entries "
                    f"for {len(self.classes)} classes"
                )

    def covers(self, categories: list[str]) -> bool:
        """True when every requested candidate is a class the head was fit on."""
        return all(c in self.class_set for c in categories)

    def predict_proba(
        self,
        text: str,
        *,
        model_id: Optional[str] = None,
        device: str = "cpu",
    ) -> dict[str, float]:
        """Embed *text* through the frozen pipeline and return per-class
        probabilities (calibrated, summing to 1 over the trained classes).

        ``model_id`` defaults to the head's own trained backbone so
        ``head.predict_proba(task)`` works standalone; ``classify_zero_shot``
        passes the caller's explicitly (the activation guard guarantees they
        are the same).
        """
        model, tokenizer = _load_model_tokenizer(
            model_id or self.model_id, device,
        )
        _, tensor_device = _resolve_device(device or "cpu")
        emb = _embed_task(
            text, model, tokenizer, model_id or self.model_id, tensor_device,
        )
        return self.probs_from_embedding(emb)

    def probs_from_embedding(self, emb: Any) -> dict[str, float]:
        """Per-class Platt-calibrated probabilities for one pooled embedding.

        Mirrors sklearn's ``CalibratedClassifierCV(method="sigmoid",
        ensemble=False)`` multiclass predict_proba exactly: the decision
        value for class k is ``x·coef_k + intercept_k``, the calibrated
        value is ``expit(-(a_k * decision + b_k))``, and the vector is
        normalized to sum to 1 (sklearn/calibration.py,
        _SigmoidCalibration.predict). Pure Python over ``emb[0]`` — no
        numpy/sklearn at serving time.
        """
        x = emb[0].tolist()
        raw: dict[str, float] = {}
        for i, cat in enumerate(self.classes):
            row = self.coef[i]
            decision = self.intercept[i]
            for j, xv in enumerate(x):
                decision += row[j] * xv
            # expit(-(a*s + b)) = 1 / (1 + exp(a*s + b)). Guard the exponent
            # against overflow: math.exp raises past ~709, and the sigmoid
            # limit is the honest answer long before that.
            z = self.calibration_a[i] * decision + self.calibration_b[i]
            if z > 700.0:
                p = 0.0
            elif z < -700.0:
                p = 1.0
            else:
                p = 1.0 / (1.0 + math.exp(z))
            raw[cat] = p
        total = sum(raw.values())
        if total <= 0.0:
            # sklearn's uniform fallback for an all-zero calibrated row.
            n = len(self.classes)
            return {cat: 1.0 / n for cat in self.classes}
        return {cat: p / total for cat, p in raw.items()}

    @classmethod
    def load(cls, path: str) -> Optional["_TrainableHead"]:
        """Load a coefficients artifact, or None when absent/unusable.

        Fail-open by design: the fitted head is an enhancement, and any
        problem with it (not committed yet, corrupt, shape-drifted) degrades
        to the zero-shot centroid path with one warning, never a failed
        request.
        """
        coefficients = Path(path)
        if not coefficients.exists():
            return None
        try:
            payload = json.loads(coefficients.read_text())
            return cls(payload, path)
        except Exception as exc:  # noqa: BLE001 — advisory artifact, degrade
            logs.warning(
                "encoder_trainable_head_unreadable",
                path=path,
                error=str(exc),
            )
            return None


def _get_trainable_head() -> Optional["_TrainableHead"]:
    """Resolve the committed coefficients file once per process.

    Returns the head, or None when no artifact exists (the zero-shot path
    is then the only classifier). Subsequent calls return the memoized
    result without touching the filesystem.
    """
    global _TRAINABLE_HEAD, _TRAINABLE_HEAD_RESOLVED, _head_mismatch_logged
    if not _TRAINABLE_HEAD_RESOLVED:
        _TRAINABLE_HEAD_RESOLVED = True
        head = _TrainableHead.load(_HEAD_COEFFICIENTS_PATH)
        if head is not None:
            logs.info(
                "encoder_trainable_head_loaded",
                path=_HEAD_COEFFICIENTS_PATH,
                model_id=head.model_id,
                classes=len(head.classes),
            )
            _TRAINABLE_HEAD = head
    return _TRAINABLE_HEAD


def _resolve_trainable_head(
    head: Optional["_TrainableHead"], model_id: str,
) -> Optional["_TrainableHead"]:
    """Gate the head to the backbone it was trained on.

    A coefficients file fitted on one model's embedding space is meaningless
    on another's (different dimensionality, different geometry), so a
    ``classifier.encoder.model`` swap silently reverts to nearest-centroid —
    logged once so the regression is visible.
    """
    global _head_mismatch_logged
    if head is None or head.model_id == model_id:
        return head
    if not _head_mismatch_logged:
        _head_mismatch_logged = True
        logs.warning(
            "encoder_trainable_head_model_mismatch",
            trained_for=head.model_id,
            configured=model_id,
            effect="falling back to zero-shot nearest-centroid",
        )
    return None


def classify_zero_shot(
    task: str,
    categories: list[str],
    *,
    model_id: str,
    device: str,
) -> tuple[str, float]:
    """One zero-shot classification via embedding similarity.

    Returns (top_category, confidence).

    Raises ``ImportError`` with an actionable message if transformers/torch
    are not installed, and propagates any other failure (a bad model id, an
    OOM, ...) rather than swallowing it — the caller treats any exception
    here identically to a local-LLM parse failure and falls through the
    existing cascade.

    Two scoring paths share one frozen embedding (``_embed_task``), chosen
    per call:

    * Fitted head (optional): when a committed coefficients artifact was
      trained for exactly this model id and covers every requested category
      (see ``_TrainableHead``), the verdict is ``argmax(W·x + b)`` over the
      same embedding and the confidence is a Platt-calibrated probability —
      a genuine P(correct), so a below-threshold answer is a real
      abstention signal.
    * Otherwise, zero-shot nearest-centroid: the embedding is compared via
      cosine similarity against each category's precomputed description
      embedding, and similarities become a [0,1] confidence via softmax
      with a temperature of 0.10 (see the module-level docstring on
      ``_SOFTMAX_TEMPERATURE`` for the rationale).

    Categories are matched under their natural-language description
    (``_CATEGORY_DESCRIPTIONS``), never the raw config identifier — see that
    map's docstring for why.

    Long inputs are first fitted to the model's real token window (see
    ``_fit_task_to_token_budget``): the tokenizer's own silent truncation
    would otherwise keep only the shared head of a long prompt and drop
    the task-specific tail. Before that, structural agent-session noise
    (fenced code blocks, tool call/result lines, system-reminder-style
    tags) is stripped (see ``_isolate_task_text``): mean pooling weights
    every token equally, so unwrapped noise outnumbers the instruction
    and the verdict scores the noise, not the intent.
    """
    import torch

    model, tokenizer = _load_model_tokenizer(model_id, device)
    _model_device, tensor_device = _resolve_device(device)

    # Fitted head first (optional): a trained-backbone match that covers the
    # candidate set replaces the hand-written description centroids with
    # argmax(W·x + b) over the SAME frozen embedding, plus Platt-calibrated
    # probabilities. The category-description pass below is skipped entirely
    # — the head path needs only the one task forward pass. Anything that
    # doesn't match (no artifact, different backbone, unseen categories)
    # falls through to the zero-shot centroid path, unchanged.
    head = _resolve_trainable_head(_get_trainable_head(), model_id)
    if head is not None and head.covers(categories):
        task_emb = _embed_task(task, model, tokenizer, model_id, tensor_device)
        calibrated = head.probs_from_embedding(task_emb)
        candidates = {c: calibrated[c] for c in categories}
        total = sum(candidates.values())
        if total > 0.0:
            candidates = {c: p / total for c, p in candidates.items()}
        top_cat = max(categories, key=lambda c: candidates[c])
        return top_cat, candidates[top_cat]

    # Build or retrieve cached description embeddings
    descriptions = _get_or_build_descs(
        model, tokenizer, tensor_device, model_id, device, categories,
    )

    # One forward pass through the shared frozen pipeline (isolate noise,
    # fit the token window, prepend the query prefix, pool, L2-normalize —
    # the full rationale lives on _embed_task).
    task_emb = _embed_task(task, model, tokenizer, model_id, tensor_device)

    # Compute cosine similarity against each category description embedding.
    # All description embeddings are already L2-normalized, so cosine
    # similarity = dot product.
    similarities: list[tuple[str, float]] = []
    for cat in categories:
        desc_emb = descriptions[cat]
        sim = (task_emb @ desc_emb.T).item()  # scalar cosine similarity
        similarities.append((cat, sim))

    # Softmax with temperature to produce [0,1] confidence scores. We use
    # the numerically stable formulation: exp((x - x_max) / T) / sum(...).
    sim_values = torch.tensor([s for _, s in similarities])
    scaled = sim_values / _SOFTMAX_TEMPERATURE
    probs = torch.softmax(scaled, dim=0)

    # Find the top category
    top_idx = int(torch.argmax(probs))
    top_cat = similarities[top_idx][0]
    confidence = float(probs[top_idx])

    return top_cat, confidence


class _TierFeatureClassifier:
    """Sketch of a tier classifier driven by non-textual request features.

    Design intent only -- a concrete stub, NOT wired into the dispatch path
    and with no training loop. This documents the idea that task_tier can be
    predicted from request *metadata* alone, without ever reading the task
    text, which is attractive here because this router never stores raw task
    text anywhere.

    The proposed feature vector (a single row in ``X``) is built purely from
    structural request properties:

    - Prompt token count          -- length of the incoming task/prompt
    - ``tools`` array length      -- how many tools the agent session exposes
    - Fenced-block count          -- number of ```...``` code blocks (a strong
      signal of a coding turn)
    - Conversation depth          -- how many prior turns are in the session
    - Whether a diff is present   -- a binary flag for diff/review turns

    ``predict`` returns ``None`` as a stub. A real implementation would learn
    a mapping from this vector to the target tier offline (e.g. via a small
    decision tree or logistic regression fit on labeled request metadata) and
    then apply it here. numpy is only needed if/when this stub is ever
    exercised; it is imported lazily and never at module import time, matching
    the ``transformers``/``torch`` rule above.
    """

    def predict(self, X: np.ndarray) -> Optional[str]:
        """Predict a tier from a feature row.

        Stub only: returns ``None``. A real implementation would map ``X``
        to a tier label string (e.g. "tier1"/"tier2"/"tier3").
        """
        return None