Creating Custom Tools for AI Agents

Defining and wrapping Flyte tasks or Python callables as tools allows LLM agents to interact with your existing infrastructure during execution. In flyte-sdk, the function_tool factory creates a bridge between Flyte's type-safe execution and the function-calling APIs of providers like Anthropic and Google Gemini.

Creating Tools from Flyte Tasks

The most common way to create a tool is by wrapping an existing Flyte task using the function_tool factory. This automatically extracts the tool's name, description (from the docstring), and input schema (from type hints).

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

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

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

# Use the tool with an agent
result = await run_agent(
    prompt="What is the weather in San Francisco?",
    tools=[weather_tool],
    model="claude-sonnet-4-20250514"
)

Automatic Schema Generation

The function_tool factory uses Flyte's internal type engine to generate the JSON schema for the LLM. This ensures that complex types—including Dataclasses, Enums, and FlyteFile—are correctly represented in the tool's input_schema.

For example, if a task uses a Literal or a dataclass, the generated schema will include the specific allowed values or the nested structure required by the LLM.

Using the function_tool Decorator

You can also use function_tool as a decorator for standard Python functions. This is useful for utility functions that do not need to be full Flyte tasks but should still be available to the agent.

from flyteplugins.gemini import function_tool

@function_tool(name="calculator", description="Perform basic arithmetic")
def add(a: int, b: int) -> int:
    return a + b

# The tool is now ready for use with Gemini

Provider-Specific Implementations

While the function_tool factory provides a consistent interface, flyte-sdk implements provider-specific FunctionTool classes to handle the unique requirements of different LLM APIs.

Anthropic (Claude)

The Anthropic implementation in flyteplugins.anthropic.agents._function_tools.FunctionTool provides a to_anthropic_tool() method that formats the schema for Claude's tool use API.

from flyteplugins.anthropic import function_tool

tool = function_tool(my_task)
anthropic_format = tool.to_anthropic_tool()
# Returns: {"name": "...", "description": "...", "input_schema": {...}}

Google Gemini

The Gemini implementation in flyteplugins.gemini.agents._function_tools.FunctionTool provides a to_gemini_tool() method that returns a google.genai.types.FunctionDeclaration.

from flyteplugins.gemini import function_tool

tool = function_tool(my_task)
gemini_format = tool.to_gemini_tool()
# Returns a FunctionDeclaration object

Execution Logic

When an agent invokes a tool, the FunctionTool.execute() method manages the execution based on the underlying object:

1. Flyte Tasks: If the tool wraps a Flyte task, it calls task.aio(). In a Flyte task context, this routes execution through the Flyte controller; otherwise, it falls back to local execution.
2. Async Functions: If the tool wraps an async def function, it is awaited directly.
3. Sync Functions: If the tool wraps a standard def function, it is run in a separate thread using asyncio.to_thread to prevent blocking the agent's event loop.

Orchestrating Multiple Tools

In complex workflows, you can pass a list of tools to run_agent. The agent will determine which tools to call and in what order to satisfy the prompt.

from flyteplugins.anthropic import function_tool, run_agent

@flyte.task
async def get_bread() -> str:
    """Get bread for making a sandwich."""
    return "sourdough bread"

@flyte.task
async def get_peanut_butter() -> str:
    """Get peanut butter."""
    return "creamy peanut butter"

tools = [
    function_tool(get_bread),
    function_tool(get_peanut_butter),
]

# The agent will call both tools to fulfill this prompt
result = await run_agent(
    prompt="I need bread and peanut butter for a sandwich.",
    tools=tools,
)

Troubleshooting and Gotchas

API Key Configuration

Agents require API keys for their respective providers. You can provide these via environment variables or directly to run_agent:

• Anthropic: ANTHROPIC_API_KEY
• Gemini: GOOGLE_API_KEY

Sync vs. Async

If you are running tools that perform heavy I/O or computation synchronously, ensure they are wrapped correctly. flyte-sdk uses asyncio.to_thread for sync functions, but for maximum performance in high-concurrency agent loops, prefer async def for your tool definitions.

Task Context

When running run_agent inside a Flyte task, ensure you use flyte.group() if you are running multiple agents in parallel. This allows Flyte to track the sub-executions correctly in the UI.

@flyte.task
async def multi_agent_task(goals: list[str]):
    for idx, goal in enumerate(goals):
        with flyte.group(f"agent-loop-{idx}"):
            await run_agent(prompt=goal, tools=my_tools)