Rollout Tools Framework
Relevant source files
The following files were used as context for generating this wiki page:
The Rollout Tools Framework provides a structured way to define, validate, and execute tools within the Osmosis remote rollout protocol. It allows agent loops to interact with external environments by exposing functions that Large Language Models (LLMs) can invoke during a trajectory. The framework ensures that these interactions are type-safe, OpenAI-compatible, and properly metered for latency and usage.
Sources: osmosis_ai/rollout/__init__.py, osmosis_ai/rollout/core/base.py:270-320
Tool Architecture and Definition
The framework utilizes OpenAI-compatible schemas to define tools. A tool is primarily a "function" composed of a name, description, and a JSON Schema representing its parameters. This structure allows the LLM to understand how to call the function and what arguments are required.
Core Schemas
The system defines several Pydantic models to enforce the structure of tool definitions and runtime tool calls:
| Class | Description |
|---|---|
OpenAIFunctionToolSchema | The top-level schema for a tool definition, requiring a type (function) and a function object. |
OpenAIFunctionSchema | Defines the function's metadata, including name, description, and parameters. |
OpenAIFunctionParametersSchema | A JSON Schema object defining the properties and required fields for function arguments. |
OpenAIFunctionPropertySchema | Defines individual argument types, descriptions, and optional enum constraints. |
Sources: osmosis_ai/rollout/core/schemas.py:59-138
Tool Definition Flow
The following diagram illustrates how tools are defined and provided to the system during a rollout initialization:
Tool Execution Logic
Tool execution involves parsing LLM output, invoking the relevant logic, and serializing results back into a format the LLM can process. The framework provides utilities to handle this lifecycle concurrently.
Parsing and Information Extraction
When an LLM generates a tool call, it arrives as a JSON structure containing a function name and arguments (often as a string). The framework uses get_tool_call_info to extract the ID, name, and parsed dictionary of arguments.
# Example of tool call info extraction
call_id, name, args = get_tool_call_info(tool_call)
# Sources: [osmosis_ai/rollout/tools.py], [tests/test_rollout_tools.py:165-180]
Execution Lifecycle
The framework supports concurrent execution of multiple tool calls using execute_tool_calls. This utility takes a list of tool calls and an asynchronous executor function, returning a list of MessageDict objects formatted for the conversation history.
Handling Results and Errors
Results must be serialized into strings for the content field of a tool message. The framework provides serialize_tool_result to handle various Python types (dicts, lists, bools, etc.) and create_tool_error_result to format exceptions into a string that informs the LLM of the failure.
Sources: osmosis_ai/rollout/tools.py, tests/test_rollout_tools.py:61-100, tests/test_rollout_tools.py:265-275
Validation and Integrity
To ensure reliability, the framework includes a validation layer that checks tool implementations before the server starts.
Validation Criteria
The validate_agent_loop utility performs several checks on the tools returned by an agent:
- Schema Compliance: Ensures tools are either
OpenAIFunctionToolSchemaobjects or correctly structured dictionaries. - Naming: Validates that function names are non-empty strings.
- Descriptions: Issues warnings if functions lack descriptions, as this degrades LLM performance.
- Parameters: Checks that the
parametersfield is a dictionary and that itstypeis set to "object".
Sources: osmosis_ai/rollout/validator.py:145-235
Error Handling
The framework defines specific exceptions for tool-related failures:
ToolExecutionError: A base class for errors occurring during tool invocation.ToolArgumentError: Specifically for failures in parsing or validating tool arguments (e.g., invalid JSON).
Sources: osmosis_ai/rollout/core/exceptions.py, tests/test_rollout_tools.py:280-300
Metrics and Tracking
The framework tracks tool usage through the RolloutContext. This allows developers to monitor performance and latency within a trajectory.
# Recording a tool call within the agent loop
start = time.monotonic()
result = await execute_tool(tool_call)
latency = (time.monotonic() - start) * 1000
ctx.record_tool_call(latency_ms=latency)
Sources: osmosis_ai/rollout/core/base.py:185-200
The metrics collected include:
num_tool_calls: Total number of tools executed during the rollout.tool_latency_ms: Cumulative time spent in tool execution.
Sources: osmosis_ai/rollout/core/schemas.py:260-285, osmosis_ai/rollout/core/base.py:250-265
Conclusion
The Rollout Tools Framework is a critical component of the Osmosis SDK, providing the necessary abstractions to build interactive AI agents. By enforcing OpenAI-compatible schemas and providing robust utilities for concurrent execution, parsing, and validation, it enables stable and measurable agent-environment interactions.