Tool

Define tools for AI function calling

Tools enable AI models to call functions in your code. ai-query provides the @tool decorator and Field helper for type-safe tool definitions.

The @tool Decorator

from ai_query import tool, Field

@tool(description="Description shown to the AI")
async def my_tool(
    param1: str = Field(description="What this parameter does"),
    param2: int = Field(description="Another parameter", default=10)
) -> str:
    return "result"

Decorator Parameters

description · str

Optional description of what the tool does. If omitted, ai-query uses the function docstring or a generated fallback.

The Field Helper

from ai_query import Field

Field(description="...", default=None)

Field Parameters

description · str · required

Description of the parameter shown to the AI.

default · Any

Default value. If provided, the parameter becomes optional.

enum · list[Any]

Optional fixed set of allowed values for the parameter.

min_value · float

Optional minimum numeric value.

max_value · float

Optional maximum numeric value.

Supported Types

Type	Description
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 via additionalProperties
tuple[T1, T2]	Fixed-length arrays using prefixItems
tuple[T, ...]	Variable-length arrays with typed items
Optional[T]	Optional parameters
Union[T1, T2]	Union types via anyOf
Literal[...]	Enum-like constrained values
TypedDict	Nested object schemas
@dataclass types	Nested object schemas

Unrecognized annotations fall back to string inference. If you need full control, use the factory form and pass parameters explicitly.

Examples

Simple Tool

from ai_query import tool, Field

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

Async Tool

import aiohttp
from ai_query import tool, Field

@tool(description="Fetch a webpage")
async def fetch_url(
    url: str = Field(description="The URL to fetch")
) -> str:
    async with aiohttp.ClientSession() as session:
        async with session.get(url) as response:
            return await response.text()

Optional Parameters

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

Nested Parameters

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 search filters")
) -> str:
    ...

Manual Schema Override

def run_query(payload: dict) -> str:
    return "ok"

custom_tool = tool(
    description="Run a query with a custom schema",
    parameters={
        "type": "object",
        "properties": {
            "payload": {
                "type": "object",
                "properties": {
                    "sql": {"type": "string"},
                    "limit": {"type": "integer"},
                },
                "required": ["sql"],
            }
        },
        "required": ["payload"],
    },
    execute=run_query,
)

Using Tools

from ai_query import generate_text
from ai_query.providers import google

result = await generate_text(
    model=google("gemini-2.0-flash"),
    prompt="What is 5 + 3?",
    tools={"add": add}  # Pass as dict
)

Related Types

ToolCall

Represents a tool invocation by the AI:

@dataclass
class ToolCall:
    id: str           # Unique call ID
    name: str         # Tool name
    arguments: dict   # Arguments passed

ToolResult

Result from tool execution:

@dataclass
class ToolResult:
    tool_call_id: str  # Matches ToolCall.id
    result: str        # The result string

Accessing Tool Calls

result = await generate_text(
    model=google("gemini-2.0-flash"),
    prompt="Calculate 10 * 5",
    tools={"multiply": multiply}
)

# See what tools were called
for call in result.tool_calls:
    print(f"Called: {call.name}")
    print(f"Args: {call.arguments}")

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