Upload folder using huggingface_hub

__init__.py

@@ -1,16 +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",
-]
+# 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

@@ -1,56 +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),
-        )
+# 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

@@ -1,180 +1,180 @@
-import os
-import sys
-import json
-from openai import OpenAI
-
-
-def load_local_env() -> None:
-    """Lightweight .env loader for local runs without requiring python-dotenv."""
-    env_path = os.path.join(os.path.dirname(__file__), ".env")
-    if not os.path.exists(env_path):
-        return
-
-    with open(env_path, "r", encoding="utf-8") as f:
-        for raw_line in f:
-            line = raw_line.strip()
-            if not line or line.startswith("#") or "=" not in line:
-                continue
-            key, value = line.split("=", 1)
-            key = key.strip()
-            value = value.strip().strip('"').strip("'")
-            if key and key not in os.environ:
-                os.environ[key] = value
-
-
-load_local_env()
-
-# 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", "https://router.huggingface.co/v1")
-MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")
-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()
+import os
+import sys
+import json
+from openai import OpenAI
+
+
+def load_local_env() -> None:
+    """Lightweight .env loader for local runs without requiring python-dotenv."""
+    env_path = os.path.join(os.path.dirname(__file__), ".env")
+    if not os.path.exists(env_path):
+        return
+
+    with open(env_path, "r", encoding="utf-8") as f:
+        for raw_line in f:
+            line = raw_line.strip()
+            if not line or line.startswith("#") or "=" not in line:
+                continue
+            key, value = line.split("=", 1)
+            key = key.strip()
+            value = value.strip().strip('"').strip("'")
+            if key and key not in os.environ:
+                os.environ[key] = value
+
+
+load_local_env()
+
+# 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", "https://router.huggingface.co/v1")
+MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")
+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

@@ -1,62 +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.")
+# 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

@@ -1,7 +1,16 @@
-spec_version: 1
-name: rag_optimizer
-type: space
-runtime: fastapi
-app: server.app:app
-port: 8000
-
+spec_version: 1
+name: rag_optimizer
+type: space
+runtime: fastapi
+app: server.app:app
+port: 8000
+tasks:
+  - id: easy
+    steps: 10
+    description: Conflict Resolution - Resolve overlapping pricing parameters by deleting the legacy pricing format.
+  - id: medium
+    steps: 10
+    description: Ontological Tagging - Route exact programatic metadata tags to support tickets.
+  - id: hard
+    steps: 15
+    description: Syntactic Splintering - Break down the monolithic onboarding blob into multiple granular chunks.

pyproject.toml

@@ -1,46 +1,46 @@
-# 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",
-    "scikit-learn>=1.3.0",
-]
-
-[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"]
+# 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",
+    "scikit-learn>=1.3.0",
+]
+
+[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" }

rag_optimizer_environment.py

@@ -1,212 +1,212 @@
-# 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)
-        self.kb = self._get_initial_kb()
-        
-        # 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_initial_kb(self) -> Dict[str, Dict]:
-        """Returns a fresh copy of the initial messy knowledge base."""
-        return {
-            "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)},
-        }
-
-    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)
-        self.kb = self._get_initial_kb()
-        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
+# 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)
+        self.kb = self._get_initial_kb()
+        
+        # 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_initial_kb(self) -> Dict[str, Dict]:
+        """Returns a fresh copy of the initial messy knowledge base."""
+        return {
+            "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)},
+        }
+
+    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)
+        self.kb = self._get_initial_kb()
+        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

run.sh

