Skeleton PR: Agentic Environment for Tool Synthesis

docs/index.md

@@ -106,6 +106,7 @@ faq/oom
 
 development/dev_setup
 development/contributing
+development/agentic_synthesis_environments
 development/code_of_conduct
 development/style_guide
 development/docs_guide

docs/user_guides/synth.md

@@ -164,6 +164,76 @@ Ready to dive deeper? The sections below cover all available options in detail.
 
 ---
 
+## Environment-First Tool Synthesis
+
+Agentic synthesis now follows an environment-first model. Tools do not declare an output strategy directly. Instead, each tool is bound to an environment, and the environment type defines the execution model.
+
+- **`stateful` environments** maintain shared JSON state. Tool calls read from or update that state, which is how consistency is preserved across turns.
+- **`stateless` environments** generate tool results with an LLM. Responses are cached by input, so the same tool input can reuse the same generated output.
+- **`deterministic` environments** behave like lookup tables. Matching inputs return responses from a predefined set without LLM generation.
+
+At the config level:
+
+- Environments own their tool definitions.
+- Reusable environment catalogs live in top-level `environment_config` or `environment_config_path`.
+- Tools do not declare an `environment` field. The parent environment owns the binding.
+- `generated_output` is only used for tools in `stateless` environments.
+- `deterministic_outputs` is only used for tools in `deterministic` environments.
+- `read_only` is only meaningful for tools in `stateful` environments.
+
+Example:
+
+```yaml
+environment_config:
+  environments:
+    - id: support_backend
+      name: Support Backend
+      description: Simulated support system state
+      type: stateful
+      system_prompt: You manage support system state.
+      tools:
+        - id: get_ticket
+          name: GetTicket
+          description: Read a ticket from the support backend.
+          read_only: true
+
+    - id: faq_lookup
+      name: FAQ Lookup
+      description: Cached LLM-backed FAQ answers
+      type: stateless
+      system_prompt: Generate concise FAQ answers grounded in the tool contract.
+      tools:
+        - id: answer_faq
+          name: AnswerFAQ
+          description: Answer common support questions.
+          generated_output:
+            instruction: Return the FAQ answer for the given question.
+
+    - id: policy_table
+      name: Policy Table
+      description: Predefined policy responses
+      type: deterministic
+      tools:
+        - id: get_refund_policy
+          name: GetRefundPolicy
+          description: Return the matching refund policy.
+          deterministic_outputs:
+            - input:
+                policy_type: standard
+              output:
+                policy: Standard 30-day refund policy
+
+strategy_params:
+  multiturn_attributes:
+    - id: support_chat
+      min_turns: 2
+      max_turns: 4
+      role_instruction_messages:
+        USER: You are a customer contacting support.
+        ASSISTANT: You are a helpful support agent.
+      available_tools: [get_ticket, answer_faq, get_refund_policy]
+```
+
 ## Complete Configuration Reference
 
 ### Top-Level Parameters

pyproject.toml

@@ -45,11 +45,12 @@ dependencies = [
     "aioresponses>=0.7,<0.8",     # User by inference engine tests
     "backoff>=2.2.1,<2.3",
     "click<8.4.0",                # Used by CLI. 8.2.0 is currently unsupported by Typer.
-    "datasets>=3.2,<4.8.5",
+    "datasets>=3.2,<5",
     "greenlet",                   # Required by skypilot 0.11+ (sqlalchemy asyncio)
     "hdrhistogram>=0.10,<0.11",
     "httpx>=0.27,<1.0",           # Used by deploy module (async HTTP client)
     "jsonlines",
+    "jsonpatch>=1.33,<2.0",
     "lm_eval[wandb]>=0.4,<0.5.0",
     "mlflow>=3.1",                # >=3.1.4 requires Python3.10>=
     "numpy>=1.26,<2.4",           # verl==0.5.0 depends on numpy<2.0.0
@@ -304,7 +305,6 @@ unsupported-operator = "warn"  # Type narrowing limitations with isinstance
 too-many-positional-arguments = "warn"  # Loose typing (Callable) doesn't capture signatures
 parameter-already-assigned = "warn"  # False positives with *args/**kwargs patterns
 
-
 [tool.pytest.ini_options]
 asyncio_default_fixture_loop_scope = "function"
 testpaths = ["tests"]

src/oumi/core/configs/__init__.py

@@ -83,6 +83,7 @@
 )
 from oumi.core.configs.async_evaluation_config import AsyncEvaluationConfig
 from oumi.core.configs.base_config import BaseConfig
