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 synthetic 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.

class oumi.environments.DeterministicEnvironment(params: EnvironmentParams, kwargs: DeterministicEnvironmentKwargs)

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) → DeterministicEnvironment

Build a DeterministicEnvironment 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 deterministic tool calls to their outputs.

tool_params_cls

alias of ToolParams

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

Bases: BaseParams

Type-specific kwargs for DeterministicEnvironment.

__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.ExecutableEnvironment(params: EnvironmentParams)

Bases: BaseEnvironment

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

Each tool declares its executor as a 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 envs that take user-supplied dotted-path 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. Deterministic envs project from tools (per-tool lookup-table entries); stateful synthetic 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 synthetic 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.StateGroundingConfig(state_path: str, fields: list[str])

Bases: BaseParams

Per-state-pool grounding for stateful synthetic 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

class oumi.environments.SyntheticEnvironment(params: EnvironmentParams, kwargs: SyntheticEnvironmentKwargs)

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) → SyntheticEnvironment

Build a SyntheticEnvironment 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.SyntheticEnvironmentKwargs(tool_persona: str = '', state_params: SyntheticStateParams | None = None, cache_by_input: bool = True)

Bases: BaseParams

Type-specific kwargs for SyntheticEnvironment.

__finalize_and_validate__() → None

Finalize and validate the kwargs.

__post_init__() → None

Coerce state_params dict into SyntheticStateParams if needed.

cache_by_input: bool = True

state_params: SyntheticStateParams | None = None

tool_persona: str = ''

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

Bases: BaseParams

Optional state configuration for a synthetic 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

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 deterministic-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 deterministic 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 DeterministicEnvironment 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 SyntheticEnvironment, 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