Skip to main content
The osmosis-ai package includes everything you need to write reward functions, build rollout agents, and interact with the Osmosis Platform — available in three install tiers depending on your use case.
Prerequisite: You need an Osmosis Platform account to use the SDK with the training platform. Install and authenticate with pip install osmosis-ai && osmosis login.

​
Requirements

  • Python 3.10 or higher
  • pip package manager

​
Install via pip

pip install osmosis-ai
This installs the core Python library and CLI tool.

​
Optional Dependencies

Install additional features based on your use case: 
# For running remote rollout servers (FastAPI + Uvicorn)
pip install "osmosis-ai[server]"

# For development (pytest, formatting tools)
pip install "osmosis-ai[dev]"

# Install everything
pip install "osmosis-ai[full]"

​
Set Up API Keys

Configure API keys for the LLM providers you’ll use: 
# OpenAI
export OPENAI_API_KEY="sk-..."

# Anthropic (Claude)
export ANTHROPIC_API_KEY="sk-ant-..."

# Google Gemini
export GOOGLE_API_KEY="..."
You only need API keys for the providers you’ll use. The SDK supports OpenAI, Anthropic, Gemini, xAI, OpenRouter, and Cerebras.

​
Using .env Files

Create a .env file in your project: 
OPENAI_API_KEY=sk-...
ANTHROPIC_API_KEY=sk-ant-...
Load it in your code: 
from dotenv import load_dotenv
load_dotenv()

​
Verify Installation

Test the Python SDK: 
from osmosis_ai import osmosis_reward

@osmosis_reward
def test_fn(solution_str: str, ground_truth: str, extra_info: dict = None, **kwargs) -> float:
    return 1.0

print(test_fn("hello", "hello"))  # Output: 1.0
Test the CLI: 
osmosis --help

​
Troubleshooting

Import Error: Make sure you installed osmosis-ai (not osmosis-sdk) 
pip install osmosis-ai
CLI Not Found: Reinstall and verify PATH 
pip install --force-reinstall osmosis-ai
which osmosis
API Key Error: Verify environment variables are set 
import os
print(os.getenv("OPENAI_API_KEY"))

​
Next Steps

API Reference

Complete API documentation

CLI Reference

Command-line tool documentation