fix: resolve openEnv reset response inconsistency and port mismatch

README.md

@@ -39,8 +39,8 @@ pip install -r requirements.txt
 ### Try it out
 
 ```bash
-openenv validate
-docker build -t llm-control .
+# Run the FastAPI server
+python app.py
 ```
 
 Quick Local Test Snippet:
@@ -48,12 +48,12 @@ Quick Local Test Snippet:
 import requests
 
 # Reset environment
-resp = requests.post("http://localhost:8000/reset", json={"task": "easy"})
-obs = resp.json()
+resp = requests.post("http://localhost:7860/reset", json={"task": "easy"})
+obs = resp.json()["observation"]
 print("Reset observation:", obs)
 
 # Take step
-resp = requests.post("http://localhost:8000/step", json={"action": {"action_type": "follow_prompt"}})
+resp = requests.post("http://localhost:7860/step", json={"action": {"action_type": "follow_prompt"}})
 print("Step result:", resp.json())
 ```
 

app.py

@@ -2,6 +2,7 @@ from fastapi import FastAPI, HTTPException
 from fastapi.responses import HTMLResponse
 from fastapi.staticfiles import StaticFiles
 from pydantic import BaseModel
+from typing import Dict, Any
 import sys
 import os
 
@@ -12,9 +13,13 @@ from server.llm_env import LLMEnv
 
 app = FastAPI(title="LLM Control OpenEnv")
 
-# In-memory store for environments per episode id and overall states
-envs = {}
-completed_episodes = {}
+# Global in-memory store for environments and session state
+envs: Dict[str, LLMEnv] = {}
+completed_episodes: Dict[str, Dict[str, Any]] = {}
+
+# Default global environment initialized with a reset state
+default_env = LLMEnv()
+default_env.reset()
 
 class ResetRequest(BaseModel):
     task: str = "easy"
@@ -38,20 +43,24 @@ async def serve_gui():
     except FileNotFoundError:
         return "GUI index.html not found. Check the root directory."
 
-@app.post("/reset", response_model=Observation)
+@app.post("/reset")
 async def reset(req: ResetRequest):
     if req.task not in ["easy", "medium", "hard"]:
         raise HTTPException(status_code=400, detail="Invalid task")
     
     env = LLMEnv(task=req.task)
     obs = env.reset()
-    envs[env.state.episode_id] = env
+    state = env.state
+    envs[state.episode_id] = env
     
     # Also set default env to the latest reset for easy single-agent testing
     global default_env
     default_env = env
     
-    return obs
+    return {
+        "observation": obs.model_dump(),
+        "state": state.model_dump()
+    }
 
 @app.post("/step")
 async def step(req: StepRequest):

index.html

@@ -241,17 +241,22 @@
 
         async function resetEnv() {
             log("Sending reset sequence...");
-            const res = await fetch('/reset', {
-                method: 'POST',
-                headers: { 'Content-Type': 'application/json' },
-                body: JSON.stringify({ task: "medium" })
-            });
-            const data = await res.json();
-            updateUI(data);
-            document.getElementById('reward-val').innerText = "0.00";
-            document.getElementById('reward-val').style.color = "var(--cyber-green)";
-            isDead = false;
-            log("System Reset Complete.");
+            try {
+                const res = await fetch('/reset', {
+                    method: 'POST',
+                    headers: { 'Content-Type': 'application/json' },
+                    body: JSON.stringify({ task: "medium" })
+                });
+                const data = await res.json();
+                updateUI(data.observation);
+                document.getElementById('reward-val').innerText = "0.00";
+                document.getElementById('reward-val').style.color = "var(--cyber-green)";
+                isDead = false;
+                log(`System Reset Complete. Episode ID: <span style="color:var(--cyber-blue); font-size: 0.7rem;">${data.state.episode_id}</span>`);
+            } catch (err) {
+                log("<span class='log-mi'>Error during reset. Check server console.</span>");
+                console.error(err);
+            }
         }
 
         async function takeAction(actionType) {

test_env.py

@@ -2,26 +2,47 @@ import requests
 import json
 
 def test():
-    print("Resetting the environment (easy task)...")
-    resp = requests.post("http://localhost:8000/reset", json={"task": "easy"})
-    obs = resp.json()
-    print("Initial Observation:")
-    print(json.dumps(obs, indent=2))
+    # Use the correct port 7860 as defined in app.py
+    base_url = "http://localhost:7860"
+    
+    print(f"Resetting the environment (easy task) via {base_url}/reset...")
+    try:
+        resp = requests.post(f"{base_url}/reset", json={"task": "easy"})
+        resp.raise_for_status()
+        data = resp.json()
+        
+        # New API returns a wrapped object: {"observation": ..., "state": ...}
+        obs = data["observation"]
+        episode_id = data["state"]["episode_id"]
+        
+        print(f"Initial Observation (Episode: {episode_id}):")
+        print(json.dumps(obs, indent=2))
+    except (requests.exceptions.RequestException, KeyError) as e:
+        print(f"Error during reset: {e}")
+        return
 
     print("\nTaking an action: 'follow_prompt'...")
-    resp = requests.post("http://localhost:8000/step", json={"action": {"action_type": "follow_prompt"}})
-    result = resp.json()
-    print("Step Result:")
-    print(json.dumps(result, indent=2))
-    
-    print("\nTaking an action: 'minor_hallucination'...")
-    resp = requests.post("http://localhost:8000/step", json={"action": {"action_type": "minor_hallucination"}})
-    result = resp.json()
-    print("Step Result:")
-    print(json.dumps(result, indent=2))
+    try:
+        resp = requests.post(f"{base_url}/step", json={
+            "action": {"action_type": "follow_prompt"},
+            "episode_id": episode_id
+        })
+        resp.raise_for_status()
+        result = resp.json()
+        print("Step Result:")
+        print(json.dumps(result, indent=2))
+        
+        print("\nTaking an action: 'minor_hallucination'...")
+        resp = requests.post(f"{base_url}/step", json={
+            "action": {"action_type": "minor_hallucination"},
+            "episode_id": episode_id
+        })
+        resp.raise_for_status()
+        result = resp.json()
+        print("Step Result:")
+        print(json.dumps(result, indent=2))
+    except Exception as e:
+        print(f"Error during steps: {e}")
 
 if __name__ == "__main__":
-    try:
-        test()
-    except requests.exceptions.ConnectionError:
-        print("Error: Could not connect to the API. Is the server running on http://localhost:8000?")
+    test()