Update server.py

server.py

@@ -1,129 +1,431 @@
 """
-CustomerSupportEnv — FastAPI server.
+server.py — FastAPI/OpenEnv server wrapper for CustomerSupportEnv.
+
+Exposes the environment as REST endpoints compatible with OpenEnv specification.
+Handles session management and action validation.
 
 Endpoints:
-  POST /reset              → Observation
-  POST /step               → StepResult
-  GET  /state              → Observation
-  GET  /tasks              → list of task specs
-  POST /grade              → GraderResult
-  GET  /health             → 200 OK
-  GET  /openenv.yaml       → spec file
+  POST   /reset              → Initialize new episode, return initial observation
+  POST   /step               → Apply action, return (obs, reward, done)
+  GET    /state              → Get current environment state
+  GET    /tasks              → List all tasks
+  POST   /grade              → Grade current episode
+  GET    /health             → Health check
+  GET    /openenv.yaml       → Spec file
 """
 from __future__ import annotations
 
+import json
 import os
 import sys
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-
-from typing import Optional
-from fastapi import FastAPI, HTTPException
-from fastapi.responses import FileResponse, JSONResponse
-from pydantic import BaseModel
-
-from env.environment import CustomerSupportEnv, TASKS
-from env.models import Action, Observation, StepResult, GraderResult
-from graders.graders import grade
+import traceback
+from typing import Any, Dict, Optional
+from pathlib import Path
+
+# FastAPI imports
+try:
+    from fastapi import FastAPI, HTTPException, Request
+    from fastapi.responses import FileResponse, JSONResponse
+    from pydantic import BaseModel
+    import uvicorn
+except ImportError as e:
+    print(f"[ERROR] Missing FastAPI dependency: {e}", flush=True)
+    print("Run: pip install fastapi uvicorn pydantic", flush=True)
+    sys.exit(1)
 
+# Local env imports
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+try:
+    from env.environment import CustomerSupportEnv, TASKS
+    from env.models import Action, ActionType, Observation, Reward
+    from graders.graders import grade
+except ImportError as e:
+    print(f"[ERROR] Missing local env module: {e}", flush=True)
+    traceback.print_exc()
+    sys.exit(1)
+
+# ── FastAPI App ──────────────────────────────────────────────────────────────
 app = FastAPI(
     title="CustomerSupportEnv",
-    description="OpenEnv-compatible RL environment for customer support agent training.",
-    version="1.0.0",
+    description="OpenEnv-compatible customer support RL environment",
+    version="1.0.0"
 )
 
-# One env instance per task (keyed by task_id)
-_envs: dict[str, CustomerSupportEnv] = {}
+# ── Session Storage (in-memory for single deployment) ───────────────────────
+_sessions: Dict[str, Dict[str, Any]] = {}
+_session_counter = 0
 
 
-def _get_env(task_id: str) -> CustomerSupportEnv:
-    if task_id not in TASKS:
-        raise HTTPException(status_code=404, detail=f"Unknown task_id: {task_id}")
-    if task_id not in _envs:
-        _envs[task_id] = CustomerSupportEnv(task_id=task_id)
-    return _envs[task_id]
+def new_session_id() -> str:
+    """Generate a unique session ID."""
+    global _session_counter
+    _session_counter += 1
+    return f"session_{_session_counter:06d}"
 
 
+# ── Pydantic Models ──────────────────────────────────────────────────────────
 class ResetRequest(BaseModel):
     task_id: str = "task_1"
+    seed: Optional[int] = None
 
 
 class StepRequest(BaseModel):
-    task_id: str = "task_1"
+    session_id: str
     action_type: str
     payload: Optional[str] = None
 
 
 class GradeRequest(BaseModel):
-    task_id: str
-
-
-@app.get("/health")
-def health():
-    return {"status": "ok", "version": CustomerSupportEnv.VERSION}
-
-
-@app.post("/reset", response_model=Observation)
-def reset(req: ResetRequest):
-    env = _get_env(req.task_id)
-    obs = env.reset()
-    return obs
-
-
-@app.post("/step", response_model=StepResult)
-def step(req: StepRequest):
-    env = _get_env(req.task_id)
+    session_id: str
+
+
+# ── Helper: Serialize observation/reward to JSON ──────────────────────────────
+def serialize_obs(obs: Observation) -> Dict[str, Any]:
+    """Convert Observation dataclass to JSON-serializable dict."""
+    return {
+        "ticket_id": obs.ticket_id,
+        "task_id": obs.task_id,
+        "status": obs.status,
+        "sentiment": obs.sentiment,
+        "priority": obs.priority,
+        "category": obs.category,
+        "turn": obs.turn,
+        "max_turns": obs.max_turns,
+        "history": obs.history,
+        "kb_results": obs.kb_results,
+        "kb_searched": obs.kb_searched,
+        "empathized": obs.empathized,
+        "clarified": obs.clarified,
+        "solution_offered": obs.solution_offered,
+        "escalated": obs.escalated,
+        "cumulative_reward": obs.cumulative_reward,
+        "done": obs.done,
+    }
+
+
+def serialize_reward(reward: Reward) -> Dict[str, Any]:
+    """Convert Reward dataclass to JSON-serializable dict."""
+    return {
+        "total": reward.total,
+        "breakdown": reward.breakdown,
+        "reason": reward.reason,
+    }
+
+
+# ── OpenEnv Endpoints ────────────────────────────────────────────────────────
+
+@app.post("/reset")
+async def reset(request: ResetRequest) -> JSONResponse:
+    """
+    Reset environment and start a new episode.
+    
+    Args:
+        task_id: One of task_1, task_2, task_3
+        seed: Optional random seed
+    
+    Returns:
+        {
+            "session_id": str,
+            "observation": {...},
+            "info": {...}
+        }
+    """
     try:
