"""Typed decision schemas. A schema declares, in advance, every field a query returns and every value each field may take. Type safety is therefore constructive: the model normalises only over declared options, so a returned struct cannot name an option that does not exist and cannot omit a declared field. Nothing is parsed or validated at inference time. Four field kinds cover the shapes decisions actually take: - ``choice`` exactly one of N unordered options (softmax over N) - ``bool`` a two-option choice, kept separate so it calibrates on its own - ``multi`` any subset of N options (independent sigmoid per option) - ``bucket`` one of N *ordered* options, scored with cumulative logits so that the order is part of the model rather than an accident of labels """ from __future__ import annotations from dataclasses import dataclass from typing import Literal Kind = Literal["choice", "bool", "multi", "bucket"] KINDS: tuple[Kind, ...] = ("choice", "bool", "multi", "bucket") BOOL_OPTIONS = ("false", "true") @dataclass(frozen=True) class Field: """One decision within a schema.""" name: str kind: Kind = "choice" options: tuple[str, ...] = BOOL_OPTIONS description: str = "" def __post_init__(self) -> None: if not self.name: raise ValueError("field name must not be empty") if self.kind not in KINDS: raise ValueError(f"{self.name}: unknown kind {self.kind!r}, expected one of {KINDS}") if self.kind == "bool": if self.options != BOOL_OPTIONS: raise ValueError(f"{self.name}: bool options are fixed to {BOOL_OPTIONS}") elif len(self.options) < 2: raise ValueError(f"{self.name}: {self.kind} needs at least two options") if len(set(self.options)) != len(self.options): raise ValueError(f"{self.name}: options must be unique") if any(not option for option in self.options): raise ValueError(f"{self.name}: options must not be empty strings") @property def cardinality(self) -> int: return len(self.options) @property def prompt(self) -> str: """The text the model reads to know what this field asks. This is what makes a schema the model never trained on usable: the field is identified by its words, not by a row index into a learned table. An earlier version stored one trainable vector per field, which meant adding a field to a schema raised a shape error and renaming one made the field unrecognisable. """ parts = [self.name.replace("_", " ")] if self.description: parts.append(self.description) parts.append(f"kind: {self.kind}") if self.kind != "choice" or len(self.options) <= 8: parts.append("options: " + ", ".join(self.options[:8])) return " | ".join(parts) @property def single(self) -> bool: """True when exactly one option is correct, so the field owns a softmax.""" return self.kind in ("choice", "bool", "bucket") def index(self, value: str) -> int: try: return self.options.index(value) except ValueError: raise ValueError( f"{self.name}: {value!r} is not a declared option; expected one of {self.options}" ) from None def encode(self, value) -> list[int] | int: """Label for one row: an option index, or a 0/1 vector for ``multi``.""" if self.kind == "bool": if not isinstance(value, bool): raise ValueError(f"{self.name}: expected a bool, got {value!r}") return int(value) if self.kind == "multi": if isinstance(value, str) or not isinstance(value, (list, tuple, set)): raise ValueError(f"{self.name}: multi expects a list of options, got {value!r}") chosen = {self.index(item) for item in value} return [int(index in chosen) for index in range(self.cardinality)] if not isinstance(value, str): raise ValueError(f"{self.name}: expected an option string, got {value!r}") return self.index(value) def decode(self, probabilities, threshold: float = 0.5): """Probabilities for this field -> the typed value plus its confidence.""" if len(probabilities) != self.cardinality: raise ValueError( f"{self.name}: expected {self.cardinality} probabilities, got {len(probabilities)}" ) if self.kind == "multi": selected = [option for option, p in zip(self.options, probabilities) if p >= threshold] confidence = min((max(p, 1.0 - p) for p in probabilities), default=1.0) return selected, float(confidence) best = max(range(self.cardinality), key=probabilities.__getitem__) value = bool(best) if self.kind == "bool" else self.options[best] return value, float(probabilities[best]) @dataclass(frozen=True) class Schema: """The full set of fields one query returns.""" fields: tuple[Field, ...] def __post_init__(self) -> None: if not self.fields: raise ValueError("a schema needs at least one field") names = [field.name for field in self.fields] if len(set(names)) != len(names): raise ValueError("field names must be unique") def __len__(self) -> int: return len(self.fields) def __iter__(self): return iter(self.fields) def __getitem__(self, name: str) -> Field: for field in self.fields: if field.name == name: return field raise KeyError(name) @property def max_cardinality(self) -> int: return max(field.cardinality for field in self.fields) def encode(self, values: dict) -> list: """One row of labels, in field order. Every declared field must be present.""" missing = [field.name for field in self.fields if field.name not in values] if missing: raise ValueError(f"missing labels for fields: {missing}") extra = set(values) - {field.name for field in self.fields} if extra: raise ValueError(f"labels for undeclared fields: {sorted(extra)}") return [field.encode(values[field.name]) for field in self.fields] @property def prompts(self) -> tuple[str, ...]: return tuple(field.prompt for field in self.fields) def to_json(self) -> dict: return { "fields": [ {"name": f.name, "kind": f.kind, "options": list(f.options), "description": f.description} for f in self.fields ] } @classmethod def from_json(cls, payload: dict) -> "Schema": return cls( tuple( Field(item["name"], item.get("kind", "choice"), tuple(item["options"]), item.get("description", "")) for item in payload["fields"] ) )