from typing import List, Annotated, Union, Literal
from pydantic import BaseModel, Field
from openai.types.shared import Reasoning
import asyncio
from agents import Agent, ModelSettings, AgentOutputSchema, Runner

class FruitArgs(BaseModel):
    color: str

class FoodArgs(BaseModel):
    price: int

class BuyFruitStep(BaseModel):
    action: Literal["buy_fruit"]
    args: FruitArgs

class BuyFoodStep(BaseModel):
    action: Literal["buy_food"]
    args: FoodArgs

Step = Annotated[
    Union[BuyFruitStep, BuyFoodStep],
    Field(discriminator="action")   
]

class Actions(BaseModel):
    steps: List[Step]

Shopper = Agent(
    name="Shopper",
    instructions="You are a shopper, when the user gives you a list, convert them to a list of actions.",
    model="o3",
    tools=[],
    output_type=AgentOutputSchema(Actions),
    model_settings=ModelSettings(reasoning=Reasoning(effort="medium", summary="auto"), truncation="auto"),
)

shopper_tool = Shopper.as_tool(
    tool_name="shopper",
    tool_description="",
)

Manager = Agent(
    name="Manager",
    instructions="You are a manager with many shoppers under you. When the user asks to buy something, make a list of these and call the shopper tool to buy the items.",
    model="o3",
    tools=[shopper_tool],
    model_settings=ModelSettings(reasoning=Reasoning(effort="medium", summary="auto"), truncation="auto"),
)

async def main():
    result = await Runner.run(
        Manager,
        [{"role":"user", "content": "I want to buy a red apple."}],
    )

    print(result.final_output)

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