External LLM Client
Relevant source files
The following files were used as context for generating this wiki page:
The External LLM Client is a core component of the Osmosis SDK's "test mode," designed to facilitate local testing of agent loops against real Large Language Model (LLM) APIs. While production rollout traffic typically flows through a centralized TrainGate server, the ExternalLLMClient allows developers to bypass this infrastructure during development to exercise agent logic against over 100 supported providers including OpenAI, Anthropic, Google Gemini, and local models via Ollama.
The client acts as a standardized wrapper around the LiteLLM library, handling message format conversion, tool call translation, and detailed metrics tracking for latency and token usage. It provides a consistent interface for both batch testing and interactive debugging sessions, ensuring that agent behavior can be verified in an environment that closely mimics production without requiring the full Osmosis server stack.
Sources: osmosis_ai/rollout/test_mode/external_llm_client.py:1-24, osmosis_ai/rollout/test_mode/init.py:1-15
Architecture and Integration
The ExternalLLMClient serves as the bridge between the high-level test runners and the underlying LLM providers. It is lazily initialized to avoid hard dependencies on LiteLLM unless test mode is explicitly invoked.
Component Relationship
The following diagram illustrates how the ExternalLLMClient fits into the Osmosis test mode ecosystem:
The client converts internal Osmosis schemas into formats compatible with the selected LLM provider and standardizes the response into a CompletionsResult.
Sources: osmosis_ai/rollout/test_mode/external_llm_client.py:44-70, osmosis_ai/rollout/test_mode/runner.py:75-90, osmosis_ai/rollout/test_mode/interactive.py:60-75
Model Naming and Providers
The client follows LiteLLM's convention for model identifiers. If a simple model name is provided without a prefix (e.g., "gpt-4o"), the client automatically prefixes it with openai/.
| Format | Example | Description |
|---|---|---|
| Simple | gpt-4o | Auto-prefixed to openai/gpt-4o |
| LiteLLM Format | anthropic/claude-3-sonnet | Explicit provider/model mapping |
| Local/Custom | ollama/llama3 | Used with api_base for local proxies |
Sources: osmosis_ai/rollout/test_mode/external_llm_client.py:87-92, tests/test_test_mode_providers.py:20-30
Core Functionality
Chat Completions Lifecycle
The primary method, chat_completions, manages the asynchronous request to the LLM. It performs several key tasks:
- Request Construction: Merges default settings with user-provided
kwargs, adding tools if they have been set for the current test row. - Execution: Calls LiteLLM's
acompletion. - Metrics Recording: Captures wall-clock latency and extracts token usage from the provider's response.
- Result Normalization: Returns a
CompletionsResultcontaining the message, usage statistics, and finish reason.
Tool Management
Tools are injected into the client per test row. The set_tools method converts Pydantic tool schemas into dictionary format, specifically using exclude_none=True to prevent LLM APIs from rejecting null values in optional fields like enum.
Sources: osmosis_ai/rollout/test_mode/external_llm_client.py:101-114, tests/test_test_mode_providers.py:38-55
Metrics Tracking
The client maintains internal state to track performance across LLM calls within a single test execution.
| Metric | Field | Description |
|---|---|---|
| Latency | _llm_latency_ms | Cumulative time spent in API calls (ms). |
| Call Count | _num_llm_calls | Total number of successful LLM requests. |
| Input Tokens | _prompt_tokens | Accumulated prompt tokens used. |
| Output Tokens | _response_tokens | Accumulated completion tokens generated. |
Sources: osmosis_ai/rollout/test_mode/external_llm_client.py:208-218, tests/test_test_mode_providers.py:75-90
Interactive Mode Wrapper
In debugging scenarios, the InteractiveLLMClient wraps the ExternalLLMClient. It intercepts the chat_completions call to trigger a callback, allowing the InteractiveRunner to pause execution, display intermediate messages, and wait for user commands (e.g., "next", "continue").
Error Handling
The client translates various provider-specific errors into a unified ProviderError. This includes handling rate limits, authentication failures, budget exhaustion, timeouts, and context window violations. If LiteLLM is not installed, the client raises a descriptive ProviderError instructing the user on how to install the necessary dependencies.
| Exception Class | Trigger Condition |
|---|---|
RateLimitError | API quota or rate limit reached. |
AuthenticationError | Invalid or missing API key. |
ContextWindowExceededError | Message history exceeds model limits. |
Timeout | Request exceeds specified timeout duration. |
ImportError | LiteLLM package is missing from the environment. |
Sources: osmosis_ai/rollout/test_mode/external_llm_client.py:155-185, osmosis_ai/rollout/test_mode/exceptions.py:10-25, tests/test_test_mode_providers.py:180-195
Summary
The ExternalLLMClient is a specialized utility that enables robust local testing within the Osmosis SDK. By standardizing provider interactions through LiteLLM and providing integrated metrics and interactive capabilities, it allows developers to iterate on agent logic quickly and reliably before deploying to production environments.