Vedi23/internet-diagnosis-env

1from fastapi import FastAPI, HTTPException2from fastapi.middleware.cors import CORSMiddleware3from models import Action4from environment import InternetDiagnosisEnvironment5from tasks import get_all_task_ids6 7# ─────────────────────────────────────────────8# What is this file?9# This is the FRONT DOOR of your environment.10# It turns your Python code into a web API.11# Judges, AI agents, and validators will12# talk to your environment through this file.13# ─────────────────────────────────────────────14 15# ── Create the FastAPI app ──16app = FastAPI(17    title       = "Internet Outage Diagnosis Environment",18    description = (19        "An OpenEnv environment where an AI agent diagnoses "20        "real-world internet connectivity failures. "21        "The agent observes network diagnostic data and must "22        "identify the root cause of the outage."23    ),24    version     = "1.0.0"25)26 27# ── Allow all origins (needed for HuggingFace) ──28app.add_middleware(29    CORSMiddleware,30    allow_origins     = ["*"],31    allow_methods     = ["*"],32    allow_headers     = ["*"],33)34 35# ── Create one environment instance ──36env = InternetDiagnosisEnvironment()37 38 39# ─────────────────────────────────────────────40# ROUTE 1: Health Check41# Judges ping this first to check if server works42# ─────────────────────────────────────────────43 44@app.get("/")45def root():46    return {47        "status"     : "ok",48        "environment": "Internet Outage Diagnosis",49        "version"    : "1.0.0",50        "endpoints"  : ["/reset", "/step", "/state", "/tasks"]51    }52 53 54# ─────────────────────────────────────────────55# ROUTE 2: Reset56# Starts a fresh episode57# ─────────────────────────────────────────────58 59@app.post("/reset")60def reset(task_id: str = "task_1_easy"):61    """62    Starts a new episode.63    Returns the first observation (network diagnostic data).64 65    task_id options:66    - task_1_easy67    - task_2_medium68    - task_3_hard69    """70    try:71        observation = env.reset(task_id=task_id)72        return {73            "observation" : observation.model_dump(),74            "task_id"     : task_id,75            "message"     : "Episode started. Analyze the network data and diagnose."76        }77    except ValueError as e:78        raise HTTPException(status_code=400, detail=str(e))79 80 81# ─────────────────────────────────────────────82# ROUTE 3: Step83# AI submits its action/diagnosis84# ─────────────────────────────────────────────85 86@app.post("/step")87def step(action: Action):88    """89    AI submits its diagnosis.90    Returns: observation, reward, done, info.91 92    Example action:93    {94        "diagnosis"         : "diagnose_router",95        "failing_component" : "router",96        "suggested_fix"     : "Restart the router",97        "confidence"        : 0.998    }99    """100    try:101        observation, reward, done, info = env.step(action)102        return {103            "observation" : observation.model_dump(),104            "reward"      : reward.model_dump(),105            "done"        : done,106            "info"        : info107        }108    except RuntimeError as e:109        raise HTTPException(status_code=400, detail=str(e))110 111 112# ─────────────────────────────────────────────113# ROUTE 4: State114# Returns current state of environment115# ─────────────────────────────────────────────116 117@app.get("/state")118def state():119    """120    Returns the full current state of the environment.121    Useful for debugging and inspection.122    """123    return env.state()124 125 126# ─────────────────────────────────────────────127# ROUTE 5: Tasks128# Lists all available tasks129# ─────────────────────────────────────────────130 131@app.get("/tasks")132def tasks():133    """134    Lists all available tasks with descriptions.135    """136    from tasks import TASKS137    return {138        "total_tasks" : len(TASKS),139        "tasks"       : {140            task_id: {141                "name"       : task["name"],142                "description": task["description"],143                "difficulty" : task["difficulty"],144                "max_steps"  : task["max_steps"],145                "pass_score" : task["pass_score"],146                "scenarios"  : task["scenarios"]147            }148            for task_id, task in TASKS.items()149        }150    }151 152 153# ─────────────────────────────────────────────154# Run the server155# ─────────────────────────────────────────────156 157if __name__ == "__main__":158    import uvicorn159    uvicorn.run(app, host="0.0.0.0", port=7860)