server.py

"""
FastAPI server for CustomerSupportTriage-v0.

Endpoints
---------
POST /reset          → ResetResponse
POST /step           → StepResult
GET  /state          → EnvState
GET  /health         → {"status": "ok"}
GET  /tasks          → list of available tasks
GET  /openenv.yaml   → spec file (served as text)
"""
from __future__ import annotations

import os
from pathlib import Path

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import PlainTextResponse

from env import CustomerSupportTriageEnv
from models import BatchTriageAction, EnvState, ResetRequest, ResetResponse, StepResult

app = FastAPI(
    title="CustomerSupportTriage-v0",
    description=(
        "OpenEnv environment: AI agent triages real-world customer support tickets "
        "by assigning priority, routing to department, and drafting replies."
    ),
    version="1.0.0",
)

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

# One env instance per server process (stateful; suitable for single-agent eval)
_env = CustomerSupportTriageEnv()


@app.get("/health")
async def health():
    return {"status": "ok", "env": CustomerSupportTriageEnv.ENV_NAME}


@app.post("/reset", response_model=ResetResponse)
async def reset(request: ResetRequest = None):
    """Reset the environment and return the initial observation."""
    try:
        return _env.reset(request)
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc))


@app.post("/step", response_model=StepResult)
async def step(action: BatchTriageAction):
    """Submit triage actions and advance the episode."""
    try:
        return _env.step(action)
    except RuntimeError as exc:
        raise HTTPException(status_code=400, detail=str(exc))
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc))


@app.get("/state", response_model=EnvState)
async def state():
    """Return a full snapshot of the current environment state."""
    return _env.state()


@app.get("/tasks")
async def list_tasks():
    """List available tasks with metadata."""
    return {
        "tasks": [
            {"name": "easy",   "tickets": 5,  "difficulty": "easy",   "description": "Clear-cut signals, unambiguous routing"},
            {"name": "medium", "tickets": 10, "difficulty": "medium", "description": "Multi-issue bodies, ambiguous routing"},
            {"name": "hard",   "tickets": 15, "difficulty": "hard",   "description": "Adversarial phrasing, legal edge-cases, misleading sentiment"},
        ]
    }


@app.get("/openenv.yaml", response_class=PlainTextResponse)
async def serve_openenv_yaml():
    path = Path(__file__).parent.parent / "openenv.yaml"
    if not path.exists():
        raise HTTPException(status_code=404, detail="openenv.yaml not found")
    return path.read_text()


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