oumi.environments

Environments for agentic tool interactions.

Importing this package populates the environment registry by triggering each concrete environment’s @register_environment(…) decorator.

class oumi.environments.BaseEnvironment

Bases: ABC

Abstract base class for tool environments.

close() → None

Release resources owned by this env. Default no-op.

Counterpart to requires_isolation: envs rebuilt per sample (e.g. a DB session) override this to release their resources at episode end.

describe_grounding(facts: list[GroundingFact]) → str

Render grounding facts as a bulleted markdown block.

requires_isolation() → bool

Whether this env must be rebuilt per sample to avoid cross-sample leakage.

Default False — the env is safe to share across samples. Override and return True for envs carrying mutable per-sample state (e.g. stateful simulated tool execution).

sample_grounding(n: int, *, rng: Random, tool_ids: set[str] | None = None) → list[GroundingFact]

Sample grounding facts from this environment. Default: [].

abstractmethod step(calls: list[tuple[str, dict[str, Any]]]) → list[ToolResult]

Execute a batch of tool calls; returns results in the same order.

tool_params_cls

alias of ToolParams

class oumi.environments.DatabaseExecutableEnvironment(params: EnvironmentParams, session: DatabaseSession)

Bases: ExecutableEnvironment

Runs SQL-executing tools against an isolated database session.

close() → None

Roll back the episode’s writes and tear down the session.

classmethod from_params(params: EnvironmentParams) → DatabaseExecutableEnvironment

Build the env, opening a session over its configured DB.

db_path shares one snapshot file across rollouts (scales to large DBs) and is safe for concurrent readers, but SQLite serializes concurrent writers on one file, so write-heavy concurrent rollouts should use schema_sql (a fresh per-rollout file) instead.

requires_isolation() → bool

Each rollout needs its own session; never share across samples.

step(calls: list[tuple[str, dict[str, Any]]]) → list[ToolResult]

Execute a batch atomically within the rollout transaction.

oumi.environments.DeterministicEnvironment

alias of LookupEnvironment

oumi.environments.DeterministicEnvironmentKwargs

alias of LookupEnvironmentKwargs

exception oumi.environments.EndpointCallError

Bases: Exception

Raised when the endpoint cannot be reached or answers unusably.

Deliberately not a ToolError: the tool never answered, so a caller can tell an endpoint failure apart from a tool refusing the call and choose whether to retry it, fail the row, or tell the model something generic.

class oumi.environments.EndpointEnvironment(params: EnvironmentParams, protocol: EndpointProtocol)

Bases: BaseEnvironment

Environment that executes each tool call against a remote endpoint.

The endpoint owns the tool’s behavior; this environment owns the contract: it validates arguments against the tool’s schema, hands the call to a protocol, and validates the answer against the tool’s output_schema. The protocol owns the wire format, so a different one leaves this class untouched.

Each call is identified by a caller-supplied session_id naming the conversation and a call_id naming the call within it, so a protocol can make retries deduplicable end to end.

Shareable across samples, so the harness never closes it. Whoever builds the protocol’s client owns releasing it.

call(tool_id: str, arguments: dict[str, Any], *, call_id: str, session_id: str) → ToolResult

Execute one tool call against the endpoint.

session_id names the conversation and call_id the call within it. Both come from the caller: this env is shared across samples, so it cannot know which conversation a call belongs to.

Raises:

• ToolLookupError – If the environment does not serve tool_id.

• ToolArgumentError – If arguments do not match the tool’s schema.

• ToolError – As raised by the protocol, when the tool itself refused the call.

• EndpointCallError – If the endpoint is unreachable or its response does not match the tool’s output schema.

classmethod from_params(params: EnvironmentParams) → EndpointEnvironment

Build the env from its configured kwargs, over a plain JSON POST.

step(calls: list[tuple[str, dict[str, Any]]]) → list[ToolResult]

Execute a batch of tool calls; results are returned in input order.

Identifies the calls itself, treating the batch as one conversation. Callers holding stable ids for a conversation and its calls should use call() so retries stay deduplicable end to end.

tool_params_cls

alias of ToolParams

