Tools

Tools enable AI models to call functions in your code. ai-query provides a type-safe way to define tools using Python decorators and type hints.

Defining Tools

Use the @tool decorator to create a tool from any function:

from ai_query import tool, Field

@tool(description="Get the current weather for a location")
async def get_weather(
    location: str = Field(description="The city name"),
    unit: str = Field(description="Temperature unit", default="celsius")
) -> str:
    # Your implementation here
    return f"Weather in {location}: 72°F"

The `@tool` Decorator

The decorator accepts this optional parameter:

Parameter	Type	Description
`description`	`str`	What the tool does (shown to the AI)

Parameter Definitions with `Field`

Use Field to describe each parameter:

from ai_query import Field

@tool(description="Search a database")
async def search(
    query: str = Field(description="The search query"),
    limit: int = Field(description="Max results to return", default=10),
    include_archived: bool = Field(description="Include archived items", default=False)
) -> str:
    ...

Field Argument	Type	Description
`description`	`str`	Parameter description for the AI
`default`	`Any`	Default value (makes parameter optional)
`enum`	`list[Any]`	Restrict the parameter to a fixed set of values
`min_value`	`float`	Minimum numeric value
`max_value`	`float`	Maximum numeric value

Supported Parameter Types

ai-query supports these Python types:

str - String values
int - Integer values
float - Floating point numbers
bool - Boolean values
list[T] - Arrays with inferred item types
dict[str, T] - Objects with typed values
tuple[T1, T2] and tuple[T, ...] - Fixed or variable-length arrays
Optional[T], Union[...], and Literal[...] - Optional, union, and enum-like values
TypedDict and dataclasses - Nested object schemas

If a type annotation is not recognized, inference falls back to string. For full control, create a tool with an explicit parameters schema.

from typing import NotRequired, TypedDict

class SearchFilters(TypedDict):
    tags: list[str]
    include_archived: NotRequired[bool]

@tool(description="Search documents")
async def search_documents(
    query: str = Field(description="Search query"),
    filters: SearchFilters = Field(description="Structured filters")
) -> str:
    ...

Using Tools with generate_text

Pass tools as a dictionary:

from ai_query import generate_text
from ai_query.providers import google, tool, Field

@tool(description="Add two numbers")
async def add(
    a: int = Field(description="First number"),
    b: int = Field(description="Second number")
) -> str:
    return str(a + b)

@tool(description="Multiply two numbers")
async def multiply(
    a: int = Field(description="First number"),
    b: int = Field(description="Second number")
) -> str:
    return str(a * b)

result = await generate_text(
    model=google("gemini-2.0-flash"),
    prompt="What is 5 + 3, then multiply the result by 2?",
    tools={
        "add": add,
        "multiply": multiply
    }
)
print(result.text)  # "The result is 16."

Sync vs Async Tools

Tools can be either synchronous or asynchronous:

# Async tool (recommended for I/O operations)
@tool(description="Fetch data from API")
async def fetch_data(url: str = Field(description="API URL")) -> str:
    async with aiohttp.ClientSession() as session:
        async with session.get(url) as response:
            return await response.text()

# Sync tool (fine for simple operations)
@tool(description="Calculate square root")
def sqrt(n: float = Field(description="Number")) -> str:
    import math
    return str(math.sqrt(n))

Tool Execution Context

Annotate one parameter as ToolExecutionContext to receive framework runtime data that must not be generated by the model. The parameter is omitted from the provider tool schema and ai-query injects a separate context for every call.

from ai_query import ToolExecutionContext, tool

@tool(description="Run a workspace command")
async def run_command(
    command: str,
    ctx: ToolExecutionContext,
) -> str:
    await ctx.emit_progress("command started", data={"command": command})
    ctx.signal.throw_if_aborted()
    result = await run_subprocess(command, signal=ctx.signal)
    await ctx.emit_progress("command finished")
    return result

The context exposes:

tool_call_id, tool_name, and step_number
the effective AbortSignal
turn_id and agent_id when execution belongs to an AgentTurn
deeply immutable application metadata
emit_progress(message, data=...) for correlated typed progress events

Pass metadata with generate_text(..., metadata={...}), stream_text(..., metadata={...}), or TurnOptions(metadata={...}). Metadata is not sent to the model. Parallel tool calls receive isolated context objects. Progress events are emitted live, so their order reflects actual concurrent execution rather than provider call order.

