qwen_test_runner/schema.py

"""
schema.py — Registry-parametric schema code generation.

Generates three representations of a *registry* (a `dict[str, SlotSpec]`):

    build_model_from_registry(name, registry)  → Pydantic model class
    build_json_schema(model)                    → JSON Schema dict
    build_gbnf_from_registry(registry)          → GBNF grammar string

The caption schema is the canonical instance, exposed under stable names:

    Caption                 — Pydantic model (validation / parsing)
    CAPTION_JSON_SCHEMA     — JSON Schema dict (Anthropic API, outlines, etc.)
    CAPTION_GRAMMAR_GBNF    — GBNF grammar string (xgrammar)

The vision subpackage reuses the SAME generators per task category (each
category owns a small `dict[str, SlotSpec]`), so numbers, bounding boxes, and
nested objects are described with the same machinery — see
`qwen_test_runner/vision/`.

All caption artifacts are generated at import time from `registry.SLOT_REGISTRY`.
To add or modify a caption slot, edit `registry.py` only — this file stays
untouched.
"""

from __future__ import annotations

from typing import Any, Literal, Mapping, Optional

from pydantic import BaseModel, Field, create_model

from .registry import SLOT_REGISTRY, SlotSpec, SubjectValue


# ──────────────────────────────────────────────────────────────────────────────
# Pydantic model — built dynamically from any registry.
# ──────────────────────────────────────────────────────────────────────────────

_OBJECT_MODEL_CACHE: dict[SlotSpec, type[BaseModel]] = {}


def _object_model_name(spec: SlotSpec) -> str:
    return "".join(part.capitalize() for part in spec.name.split("_")) + "Obj"


def _build_object_model(spec: SlotSpec) -> type[BaseModel]:
    """Build (and cache) a Pydantic model from a spec's `nested_fields`."""
    cached = _OBJECT_MODEL_CACHE.get(spec)
    if cached is not None:
        return cached
    fields: dict[str, Any] = {}
    for f in spec.nested_fields:
        fields[f.name] = (_python_type_for_slot(f), _field_for_slot(f))
    model = create_model(_object_model_name(spec), **fields)
    _OBJECT_MODEL_CACHE[spec] = model
    return model


def _item_type_for_slot(spec: SlotSpec) -> Any:
    """Python type of a single value of this slot (before list/Optional wrapping)."""
    if spec.nested_fields:
        return _build_object_model(spec)
    if spec.nested_model is not None:
        return spec.nested_model
    if spec.vocabulary == "closed":
        # Literal[("a", "b", "c")] parses identically to Literal["a", "b", "c"].
        return Literal[spec.closed_values]
    vk = spec.value_kind
    if vk == "number":
        return float
    if vk == "integer":
        return int
    if vk in ("bbox", "point"):
        return list[float]
    return str


def _python_type_for_slot(spec: SlotSpec) -> Any:
    """Compute the Python type annotation for a slot's value.

    List cardinality wraps the item type in list[...].
    Single + optional wraps in Optional[...].
    """
    item_type = _item_type_for_slot(spec)
    if spec.cardinality == "list":
        return list[item_type]
    if spec.optional:
        return Optional[item_type]
    return item_type


def _default_for_slot(spec: SlotSpec) -> Any:
    if spec.cardinality == "list":
        return []  # default_factory handled by Field below
    if spec.optional:
        return None
    # Required single value with no default. For closed vocab, default to
    # the last value (usually "unknown") so partial outputs don't blow up.
    if spec.vocabulary == "closed":
        return spec.closed_values[-1]
    return ...  # required, no default


