qwen_test_runner/registry.py

"""
registry.py — The slot registry.

This is the source of truth for the caption schema. Every slot the system
knows about lives here as a SlotSpec entry. The Pydantic Caption model, the
JSON Schema export, the GBNF grammar, and the evaluator's grounding rules
are all derived from this registry at import time.

Adding a slot is one dict entry. Adding a category is one Literal expansion.
No code outside this file should hardcode slot names or category logic.

Slot taxonomy (the three categories that came out of the baseline analysis):
  - descriptive : grounded in the input caption. Hallucination forbidden.
                  Examples: subjects, actions, setting.
  - aesthetic   : how the scene should look. Often empty in input;
                  legitimate inference (or null) in enhancement mode.
                  Examples: style, lighting, palette.
  - semantic    : interpretive meaning. Inferential by definition.
                  Examples: mood, implication, narrative_function.

Groundedness rules (drive the evaluator):
  - must_ground   : every leaf MUST trace to the input caption.
  - may_infer     : leaf may be grounded OR inferred; both are acceptable.
  - derived_only  : leaf is expected to be inferred. Grounding check skipped.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Literal, Optional, Type

from pydantic import BaseModel, Field


# ──────────────────────────────────────────────────────────────────────────────
# Slot-level enums. Adding a value here is a registry-only change; no
# code outside this file matches on these strings directly (the helpers below
# encapsulate all behavior).
# ──────────────────────────────────────────────────────────────────────────────

Category = Literal["descriptive", "aesthetic", "semantic"]
Cardinality = Literal["single", "list"]
Vocabulary = Literal["closed", "open"]
Groundedness = Literal["must_ground", "may_infer", "derived_only"]

# value_kind selects the leaf's primitive type for code generation. "string" is
# the caption default (every existing slot). The numeric kinds exist so the SAME
# registry→Pydantic→JSON-Schema→GBNF machinery can describe vision outputs
# (bounding boxes, confidences, depths) without a second codegen path.
#   string  → str
#   number  → float   (optionally bounded via number_range)
#   integer → int
#   bbox    → list[float] of length 4   (x1,y1,x2,y2 or x,y,w,h — see coords.py)
#   point   → list[float] of length 2
ValueKind = Literal["string", "number", "integer", "bbox", "point"]


# ──────────────────────────────────────────────────────────────────────────────
# Nested value models. Used by slots whose value is structured (e.g. subjects
# have a name and a list of attributes). New nested types go here and are
# referenced from the SlotSpec via `nested_model=`.
# ──────────────────────────────────────────────────────────────────────────────

class SubjectValue(BaseModel):
    """A single entity in the caption."""
    name: str = Field(..., min_length=1, max_length=64)
    # No max_length on attributes: rich captions (JoyCaption prose, booru tag
    # strings) legitimately carry >8 per subject, and the cap was rejecting 44%
    # of otherwise-valid structs in the 100-row bench (2026-07).
    attributes: list[str] = Field(default_factory=list)


# ──────────────────────────────────────────────────────────────────────────────
# SlotSpec — the unit of the registry.
# ──────────────────────────────────────────────────────────────────────────────

@dataclass(frozen=True)
class SlotSpec:
    """Declarative description of one schema slot.

    The structural axes (cardinality, vocabulary, value_kind, nested_fields,
    nested_model) are what drive code generation in schema.py — these apply
    equally to caption slots and to vision-task fields. The caption-only axes
    (category, groundedness) drive the text evaluator and default to neutral
    values so vision per-category registries can omit them.

        cardinality    — single value vs list
        vocabulary     — open (any string) vs closed (one of `closed_values`)
        value_kind     — leaf primitive type (string / number / integer / bbox / point)
        nested_model   — BaseModel subclass for a structured value (caption: SubjectValue)
        nested_fields  — declarative nested-object fields; the generalized form of
                         nested_model — schema.py builds both the Pydantic model and
                         the GBNF object rule recursively from these
        category       — taxonomy bucket (caption prompts); ignored by vision
        groundedness   — strict / soft / never (text evaluator); ignored by vision
        optional       — may the model emit null/[] when empty
        number_range   — (min, max) bound for numeric value_kinds (Pydantic ge/le)
    """
    name: str
    cardinality: Cardinality
    vocabulary: Vocabulary
    category: Category = "descriptive"
    groundedness: Groundedness = "may_infer"
    value_kind: ValueKind = "string"
    closed_values: tuple[str, ...] = ()
    nested_model: Optional[Type[BaseModel]] = None
    nested_fields: tuple["SlotSpec", ...] = ()
    optional: bool = True
    max_items: int = 8           # only for cardinality == "list"
    max_str_length: int = 64     # for open-vocab strings
    number_range: Optional[tuple[float, float]] = None

    def __post_init__(self):
        # Lightweight validation — catch registry mistakes at import time
        if self.vocabulary == "closed" and not self.closed_values:
            raise ValueError(f"slot {self.name!r}: closed vocab requires closed_values")
        if self.vocabulary == "open" and self.closed_values:
            raise ValueError(f"slot {self.name!r}: open vocab cannot have closed_values")
        if self.nested_model is not None and self.vocabulary == "closed":
            raise ValueError(f"slot {self.name!r}: nested_model is incompatible with closed vocab")
        if self.nested_fields and self.nested_model is not None:
            raise ValueError(f"slot {self.name!r}: nested_fields and nested_model are mutually exclusive")
        if self.nested_fields and self.vocabulary == "closed":
            raise ValueError(f"slot {self.name!r}: nested_fields is incompatible with closed vocab")


# ──────────────────────────────────────────────────────────────────────────────
# THE REGISTRY.
#
# Starter set: 5 slots that exercise all three categories and both
# groundedness extremes. Adding a slot is a single entry below.
# ──────────────────────────────────────────────────────────────────────────────

SLOT_REGISTRY: dict[str, SlotSpec] = {
    "subjects": SlotSpec(
        name="subjects",
        category="descriptive",
        cardinality="list",
        vocabulary="open",
        groundedness="must_ground",
        nested_model=SubjectValue,
        max_items=8,
    ),
    "actions": SlotSpec(
        name="actions",
        category="descriptive",
        cardinality="list",
        vocabulary="open",
        groundedness="must_ground",
        max_items=8,
    ),
    "setting": SlotSpec(
        name="setting",
        category="descriptive",
        cardinality="single",
        vocabulary="closed",
        # `may_infer` because Qwen reliably guesses indoor/outdoor from cues
        # even when the caption doesn't say. The grammar pins the value to
        # the enum anyway.
        groundedness="may_infer",
        closed_values=("indoor", "outdoor", "unknown"),
        optional=False,   # always required; the enum includes "unknown" as escape
    ),
    "style": SlotSpec(
        name="style",
        category="aesthetic",
        cardinality="single",
        vocabulary="open",
        groundedness="may_infer",
    ),
    "mood": SlotSpec(
        name="mood",
        category="semantic",
        cardinality="single",
        vocabulary="open",
        # Baseline finding: mood is 73% of all hallucinations under the old
        # rule. Reclassifying it as derived_only stops penalizing the model
        # for inferring; it's correct behavior now, not error.
        groundedness="derived_only",
    ),
}


# ──────────────────────────────────────────────────────────────────────────────
# Query helpers. Use these instead of poking SLOT_REGISTRY directly so behavior
# stays centralized.
# ──────────────────────────────────────────────────────────────────────────────

def slots_by_category(category: Category) -> list[SlotSpec]:
    return [s for s in SLOT_REGISTRY.values() if s.category == category]


def slot_names() -> list[str]:
    """Slot names in registry-declaration order. JSON output uses this order."""
    return list(SLOT_REGISTRY.keys())


def get_slot(name: str) -> SlotSpec:
    if name not in SLOT_REGISTRY:
        raise KeyError(f"unknown slot: {name!r}")
    return SLOT_REGISTRY[name]


# Set of closed-vocab values across all slots — used by the evaluator as the
# "always grounded" allowlist for the `may_infer` closed-vocab case.
def all_closed_vocab() -> set[str]:
    out: set[str] = set()
    for s in SLOT_REGISTRY.values():
        out.update(s.closed_values)
    return out