"""
Example: Using Latch with pydantic-ai Agent

This shows how to integrate Latch as a model provider for pydantic-ai.
Latch will handle model selection and routing automatically.

Prerequisites:
1. Start the Latch API service: task dev:api
2. Start the Latch Engine service: task dev:engine
3. Set LATCH_API_KEY environment variable

Run:
    cd packages/sdk-python
    uv run python examples/pydantic_ai_example.py
"""

import asyncio
import os

from pydantic_ai import Agent

from latch.integrations.pydantic_ai import LatchModel, LatchModelSettings


async def main():
    api_key = os.environ.get("LATCH_API_KEY", "test_key")

    model = LatchModel(
        api_key=api_key,
        base_url="http://localhost:8081",
        task="qa",
        budget=0.10,
    )

    agent = Agent(
        model=model,
        system_prompt="You are a helpful assistant. Keep responses concise.",
    )

    print("Testing Latch + pydantic-ai integration...")
    print("-" * 50)

    result = await agent.run("What is the capital of France?")

    print(f"Response: {result.output}")
    print(f"Usage: {result.usage}")
    print("-" * 50)

    result2 = await agent.run(
        "Explain quantum computing in one sentence.",
        model_settings=LatchModelSettings(
            task="explanation",
            budget=0.05,
        ),
    )

    print(f"Response: {result2.output}")
    print(f"Usage: {result2.usage}")
    print("-" * 50)

    await model.close()
    print("Done!")


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