Agent tools

In this example, we will use the openai-agents library to create a simple agent that can use tools to perform tasks. This example is based on the basic tools example example from the openai-agents-python repo.

First, create an OpenAI API key, which you can get from the OpenAI website. Then, create a secret on your Flyte cluster with:

flyte create secret OPENAI_API_KEY --value <your-api-key>

Then, we’ll use uv script to specify our dependencies.

        
    

# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b7",
#    "flyteplugins-openai>=2.0.0b7",
#    "openai-agents>=0.2.4",
#    "pydantic>=2.10.6",
# ]
# ///

Next, we’ll import the libraries and create a TaskEnvironment, which we need to run the example:

        
    
from agents import Agent, Runner
from pydantic import BaseModel

import flyte
from flyteplugins.openai.agents import function_tool


env = flyte.TaskEnvironment(
    name="openai_agents_tools",
    resources=flyte.Resources(cpu=1, memory="250Mi"),
    image=flyte.Image.from_uv_script(__file__, name="openai_agents_image"),
    secrets=flyte.Secret("OPENAI_API_KEY", as_env_var="OPENAI_API_KEY"),
)

Define the tools

We’ll define a tool that can get weather information for a given city. In this case, we’ll use a toy function that returns a hard-coded Weather object.

        
    
class Weather(BaseModel):
    city: str
    temperature_range: str
    conditions: str


@function_tool
@env.task
async def get_weather(city: str) -> Weather:
    """Get the weather for a given city."""
    return Weather(city=city, temperature_range="14-20C", conditions="Sunny with wind.")

In this code snippet, the @function_tool decorator is imported from flyteplugins.openai.agents, which is a drop-in replacement for the @function_tool decorator from openai-agents library.

Define the agent

Then, we’ll define the agent, which calls the tool:

        
    
agent = Agent(
    name="Hello world",
    instructions="You are a helpful agent.",
    tools=[get_weather],
)


@env.task
async def main() -> str:
    result = await Runner.run(agent, input="What's the weather in Tokyo?")
    print(result.final_output)
    return result.final_output

Run the agent

Finally, we’ll run the agent. Create config.yaml file, which the flyte.init_from_config() function will use to connect to the Flyte cluster:

        
    
flyte create config \
--output ~/.flyte/config.yaml \
--endpoint demo.hosted.unionai.cloud/ \
--project flytesnacks \
--domain development \
--builder remote

        
    
if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main)
    print(run.url)
    run.wait(run)

Conclusion

In this example, we’ve seen how to use the openai-agents library to create a simple agent that can use tools to perform tasks.

The full code is available here.