Introduction & Philosophy
Relevant source files
The following files were used as context for generating this wiki page:
The Osmosis Python SDK is designed to streamline Large Language Model (LLM) training workflows by providing robust tools for reward validation and remote rollout integration. Its primary purpose is to bridge the gap between agent frameworks and the Osmosis training infrastructure, specifically through strict type enforcement and standardized communication protocols.
Sources: README.md:1-6, pyproject.toml:7-10
The project philosophy centers on developer experience and reliability. It achieves this by providing high-level decorators for local logic, a specialized CLI for lifecycle management, and a "Remote Rollout" architecture that allows agent loops to be hosted as independent services. This modularity ensures that users can validate their models and agents in isolation before scaling to full training pipelines.
Sources: README.md:15-20, osmosis_ai/rollout/init.py:1-10
Core Pillars of the SDK
The SDK is built upon three main functional pillars: Reward Validation, Remote Rubric Evaluation, and the Remote Rollout protocol.
Reward Validation & Type Enforcement
The SDK emphasizes "strict type enforcement" for reward functions. By using decorators like @osmosis_reward and @osmosis_rubric, the system ensures that user-defined evaluation logic adheres to precise signatures required by the training runtime.
Sources: README.md:3-4, README.md:88-95
Remote Rollout Architecture
The Rollout SDK provides a lightweight layer for integrating agent frameworks with the Osmosis remote rollout protocol. It allows developers to define complex agent logic within a RolloutAgentLoop class, which the SDK then exposes via an HTTP server (TrainGate) for integration with training loops.
Sources: osmosis_ai/rollout/init.py:3-8, osmosis_ai/rollout/core/init.py:3-10
Unified Provider Interface
For rubric evaluations, the SDK utilizes LiteLLM to provide a unified interface supporting over 100 LLM providers. This abstract layer allows developers to switch between providers (OpenAI, Anthropic, Gemini, etc.) without changing code, merely by updating configuration.
Sources: README.md:43-50, pyproject.toml:23
Project Structure & Components
The SDK is organized into several key modules that handle everything from CLI interaction to network utilities.
| Component | Responsibility |
|---|---|
CLI (osmosis_ai/cli.py) | Entry point for authentication, workspace management, and server orchestration. |
Rollout Core (osmosis_ai/rollout/core) | Base classes (RolloutAgentLoop, RolloutContext) and Pydantic schemas for protocol messages. |
Server (osmosis_ai/rollout/server) | FastAPI-based server factory for hosting agent implementations. |
Validator (osmosis_ai/rollout/validator.py) | Ensures implementations conform to protocol requirements before execution. |
Auth (osmosis_ai/auth) | Manages OAuth callbacks and local credential storage. |
Sources: osmosis_ai/cli.py:20-50, osmosis_ai/rollout/core/init.py:22-35, osmosis_ai/rollout/validator.py:127-140
System Workflow
The SDK facilitates a flow where developers implement logic locally, validate it using built-in tools, and then deploy it as a remote service for training.
This diagram illustrates the lifecycle of a component from local definition to remote execution and evaluation.
Sources: README.md:144-155, osmosis_ai/cli.py:53-70, osmosis_ai/rollout/init.py:10-25
Developer Philosophy: Contract-First Integration
The SDK adopts a "contract-first" approach. Before a server is started or a reward is calculated, the SDK performs rigorous checks on signatures and schemas.
Validation Logic
The validate_agent_loop function inspects RolloutAgentLoop implementations for:
- Name Attribute: A unique identifier for the agent.
- Tool Definitions: Proper OpenAI-compatible tool schemas.
- Async Run Method: Ensuring the core logic is non-blocking.
Sources: osmosis_ai/rollout/validator.py:270-300
Signature Enforcement
For reward functions, the SDK requires exact parameter naming and type hinting. For example, a reward function must accept solution_str, ground_truth, and an optional extra_info dictionary.
@osmosis_reward
def your_function(solution_str: str, ground_truth: str, extra_info: dict = None) -> float:
return 1.0
Sources: README.md:113-118
Summary & Significance
The Osmosis Python SDK provides the necessary plumbing for modern AI training workflows. By abstracting the complexities of HTTP communication, provider-specific API formats, and authentication, it allows researchers and developers to focus on the core logic of their agents and reward models. Its significance lies in its ability to turn local Python code into production-ready training components with minimal friction.