Upload folder using huggingface_hub

inference.py

@@ -35,6 +35,7 @@ MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")
 HF_TOKEN = os.getenv("HF_TOKEN")
 
 MAX_STEPS = 30
+TASK_IDS = ["easy", "medium", "hard"]
 
 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.
@@ -64,117 +65,142 @@ def format_action_str(action: RagOptimizerAction) -> str:
         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
+
+def _safe_reset(env: RagOptimizerEnvClient, task_id: str):
+    """Reset env for a specific task with compatibility fallbacks."""
+    try:
+        return env.reset(task_id=task_id)
+    except TypeError:
+        try:
+            return env.reset(task=task_id)
+        except TypeError:
+            return env.reset()
+
+
+def _clamp_score(value: float) -> float:
+    if value < 0.01:
+        return 0.01
+    if value > 0.99:
+        return 0.99
+    return value
+
+
+def run_task_episode(
+    env: RagOptimizerEnvClient,
+    llm_client: OpenAI,
+    task_id: str,
+) -> None:
     step_rewards = []
     success = False
     error_msg = "null"
-    score = 0.0
-    
-    print(f"[START] task=rag_optimizer_env env=OpenEnv model={MODEL_NAME}")
-    
+    score = 0.01
+    step = 0
+
+    print(f"[START] task={task_id} 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 = _safe_reset(env, task_id)
+            observation = result.observation
+        except Exception as e:
+            error_msg = str(e).replace('\n', ' ')
+            print(f"[END] success=false steps=0 score=0.01 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)})
+
+    for i in range(1, MAX_STEPS + 1):
+        step = i
+        messages = list(history)
+
+        action_str = "unknown"
+        error_msg = "null"
+
+        try:
+            completion = llm_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.reset()
+                result = env.step(action)
                 observation = result.observation
+                reward = _clamp_score(float(result.reward))
             except Exception as e:
                 error_msg = str(e).replace('\n', ' ')
-                print(f"[END] success=false steps=0 score=0.00 rewards=")
-                return
+                reward = 0.01
+                result = type("obj", (object,), {"done": True})()
+                observation = type("obj", (object,), {"message": "error", "current_docs": {}})()
 
-        history = [{"role": "system", "content": SYSTEM_PROMPT}]
-        
-        init_obs = {
+        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
+            score = _clamp_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
+            "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)
+        history.append({"role": "user", "content": json.dumps(next_obs, indent=2)})
+    else:
+        # Reached max steps
+        success = False
+        score = _clamp_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}")
 
+def main():
+    # Setup OpenAI Client
+    client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN)
+
+    with RagOptimizerEnvClient(base_url="http://localhost:8000").sync() as env:
+        for task_id in TASK_IDS:
+            run_task_episode(env=env, llm_client=client, task_id=task_id)
+
 if __name__ == "__main__":
     main()

openenv.yaml

@@ -1,16 +1,30 @@
-spec_version: 1
-name: rag_optimizer
-type: space
-runtime: fastapi
-app: server.app:app
-port: 8000
+spec_version: "1.0"
+name: "rag_optimizer"
+environment:
+  name: "RAG Optimizer"
+  description: "RAG pipeline optimization benchmark. Agents resolve pricing conflicts, route metadata tags, and splinter monolithic content blobs."
+  version: "1.0.0"
+  dockerfile: "Dockerfile"
+  endpoints:
+    reset: "/reset"
+    step: "/step/{session_id}"
+    state: "/state/{session_id}"
 tasks:
   - id: easy
-    steps: 10
-    description: Conflict Resolution - Resolve overlapping pricing parameters by deleting the legacy pricing format.
+    name: "Conflict Resolution"
+    description: "Resolve overlapping pricing parameters by deleting the legacy pricing format."
+    reward_range: [0.1, 0.99]
+    grader: "rag_optimizer.server.grader:grade_easy"
   - id: medium
-    steps: 10
-    description: Ontological Tagging - Route exact programatic metadata tags to support tickets.
+    name: "Ontological Tagging"
+    description: "Route exact programmatic metadata tags to support tickets."
+    reward_range: [0.1, 0.99]
+    grader: "rag_optimizer.server.grader:grade_medium"
   - id: hard
-    steps: 15
-    description: Syntactic Splintering - Break down the monolithic onboarding blob into multiple granular chunks.
+    name: "Syntactic Splintering"
+    description: "Break down the monolithic onboarding blob into multiple granular chunks."
+    reward_range: [0.1, 0.99]
+    grader: "rag_optimizer.server.grader:grade_hard"
+agent:
+  inference_entrypoint: "inference.py"
+  requirements: "requirements.txt"

server/grader.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+from typing import Any
+
+
+def _clamp_score(value: Any) -> float:
+    try:
+        score = float(value)
+    except (TypeError, ValueError):
+        return 0.01
+    if score < 0.01:
+        return 0.01
+    if score > 0.99:
+        return 0.99
+    return score
+
+
+def _extract_score_from_trajectory(trajectory: Any) -> float:
+    if trajectory is None:
+        return 0.01
+
+    if isinstance(trajectory, (int, float)):
+        return _clamp_score(trajectory)
+
+    if isinstance(trajectory, dict):
+        for key in ("score", "reward", "final_score", "final_reward"):
+            if key in trajectory:
+                return _clamp_score(trajectory.get(key))
+
+        observation = trajectory.get("observation")
+        if isinstance(observation, dict):
+            for key in ("reward", "score"):
+                if key in observation:
+                    return _clamp_score(observation.get(key))
+
+        steps = trajectory.get("steps")
+        if isinstance(steps, list) and steps:
+            last_step = steps[-1]
+            if isinstance(last_step, dict):
+                for key in ("reward", "score"):
+                    if key in last_step:
+                        return _clamp_score(last_step.get(key))
+
+    if isinstance(trajectory, (list, tuple)) and trajectory:
+        return _extract_score_from_trajectory(trajectory[-1])
+
+    return 0.01
+
+
+def grade_easy(trajectory: Any = None) -> float:
+    return _extract_score_from_trajectory(trajectory)
+
+
+def grade_medium(trajectory: Any = None) -> float:
+    return _extract_score_from_trajectory(trajectory)
+
+
+def grade_hard(trajectory: Any = None) -> float:
+    return _extract_score_from_trajectory(trajectory)