server/app.py

"""
FastAPI server for the Customer Support Inbox OpenEnv environment.
Compatible with Hugging Face Spaces deployment.

Endpoints:
  GET  /           → Health check + info
  GET  /health     → Ping
  POST /reset      → Reset environment, returns Observation
  POST /step       → Take action, returns {observation, reward, done, info}
  GET  /state      → Current TicketState
  GET  /tasks      → List all task definitions
  GET  /tasks/{id} → Single task info
  GET  /knowledge/{key} → Knowledge base article
  GET  /summary    → Episode summary
  POST /validate   → Run openenv validate check
"""
from __future__ import annotations

import os
import uuid
from typing import Any, Dict, Optional

from fastapi import FastAPI, HTTPException, Request
from fastapi import Body
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from pydantic import BaseModel

from environment import CustomerSupportEnv
from environment.models import Action, Observation, Reward, TicketState
from environment.tasks import list_tasks, get_task

# ─── App setup ────────────────────────────────────────────────────────────────

app = FastAPI(
    title="Customer Support Inbox — OpenEnv",
    description=(
        "A real-world OpenEnv environment simulating a customer support inbox. "
        "Three tasks: Ticket Triage (easy), Guided Resolution (medium), "
        "VIP Retention (hard). Compatible with OpenEnv spec."
    ),
    version="1.0.0",
    docs_url="/docs",
    redoc_url="/redoc",
)

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

# ─── Session management ───────────────────────────────────────────────────────

# In-memory session store (for demo/single-instance use)
# In production, use Redis or database-backed sessions
_sessions: Dict[str, CustomerSupportEnv] = {}
_DEFAULT_SESSION = "default"


def _get_env(session_id: str = _DEFAULT_SESSION) -> CustomerSupportEnv:
    if session_id not in _sessions:
        _sessions[session_id] = CustomerSupportEnv()
    return _sessions[session_id]


# ─── Request models ───────────────────────────────────────────────────────────

class ResetRequest(BaseModel):
    task_id: str = "task1"
    ticket_id: Optional[str] = None
    seed: Optional[int] = None
    session_id: str = _DEFAULT_SESSION


class StepRequest(BaseModel):
    action: Action
    session_id: str = _DEFAULT_SESSION


class SessionRequest(BaseModel):
    session_id: str = _DEFAULT_SESSION


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

@app.get("/")
async def root():
    return {
        "name": "Customer Support Inbox — OpenEnv",
        "version": "1.0.0",
        "description": "Real-world customer support inbox simulation for agent training and evaluation.",
        "tasks": ["task1 (easy)", "task2 (medium)", "task3 (hard)"],
        "endpoints": {
            "reset": "POST /reset",
            "step": "POST /step",
            "state": "GET /state",
            "tasks": "GET /tasks",
            "health": "GET /health",
            "docs": "GET /docs",
        },
        "openenv_spec": "1.0",
        "tags": ["customer-support", "NLP", "multi-turn", "real-world"],
    }


@app.get("/health")
async def health():
    return {"status": "ok", "environment": "customer-support-inbox"}


@app.post("/reset", response_model=Observation)
async def reset(request: Optional[ResetRequest] = Body(default=None)):
    """
    Reset the environment for a new episode.
    Returns the initial Observation.
    """
    if request is None:
        request = ResetRequest()
    env = _get_env(request.session_id)
    try:
        obs = env.reset(
            task_id=request.task_id,
            ticket_id=request.ticket_id,
            seed=request.seed,
        )
        return obs
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))


@app.post("/step")
async def step(request: StepRequest):
    """
    Execute one action. Returns observation, reward, done, info.
    """
    env = _get_env(request.session_id)
    try:
        obs, reward, done, info = env.step(request.action)
        return {
            "observation": obs.model_dump(),
            "reward": reward.model_dump(),
            "done": done,
            "info": info,
        }
    except RuntimeError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Step error: {str(e)}")


