supportbench/server.py

"""FastAPI server exposing the SupportBench environment over HTTP."""

from __future__ import annotations

from typing import Any, Dict, Optional

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

from .env import SupportBenchEnv
from .models import Action
from .tasks import list_tasks, get_default_task_id

app = FastAPI(
    title="SupportBench",
    description="Policy-grounded customer support environment for AI agents",
    version="1.0.0",
)

# One global env instance (single-session; for multi-session use a session map)
_env = SupportBenchEnv()


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


class StepRequest(BaseModel):
    action: Action


@app.get("/health")
def health() -> Dict[str, str]:
    return {"status": "ok"}


@app.post("/reset")
def reset(req: ResetRequest = ResetRequest()) -> Dict[str, Any]:
    task_id = req.task_id or get_default_task_id()
    obs = _env.reset(task_id=task_id)
    return {"observation": obs.model_dump(), "task_id": task_id}


@app.post("/step")
def step(req: StepRequest) -> Dict[str, Any]:
    try:
        obs, reward, done, info = _env.step(req.action)
    except RuntimeError as e:
        raise HTTPException(status_code=400, detail=str(e))
    return {
        "observation": obs.model_dump(),
        "reward": reward.model_dump(),
        "done": done,
        "info": info,
    }


@app.get("/state")
def state() -> Dict[str, Any]:
    return _env.state()


@app.post("/close")
def close() -> Dict[str, Any]:
    return _env.close()


@app.get("/tasks")
def tasks() -> Dict[str, Any]:
    return {"tasks": list_tasks()}