class oumi.environments.EndpointEnvironmentKwargs(endpoint_url: str = '', timeout_seconds: float = 30.0, max_retries: int = 3)

Bases: BaseParams

Type-specific kwargs for EndpointEnvironment.

__finalize_and_validate__() → None

Validate the endpoint configuration.

endpoint_url: str = ''

The endpoint every tool call is sent to.

max_retries: int = 3

Retries for a connection failure or a retryable status.

A retry re-sends the identical body, so an endpoint deduplicating on {session_id}:{call_id} performs the side effect once. Set to 0 for an endpoint that does not deduplicate.

timeout_seconds: float = 30.0

How long to wait for the endpoint to answer one call.

class oumi.environments.EndpointProtocol(*args, **kwargs)

Bases: Protocol

Turns one tool call into a request, and its answer into the tool’s output.

Implementations own the wire format: the shape of the request, and which answers count as the tool refusing rather than the endpoint failing.

Raise ToolError (or a subclass) when the tool itself refused the call — an in-band answer that reaches the model verbatim. Raise anything else when the endpoint could not answer at all.

A protocol also brings the client that suits it. MCP over Streamable HTTP answers a call with either JSON or an event stream, so it declares its own client rather than reusing JsonHttpClient. The client holds the connection, so releasing it is the job of whoever built it.

Example

class McpProtocol:
    def __init__(self, mcp_client):
        self._mcp_client = mcp_client

    def call(self, request):
        result = self._mcp_client.call_tool(
            name=request.name, arguments=request.arguments
        )
        if result.get("isError"):
            raise ToolError(result["content"][0]["text"])
        return result["structuredContent"]

call(request: RemoteToolCall) → JsonValue

Execute one tool call and return the tool’s output.

exception oumi.environments.EndpointStatusError(status_code: int, message: str)

Bases: Exception

Raised by a client when the endpoint answered with a non-2xx status.

Carries the status so a protocol can tell the tool refusing the call apart from the endpoint failing to serve it. Clients raise this instead of their own HTTP error type, which is what keeps that decision in the protocol.

class oumi.environments.ExecutableEnvironment(params: EnvironmentParams)

Bases: BaseEnvironment

Abstract base for envs that dispatch tool calls to Python executors.

Each tool declares its executor as a registry name or dotted import path; the base resolves them into _executors at construction. Subclasses supply the per-call execution context (DB connection, HTTP client, FS root, …) by implementing _build_execution_context. The base owns tool lookup, argument and result validation, the _absorb_result post-hook, and the close lifecycle. Executors are invoked as executor(arguments=<dict>, context=<ctx>) and must return a ToolResult. Result validation runs inside the execution context so a transactional context manager sees a validation failure and can roll back; _absorb_result runs only after the context exits cleanly.

step(calls: list[tuple[str, dict[str, Any]]]) → list[ToolResult]

Execute a batch of tool calls; results are returned in input order.

tool_params_cls

alias of ExecutableTool

class oumi.environments.ExecutableTool(id: str, name: str, description: str, parameters: dict[str, Any]=<factory>, output_schema: dict[str, Any] | None=None, read_only: bool = True, executor: str = '')

Bases: ToolParams

ToolParams variant for environments with user-supplied executors.

__post_init__() → None

Validate inherited fields and enforce non-empty executor.

class oumi.environments.GroundingConfig(sample_size: int = 3, seed: int | None = None, tools: dict[str, ~oumi.core.configs.params.grounding_params.ToolGroundingConfig]=<factory>, state: list[StateGroundingConfig] = <factory>)

Bases: BaseParams

Per-environment grounding configuration.

Both sub-blocks are optional; envs read whichever applies. Lookup envs project from tools (per-tool lookup-table entries); stateful simulated envs project from state (per-pool initial_state rows).

__post_init__() → None

Validate sample_size and coerce tools/state entries.

sample_size: int = 3

Number of grounding facts sampled per conversation.

seed: int | None = None

Optional seed for reproducible grounding sampling.

state: list[StateGroundingConfig]

Per-state-pool projections for stateful simulated envs.

tools: dict[str, ToolGroundingConfig]

Per-tool field whitelists, keyed by tool id.

