flyteplugins.gemini

Google Gemini plugin for Flyte.

This plugin provides integration between Flyte tasks and Google’s Gemini API, enabling you to use Flyte tasks as tools for Gemini agents. Tool calls run with full Flyte observability, retries, and caching.

Key features:

• Use any Flyte task as a Gemini tool via function_tool
• Full agent loop with automatic tool dispatch via run_agent
• Configurable agent via Agent (model, system prompt, tools, iteration limits)

Basic usage example:

import flyte
from flyteplugins.gemini import Agent, function_tool, run_agent

env = flyte.TaskEnvironment(
    name="agent_env",
    image=flyte.Image.from_debian_base(name="agent").with_pip_packages(
        "flyteplugins-gemini"
    ),
)

@env.task
async def get_weather(city: str) -> str:
    '''Get the current weather for a city.'''
    return f"Weather in {city}: sunny, 22°C"

weather_tool = function_tool(get_weather)

@env.task
async def run_weather_agent(question: str) -> str:
    return await run_agent(
        prompt=question,
        tools=[weather_tool],
        model="gemini-2.5-flash",
    )

Directory

Classes

Methods

Methods

function_tool()

def function_tool(
    func: typing.Union[flyte._task.AsyncFunctionTaskTemplate, typing.Callable, NoneType],
    name: str | None,
    description: str | None,
) -> FunctionTool | partial[FunctionTool]

Convert a function or Flyte task to a Gemini-compatible tool.

This function converts a Python function, @flyte.trace decorated function, or Flyte task into a FunctionTool that can be used with Gemini’s function calling API.

The input_schema is derived via the Flyte type engine, producing JSON schema. This ensures that Literal types, dataclasses, FlyteFile, and other Flyte-native types are represented correctly.

For @flyte.trace decorated functions, the tracing context is preserved automatically since functools.wraps maintains the original function’s metadata.

Example: ```python @env.task async def get_weather(city: str) -> str: ‘‘‘Get the current weather for a city.’’’ return f"Weather in {city}: sunny"

tool = function_tool(get_weather)
```

Returns

A FunctionTool instance that can be used with run_agent().

run_agent()

def run_agent(
    prompt: str,
    tools: list[flyteplugins.gemini.agents._function_tools.FunctionTool] | None,
    agent: flyteplugins.gemini.agents._function_tools.Agent | None,
    model: str,
    system: str | None,
    max_output_tokens: int,
    max_iterations: int,
    api_key: str | None,
) -> str

Run a Gemini agent with the given tools and prompt.

This function creates a Gemini conversation loop that can use tools to accomplish tasks. It handles the back-and-forth of function calls and responses until the agent produces a final text response.

Example: python result = await run_agent( prompt="What's the weather in SF?", tools=[function_tool(get_weather)], )

Returns

The final text response from the agent.