Test Mode Overview

Relevant source files

The following files were used as context for generating this wiki page:

• osmosis_ai/rollout/test_mode/cli.py
• osmosis_ai/rollout/test_mode/runner.py (Inferred from test coverage and module structure)
• osmosis_ai/rollout/test_mode/external_llm_client.py (Inferred from test coverage)
• osmosis_ai/rollout/test_mode/dataset.py (Cited via reader implementation)
• osmosis_ai/cli.py
• tests/test_test_mode_runner.py
• tests/test_test_mode_providers.py

Test Mode is a specialized subsystem within the Osmosis SDK designed for local validation and batch testing of RolloutAgentLoop implementations. It allows developers to execute agent logic against structured datasets using external LLM providers (via LiteLLM) without requiring a connection to the full Osmosis TrainGate platform. This environment supports two primary execution paths: a high-throughput Batch Mode for automated evaluation and an Interactive Mode for step-by-step debugging of agent turns.

Sources: osmosis_ai/rollout/test_mode/cli.py:12-25, osmosis_ai/cli.py:84-88

System Architecture

The Test Mode architecture decouples the agent logic from the production rollout infrastructure by providing local mock implementations of the protocol. It utilizes a LocalTestRunner to orchestrate the lifecycle of a test, which involves loading datasets, initializing an ExternalLLMClient, and managing the state of a RolloutContext for each row in the dataset.

The following diagram illustrates the relationship between the CLI, the Runner, and the external dependencies:

Core Components

TestCommand (CLI)

The TestCommand class handles argument parsing and environment setup. It supports extensive configuration for model selection, sampling parameters (temperature, max tokens), and dataset range control (limit/offset).

Argument	Description	Default
--agent / -m	Module path to the agent loop (module:attribute)	Required
--dataset / -d	Path to .json, .jsonl, or .parquet file	Required
--model	LiteLLM model string (e.g., gpt-4o, anthropic/claude-3)	gpt-4o
--interactive	Enables step-by-step execution	False
--max-turns	Maximum agent turns allowed per dataset row	10

Sources: osmosis_ai/rollout/test_mode/cli.py:58-135

LocalTestRunner

The LocalTestRunner is responsible for the execution logic of individual test rows. It performs the following sequence for every row:

1. Metric Reset: Clears LLM usage metrics to ensure per-row accuracy.
2. Tool Injection: Calls agent.get_tools() and configures the ExternalLLMClient with the returned schemas.
3. Context Creation: Initializes a RolloutContext containing the dataset's user_prompt, system_prompt, and ground_truth.
4. Execution: Invokes agent.run(ctx) and captures the result or any runtime exceptions.

Sources: tests/test_test_mode_runner.py:133-170, tests/test_test_mode_runner.py:317-340

ExternalLLMClient

This client wraps LiteLLM to provide a unified interface for cloud providers. It automatically prefixes simple model names (e.g., gpt-4o becomes openai/gpt-4o) and tracks latency and token usage (prompt and response) for reporting.

Sources: tests/test_test_mode_providers.py:22-35, tests/test_test_mode_providers.py:61-90

Data Flow and Execution Modes

Batch Execution Flow

In Batch Mode, the system processes multiple rows concurrently or sequentially, providing real-time progress updates.

Interactive Mode

Interactive mode allows developers to step through an agent's reasoning process. It is initiated via the --interactive flag and can be targeted to a specific row using --row. It utilizes an InteractiveRunner to manage a session where the developer can inspect messages and tool calls at each turn.
Sources: osmosis_ai/rollout/test_mode/cli.py:210-244

Validation and Error Handling

Test Mode includes strict validation for tool schemas used during testing. If an agent returns an invalid tool definition, the runner will fail the test row with a ToolValidationError.

Common Validation Checks:

• Function Name: Must start with a letter or underscore and contain only alphanumeric characters/underscores.
• Parameter Types: The parameters.type field must be exactly "object".
• Structure: Ensures the tool contains the required function field.

Sources: tests/test_test_mode_runner.py:356-420

Output and Reporting

Upon completion of a batch run, the system generates a summary including:

• Pass/Fail Count: Number of successful agent completions vs. errors.
• Token Usage: Aggregate and per-row token consumption.
• Duration: Human-readable execution time (e.g., 1.5s, 2m10.3s).
• JSON Report: If the --output flag is provided, a detailed JSON file is written containing row-level results, error messages, and finish reasons.

Sources: osmosis_ai/rollout/test_mode/cli.py:307-350

Summary

The Test Mode provides a critical development environment for Osmosis agent loops, enabling rapid iteration through local datasets. By mimicking the platform's rollout protocol while leveraging standard cloud LLM providers, it ensures that agent logic is robust and accurately consumes tokens before being deployed to a production training workflow.