Upload 53 files

Dockerfile

@@ -0,0 +1,14 @@
+FROM python:3.11-slim
+
+WORKDIR /app
+
+ENV PYTHONUNBUFFERED=1
+
+COPY . /app
+
+RUN pip install --no-cache-dir --upgrade pip && \
+    pip install --no-cache-dir .
+
+EXPOSE 8000
+
+CMD ["uvicorn", "server.app:app", "--host", "0.0.0.0", "--port", "8000"]

README.md

@@ -1,10 +1,127 @@
----
-title: ToolUseEnv
-emoji: 📊
-colorFrom: gray
-colorTo: yellow
-sdk: docker
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+## Why this environment
+
+Modern tool-using agents often fail on operational workflows that require evidence gathering, policy compliance, and safe escalation. This environment targets that gap with deterministic tasks that resemble what ecommerce support, trust-and-safety, and operations agents do every day.
+
+## Task set
+
+The benchmark ships with three deterministic tasks and matching deterministic graders:
+
+1. `damaged-mug-replacement` (`easy`)
+   Resolve a damaged-item replacement request.
+2. `duplicate-charge-refund` (`medium`)
+   Investigate a duplicate billing complaint and refund the extra capture.
+3. `account-takeover-fraud` (`hard`)
+   Handle a suspected account takeover with a security-first fraud escalation.
+
+Each task has a fixed expected resolution, required evidence, and reply keywords. The grader returns a score in `[0.0, 1.0]` from weighted resolution accuracy, evidence coverage, reply quality, and efficiency.
+
+## Action space
+
+The environment uses a typed `ToolUseAction` model with these actions:
+
+- `review_ticket`
+- `inspect_artifact`
+- `search_policy`
+- `draft_reply`
+- `submit_resolution`
+
+Optional fields on the action are `artifact_id`, `query`, `message`, and `resolution_code`.
+
+## Observation space
+
+The typed `ToolUseObservation` includes:
+
+- `task_id`, `difficulty`, `objective`
+- `customer_message`
+- `workspace_summary`
+- `available_actions`
+- `available_resolution_codes`
+- `collected_evidence`
+- `last_tool_result`
+- `last_action_error`
+- `remaining_steps`
+- `current_score`
+
+The typed `ToolUseState` exposes internal progress such as `final_score`, `drafted_reply`, `resolution_code`, `required_evidence`, `collected_evidence`, and action history.
+
+## Reward design
+
+The reward is shaped over the full trajectory:
+
+- Positive reward for first-time collection of relevant artifacts and policies
+- Smaller reward for drafting a reply that includes required customer-facing details
+- Very small or zero reward for repeated or invalid actions
+- Final step reward equal to the deterministic grader score
+
+This gives agents signal before the final submission while still anchoring the episode outcome to task completion quality.
+
+## Setup
+
+### Local Python
+
+```bash
+UV_CACHE_DIR=/tmp/uv-cache uv sync
+.venv/bin/pip install -e .
+```
+
+### Run the server
+
+```bash
+.venv/bin/python -m uvicorn server.app:app --host 0.0.0.0 --port 8000
+```
+
+### Docker
+
+```bash
+docker build -t support-ops-openenv .
+docker run --rm -p 8000:8000 support-ops-openenv
+```
+
+## Baseline inference
+
+The required root `inference.py` uses the OpenAI client for model calls and emits the mandatory `[START]`, `[STEP]`, and `[END]` logs.
+
+Environment variables:
+
+- `HF_TOKEN` or `OPENAI_API_KEY`
+- `API_BASE_URL`
+- `MODEL_NAME`
+- `LOCAL_IMAGE_NAME` if you want to run via `from_docker_image()`
+- `ENV_BASE_URL` if you want to connect to a running server
+
+Example:
+
+```bash
+export HF_TOKEN=...
+export API_BASE_URL=https://router.huggingface.co/v1
+export MODEL_NAME=meta-llama/Llama-3.1-8B-Instruct
+python inference.py
+```
+
+The script evaluates all three tasks in a fixed order for reproducible scoring. If no API key is available, it falls back to a deterministic scripted policy so the benchmark remains runnable offline.
+
+## Expected baseline behavior
+
+The bundled fallback policy should solve all three tasks with high scores because it follows the intended evidence path exactly. Frontier LLMs should also perform well on the easy and medium tasks and show larger variance on the hard fraud-escalation task if they over-index on issuing refunds instead of following policy.
+
+## Project structure
+
+```text
+.
+├── Dockerfile
+├── README.md
+├── inference.py
+├── openenv.yaml
+├── pyproject.toml
+├── server/
+│   └── app.py
+└── tool_use_env/
+    ├── client.py
+    ├── grader.py
+    ├── models.py
+    ├── tasks.py
+    └── server/
+        ├── app.py
+        └── tool_use_env_environment.py
+```

inference.py

