from pydantic import BaseModel, Field
from typing import Annotated

from ag2 import Agent
from ag2.config import OpenAIConfig

class TicketTriage(BaseModel):
    category: Annotated[str, Field(description="e.g. billing, bug, account_access")]
    urgency: Annotated[str, Field(description="low, medium, or high")]
    summary_one_line: Annotated[str, Field(description="Max 120 characters", max_length=120)]

agent = Agent(
    "triage",
    prompt="You triage support messages. Be conservative with urgency.",
    config=OpenAIConfig(model="gpt-4o-mini"),
    response_schema=TicketTriage,
)

reply = await agent.ask("I was charged twice and can't export reports. Quarter close blocked.")
triage = await reply.content()      # → typed TicketTriage
print(triage.category, triage.urgency)

from ag2 import Agent, ResponseSchema

schema = ResponseSchema(int | str, name="ByteWidth", description="Number of bits in one byte.")
agent = Agent("assistant", config=config, response_schema=schema)

from ag2 import Agent, response_schema

@response_schema
def parse_rating(content: str) -> int:
    """Parse a rating and clamp to 1–5."""
    return max(1, min(5, int(content)))

agent = Agent("assistant", config=config, response_schema=parse_rating)

from typing import Annotated
from pydantic import Field
from ag2 import response_schema

@response_schema
def extract_listing(
    title: Annotated[str, Field(description="Product name")],
    price_usd: Annotated[float, Field(description="Price in USD", ge=0)],
    in_stock: Annotated[bool, Field(description="True if it ships now")],
) -> dict:
    return {"title": title, "price_usd": price_usd, "in_stock": in_stock}

import json

@response_schema
async def fetch_and_validate(content: str) -> dict:
    data = json.loads(content)
    data["validated"] = True
    return data

from ag2 import Agent, PromptedSchema

agent = Agent("assistant", config=config, response_schema=PromptedSchema(int))

PromptedSchema(int, prompt_template="Reply with JSON matching this schema:\n{schema}")

agent = Agent("assistant", config=config)

turn = await agent.ask("How many seconds in a minute?", response_schema=int)
print(await turn.content())   # 60

turn2 = await turn.ask("Say hello.")     # back to default (no schema)

result = await reply.content(retries=3)   # initial + up to 3 re-asks
result = await reply.content(retries=math.inf)   # interactive only — could loop forever

ResponseSchema(int, name="RawInt", embed=False)              # model must produce a bare 42
@response_schema(embed=False)
def parse_rating(value: int) -> int: ...