class oumi.environments.GroundingFact(data: dict[str, ~typing.Any]=<factory>)

Bases: BaseParams

Env-agnostic representation of a single grounding fact.

Environments produce these during sampling; the planner prompt renders each fact’s data dict as one bullet line. Values are expected to be JSON-serializable scalars.

data: dict[str, Any]

class oumi.environments.JSONSchema(*, type: Literal['object', 'string', 'number', 'integer', 'boolean', 'array', 'null'] | list[Literal['object', 'string', 'number', 'integer', 'boolean', 'array', 'null']] | None = None, description: str | None = None, title: str | None = None, properties: dict[str, JSONSchema] | None = None, required: list[str] | None = None, items: JSONSchema | None = None, enum: list[Any] | None = None, default: Any = None, format: str | None = None, **extra_data: Any)

Bases: BaseModel

A JSON Schema object describing the shape of a value.

Models the subset of JSON Schema commonly used in LLM tool definitions. extra="allow" lets less-common keywords ($ref, $defs, additionalProperties, anyOf, numeric constraints, etc.) round-trip unchanged, matching the rest of this module — see the module docstring for why round-tripping matters.

default: Any

Default value used when the field is omitted.

description: str | None

Human-readable description, used by the model to choose values.

enum: list[Any] | None

Restricts the value to a fixed set of allowed values.

format: str | None

Semantic format hint (e.g., "date-time", "email").

items: JSONSchema | None

schema for array elements.

Type:

For type="array"