def _field_for_slot(spec: SlotSpec):
    """Construct a Pydantic Field with the right constraints for this slot."""
    kwargs: dict[str, Any] = {}

    if spec.cardinality == "list":
        kwargs["default_factory"] = list
        kwargs["max_length"] = spec.max_items
        return Field(**kwargs)

    default = _default_for_slot(spec)

    # Fixed-length numeric arrays (bbox/point): exactly 4 / 2 elements.
    if spec.value_kind in ("bbox", "point"):
        n = 4 if spec.value_kind == "bbox" else 2
        if default is ...:
            return Field(..., min_length=n, max_length=n)
        return Field(default=default, min_length=n, max_length=n)

    # Scalar numerics, optionally bounded.
    if spec.value_kind in ("number", "integer"):
        rng: dict[str, Any] = {}
        if spec.number_range is not None:
            rng["ge"] = spec.number_range[0]
            rng["le"] = spec.number_range[1]
        if default is ...:
            return Field(..., **rng)
        return Field(default=default, **rng)

    # Strings / enums / nested objects.
    if default is ...:
        # Only plain open strings get a length cap; nested models / enums don't.
        if spec.vocabulary == "open" and spec.nested_model is None and not spec.nested_fields:
            return Field(..., max_length=spec.max_str_length)
        return Field(...)
    kwargs["default"] = default
    if (
        spec.vocabulary == "open"
        and spec.nested_model is None
        and not spec.nested_fields
        and spec.value_kind == "string"
    ):
        kwargs["max_length"] = spec.max_str_length
    return Field(**kwargs)


def build_model_from_registry(model_name: str, registry: Mapping[str, SlotSpec]) -> type[BaseModel]:
    """Build a Pydantic model with one field per registry entry."""
    fields: dict[str, Any] = {}
    for name, spec in registry.items():
        fields[name] = (_python_type_for_slot(spec), _field_for_slot(spec))
    return create_model(model_name, **fields)


def build_json_schema(model: type[BaseModel]) -> dict:
    """JSON Schema for a generated model (thin wrapper for symmetry)."""
    return model.model_json_schema()


Caption = build_model_from_registry("Caption", SLOT_REGISTRY)

# Re-export SubjectValue under the old name "Subject" for callers that
# imported it from schema previously.
Subject = SubjectValue


# ──────────────────────────────────────────────────────────────────────────────
# JSON Schema — derived from the Pydantic model.
# ──────────────────────────────────────────────────────────────────────────────

CAPTION_JSON_SCHEMA: dict = build_json_schema(Caption)


# ──────────────────────────────────────────────────────────────────────────────
# GBNF grammar — built from the registry. Independent of pydantic.
#
# xgrammar's auto-converter from JSON schema sometimes adds unwanted slack
# (e.g. permissive whitespace patterns that hurt parse rates). Generating GBNF
# by hand from the registry gives tighter control and stays consistent with
# the Pydantic model.
#
# The four base primitives (str_array, string, char, ws) are always emitted, as
# in v0.2. Numeric primitives (number, bbox4, …) are emitted ONLY when a slot
# references them, so the caption grammar is byte-for-byte the v0.2 grammar.
# ──────────────────────────────────────────────────────────────────────────────

# Primitive rule definitions, emitted on demand.
_PRIMITIVE_DEFS: dict[str, str] = {
    "uint": 'uint ::= "0" | [1-9] [0-9]*',
    "frac": 'frac ::= "." [0-9]+',
    "exp": 'exp ::= ("e" | "E") ("+" | "-")? [0-9]+',
    "integer": 'integer ::= "-"? uint',
    "number": 'number ::= "-"? uint frac? exp?',
    "num_array": 'num_array ::= "[" ws "]" | "[" ws number (ws "," ws number)* ws "]"',
    "bbox4": 'bbox4 ::= "[" ws number ws "," ws number ws "," ws number ws "," ws number ws "]"',
    "point2": 'point2 ::= "[" ws number ws "," ws number ws "]"',
}

# Transitive dependencies between numeric primitives (the four base primitives
# are always present, so they are never listed here).
_PRIMITIVE_DEPS: dict[str, set[str]] = {
    "uint": set(),
    "frac": set(),
    "exp": set(),
    "integer": {"uint"},
    "number": {"uint", "frac", "exp"},
    "num_array": {"number"},
    "bbox4": {"number"},
    "point2": {"number"},
}

# Stable emission order so the grammar regenerates deterministically.
_PRIMITIVE_ORDER = ["integer", "number", "num_array", "bbox4", "point2", "uint", "frac", "exp"]


