Spaces:
Sleeping
Sleeping
| """ | |
| FastAPI server for SupportOps Arena. | |
| Implements OpenEnv HTTP specification for remote environment access. | |
| """ | |
| import os | |
| import uuid | |
| import logging | |
| from datetime import datetime, timedelta | |
| from typing import Any, Dict | |
| from fastapi import FastAPI, HTTPException, Header, Request | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from fastapi.responses import HTMLResponse, FileResponse, JSONResponse | |
| from fastapi.staticfiles import StaticFiles | |
| from pydantic import BaseModel | |
| from env.environment import SupportOpsArena | |
| from env.state import EnvObservation, StepResult, TaskLevel | |
| from env.actions import ActionType | |
| from adversary.adversary import AdaptiveAdversary | |
| from baseline.baseline_agent import BaselineAgent | |
| from graders.programmatic import ProgrammaticGrader | |
| from graders.llm_grader import LLMGrader | |
| from graders.adversarial_grader import AdversarialGrader, aggregate_scores | |
| # Setup logging | |
| logging.basicConfig( | |
| level=logging.INFO, | |
| format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' | |
| ) | |
| logger = logging.getLogger(__name__) | |
| # Create FastAPI app | |
| app = FastAPI( | |
| title="SupportOps Arena", | |
| description="OpenEnv-compliant IT incident triage RL environment", | |
| version="1.0.0" | |
| ) | |
| # Add CORS middleware | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], | |
| allow_methods=["*"], | |
| allow_headers=["*"], | |
| ) | |
| # Session management | |
| sessions: Dict[str, Dict[str, Any]] = {} | |
| SESSION_TTL_SECONDS = 1800 # 30 minutes | |
| # Baseline scores cache | |
| baseline_scores: Dict[str, float] = {} | |
| # ─── Request/Response Models ────────────────────────────────── | |
| class ResetRequest(BaseModel): | |
| """Request to reset environment.""" | |
| task_level: str = "easy" | |
| seed: int | None = None | |
| class StepRequest(BaseModel): | |
| """Request to take a step.""" | |
| action_name: str | |
| rationale: str | None = None | |
| class GradeRequest(BaseModel): | |
| """Request to grade episode.""" | |
| use_llm: bool = True | |
| use_adversarial: bool = False | |
| # ─── Session Management ─────────────────────────────────────── | |
| def get_or_create_session(session_id: str | None = None) -> tuple[str, SupportOpsArena]: | |
| """Get existing session or create new one.""" | |
| # Clean up expired sessions | |
| now = datetime.utcnow() | |
| expired = [ | |
| sid for sid, data in sessions.items() | |
| if now - data["created_at"] > timedelta(seconds=SESSION_TTL_SECONDS) | |
| ] | |
| for sid in expired: | |
| del sessions[sid] | |
| logger.info(f"Expired session {sid}") | |
| # Create new session if needed | |
| if not session_id or session_id not in sessions: | |
| session_id = str(uuid.uuid4()) | |
| adversary = AdaptiveAdversary() | |
| env = SupportOpsArena(adversary=adversary) | |
| sessions[session_id] = { | |
| "env": env, | |
| "adversary": adversary, | |
| "created_at": now, | |
| "last_accessed": now, | |
| } | |
| logger.info(f"Created new session {session_id}") | |
| else: | |
| sessions[session_id]["last_accessed"] = now | |
| return session_id, sessions[session_id]["env"] | |
| # ─── API Endpoints ──────────────────────────────────────────── | |
| async def root(): | |
| """Serve dashboard UI.""" | |
| try: | |
| with open("app/static/index.html", "r") as f: | |
| return HTMLResponse(content=f.read()) | |
| except FileNotFoundError: | |
| return HTMLResponse( | |
| content="<h1>SupportOps Arena</h1><p>Dashboard coming soon...</p>" | |
| ) | |
| async def health(): | |
| """Health check endpoint.""" | |
| return { | |
| "status": "ok", | |
| "version": "1.0.0", | |
| "active_sessions": len(sessions) | |
| } | |
| async def reset( | |
| request: ResetRequest, | |
| x_session_id: str | None = Header(None) | |
| ) -> Dict[str, Any]: | |
| """ | |
| Reset environment and start new episode. | |
| Returns: | |
| Initial observation and session ID | |
| """ | |
| try: | |
| session_id, env = get_or_create_session(x_session_id) | |
| observation = await env.reset( | |
| task_level=request.task_level, | |
| seed=request.seed | |
| ) | |
| logger.info(f"Reset session {session_id} to {request.task_level}") | |
| return { | |
| "session_id": session_id, | |
| "observation": observation.model_dump() | |
| } | |
| except ValueError as e: | |
| raise HTTPException(status_code=400, detail=str(e)) | |
| except Exception as e: | |
| logger.error(f"Reset failed: {e}", exc_info=True) | |
| raise HTTPException(status_code=500, detail=f"Reset failed: {str(e)}") | |
| async def step( | |
| request: StepRequest, | |
| x_session_id: str | None = Header(None) | |
| ) -> Dict[str, Any]: | |
| """ | |
| Take a step in the environment. | |
| Returns: | |
| Step result with observation, reward, done, info | |
| """ | |
| try: | |
| if not x_session_id or x_session_id not in sessions: | |
| raise HTTPException( | |
| status_code=400, | |
| detail="No active session. Call /reset first." | |
| ) | |
| session_id, env = get_or_create_session(x_session_id) | |
| # Parse action | |
| try: | |
| action = ActionType(request.action_name) | |
| except ValueError: | |
| raise HTTPException( | |
| status_code=400, | |
| detail=f"Invalid action: {request.action_name}" | |
| ) | |
| # Take step | |
| result = await env.step(action) | |
| logger.debug( | |
| f"Session {session_id} step: {action.value} -> " | |
| f"reward={result.reward:.2f}, done={result.done}" | |
| ) | |
| return { | |
| "observation": result.observation.model_dump(), | |
| "reward": result.reward, | |
| "done": result.done, | |
| "info": result.info | |
| } | |
| except HTTPException: | |
| raise | |
| except Exception as e: | |
| logger.error(f"Step failed: {e}", exc_info=True) | |
| raise HTTPException(status_code=500, detail=f"Step failed: {str(e)}") | |
| async def get_state(x_session_id: str | None = Header(None)) -> Dict[str, Any]: | |
| """ | |
| Get full environment state (for graders). | |
| Returns: | |
| Complete EnvState including hidden information | |
| """ | |
| try: | |
| if not x_session_id or x_session_id not in sessions: | |
| raise HTTPException( | |
| status_code=400, | |
| detail="No active session" | |
| ) | |
| session_id, env = get_or_create_session(x_session_id) | |
| state = await env.state() | |
| return state.model_dump() | |
| except HTTPException: | |
| raise | |
| except Exception as e: | |
| logger.error(f"Get state failed: {e}", exc_info=True) | |
| raise HTTPException(status_code=500, detail=f"Get state failed: {str(e)}") | |
| async def grade_episode( | |
| request: GradeRequest, | |
| x_session_id: str | None = Header(None) | |
| ) -> Dict[str, Any]: | |
| """ | |
| Grade current episode. | |
| Returns: | |
| Grading results from requested graders | |
| """ | |
| try: | |
| if not x_session_id or x_session_id not in sessions: | |
| raise HTTPException( | |
| status_code=400, | |
| detail="No active session" | |
| ) | |
| session_id, env = get_or_create_session(x_session_id) | |
| state = await env.state() | |
| if not state.done: | |
| raise HTTPException( | |
| status_code=400, | |
| detail="Episode not complete. Cannot grade active episode." | |
| ) | |
| # Run programmatic grader | |
| prog_grader = ProgrammaticGrader() | |
| prog_result = await prog_grader.grade(state) | |
| results = { | |
| "programmatic": prog_result.model_dump() | |
| } | |
| # Run LLM grader if requested | |
| if request.use_llm: | |
| llm_grader = LLMGrader() | |
| llm_result = await llm_grader.grade(state) | |
| results["llm"] = llm_result.model_dump() | |
| # Run adversarial grader if requested | |
| if request.use_adversarial: | |
| adv_grader = AdversarialGrader() | |
| adversary = sessions[session_id].get("adversary") | |
| adv_score = await adv_grader.grade(state, adversary) | |
| results["adversarial"] = {"score": adv_score} | |
| # Calculate aggregate score | |
| aggregate = aggregate_scores( | |
| programmatic=prog_result.score, | |
| llm=results.get("llm", {}).get("overall"), | |
| adversarial=results.get("adversarial", {}).get("score"), | |
| task_level=state.task_level.value | |
| ) | |
| results["aggregate_score"] = aggregate | |
| return results | |
| except HTTPException: | |
| raise | |
| except Exception as e: | |
| logger.error(f"Grading failed: {e}", exc_info=True) | |
| raise HTTPException(status_code=500, detail=f"Grading failed: {str(e)}") | |
| async def list_tasks() -> Dict[str, Any]: | |
| """List all available tasks with metadata.""" | |
| return { | |
| "tasks": [ | |
| { | |
| "id": "easy", | |
| "name": "Wi-Fi Connectivity Failure", | |
| "description": "Basic network connectivity issues", | |
| "max_steps": 10, | |
| "scenarios": 6 | |
| }, | |
| { | |
| "id": "medium", | |
| "name": "VPN Access Failure", | |
| "description": "Remote access and authentication issues", | |
| "max_steps": 16, | |
| "scenarios": 6 | |
| }, | |
| { | |
| "id": "hard", | |
| "name": "Cross-System Access Failure", | |
| "description": "Complex multi-service failures with shared root causes", | |
| "max_steps": 24, | |
| "scenarios": 6 | |
| } | |
| ] | |
| } | |
| async def run_baseline() -> Dict[str, Any]: | |
| """ | |
| Run baseline agent on all three task levels. | |
| Returns: | |
| Scores for each task level | |
| """ | |
| try: | |
| logger.info("Starting baseline agent evaluation") | |
| agent = BaselineAgent() | |
| results = {} | |
| for task_level in ["easy", "medium", "hard"]: | |
| # Create fresh environment | |
| env = SupportOpsArena() | |
| # Run episode | |
| episode_result = await agent.run_episode(env, task_level) | |
| # Get final state for scoring | |
| state = await env.state() | |
| # Grade with programmatic grader | |
| grader = ProgrammaticGrader() | |
| grade_result = await grader.grade(state) | |
| results[task_level] = { | |
| "steps": episode_result["steps"], | |
| "reward": episode_result["cumulative_reward"], | |
| "grade_score": grade_result.score, | |
| "root_cause_identified": grade_result.root_cause_identified | |
| } | |
| # Cache score | |
| baseline_scores[task_level] = grade_result.score | |
| logger.info( | |
| f"Baseline {task_level}: " | |
| f"steps={episode_result['steps']}, " | |
| f"score={grade_result.score:.3f}" | |
| ) | |
| return { | |
| "results": results, | |
| "timestamp": datetime.utcnow().isoformat() | |
| } | |
| except Exception as e: | |
| logger.error(f"Baseline run failed: {e}", exc_info=True) | |
| raise HTTPException( | |
| status_code=500, | |
| detail=f"Baseline run failed: {str(e)}" | |
| ) | |
| async def get_scores() -> Dict[str, Any]: | |
| """Get cached baseline scores.""" | |
| if not baseline_scores: | |
| return { | |
| "scores": {}, | |
| "message": "No baseline scores available. Run /baseline/run first." | |
| } | |
| return { | |
| "scores": baseline_scores, | |
| "timestamp": datetime.utcnow().isoformat() | |
| } | |
| async def get_openenv_spec(): | |
| """Serve openenv.yaml specification file.""" | |
| try: | |
| return FileResponse("openenv.yaml", media_type="text/yaml") | |
| except FileNotFoundError: | |
| raise HTTPException(status_code=404, detail="openenv.yaml not found") | |
| # ─── Server Entry Point ─────────────────────────────────────── | |
| if __name__ == "__main__": | |
| import uvicorn | |
| host = os.environ.get("HOST", "0.0.0.0") | |
| port = int(os.environ.get("PORT", "7860")) | |
| logger.info(f"Starting SupportOps Arena server on {host}:{port}") | |
| uvicorn.run( | |
| "app.server:app", | |
| host=host, | |
| port=port, | |
| log_level="info" | |
| ) | |