Structured Outputs

# ===== The problem =====
# Models return prose by default. For tools / pipelines you want typed objects.

from pydantic import BaseModel, Field
from langchain_openai import ChatOpenAI

# ===== Define the shape =====
class Answer(BaseModel):
    answer: str = Field(..., description='The answer in one sentence.')
    confidence: float = Field(..., ge=0, le=1, description='0-1.')
    sources: list[str] = Field(default_factory=list)

# ===== with_structured_output =====
llm = ChatOpenAI(model='gpt-4o-mini', temperature=0)
structured = llm.with_structured_output(Answer)

result: Answer = structured.invoke('What is the capital of Japan? Provide confidence.')
print(result.answer, result.confidence)

# ===== JSON Schema (raw) =====
schema = {
    'name': 'order_summary',
    'description': 'Summary of an order',
    'parameters': {
        'type': 'object',
        'properties': {
            'order_id': {'type': 'string'},
            'total':    {'type': 'number'},
            'items':    {'type': 'array', 'items': {'type': 'string'}},
        },
        'required': ['order_id', 'total'],
    },
}
structured2 = llm.with_structured_output(schema)
structured2.invoke('Order O-42 for AUD 49.95, two items: A-1 and B-2.')

# ===== TypedDict =====
from typing_extensions import TypedDict
class Result(TypedDict):
    answer: str
    confidence: float

structured3 = llm.with_structured_output(Result)
out: Result = structured3.invoke('What is 1+1?')

# ===== Choosing the method =====
# Modern providers (OpenAI, Anthropic, Gemini) support structured output via:
#   - JSON mode  (response_format)
#   - Tool calling (function calling)
#   - Schema-aware decoding
# LangChain picks based on provider. You can force a specific method:
structured4 = llm.with_structured_output(Answer, method='json_schema')
structured5 = llm.with_structured_output(Answer, method='function_calling')

# ===== Validation + retries =====
# Pydantic validates the response and raises if the model misbehaves.
# Use LangChain's OutputParser + RetryWithErrorOutputParser for self-correcting loops:
from langchain.output_parsers import RetryWithErrorOutputParser, PydanticOutputParser

parser = PydanticOutputParser(pydantic_object=Answer)
retrying = RetryWithErrorOutputParser.from_llm(llm, parser)

# Pair with a prompt that includes parser.get_format_instructions().

# ===== Use it in a chain =====
from langchain_core.prompts import ChatPromptTemplate
prompt = ChatPromptTemplate.from_messages([
    ('system', 'Answer with the JSON shape only.'),
    ('human', '{question}'),
])
chain = prompt | structured
ans: Answer = chain.invoke({'question': 'What is the speed of light?'})

# ===== Why bother =====
# - Downstream code can rely on shapes; no regex on prose
# - Less prompt engineering for format; the schema is the format
# - Easier to evaluate (compare values, not strings)
# - Smaller responses (no padding prose)

# ===== Patterns to internalise =====
# - Use Pydantic models for any non-prose response
# - temperature=0 for structured tasks
# - Provide concise descriptions on each field (the model reads them)
# - Validate on receipt; reject + retry on failure

# ===== Pitfalls =====
# - Schemas with overly-permissive types ('any') -> garbage in, garbage out
# - Long descriptions inflate the prompt
# - Providers fall back to JSON mode if function calling fails -> shapes can drift
# - Forgetting timeouts -> stuck pipelines on malformed responses

from pydantic import BaseModel
class Person(BaseModel):
    name: str
    age: int
struct_llm = llm.with_structured_output(Person)
print(struct_llm.invoke('Ada Lovelace, mathematician, born 1815').age)