model_config = {'extra': 'allow', 'frozen': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

properties: dict[str, JSONSchema] | None

schema for each named property.

Type:

For type="object"

required: list[str] | None

names of properties that must be present.

Type:

For type="object"

title: str | None

Short human-readable label.

type: Literal['object', 'string', 'number', 'integer', 'boolean', 'array', 'null'] | list[Literal['object', 'string', 'number', 'integer', 'boolean', 'array', 'null']] | None

JSON type(s) of this value. A list expresses a union (e.g., ["string", "null"] for a nullable string).

class oumi.environments.JsonHttpClient(*args, **kwargs)

Bases: Protocol

POSTs a JSON body to one fixed endpoint and decodes the JSON answer.

What JsonHttpProtocol sends over. The client owns the URL, the credential, and the egress policy, so a protocol never chooses where a call goes or what it is allowed to reach. A protocol needing more than a JSON POST declares its own client rather than widening this one.

post_json(payload: JsonValue) → JsonValue

POST payload and return the decoded response body.

class oumi.environments.JsonHttpProtocol(http_client: JsonHttpClient)

Bases: object

Sends each tool call as one JSON POST whose response body is the output.

The request carries the call’s name and arguments alongside the ids identifying it. A retry resends the identical body, so an endpoint that deduplicates on {session_id}:{call_id} performs a side effect once no matter how often it is re-sent.

call(request: RemoteToolCall) → JsonValue

POST one call and return the response body as the tool’s output.

Raises:

• ToolError – If the endpoint answered 4xx, which is the tool rejecting the call rather than the endpoint failing to serve it.

• EndpointStatusError – For any other non-2xx status.

class oumi.environments.LookupEnvironment(params: EnvironmentParams, kwargs: LookupEnvironmentKwargs)

Bases: BaseEnvironment

Environment that resolves tools from a per-tool lookup table.

The env’s env_kwargs.lookup_table is the source of truth for tool behavior. Tools listed in params.tools declare contracts only; their data lives on the env.

classmethod from_params(params: EnvironmentParams) → LookupEnvironment

Build a LookupEnvironment from its params object.

sample_grounding(n: int, *, rng: Random, tool_ids: set[str] | None = None) → list[GroundingFact]

Sample grounding facts from per-tool projected pools.

Walks every tool that has a per-tool entry in params.grounding.tools. Each entry in that tool’s lookup table is projected to its input fields (merged with output when the output is a dict), filtered through the configured fields whitelist. Tools without a grounding entry contribute nothing.

step(calls: list[tuple[str, dict[str, Any]]]) → list[ToolResult]

Resolve a batch of lookup tool calls to their outputs.

tool_params_cls

alias of ToolParams

class oumi.environments.LookupEnvironmentKwargs(lookup_table: dict[str, list[~oumi.environments.lookup_environment.ToolLookupEntry]]=<factory>)

Bases: BaseParams

Type-specific kwargs for LookupEnvironment.

__post_init__() → None

Coerce raw entry dicts into ToolLookupEntry instances.

lookup_table: dict[str, list[ToolLookupEntry]]

Per-tool list of (input, output) entries, keyed by tool id.

class oumi.environments.RemoteToolCall(name: str, arguments: dict[str, Any], call_id: str, session_id: str)

Bases: object

One tool call to send, identified within its conversation.

arguments: dict[str, Any]

Arguments, already validated against the tool’s schema.

call_id: str

Identifies this call. Stable across retries of the same call.

name: str

Id of the tool being called.

session_id: str

Identifies the conversation the call belongs to.

class oumi.environments.RequestsJsonClient(url: str, timeout_seconds: float, max_retries: int = 3)

Bases: object

Default client: a plain JSON POST with no egress policy of its own.

close() → None

Release the pooled connections. Not part of JsonHttpClient.

post_json(payload: JsonValue) → JsonValue

POST payload and return the decoded response body.

class oumi.environments.SimulatedEnvironment(params: EnvironmentParams, kwargs: SimulatedEnvironmentKwargs)

Bases: BaseEnvironment

LLM-simulated environment with optional mutable state.

See the module docstring for the stateless vs stateful contract.

attach_inference(engine: BaseInferenceEngine, base_config: InferenceConfig) → None

Inject the orchestrator’s inference engine + base config.

build_call_conversation(tool_id: str, arguments: dict[str, Any]) → Conversation

Build the simulator conversation for one tool call.

property current_state: dict[str, Any] | None

Return the current in-memory state snapshot.

classmethod from_params(params: EnvironmentParams) → SimulatedEnvironment

Build a SimulatedEnvironment from its params object.

parse_tool_response(tool_id: str, response: Conversation) → ToolResult

Extract a ToolResult from a simulator response conversation.

requires_isolation() → bool

Stateful synth envs need per-sample isolation; stateless do not.

sample_grounding(n: int, *, rng: Random, tool_ids: set[str] | None = None) → list[GroundingFact]

Project grounding facts from grounding.state pools.

No-op for stateless envs or envs without grounding.state entries. tool_ids is accepted for BaseEnvironment signature compatibility but ignored — state grounding is pool-scoped, not tool-scoped.

_validate_state_grounding at init guarantees each state_path resolves to a list in self._state, and state_schema validation on every commit keeps it that way, so the projection loop trusts the shape.

step(calls: list[tuple[str, dict[str, Any]]]) → list[ToolResult]

Execute tool calls. See module docstring for routing rules.

Raises:

• ValueError – If any tool id is unknown.

• RuntimeError – If an LLM-simulated tool is invoked before attach_inference was called.

• ToolError – On simulator parse failure or schema mismatch.

class oumi.environments.SimulatedEnvironmentKwargs(tool_persona: str = '', state_params: SimulatedStateParams | None = None, cache_by_input: bool = True, use_guided_decoding: bool = True)

Bases: BaseParams

Type-specific kwargs for SimulatedEnvironment.

__finalize_and_validate__() → None

Finalize and validate the kwargs.

__post_init__() → None

Coerce state_params dict into SimulatedStateParams if needed.

cache_by_input: bool = True

state_params: SimulatedStateParams | None = None

tool_persona: str = ''

use_guided_decoding: bool = True

Constrain simulator output to each tool’s output_schema.

Applies to every tool in this environment. Set to False to generate freely and rely on the jsonschema post-validation instead: large output_schema values can exceed a provider’s grammar-compiler limits or make constrained decoding several times slower.

class oumi.environments.SimulatedStateParams(state_schema: dict[str, Any] | None = None, initial_state: dict[str, Any] | None = None)

Bases: BaseParams

Optional state configuration for a simulated environment.

State grounding for these pools is declared at the env level via EnvironmentParams.grounding.state — each entry’s state_path must resolve to a list[dict] in initial_state.

__post_init__()

Validate state config consistency.

initial_state: dict[str, Any] | None = None

state_schema: dict[str, Any] | None = None

class oumi.environments.StateGroundingConfig(state_path: str, fields: list[str])

Bases: BaseParams

Per-state-pool grounding for stateful simulated environments.

Projects rows from initial_state[state_path] through fields.

__post_init__() → None

Validate state_path and fields invariants.

fields: list[str]

state_path: str

oumi.environments.SyntheticEnvironment

alias of SimulatedEnvironment

oumi.environments.SyntheticEnvironmentKwargs

alias of SimulatedEnvironmentKwargs

oumi.environments.SyntheticStateParams

alias of SimulatedStateParams

exception oumi.environments.ToolArgumentError

Bases: ToolError

Raised when tool-call arguments fail schema validation.

exception oumi.environments.ToolError

Bases: Exception

Base class for tool errors surfaced back to the LLM.

Subclasses of this exception are caught by the tool-call loop and re-emitted as structured tool messages so the model can self-correct on the next iteration.

class oumi.environments.ToolGroundingConfig(fields: list[str])

Bases: BaseParams

Per-tool field whitelist for lookup-env grounding projection.

__post_init__() → None

Validate fields is non-empty and de-duplicated.

fields: list[str]

class oumi.environments.ToolLookupEntry(input: dict[str, ~typing.Any]=<factory>, output: JsonValue = None)

Bases: BaseParams

One (input, output) pair in a lookup env’s lookup table.

output may be any JSON value (scalar, list, object, or null).

input: dict[str, Any]

input_key() → str

Canonical JSON form of input for matching and dedup.

matches(arguments: dict[str, Any]) → bool

Check if the input matches the given arguments.

output: JsonValue = None

exception oumi.environments.ToolLookupError

Bases: ToolError

Raised when a tool call cannot be resolved.

Covers a tool id that isn’t registered in the environment (ExecutableEnvironment) and a LookupEnvironment LookupEntry that matches none of the provided arguments.

class oumi.environments.ToolParams(id: str, name: str, description: str, parameters: dict[str, ~typing.Any]=<factory>, output_schema: dict[str, ~typing.Any] | None=None, read_only: bool = True, executor: str = '')

Bases: BaseParams

Tool schema owned by an environment.

parameters and output_schema are stored as plain JSON-Schema dicts so OmegaConf can carry them through YAML round-trips. They are converted to a Pydantic JSONSchema only at the wire-format boundary in to_tool_definition().

__post_init__()

Validate common tool fields.

Accepts JSONSchema instances on parameters / output_schema for callers that build a tool with Pydantic types directly; converts them to dicts so the canonical in-memory shape stays JSON-Schema-shaped.

classmethod create(raw: Any) → ToolParams

Create a tool from raw config data.

description: str

executor: str = ''

Optional dotted import path to a callable that executes this tool.

When set, the host env dispatches calls to it instead of LLM-simulating. The callable is invoked with keyword arguments — always arguments (the tool-call args), plus a per-env context keyword: state for a stateful SimulatedEnvironment, context for an ExecutableEnvironment — and must return a ToolResult.

id: str

name: str

output_schema: dict[str, Any] | None = None

parameters: dict[str, Any]

read_only: bool = True

to_llm_schema() → dict[str, Any]

Export a provider-agnostic schema for LLM tool registration.

to_tool_definition() → ToolDefinition

Project to OpenAI-wire-format ToolDefinition.

Drops chain-internal fields (output_schema, read_only, name display label) that have no slot in the OpenAI contract. Coerces parameters to JSONSchema at the boundary.

validate_arguments(arguments: dict[str, Any]) → None

Validate call-time arguments against this tool’s parameters schema.

Raises:

ToolArgumentError – If arguments do not conform.

class oumi.environments.ToolResult(*, output: JsonValue, updated_state: dict[str, Any] | None = None)

Bases: BaseModel

Result returned by an environment step().

Runtime value (not an OpenAI wire-format type) — projected by the synthesizer into Message(role=TOOL, content=...) before output. output may be any JSON value — a string is used as-is, everything else is json-encoded at the message boundary.

model_config = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

output: JsonValue

updated_state: dict[str, Any] | None