server.py

"""
server.py — FastAPI server exposing the OpenEnv HTTP API.

Endpoints:
  POST /reset          Reset environment, begin new episode
  POST /step           Submit a triage action, get StepResult
  GET  /state          Get current environment state
  GET  /tasks          List all available tasks
  GET  /health         Health check
"""

import os
import traceback
from fastapi import FastAPI, HTTPException, Body
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, ValidationError
from typing import Optional

from env_models import TriageAction, StepResult, EnvState, TicketObservation
from env_core import ITSupportEnv

app = FastAPI(
    title="IT Support Triage — OpenEnv",
    description=(
        "An OpenEnv-compliant RL environment for training and evaluating agents "
        "on IT helpdesk ticket triage tasks. Includes 3 tasks (easy → medium → hard) "
        "with deterministic graders and safety-aware reward functions."
    ),
    version="1.0.0",
)

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

# One environment instance per server (stateful)
env = ITSupportEnv()


# ─── Request schemas ──────────────────────────────────────────────────────────

class ResetRequest(BaseModel):
    task_id: Optional[str] = "task_easy"


class StepRequest(BaseModel):
    action: TriageAction


# ─── Endpoints ────────────────────────────────────────────────────────────────

@app.get("/health")
def health():
    return {"status": "ok", "environment": "it-support-triage", "version": "1.0.0"}


@app.post("/reset", response_model=TicketObservation)
async def reset(request: Optional[ResetRequest] = Body(default=None)):
    """
    Reset the environment and return the initial observation (ticket).
    Pass task_id to select difficulty: task_easy | task_medium | task_hard
    """
    task_id = (request.task_id if request else None) or "task_easy"
    try:
        obs = env.reset(task_id)
        return obs
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=traceback.format_exc())


@app.post("/step", response_model=StepResult)
def step(request: StepRequest):
    """
    Submit a TriageAction and receive a StepResult with reward and grader breakdown.
    """
    try:
        result = env.step(request.action)
        return result
    except RuntimeError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except ValidationError as e:
        raise HTTPException(status_code=422, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=traceback.format_exc())


@app.get("/state", response_model=EnvState)
def state():
    """Return the full current environment state."""
    return env.state()


@app.get("/tasks")
def list_tasks():
    """List all available tasks with metadata."""
    return {"tasks": env.list_tasks()}


# ─── Entry point ──────────────────────────────────────────────────────────────

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