@@ -1,180 +1,180 @@
-#
-# Checks that your HF Space is live, Docker image builds, and openenv validate passes.
-#
-# Prerequisites:
-#   - Docker:       https://docs.docker.com/get-docker/
-#   - openenv-core: pip install openenv-core
-#   - curl (usually pre-installed)
-#
-# Run:
-#   curl -fsSL https://raw.githubusercontent.com/<owner>/<repo>/main/scripts/validate-submission.sh | bash -s -- <ping_url> [repo_dir]
-#
-#   Or download and run locally:
-#     chmod +x validate-submission.sh
-#     ./validate-submission.sh <ping_url> [repo_dir]
-#
-# Arguments:
-#   ping_url   Your HuggingFace Space URL (e.g. https://your-space.hf.space)
-#   repo_dir   Path to your repo (default: current directory)
-#
-# Examples:
-#   ./validate-submission.sh https://my-team.hf.space
-#   ./validate-submission.sh https://my-team.hf.space ./my-repo
-#
-
-set -uo pipefail
-
-DOCKER_BUILD_TIMEOUT=600
-if [ -t 1 ]; then
-  RED='\033[0;31m'
-  GREEN='\033[0;32m'
-  YELLOW='\033[1;33m'
-  BOLD='\033[1m'
-  NC='\033[0m'
-else
-  RED='' GREEN='' YELLOW='' BOLD='' NC=''
-fi
-
-run_with_timeout() {
-  local secs="$1"; shift
-  if command -v timeout &>/dev/null; then
-    timeout "$secs" "$@"
-  elif command -v gtimeout &>/dev/null; then
-    gtimeout "$secs" "$@"
-  else
-    "$@" &
-    local pid=$!
-    ( sleep "$secs" && kill "$pid" 2>/dev/null ) &
-    local watcher=$!
-    wait "$pid" 2>/dev/null
-    local rc=$?
-    kill "$watcher" 2>/dev/null
-    wait "$watcher" 2>/dev/null
-    return $rc
-  fi
-}
-
-portable_mktemp() {
-  local prefix="${1:-validate}"
-  mktemp "${TMPDIR:-/tmp}/${prefix}-XXXXXX" 2>/dev/null || mktemp
-}
-
-CLEANUP_FILES=()
-cleanup() { rm -f "${CLEANUP_FILES[@]+"${CLEANUP_FILES[@]}"}"; }
-trap cleanup EXIT
-
-PING_URL="${1:-}"
-REPO_DIR="${2:-.}"
-
-if [ -z "$PING_URL" ]; then
-  printf "Usage: %s <ping_url> [repo_dir]\n" "$0"
-  printf "\n"
-  printf "  ping_url   Your HuggingFace Space URL (e.g. https://your-space.hf.space)\n"
-  printf "  repo_dir   Path to your repo (default: current directory)\n"
-  exit 1
-fi
-
-if ! REPO_DIR="$(cd "$REPO_DIR" 2>/dev/null && pwd)"; then
-  printf "Error: directory '%s' not found\n" "${2:-.}"
-  exit 1
-fi
-PING_URL="${PING_URL%/}"
-export PING_URL
-PASS=0
-
-log()  { printf "[%s] %b\n" "$(date -u +%H:%M:%S)" "$*"; }
-pass() { log "${GREEN}PASSED${NC} -- $1"; PASS=$((PASS + 1)); }
-fail() { log "${RED}FAILED${NC} -- $1"; }
-hint() { printf "  ${YELLOW}Hint:${NC} %b\n" "$1"; }
-stop_at() {
-  printf "\n"
-  printf "${RED}${BOLD}Validation stopped at %s.${NC} Fix the above before continuing.\n" "$1"
-  exit 1
-}
-
-printf "\n"
-printf "${BOLD}========================================${NC}\n"
-printf "${BOLD}  OpenEnv Submission Validator${NC}\n"
-printf "${BOLD}========================================${NC}\n"
-log "Repo:     $REPO_DIR"
-log "Ping URL: $PING_URL"
-printf "\n"
-
-log "${BOLD}Step 1/3: Pinging HF Space${NC} ($PING_URL/reset) ..."
-
-CURL_OUTPUT=$(portable_mktemp "validate-curl")
-CLEANUP_FILES+=("$CURL_OUTPUT")
-HTTP_CODE=$(curl -s -o "$CURL_OUTPUT" -w "%{http_code}" -X POST \
-  -H "Content-Type: application/json" -d '{}' \
-  "$PING_URL/reset" --max-time 30 2>"$CURL_OUTPUT" || printf "000")
-
-if [ "$HTTP_CODE" = "200" ]; then
-  pass "HF Space is live and responds to /reset"
-elif [ "$HTTP_CODE" = "000" ]; then
-  fail "HF Space not reachable (connection failed or timed out)"
-  hint "Check your network connection and that the Space is running."
-  hint "Try: curl -s -o /dev/null -w '%%{http_code}' -X POST $PING_URL/reset"
-  stop_at "Step 1"
-else
-  fail "HF Space /reset returned HTTP $HTTP_CODE (expected 200)"
-  hint "Make sure your Space is running and the URL is correct."
-  hint "Try opening $PING_URL in your browser first."
-  stop_at "Step 1"
-fi
-
-log "${BOLD}Step 2/3: Running docker build${NC} ..."
-
-if ! command -v docker &>/dev/null; then
-  fail "docker command not found"
-  hint "Install Docker: https://docs.docker.com/get-docker/"
-  stop_at "Step 2"
-fi
-
-if [ -f "$REPO_DIR/Dockerfile" ]; then
-  DOCKER_CONTEXT="$REPO_DIR"
-elif [ -f "$REPO_DIR/server/Dockerfile" ]; then
-  DOCKER_CONTEXT="$REPO_DIR/server"
-else
-  fail "No Dockerfile found in repo root or server/ directory"
-  stop_at "Step 2"
-fi
-
-log "  Found Dockerfile in $DOCKER_CONTEXT"
-
-BUILD_OK=false
-BUILD_OUTPUT=$(run_with_timeout "$DOCKER_BUILD_TIMEOUT" docker build "$DOCKER_CONTEXT" 2>&1) && BUILD_OK=true
-
-if [ "$BUILD_OK" = true ]; then
-  pass "Docker build succeeded"
-else
-  fail "Docker build failed (timeout=${DOCKER_BUILD_TIMEOUT}s)"
-  printf "%s\n" "$BUILD_OUTPUT" | tail -20
-  stop_at "Step 2"
-fi
-
-log "${BOLD}Step 3/3: Running openenv validate${NC} ..."
-
-if ! command -v openenv &>/dev/null; then
-  fail "openenv command not found"
-  hint "Install it: pip install openenv-core"
-  stop_at "Step 3"
-fi
-
-VALIDATE_OK=false
-VALIDATE_OUTPUT=$(cd "$REPO_DIR" && openenv validate 2>&1) && VALIDATE_OK=true
-
-if [ "$VALIDATE_OK" = true ]; then
-  pass "openenv validate passed"
-  [ -n "$VALIDATE_OUTPUT" ] && log "  $VALIDATE_OUTPUT"
-else
-  fail "openenv validate failed"
-  printf "%s\n" "$VALIDATE_OUTPUT"
-  stop_at "Step 3"
-fi
-
-printf "\n"
-printf "${BOLD}========================================${NC}\n"
-printf "${GREEN}${BOLD}  All 3/3 checks passed!${NC}\n"
-printf "${GREEN}${BOLD}  Your submission is ready to submit.${NC}\n"
-printf "${BOLD}========================================${NC}\n"
+#
+# Checks that your HF Space is live, Docker image builds, and openenv validate passes.
+#
+# Prerequisites:
+#   - Docker:       https://docs.docker.com/get-docker/
+#   - openenv-core: pip install openenv-core
+#   - curl (usually pre-installed)
+#
+# Run:
+#   curl -fsSL https://raw.githubusercontent.com/<owner>/<repo>/main/scripts/validate-submission.sh | bash -s -- <ping_url> [repo_dir]
+#
+#   Or download and run locally:
+#     chmod +x validate-submission.sh
+#     ./validate-submission.sh <ping_url> [repo_dir]
+#
+# Arguments:
+#   ping_url   Your HuggingFace Space URL (e.g. https://your-space.hf.space)
+#   repo_dir   Path to your repo (default: current directory)
+#
+# Examples:
+#   ./validate-submission.sh https://my-team.hf.space
+#   ./validate-submission.sh https://my-team.hf.space ./my-repo
+#
+
+set -uo pipefail
+
+DOCKER_BUILD_TIMEOUT=600
+if [ -t 1 ]; then
+  RED='\033[0;31m'
+  GREEN='\033[0;32m'
+  YELLOW='\033[1;33m'
+  BOLD='\033[1m'
+  NC='\033[0m'
+else
+  RED='' GREEN='' YELLOW='' BOLD='' NC=''
+fi
+
+run_with_timeout() {
+  local secs="$1"; shift
+  if command -v timeout &>/dev/null; then
+    timeout "$secs" "$@"
+  elif command -v gtimeout &>/dev/null; then
+    gtimeout "$secs" "$@"
+  else
+    "$@" &
+    local pid=$!
+    ( sleep "$secs" && kill "$pid" 2>/dev/null ) &
+    local watcher=$!
+    wait "$pid" 2>/dev/null
+    local rc=$?
+    kill "$watcher" 2>/dev/null
+    wait "$watcher" 2>/dev/null
+    return $rc
+  fi
+}
+
+portable_mktemp() {
+  local prefix="${1:-validate}"
+  mktemp "${TMPDIR:-/tmp}/${prefix}-XXXXXX" 2>/dev/null || mktemp
+}
+
+CLEANUP_FILES=()
+cleanup() { rm -f "${CLEANUP_FILES[@]+"${CLEANUP_FILES[@]}"}"; }
+trap cleanup EXIT
+
+PING_URL="${1:-}"
+REPO_DIR="${2:-.}"
+
+if [ -z "$PING_URL" ]; then
+  printf "Usage: %s <ping_url> [repo_dir]\n" "$0"
+  printf "\n"
+  printf "  ping_url   Your HuggingFace Space URL (e.g. https://your-space.hf.space)\n"
+  printf "  repo_dir   Path to your repo (default: current directory)\n"
+  exit 1
+fi
+
+if ! REPO_DIR="$(cd "$REPO_DIR" 2>/dev/null && pwd)"; then
+  printf "Error: directory '%s' not found\n" "${2:-.}"
+  exit 1
+fi
+PING_URL="${PING_URL%/}"
+export PING_URL
+PASS=0
+
+log()  { printf "[%s] %b\n" "$(date -u +%H:%M:%S)" "$*"; }
+pass() { log "${GREEN}PASSED${NC} -- $1"; PASS=$((PASS + 1)); }
+fail() { log "${RED}FAILED${NC} -- $1"; }
+hint() { printf "  ${YELLOW}Hint:${NC} %b\n" "$1"; }
+stop_at() {
+  printf "\n"
+  printf "${RED}${BOLD}Validation stopped at %s.${NC} Fix the above before continuing.\n" "$1"
+  exit 1
+}
+
+printf "\n"
+printf "${BOLD}========================================${NC}\n"
+printf "${BOLD}  OpenEnv Submission Validator${NC}\n"
+printf "${BOLD}========================================${NC}\n"
+log "Repo:     $REPO_DIR"
+log "Ping URL: $PING_URL"
+printf "\n"
+
+log "${BOLD}Step 1/3: Pinging HF Space${NC} ($PING_URL/reset) ..."
+
+CURL_OUTPUT=$(portable_mktemp "validate-curl")
+CLEANUP_FILES+=("$CURL_OUTPUT")
+HTTP_CODE=$(curl -s -o "$CURL_OUTPUT" -w "%{http_code}" -X POST \
+  -H "Content-Type: application/json" -d '{}' \
+  "$PING_URL/reset" --max-time 30 2>"$CURL_OUTPUT" || printf "000")
+
+if [ "$HTTP_CODE" = "200" ]; then
+  pass "HF Space is live and responds to /reset"
+elif [ "$HTTP_CODE" = "000" ]; then
+  fail "HF Space not reachable (connection failed or timed out)"
+  hint "Check your network connection and that the Space is running."
+  hint "Try: curl -s -o /dev/null -w '%%{http_code}' -X POST $PING_URL/reset"
+  stop_at "Step 1"
+else
+  fail "HF Space /reset returned HTTP $HTTP_CODE (expected 200)"
+  hint "Make sure your Space is running and the URL is correct."
+  hint "Try opening $PING_URL in your browser first."
+  stop_at "Step 1"
+fi
+
+log "${BOLD}Step 2/3: Running docker build${NC} ..."
+
+if ! command -v docker &>/dev/null; then
+  fail "docker command not found"
+  hint "Install Docker: https://docs.docker.com/get-docker/"
+  stop_at "Step 2"
+fi
+
+if [ -f "$REPO_DIR/Dockerfile" ]; then
+  DOCKER_CONTEXT="$REPO_DIR"
+elif [ -f "$REPO_DIR/server/Dockerfile" ]; then
+  DOCKER_CONTEXT="$REPO_DIR/server"
+else
+  fail "No Dockerfile found in repo root or server/ directory"
+  stop_at "Step 2"
+fi
+
+log "  Found Dockerfile in $DOCKER_CONTEXT"
+
+BUILD_OK=false
+BUILD_OUTPUT=$(run_with_timeout "$DOCKER_BUILD_TIMEOUT" docker build "$DOCKER_CONTEXT" 2>&1) && BUILD_OK=true
+
+if [ "$BUILD_OK" = true ]; then
+  pass "Docker build succeeded"
+else
+  fail "Docker build failed (timeout=${DOCKER_BUILD_TIMEOUT}s)"
+  printf "%s\n" "$BUILD_OUTPUT" | tail -20
+  stop_at "Step 2"
+fi
+
+log "${BOLD}Step 3/3: Running openenv validate${NC} ..."
+
+if ! command -v openenv &>/dev/null; then
+  fail "openenv command not found"
+  hint "Install it: pip install openenv-core"
+  stop_at "Step 3"
+fi
+
+VALIDATE_OK=false
+VALIDATE_OUTPUT=$(cd "$REPO_DIR" && openenv validate 2>&1) && VALIDATE_OK=true
+
+if [ "$VALIDATE_OK" = true ]; then
+  pass "openenv validate passed"
+  [ -n "$VALIDATE_OUTPUT" ] && log "  $VALIDATE_OUTPUT"
+else
+  fail "openenv validate failed"
+  printf "%s\n" "$VALIDATE_OUTPUT"
+  stop_at "Step 3"
+fi
+
+printf "\n"
+printf "${BOLD}========================================${NC}\n"
+printf "${GREEN}${BOLD}  All 3/3 checks passed!${NC}\n"
+printf "${GREEN}${BOLD}  Your submission is ready to submit.${NC}\n"
+printf "${BOLD}========================================${NC}\n"
 printf "\n"