-        action = Action(action_type=req.action_type, payload=req.payload)
+        task_id = request.task_id
+        seed = request.seed or 42
+        
+        if task_id not in TASKS:
+            raise HTTPException(
+                status_code=400,
+                detail=f"Invalid task_id. Must be one of: {list(TASKS.keys())}"
+            )
+        
+        # Create new environment
+        env = CustomerSupportEnv(task_id=task_id, seed=seed)
+        obs = env.reset()
+        
+        # Store session
+        session_id = new_session_id()
+        _sessions[session_id] = {
+            "env": env,
+            "task_id": task_id,
+            "observation": obs,
+            "steps": 0,
+            "done": False,
+        }
+        
+        return JSONResponse(
+            status_code=200,
+            content={
+                "session_id": session_id,
+                "observation": serialize_obs(obs),
+                "info": {
+                    "task_id": task_id,
+                    "difficulty": TASKS[task_id].difficulty,
+                    "description": TASKS[task_id].description,
+                }
+            }
+        )
+    
+    except Exception as e:
+        traceback.print_exc()
+        raise HTTPException(status_code=500, detail=f"Reset failed: {str(e)}")
+
+
+@app.post("/step")
+async def step(request: StepRequest) -> JSONResponse:
+    """
+    Apply an action and step the environment.
+    
+    Args:
+        session_id: Session ID from /reset
+        action_type: One of [search_kb, empathize, ask_clarify, offer_solution, escalate, resolve, send_message]
+        payload: Optional action payload (required for some action types)
+    
+    Returns:
+        {
+            "observation": {...},
+            "reward": {...},
+            "done": bool,
+            "info": {...}
+        }
+    """
+    try:
+        session_id = request.session_id
+        action_type = request.action_type
+        payload = request.payload
+        
+        if session_id not in _sessions:
+            raise HTTPException(status_code=404, detail=f"Session not found: {session_id}")
+        
+        session = _sessions[session_id]
+        env = session["env"]
+        
+        if session["done"]:
+            raise HTTPException(status_code=400, detail="Episode already done. Call /reset to start new episode.")
+        
+        # Create action
+        action = Action(action_type=action_type, payload=payload)
+        
+        # Step environment
         result = env.step(action)
-        return result
-    except RuntimeError as e:
-        raise HTTPException(status_code=400, detail=str(e))
+        
+        # Update session
+        session["observation"] = result.observation
+        session["steps"] += 1
+        session["done"] = result.observation.done
+        
+        return JSONResponse(
+            status_code=200,
+            content={
+                "observation": serialize_obs(result.observation),
+                "reward": serialize_reward(result.reward),
+                "done": result.observation.done,
+                "info": {
+                    "step": session["steps"],
+                    "action": action_type,
+                }
+            }
+        )
+    
+    except HTTPException:
+        raise
     except Exception as e:
-        raise HTTPException(status_code=422, detail=str(e))
-
-
-@app.get("/state", response_model=Observation)
-def state(task_id: str = "task_1"):
-    env = _get_env(task_id)
-    return env.state()
+        traceback.print_exc()
+        raise HTTPException(status_code=500, detail=f"Step failed: {str(e)}")
+
+
+@app.get("/state")
+async def state_endpoint(session_id: str) -> JSONResponse:
+    """
+    Get current environment state without stepping.
+    
+    Args:
+        session_id: Session ID from /reset
+    
+    Returns:
+        {
+            "observation": {...},
+            "info": {...}
+        }
+    """
+    try:
+        if session_id not in _sessions:
+            raise HTTPException(status_code=404, detail=f"Session not found: {session_id}")
+        
+        session = _sessions[session_id]
+        obs = session["observation"]
+        
+        return JSONResponse(
+            status_code=200,
+            content={
+                "observation": serialize_obs(obs),
+                "info": {
+                    "task_id": session["task_id"],
+                    "steps": session["steps"],
+                    "done": session["done"],
+                }
+            }
+        )
+    
+    except HTTPException:
+        raise
+    except Exception as e:
+        traceback.print_exc()
+        raise HTTPException(status_code=500, detail=f"State query failed: {str(e)}")
 
 
 @app.get("/tasks")