Automatic Tool Execution

When you provide tools, ai-query automatically:

Sends the tool definitions to the AI
Executes any tools the AI calls
Returns the results to the AI
Repeats until the AI responds without tool calls

@tool(description="Get stock price")
async def get_stock_price(symbol: str = Field(description="Stock symbol")) -> str:
    # Simulated API call
    prices = {"AAPL": 150.00, "GOOGL": 140.00, "MSFT": 380.00}
    return f"${prices.get(symbol, 0):.2f}"

result = await generate_text(
    model=google("gemini-2.0-flash"),
    prompt="What's the stock price of Apple?",
    tools={"get_stock_price": get_stock_price}
)

# The AI called get_stock_price("AAPL") automatically
print(result.text)  # "Apple (AAPL) is currently trading at $150.00."

Accessing Tool Calls and Results

Inspect what tools were called:

result = await generate_text(
    model=google("gemini-2.0-flash"),
    prompt="What's the weather in Tokyo and New York?",
    tools={"get_weather": get_weather}
)

# All tool calls made
for call in result.tool_calls:
    print(f"Called: {call.name}({call.arguments})")

# All tool results
for tr in result.tool_results:
    print(f"Result: {tr.result}")

Returning Images and Files

Return ToolOutput when a tool needs to send rich content back to the model. The same normalized value flows through ToolResult, hooks, streaming events, and persisted agent history. Provider adapters handle the wire format.

from ai_query import FilePart, ImagePart, TextPart, ToolOutput, tool

@tool(description="Render the requested report")
def render_report() -> ToolOutput:
    return ToolOutput(content=[
        TextPart(text="The report and preview are attached."),
        ImagePart(image=preview_png, media_type="image/png"),
        FilePart(
            data=report_pdf,
            media_type="application/pdf",
            filename="report.pdf",
        ),
    ])

Use ordinary strings, numbers, lists, or dictionaries for normal tool results. A list is not treated as rich content unless it is wrapped in ToolOutput. Rich tool outputs require upstream support:

OpenAI uses the Responses API. Set provider_options={"openai": {"api": "responses"}} to use it from the first request; ai-query also selects it when rich output is present in a follow-up request.
Anthropic Messages supports nested text, image, and document tool-result blocks.
Gemini multimodal function responses require a Gemini 3 or newer model.
Bedrock image tool results require Amazon Nova or Anthropic Claude 3/4.
OpenAI-compatible chat-completions endpoints do not have a lossless rich tool-result representation.

Unsupported endpoint and model combinations raise UnsupportedToolOutputError before the request is sent. Binary values are encoded for persistence and wire transport; ToolOutput representations redact their payloads.

Limiting Execution Steps

Control the maximum number of tool execution loops using stop_when:

from ai_query import generate_text
from ai_query.providers import google, step_count_is

result = await generate_text(
    model=google("gemini-2.0-flash"),
    prompt="Research this topic thoroughly.",
    tools={"search": search},
    stop_when=step_count_is(3)  # Stop after 3 steps
)

Next Steps

Agents

Build autonomous agents with advanced loop control

Start Here

Foundations

Agents

Deployment

Advanced

Defining Tools

The `@tool` Decorator

Parameter Definitions with `Field`

Supported Parameter Types

Using Tools with generate_text

Sync vs Async Tools

Tool Execution Context

Automatic Tool Execution

Accessing Tool Calls and Results

Returning Images and Files

Limiting Execution Steps

Next Steps

Agents

​Defining Tools

​The @tool Decorator

​Parameter Definitions with Field

​Supported Parameter Types

​Using Tools with generate_text

​Sync vs Async Tools

​Tool Execution Context

​Automatic Tool Execution

​Accessing Tool Calls and Results

​Returning Images and Files

​Limiting Execution Steps

​Next Steps

Agents

Defining Tools

The `@tool` Decorator

Parameter Definitions with `Field`

Supported Parameter Types

Using Tools with generate_text

Sync vs Async Tools

Tool Execution Context

Automatic Tool Execution

Accessing Tool Calls and Results

Returning Images and Files

Limiting Execution Steps

Next Steps