import cascadeflow

cascadeflow.init(mode="enforce")

@cascadeflow.agent(budget=0.20)
async def my_agent(query: str):
    """This agent cannot spend more than $0.20."""
    return await llm.complete(query)

@cascadeflow.agent(budget=0.50, compliance="gdpr")
async def eu_agent(query: str):
    """Process EU data — only GDPR-approved models, $0.50 max."""
    return await llm.complete(query)

@cascadeflow.agent(
    budget=1.00,
    kpi_weights={"quality": 0.8, "cost": 0.2},
    kpi_targets={"quality": 0.9},
)
async def premium_agent(query: str):
    """High-quality responses — prioritize quality over cost."""
    return await llm.complete(query)

@cascadeflow.agent(
    budget=0.10,
    kpi_weights={"cost": 0.9, "quality": 0.1},
)
async def triage_agent(query: str):
    """Quick classification — optimize for cost."""
    return await llm.complete(query)

@cascadeflow.agent(
    budget=2.00,
    compliance="hipaa",
    kpi_weights={"quality": 0.9, "cost": 0.1},
    kpi_targets={"quality": 0.95},
)
async def medical_agent(query: str):
    """Patient data — strict compliance, high quality, higher budget."""
    return await llm.complete(query)

@cascadeflow.agent(
    budget=0.50,
    max_tool_calls=5,
)
async def research_agent(query: str):
    """Research with tools — capped at 5 tool calls and $0.50."""
    return await llm.complete(query)

@cascadeflow.agent(budget=0.50, compliance="gdpr")
async def my_agent(query: str):
    return await llm.complete(query)

# The run adds an outer budget cap on top of the agent's own cap
with cascadeflow.run(budget=2.00) as session:
    await my_agent("First query")   # Agent cap: $0.50, Run cap: $2.00
    await my_agent("Second query")  # Agent cap: $0.50, Run cap: $2.00 minus first query
    print(session.summary())