app.py

"""Local FastAPI app for a pharmacovigilance demo with managed ChatKit sessions."""

from __future__ import annotations

import json
import os
import uuid
from pathlib import Path
from typing import Any, Mapping

import httpx
from dotenv import load_dotenv
from fastapi import FastAPI, Request
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles

load_dotenv()

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

DEFAULT_OPENAI_BASE = "https://api.openai.com"
SESSION_COOKIE_NAME = "chatkit_user_id"
SESSION_COOKIE_MAX_AGE_SECONDS = 60 * 60 * 24 * 30  # 30 days
OPENAI_BETA_HEADER = "chatkit_beta=v1"

app = FastAPI(title="Pharmacovigilance MIT ChatKit Demo")
app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")


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


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


@app.post("/api/chatkit/session")
async def create_chatkit_session(request: Request) -> JSONResponse:
    api_key = os.getenv("OPENAI_API_KEY", "").strip()
    workflow_id = os.getenv("WORKFLOW_ID", "").strip()

    missing_env = [
        name
        for name, value in [
            ("OPENAI_API_KEY", api_key),
            ("WORKFLOW_ID", workflow_id),
        ]
        if not value
    ]
    if missing_env:
        return response(
            {"error": f"Missing required environment variable(s): {', '.join(missing_env)}"},
            500,
        )

    user_id, cookie_value = resolve_user_id(request.cookies)
    openai_base = os.getenv("OPENAI_BASE_URL", DEFAULT_OPENAI_BASE).strip()

    payload = {
        "workflow": {"id": workflow_id},
        "user": user_id,
    }

    try:
        upstream = await create_openai_session(api_key, openai_base, payload)
    except httpx.RequestError as exc:
        return response({"error": f"Failed to reach OpenAI API: {exc}"}, 502, cookie_value)

    parsed = safe_json(upstream)
    workflow_version_used: str | None = None

    if not upstream.is_success:
        message = "Failed to create ChatKit session"
        if isinstance(parsed.get("error"), str):
            message = parsed["error"]
        elif isinstance(parsed.get("error"), Mapping):
            message = str(parsed["error"].get("message") or message)
        return response({"error": message}, upstream.status_code, cookie_value)

    client_secret = parsed.get("client_secret")
    expires_after = parsed.get("expires_after")
    if not isinstance(client_secret, str) or not client_secret.strip():
        return response(
            {"error": "Missing client_secret in OpenAI response"},
            502,
            cookie_value,
        )

    return response(
        {
            "client_secret": client_secret,
            "expires_after": expires_after,
            "workflow_version_used": workflow_version_used,
        },
        200,
        cookie_value,
    )


def resolve_user_id(cookies: Mapping[str, str]) -> tuple[str, str | None]:
    existing = cookies.get(SESSION_COOKIE_NAME)
    if existing and existing.strip():
        return existing.strip(), None

    new_user_id = str(uuid.uuid4())
    return new_user_id, new_user_id


def safe_json(http_response: httpx.Response) -> dict[str, Any]:
    try:
        parsed = http_response.json()
    except (json.JSONDecodeError, httpx.DecodingError):
        return {}
    return parsed if isinstance(parsed, dict) else {}


async def create_openai_session(
    api_key: str, openai_base: str, payload: Mapping[str, Any]
) -> httpx.Response:
    async with httpx.AsyncClient(base_url=openai_base, timeout=20.0) as client:
        return await client.post(
            "/v1/chatkit/sessions",
            headers={
                "Authorization": f"Bearer {api_key}",
                "OpenAI-Beta": OPENAI_BETA_HEADER,
                "Content-Type": "application/json",
            },
            json=payload,
        )


def is_production() -> bool:
    env = (os.getenv("ENVIRONMENT") or os.getenv("NODE_ENV") or "").strip().lower()
    return env == "production"


def response(
    payload: Mapping[str, Any],
    status_code: int,
    cookie_value: str | None = None,
) -> JSONResponse:
    res = JSONResponse(payload, status_code=status_code)
    if cookie_value:
        res.set_cookie(
            key=SESSION_COOKIE_NAME,
            value=cookie_value,
            max_age=SESSION_COOKIE_MAX_AGE_SECONDS,
            httponly=True,
            samesite="lax",
            secure=is_production(),
            path="/",
        )
    return res