-def list_tasks():
-    return {tid: spec.dict() for tid, spec in TASKS.items()}
+async def tasks_endpoint() -> JSONResponse:
+    """
+    List all available tasks.
+    
+    Returns:
+        {
+            "tasks": [
+                {
+                    "id": "task_1",
+                    "name": "...",
+                    "difficulty": "easy|medium|hard",
+                    "description": "...",
+                    "max_turns": int
+                },
+                ...
+            ]
+        }
+    """
+    try:
+        task_list = []
+        for task_id, task_obj in TASKS.items():
+            task_list.append({
+                "id": task_id,
+                "name": task_obj.name,
+                "difficulty": task_obj.difficulty,
+                "description": task_obj.description,
+                "max_turns": task_obj.max_turns,
+            })
+        
+        return JSONResponse(
+            status_code=200,
+            content={"tasks": task_list}
+        )
+    
+    except Exception as e:
+        traceback.print_exc()
+        raise HTTPException(status_code=500, detail=f"Tasks query failed: {str(e)}")
+
+
+@app.post("/grade")
+async def grade_endpoint(request: GradeRequest) -> JSONResponse:
+    """
+    Grade the current episode.
+    
+    Args:
+        session_id: Session ID from /reset
+    
+    Returns:
+        {
+            "score": float (0.0 to 1.0),
+            "passed": bool,
+            "breakdown": {...},
+            "reason": str
+        }
+    """
+    try:
+        session_id = request.session_id
+        
+        if session_id not in _sessions:
+            raise HTTPException(status_code=404, detail=f"Session not found: {session_id}")
+        
+        session = _sessions[session_id]
+        env = session["env"]
+        task_id = session["task_id"]
+        
+        # Get final state
+        final_obs = env.state()
+        
+        # Grade
+        grader_result = grade(task_id, final_obs)
+        
+        return JSONResponse(
+            status_code=200,
+            content={
+                "score": grader_result.score,
+                "passed": grader_result.passed,
+                "breakdown": grader_result.breakdown,
+                "reason": grader_result.reason,
+            }
+        )
+    
+    except HTTPException:
+        raise
+    except Exception as e:
+        traceback.print_exc()
+        raise HTTPException(status_code=500, detail=f"Grading failed: {str(e)}")
 
 
-@app.post("/grade", response_model=GraderResult)
-def grade_endpoint(req: GradeRequest):
-    env = _get_env(req.task_id)
-    obs = env.state()
-    result = grade(req.task_id, obs)
-    return result
+@app.get("/health")
+async def health() -> JSONResponse:
+    """Health check endpoint."""
+    return JSONResponse(
+        status_code=200,
+        content={
+            "status": "healthy",
+            "service": "CustomerSupportEnv",
+            "version": "1.0.0",
+            "sessions_active": len(_sessions),
+        }
+    )
 
 
 @app.get("/openenv.yaml")
-def get_yaml():
-    yaml_path = os.path.join(os.path.dirname(__file__), "openenv.yaml")
-    if os.path.exists(yaml_path):
-        return FileResponse(yaml_path, media_type="text/yaml")
-    return JSONResponse({"error": "openenv.yaml not found"}, status_code=404)
-
-
+async def openenv_spec() -> FileResponse:
+    """Serve OpenEnv specification."""
+    spec_path = Path(__file__).parent / "openenv.yaml"
+    if not spec_path.exists():
+        raise HTTPException(status_code=404, detail="openenv.yaml not found")
+    return FileResponse(spec_path, media_type="text/yaml")
+
+
+# ── Root endpoint ────────────────────────────────────────────────────────────
+@app.get("/")
+async def root() -> JSONResponse:
+    """Root endpoint."""
+    return JSONResponse(
+        status_code=200,
+        content={
+            "service": "CustomerSupportEnv OpenEnv Server",
+            "version": "1.0.0",
+            "endpoints": {
+                "POST /reset": "Initialize new episode",
+                "POST /step": "Apply action",
+                "GET /state": "Get current state",
+                "GET /tasks": "List tasks",
+                "POST /grade": "Grade episode",
+                "GET /health": "Health check",
+                "GET /openenv.yaml": "Specification",
+            }
+        }
+    )
+
+
+# ── Startup/Shutdown ─────────────────────────────────────────────────────────
+@app.on_event("startup")
+async def startup_event():
+    """Log startup."""
+    print("[INFO] CustomerSupportEnv server started", flush=True)
+
+
+@app.on_event("shutdown")
+async def shutdown_event():
+    """Log shutdown."""
+    print("[INFO] CustomerSupportEnv server shutdown", flush=True)
+
+
+# ── Main ─────────────────────────────────────────────────────────────────────
 if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run("server:app", host="0.0.0.0", port=7860, reload=False)
-
-from fastapi.responses import HTMLResponse
-
-@app.get("/", response_class=HTMLResponse)
-def home():
-    return """
-    <h1>Customer Support Env 🚀</h1>
-    <p>API is running successfully.</p>
-    <ul>
-        <li>/reset</li>
-        <li>/step</li>
-        <li>/state</li>
-        <li>/grade</li>
-    </ul>
-    """
+    port = int(os.environ.get("PORT", 7860))
+    host = os.environ.get("HOST", "0.0.0.0")
+    
+    print(f"[INFO] Starting server on {host}:{port}", flush=True)
+    uvicorn.run(app, host=host, port=port, log_level="info")