server/app.py

from __future__ import annotations

import json
from contextlib import asynccontextmanager
from typing import Any

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

from dbre.environment import DBREAction, DBREEnvironment, DBREObservation


env_instance: DBREEnvironment | None = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    global env_instance
    env_instance = DBREEnvironment(config={"max_steps": 20, "latency_threshold_pct": 0.6})
    yield
    if env_instance and env_instance.db and env_instance.db.conn:
        env_instance.db.close()


app = FastAPI(title="DBRE API", lifespan=lifespan)

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


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


@app.post("/reset", response_model=dict)
async def reset_environment() -> dict[str, Any]:
    if not env_instance:
        raise HTTPException(status_code=503, detail="Environment not initialized")

    observation, info = env_instance.reset()

    return {
        "observation": observation.dict(),
        "info": info,
    }


@app.post("/step", response_model=dict)
async def step_environment(action_data: dict) -> dict[str, Any]:
    if not env_instance:
        raise HTTPException(status_code=503, detail="Environment not initialized")

    try:
        action = DBREAction(**action_data)
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Invalid action: {str(e)}")

    observation, reward, terminated, truncated, info = env_instance.step(action)

    return {
        "observation": observation.dict(),
        "reward": reward,
        "terminated": terminated,
        "truncated": truncated,
        "info": info,
    }


@app.get("/state", response_model=dict)
async def get_state() -> dict[str, Any]:
    if not env_instance:
        raise HTTPException(status_code=503, detail="Environment not initialized")

    observation = env_instance.state()

    return {
        "observation": observation.dict(),
    }


@app.get("/elo_history", response_model=dict)
async def get_elo_history() -> dict[str, Any]:
    if not env_instance:
        raise HTTPException(status_code=503, detail="Environment not initialized")

    history = env_instance.elo_tracker.get_elo_history()

    return {
        "history": history,
    }


@app.get("/current_playbook", response_model=dict)
async def get_current_playbook() -> dict[str, Any]:
    if not env_instance:
        raise HTTPException(status_code=503, detail="Environment not initialized")

    playbook = env_instance.playbook_manager.get_current()

    return {
        "playbook": playbook,
    }


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)


@app.post("/train")
def start_training():
    import subprocess
    subprocess.Popen(["python3", "train.py"])
    return {"status": "training started"}