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 字典用来传 temperature、max_tokens 等采样参数(完整列表见 LiteLLM input params)。请把 agent 放在 AgentWorkflow.run() 里构造(或者放在其他已经由执行后端建立好活跃 RolloutContext 的代码路径里)。
工作原理
OsmosisRolloutModel
OsmosisRolloutModel 是 LiteLLMModel 的子类,在构造时充当占位符。它不会直接发起 LLM 调用 —— 而是在运行时委托给训练集群。
- 在运行时,
for_sample(sample_id, rollout_ctx) 会创建一个连接到训练集群 chat completions 端点的具体 LiteLLMModel 实例
- 添加
x-sample-id 和 x-rollout-id 请求头,以便训练集群正确路由请求
- 不能直接调用 —— 如果在运行时解析之前调用,会抛出
NotImplementedError
OsmosisStrandsAgent
当使用 OsmosisRolloutModel 构造 OsmosisStrandsAgent 时,会执行以下步骤:
读取 RolloutContext
从当前执行作用域中读取活跃的 RolloutContext。如果没有可用的上下文,则抛出 RuntimeError。
生成 sample ID
根据 Agent 的 name、agent_id 或随机 UUID 生成 sample ID。
解析模型
调用 model.for_sample() 创建连接到训练集群的真实 LiteLLMModel。
注册到上下文
通过 rollout_ctx.register_agent(sample_id, self) 将自身注册到 RolloutContext,实现自动样本收集。
初始化 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_prompt 和 user_prompt,SDK 已经帮您把这部分输入整理好了。在 async 的 run() 里,优先使用 await agent.invoke_async() —— 这和 SDK 的规范示例保持一致,也避免绕经 Strands 的同步包装层。同步形式 agent(...) 仍然可能运行,但这里没有必要。
从 StrandsAgent 迁移
如果您已有 Strands Agent,迁移到 Osmosis 只需四个步骤:
更新 Agent 导入
将 from strands import Agent 替换为:from osmosis_ai.rollout.integrations.agents.strands import OsmosisStrandsAgent
更新 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 —— 由训练集群决定要服务的策略。把所有采样参数(temperature、max_tokens…)挪到 params 字典里。 替换 Agent 类
将 Agent(...) 替换为 OsmosisStrandsAgent(...)。所有其他构造函数参数(tools、system prompt 等)保持不变。
包装到 AgentWorkflow 中
将您的 Agent 逻辑放入 AgentWorkflow.run() 方法中,并用 await agent.invoke_async() 驱动。RolloutContext 由执行后端自动管理。
其他所有内容 —— 工具定义、system prompt、Agent 行为 —— 完全保持不变。
下一步
构建 AgentWorkflow
深入了解 AgentWorkflow 类和实现模式。
本地评估
使用 eval 模式测试基于 Strands 的工作流。