server/app.py

"""
FastAPI server for Customer Support Email Triage Environment.
Exposes OpenEnv-compliant API endpoints.
"""

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from typing import Dict, Any
import sys
import os

# Add parent directory to path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from models import EmailAction, EmailObservation, EmailState, StepReturn, ResetReturn
from .environment import CustomerSupportEnv

# Initialize FastAPI app
app = FastAPI(
    title="Customer Support Email Triage Environment",
    description="OpenEnv-compliant environment for email classification and response generation",
    version="1.0.0"
)

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

# Initialize environment
env = CustomerSupportEnv()


@app.get("/health")
def health_check() -> Dict[str, str]:
    """
    Health check endpoint.
    
    Returns:
        Status indicator
    """
    return {"status": "healthy"}


@app.get("/info")
def info() -> Dict[str, Any]:
    """
    Get environment information.
    
    Returns:
        Environment metadata
    """
    return {
        "name": "customer_support_env",
        "version": "1.0.0",
        "description": "Customer Support Email Triage and Response System",
        "action_space": "EmailAction",
        "observation_space": "EmailObservation",
        "reward_range": [0.0, 1.0],
        "tasks": 12,
        "episode_type": "multi-step"
    }


@app.post("/reset", response_model=ResetReturn)
def reset() -> ResetReturn:
    """
    Reset the environment and return initial observation.
    
    Returns:
        Dict with observation and info
    """
    try:
        result = env.reset()
        return ResetReturn(
            observation=result["observation"],
            info=result["info"]
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/step", response_model=StepReturn)
def step(action: EmailAction) -> StepReturn:
    """
    Execute one step in the environment.
    
    Args:
        action: EmailAction with category, priority, response
    
    Returns:
        Dict with observation, reward, done, info
    """
    try:
        result = env.step(action)
        return StepReturn(
            observation=result["observation"],
            reward=result["reward"],
            done=result["done"],
            info=result["info"],
            step_reward_breakdown=result.get("step_reward_breakdown", {})
        )
    except RuntimeError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/state")
def get_state() -> Dict[str, Any]:
    """
    Get current environment state.
    
    Returns:
        Current state dictionary
    """
    try:
        return env.get_state()
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/stats")
def get_stats() -> Dict[str, Any]:
    """
    Get environment statistics.
    
    Returns:
        Statistics dictionary
    """
    try:
        return env.get_stats()
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/")
def root() -> Dict[str, str]:
    """
    Root endpoint with API documentation link.
    
    Returns:
        API info
    """
    return {
        "name": "Customer Support Email Triage Environment",
        "version": "1.0.0",
        "docs": "/docs",
        "openapi": "/openapi.json"
    }


def main():
    """Main entry point for running the server."""
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=5001)


if __name__ == "__main__":
    main()