@@ -0,0 +1,260 @@
+import asyncio
+import json
+import os
+import textwrap
+from typing import Any, List, Optional
+
+from openai import OpenAI
+
+from tool_use_env.client import ToolUseEnv
+from tool_use_env.models import ToolUseAction
+from tool_use_env.tasks import TASK_SEQUENCE
+
+LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME")
+API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
+MODEL_NAME = os.getenv("MODEL_NAME", "meta-llama/Llama-3.1-8B-Instruct")
+API_KEY = os.getenv("HF_TOKEN") or os.getenv("OPENAI_API_KEY") or os.getenv("API_KEY")
+ENV_BASE_URL = os.getenv("ENV_BASE_URL", "http://127.0.0.1:8000")
+BENCHMARK = os.getenv("MY_ENV_V4_BENCHMARK", "support_ops_env")
+MAX_STEPS = 6
+TEMPERATURE = 0.0
+MAX_TOKENS = 220
+
+SYSTEM_PROMPT = textwrap.dedent(
+    """
+    You are operating a customer-support workflow environment.
+    Your job is to gather the minimum necessary evidence, draft a short customer reply,
+    and submit the correct final resolution code.
+
+    Reply with JSON only using this schema:
+    {
+      "action_type": "review_ticket|inspect_artifact|search_policy|draft_reply|submit_resolution",
+      "artifact_id": "optional string",
+      "query": "optional string",
+      "message": "optional string",
+      "resolution_code": "optional string"
+    }
+
+    Use concise messages. Prefer exact artifact ids and exact resolution codes shown in the observation.
+    """
+).strip()
+
+
+def log_start(task: str, env: str, model: str) -> None:
+    print(f"[START] task={task} env={env} model={model}", flush=True)
+
+
+def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None:
+    error_val = error if error else "null"
+    print(
+        f"[STEP] step={step} action={action} reward={reward:.2f} "
+        f"done={str(done).lower()} error={error_val}",
+        flush=True,
+    )
+
+
+def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:
+    rewards_str = ",".join(f"{reward:.2f}" for reward in rewards)
+    print(
+        f"[END] success={str(success).lower()} steps={steps} score={score:.3f} rewards={rewards_str}",
+        flush=True,
+    )
+
+
+def _serialize_action(action: ToolUseAction) -> str:
+    payload = {"action_type": action.action_type}
+    if action.artifact_id:
+        payload["artifact_id"] = action.artifact_id
+    if action.query:
+        payload["query"] = action.query
+    if action.message:
+        payload["message"] = action.message.replace("\n", " ").strip()
+    if action.resolution_code:
+        payload["resolution_code"] = action.resolution_code
+    return json.dumps(payload, ensure_ascii=True, separators=(",", ":"))
+
+
+def _fallback_action(observation: Any) -> ToolUseAction:
+    evidence = set(observation.collected_evidence)
+    task_id = observation.task_id
+
+    if "ticket" not in evidence:
+        return ToolUseAction(action_type="review_ticket")
+
+    task_plans = {
+        "damaged-mug-replacement": [
+            ToolUseAction(action_type="inspect_artifact", artifact_id="order"),
+            ToolUseAction(action_type="search_policy", query="damaged_items"),
+            ToolUseAction(
+                action_type="draft_reply",
+                message=(
+                    "We are sending a replacement within 48 hours. "
+                    "There is no need to return the broken mug."
+                ),
+            ),
+            ToolUseAction(action_type="submit_resolution", resolution_code="send_replacement"),
+        ],
+        "duplicate-charge-refund": [
+            ToolUseAction(action_type="inspect_artifact", artifact_id="order"),
+            ToolUseAction(action_type="inspect_artifact", artifact_id="payment"),
+            ToolUseAction(action_type="search_policy", query="duplicate_charge"),
+            ToolUseAction(
+                action_type="draft_reply",
+                message=(
+                    "We confirmed the duplicate charge and issued a refund. "
+                    "You should see the refund in 3-5 business days."
+                ),
+            ),
+            ToolUseAction(
+                action_type="submit_resolution",
+                resolution_code="refund_duplicate_charge",
+            ),
+        ],
+        "account-takeover-fraud": [
+            ToolUseAction(action_type="inspect_artifact", artifact_id="account"),
+            ToolUseAction(action_type="inspect_artifact", artifact_id="risk_log"),
+            ToolUseAction(action_type="search_policy", query="account_takeover"),
+            ToolUseAction(
+                action_type="draft_reply",
+                message=(
+                    "We locked your account immediately and escalated this to our fraud team. "
+                    "You will receive an update within 24 hours."
+                ),
+            ),
+            ToolUseAction(
+                action_type="submit_resolution",
+                resolution_code="lock_account_and_escalate_fraud",
+            ),
+        ],
+    }
+
+    plan = task_plans[task_id]
+    for candidate in plan:
+        if candidate.action_type == "inspect_artifact":
+            if f"artifact:{candidate.artifact_id}" not in evidence:
+                return candidate
+        elif candidate.action_type == "search_policy":
+            if f"policy:{candidate.query}" not in evidence:
+                return candidate
+        elif candidate.action_type == "draft_reply" and not observation.last_tool_result.startswith("Draft saved"):
+            return candidate
+        elif candidate.action_type == "submit_resolution":
+            return candidate
+
+    return ToolUseAction(action_type="submit_resolution", resolution_code=observation.available_resolution_codes[0])
+
+
+def _prompt_for_observation(step: int, observation: Any) -> str:
+    return textwrap.dedent(
+        f"""
+        Step: {step}
+        Task ID: {observation.task_id}
+        Difficulty: {observation.difficulty}
+        Objective: {observation.objective}
+        Customer message: {observation.customer_message}
+        Workspace summary: {observation.workspace_summary}
+        Collected evidence: {observation.collected_evidence}
+        Available resolution codes: {observation.available_resolution_codes}
+        Last tool result: {observation.last_tool_result}
+        Last action error: {observation.last_action_error}
+        Remaining steps: {observation.remaining_steps}
+
+        Return the single best next action as JSON.
+        """
+    ).strip()
+
+
+def _model_action(client: OpenAI, step: int, observation: Any) -> ToolUseAction:
+    fallback = _fallback_action(observation)
+
+    if not API_KEY:
+        return fallback
+
+    try:
+        completion = client.chat.completions.create(
+            model=MODEL_NAME,
+            messages=[
+                {"role": "system", "content": SYSTEM_PROMPT},
+                {"role": "user", "content": _prompt_for_observation(step, observation)},
+            ],
+            temperature=TEMPERATURE,
+            max_tokens=MAX_TOKENS,
+            response_format={"type": "json_object"},
+        )
+        raw = (completion.choices[0].message.content or "").strip()
+        data = json.loads(raw)
+        return ToolUseAction(
+            action_type=data.get("action_type", fallback.action_type),
+            artifact_id=data.get("artifact_id"),
+            query=data.get("query"),
+            message=data.get("message"),
+            resolution_code=data.get("resolution_code"),
+        )
+    except Exception:
+        return fallback
+
+
+async def _connect_env() -> ToolUseEnv:
+    if LOCAL_IMAGE_NAME:
+        return await ToolUseEnv.from_docker_image(LOCAL_IMAGE_NAME)
+
+    env = ToolUseEnv(base_url=ENV_BASE_URL)
+    await env.connect()
+    return env
+
+
+async def run_task(client: OpenAI, env: ToolUseEnv, task_id: str) -> float:
+    rewards: List[float] = []
+    steps_taken = 0
+    score = 0.0
+    success = False
+
+    log_start(task=task_id, env=BENCHMARK, model=MODEL_NAME)
+
+    try:
+        result = await env.reset(task_id=task_id, seed=7)
+        observation = result.observation
+
+        for step in range(1, MAX_STEPS + 1):
+            if result.done:
+                break
+
+            action = _model_action(client, step, observation)
+            action_str = _serialize_action(action)
+            result = await env.step(action)
+            observation = result.observation
+
+            reward = float(result.reward or 0.0)
+            done = bool(result.done)
+            error = observation.last_action_error
+            rewards.append(reward)
+            steps_taken = step
+
+            log_step(step=step, action=action_str, reward=reward, done=done, error=error)
+
+            if done:
+                break
+
+        state = await env.state()
+        score = float(state.final_score)
+        success = score >= 0.8
+    finally:
+        log_end(success=success, steps=steps_taken, score=score, rewards=rewards)
+
+    return score
+
+
+async def main() -> None:
+    client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY or "missing")
+    env = await _connect_env()
+    try:
+        scores = []
+        for task_id in TASK_SEQUENCE:
+            score = await run_task(client, env, task_id)
+            scores.append(score)
+    finally:
+        await env.close()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

openenv.yaml

@@ -0,0 +1,84 @@
+name: support_ops_env
+description: Customer support operations environment for multi-step tool-using agents
+version: 1.0.0
+
+entrypoint: server.app:app
+
+actions:
+  type: object
+  properties:
+    action_type:
+      type: string
+      enum:
+        - review_ticket
+        - inspect_artifact
+        - search_policy
+        - draft_reply
+        - submit_resolution
+    artifact_id:
+      type: string
+      nullable: true
+    query:
+      type: string
+      nullable: true
+    message:
+      type: string
+      nullable: true
+    resolution_code:
+      type: string
+      nullable: true
+  required:
+    - action_type
+
+observations:
+  type: object
+  properties:
+    task_id:
+      type: string
+    difficulty:
+      type: string
+      enum: [easy, medium, hard]
+    objective:
+      type: string
+    customer_message:
+      type: string
+    workspace_summary:
+      type: string
+    available_actions:
+      type: array
+      items:
+        type: string
+    available_resolution_codes:
+      type: array
+      items:
+        type: string
+    collected_evidence:
+      type: array
+      items:
+        type: string
+    last_tool_result:
+      type: string
+      nullable: true
+    last_action_error:
+      type: string
+      nullable: true
+    remaining_steps:
+      type: integer
+    current_score:
+      type: number
+
+reward_range: [0.0, 1.0]
+
+metadata:
+  benchmark: support_ops_env
+  domain: customer_support
+  difficulty_levels:
+    - easy
+    - medium
+    - hard
+  features:
+    - multi_step_reasoning
+    - tool_selection
+    - policy_lookup
+    - customer_support_triage
+    - shaped_rewards

openenv_support_ops_env.egg-info/PKG-INFO

@@ -0,0 +1,11 @@
+Metadata-Version: 2.4
+Name: openenv-support-ops-env
+Version: 0.1.0
+Summary: OpenEnv customer support operations environment
+Requires-Python: >=3.10
+Requires-Dist: openenv-core[core]>=0.2.1
+Requires-Dist: openai>=1.40.0
+Requires-Dist: python-dotenv>=1.0.1
+Requires-Dist: uvicorn>=0.30.0
+Provides-Extra: dev
+Requires-Dist: pytest>=8.0.0; extra == "dev"

openenv_support_ops_env.egg-info/SOURCES.txt

@@ -0,0 +1,18 @@
+README.md
+pyproject.toml
+openenv_support_ops_env.egg-info/PKG-INFO
+openenv_support_ops_env.egg-info/SOURCES.txt
+openenv_support_ops_env.egg-info/dependency_links.txt
+openenv_support_ops_env.egg-info/entry_points.txt
+openenv_support_ops_env.egg-info/requires.txt
+openenv_support_ops_env.egg-info/top_level.txt
+server/__init__.py
+server/app.py
+tool_use_env/__init__.py
+tool_use_env/client.py
+tool_use_env/grader.py
+tool_use_env/models.py
+tool_use_env/tasks.py
+tool_use_env/server/__init__.py
+tool_use_env/server/app.py
+tool_use_env/server/tool_use_env_environment.py

openenv_support_ops_env.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

openenv_support_ops_env.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+server = server.app:main

openenv_support_ops_env.egg-info/requires.txt

@@ -0,0 +1,7 @@
+openenv-core[core]>=0.2.1
+openai>=1.40.0
+python-dotenv>=1.0.1
+uvicorn>=0.30.0
+
+[dev]
+pytest>=8.0.0

openenv_support_ops_env.egg-info/top_level.txt

@@ -0,0 +1,2 @@
+server
+tool_use_env

pyproject.toml

