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.

describe_grounding(facts: list[GroundingFact]) → str

Render grounding facts as a bulleted markdown block.

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.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 via {**input, **output} 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.

Raises:

• ValueError – If any tool_id is not declared in this env’s tools list.

• ToolLookupError – If no entry in the env’s lookup table matches the provided arguments for any call.

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.GroundingConfig(sample_size: int = 3, seed: int | None = None, tools: dict[str, ~oumi.core.configs.params.grounding_params.ToolGroundingConfig]=<factory>)

Bases: BaseParams

Per-environment grounding configuration.

__post_init__() → None

Validate sample_size and coerce tools entries.

sample_size: int = 3

Number of grounding facts sampled per conversation.

seed: int | None = None

Optional seed for reproducible grounding sampling.

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.SyntheticEnvironment(params: EnvironmentParams, kwargs: SyntheticEnvironmentKwargs)

Bases: BaseEnvironment

LLM-simulated environment with optional mutable state.

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

Inject the orchestrator’s inference engine + base config.

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.

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

Execute synthetic tool calls. Cache-misses batched per tool_id.

Raises:

• RuntimeError – If attach_inference was not called.

• ValueError – If any tool id is unknown.

• ToolError – On simulator parse failure or output_schema mismatch.

class oumi.environments.SyntheticEnvironmentKwargs(system_prompt: 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

system_prompt: 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.

__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 grounding projection.

Tools listed in GroundingConfig.tools are projected through fields when the env samples grounding facts. Tools not listed contribute nothing.

__post_init__() → None

Validate fields invariants.

fields: list[str]

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

Bases: BaseParams

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

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: dict[str, Any]

exception oumi.environments.ToolLookupError

Bases: ToolError

Raised when a tool cannot resolve an output for the given arguments.

Currently used by DeterministicEnvironment when no configured LookupEntry matches 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)

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

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: str | dict[str, Any], updated_state: dict[str, Any] | None = None)

Bases: object

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 a string or a JSON-serializable dict; the synthesizer json-encodes dicts at the message boundary.

output: str | dict[str, Any]

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