> ## Documentation Index > Fetch the complete documentation index at: https://docs.osmosis.ai/llms.txt > Use this file to discover all available pages before exploring further. # Building AgentWorkflows > Implement the AgentWorkflow class to define your agent behavior for training `AgentWorkflow` is the SDK contract for rollout behavior. You subclass it, implement one async `run()` method, and create samples by calling the current policy through an Osmosis-supported agent integration. The workflow should answer one question: **given this dataset prompt, what should the agent do before the grader scores the result?** ## Base Class ```python theme={"theme":{"light":"github-light","dark":"github-dark"},"languages":{"custom":["/languages/cli.json"]}} from osmosis_ai.rollout import AgentWorkflow, AgentWorkflowContext class MyWorkflow(AgentWorkflow): async def run(self, ctx: AgentWorkflowContext) -> None: # Build and run your agent here. pass ``` The SDK shape is: ```python theme={"theme":{"light":"github-light","dark":"github-dark"},"languages":{"custom":["/languages/cli.json"]}} class AgentWorkflow(Generic[TConfig], ABC): def __init__(self, config: TConfig | None = None): self.config = config @abstractmethod async def run(self, ctx: AgentWorkflowContext[TConfig]) -> Any: raise NotImplementedError ``` `run()` is called once for each workflow execution. It should construct any per-execution agent/session objects inside the method, run the agent, and let the integration register the resulting conversation with the active `RolloutContext`. ## AgentWorkflowContext The `ctx` object gives the workflow its input and config: | Field | Type | Description | | -------------- | ------------------------ | ---------------------------------------------------------------------------------------------------- | | `ctx.prompt` | `list[dict[str, Any]]` | Input messages for the current dataset row | | `ctx.config` | `TConfig \| None` | Custom workflow config object, if one was provided | | `ctx.metadata` | `dict[str, Any] \| None` | Per-row metadata from the dataset's optional `metadata` column. `None` when the row has no metadata. | If your dataset row contains `system_prompt`, `user_prompt`, and `ground_truth`, the prompt fields are assembled into `ctx.prompt`. The reference answer is not passed to the workflow; it is exposed to your grader as `ctx.label`. The same `metadata` object is available on both `AgentWorkflowContext` and `GraderContext`, so workflows and graders can read the same per-row context. Keep task answers out of `AgentWorkflow.run()`. The workflow should produce behavior; the `Grader` should decide whether that behavior deserves reward. ## Model Routing Requirement LLM calls inside `run()` **must** route through the `RolloutContext` installed by the execution backend. The training cluster uses this context to serve the current policy, attach `x-sample-id` / `x-rollout-id` headers, collect traces, and connect rewards to samples. Use one of the supported integrations: | Framework | Use | Integration objects | | ----------------- | ---------------------------------------------------------------------- | ------------------------------------------------------------- | | Strands Agents | Strands tools, Strands message history, migration from `strands.Agent` | `OsmosisStrandsAgent`, `OsmosisRolloutModel` | | OpenAI Agents SDK | `Runner.run`, sessions, handoffs, OpenAI-style tools | `OsmosisAgent`, `OsmosisRolloutModel`, `OsmosisMemorySession` | Do not call `litellm`, the OpenAI SDK, or another provider SDK directly with a hard-coded policy model from `run()`. Direct calls bypass the rollout context and are not compatible with training. ## Strands Pattern For Strands, pass `ctx.prompt` directly as `messages` and call `invoke_async()`: ```python theme={"theme":{"light":"github-light","dark":"github-dark"},"languages":{"custom":["/languages/cli.json"]}} from osmosis_ai.rollout import AgentWorkflow, AgentWorkflowContext from osmosis_ai.rollout.integrations.agents.strands import ( OsmosisRolloutModel, OsmosisStrandsAgent, ) class SimpleStrandsWorkflow(AgentWorkflow): async def run(self, ctx: AgentWorkflowContext) -> None: agent = OsmosisStrandsAgent( name="simple-strands-agent", model=OsmosisRolloutModel(params={"temperature": 1.0}), messages=ctx.prompt, callback_handler=None, ) await agent.invoke_async() ``` Constructing `OsmosisStrandsAgent` inside `run()` binds it to the active rollout context and registers the agent as a sample source. See [Strands Integration](/cli/rollout/strands-integration) for tool examples, migration steps, and details about `OsmosisRolloutModel`. ## OpenAI Agents Pattern For OpenAI Agents, construct an `OsmosisAgent`, attach OpenAI Agents `ModelSettings`, create one `OsmosisMemorySession`, and pass that session to `Runner.run()`: ```python theme={"theme":{"light":"github-light","dark":"github-dark"},"languages":{"custom":["/languages/cli.json"]}} from agents import ModelSettings, Runner from osmosis_ai.rollout import AgentWorkflow, AgentWorkflowContext from osmosis_ai.rollout.integrations.agents.openai_agents import ( OsmosisAgent, OsmosisMemorySession, OsmosisRolloutModel, ) class SimpleOpenAIWorkflow(AgentWorkflow): async def run(self, ctx: AgentWorkflowContext) -> None: agent = OsmosisAgent( name="simple-openai-agent", instructions="Answer the user's request clearly.", model=OsmosisRolloutModel(), model_settings=ModelSettings(temperature=1.0, max_tokens=4096), ) session = OsmosisMemorySession(name="simple-openai-agent") await Runner.run( agent, ctx.prompt, session=session, ) ``` The session is what records the OpenAI Agents SDK conversation for grading. Create it inside `run()` so it registers with the current `RolloutContext`. See [OpenAI Agents Integration](/cli/rollout/openai-agents-integration) for session behavior, tracing notes, and migration steps. ## Custom Configuration Custom configs extend `AgentWorkflowConfig`. Define a module-level config instance in your rollout entrypoint and pass it to the backend: ```python theme={"theme":{"light":"github-light","dark":"github-dark"},"languages":{"custom":["/languages/cli.json"]}} from osmosis_ai.rollout import ( AgentWorkflow, AgentWorkflowConfig, AgentWorkflowContext, ConcurrencyConfig, ) class SearchWorkflowConfig(AgentWorkflowConfig): name: str = "search-workflow" max_iterations: int = 8 temperature: float = 1.0 concurrency: ConcurrencyConfig = ConcurrencyConfig(max_concurrent=4) class SearchWorkflow(AgentWorkflow[SearchWorkflowConfig]): async def run(self, ctx: AgentWorkflowContext[SearchWorkflowConfig]) -> None: config = ctx.config or SearchWorkflowConfig() max_iterations = config.max_iterations temperature = config.temperature # Use these values when constructing your agent. search_workflow_config = SearchWorkflowConfig() ``` `osmosis train submit` preflight auto-discovers at most one module-level `AgentWorkflowConfig` instance from the entrypoint module. Eval and training TOML files do not currently set workflow config fields directly. `BaseConfig` allows extra fields, so simple rollout configs usually do not need additional Pydantic boilerplate. | Field | Type | Default | Description | | ------------- | ------------------- | --------- | -------------------------------------- | | `name` | `str` | required | Identifier for the workflow | | `description` | `str \| None` | `None` | Optional description | | `concurrency` | `ConcurrencyConfig` | unlimited | Maximum concurrent workflow executions | ## Tool-Using Workflows Tool use belongs inside your agent framework, not in the backend. Define tools the way your framework expects, pass them into the Osmosis-wrapped agent, and let the integration record the resulting messages. For example, a Strands workflow can keep its tool list in config: ```python theme={"theme":{"light":"github-light","dark":"github-dark"},"languages":{"custom":["/languages/cli.json"]}} from typing import Any from strands import tool from osmosis_ai.rollout import ( AgentWorkflow, AgentWorkflowConfig, AgentWorkflowContext, ) from osmosis_ai.rollout.integrations.agents.strands import ( OsmosisRolloutModel, OsmosisStrandsAgent, ) @tool(name="search") def search_tool(query: str) -> str: """Search for information.""" return f"results for {query}" class ToolWorkflowConfig(AgentWorkflowConfig): name: str = "tool-workflow" model: Any = OsmosisRolloutModel(params={"temperature": 1.0}) tools: Any = [search_tool] max_iterations: int = 8 tool_workflow_config = ToolWorkflowConfig() class ToolWorkflow(AgentWorkflow[ToolWorkflowConfig]): async def run(self, ctx: AgentWorkflowContext[ToolWorkflowConfig]) -> None: config = ctx.config or ToolWorkflowConfig() agent = OsmosisStrandsAgent( name="search-agent", model=config.model, tools=config.tools, messages=ctx.prompt, callback_handler=None, ) for _ in range(config.max_iterations): result = await agent.invoke_async() content = result.message.get("content", []) if not any("toolUse" in block for block in content): break ``` ## Auto-Discovery `osmosis train submit` and `osmosis eval submit` both scan the rollout entrypoint module for concrete `AgentWorkflow` subclasses. You do not need decorators or registration functions, but the entrypoint still needs to construct a backend and serve it with `create_rollout_server()` so the rollout server can run on the platform. For training preflight, your entrypoint file must contain exactly **one** concrete `AgentWorkflow` subclass. If the SDK finds zero or more than one, `osmosis train submit` fails during discovery. Use helpers, base classes, tools, and configs freely, but keep only the workflow class you want to run as a concrete `AgentWorkflow` in the entrypoint. ## Next Steps Build a Strands-based rollout with tools and `OsmosisStrandsAgent`. Build an OpenAI Agents SDK rollout with `OsmosisAgent` and `OsmosisMemorySession`. Define reward logic for the samples your workflow produces. Submit an evaluation run to test your workflow and grader before a training run.