Implementing Agents

Relevant source files

The following files were used as context for generating this wiki page:

• osmosis_ai/rollout/core/base.py
• osmosis_ai/rollout/server/app.py
• osmosis_ai/rollout/validator.py
• osmosis_ai/rollout/core/schemas.py
• osmosis_ai/rollout/init.py
• osmosis_ai/rollout/test_mode/interactive.py
• README.md

Implementing agents in the Osmosis SDK involves subclassing the RolloutAgentLoop to define custom logic for agent trajectories. The SDK provides a remote rollout protocol that allows these agents to integrate with Osmosis LLM training workflows, specifically interacting with a TrainGate server to perform autonomous loops, tool executions, and state reporting.

The core architecture separates the agent's logic (the "loop") from the transport and communication details. Developers focus on defining tools and the step-by-step execution flow, while the SDK handles metrics collection, protocol adherence, and server hosting via FastAPI.

Sources: osmosis_ai/rollout/core/base.py:15-35, osmosis_ai/rollout/init.py:1-15, README.md:1-10

Core Agent Components

The implementation of an agent relies on three primary classes that manage the execution lifecycle.

RolloutAgentLoop

This is the abstract base class that all agents must inherit from. It requires the definition of a unique name and the implementation of two mandatory methods: get_tools() and run().

RolloutContext

An instance of RolloutContext is passed to the agent's run() method. It provides access to the LLM client, the list of tools, and utility methods for completing the rollout or logging errors.

RolloutResult

The final object returned by the agent loop, containing the status of the execution, the final message history, and optional metrics or rewards.

Sources: osmosis_ai/rollout/core/base.py:53-150, osmosis_ai/rollout/core/base.py:356-370

Agent Class Hierarchy

The following diagram illustrates the relationship between the base classes and a concrete implementation.

The diagram shows the standard inheritance and return types for implementing a custom agent.
Sources: osmosis_ai/rollout/core/base.py:53-415

The Execution Lifecycle

The rollout execution follows a specific sequence where the SDK initializes the environment before calling the agent's logic.

1. Initialization: The server receives a RolloutRequest and calls get_tools() to determine the available capabilities.
2. Context Preparation: A RolloutContext is created, wrapping the LLM client and tracking metrics.
3. Loop Execution: The run() method is invoked. The agent typically performs multiple turns of LLM interaction and tool execution.
4. Completion: The agent returns a RolloutResult via ctx.complete() or ctx.error().
5. Callback: The SDK automatically sends the final state back to the TrainGate server.

Sources: osmosis_ai/rollout/server/app.py:255-320, osmosis_ai/rollout/core/base.py:120-145

Rollout Flow Diagram

The following flowchart describes the internal processing of a rollout request within the server.

This diagram depicts the flow from the initial request through the iterative agent logic to the final reporting.
Sources: osmosis_ai/rollout/server/app.py:255-330, osmosis_ai/rollout/core/base.py:440-475

Defining Tools

Agents interact with their environment through tools defined using OpenAI-compatible function schemas. These schemas must be returned by the get_tools() method.

Schema Component	Type	Description
type	string	Must be "function".
function.name	string	Unique name of the tool.
function.description	string	Description for the LLM to understand when to use the tool.
function.parameters	object	JSON Schema defining the required arguments.

Sources: osmosis_ai/rollout/core/schemas.py:110-150, osmosis_ai/rollout/validator.py:140-180

Tool Validation Requirements

The SDK enforces strict validation on tool definitions to ensure compatibility with LLM providers:

• name must be a non-empty string.
• parameters should be of type "object".
• get_tools() must return a list of OpenAIFunctionToolSchema or compatible dictionaries.

Sources: osmosis_ai/rollout/validator.py:185-220

Implementation Example

A basic agent implementation involves managing a message list and checking for tool calls in the LLM's response.

class CalculatorAgent(RolloutAgentLoop):
    name = "calculator"

    def get_tools(self, request: RolloutRequest) -> list:
        return [
            {
                "type": "function",
                "function": {
                    "name": "add",
                    "description": "Add two numbers",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "a": {"type": "number"},
                            "b": {"type": "number"}
                        },
                        "required": ["a", "b"]
                    }
                }
            }
        ]

    async def run(self, ctx: RolloutContext) -> RolloutResult:
        messages = list(ctx.request.messages)
        
        # Iterative loop for multi-turn reasoning
        for _ in range(ctx.request.max_turns):
            result = await ctx.chat(messages)
            messages.append(result.message)

            if not result.has_tool_calls:
                break

            # Process tool calls (logic defined by implementer)
            for tc in result.tool_calls:
                # Execution logic here...
                ctx.record_tool_call(latency_ms=10.5)
        
        return ctx.complete(messages)

Sources: osmosis_ai/rollout/core/base.py:415-475, osmosis_ai/rollout/init.py:20-35

Debugging and Monitoring

The SDK provides built-in mechanisms for monitoring agent performance and debugging execution steps.

Metrics Collection

RolloutContext automatically tracks the following metrics during execution:

• total_latency_ms: Total time from start to finish.
• llm_latency_ms: Cumulative time spent in LLM calls.
• tool_latency_ms: Cumulative time spent in tool execution (must be recorded manually via ctx.record_tool_call).
• num_llm_calls and num_tool_calls: Frequency of interactions.
• prompt_tokens and response_tokens: Token usage for billing and optimization.

Sources: osmosis_ai/rollout/core/base.py:285-310, osmosis_ai/rollout/core/schemas.py:240-255

Debug Logging

If a debug_dir is configured in the server, the agent can log custom events to a JSONL file unique to each rollout_id.

# Inside run()
ctx.log_event("pre_llm", turn=1, message_count=len(messages))

This writes an entry to {debug_dir}/{rollout_id}.jsonl.
Sources: osmosis_ai/rollout/core/base.py:220-275, osmosis_ai/rollout/server/app.py:280-290

Validation

Before serving or testing, agents can be validated using the validate_agent_loop utility. This checks for:

1. Existence and type of the name attribute.
2. Asynchronous nature of the run method.
3. Correctness and serializability of tool schemas.

Sources: osmosis_ai/rollout/validator.py:290-345

The implementation of agents provides a structured way to build complex, tool-using LLM workflows while maintaining strict metrics and protocol compliance required for training feedback loops.