#!/usr/bin/env python3
"""02 — Structured output with Pydantic.

Ask the LLM to respond with JSON, then parse it into a typed Pydantic model.

Requirements:
    pip install koan[llm]
    export OPENAI_API_KEY=sk-...
"""

from __future__ import annotations

import asyncio

from pydantic import BaseModel, Field

from koan.llm.client import LLMClient
from koan.llm.config import LLMConfig
from koan.output.parser import parse_response
from koan.types import Message


# Define the structure you want back from the LLM
class MovieReview(BaseModel):
    title: str
    year: int
    rating: float = Field(ge=0, le=10)
    summary: str


async def main() -> None:
    async with LLMClient() as client:
        config = LLMConfig(provider="openai", model="gpt-4o-mini")

        messages = [
            Message(
                role="system",
                content="You are a movie critic. Respond ONLY with valid JSON matching "
                "this schema: {title: str, year: int, rating: float, summary: str}",
            ),
            Message(role="user", content="Review the movie 'Inception'"),
        ]

        response = await client.complete(config, messages)

        # Parse the raw LLM text into a typed object
        review = parse_response(response.text, MovieReview)

        print(f"Title: {review.title} ({review.year})")
        print(f"Rating: {review.rating}/10")
        print(f"Summary: {review.summary}")


if __name__ == "__main__":
    asyncio.run(main())
