Tutorial: Your First Gemini Agent

This tutorial walks you through building a Gemini-based agent that can execute Flyte tasks as tools. By the end, you will have a deployed agent capable of answering questions by dynamically calling your custom logic.

Prerequisites

To follow this tutorial, you need:

• The flyteplugins-gemini package installed.
• A Google Gemini API key.
• A Flyte environment configured to run tasks.

Step 1: Define Your Tools

In the Gemini plugin, any Flyte task can be converted into a tool. The plugin uses Flyte's type engine to automatically generate the JSON schema that Gemini needs to understand your function's parameters.

Create a file named agent.py and define a simple task:

import flyte
from flyteplugins.gemini import function_tool

# Define a standard Flyte task
@flyte.task
async def get_weather(city: str) -> str:
    """
    Get the current weather for a specific city.
    
    Args:
        city: The name of the city to check weather for.
    """
    # In a real app, you'd call a weather API here
    return f"The weather in {city} is currently 22°C and sunny."

# Wrap the task as a Gemini tool
weather_tool = function_tool(get_weather)

The function_tool factory extracts the tool's name and description directly from the function name and docstring. It also maps Python type hints (like str) to the appropriate JSON schema types.

Step 2: Configure the Agent

The Agent class allows you to define the personality and constraints of your assistant. You can specify the model, system instructions, and the list of tools available to it.

from flyteplugins.gemini import Agent

# Configure the agent's behavior
weather_agent_config = Agent(
    name="weather_assistant",
    instructions="You are a helpful weather assistant. Use the provided tools to answer questions.",
    model="gemini-2.5-flash",
    tools=[weather_tool],
    max_iterations=5  # Limit the number of tool-call cycles
)

Step 3: Run the Agent Loop

To execute the agent, use the run_agent function. This function manages the conversation loop: it sends your prompt to Gemini, receives tool calls, executes the corresponding Flyte tasks, and feeds the results back to the model until a final answer is reached.

You typically call run_agent from within another Flyte task:

from flyteplugins.gemini import run_agent

@flyte.task
async def ask_weather_agent(question: str) -> str:
    # run_agent handles the back-and-forth with the LLM
    return await run_agent(
        prompt=question,
        agent=weather_agent_config
    )

Step 4: Deploy with Secrets

To run this on a Flyte cluster, you must provide your GOOGLE_API_KEY. The plugin looks for this key in the environment variables. You can use Flyte's Secret and TaskEnvironment to securely inject it.

# Define the environment where the agent will run
agent_env = flyte.TaskEnvironment(
    name="gemini_env",
    image=flyte.Image.from_debian_base(name="gemini-agent").with_pip_packages(
        ["flyteplugins-gemini"]
    ),
    secrets=[
        flyte.Secret(group="google", key="api_key", env_var="GOOGLE_API_KEY")
    ]
)

# Apply the environment to your task
@agent_env.task
async def weather_bot(question: str) -> str:
    return await run_agent(
        prompt=question,
        tools=[weather_tool],
        system="You are a helpful assistant."
    )

Complete Working Result

Combining these steps, here is a complete script you can run locally or deploy to a Flyte cluster:

import flyte
from flyteplugins.gemini import function_tool, run_agent

# 1. Define the execution environment
env = flyte.TaskEnvironment(
    name="agent_env",
    image=flyte.Image.from_debian_base(name="agent").with_pip_packages(
        ["flyteplugins-gemini"]
    ),
)

# 2. Define a tool task
@env.task
async def get_current_time(location: str) -> str:
    """Get the current time for a given location."""
    return f"The time in {location} is 10:00 AM."

# 3. Define the agent task
@env.task
async def run_my_first_agent(prompt: str) -> str:
    # Convert task to tool
    time_tool = function_tool(get_current_time)
    
    # Execute the agent loop
    return await run_agent(
        prompt=prompt,
        tools=[time_tool],
        model="gemini-2.5-flash",
        system="You are a time-telling assistant."
    )

if __name__ == "__main__":
    # For local testing, ensure GOOGLE_API_KEY is in your shell environment
    import asyncio
    result = asyncio.run(run_my_first_agent(prompt="What time is it in London?"))
    print(f"Agent Response: {result}")

Next Steps

• Complex Types: Try using dataclasses or FlyteFile as inputs to your tool tasks; the plugin's integration with the Flyte type engine handles the serialization automatically.
• Parallel Execution: Gemini may request multiple tool calls in a single turn. The run_agent function handles these sequentially by default, but you can customize the execution logic by interacting with the FunctionTool class directly.
• Observability: Because your tools are Flyte tasks, every tool call made by the agent is tracked in the Flyte console with its own inputs, outputs, and logs.