app.py

"""
FastAPI server for DevOpsEnv - Linux DevOps/SRE troubleshooting environment.

Endpoints:
---------
POST   /reset          Create a new episode
POST   /step           Advance the episode
GET    /state          Current episode state
GET    /tasks          List tasks and action schema
POST   /grader         Grade a finished episode
GET    /health         Liveness check
GET    /               Info / spec link
"""
from __future__ import annotations

import os
import uuid
import json
from typing import Any, Dict, List, Optional
from datetime import datetime

from fastapi import FastAPI, HTTPException, Query
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

import environment as env
from data import TASK_META
from models import (
    Action,
    Observation,
    State,
    StepResult,
    TaskInfo,
    Reward,
    GraderResponse,
)

app = FastAPI(
    title="DevOpsEnv",
    description="An OpenEnv-compliant Linux DevOps/SRE troubleshooting environment.",
    version="1.0.0",
    docs_url="/docs",
    redoc_url="/redoc",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


class ResetRequest(BaseModel):
    task_id: str


class StepRequest(BaseModel):
    episode_id: str
    action: Action


class GraderRequest(BaseModel):
    episode_id: str


# ---------------------------------------------------------------------------
# Endpoints
# ---------------------------------------------------------------------------

@app.get("/", tags=["meta"])
def root():
    return {
        "name": "DevOpsEnv",
        "version": "1.0.0",
        "description": "OpenEnv DevOps/SRE troubleshooting environment",
        "openenv_spec": "https://github.com/meta-pytorch/OpenEnv",
        "tasks": list(TASK_META.keys()),
        "endpoints": {
            "reset": "POST /reset",
            "step": "POST /step",
            "state": "GET /state?episode_id=...",
            "tasks": "GET /tasks",
            "grader": "POST /grader",
            "health": "GET /health",
            "docs": "GET /docs",
        },
    }


@app.get("/health", tags=["meta"])
def health():
    return {"status": "healthy", "timestamp": datetime.now().isoformat()}


@app.get("/tasks", tags=["tasks"])
def tasks():
    result = []
    for task_id, meta in TASK_META.items():
        result.append(
            TaskInfo(
                task_id=task_id,
                name=meta["name"],
                description=meta["description"],
                difficulty=meta["difficulty"],
                max_steps=meta["max_steps"],
            )
        )
    return result


@app.post("/reset", tags=["control"])
def reset(req: ResetRequest):
    try:
        obs = env.reset(req.task_id)
        return obs.model_dump()
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))


@app.post("/step", tags=["control"])
def step(req: StepRequest):
    try:
        result = env.step(req.episode_id, req.action)
        return result.model_dump()
    except (KeyError, ValueError) as e:
        raise HTTPException(status_code=400, detail=str(e))


@app.get("/state", tags=["observation"])
def state(episode_id: str = Query(...)):
    try:
        st = env.get_state(episode_id)
        return st.model_dump()
    except KeyError as e:
        raise HTTPException(status_code=404, detail=str(e))


@app.post("/grader", tags=["evaluation"])
def grader(req: GraderRequest):
    try:
        score, breakdown, feedback = env.grade(req.episode_id)
        return GraderResponse(
            episode_id=req.episode_id,
            task_id=env._EPISODES[req.episode_id]["task_id"],
            score=score,
            breakdown=breakdown,
            feedback=feedback,
        ).model_dump()
    except KeyError as e:
        raise HTTPException(status_code=404, detail=str(e))


if __name__ == "__main__":
    import uvicorn
    port = int(os.environ.get("PORT", 7860))
    uvicorn.run(app, host="0.0.0.0", port=port, workers=1)


def run_server() -> None:
    """Console-script entry point used by packaging validators."""
    import uvicorn

    port = int(os.environ.get("PORT", 7860))
    uvicorn.run(app, host="0.0.0.0", port=port, workers=1)