from a1 import Agent, tool, LLM
from pydantic import BaseModel
# Define a tool
@tool(name="add", description="Add two numbers")
async def add(a: int, b: int) -> int:
return a + b
# Define schemas
class MathInput(BaseModel):
problem: str
class MathOutput(BaseModel):
answer: int
# Create an agent
agent = Agent(
name="math_agent",
description="Solves simple math problems",
input_schema=MathInput,
output_schema=MathOutput,
tools=[add, LLM(model="gpt-4.1")],
)
async def main():
# Compile ahead-of-time
compiled = await agent.aot(
strategy=Strategy(
generate=BaseGenerate(
llm_tool=LLM(model="gpt-4.1"),
)
)
)
result = await compiled.execute(problem="What is 2 + 2?")
print(f"AOT result: {result}")
# Or execute just-in-time
result = await agent.jit(problem="What is 5 + 3?")
print(f"JIT result: {result}")
import asyncio
asyncio.run(main())