Skip to main content
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

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: 
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:
FieldTypeDescription
ctx.promptlist[dict[str, Any]]Input messages for the current dataset row
ctx.configTConfig | NoneCustom workflow config object, if one was provided
ctx.metadatadict[str, Any] | NonePer-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:
FrameworkUseIntegration objects
Strands AgentsStrands tools, Strands message history, migration from strands.AgentOsmosisStrandsAgent, OsmosisRolloutModel
OpenAI Agents SDKRunner.run, sessions, handoffs, OpenAI-style toolsOsmosisAgent, 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(): 
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 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(): 
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 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: 
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.
FieldTypeDefaultDescription
namestrrequiredIdentifier for the workflow
descriptionstr | NoneNoneOptional description
concurrencyConcurrencyConfigunlimitedMaximum 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: 
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

Strands Integration

Build a Strands-based rollout with tools and OsmosisStrandsAgent.

OpenAI Agents Integration

Build an OpenAI Agents SDK rollout with OsmosisAgent and OsmosisMemorySession.

Building Graders

Define reward logic for the samples your workflow produces.

Evaluation

Submit an evaluation run to test your workflow and grader before a training run.