06_structured_output.py

"""06: Structured output validated against a Pydantic schema."""

import asyncio
import json

from pydantic import BaseModel

from claude_agent_sdk import ClaudeAgentOptions, ResultMessage, query
from demo_util import PROJECT


class Finding(BaseModel):
    file: str
    line: int
    severity: str  # "low" | "medium" | "high"
    description: str


class SecurityReport(BaseModel):
    findings: list[Finding]
    total_count: int
    overall_risk: str


async def main():
    try:
        async for message in query(
            prompt="Audit every file under src/ for security problems and report them.",
            options=ClaudeAgentOptions(
                cwd=PROJECT,
                allowed_tools=["Read", "Glob", "Grep"],
                max_turns=8,
                output_format={"type": "json_schema", "schema": SecurityReport.model_json_schema()},
            ),
        ):
            if isinstance(message, ResultMessage):
                if message.subtype == "success" and message.structured_output:
                    report = SecurityReport.model_validate(message.structured_output)
                    print("Raw structured_output:")
                    print(json.dumps(message.structured_output, indent=2))
                    print(f"\nTyped access: {report.total_count} findings, risk={report.overall_risk}")
                    for f in report.findings:
                        print(f"  {f.file}:{f.line} [{f.severity}] {f.description}")
                else:
                    print("No structured output:", message.subtype, message.errors)
    except Exception as error:
        print(f"Session ended with an error: {error}")


asyncio.run(main())