@@ -0,0 +1,28 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "openenv-support-ops-env"
+version = "0.1.0"
+description = "OpenEnv customer support operations environment"
+requires-python = ">=3.10"
+dependencies = [
+    "openenv-core[core]>=0.2.1",
+    "openai>=1.40.0",
+    "python-dotenv>=1.0.1",
+    "uvicorn>=0.30.0",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=8.0.0",
+]
+
+[project.scripts]
+server = "server.app:main"
+
+[tool.setuptools]
+include-package-data = true
+packages = ["tool_use_env", "tool_use_env.server", "server"]
+package-dir = { "tool_use_env" = "tool_use_env", "tool_use_env.server" = "tool_use_env/server", "server" = "server" }

server/__init__.py

@@ -0,0 +1 @@
+"""Root server shim for OpenEnv validation and uv run server."""

server/app.py

@@ -0,0 +1,10 @@
+from tool_use_env.server.app import app as app
+from tool_use_env.server.app import main as _package_main
+
+
+def main(host: str = "0.0.0.0", port: int = 8000):
+    _package_main(host=host, port=port)
+
+
+if __name__ == "__main__":
+    main()

tool_use_env/README.md

@@ -0,0 +1,256 @@
+---
+title: Tool Use Env Environment Server
+emoji: 📀
+colorFrom: purple
+colorTo: gray
+sdk: docker
+pinned: false
+app_port: 8000
+base_path: /web
+tags:
+  - openenv
+---
+
+# Tool Use Env Environment
+
+A simple test environment that echoes back messages. Perfect for testing the env APIs as well as demonstrating environment usage patterns.
+
+## Quick Start
+hi
+
+The simplest way to use the Tool Use Env environment is through the `ToolUseEnv` class:
+
+```python
+from tool_use_env import ToolUseAction, ToolUseEnv
+
+try:
+    # Create environment from Docker image
+    tool_use_envenv = ToolUseEnv.from_docker_image("tool_use_env-env:latest")
+
+    # Reset
+    result = tool_use_envenv.reset()
+    print(f"Reset: {result.observation.echoed_message}")
+
+    # Send multiple messages
+    messages = ["Hello, World!", "Testing echo", "Final message"]
+
+    for msg in messages:
+        result = tool_use_envenv.step(ToolUseAction(message=msg))
+        print(f"Sent: '{msg}'")
+        print(f"  → Echoed: '{result.observation.echoed_message}'")
+        print(f"  → Length: {result.observation.message_length}")
+        print(f"  → Reward: {result.reward}")
+
+finally:
+    # Always clean up
+    tool_use_envenv.close()
+```
+
+That's it! The `ToolUseEnv.from_docker_image()` method handles:
+- Starting the Docker container
+- Waiting for the server to be ready
+- Connecting to the environment
+- Container cleanup when you call `close()`
+
+## Building the Docker Image
+
+Before using the environment, you need to build the Docker image:
+
+```bash
+# From project root
+docker build -t tool_use_env-env:latest -f server/Dockerfile .
+```
+
+## Deploying to Hugging Face Spaces
+
+You can easily deploy your OpenEnv environment to Hugging Face Spaces using the `openenv push` command:
+
+```bash
+# From the environment directory (where openenv.yaml is located)
+openenv push
+
+# Or specify options
+openenv push --namespace my-org --private
+```
+
+The `openenv push` command will:
+1. Validate that the directory is an OpenEnv environment (checks for `openenv.yaml`)
+2. Prepare a custom build for Hugging Face Docker space (enables web interface)
+3. Upload to Hugging Face (ensuring you're logged in)
+
+### Prerequisites
+
+- Authenticate with Hugging Face: The command will prompt for login if not already authenticated
+
+### Options
+
+- `--directory`, `-d`: Directory containing the OpenEnv environment (defaults to current directory)
+- `--repo-id`, `-r`: Repository ID in format 'username/repo-name' (defaults to 'username/env-name' from openenv.yaml)
+- `--base-image`, `-b`: Base Docker image to use (overrides Dockerfile FROM)
+- `--private`: Deploy the space as private (default: public)
+
+### Examples
+
+```bash
+# Push to your personal namespace (defaults to username/env-name from openenv.yaml)
+openenv push
+
+# Push to a specific repository
+openenv push --repo-id my-org/my-env
+
+# Push with a custom base image
+openenv push --base-image ghcr.io/meta-pytorch/openenv-base:latest
+
+# Push as a private space
+openenv push --private
+
+# Combine options
+openenv push --repo-id my-org/my-env --base-image custom-base:latest --private
+```
+
+After deployment, your space will be available at:
+`https://huggingface.co/spaces/<repo-id>`
+
+The deployed space includes:
+- **Web Interface** at `/web` - Interactive UI for exploring the environment
+- **API Documentation** at `/docs` - Full OpenAPI/Swagger interface
+- **Health Check** at `/health` - Container health monitoring
+- **WebSocket** at `/ws` - Persistent session endpoint for low-latency interactions
+
+## Environment Details
+
+### Action
+**ToolUseAction**: Contains a single field
+- `message` (str) - The message to echo back
+
+### Observation
+**ToolUseObservation**: Contains the echo response and metadata
+- `echoed_message` (str) - The message echoed back
+- `message_length` (int) - Length of the message
+- `reward` (float) - Reward based on message length (length × 0.1)
+- `done` (bool) - Always False for echo environment
+- `metadata` (dict) - Additional info like step count
+
+### Reward
+The reward is calculated as: `message_length × 0.1`
+- "Hi" → reward: 0.2
+- "Hello, World!" → reward: 1.3
+- Empty message → reward: 0.0
+
+## Advanced Usage
+
+### Connecting to an Existing Server
+
+If you already have a Tool Use Env environment server running, you can connect directly:
+
+```python
+from tool_use_env import ToolUseEnv
+
+# Connect to existing server
+tool_use_envenv = ToolUseEnv(base_url="<ENV_HTTP_URL_HERE>")
+
+# Use as normal
+result = tool_use_envenv.reset()
+result = tool_use_envenv.step(ToolUseAction(message="Hello!"))
+```
+
+Note: When connecting to an existing server, `tool_use_envenv.close()` will NOT stop the server.
+
+### Using the Context Manager
+
+The client supports context manager usage for automatic connection management:
+
+```python
+from tool_use_env import ToolUseAction, ToolUseEnv
+
+# Connect with context manager (auto-connects and closes)
+with ToolUseEnv(base_url="http://localhost:8000") as env:
+    result = env.reset()
+    print(f"Reset: {result.observation.echoed_message}")
+    # Multiple steps with low latency
+    for msg in ["Hello", "World", "!"]:
+        result = env.step(ToolUseAction(message=msg))
+        print(f"Echoed: {result.observation.echoed_message}")
+```
+
+The client uses WebSocket connections for:
+- **Lower latency**: No HTTP connection overhead per request
+- **Persistent session**: Server maintains your environment state
+- **Efficient for episodes**: Better for many sequential steps
+
+### Concurrent WebSocket Sessions
+
+The server supports multiple concurrent WebSocket connections. To enable this,
+modify `server/app.py` to use factory mode:
+
+```python
+# In server/app.py - use factory mode for concurrent sessions
+app = create_app(
+    ToolUseEnvironment,  # Pass class, not instance
+    ToolUseAction,
+    ToolUseObservation,
+    max_concurrent_envs=4,  # Allow 4 concurrent sessions
+)
+```
+
+Then multiple clients can connect simultaneously:
+
+```python
+from tool_use_env import ToolUseAction, ToolUseEnv
+from concurrent.futures import ThreadPoolExecutor
+
+def run_episode(client_id: int):
+    with ToolUseEnv(base_url="http://localhost:8000") as env:
+        result = env.reset()
+        for i in range(10):
+            result = env.step(ToolUseAction(message=f"Client {client_id}, step {i}"))
+        return client_id, result.observation.message_length
+
+# Run 4 episodes concurrently
+with ThreadPoolExecutor(max_workers=4) as executor:
+    results = list(executor.map(run_episode, range(4)))
+```
+
+## Development & Testing
+
+### Direct Environment Testing
+
+Test the environment logic directly without starting the HTTP server:
+
+```bash
+# From the server directory
+python3 server/tool_use_env_environment.py
+```
+
+This verifies that:
+- Environment resets correctly
+- Step executes actions properly
+- State tracking works
+- Rewards are calculated correctly
+
+### Running Locally
+
+Run the server locally for development:
+
+```bash
+uvicorn server.app:app --reload
+```
+
+## Project Structure
+
+```
+tool_use_env/
+├── .dockerignore         # Docker build exclusions
+├── __init__.py            # Module exports
+├── README.md              # This file
+├── openenv.yaml           # OpenEnv manifest
+├── pyproject.toml         # Project metadata and dependencies
+├── uv.lock                # Locked dependencies (generated)
+├── client.py              # ToolUseEnv client
+├── models.py              # Action and Observation models
+└── server/
+    ├── __init__.py        # Server module exports
+    ├── tool_use_env_environment.py  # Core environment logic
+    ├── app.py             # FastAPI application (HTTP + WebSocket endpoints)
+    └── Dockerfile         # Container image definition
+```

tool_use_env/__init__.py

@@ -0,0 +1,17 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Tool Use Env Environment."""
+
+from .client import ToolUseEnv
+from .models import ToolUseAction, ToolUseObservation, ToolUseState
+
+__all__ = [
+    "ToolUseAction",
+    "ToolUseObservation",
+    "ToolUseState",
+    "ToolUseEnv",
+]

tool_use_env/agents/baseline.py

@@ -0,0 +1,267 @@
+# from tool_use_env.client import ToolUseEnv
+# from tool_use_env.models import ToolUseAction
+# import random
+
+# def rule_based_policy(query: str):
+#     query = query.lower()
+
+#     # --- Introduce slight imperfection ---
+#     if random.random() < 0.1:
+#         return "answer_directly"
+
+#     if "what is" in query and any(op in query for op in ["+", "-", "*", "/"]):
+#         return "use_calculator"
+
+#     if "capital" in query or "who is" in query:
+#         return "use_search"
+
+#     return "answer_directly"
+
+
+# def run_single_episode(env):
+#     result = env.reset()
+#     obs = result.observation
+
+#     query = obs.query
+#     action_type = rule_based_policy(query)
+
+#     action = ToolUseAction(action_type=action_type)
+
+#     result = env.step(action)
+#     obs = result.observation
+
+#     return {
+#         "query": query,
+#         "action": action_type,
+#         "reward": result.reward,
+#         "message": obs.message
+#     }
+
+# def run_evaluation(num_episodes=20):
+#     results = []
+
+#     difficulty_scores = {
+#         "easy": [],
+#         "medium": [],
+#         "hard": []
+#     }
+
+#     total_score = 0
+
+#     with ToolUseEnv(base_url="http://localhost:8000").sync() as env:
+#         for _ in range(num_episodes):
+#             result = env.reset()
+#             obs = result.observation
+#             query = obs.query
+#             state = env.state()
+#             difficulty = state.difficulty
+
+#             action_type = rule_based_policy(query)
+#             action = ToolUseAction(action_type=action_type)
+
+#             result = env.step(action)
+
+#             score = result.reward
+#             total_score += score
+
+#             difficulty_scores[difficulty].append(score)
+
+#             results.append({
+#                 "query": query,
+#                 "difficulty": difficulty,
+#                 "action": action_type,
+#                 "score": score,
+#                 "message": result.observation.message
+#             })
+
+#     avg_score = total_score / num_episodes
+
+#     print("\n=== OVERALL PERFORMANCE ===")
+#     print(f"Average Score: {avg_score:.2f}")
+
+#     print("\n=== DIFFICULTY BREAKDOWN ===")
+#     for level in difficulty_scores:
+#         if difficulty_scores[level]:
+#             avg = sum(difficulty_scores[level]) / len(difficulty_scores[level])
+#             print(f"{level.capitalize()}: {avg:.2f}")
+    
+#     print("\n=== SAMPLE CASES ===")
+#     for r in results[:5]:
+#         print(f"\nQuery: {r['query']}")
+#         print(f"Action: {r['action']}")
+#         print(f"Score: {r['score']:.2f}")
+#         print(f"Details: {r['message']}")
+
+#     return results
+
+# def analyze_failures(results):
+#     wrong_decisions = 0
+#     tool_failures = 0
+#     total = len(results)
+
+#     for r in results:
+#         msg = r["message"]
+
+#         if "Correct: False" in msg:
+#             if "use_" in msg:
+#                 tool_failures += 1
+#             else:
+#                 wrong_decisions += 1
+
+#     print("\n=== FAILURE ANALYSIS ===")
+#     print(f"Tool failures: {tool_failures}/{total} ({(tool_failures/total)*100:.1f}%)")
+#     print(f"Wrong decisions: {wrong_decisions}/{total} ({(wrong_decisions/total)*100:.1f}%)")
+
+
+# if __name__ == "__main__":
+#     results = run_evaluation(50)
+#     analyze_failures(results)
+
+import os
+import random
+from collections import defaultdict
+
+from dotenv import load_dotenv
+from openai import OpenAI
+
+from tool_use_env.client import ToolUseEnv
+from tool_use_env.models import ToolUseAction
+
+# --- Load environment variables ---
+load_dotenv()
+
+# --- Initialize OpenAI client ---
+client = OpenAI()
+
+# --- Reproducibility ---
+random.seed(42)
+
+
+# 🧠 LLM Policy (CORE)
+def llm_policy(query: str):
+    prompt = f"""
+You are an AI agent choosing the best tool.
+
+Available actions:
+- use_calculator (for math problems)
+- use_search (for factual questions)
+- answer_directly (if neither tool is needed)
+
+Query: {query}
+
+Respond with ONLY one of:
+use_calculator
+use_search
+answer_directly
+"""
+
+    try:
+        response = client.chat.completions.create(
+            model="gpt-4o-mini",
+            messages=[{"role": "user", "content": prompt}],
+            temperature=0
+        )
+
+        action = response.choices[0].message.content.strip()
+
+        # --- Safety check ---
+        if action not in ["use_calculator", "use_search", "answer_directly"]:
+            return "answer_directly"
+
+        return action
+
+    except Exception as e:
+        print(f"[ERROR] LLM call failed: {e}")
+        return "answer_directly"
+
+
+# 🧪 Evaluation Loop
+def run_evaluation(num_episodes=50):
+    results = []
+    total_score = 0
+
+    difficulty_scores = defaultdict(list)
+
+    with ToolUseEnv(base_url="http://localhost:8000").sync() as env:
+        for _ in range(num_episodes):
+            # --- Reset ---
+            result = env.reset()
+            obs = result.observation
+
+            query = obs.query
+
+            # --- Get difficulty ---
+            state = env.state()
+            difficulty = state.difficulty
+
+            # --- LLM decides action ---
+            action_type = llm_policy(query)
+            action = ToolUseAction(action_type=action_type)
+
+            # --- Step ---
+            result = env.step(action)
+            obs = result.observation
+
+            score = result.reward
+            total_score += score
+
+            difficulty_scores[difficulty].append(score)
+
+            results.append({
+                "query": query,
+                "difficulty": difficulty,
+                "action": action_type,
+                "score": score,
+                "message": obs.message
+            })
+
+            print(f"Score: {score:.2f}")
+
+    # --- Overall ---
+    avg_score = total_score / num_episodes
+
+    print("\n=== OVERALL PERFORMANCE ===")
+    print(f"Average Score: {avg_score:.2f}")
+
+    # --- Breakdown ---
+    print("\n=== DIFFICULTY BREAKDOWN ===")
+    for level in ["easy", "medium", "hard"]:
+        if difficulty_scores[level]:
+            avg = sum(difficulty_scores[level]) / len(difficulty_scores[level])
+            print(f"{level.capitalize()}: {avg:.2f}")
+
+    # --- Sample Cases ---
+    print("\n=== SAMPLE CASES ===")
+    for r in results[:5]:
+        print(f"\nQuery: {r['query']}")
+        print(f"Action: {r['action']}")
+        print(f"Score: {r['score']:.2f}")
+        print(f"Details: {r['message']}")
+
+    return results
+
+
+# 📊 Failure Analysis
+def analyze_failures(results):
+    total = len(results)
+    tool_failures = 0
+    wrong_decisions = 0
+
+    for r in results:
+        msg = r["message"]
+
+        if "Correct: False" in msg:
+            if "use_" in msg:
+                tool_failures += 1
+            else:
+                wrong_decisions += 1
+
+    print("\n=== FAILURE ANALYSIS ===")
+    print(f"Tool failures: {tool_failures}/{total} ({(tool_failures/total)*100:.1f}%)")
+    print(f"Wrong decisions: {wrong_decisions}/{total} ({(wrong_decisions/total)*100:.1f}%)")
+
+
+# 🚀 Main
+if __name__ == "__main__":
+    results = run_evaluation(50)
+    analyze_failures(results)

tool_use_env/client.py

@@ -0,0 +1,165 @@
+# # Copyright (c) Meta Platforms, Inc. and affiliates.
+# # All rights reserved.
+# #
+# # This source code is licensed under the BSD-style license found in the
+# # LICENSE file in the root directory of this source tree.
+
+# """Tool Use Env Environment Client."""
+
+# from typing import Dict
+
+# from openenv.core import EnvClient
+# from openenv.core.client_types import StepResult
+# from openenv.core.env_server.types import State
+
+# from .models import ToolUseAction, ToolUseObservation
+
+
+# class ToolUseEnv(
+#     EnvClient[ToolUseAction, ToolUseObservation, State]
+# ):
+#     """
+#     Client for the Tool Use Env Environment.
+
+#     This client maintains a persistent WebSocket connection to the environment server,
+#     enabling efficient multi-step interactions with lower latency.
+#     Each client instance has its own dedicated environment session on the server.
+
+#     Example:
+#         >>> # Connect to a running server
+#         >>> with ToolUseEnv(base_url="http://localhost:8000") as client:
+#         ...     result = client.reset()
+#         ...     print(result.observation.echoed_message)
+#         ...
+#         ...     result = client.step(ToolUseAction(message="Hello!"))
+#         ...     print(result.observation.echoed_message)
+
+#     Example with Docker:
+#         >>> # Automatically start container and connect
+#         >>> client = ToolUseEnv.from_docker_image("tool_use_env-env:latest")
+#         >>> try:
+#         ...     result = client.reset()
+#         ...     result = client.step(ToolUseAction(message="Test"))
+#         ... finally:
+#         ...     client.close()
+#     """
+
+#     def _step_payload(self, action: ToolUseAction) -> Dict:
+#         """
+#         Convert ToolUseAction to JSON payload for step message.
+
+#         Args:
+#             action: ToolUseAction instance
+
+#         Returns:
+#             Dictionary representation suitable for JSON encoding
+#         """
+#         return {
+#             "message": action.message,
+#         }
+
+#     def _parse_result(self, payload: Dict) -> StepResult[ToolUseObservation]:
+#         """
+#         Parse server response into StepResult[ToolUseObservation].
+
+#         Args:
+#             payload: JSON response data from server
+
+#         Returns:
+#             StepResult with ToolUseObservation
+#         """
+#         obs_data = payload.get("observation", {})
+#         observation = ToolUseObservation(
+#             echoed_message=obs_data.get("echoed_message", ""),
+#             message_length=obs_data.get("message_length", 0),
+#             done=payload.get("done", False),
+#             reward=payload.get("reward"),
+#             metadata=obs_data.get("metadata", {}),
+#         )
+
+#         return StepResult(
+#             observation=observation,
+#             reward=payload.get("reward"),
+#             done=payload.get("done", False),
+#         )
+
+#     def _parse_state(self, payload: Dict) -> State:
+#         """
+#         Parse server response into State object.
+
+#         Args:
+#             payload: JSON response from state request
+
+#         Returns:
+#             State object with episode_id and step_count
+#         """
+#         return State(
+#             episode_id=payload.get("episode_id"),
+#             step_count=payload.get("step_count", 0),
+#         )
+
+from openenv.core.env_client import EnvClient
+from openenv.core.client_types import StepResult
+
+from tool_use_env.models import ToolUseAction, ToolUseObservation, ToolUseState
+
+
+class ToolUseEnv(EnvClient[ToolUseAction, ToolUseObservation, ToolUseState]):
+
+    def _step_payload(self, action: ToolUseAction) -> dict:
+        return {
+            "action_type": action.action_type,
+            "artifact_id": action.artifact_id,
+            "query": action.query,
+            "message": action.message,
+            "resolution_code": action.resolution_code,
+        }
+
+    def _parse_result(self, payload: dict) -> StepResult:
+        obs_data = payload.get("observation", {})
+
+        observation = ToolUseObservation(
+            done=payload.get("done", False),
+            reward=payload.get("reward"),
+            task_id=obs_data.get("task_id", ""),
+            difficulty=obs_data.get("difficulty", "easy"),
+            objective=obs_data.get("objective", ""),
+            customer_message=obs_data.get("customer_message", ""),
+            workspace_summary=obs_data.get("workspace_summary", ""),
+            available_actions=obs_data.get("available_actions", []),
+            available_resolution_codes=obs_data.get("available_resolution_codes", []),
+            collected_evidence=obs_data.get("collected_evidence", []),
+            last_tool_result=obs_data.get("last_tool_result"),
+            last_action_error=obs_data.get("last_action_error"),
+            remaining_steps=obs_data.get("remaining_steps", 0),
+            current_score=obs_data.get("current_score", 0.0),
+            metadata=obs_data.get("metadata", {}),
+        )
+
+        return StepResult(
+            observation=observation,
+            reward=payload.get("reward"),
+            done=payload.get("done", False),
+        )
+
+    def _parse_state(self, payload: dict) -> ToolUseState:
+        return ToolUseState(
+            episode_id=payload.get("episode_id"),
+            step_count=payload.get("step_count", 0),
+            task_id=payload.get("task_id", ""),
+            task_name=payload.get("task_name", ""),
+            difficulty=payload.get("difficulty", ""),
+            objective=payload.get("objective", ""),
+            cumulative_reward=payload.get("cumulative_reward", 0.0),
+            final_score=payload.get("final_score", 0.0),
+            drafted_reply=payload.get("drafted_reply"),
+            resolution_code=payload.get("resolution_code"),
+            expected_resolution_code=payload.get("expected_resolution_code", ""),
+            required_evidence=payload.get("required_evidence", []),
+            collected_evidence=payload.get("collected_evidence", []),
+            action_history=payload.get("action_history", []),
+            repeat_action_count=payload.get("repeat_action_count", 0),
+            last_action_error=payload.get("last_action_error"),
+            known_artifacts=payload.get("known_artifacts", {}),
+            known_policies=payload.get("known_policies", {}),
+        )

tool_use_env/grader.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+from typing import Any
+
+
+def _keyword_score(reply: str | None, keywords: list[str]) -> float:
+    if not reply or not keywords:
+        return 0.0
+
+    lowered = reply.lower()
+    hits = sum(1 for keyword in keywords if keyword.lower() in lowered)
+    return hits / len(keywords)
+
+
+def grade_task(
+    task: dict[str, Any],
+    collected_evidence: list[str],
+    drafted_reply: str | None,
+    resolution_code: str | None,
+    step_count: int,
+    repeat_action_count: int,
+) -> dict[str, float]:
+    required_evidence = task["required_evidence"]
+    evidence_hits = sum(1 for key in required_evidence if key in collected_evidence)
+    evidence_score = evidence_hits / len(required_evidence)
+
+    resolution_score = 1.0 if resolution_code == task["expected_resolution_code"] else 0.0
+    reply_score = _keyword_score(drafted_reply, task["reply_keywords"])
+
+    optimal_steps = task.get("optimal_steps", len(required_evidence) + 2)
+    extra_steps = max(0, step_count - optimal_steps)
+    efficiency_penalty = min(0.25, (extra_steps * 0.05) + (repeat_action_count * 0.04))
+    efficiency_score = max(0.0, 1.0 - efficiency_penalty)
+
+    final_score = (
+        0.5 * resolution_score
+        + 0.25 * evidence_score
+        + 0.2 * reply_score
+        + 0.05 * efficiency_score
+    )
+
+    return {
+        "resolution_score": round(resolution_score, 3),
+        "evidence_score": round(evidence_score, 3),
+        "reply_score": round(reply_score, 3),
+        "efficiency_score": round(efficiency_score, 3),
+        "final_score": round(min(max(final_score, 0.0), 1.0), 3),
+    }

tool_use_env/models.py

@@ -0,0 +1,86 @@
+from typing import Dict, List, Literal, Optional
+
+from openenv.core.env_server import Action, Observation, State
+from pydantic import Field
+
+
+class ToolUseAction(Action):
+    action_type: Literal[
+        "review_ticket",
+        "inspect_artifact",
+        "search_policy",
+        "draft_reply",
+        "submit_resolution",
+    ] = Field(..., description="The action the agent wants to execute.")
+    artifact_id: Optional[str] = Field(
+        default=None,
+        description="Artifact identifier for inspect_artifact, such as order or risk_log.",
+    )
+    query: Optional[str] = Field(
+        default=None,
+        description="Policy name or search query for search_policy.",
+    )
+    message: Optional[str] = Field(
+        default=None,
+        description="Customer-facing reply draft used with draft_reply.",
+    )
+    resolution_code: Optional[str] = Field(
+        default=None,
+        description="Final resolution code used with submit_resolution.",
+    )
+
+
+class ToolUseObservation(Observation):
+    task_id: str = Field(..., description="Deterministic task identifier.")
+    difficulty: Literal["easy", "medium", "hard"] = Field(
+        ..., description="Difficulty tier for the active task."
+    )
+    objective: str = Field(..., description="Concrete task objective for the agent.")
+    customer_message: str = Field(..., description="The raw customer support ticket.")
+    workspace_summary: str = Field(
+        ..., description="Short summary of known evidence and remaining work."
+    )
+    available_actions: List[str] = Field(
+        default_factory=list, description="Available environment actions."
+    )
+    available_resolution_codes: List[str] = Field(
+        default_factory=list,
+        description="Resolution codes accepted by submit_resolution.",
+    )
+    collected_evidence: List[str] = Field(
+        default_factory=list,
+        description="Evidence keys collected so far, such as ticket or payment.",
+    )
+    last_tool_result: Optional[str] = Field(
+        default=None,
+        description="Most recent tool or grader output shown to the agent.",
+    )
+    last_action_error: Optional[str] = Field(
+        default=None, description="Validation error for the last action, if any."
+    )
+    remaining_steps: int = Field(
+        ..., description="How many steps are left before the episode ends."
+    )
+    current_score: float = Field(
+        default=0.0,
+        description="Current deterministic grader score in the [0, 1] range.",
+    )
+
+
+class ToolUseState(State):
+    task_id: str = ""
+    task_name: str = ""
+    difficulty: str = ""
+    objective: str = ""
+    cumulative_reward: float = 0.0
+    final_score: float = 0.0
+    drafted_reply: Optional[str] = None
+    resolution_code: Optional[str] = None
+    expected_resolution_code: str = ""
+    required_evidence: List[str] = Field(default_factory=list)
+    collected_evidence: List[str] = Field(default_factory=list)
+    action_history: List[str] = Field(default_factory=list)
+    repeat_action_count: int = 0
+    last_action_error: Optional[str] = None
+    known_artifacts: Dict[str, str] = Field(default_factory=dict)
+    known_policies: Dict[str, str] = Field(default_factory=dict)

tool_use_env/openenv_tool_use_env.egg-info/PKG-INFO

@@ -0,0 +1,9 @@
+Metadata-Version: 2.4
+Name: openenv-tool_use_env
+Version: 0.1.0
+Summary: Tool Use Env environment for OpenEnv
+Requires-Python: >=3.10
+Requires-Dist: openenv-core[core]>=0.2.1
+Provides-Extra: dev
+Requires-Dist: pytest>=8.0.0; extra == "dev"
+Requires-Dist: pytest-cov>=4.0.0; extra == "dev"

tool_use_env/openenv_tool_use_env.egg-info/SOURCES.txt

@@ -0,0 +1,20 @@
+README.md
+__init__.py
+client.py
+grader.py
+models.py
+pyproject.toml
+./__init__.py
+./client.py
+./grader.py
+./models.py
+openenv_tool_use_env.egg-info/PKG-INFO
+openenv_tool_use_env.egg-info/SOURCES.txt
+openenv_tool_use_env.egg-info/dependency_links.txt
+openenv_tool_use_env.egg-info/entry_points.txt
+openenv_tool_use_env.egg-info/requires.txt
+openenv_tool_use_env.egg-info/top_level.txt
+server/__init__.py
+server/app.py
+server/tool_use_env_environment.py
+tests/test_tools.py

tool_use_env/openenv_tool_use_env.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

tool_use_env/openenv_tool_use_env.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+server = tool_use_env.server.app:main

tool_use_env/openenv_tool_use_env.egg-info/requires.txt

@@ -0,0 +1,5 @@
+openenv-core[core]>=0.2.1
+
+[dev]
+pytest>=8.0.0
+pytest-cov>=4.0.0

tool_use_env/openenv_tool_use_env.egg-info/top_level.txt

@@ -0,0 +1 @@
+tool_use_env

tool_use_env/pyproject.toml

@@ -0,0 +1,45 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+[build-system]
+requires = ["setuptools>=45", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "openenv-tool_use_env"
+version = "0.1.0"
+description = "Tool Use Env environment for OpenEnv"
+requires-python = ">=3.10"
+dependencies = [
+    # Core OpenEnv runtime (provides FastAPI server + HTTP client types)
+    # install from github
+    # "openenv-core[core] @ git+https://github.com/meta-pytorch/OpenEnv.git",
+    "openenv-core[core]>=0.2.1",
+    # Environment-specific dependencies
+    # Add all dependencies needed for your environment here
+    # Examples:
+    # "numpy>=1.19.0",
+    # "torch>=2.0.0",
+    # "gymnasium>=0.29.0",
+    # "openspiel>=1.0.0",
+    # "smolagents>=1.22.0,<2",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=8.0.0",
+    "pytest-cov>=4.0.0",
+]
+
+[project.scripts]
+# Server entry point - enables running via: uv run --project . server
+# or: python -m tool_use_env.server.app
+server = "tool_use_env.server.app:main"
+
+[tool.setuptools]
+include-package-data = true
+packages = ["tool_use_env", "tool_use_env.server"]
+package-dir = { "tool_use_env" = ".", "tool_use_env.server" = "server" }

tool_use_env/server/Dockerfile

@@ -0,0 +1,80 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Multi-stage build using openenv-base
+# This Dockerfile is flexible and works for both:
+# - In-repo environments (with local OpenEnv sources)
+# - Standalone environments (with openenv from PyPI/Git)
+# The build script (openenv build) handles context detection and sets appropriate build args.
+
+ARG BASE_IMAGE=ghcr.io/meta-pytorch/openenv-base:latest
+FROM ${BASE_IMAGE} AS builder
+
+WORKDIR /app
+
+# Ensure git is available (required for installing dependencies from VCS)
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends git && \
+    rm -rf /var/lib/apt/lists/*
+
+# Build argument to control whether we're building standalone or in-repo
+ARG BUILD_MODE=in-repo
+ARG ENV_NAME=tool_use_env
+
+# Copy environment code (always at root of build context)
+COPY . /app/env
+
+# For in-repo builds, openenv is already vendored in the build context
+# For standalone builds, openenv will be installed via pyproject.toml
+WORKDIR /app/env
+
+# Ensure uv is available (for local builds where base image lacks it)
+RUN if ! command -v uv >/dev/null 2>&1; then \
+        curl -LsSf https://astral.sh/uv/install.sh | sh && \
+        mv /root/.local/bin/uv /usr/local/bin/uv && \
+        mv /root/.local/bin/uvx /usr/local/bin/uvx; \
+    fi
+    
+# Install dependencies using uv sync
+# If uv.lock exists, use it; otherwise resolve on the fly
+RUN --mount=type=cache,target=/root/.cache/uv \
+    if [ -f uv.lock ]; then \
+        uv sync --frozen --no-install-project --no-editable; \
+    else \
+        uv sync --no-install-project --no-editable; \
+    fi
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    if [ -f uv.lock ]; then \
+        uv sync --frozen --no-editable; \
+    else \
+        uv sync --no-editable; \
+    fi
+
+# Final runtime stage
+FROM ${BASE_IMAGE}
+
+WORKDIR /app
+
+# Copy the virtual environment from builder
+COPY --from=builder /app/env/.venv /app/.venv
+
+# Copy the environment code
+COPY --from=builder /app/env /app/env
+
+# Set PATH to use the virtual environment
+ENV PATH="/app/.venv/bin:$PATH"
+
+# Set PYTHONPATH so imports work correctly
+ENV PYTHONPATH="/app/env:$PYTHONPATH"
+
+# Health check
+HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --retries=3 \
+    CMD curl -f http://localhost:8000/health || exit 1
+
+# Run the FastAPI server
+# The module path is constructed to work with the /app/env structure
+CMD ["sh", "-c", "cd /app/env && uvicorn server.app:app --host 0.0.0.0 --port 8000"]

tool_use_env/server/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Tool Use Env environment server components."""
+
+from .tool_use_env_environment import ToolUseEnvironment
+
+__all__ = ["ToolUseEnvironment"]

tool_use_env/server/app.py

@@ -0,0 +1,29 @@
+from openenv.core.env_server.http_server import create_app
+
+from tool_use_env.models import ToolUseAction, ToolUseObservation
+from tool_use_env.server.tool_use_env_environment import ToolUseEnvironment
+
+
+app = create_app(
+    ToolUseEnvironment,
+    ToolUseAction,
+    ToolUseObservation,
+    env_name="support_ops_env",
+    max_concurrent_envs=4,
+)
+
+
+import uvicorn
+
+
+def main(host: str = "0.0.0.0", port: int = 8000):
+    uvicorn.run("tool_use_env.server.app:app", host=host, port=port)
+
+
+@app.get("/")
+def root():
+    return {"status": "running"}
+
+
+if __name__ == "__main__":
+    main()

tool_use_env/server/requirements.txt

@@ -0,0 +1,7 @@
+openenv
+fastapi
+dotenv
+uvicorn
+pydantic
+python-dotenv
+openai

tool_use_env/server/tool_use_env_environment.py

@@ -0,0 +1,351 @@
+from __future__ import annotations
+
+import random
+import uuid
+from typing import Any
+
+from openenv.core.env_server import Environment
+
+from tool_use_env.grader import grade_task
+from tool_use_env.models import ToolUseAction, ToolUseObservation, ToolUseState
+from tool_use_env.tasks import TASKS, TASK_SEQUENCE
+
+
+class ToolUseEnvironment(Environment):
+    SUPPORTS_CONCURRENT_SESSIONS = True
+    MAX_STEPS = 6
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._state = ToolUseState()
+        self._active_task: dict[str, Any] | None = None
+        self._task_cursor = 0
+
+    def _select_task(self, seed: int | None = None, task_id: str | None = None) -> dict[str, Any]:
+        if task_id:
+            if task_id not in TASKS:
+                raise ValueError(f"Unknown task_id '{task_id}'")
+            return TASKS[task_id]
+
+        if seed is not None:
+            rng = random.Random(seed)
+            return TASKS[TASK_SEQUENCE[rng.randrange(len(TASK_SEQUENCE))]]
+
+        selected = TASKS[TASK_SEQUENCE[self._task_cursor % len(TASK_SEQUENCE)]]
+        self._task_cursor += 1
+        return selected
+
+    def reset(
+        self,
+        seed: int | None = None,
+        episode_id: str | None = None,
+        **kwargs: Any,
+    ) -> ToolUseObservation:
+        task = self._select_task(seed=seed, task_id=kwargs.get("task_id"))
+        self._active_task = task
+
+        self._state = ToolUseState(
+            episode_id=episode_id or str(uuid.uuid4()),
+            step_count=0,
+            task_id=task["task_id"],
+            task_name=task["task_name"],
+            difficulty=task["difficulty"],
+            objective=task["objective"],
+            cumulative_reward=0.0,
+            final_score=0.0,
+            drafted_reply=None,
+            resolution_code=None,
+            expected_resolution_code=task["expected_resolution_code"],
+            required_evidence=list(task["required_evidence"]),
+            collected_evidence=["ticket"],
+            action_history=[],
+            repeat_action_count=0,
+            last_action_error=None,
+            known_artifacts={},
+            known_policies={},
+        )
+
+        return self._build_observation(
+            reward=0.0,
+            done=False,
+            last_tool_result=(
+                "Ticket loaded. Start by reviewing the ticket, then inspect the most relevant "
+                "artifacts and policy before submitting a resolution."
+            ),
+        )
+
+    def _normalize_artifact_id(self, artifact_id: str | None) -> str | None:
+        if not artifact_id:
+            return None
+        normalized = artifact_id.strip().lower().replace(" ", "_")
+        aliases = {
+            "payments": "payment",
+            "billing": "payment",
+            "risk": "risk_log",
+            "risklog": "risk_log",
+            "profile": "account",
+        }
+        return aliases.get(normalized, normalized)
+
+    def _resolve_policy_key(self, query: str | None) -> str | None:
+        if not query or not self._active_task:
+            return None
+
+        normalized = query.strip().lower().replace(" ", "_")
+        policies = self._active_task["policies"]
+
+        if normalized in policies:
+            return normalized
+
+        alias_map = {
+            "damaged": "damaged_items",
+            "damage": "damaged_items",
+            "replacement": "damaged_items",
+            "duplicate": "duplicate_charge",
+            "duplicate_charge": "duplicate_charge",
+            "billing": "duplicate_charge",
+            "fraud": "account_takeover",
+            "takeover": "account_takeover",
+            "account_takeover": "account_takeover",
+            "security": "account_takeover",
+        }
+        mapped = alias_map.get(normalized)
+        if mapped in policies:
+            return mapped
+
+        for key in policies:
+            if normalized in key:
+                return key
+        return None
+
+    def _record_repeat_if_needed(self, evidence_key: str) -> bool:
+        if evidence_key in self._state.collected_evidence:
+            self._state.repeat_action_count += 1
+            return True
+        return False
+
+    def _partial_score(self) -> float:
+        if not self._active_task:
+            return 0.0
+        return grade_task(
+            self._active_task,
+            self._state.collected_evidence,
+            self._state.drafted_reply,
+            self._state.resolution_code,
+            self._state.step_count,
+            self._state.repeat_action_count,
+        )["final_score"]
+
+    def _append_history(self, action: ToolUseAction) -> None:
+        parts = [action.action_type]
+        if action.artifact_id:
+            parts.append(f"artifact={action.artifact_id}")
+        if action.query:
+            parts.append(f"query={action.query}")
+        if action.resolution_code:
+            parts.append(f"resolution={action.resolution_code}")
+        self._state.action_history.append(" | ".join(parts))
+
+    def _build_observation(
+        self,
+        reward: float,
+        done: bool,
+        last_tool_result: str | None,
+        last_action_error: str | None = None,
+    ) -> ToolUseObservation:
+        task = self._active_task
+        if not task:
+            raise RuntimeError("Environment has no active task.")
+
+        score = self._state.final_score if done else self._partial_score()
+        remaining_steps = max(0, self.MAX_STEPS - self._state.step_count)
+        known_items = self._state.collected_evidence or ["ticket"]
+        draft_status = "present" if self._state.drafted_reply else "missing"
+        resolution_status = self._state.resolution_code or "not submitted"
+
+        summary = (
+            f"Known evidence: {', '.join(known_items)}. "
+            f"Draft reply: {draft_status}. "
+            f"Resolution: {resolution_status}. "
+            f"Submit the best supported resolution before steps run out."
+        )
+
+        return ToolUseObservation(
+            done=done,
+            reward=round(min(max(reward, 0.0), 1.0), 3),
+            task_id=task["task_id"],
+            difficulty=task["difficulty"],
+            objective=task["objective"],
+            customer_message=task["customer_message"],
+            workspace_summary=summary,
+            available_actions=[
+                "review_ticket",
+                "inspect_artifact",
+                "search_policy",
+                "draft_reply",
+                "submit_resolution",
+            ],
+            available_resolution_codes=list(task["available_resolution_codes"]),
+            collected_evidence=list(self._state.collected_evidence),
+            last_tool_result=last_tool_result,
+            last_action_error=last_action_error,
+            remaining_steps=remaining_steps,
+            current_score=round(score, 3),
+            metadata={
+                "task_name": task["task_name"],
+                "action_history": list(self._state.action_history),
+            },
+        )
+
+    def _finish_episode(self, resolution_code: str | None, feedback: str) -> ToolUseObservation:
+        if not self._active_task:
+            raise RuntimeError("Environment has no active task.")
+
+        self._state.resolution_code = resolution_code
+        breakdown = grade_task(
+            self._active_task,
+            self._state.collected_evidence,
+            self._state.drafted_reply,
+            self._state.resolution_code,
+            self._state.step_count,
+            self._state.repeat_action_count,
+        )
+        self._state.final_score = breakdown["final_score"]
+        self._state.last_action_error = None
+
+        result_text = (
+            f"{feedback} | final_score={breakdown['final_score']:.3f} | "
+            f"resolution_score={breakdown['resolution_score']:.3f} | "
+            f"evidence_score={breakdown['evidence_score']:.3f} | "
+            f"reply_score={breakdown['reply_score']:.3f} | "
+            f"efficiency_score={breakdown['efficiency_score']:.3f}"
+        )
+
+        return self._build_observation(
+            reward=breakdown["final_score"],
+            done=True,
+            last_tool_result=result_text,
+        )
+
+    def step(
+        self,
+        action: ToolUseAction,
+        timeout_s: float | None = None,
+        **kwargs: Any,
+    ) -> ToolUseObservation:
+        if not self._active_task:
+            raise RuntimeError("Call reset() before step().")
+
+        if self._state.final_score > 0 and self._state.resolution_code:
+            return self._build_observation(
+                reward=0.0,
+                done=True,
+                last_tool_result="Episode already finished.",
+                last_action_error="episode_already_done",
+            )
+
+        self._state.step_count += 1
+        self._append_history(action)
+
+        reward = 0.0
+        last_tool_result = None
+        error = None
+
+        if action.action_type == "review_ticket":
+            repeated = self._record_repeat_if_needed("ticket")
+            reward = 0.02 if repeated else 0.10
+            last_tool_result = self._active_task["customer_message"]
+
+        elif action.action_type == "inspect_artifact":
+            artifact_id = self._normalize_artifact_id(action.artifact_id)
+            artifacts = self._active_task["artifacts"]
+            if not artifact_id or artifact_id not in artifacts:
+                error = "invalid_artifact_id"
+                last_tool_result = (
+                    "Unknown artifact. Valid artifacts: "
+                    + ", ".join(sorted(artifacts.keys()))
+                )
+            else:
+                evidence_key = f"artifact:{artifact_id}"
+                repeated = self._record_repeat_if_needed(evidence_key)
+                if not repeated:
+                    self._state.collected_evidence.append(evidence_key)
+                    self._state.known_artifacts[artifact_id] = artifacts[artifact_id]
+                    reward = 0.14 if evidence_key in self._state.required_evidence else 0.04
+                else:
+                    reward = 0.01
+                last_tool_result = artifacts[artifact_id]
+
+        elif action.action_type == "search_policy":
+            policy_key = self._resolve_policy_key(action.query)
+            policies = self._active_task["policies"]
+            if not policy_key:
+                error = "policy_not_found"
+                last_tool_result = (
+                    "No matching policy found. Available policies: "
+                    + ", ".join(sorted(policies.keys()))
+                )
+            else:
+                evidence_key = f"policy:{policy_key}"
+                repeated = self._record_repeat_if_needed(evidence_key)
+                if not repeated:
+                    self._state.collected_evidence.append(evidence_key)
+                    self._state.known_policies[policy_key] = policies[policy_key]
+                    reward = 0.14 if evidence_key in self._state.required_evidence else 0.04
+                else:
+                    reward = 0.01
+                last_tool_result = policies[policy_key]
+
+        elif action.action_type == "draft_reply":
+            if not action.message or not action.message.strip():
+                error = "empty_reply"
+                last_tool_result = "Draft reply cannot be empty."
+            else:
+                self._state.drafted_reply = action.message.strip()
+                keywords = self._active_task["reply_keywords"]
+                hits = sum(
+                    1 for keyword in keywords if keyword.lower() in self._state.drafted_reply.lower()
+                )
+                reward = round(0.05 + (0.15 * (hits / len(keywords))), 3)
+                last_tool_result = (
+                    f"Draft saved. Included {hits}/{len(keywords)} required reply cues."
+                )
+
+        elif action.action_type == "submit_resolution":
+            if not action.resolution_code:
+                error = "missing_resolution_code"
+                last_tool_result = "submit_resolution requires a resolution_code."
+            elif action.resolution_code not in self._active_task["available_resolution_codes"]:
+                error = "invalid_resolution_code"
+                last_tool_result = (
+                    "Unsupported resolution code. Valid codes: "
+                    + ", ".join(self._active_task["available_resolution_codes"])
+                )
+            else:
+                return self._finish_episode(
+                    resolution_code=action.resolution_code,
+                    feedback=f"Resolution submitted: {action.resolution_code}",
+                )
+
+        else:
+            error = "invalid_action_type"
+            last_tool_result = "Unsupported action_type."
+
+        self._state.last_action_error = error
+        if self._state.step_count >= self.MAX_STEPS:
+            return self._finish_episode(
+                resolution_code=self._state.resolution_code,
+                feedback="Episode ended because the step limit was reached.",
+            )
+
+        self._state.cumulative_reward = round(self._state.cumulative_reward + reward, 3)
+        return self._build_observation(
+            reward=reward,
+            done=False,
+            last_tool_result=last_tool_result,
+            last_action_error=error,
+        )
+
+    @property
+    def state(self) -> ToolUseState:
+        return self._state

tool_use_env/tasks.py

@@ -0,0 +1,141 @@
+TASKS = {
+    "damaged-mug-replacement": {
+        "task_id": "damaged-mug-replacement",
+        "task_name": "Damaged Mug Replacement",
+        "difficulty": "easy",
+        "objective": (
+            "Resolve a damaged-item support ticket by gathering the right evidence, "
+            "drafting a concise customer reply, and submitting the correct resolution."
+        ),
+        "customer_message": (
+            "Hi support, my Northwind ceramic mug from order O-1001 arrived shattered. "
+            "I uploaded a photo. I still want the mug if you can replace it quickly."
+        ),
+        "artifacts": {
+            "order": (
+                "Order O-1001 | Item: Northwind ceramic mug | Delivered: 2026-04-01 | "
+                "Photo evidence attached: yes | Carrier note: box dented on arrival."
+            ),
+            "account": (
+                "Customer since 2023 | No prior claims abuse | Shipping address verified."
+            ),
+        },
+        "policies": {
+            "damaged_items": (
+                "Damaged items reported within 7 days with photo evidence qualify for a "
+                "free replacement. Low-cost broken items do not need to be returned."
+            )
+        },
+        "required_evidence": ["ticket", "artifact:order", "policy:damaged_items"],
+        "expected_resolution_code": "send_replacement",
+        "available_resolution_codes": [
+            "send_replacement",
+            "issue_refund",
+            "request_more_info",
+            "deny_request",
+        ],
+        "reply_keywords": ["replacement", "48 hours", "no need to return"],
+        "optimal_steps": 4,
+    },
+    "duplicate-charge-refund": {
+        "task_id": "duplicate-charge-refund",
+        "task_name": "Duplicate Charge Refund",
+        "difficulty": "medium",
+        "objective": (
+            "Investigate a billing complaint, confirm whether a duplicate charge occurred, "
+            "and choose the correct refund resolution."
+        ),
+        "customer_message": (
+            "I was charged twice for the same blender order and only received one item. "
+            "Please fix this. Order number is O-2044."
+        ),
+        "artifacts": {
+            "order": (
+                "Order O-2044 | Item: Pulse blender | Fulfilled once on 2026-03-28 | "
+                "Only one shipment and one invoice should exist."
+            ),
+            "payment": (
+                "Payments for O-2044 | Charge 1: CAPTURED $89.00 | Charge 2: CAPTURED $89.00 | "
+                "Gateway note: retry after timeout caused duplicate capture."
+            ),
+            "account": "Customer in good standing | No prior charge disputes.",
+        },
+        "policies": {
+            "duplicate_charge": (
+                "If a duplicate capture is confirmed, refund the extra charge to the original "
+                "payment method within 3-5 business days."
+            )
+        },
+        "required_evidence": [
+            "ticket",
+            "artifact:order",
+            "artifact:payment",
+            "policy:duplicate_charge",
+        ],
+        "expected_resolution_code": "refund_duplicate_charge",
+        "available_resolution_codes": [
+            "refund_duplicate_charge",
+            "issue_refund",
+            "request_more_info",
+            "deny_request",
+        ],
+        "reply_keywords": ["duplicate charge", "refund", "3-5 business days"],
+        "optimal_steps": 5,
+    },
+    "account-takeover-fraud": {
+        "task_id": "account-takeover-fraud",
+        "task_name": "Account Takeover Fraud Escalation",
+        "difficulty": "hard",
+        "objective": (
+            "Handle a potential account takeover involving digital gift cards. Gather fraud "
+            "signals, apply the correct security-first resolution, and communicate next steps."
+        ),
+        "customer_message": (
+            "I did not place the gift card order on my account last night. Please stop any "
+            "more charges and secure my account immediately."
+        ),
+        "artifacts": {
+            "account": (
+                "Account A-77 | Password reset disabled | New login from unfamiliar IP at 02:14 UTC | "
+                "2FA not enabled."
+            ),
+            "order": (
+                "Order O-7781 | Digital gift cards totaling $250 | Delivered instantly by email | "
+                "Marked high-risk by payments."
+            ),
+            "risk_log": (
+                "Risk engine score: 0.98 | Device mismatch: true | Velocity spike: true | "
+                "Recommendation: lock account and escalate to fraud operations."
+            ),
+        },
+        "policies": {
+            "account_takeover": (
+                "Suspected account takeover with digital goods requires immediate account lock, "
+                "fraud team escalation, and customer follow-up within 24 hours. Do not promise "
+                "an automatic refund before investigation."
+            )
+        },
+        "required_evidence": [
+            "ticket",
+            "artifact:account",
+            "artifact:risk_log",
+            "policy:account_takeover",
+        ],
+        "expected_resolution_code": "lock_account_and_escalate_fraud",
+        "available_resolution_codes": [
+            "lock_account_and_escalate_fraud",
+            "issue_refund",
+            "request_more_info",
+            "deny_request",
+        ],
+        "reply_keywords": ["account locked", "fraud team", "24 hours"],
+        "optimal_steps": 5,
+    },
+}
+
+
+TASK_SEQUENCE = [
+    "damaged-mug-replacement",
+    "duplicate-charge-refund",
+    "account-takeover-fraud",
+]