Upload folder using huggingface_hub

Dockerfile

@@ -0,0 +1,81 @@
+# 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=rag_optimizer
+
+# 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
+ENV ENABLE_WEB_INTERFACE=true
+CMD ["sh", "-c", "cd /app/env && uvicorn server.app:app --host 0.0.0.0 --port 8000"]

LICENSE

@@ -0,0 +1,28 @@
+BSD 3-Clause License
+
+Copyright (c) 2026, Jayadev D
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

README.md

@@ -1,10 +1,255 @@
 ---
-title: Rag Optimizer
-emoji: 👁
-colorFrom: purple
-colorTo: indigo
+title: Rag Optimizer Environment Server
+emoji: 📸
+colorFrom: blue
+colorTo: gray
 sdk: docker
 pinned: false
+app_port: 8000
+base_path: /web
+tags:
+  - openenv
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Rag Optimizer Environment
+
+A simple test environment that echoes back messages. Perfect for testing the env APIs as well as demonstrating environment usage patterns.
+
+## Quick Start
+
+The simplest way to use the Rag Optimizer environment is through the `RagOptimizerEnv` class:
+
+```python
+from rag_optimizer import RagOptimizerAction, RagOptimizerEnv
+
+try:
+    # Create environment from Docker image
+    rag_optimizerenv = RagOptimizerEnv.from_docker_image("rag_optimizer-env:latest")
+
+    # Reset
+    result = rag_optimizerenv.reset()
+    print(f"Reset: {result.observation.echoed_message}")
+
+    # Send multiple messages
+    messages = ["Hello, World!", "Testing echo", "Final message"]
+
+    for msg in messages:
+        result = rag_optimizerenv.step(RagOptimizerAction(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
+    rag_optimizerenv.close()
+```
+
+That's it! The `RagOptimizerEnv.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 rag_optimizer-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
+**RagOptimizerAction**: Contains a single field
+- `message` (str) - The message to echo back
+
+### Observation
+**RagOptimizerObservation**: 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 Rag Optimizer environment server running, you can connect directly:
+
+```python
+from rag_optimizer import RagOptimizerEnv
+
+# Connect to existing server
+rag_optimizerenv = RagOptimizerEnv(base_url="<ENV_HTTP_URL_HERE>")
+
+# Use as normal
+result = rag_optimizerenv.reset()
+result = rag_optimizerenv.step(RagOptimizerAction(message="Hello!"))
+```
+
+Note: When connecting to an existing server, `rag_optimizerenv.close()` will NOT stop the server.
+
+### Using the Context Manager
+
+The client supports context manager usage for automatic connection management:
+
+```python
+from rag_optimizer import RagOptimizerAction, RagOptimizerEnv
+
+# Connect with context manager (auto-connects and closes)
+with RagOptimizerEnv(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(RagOptimizerAction(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(
+    RagOptimizerEnvironment,  # Pass class, not instance
+    RagOptimizerAction,
+    RagOptimizerObservation,
+    max_concurrent_envs=4,  # Allow 4 concurrent sessions
+)
+```
+
+Then multiple clients can connect simultaneously:
+
+```python
+from rag_optimizer import RagOptimizerAction, RagOptimizerEnv
+from concurrent.futures import ThreadPoolExecutor
+
+def run_episode(client_id: int):
+    with RagOptimizerEnv(base_url="http://localhost:8000") as env:
+        result = env.reset()
+        for i in range(10):
+            result = env.step(RagOptimizerAction(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/rag_optimizer_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
+
+```
+rag_optimizer/
+├── .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              # RagOptimizerEnv client
+├── models.py              # Action and Observation models
+└── server/
+    ├── __init__.py        # Server module exports
+    ├── rag_optimizer_environment.py  # Core environment logic
+    ├── app.py             # FastAPI application (HTTP + WebSocket endpoints)
+    └── Dockerfile         # Container image definition
+```

__init__.py

@@ -0,0 +1,16 @@
+# 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.
+
+"""Rag Optimizer Environment."""
+
+from .client import RagOptimizerEnv
+from .models import RagOptimizerAction, RagOptimizerObservation
+
+__all__ = [
+    "RagOptimizerAction",
+    "RagOptimizerObservation",
+    "RagOptimizerEnv",
+]

client.py

@@ -0,0 +1,56 @@
+# 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.
+
+"""Rag Optimizer 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 RagOptimizerAction, RagOptimizerObservation
+
+
+class RagOptimizerEnvClient(EnvClient[RagOptimizerAction, RagOptimizerObservation, State]):
+    """
+    Client for the Rag Optimizer Environment.
+    Translates local Pydantic objects to JSON for the OpenEnv WebSocket.
+    """
+
+    def _step_payload(self, action: RagOptimizerAction) -> Dict:
+        """Convert RagOptimizerAction to JSON payload."""
+        return {
+            "action_type": action.action_type,
+            "doc_id": action.doc_id,
+            "text": action.text,
+            "metadata_key": action.metadata_key,
+            "metadata_value": action.metadata_value,
+        }
+
+    def _parse_result(self, payload: Dict) -> StepResult[RagOptimizerObservation]:
+        """Parse server response back into RagOptimizerObservation."""
+        obs_data = payload.get("observation", {})
+        observation = RagOptimizerObservation(
+            message=obs_data.get("message", ""),
+            current_docs=obs_data.get("current_docs", {}),
+            done=payload.get("done", False),
+            reward=payload.get("reward", 0.0),
+            metadata=obs_data.get("metadata", {}),
+        )
+
+        return StepResult(
+            observation=observation,
+            reward=payload.get("reward", 0.0),
+            done=payload.get("done", False),
+        )
+
+    def _parse_state(self, payload: Dict) -> State:
+        """Parse the hidden tracking state."""
+        return State(
+            episode_id=payload.get("episode_id"),
+            step_count=payload.get("step_count", 0),
+        )

inference.py

@@ -0,0 +1,159 @@
+import os
+import sys
+import json
+from openai import OpenAI
+
+# Add the envs module to path so we can import client and models
+sys.path.append(os.path.join(os.path.dirname(__file__), "envs", "rag_optimizer_env"))
+from client import RagOptimizerEnvClient
+from models import RagOptimizerAction
+
+# Load environment variables
+API_BASE_URL = os.getenv("API_BASE_URL")
+MODEL_NAME = os.getenv("MODEL_NAME")
+HF_TOKEN = os.getenv("HF_TOKEN")
+
+MAX_STEPS = 30
+
+SYSTEM_PROMPT = """You are an automated Data Engineer managing an AI Knowledge Base.
+Your goal is to optimize the messy chunks of text in the database so that a TF-IDF Search Algorithm can find answers easily.
+You must resolve contradictions, categorize documents, and delete unnecessary documents.
+
+After each action you will receive a "current_reward" score (0.0 to 1.0) indicating how well the KB currently performs. Use this to guide your strategy.
+
+You have the following actions:
+- {"action_type": "read_document", "doc_id": "..."}
+- {"action_type": "update_document", "doc_id": "...", "text": "..."}
+- {"action_type": "delete_document", "doc_id": "..."}
+- {"action_type": "add_metadata", "doc_id": "...", "metadata_key": "...", "metadata_value": "..."}
+- {"action_type": "submit"}
+
+You must return ONLY a raw JSON object detailing the action you want to take!"""
+
+def format_action_str(action: RagOptimizerAction) -> str:
+    if action.action_type == "read_document":
+        return f"read('{action.doc_id}')"
+    elif action.action_type == "update_document":
+        return f"update('{action.doc_id}')"
+    elif action.action_type == "delete_document":
+        return f"delete('{action.doc_id}')"
+    elif action.action_type == "add_metadata":
+        return f"add_metadata('{action.doc_id}','{action.metadata_key}')"
+    elif action.action_type == "submit":
+        return "submit()"
+    return f"{action.action_type}()"
+
+def main():
+    # Setup OpenAI Client
+    client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN)
+    
+    # Track metrics for the final output
+    step_rewards = []
+    success = False
+    error_msg = "null"
+    score = 0.0
+    
+    print(f"[START] task=rag_optimizer_env env=OpenEnv model={MODEL_NAME}")
+    
+    # We suppress any other custom prints to respect the STDOUT format strictly
+    import contextlib
+    import io
+    
+    with RagOptimizerEnvClient(base_url="http://localhost:8000").sync() as env:
+        # Suppress prints from client or env reset
+        with contextlib.redirect_stdout(io.StringIO()):
+            try:
+                result = env.reset()
+                observation = result.observation
+            except Exception as e:
+                error_msg = str(e).replace('\n', ' ')
+                print(f"[END] success=false steps=0 score=0.00 rewards=")
+                return
+
+        history = [{"role": "system", "content": SYSTEM_PROMPT}]
+        
+        init_obs = {
+            "server_feedback": observation.message,
+            "current_reward": observation.reward,
+            "current_knowledge_base": observation.current_docs
+        }
+        history.append({"role": "user", "content": json.dumps(init_obs, indent=2)})
+        
+        step = 0
+        for i in range(1, MAX_STEPS + 1):
+            step = i
+            messages = list(history)
+            
+            action_str = "unknown"
+            error_msg = "null"
+
+            try:
+                completion = client.chat.completions.create(
+                    model=MODEL_NAME,
+                    messages=messages,
+                    response_format={"type": "json_object"},
+                    max_tokens=1000
+                )
+                response_text = completion.choices[0].message.content or ""
+                action_data = json.loads(response_text)
+                
+                # Normalize fields if model returns lists instead of strings
+                for field in ("doc_id", "text", "metadata_key", "metadata_value"):
+                    val = action_data.get(field)
+                    if isinstance(val, list):
+                        if val and isinstance(val[0], str):
+                            action_data[field] = " ".join(val)
+                        elif val and isinstance(val[0], dict):
+                            action_data[field] = json.dumps(val[0])
+                        else:
+                            action_data[field] = str(val[0]) if val else ""
+                
+                action = RagOptimizerAction(**action_data)
+                action_str = format_action_str(action)
+                
+            except Exception as exc: 
+                error_msg = str(exc).replace('\n', ' ')
+                action = RagOptimizerAction(action_type="submit")
+                action_str = format_action_str(action)
+
+            # Suppress normal prints during step
+            with contextlib.redirect_stdout(io.StringIO()):
+                try:
+                    result = env.step(action)
+                    observation = result.observation
+                    reward = result.reward
+                except Exception as e:
+                    error_msg = str(e).replace('\n', ' ')
+                    reward = 0.0
+                    result = type('obj', (object,), {'done': True})()
+                    observation = type('obj', (object,), {'message': 'error', 'current_docs': {}})()
+
+            step_rewards.append(reward)
+            done = "true" if result.done else "false"
+            
+            print(f"[STEP] step={step} action={action_str} reward={reward:.2f} done={done} error={error_msg}")
+
+            if result.done:
+                success = True if reward > 0.5 else False  # Or however you define success
+                score = float(reward)
+                break
+
+            history.append({"role": "assistant", "content": json.dumps(action.model_dump(), default=str)})
+            next_obs = {
+                "server_feedback": observation.message,
+                "current_reward": observation.reward,
+                "current_knowledge_base": observation.current_docs
+            }
+            history.append({"role": "user", "content": json.dumps(next_obs, indent=2)})
+
+        else:
+            # Reached max steps
+            success = False
+            score = float(result.reward)
+
+    rewards_str = ",".join([f"{r:.2f}" for r in step_rewards])
+    done_str = "true" if success else "false"
+    print(f"[END] success={done_str} steps={step} score={score:.2f} rewards={rewards_str}")
+
+if __name__ == "__main__":
+    main()

models.py

@@ -0,0 +1,62 @@
+# 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.
+
+"""
+Pydantic schemas for the Rag Optimizer environment.
+These define the API contract between the client (agent) and the server.
+"""
+
+from typing import Dict, Literal, Optional
+from pydantic import BaseModel, Field
+
+
+class RagOptimizerAction(BaseModel):
+    """
+    Actions the agent can take to interact with the Knowledge Base.
+    """
+    
+    action_type: Literal["read_document", "update_document", "delete_document", "add_metadata", "submit"] = Field(
+        ..., 
+        description="The RAG optimization tool to execute."
+    )
+    
+    doc_id: Optional[str] = Field(
+        None, 
+        description="The ID of the document to target."
+    )
+    text: Optional[str] = Field(
+        None, 
+        description="The text content (used for update_document)."
+    )
+    metadata_key: Optional[str] = Field(
+        None, 
+        description="The key of the metadata tag (used for add_metadata)."
+    )
+    metadata_value: Optional[str] = Field(
+        None, 
+        description="The value of the metadata tag (used for add_metadata)."
+    )
+
+
+class RagOptimizerObservation(BaseModel):
+    """
+    The environment's response to an action, including the state of the KB.
+    """
+    
+    message: str = Field(
+        ..., 
+        description="Feedback from the last action (e.g., success/error messages)."
+    )
+    
+    current_docs: Dict[str, Dict] = Field(
+        ..., 
+        description="A live summary of the documents currently inside the KB (doc_id -> metadata/length)."
+    )
+    
+    # Required OpenEnv standard fields
+    done: bool = Field(False, description="Whether the episode has finished.")
+    reward: float = Field(0.0, description="The reward obtained from the last step.")
+    metadata: Dict = Field(default_factory=dict, description="Additional optional information.")

openenv.yaml

@@ -0,0 +1,7 @@
+spec_version: 1
+name: rag_optimizer
+type: space
+runtime: fastapi
+app: server.app:app
+port: 8000
+

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-rag_optimizer"
+version = "0.1.0"
+description = "Rag Optimizer 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.2",
+    # 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 rag_optimizer.server.app
+server = "rag_optimizer.server.app:main"
+
+[tool.setuptools]
+include-package-data = true
+packages = ["rag_optimizer", "rag_optimizer.server"]
+package-dir = { "rag_optimizer" = ".", "rag_optimizer.server" = "server" }

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.
+
+"""Rag Optimizer environment server components."""
+
+from .rag_optimizer_environment import RagOptimizerEnvironment
+
+__all__ = ["RagOptimizerEnvironment"]

server/app.py

@@ -0,0 +1,79 @@
+# 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.
+
+"""
+FastAPI application for the Rag Optimizer Environment.
+
+This module creates an HTTP server that exposes the RagOptimizerEnvironment
+over HTTP and WebSocket endpoints, compatible with EnvClient.
+
+Endpoints:
+    - POST /reset: Reset the environment
+    - POST /step: Execute an action
+    - GET /state: Get current environment state
+    - GET /schema: Get action/observation schemas
+    - WS /ws: WebSocket endpoint for persistent sessions
+
+Usage:
+    # Development (with auto-reload):
+    uvicorn server.app:app --reload --host 0.0.0.0 --port 8000
+
+    # Production:
+    uvicorn server.app:app --host 0.0.0.0 --port 8000 --workers 4
+
+    # Or run directly:
+    python -m server.app
+"""
+
+try:
+    from openenv.core.env_server.http_server import create_app
+except Exception as e:  # pragma: no cover
+    raise ImportError(
+        "openenv is required for the web interface. Install dependencies with '\n    uv sync\n'"
+    ) from e
+
+try:
+    from models import RagOptimizerAction, RagOptimizerObservation
+    from .rag_optimizer_environment import RagOptimizerEnvironment
+except ModuleNotFoundError:
+    from models import RagOptimizerAction, RagOptimizerObservation
+    from server.rag_optimizer_environment import RagOptimizerEnvironment
+
+
+# Create the app with web interface and README integration
+app = create_app(
+    RagOptimizerEnvironment,
+    RagOptimizerAction,
+    RagOptimizerObservation,
+    env_name="rag_optimizer",
+    max_concurrent_envs=1,  # increase this number to allow more concurrent WebSocket sessions
+)
+
+
+def main(host: str = "0.0.0.0", port: int = 8002):
+    """
+    Entry point for direct execution via uv run or python -m.
+
+    This function enables running the server without Docker:
+        uv run --project . server
+        uv run --project . server --port 8001
+        python -m rag_optimizer.server.app
+
+    Args:
+        host: Host address to bind to (default: "0.0.0.0")
+        port: Port number to listen on (default: 8000)
+
+    For production deployments, consider using uvicorn directly with
+    multiple workers:
+        uvicorn rag_optimizer.server.app:app --workers 4
+    """
+    import uvicorn
+
+    uvicorn.run(app, host=host, port=port)
+
+
+if __name__ == '__main__':
+    main()

server/rag_optimizer_environment.py

@@ -0,0 +1,209 @@
+# 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.
+
+"""
+Rag Optimizer Environment Implementation.
+The agent acts as a Data Engineer to un-block a broken RAG pipeline.
+"""
+
+from uuid import uuid4
+from typing import Dict, Any, List
+
+from openenv.core.env_server.interfaces import Environment
+from openenv.core.env_server.types import State
+
+# Import scikit-learn for our Grader
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.metrics.pairwise import cosine_similarity
+import numpy as np
+
+try:
+    from models import RagOptimizerAction, RagOptimizerObservation
+except ImportError:
+    from models import RagOptimizerAction, RagOptimizerObservation
+
+
+class RagOptimizerEnvironment(Environment):
+    """
+    RAG Optimizer Engine.
+    Maintains a simulated Knowledge Base and grades it using TF-IDF.
+    """
+
+    SUPPORTS_CONCURRENT_SESSIONS: bool = True
+
+    def __init__(self):
+        self._state = State(episode_id=str(uuid4()), step_count=0)
+        
+        # Initial messy knowledge base
+        self.kb = {
+            "doc_pricing_legacy": {
+                "text": "Pricing for 2021: Enterprise tier is $1000/mo. Standard is $500/mo. All plans include 10 users.",
+                "metadata": {"type": "pricing"}
+            },
+            "doc_pricing_current_v2": {
+                "text": "Current Pricing 2024: Enterprise is $1500/mo. Standard is $750/mo. Refunds are not permitted on the enterprise tier.",
+                "metadata": {}
+            },
+            "doc_shipping_policy": {
+                "text": "All internal shipments to remote branch offices take 5-7 business days. Overnight shipping is only available for C-suite.",
+                "metadata": {"department": "logistics"}
+            },
+            "doc_messy_support_ticket_1": {
+                "text": "User complained the button disappeared on the frontend. Another user said the database latency was high. The frontend team fixed the button by updating CSS.",
+                "metadata": {}
+            },
+            "doc_messy_support_ticket_2": {
+                "text": "Email integration is failing with error 401 Unauthorized. The API key was rotated on Tuesday.",
+                "metadata": {}
+            },
+            "doc_monolithic_onboarding": {
+                "text": "Welcome to the company! Here are some rules. 1) VPN access requires DUO. 2) The cafetaria opens at 8 AM. 3) For HR issues, email hr@company.com. 4) The 2024 holiday schedule includes Dec 25, Jan 1, and July 4. 5) Parking passes must be renewed annually in March.",
+                "metadata": {}
+            },
+            # Add distractor files
+            **{f"doc_distractor_hr_{i}": {"text": f"This is an old HR policy document regarding {['pto', 'sick leave', 'travel', 'expenses'][i%4]} from 201{i%10}.", "metadata":{}} for i in range(10)},
+            **{f"doc_distractor_eng_{i}": {"text": f"Engineering architecture decision record {i}. We decided to use {['React', 'Postgres', 'Redis', 'Kafka'][i%4]} because of scaling concerns.", "metadata":{}} for i in range(10)},
+            **{f"doc_distractor_random_{i}": {"text": f"Weekly team update notes. Nothing important here, just discussed the weather and the upcoming launch {i}.", "metadata":{}} for i in range(10)},
+        }
+        
+        # Hidden test suite for the grader
+        self.test_suite = [
+            {
+                "query": "What is the current 2024 price for standard?",
+                "target_concept": "750/mo"
+            },
+            {
+                "query": "What is the refund policy for enterprise?",
+                "target_concept": "Refunds are not permitted"
+            },
+            {
+                "query": "UI issues frontend CSS missing button",
+                "target_concept": "frontend team fixed the button"
+            },
+            {
+                "query": "How long does shipping take to branch offices?",
+                "target_concept": "5-7 business days"
+            },
+            {
+                "query": "What months do parking passes need to be renewed?",
+                "target_concept": "March"
+            },
+            {
+                "query": "What holidays are we off in 2024?",
+                "target_concept": "July 4"
+            }
+        ]
+
+    def _get_kb_summary(self) -> Dict[str, Dict]:
+        """Returns a summary of the KB for the observation."""
+        summary = {}
+        for k, v in self.kb.items():
+            summary[k] = {"metadata": v.get("metadata", {}), "length": len(v.get("text", ""))}
+        return summary
+
+    def reset(self) -> RagOptimizerObservation:
+        self._state = State(episode_id=str(uuid4()), step_count=0)
+        return RagOptimizerObservation(
+            message="RagOptimizerEnv Initialized. You have messy chunks in the KB. Resolve conflicts, add metadata tags to short tickets, and splinter monolithic files to win.",
+            current_docs=self._get_kb_summary(),
+            done=False,
+            reward=self._evaluate_kb()
+        )
+
+    def _evaluate_kb(self) -> float:
+        """The Grader: Evaluates the agent's current KB using TF-IDF."""
+        if not self.kb:
+            return 0.0
+            
+        doc_texts = [doc["text"] for doc in self.kb.values()]
+        
+        vectorizer = TfidfVectorizer(stop_words='english')
+        try:
+            doc_vectors = vectorizer.fit_transform(doc_texts)
+        except ValueError:
+            return 0.0
+            
+        score = 0.0
+        
+        for case in self.test_suite:
+            query_vec = vectorizer.transform([case["query"]])
+            similarities = cosine_similarity(query_vec, doc_vectors)[0]
+            
+            # Get top 3
+            top_k_indices = similarities.argsort()[-3:][::-1]
+            
+            found = False
+            for idx in top_k_indices:
+                if similarities[idx] > 0.01:
+                    if case["target_concept"].lower() in doc_texts[idx].lower():
+                        found = True
+                        break
+            if found:
+                score += 1.0
+                
+        return float(score / len(self.test_suite))
+
+    def step(self, action: RagOptimizerAction) -> RagOptimizerObservation:  # type: ignore[override]
+        self._state.step_count += 1
+        
+        msg = ""
+        done = False
+        reward = 0.0
+        
+        try:
+            if action.action_type == "read_document":
+                if action.doc_id in self.kb:
+                    msg = f"Content of {action.doc_id}: {self.kb[action.doc_id]['text']}"
+                else:
+                    msg = f"Error: doc_id {action.doc_id} not found."
+            
+            elif action.action_type == "delete_document":
+                if action.doc_id in self.kb:
+                    del self.kb[action.doc_id]
+                    msg = f"Deleted {action.doc_id}."
+                else:
+                    msg = f"Error: doc_id {action.doc_id} not found."
+                    
+            elif action.action_type == "update_document":
+                if not action.doc_id or not action.text:
+                    msg = "Error: doc_id and text required for update_document."
+                else:
+                    if action.doc_id not in self.kb:
+                        self.kb[action.doc_id] = {"text": "", "metadata": {}}
+                    self.kb[action.doc_id]["text"] = action.text
+                    msg = f"Updated text for {action.doc_id}."
+                    
+            elif action.action_type == "add_metadata":
+                if not action.doc_id or not action.metadata_key or not action.metadata_value:
+                    msg = "Error: doc_id, metadata_key, and metadata_value required."
+                else:
+                    if action.doc_id not in self.kb:
+                        msg = f"Error: doc_id {action.doc_id} not found."
+                    else:
+                        self.kb[action.doc_id]["metadata"][action.metadata_key] = action.metadata_value
+                        msg = f"Added metadata to {action.doc_id}."
+                        
+            elif action.action_type == "submit":
+                done = True
+                reward = self._evaluate_kb()
+                msg = f"Evaluation complete. Final reward: {reward:.2f}"
+                
+        except Exception as e:
+            msg = f"Action failed: {str(e)}"
+
+        if not done:
+            reward = self._evaluate_kb()
+
+        return RagOptimizerObservation(
+            message=msg,
+            current_docs=self._get_kb_summary(),
+            done=done,
+            reward=reward,
+        )
+
+    @property
+    def state(self) -> State:
+        return self._state

server/requirements.txt

@@ -0,0 +1,4 @@
+openenv[core]>=0.2.0
+fastapi>=0.115.0
+uvicorn>=0.24.0
+scikit-learn>=1.3.0