+from oumi.core.configs.environment_config import EnvironmentConfig
 from oumi.core.configs.evaluation_config import EvaluationConfig
 from oumi.core.configs.inference_config import InferenceConfig
 from oumi.core.configs.inference_engine_type import InferenceEngineType
@@ -158,12 +159,23 @@
 from oumi.core.configs.synthesis_config import SynthesisConfig
 from oumi.core.configs.training_config import TrainingConfig
 from oumi.core.configs.tuning_config import TuningConfig
+from oumi.environments import (
+    BaseEnvironment,
+    BaseTool,
+    DeterministicEnvironment,
+    DeterministicToolOutput,
+    GeneratedToolOutput,
+    StatefulEnvironment,
+    StatelessEnvironment,
+    ToolEnvironmentType,
+)
 
 __all__ = [
     "AsyncEvaluationConfig",
     "AutoWrapPolicy",
     "BackwardPrefetch",
     "BaseConfig",
+    "BaseEnvironment",
     "DataParams",
     "DatasetParams",
     "DatasetSplit",
@@ -176,6 +188,7 @@
     "EvaluationBackend",
     "EvaluationConfig",
     "EvaluationTaskParams",
+    "EnvironmentConfig",
     "FSDPParams",
     "GenerationParams",
     "GrpoParams",
@@ -211,17 +224,24 @@
     "TuningParams",
     "AttributeCombination",
     "DatasetSourceParam",
+    "DeterministicToolOutput",
+    "DeterministicEnvironment",
     "DocumentSegmentationParams",
     "DocumentSource",
     "ExampleSource",
+    "GeneratedToolOutput",
     "GeneratedAttributePostprocessingParams",
     "GeneralSynthesisParams",
     "GeneratedAttribute",
     "SampledAttribute",
     "SampledAttributeValue",
     "SegmentationStrategy",
+    "StatefulEnvironment",
+    "StatelessEnvironment",
     "TextConversation",
     "TextMessage",
+    "BaseTool",
+    "ToolEnvironmentType",
     "TransformationStrategy",
     "TransformationType",
     "TransformedAttribute",

src/oumi/core/configs/environment_config.py

@@ -0,0 +1,132 @@
+# Copyright 2025 - Oumi
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Configuration for agentic environments."""
+
+from dataclasses import dataclass, field
+from typing import Any
+
+from oumi.core.configs.base_config import BaseConfig
+from oumi.environments import BaseEnvironment, BaseTool
+
+
+@dataclass
+class EnvironmentConfig(BaseConfig):
+    """Top-level config for environment-first tool definitions."""
+
+    environments: list[Any] = field(default_factory=list)
+    """Reusable environments and their owned tools."""
+
+    def __post_init__(self):
+        """Verifies/populates params."""
+        self.environments = [
+            self._coerce_environment(environment) for environment in self.environments
+        ]
+
+        env_ids: set[str] = set()
+        tool_ids: set[str] = set()
+
+        for environment in self.environments:
+            if environment.id in env_ids:
+                raise ValueError(
+                    f"EnvironmentConfig.environments contains duplicate "
+                    f"environment id '{environment.id}'."
+                )
+            env_ids.add(environment.id)
+
+            for tool in environment.tools:
+                if tool.id in tool_ids:
+                    raise ValueError(
+                        f"EnvironmentConfig.environments contains duplicate "
+                        f"tool id '{tool.id}'."
+                    )
+                tool_ids.add(tool.id)
+
+    @property
+    def all_tools(self) -> list[BaseTool]:
+        """Flatten all tools across environments."""
+        return [tool for environment in self.environments for tool in environment.tools]
+
+    @property
+    def tool_environment_map(self) -> dict[str, str]:
+        """Map each tool id to the environment that owns it."""
+        return {
+            tool.id: environment.id
+            for environment in self.environments
+            for tool in environment.tools
+        }
+
+    def get_environment(self, environment_id: str) -> BaseEnvironment | None:
+        """Look up an environment by id."""
+        for environment in self.environments:
+            if environment.id == environment_id:
+                return environment
+        return None
+
+    def get_tool(self, tool_id: str) -> BaseTool | None:
+        """Look up a tool by id."""
+        for tool in self.all_tools:
+            if tool.id == tool_id:
+                return tool
+        return None
+
+    def resolve_tools(
+        self,
+        environment_ids: list[str] | None = None,
+        tool_ids: list[str] | None = None,
+    ) -> list[BaseTool]:
+        """Resolve tools from selected environments and optional tool ids.
+
+        Raises:
+            ValueError: If any environment_id or tool_id is not found.
+        """
+        all_env_ids = {env.id for env in self.environments}
+
+        if environment_ids:
+            unknown_envs = set(environment_ids) - all_env_ids
+            if unknown_envs:
+                raise ValueError(
+                    f"Unknown environment id(s): {sorted(unknown_envs)}. "
+                    f"Defined: {sorted(all_env_ids)}"
+                )
+            selected_environment_ids = environment_ids
+        else:
+            selected_environment_ids = list(all_env_ids)
+
+        selected_environments = [
+            environment
+            for environment in self.environments
+            if environment.id in set(selected_environment_ids)
+        ]
+        tools = [
+            tool for environment in selected_environments for tool in environment.tools
+        ]
+
+        if tool_ids:
+            available_tool_ids = {tool.id for tool in tools}
+            unknown_tools = set(tool_ids) - available_tool_ids
+            if unknown_tools:
+                raise ValueError(
+                    f"Unknown tool id(s): {sorted(unknown_tools)}. "
+                    f"Available in selected environments: "
+                    f"{sorted(available_tool_ids)}"
+                )
+            allowed_tool_ids = set(tool_ids)
+            tools = [tool for tool in tools if tool.id in allowed_tool_ids]
+
+        return tools
+
+    def _coerce_environment(self, environment: Any) -> BaseEnvironment:
+        """Coerce a raw dict or environment instance into a concrete environment."""
+        return BaseEnvironment.create(environment)

src/oumi/core/configs/params/synthesis_params.py

@@ -474,6 +474,16 @@ class MultiTurnAttribute:
     Allows user to specify custom instructions for the planner while planning
     out the conversation."""
 
+    available_environments: list[str] = field(default_factory=list)
+    """List of environment ids availabe in this conversation."""
+
+    available_tools: list[str] = field(default_factory=list)
+    """List of tool ids available in this conversation."""
+
+    max_tool_calls_per_turn: int = 50
+    """Safety ceiling for tool calls per ASSISTANT turn. The agent naturally stops
+    when it decides no more tools are needed. This only prevents runaway loops."""
+
     def __post_init__(self):
         """Verifies/populates params."""
         if not self.id:
@@ -543,6 +553,29 @@ def __post_init__(self):
                     "string."
                 )
 
+        if self.available_tools is not None:
+            if not isinstance(self.available_tools, list):
+                raise ValueError(
+                    "MultiTurnAttribute.available_tools must be a list of tool names."
+                )
+            for tool in self.available_tools:
+                if not isinstance(tool, str):
+                    raise ValueError(
+                        "MultiTurnAttribute.available_tools must be a list of strings."
+                    )
+        if self.available_environments is not None:
+            if not isinstance(self.available_environments, list):
+                raise ValueError(
+                    "MultiTurnAttribute.available_environments must be a list of "
+                    "environment ids."
+                )
+            for environment in self.available_environments:
+                if not isinstance(environment, str):
+                    raise ValueError(
+                        "MultiTurnAttribute.available_environments must be a list "
+                        "of strings."
+                    )
+
 
 class TransformationType(str, Enum):
     """Types of transformation strategies."""