@app.get("/state")
async def state(session_id: str = _DEFAULT_SESSION):
    """Return the current full TicketState."""
    env = _get_env(session_id)
    try:
        s = env.state()
        return s.model_dump()
    except RuntimeError as e:
        raise HTTPException(status_code=400, detail=str(e))


@app.get("/tasks")
async def get_tasks():
    """List all available tasks with descriptions and objectives."""
    return {"tasks": list_tasks()}


@app.get("/tasks/{task_id}")
async def get_single_task(task_id: str):
    """Get a specific task definition."""
    try:
        t = get_task(task_id)
        return {
            "task_id": t.task_id,
            "name": t.name,
            "difficulty": t.difficulty,
            "description": t.description,
            "objectives": t.objectives,
            "max_turns": t.max_turns,
            "min_score_to_pass": t.min_score_to_pass,
            "ticket_pool": t.ticket_pool,
        }
    except ValueError as e:
        raise HTTPException(status_code=404, detail=str(e))


@app.get("/knowledge/{key}")
async def knowledge_base(key: str, session_id: str = _DEFAULT_SESSION):
    """Look up a knowledge base article."""
    env = _get_env(session_id)
    article = env.get_knowledge_article(key)
    return {"key": key, "content": article}


@app.get("/knowledge")
async def list_knowledge():
    """List all available knowledge base keys."""
    from environment.data import KNOWLEDGE_BASE
    return {"keys": list(KNOWLEDGE_BASE.keys())}


@app.get("/summary")
async def episode_summary(session_id: str = _DEFAULT_SESSION):
    """Return episode summary statistics."""
    env = _get_env(session_id)
    return env.get_episode_summary()


@app.post("/validate")
async def validate():
    """
    OpenEnv validation endpoint.
    Runs a quick smoke test of all three tasks.
    """
    results = {}
    errors = []

    for task_id in ["task1", "task2", "task3"]:
        try:
            env = CustomerSupportEnv(seed=42)
            obs = env.reset(task_id=task_id, seed=42)

            # Validate Observation structure
            assert obs.ticket_id, "Missing ticket_id"
            assert obs.task_id == task_id, "task_id mismatch"
            assert obs.available_actions, "No available_actions"
            assert obs.task_description, "Missing task_description"

            # Take a classify action
            from environment.models import Action, ActionType, TicketCategory, TicketPriority
            action = Action(
                action_type=ActionType.CLASSIFY,
                category=TicketCategory.BILLING,
                priority=TicketPriority.HIGH,
            )
            step_obs, reward, done, info = env.step(action)

            # Validate Reward structure
            assert 0.0 <= reward.score <= 1.0, f"Reward out of range: {reward.score}"
            assert isinstance(done, bool), "done must be bool"

            # Validate state()
            s = env.state()
            assert s.ticket_id == obs.ticket_id

            results[task_id] = {
                "status": "pass",
                "obs_keys": list(obs.model_fields.keys()),
                "reward_score": reward.score,
                "done": done,
            }
        except Exception as e:
            results[task_id] = {"status": "fail", "error": str(e)}
            errors.append(f"{task_id}: {e}")

    return {
        "validation": "pass" if not errors else "fail",
        "errors": errors,
        "task_results": results,
        "spec_version": "1.0",
    }


@app.get("/sessions")
async def list_sessions():
    """List active session IDs."""
    return {"sessions": list(_sessions.keys()), "count": len(_sessions)}


@app.delete("/sessions/{session_id}")
async def delete_session(session_id: str):
    """Delete a session."""
    if session_id in _sessions:
        del _sessions[session_id]
        return {"deleted": session_id}
    raise HTTPException(status_code=404, detail="Session not found")


# ─── Error handlers ───────────────────────────────────────────────────────────

@app.exception_handler(Exception)
async def global_exception_handler(request: Request, exc: Exception):
    return JSONResponse(
        status_code=500,
        content={"error": str(exc), "type": type(exc).__name__},
    )


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

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


if __name__ == "__main__":
    main()