LiteLLM Provider Integration
Relevant source files
The following files were used as context for generating this wiki page:
The LiteLLM Provider Integration serves as the unified interface for interacting with over 100 Large Language Model (LLM) providers within the Osmosis SDK. It acts as a bridge between the Osmosis training workflows—specifically rubric evaluation and agent rollout testing—and various hosted APIs like OpenAI, Anthropic, Google Gemini, and xAI. By leveraging LiteLLM, the SDK provides a consistent request/response format (OpenAI-compatible) regardless of the underlying vendor's specific API requirements.
Sources: README.md:1-10, pyproject.toml:30-35, osmosis_ai/rubric_eval.py:10-15
Core Architecture
The integration is primarily utilized in two distinct contexts: Remote Rubric Evaluation and the Local Test Mode for rollouts. The rubric_eval module uses it for scoring model outputs against natural language criteria, while ExternalLLMClient uses it to simulate LLM interactions during local agent development.
Data Flow for Rubric Evaluation
The following diagram illustrates how a rubric evaluation request is processed and routed through LiteLLM to a provider.
Rubric Evaluation Logic
The evaluate_rubric function is the primary entry point for LLM-based judging. It centralizes prompt construction and strict JSON enforcement.
Model String Normalization
The SDK handles specific vendor requirements, such as the openai/responses/ prefix required for the GPT-5 family of models. Other providers use a standard provider/model format.
| Provider | Input Model | LiteLLM Format |
|---|---|---|
| OpenAI | gpt-5-mini | openai/responses/gpt-5-mini |
| OpenAI | gpt-4o | openai/gpt-4o |
| Anthropic | claude-sonnet-4 | anthropic/claude-sonnet-4 |
| xAI | grok-2 | xai/grok-2 |
Sources: osmosis_ai/rubric_eval.py:48-65, tests/test_litellm_provider.py:25-45
Response Validation and Sanitization
The judge model is instructed to return a strict JSON object: {"score": <float>, "explanation": "<string>"}. The integration includes a sanitization layer that:
- Strips
<think>blocks from reasoning models. - Removes Markdown code fences (e.g.,
```json). - Validates the existence and type of the
scoreandexplanationfields.
Sources: osmosis_ai/rubric_eval.py:105-135, tests/test_litellm_provider.py:90-110
External LLM Client for Rollouts
In test_mode, the ExternalLLMClient class wraps LiteLLM to provide local LLM testing for agent loops. Unlike the rubric evaluator, this client handles tool/function calling translation and metrics tracking.
Components and Metrics
The client tracks usage across multiple turns in a rollout session to provide feedback on latency and token consumption.
Credential Management
API keys are resolved through a hierarchy of sources. The system can pull from explicit strings, custom environment variables, or provider-specific defaults.
Default Environment Variables
| Provider | Default Environment Variable |
|---|---|
| OpenAI | OPENAI_API_KEY |
| Anthropic | ANTHROPIC_API_KEY |
| xAI | XAI_API_KEY |
| Gemini | GEMINI_API_KEY |
| OpenRouter | OPENROUTER_API_KEY |
| Cerebras | CEREBRAS_API_KEY |
Sources: osmosis_ai/rubric_eval.py:25-38, README.md:65-75
Interactive Debugging Flow
The InteractiveRunner uses the LiteLLM integration to allow developers to step through an agent's execution. It pauses after each chat_completions call, allowing inspection of the conversation state before the next turn is processed.
Error Handling
The integration maps LiteLLM and provider-specific errors to internal Osmosis exceptions:
MissingAPIKeyError: Raised when no credential is found for the requested provider.ModelNotFoundError: Raised specifically when the provider returns a 404 (mapped fromlitellm.NotFoundError).ProviderRequestError: A general wrapper for API failures, timeouts, rate limits, and authentication errors.
Sources: osmosis_ai/rubric_eval.py:255-275, osmosis_ai/rollout/test_mode/external_llm_client.py:165-185
Summary
The LiteLLM Provider Integration is the backbone for all external model communication in the Osmosis SDK. It ensures that whether a user is evaluating a rubric with OpenAI or testing an agent loop with Anthropic, the interaction remains consistent and type-safe. By centralizing model-specific logic like timeout defaults and prefix normalization, it simplifies the addition of new LLM providers to the ecosystem.