Quick Start Guide
Relevant source files
The following files were used as context for generating this wiki page:
The Osmosis SDK provides a Python-based framework for LLM training workflows, focusing on reward and rubric validation and remote agent rollout. This guide outlines the core steps to install the SDK, authenticate with the Osmosis Platform, and implement evaluation logic for machine learning training loops.
Sources: README.md:1-6, pyproject.toml:6-10
Installation and Setup
The SDK requires Python 3.10 or newer and can be installed via pip. It includes dependencies for LLM interfacing (LiteLLM), data handling (PyYAML), and networking (httpx, requests).
# Standard installation
pip install osmosis-ai
# Full installation with server support
pip install "osmosis-ai[server]"
Sources: README.md:8-13, pyproject.toml:23-45
Authentication
Before using remote features or workspace-linked tools, you must authenticate using the CLI. The login process uses a local OAuth callback server to capture credentials.
The sequence of events for authenticating the SDK CLI with the Osmosis Platform.
Sources: osmosis_ai/cli_commands.py:166-200, osmosis_ai/auth/local_server.py:15-30
Core Features
The SDK is divided into three primary functional areas: Reward Functions, Rubric Evaluation, and Remote Rollouts.
1. Reward Functions
Reward functions are used to score model outputs against ground truth data. The @osmosis_reward decorator enforces a strict function signature to ensure compatibility with the training runtime.
| Parameter | Type | Description |
|---|---|---|
solution_str | str | The generated solution to evaluate. |
ground_truth | str | The expected correct answer. |
extra_info | dict | Optional configuration or metadata (defaults to None). |
Sources: README.md:118-132, examples/reward_functions.py:11-15
from osmosis_ai import osmosis_reward
@osmosis_reward
def simple_reward(solution_str: str, ground_truth: str, extra_info: dict = None) -> float:
return 1.0 if solution_str.strip() == ground_truth.strip() else 0.0
Sources: README.md:27-33, examples/reward_functions.py:18-21
2. Rubric Evaluation
Remote rubric evaluation uses LLM judges (via LiteLLM) to provide structured scoring and explanations.
| Model Info Key | Type | Requirement | Description |
|---|---|---|---|
provider | str | Required | LLM provider (e.g., openai, anthropic). |
model | str | Required | Specific model name. |
api_key_env | str | Required* | Environment variable holding the API key. |
score_min | float | Optional | Minimum score bound (default 0.0). |
score_max | float | Optional | Maximum score bound (default 1.0). |
*Note: api_key can be used as an inline alternative to api_key_env.
Sources: README.md:144-155, examples/rubric_configs.yaml:5-22
3. Remote Rollout SDK
For agent-based training, the SDK provides a RolloutServer to host agent loops. This allows the Osmosis Platform to trigger agent trajectories and collect results.
Architecture of the Remote Rollout communication flow.
Sources: README.md:214-230, osmosis_ai/rollout/init.py:15-30
CLI Usage Summary
The osmosis command-line tool manages workspaces, previews configuration files, and runs batch evaluations.
| Command | Description | Example |
|---|---|---|
login | Authenticate the CLI with Osmosis AI. | osmosis login |
workspace switch | Change the active workspace. | osmosis workspace switch my-team |
preview | Validate and display YAML/JSONL files. | osmosis preview --path rubric.yaml |
eval | Run a batch evaluation dataset. | osmosis eval --rubric r1 --data d1.jsonl |
serve | Start a RolloutServer for an agent. | osmosis serve -m my_agent:loop |
Sources: osmosis_ai/cli.py:38-80, osmosis_ai/cli_commands.py:46-60, README.md:188-210
Summary
The Osmosis SDK provides the necessary tools to bridge local development and model training with the Osmosis Platform. By implementing decorated reward functions and utilizing the CLI for workspace and rollout management, developers can integrate LLM-based evaluation and agent trajectories into their training workflows.