app.py

"""FastAPI web app for Antern Bot: serves a chat UI and a /api/chat endpoint."""
from __future__ import annotations

import time
import uuid
from contextlib import asynccontextmanager

from fastapi import FastAPI, Header
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from pathlib import Path

import config
import monitoring
from agent import AnternBot

STATIC_DIR = Path(__file__).parent / "static"

# In-memory conversation store: session_id -> list of API message dicts.
# Fine for a single-process app; swap for Redis/DB if you scale out.
SESSIONS: dict[str, list[dict]] = {}
MAX_HISTORY_MESSAGES = 24  # cap stored turns to bound token growth

bot: AnternBot | None = None
INIT_ERROR: str | None = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    global bot, INIT_ERROR
    try:
        bot = AnternBot()  # introspects schema + checks API key at startup
        monitoring.log_event("startup", status="ready", model=bot.model)
    except Exception as exc:  # missing API key, DB unreachable, etc.
        INIT_ERROR = str(exc)
        monitoring.log_event("startup", status="failed", error=str(exc))
    yield


app = FastAPI(title="Antern Bot", lifespan=lifespan)

# Allow the widget to call the API from other web pages (configurable origins).
app.add_middleware(
    CORSMiddleware,
    allow_origins=config.ALLOWED_ORIGINS,
    allow_methods=["GET", "POST"],
    allow_headers=["*"],
)


class ChatRequest(BaseModel):
    message: str
    session_id: str | None = None


class ResetRequest(BaseModel):
    session_id: str


@app.get("/")
def index() -> FileResponse:
    return FileResponse(STATIC_DIR / "index.html")


@app.get("/metrics")
def metrics_dashboard() -> FileResponse:
    """Human-friendly metrics dashboard (fetches /api/metrics with ?token=...)."""
    return FileResponse(STATIC_DIR / "metrics.html")


@app.get("/api/health")
def health() -> JSONResponse:
    return JSONResponse(
        {"ready": bot is not None, "error": INIT_ERROR, "model": getattr(bot, "model", None)}
    )


@app.post("/api/chat")
def chat(req: ChatRequest) -> JSONResponse:
    request_id = uuid.uuid4().hex
    received = time.time()

    if bot is None:
        monitoring.log_event("api_failure", request_id=request_id,
                             reason="bot_not_initialised", detail=INIT_ERROR)
        return JSONResponse(
            {"error": f"Bot not initialised: {INIT_ERROR}"}, status_code=503
        )
    if not req.message.strip():
        return JSONResponse({"error": "Empty message."}, status_code=400)

    session_id = req.session_id or uuid.uuid4().hex
    history = SESSIONS.get(session_id, [])

    try:
        result = bot.chat(history, req.message)
    except Exception as exc:
        err = f"{type(exc).__name__}: {exc}"
        monitoring.log_request(
            request_id=request_id, session_id=session_id, question=req.message,
            answer="", latency_ms=(time.time() - received) * 1000,
            model=getattr(bot, "model", None), llm_calls=0, tokens=None,
            queries=[], presentation=None, error=err,
        )
        monitoring.log_event("api_failure", request_id=request_id,
                             session_id=session_id, error=err)
        return JSONResponse({"error": err}, status_code=500)

    SESSIONS[session_id] = result["messages"][-MAX_HISTORY_MESSAGES:]
    pres = result.get("presentation")
    usage = result.get("usage") or {}
    chat_cost = monitoring.cost(usage.get("prompt", 0), usage.get("completion", 0))

    monitoring.log_request(
        request_id=request_id, session_id=session_id, question=req.message,
        answer=result["answer"], latency_ms=(time.time() - received) * 1000,
        model=getattr(bot, "model", None), llm_calls=result.get("llm_calls", 0),
        tokens=result.get("usage"), queries=result["queries"],
        presentation=(pres or {}).get("format"), error=None,
    )

    return JSONResponse(
        {
            "session_id": session_id,
            "request_id": request_id,
            "answer": result["answer"],
            "queries": result["queries"],
            "presentation": pres,
            "tokens": usage.get("total", 0),
            "cost_usd": chat_cost,
        }
    )


@app.get("/api/metrics")
def metrics(
    x_metrics_token: str | None = Header(default=None),
    token: str | None = None,  # query param, for viewing in a browser
) -> JSONResponse:
    """Aggregate metrics (uptime, requests, errors, tokens, latency, cost, sessions).
    If METRICS_TOKEN is set, supply it via the X-Metrics-Token header OR a
    ?token=... query parameter (the query param is handy in a browser)."""
    supplied = x_metrics_token or token
    if config.METRICS_TOKEN and supplied != config.METRICS_TOKEN:
        return JSONResponse({"error": "Unauthorized"}, status_code=401)
    return JSONResponse(monitoring.get_metrics())


@app.get("/api/metrics/recent")
def metrics_recent(
    limit: int = 20,
    x_metrics_token: str | None = Header(default=None),
    token: str | None = None,
) -> JSONResponse:
    """Recent chats from the durable log (survives restarts). Token-protected."""
    supplied = x_metrics_token or token
    if config.METRICS_TOKEN and supplied != config.METRICS_TOKEN:
        return JSONResponse({"error": "Unauthorized"}, status_code=401)
    return JSONResponse({"recent": monitoring.recent_chats(min(max(limit, 1), 100))})


@app.post("/api/reset")
def reset(req: ResetRequest) -> JSONResponse:
    SESSIONS.pop(req.session_id, None)
    return JSONResponse({"ok": True})


# Serve any other static assets (none required, but handy for extension).
if STATIC_DIR.exists():
    app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")