def _gbnf_string_alternation(values: tuple[str, ...]) -> str:
    """Emit `"\"a\"" | "\"b\"" | ...` for a closed enum."""
    return " | ".join(f'"\\"{v}\\""' for v in values)


def _resolve_primitive_deps(deps: set[str]) -> set[str]:
    """Expand a set of primitive names with all transitive dependencies."""
    out: set[str] = set()
    stack = list(deps)
    while stack:
        d = stack.pop()
        if d in out:
            continue
        out.add(d)
        stack.extend(_PRIMITIVE_DEPS.get(d, set()))
    return out


def _gbnf_object_rule(spec: SlotSpec) -> tuple[str, list[str], set[str]]:
    """Build the GBNF object rule for a spec's nested_fields. Returns
    (object_rule_name, extra_rules, primitive_deps)."""
    rule_name = f"obj_{spec.name}"
    parts: list[str] = ['"{"', "ws"]
    extras: list[str] = []
    deps: set[str] = set()
    for i, f in enumerate(spec.nested_fields):
        if i > 0:
            parts += ['","', "ws"]
        frhs, fextras, fdeps = _gbnf_slot_value_rule(f)
        extras += fextras
        deps |= fdeps
        # Wrap the field value in parens so alternations (e.g. "null" | bbox4)
        # compose correctly inside the object.
        parts += [f'"\\"{f.name}\\":"', "ws", f"( {frhs} )", "ws"]
    parts.append('"}"')
    extras.append(f"{rule_name} ::= " + " ".join(parts))
    return rule_name, extras, deps


def _gbnf_slot_value_rule(spec: SlotSpec) -> tuple[str, list[str], set[str]]:
    """Return (right-hand-side, extra_rules, primitive_deps) for this slot's value.

    The RHS is what appears after `slot_<name> ::=`. It is either a rule name or
    a small alternation (e.g. `"null" | number`). Extra rules are helper rules
    this slot needs; primitive_deps names numeric primitives to emit globally.
    """
    extras: list[str] = []
    deps: set[str] = set()

    if spec.cardinality == "list":
        if spec.nested_model is SubjectValue:
            # SubjectValue is the caption's one hand-written nested type; keep the
            # exact v2 rules so the caption grammar is unchanged.
            extras.append(
                'subject ::= "{" ws "\\"name\\":" ws string ws "," ws '
                '"\\"attributes\\":" ws str_array ws "}"'
            )
            extras.append(
                'subject_list ::= "[" ws "]" | '
                '"[" ws subject (ws "," ws subject)* ws "]"'
            )
            return "subject_list", extras, deps
        if spec.nested_fields:
            obj_name, obj_extras, obj_deps = _gbnf_object_rule(spec)
            extras += obj_extras
            deps |= obj_deps
            list_name = f"{spec.name}_list"
            extras.append(
                f'{list_name} ::= "[" ws "]" | '
                f'"[" ws {obj_name} (ws "," ws {obj_name})* ws "]"'
            )
            return list_name, extras, deps
        if spec.value_kind in ("number", "integer"):
            deps.add("num_array")
            return "num_array", extras, deps
        # Primitive open-vocab list — array of strings
        return "str_array", extras, deps

    # Single value
    if spec.nested_fields:
        # A single nested object (e.g. subject_fixation.primary_subject). Without
        # this, the grammar would fall through to the string rule and force the
        # object to serialize as a string — breaking constrained decoding.
        obj_name, obj_extras, obj_deps = _gbnf_object_rule(spec)
        extras += obj_extras
        deps |= obj_deps
        if spec.optional:
            return f'"null" | {obj_name}', extras, deps
        return obj_name, extras, deps

    if spec.vocabulary == "closed":
        alts = _gbnf_string_alternation(spec.closed_values)
        rule_name = f"closed_{spec.name}"
        extras.append(f"{rule_name} ::= {alts}")
        return rule_name, extras, deps

    if spec.value_kind == "bbox":
        deps.add("bbox4")
        base = "bbox4"
    elif spec.value_kind == "point":
        deps.add("point2")
        base = "point2"
    elif spec.value_kind == "number":
        deps.add("number")
        base = "number"
    elif spec.value_kind == "integer":
        deps.add("integer")
        base = "integer"
    else:
        base = "string"

    # Optional single → allow null literal.
    if spec.optional:
        return f'"null" | {base}', extras, deps
    return base, extras, deps


