from agents import (
    Agent,
    function_tool,
    RunContextWrapper,
    Runner,
    ModelSettings,
)
from dataclasses import dataclass
import asyncio
from agents.extensions.models.litellm_model import LitellmModel
import os
from openai.types import Reasoning


from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()


@dataclass
class Count:
    count: int


@function_tool
def count(ctx: RunContextWrapper[Count]) -> str:
    """
    Increments the count by 1 and returns the count

    Returns:
        A string with the count
    """
    ctx.context.count += 1
    return f"Counted to {ctx.context.count}"


count_ctx = Count(count=0)

agent = Agent[Count](
    name="Counter Agent",
    instructions="Count until the number the user tells you to stop using count tool",
    tools=[count],
    model=LitellmModel(
        model="anthropic/claude-sonnet-4-20250514",
        api_key=os.getenv("ANTHROPIC_API_KEY"),
    ),
    model_settings=ModelSettings(
        reasoning=Reasoning(effort="high", summary="detailed")
    ),
)


async def main():
    results = await Runner.run(
        agent, input="Count to 10", context=count_ctx, max_turns=30
    )
    print(results)


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