schemas.py

"""Demo schemas for the Structured Output Playground.

Each preset is a Pydantic model. We use Pydantic only to *generate* a clean
JSON Schema (`model_json_schema`); the actual generation is constrained by
that schema at decode time, and validation is done with `jsonschema`.
"""

from enum import Enum
from typing import List, Optional

from pydantic import BaseModel, Field


class Seniority(str, Enum):
    junior = "junior"
    mid = "mid"
    senior = "senior"
    lead = "lead"
    unknown = "unknown"


class SalaryRange(BaseModel):
    min: Optional[int] = None
    max: Optional[int] = None
    currency: Optional[str] = None


class JobPosting(BaseModel):
    title: str
    company: Optional[str] = None
    location: Optional[str] = None
    remote: bool = False
    seniority: Seniority = Seniority.unknown
    skills: List[str] = Field(default_factory=list)
    salary: Optional[SalaryRange] = None


class ContactCard(BaseModel):
    name: str
    email: Optional[str] = None
    phone: Optional[str] = None
    company: Optional[str] = None
    role: Optional[str] = None


class Product(BaseModel):
    name: str
    category: Optional[str] = None
    price: Optional[float] = None
    currency: Optional[str] = None
    in_stock: bool = True
    features: List[str] = Field(default_factory=list)


class Priority(str, Enum):
    low = "low"
    medium = "medium"
    high = "high"


class Event(BaseModel):
    # enum + int + bool on purpose: this is where free-form prompting drifts
    # (a string where an int is required, a value outside the enum…), and where
    # constrained decoding earns its keep.
    title: str
    attendees: int
    priority: Priority
    online: bool


# preset label -> Pydantic model
PRESETS = {
    "Contact card": ContactCard,
    "Product": Product,
    "Job posting": JobPosting,
    "Event": Event,
}

CUSTOM_LABEL = "Custom (edit the schema)"


def preset_schema(label: str) -> dict:
    """Return the JSON Schema dict for a preset label."""
    return PRESETS[label].model_json_schema()