Skip to main content

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.

Strands is an agent framework by AWS. The osmosis-ai package includes strands-agents[litellm] as a built-in dependency, so no additional installation is needed.

​
OsmosisStrandsAgent

OsmosisStrandsAgent is a drop-in replacement for the Strands Agent class that automatically:
  • Swaps the model to connect to the training cluster’s LLM endpoint
  • Registers itself with the RolloutContext for automatic sample collection
from strands import tool
from osmosis_ai.rollout import AgentWorkflow, AgentWorkflowContext
from osmosis_ai.rollout.integrations.agents.strands import (
    OsmosisRolloutModel,
    OsmosisStrandsAgent,
)

@tool
def search(query: str) -> str:
    """Search for information."""
    return f"Results for: {query}"


class SearchWorkflow(AgentWorkflow):
    async def run(self, ctx: AgentWorkflowContext) -> None:
        agent = OsmosisStrandsAgent(
            name="search-agent",
            model=OsmosisRolloutModel(params={"temperature": 1.0}),
            tools=[search],
            system_prompt="You are a helpful assistant.",
            messages=ctx.prompt,
            callback_handler=None,
        )
        await agent.invoke_async()
OsmosisRolloutModel does not take a model_id — the SDK hardcodes the placeholder model id openai/osmosis-rollout inside for_sample(), and the training cluster routes that to the current policy. Use the params dict to pass sampling parameters like temperature, max_tokens, etc. (see LiteLLM input params for the full list). Construct the agent inside AgentWorkflow.run() (or another code path where the execution backend has already installed an active RolloutContext).

​
How It Works

​
OsmosisRolloutModel

OsmosisRolloutModel is a LiteLLMModel subclass that acts as a placeholder during construction. It does not make LLM calls directly — instead, it defers to the training cluster at runtime.
  • At runtime, for_sample(sample_id, rollout_ctx) creates a concrete LiteLLMModel instance connected to the training cluster’s chat completions endpoint
  • Adds x-sample-id and x-rollout-id headers so the training cluster can route requests correctly
  • Cannot be called directly — raises NotImplementedError if invoked before runtime resolution

​
OsmosisStrandsAgent

When constructed with an OsmosisRolloutModel, OsmosisStrandsAgent performs the following steps:
1

Read RolloutContext

Reads the active RolloutContext from the current execution scope. Raises RuntimeError if no context is available.
2

Generate sample ID

Generates a sample ID from the agent’s name, agent_id, or a random UUID.
3

Resolve model

Calls model.for_sample() to create a real LiteLLMModel connected to the training cluster.
4

Register with context

Registers itself with the RolloutContext via rollout_ctx.register_agent(sample_id, self) for automatic sample collection.
5

Initialize StrandsAgent

Delegates to StrandsAgent.__init__() with the resolved model. All other parameters (tools, system prompt, etc.) pass through unchanged.

​
Complete Example

Here is a full AgentWorkflow and Grader using Strands: 
from strands import tool
from osmosis_ai.rollout import (
    AgentWorkflow,
    AgentWorkflowContext,
    Grader,
    GraderContext,
)
from osmosis_ai.rollout.integrations.agents.strands import (
    OsmosisRolloutModel,
    OsmosisStrandsAgent,
)

@tool
def calculator(expression: str) -> str:
    """Evaluate a math expression."""
    return str(eval(expression))


class MathAgentWorkflow(AgentWorkflow):
    async def run(self, ctx: AgentWorkflowContext) -> None:
        agent = OsmosisStrandsAgent(
            name="math-agent",
            model=OsmosisRolloutModel(params={"temperature": 1.0}),
            tools=[calculator],
            system_prompt="You are a math assistant. Use the calculator tool.",
            messages=ctx.prompt,
            callback_handler=None,
        )
        await agent.invoke_async()


def _last_text(sample) -> str:
    if not sample.messages:
        return ""
    content = sample.messages[-1].get("content", "")
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        return next((b["text"] for b in content if isinstance(b, dict) and "text" in b), "")
    return ""


class MathGrader(Grader):
    async def grade(self, ctx: GraderContext) -> None:
        for sample_id, sample in ctx.samples.items():
            answer = _last_text(sample)
            reward = 1.0 if ctx.label and ctx.label.strip() in answer else 0.0
            ctx.set_sample_reward(sample_id, reward)
ctx.prompt is already the ready-to-use input for the current sample, so you can pass it straight into OsmosisStrandsAgent(messages=...) with no extra reshaping. If your dataset row contains system_prompt and user_prompt, the SDK has already assembled that input for you. Inside an async run(), prefer await agent.invoke_async() — it matches the canonical SDK pattern and avoids bouncing through Strands’ synchronous wrapper. The synchronous agent(...) form may still run, but it is unnecessary here.

​
Migrating from StrandsAgent

If you already have a Strands agent, migrating to Osmosis takes four changes:
1

Update Agent import

Replace from strands import Agent with:
from osmosis_ai.rollout.integrations.agents.strands import OsmosisStrandsAgent
2

Update Model import

Replace LiteLLMModel(model_id="openai/gpt-5.2", ...) with:
from osmosis_ai.rollout.integrations.agents.strands import OsmosisRolloutModel

model = OsmosisRolloutModel(params={"temperature": 1.0})
Drop the model_id — the training cluster decides which policy to serve. Move any sampling kwargs (temperature, max_tokens, …) into the params dict.
3

Swap Agent class

Replace Agent(...) with OsmosisStrandsAgent(...). All other constructor parameters (tools, system prompt, etc.) stay the same.
4

Wrap in AgentWorkflow

Place your agent logic inside an AgentWorkflow.run() method and drive it with await agent.invoke_async(). The RolloutContext is managed automatically by the execution backend.
Everything else — tool definitions, system prompts, agent behavior — stays exactly the same.

​
Next Steps

Building AgentWorkflows

Learn more about the AgentWorkflow class and implementation patterns.

Local Evaluation

Test your Strands-based workflow with eval mode.