Skip to content

Define agents

In Genkit, agents are actions with conversation state. A standard agent renders a prompt on each turn, appends conversation history, calls the model, streams chunks, updates state, and optionally persists a snapshot. A custom agent keeps that runtime shell, but replaces the prompt-backed loop with your own code.

This page covers defining the agent itself. For when to choose an agent over a plain flow, see Full-stack agents.

  • ai.define_agent() defines the prompt instructions, tools, and the agent shell in one place. This is the common path for chat assistants, tool-calling agents, and server-backed frontend features.
  • ai.define_prompt_agent() wraps a prompt that already exists (registered with ai.define_prompt() or loaded from a Dotprompt file). It keeps prompt copy, model settings, schemas, and tool lists in the prompt layer while the agent adds conversation state and transport.
  • ai.define_custom_agent() replaces the prompt loop with your own code, allowing multiple model calls in one turn, custom planning loops, direct session control, or custom streaming.

All three constructors produce an Agent that supports the transport-agnostic chat(), load_chat(), get_snapshot(), and abort() APIs. The agent is also a bidirectional action that can be served over HTTP.

define_agent() registers a prompt-backed agent. It accepts normal prompt-generation options plus agent-specific options such as state_schema and store.

from pydantic import BaseModel
from genkit import Genkit
from genkit.agent import FileSessionStore
from genkit_google_genai import GoogleAI
ai = Genkit(plugins=[GoogleAI()])
class WeatherInput(BaseModel):
location: str
class WeatherOutput(BaseModel):
temperature_f: float
conditions: str
@ai.tool()
async def get_weather(input: WeatherInput) -> WeatherOutput:
"""Get the current weather for a location."""
return WeatherOutput(temperature_f=72, conditions=f'sunny in {input.location}')
class WeatherState(BaseModel):
last_location: str | None = None
store = FileSessionStore('./.genkit/snapshots/weather')
weather_agent = ai.define_agent(
name='weatherAgent',
description='Answers weather questions for a location.',
model='googleai/gemini-flash-latest',
system='Answer weather questions. Ask for a location when one is missing.',
tools=[get_weather],
state_schema=WeatherState,
store=store,
)
  • name registers the prompt and agent action. Use a stable, descriptive name because it appears in action metadata, route helpers, and Dev UI listings.
  • description is surfaced in action metadata and the developer UI.
  • system provides the instructions or system prompt for the model loop.
  • tools lists the tools the agent is permitted to call during its turn.
  • max_turns caps how many model/tool iterations one turn may run before stopping.
  • state_schema is a Pydantic model for custom session state. When set, chat.state, response.state, and streamed chunk.custom come back as that model.
  • store switches the agent to server-managed state. Add a store when you want snapshots, branching, background execution, load_chat(), or smaller client payloads.
  • use lists middleware applied to the prompt model loops (for example ToolApproval from genkit_middleware).
  • state_transform / chunk_transform shape state and stream chunks before they leave the server. Use them for redaction or client-specific projections.

Tools can read and update the active session by calling ai.current_session(). The session exposes update_custom(fn), get_messages(), add_messages(), add_artifacts(), and related helpers.

from pydantic import BaseModel
from genkit import Genkit
class TaskItem(BaseModel):
id: int
title: str
done: bool = False
class TaskState(BaseModel):
tasks: list[TaskItem] = []
next_id: int = 1
class AddTaskInput(BaseModel):
title: str
@ai.tool()
async def add_task(input: AddTaskInput) -> TaskItem:
"""Add a new task to the list."""
created: TaskItem | None = None
def mutate(custom: dict | None) -> dict:
nonlocal created
state = custom or {}
next_id = state.get('next_id') or state.get('nextId') or 1
tasks = list(state.get('tasks') or [])
created = TaskItem(id=next_id, title=input.title)
tasks.append(created.model_dump())
return {'tasks': tasks, 'next_id': next_id + 1}
if sess := ai.current_session():
await sess.update_custom(mutate)
return created # type: ignore[return-value]

Provide the matching state_schema when defining the agent (here state_schema=TaskState) so custom state is returned as that model. Custom-state mutations emit streamed JSON Patch chunks, which keeps chat.state and chunk.custom current while a turn is still running.

Use define_prompt_agent() when the prompt already exists under the same name. This is common with Dotprompt files because prompt authors can tune the model, schemas, tool list, and template content without touching agent wiring.

ai.define_prompt(
name='tripPlanner',
model='googleai/gemini-flash-latest',
system='Plan short trips. Keep recommendations concise.',
)
trip_planner_agent = ai.define_prompt_agent(
name='tripPlanner',
description='Plans short trips with weather-aware recommendations.',
store=store,
)

The referenced prompt is looked up when the agent is invoked. If the prompt is not registered, the turn fails with an error telling you which prompt name was missing.

Use define_custom_agent() when the standard prompt loop is too narrow, such as for multiple model calls in one turn, custom planning loops, or manual history management. A custom agent receives a SessionRunner and an ActionRunContext, and handles turn processing directly.

from genkit import ActionRunContext
from genkit.agent import AgentResult, SessionRunner
async def research_fn(sess: SessionRunner, ctx: ActionRunContext) -> AgentResult:
# Your own per-turn loop. See Custom orchestration for the full pattern.
...
research_agent = ai.define_custom_agent(
name='researchAgent',
fn=research_fn,
store=store,
)

See Custom orchestration for the runtime contract, a complete example, and failure handling.