def build_gbnf_from_registry(registry: Mapping[str, SlotSpec], root_name: str = "root") -> str:
    """Generate a GBNF grammar that produces JSON conforming to `registry`."""
    slot_rules: list[str] = []
    helper_rules: list[str] = []
    helper_seen: set[str] = set()
    deps: set[str] = set()

    for name, spec in registry.items():
        rhs, extras, sdeps = _gbnf_slot_value_rule(spec)
        slot_rules.append(f"slot_{name} ::= {rhs}")
        deps |= sdeps
        for r in extras:
            head = r.split("::=", 1)[0].strip()
            if head not in helper_seen:
                helper_rules.append(r)
                helper_seen.add(head)

    # Root rule: opening brace, slot1, comma, slot2, ..., closing brace.
    parts: list[str] = ['"{"', "ws"]
    for i, name in enumerate(registry.keys()):
        if i > 0:
            parts += ['","', "ws"]
        parts += [f'"\\"{name}\\":"', "ws", f"slot_{name}", "ws"]
    parts.append('"}"')
    root_rule = f"{root_name} ::= " + " ".join(parts)

    # Base primitives — always present (str_array is needed by open-vocab lists
    # and SubjectValue.attributes).
    common = [
        'str_array ::= "[" ws "]" | "[" ws string (ws "," ws string)* ws "]"',
        'string ::= "\\"" char* "\\""',
        'char ::= [^"\\\\] | "\\\\" ["\\\\/bfnrt]',
        'ws ::= [ \\t\\n]*',
    ]

    # Numeric primitives — only those actually referenced (keeps caption grammar
    # identical to v2 and vision grammars minimal).
    resolved = _resolve_primitive_deps(deps)
    numeric = [_PRIMITIVE_DEFS[k] for k in _PRIMITIVE_ORDER if k in resolved]

    return "\n".join([root_rule] + slot_rules + helper_rules + common + numeric)


def build_gbnf_grammar() -> str:
    """Generate the caption GBNF grammar (back-compat wrapper)."""
    return build_gbnf_from_registry(SLOT_REGISTRY)


CAPTION_GRAMMAR_GBNF: str = build_gbnf_grammar()


# ──────────────────────────────────────────────────────────────────────────────
# Smoke test — `python -m qwen_test_runner.schema` validates the three reps.
# ──────────────────────────────────────────────────────────────────────────────

def _smoke_test() -> None:
    example = Caption(
        subjects=[Subject(name="dog", attributes=["golden"])],
        actions=["catching"],
        setting="outdoor",
        style="photorealistic",
        mood="energetic",
    )

    as_dict = example.model_dump()
    rebuilt = Caption.model_validate(as_dict)
    assert rebuilt == example, "pydantic round-trip failed"

    as_json = example.model_dump_json()
    reparsed = Caption.model_validate_json(as_json)
    assert reparsed == example, "JSON round-trip failed"

    schema = CAPTION_JSON_SCHEMA
    assert "properties" in schema
    assert set(schema["properties"].keys()) == set(SLOT_REGISTRY.keys())

    g = CAPTION_GRAMMAR_GBNF
    for slot in SLOT_REGISTRY:
        assert f'\\"{slot}\\"' in g, f"GBNF missing slot {slot}"

    print("schema.py smoke test: OK")
    print(f"  slots: {list(SLOT_REGISTRY.keys())}")
    print(f"  example JSON length: {len(as_json)}")
    print(f"  JSON Schema fields: {list(schema['properties'].keys())}")
    print(f"  GBNF length: {len(g)} chars")


if __name__ == "__main__":
    _smoke_test()