How to Define Agent Tools

To define tools for LLM agents in this project, you use the function_tool factory. This utility bridges standard Python functions or Flyte tasks with the tool-calling APIs of models like Anthropic Claude, Google Gemini, and OpenAI.

Define a Tool from a Flyte Task

The most common pattern is wrapping an existing Flyte task. The function_tool factory automatically extracts the task's name, docstring (for the description), and input schema using Flyte's type engine.

import asyncio
import flyte
from flyteplugins.anthropic import function_tool, run_agent

# Define a Flyte task
@flyte.task
async def get_weather(city: str) -> str:
    """Get the current weather for a city."""
    await asyncio.sleep(0.1)
    return f"The weather in {city} is sunny."

# Convert the task to a tool
weather_tool = function_tool(get_weather)

# Use the tool with an agent
async def main():
    result = await run_agent(
        prompt="What is the weather in San Francisco?",
        tools=[weather_tool],
    )
    print(result)

Define a Tool from a Python Function

You can also wrap standard Python functions. The function_tool factory will still use Flyte's type engine to generate the JSON schema for the LLM.

from flyteplugins.gemini import function_tool

def add_numbers(a: int, b: int) -> int:
    """Adds two numbers together."""
    return a + b

# Create the tool
math_tool = function_tool(add_numbers)

Use as a Decorator

For OpenAI agents, function_tool can be used directly as a decorator on top of a Flyte task.

from flyteplugins.openai import function_tool
import flyte

@function_tool
@flyte.task
async def search_database(query: str) -> str:
    """Search the internal database for information."""
    return f"Results for {query}: ..."

Customize Tool Metadata

If you need to override the name or description extracted from the function, pass them as arguments to function_tool.

from flyteplugins.anthropic import function_tool

@flyte.task
def fetch_data(id: int) -> dict:
    """Internal function to fetch data."""
    return {"id": id, "data": "sample"}

# Override name and description for the LLM
custom_tool = function_tool(
    fetch_data, 
    name="get_user_record", 
    description="Retrieves a specific user record by their numeric ID."
)

Handle Complex Types

Because function_tool uses the Flyte type engine, it supports complex types like Literal, Dataclass, and Optional out of the box. These are correctly converted to JSON schema for the LLM.

from typing import Literal
from dataclasses import dataclass
from flyteplugins.anthropic import function_tool

@dataclass
class User:
    name: str
    age: int

@flyte.task
def process_user(user: User, action: Literal["update", "delete"]) -> str:
    """Process a user record."""
    return f"Performed {action} on {user.name}"

# The generated schema will include the dataclass structure and the Literal enum
tool = function_tool(process_user)

Execution Logic

When an agent invokes a tool, the FunctionTool.execute method handles the call:

• Flyte Tasks: The tool uses task.aio(**kwargs), which routes the execution through the Flyte controller if running in a Flyte context, or executes it locally.
• Async Functions: Awaited directly.
• Sync Functions: Run in a thread executor using asyncio.to_thread to prevent blocking the event loop.

Troubleshooting

• Missing Descriptions: If a function lacks a docstring, the tool description defaults to "Execute <function_name>". Always provide a descriptive docstring so the LLM knows when to use the tool.
• Serialization: Tool results that are not strings are automatically JSON serialized before being returned to the LLM. Ensure your return types are JSON-serializable or provide a string return.
• API Keys: Ensure the appropriate environment variable is set (ANTHROPIC_API_KEY or GOOGLE_API_KEY) or passed directly to run_agent.