server/__init__.py

@@ -1,11 +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"]
+# 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

@@ -1,79 +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()
+# 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

@@ -1,214 +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 copy import deepcopy
-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
-
-    @staticmethod
-    def _build_initial_kb() -> Dict[str, Dict[str, Any]]:
-        """Construct a fresh KB snapshot for each new episode/reset."""
-        return {
-            "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)},
-        }
-
-    def __init__(self):
-        self._state = State(episode_id=str(uuid4()), step_count=0)
-        self.kb = deepcopy(self._build_initial_kb())
-        
-        # 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)
-        self.kb = deepcopy(self._build_initial_kb())
-        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
+# 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 copy import deepcopy
+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)
+        self.kb = {}
+        self.test_suite = []
+        self._setup_task("easy")
+        
+    def _setup_task(self, task_id: str):
+        if task_id == "easy":
+            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": {}
+                },
+                **{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)}
+            }
+            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"}
+            ]
+        elif task_id == "medium":
+            self.kb = {
+                "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": {}
+                },
+                **{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)}
+            }
+            self.test_suite = [
+                {"query": "UI issues frontend CSS missing button", "target_concept": "frontend team fixed the button"},
+                {"query": "Email integration 401", "target_concept": "API key was rotated"}
+            ]
+        elif task_id == "hard":
+            self.kb = {
+                "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_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": {}
+                },
+                **{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)}
+            }
+            self.test_suite = [
+                {"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"}
+            ]
+        else:
+            self._setup_task("easy")
+
+    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, **kwargs) -> RagOptimizerObservation:
+        self._state = State(episode_id=str(uuid4()), step_count=0)
+        task_id = kwargs.get("task_id") or kwargs.get("task") or "easy"
+        self._setup_task(task_id)
+        
+        return RagOptimizerObservation(
+            message=f"RagOptimizerEnv Initialized for task: {task_id}. Resolve conflicts, append metadata, or splinter chunks 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

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