Developer Guide
Relevant source files
The following files were used as context for generating this wiki page:
The Osmosis AI Python SDK provides a structured framework for integrating Large Language Model (LLM) training workflows, specifically focusing on reward/rubric validation and remote rollout agent implementations. This guide outlines the architectural patterns, core components, and CLI tools required for developers to build, validate, and serve agentic loops within the Osmosis ecosystem.
The SDK is divided into two primary functional areas: Reward/Rubric Evaluation, which uses unified LLM interfaces to score model outputs, and the Remote Rollout SDK, which facilitates the execution of complex agent loops that can interact with external tools and the Osmosis platform.
Sources: README.md:1-10, osmosis_ai/rollout/init.py:1-20
Core Architectural Patterns
Agent Loop Implementation
The fundamental building block for remote rollouts is the RolloutAgentLoop. Developers subclass this base class to define the agent's identity, the tools it can access, and the asynchronous logic for its execution cycle.
The diagram above shows the relationship between the core agent loop components and the context provided during execution.
Sources: osmosis_ai/rollout/core/base.py, osmosis_ai/rollout/init.py:29-43
Remote Rollout Lifecycle
A remote rollout involves a handshake between the Osmosis TrainGate and the developer's hosted RolloutServer. This flow ensures that tools are registered before the agent loop begins processing messages.
The sequence diagram illustrates the lifecycle from initialization to the final callback.
Sources: osmosis_ai/rollout/server/app.py, tests/test_rollout_server.py:450-480
Reward and Rubric Evaluation
The SDK supports strict type enforcement for reward functions and hosted rubric evaluations via LiteLLM.
Reward Function Decorator
The @osmosis_reward decorator enforces a specific function signature to ensure compatibility with Osmosis training pipelines.
| Parameter | Type | Description |
|---|---|---|
solution_str | str | The solution string to evaluate. |
ground_truth | str | The expected correct answer. |
extra_info | dict | Optional configuration (e.g., tolerance, partial credit). |
Sources: README.md:104-124, examples/reward_functions.py:12-40
Rubric Configuration
Rubrics are defined in YAML files and support model-specific parameters and provider-unified API keys.
| Field | Type | Description |
|---|---|---|
id | str | Unique identifier for the rubric. |
rubric | str | Natural language instructions for the judge model. |
model_info | dict | Contains provider, model, and api_key_env. |
score_min/max | float | Bounds for the evaluation score (default 0.0 to 1.0). |
Sources: osmosis_ai/cli_services/config.py:18-35, examples/rubric_configs.yaml:5-20
CLI Tooling and Service Management
The osmosis CLI is the primary interface for managing authentication, workspaces, and rollout servers.
Authentication and Workspaces
Developers authenticate via an OAuth flow that stores credentials locally at ~/.config/osmosis/credentials.json. The SDK supports multi-workspace environments.
# Authentication flow
osmosis login
osmosis whoami
# Workspace switching
osmosis workspace list
osmosis workspace switch <workspace-name>
Sources: osmosis_ai/cli_commands.py:150-250, osmosis_ai/auth/local_server.py:10-40
Rollout Server Deployment
The serve command hosts a RolloutAgentLoop implementation. It includes built-in validation to ensure the module path and agent logic are correct before startup.
| Option | Description |
|---|---|
-m, --module | Path to agent loop (e.g., pkg.module:AgentClass). |
-p, --port | The port to listen on (default: 8000). |
--skip-register | Skip registration with the Osmosis Platform (local mode). |
--api-key | Set a static API key for TrainGate-to-Server authentication. |
Sources: README.md:200-230, osmosis_ai/rollout/server/serve.py
Implementation Details
Tool Execution Logic
The SDK provides utilities for parsing and executing tool calls concurrently. This is critical for agents that perform multiple side effects (e.g., calculations or API lookups) in a single turn.
# From osmosis_ai/rollout/tools.py
async def execute_tool_calls(tool_calls, executor):
"""Executes multiple tool calls concurrently using the provided executor."""
tasks = [executor(tc) for tc in tool_calls]
return await asyncio.gather(*tasks)
Sources: osmosis_ai/rollout/tools.py, tests/test_rollout_tools.py:220-250
Message Utilities
Handling LLM message dictionaries requires consistent parsing of roles and content. The SDK provides helper functions to abstract these common operations.
| Function | Purpose |
|---|---|
parse_tool_calls(msg) | Safely extracts tool_calls list from an assistant message. |
normalize_stop(stop) | Converts various stop sequence formats into a standard List[str]. |
count_messages_by_role(msgs) | Returns a dictionary mapping roles (user, assistant, tool) to counts. |
Sources: osmosis_ai/rollout/utils.py:15-80
Testing and Mocking
For local development, the SDK includes a mock_trainer that simulates the behavior of the Osmosis platform.
Mock Trainer Application
The create_mock_trainer_app function generates a FastAPI instance that implements the /v1/chat/completions and /v1/rollout/completed endpoints.
# From osmosis_ai/rollout/testing.py
tracker = RolloutCompletionTracker()
app = create_mock_trainer_app(tracker=tracker)
# Use with FastAPI TestClient for integration tests
Sources: osmosis_ai/rollout/testing.py:115-150
Conclusion
The Osmosis SDK is designed to provide a high-degree of structure for developers building LLM agents. By enforcing strict schemas through Pydantic and providing a comprehensive set of CLI tools for authentication and deployment, it reduces the friction between local agent development and production training workflows. Developers should prioritize implementing the RolloutAgentLoop correctly and using the validate CLI command to ensure protocol compliance.