Build custom agent loops for Osmosis training infrastructure
Remote Rollout enables you to create custom agent implementations that integrate seamlessly with Osmosis training infrastructure. Your agent runs as an HTTP server while the training cluster handles LLM inference and trajectory collection.
Here’s a minimal agent that uses calculator tools:
Copy
Ask AI
from osmosis_ai.rollout import ( RolloutAgentLoop, RolloutContext, RolloutResult, RolloutRequest, create_app,)class CalculatorAgent(RolloutAgentLoop): name = "calculator" # REQUIRED: unique identifier # REQUIRED: Called when /v1/rollout/init is received # Returns tools to include in the InitResponse def get_tools(self, request: RolloutRequest): return [ { "type": "function", "function": { "name": "add", "description": "Add two numbers", "parameters": { "type": "object", "properties": { "a": {"type": "number"}, "b": {"type": "number"} }, "required": ["a", "b"] } } } ] # REQUIRED: Main agent loop executed after init async def run(self, ctx: RolloutContext) -> RolloutResult: messages = list(ctx.request.messages) for _ in range(ctx.request.max_turns): # Call LLM through training cluster result = await ctx.chat(messages) messages.append(result.message) if not result.has_tool_calls: break # Execute tools for tool_call in result.tool_calls: tool_result = await self.execute_tool(tool_call) messages.append(tool_result) ctx.record_tool_call() # OPTIONAL: for metrics # reward=None means no reward computed here # Set reward=<float> if platform expects reward from rollout server return ctx.complete(messages)# Create FastAPI appapp = create_app(CalculatorAgent())
The get_tools() method is called automatically by the SDK when /v1/rollout/init is received. The returned tools are included in the InitResponse sent back to the training cluster.