跳转到主要内容

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 是 AWS 推出的 Agent 框架。osmosis-ai 包已内置 strands-agents[litellm] 作为依赖项,无需额外安装。

OsmosisStrandsAgent

OsmosisStrandsAgent 是 Strands Agent 类的直接替代品,它会自动:
  • 切换模型以连接到训练集群的 LLM 端点
  • 将自身注册到 RolloutContext 以实现自动样本收集
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 不接受 model_id 参数 —— SDK 在 for_sample() 内部硬编码了占位符 model id openai/osmosis-rollout,训练集群会把它路由到当前策略。params 字典用来传 temperaturemax_tokens 等采样参数(完整列表见 LiteLLM input params)。请把 agent 放在 AgentWorkflow.run() 里构造(或者放在其他已经由执行后端建立好活跃 RolloutContext 的代码路径里)。

工作原理

OsmosisRolloutModel

OsmosisRolloutModelLiteLLMModel 的子类,在构造时充当占位符。它不会直接发起 LLM 调用 —— 而是在运行时委托给训练集群。
  • 在运行时,for_sample(sample_id, rollout_ctx) 会创建一个连接到训练集群 chat completions 端点的具体 LiteLLMModel 实例
  • 添加 x-sample-idx-rollout-id 请求头,以便训练集群正确路由请求
  • 不能直接调用 —— 如果在运行时解析之前调用,会抛出 NotImplementedError

OsmosisStrandsAgent

当使用 OsmosisRolloutModel 构造 OsmosisStrandsAgent 时,会执行以下步骤:
1

读取 RolloutContext

从当前执行作用域中读取活跃的 RolloutContext。如果没有可用的上下文,则抛出 RuntimeError
2

生成 sample ID

根据 Agent 的 nameagent_id 或随机 UUID 生成 sample ID。
3

解析模型

调用 model.for_sample() 创建连接到训练集群的真实 LiteLLMModel
4

注册到上下文

通过 rollout_ctx.register_agent(sample_id, self) 将自身注册到 RolloutContext,实现自动样本收集。
5

初始化 StrandsAgent

使用解析后的模型委托给 StrandsAgent.__init__()。所有其他参数(tools、system prompt 等)保持不变传递。

完整示例

以下是使用 Strands 的完整 AgentWorkflow 和 Grader 示例:
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 已经是当前 sample 可直接使用的输入,因此可以直接作为 OsmosisStrandsAgent(messages=...) 传入,而不需要额外转换。如果数据集中的这一行包含 system_promptuser_prompt,SDK 已经帮您把这部分输入整理好了。在 async 的 run() 里,优先使用 await agent.invoke_async() —— 这和 SDK 的规范示例保持一致,也避免绕经 Strands 的同步包装层。同步形式 agent(...) 仍然可能运行,但这里没有必要。

从 StrandsAgent 迁移

如果您已有 Strands Agent,迁移到 Osmosis 只需四个步骤:
1

更新 Agent 导入

from strands import Agent 替换为:
from osmosis_ai.rollout.integrations.agents.strands import OsmosisStrandsAgent
2

更新 Model 导入

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

model = OsmosisRolloutModel(params={"temperature": 1.0})
不要再传 model_id —— 由训练集群决定要服务的策略。把所有采样参数(temperaturemax_tokens…)挪到 params 字典里。
3

替换 Agent 类

Agent(...) 替换为 OsmosisStrandsAgent(...)。所有其他构造函数参数(tools、system prompt 等)保持不变。
4

包装到 AgentWorkflow 中

将您的 Agent 逻辑放入 AgentWorkflow.run() 方法中,并用 await agent.invoke_async() 驱动。RolloutContext 由执行后端自动管理。
其他所有内容 —— 工具定义、system prompt、Agent 行为 —— 完全保持不变。

下一步

构建 AgentWorkflow

深入了解 AgentWorkflow 类和实现模式。

本地评估

使用 eval 模式测试基于 Strands 的工作流。