src/secrag/api/app.py

"""FastAPI service.

Design notes worth stating, because they are the parts that would otherwise
look arbitrary:

- The engine is built once in the lifespan handler and warmed before the server
  accepts traffic. Lazy loading would push a 5 second model initialisation onto
  whichever user arrives first.
- Retrieval and embedding are synchronous CPU work. Running them directly in
  the event loop would stall every other in-flight request, so the blocking
  path is dispatched to a worker thread.
- Errors are translated in one exception handler from the typed hierarchy in
  core.errors, so no route constructs a status code by hand.
"""

from __future__ import annotations

import asyncio
import contextlib
import json
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Any

from fastapi import APIRouter, FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import HTMLResponse, JSONResponse, PlainTextResponse, StreamingResponse

from secrag import __version__
from secrag.core.config import get_settings
from secrag.core.errors import SecRagError
from secrag.core.logging import configure_logging, get_logger
from secrag.core.types import QueryRequest, QueryResponse
from secrag.engine import QueryEngine
from secrag.observability.metrics import REGISTRY

log = get_logger(__name__)

router = APIRouter()

UI_DIR = Path(__file__).resolve().parents[3] / "ui"


@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncIterator[None]:
    settings = get_settings()
    configure_logging(settings.log_level, json_output=settings.log_json)

    engine = QueryEngine(settings)
    app.state.engine = engine
    app.state.ready = False

    # Warmup is blocking model initialisation, so it runs in a thread to keep
    # the health endpoint responsive while it happens.
    async def warm() -> None:
        try:
            await asyncio.to_thread(engine.warmup)
            app.state.ready = True
            REGISTRY.gauge("secrag_corpus_chunks", engine.retriever.corpus_size)
            log.info("service_ready", version=__version__)
        except Exception as exc:
            log.warning("warmup_failed", error=str(exc))
            app.state.ready = True

    task = asyncio.create_task(warm())
    try:
        yield
    finally:
        # Wait for warmup rather than cancelling it. Warmup runs inside
        # asyncio.to_thread, and cancelling the awaiting coroutine does not
        # stop the thread: it carries on and reopens the vector store after
        # shutdown has closed it. The index lock is then held by a thread
        # nobody is waiting on, and the next startup fails with
        # "already accessed by another instance".
        with contextlib.suppress(TimeoutError, asyncio.CancelledError, Exception):
            await asyncio.wait_for(task, timeout=60.0)
        await engine.aclose()


def create_app() -> FastAPI:
    """Build the application.

    A factory rather than a bare module-level object so configuration is read
    when the app is constructed instead of when the module is first imported.
    Import-time configuration is invisible to tests and to anything that sets
    environment variables after import.
    """
    application = FastAPI(
        title="sec-rag-platform",
        version=__version__,
        description=(
            "Evaluation-driven retrieval-augmented generation over SEC filings. "
            "Hybrid retrieval, cross-encoder reranking, XBRL-verified figures, "
            "and groundedness verification on every answer."
        ),
        lifespan=lifespan,
    )
    application.add_middleware(
        CORSMiddleware,
        allow_origins=get_settings().cors_origin_list,
        allow_credentials=False,
        allow_methods=["GET", "POST"],
        allow_headers=["*"],
    )
    application.include_router(router)
    application.add_exception_handler(SecRagError, handle_secrag_error)  # type: ignore[arg-type]
    return application


def engine_of(request: Request) -> QueryEngine:
    return request.app.state.engine  # type: ignore[no-any-return]


async def handle_secrag_error(_: Request, exc: SecRagError) -> JSONResponse:
    REGISTRY.increment("secrag_errors_total", code=exc.code)
    log.warning("request_failed", code=exc.code, message=exc.message)
    return JSONResponse(status_code=exc.status_code, content=exc.to_dict())


# ---------------------------------------------------------------------------
# Health and introspection
# ---------------------------------------------------------------------------


@router.get("/health", tags=["ops"])
async def health() -> dict[str, str]:
    """Liveness. Answers even while models are still loading."""
    return {"status": "ok", "version": __version__}


@router.get("/ready", tags=["ops"])
async def ready(request: Request) -> JSONResponse:
    """Readiness. Only true once models are warm and the corpus is loaded."""
    is_ready = bool(getattr(request.app.state, "ready", False))
    engine = engine_of(request)
    return JSONResponse(
        status_code=200 if is_ready else 503,
        content={
            "ready": is_ready,
            "corpus_chunks": engine.retriever.corpus_size,
            "bm25_documents": engine.retriever.bm25.size,
        },
    )


@router.get("/v1/stats", tags=["ops"])
async def stats(request: Request) -> dict[str, Any]:
    """Everything the service knows about its own configuration and corpus."""
    return engine_of(request).stats()


@router.get("/metrics", response_class=PlainTextResponse, tags=["ops"])
async def metrics() -> str:
    """Prometheus text exposition."""
    return REGISTRY.render_prometheus()


@router.get("/v1/metrics", tags=["ops"])
async def metrics_json() -> dict[str, object]:
    return REGISTRY.snapshot()


# ---------------------------------------------------------------------------
# Query
# ---------------------------------------------------------------------------


@router.post("/v1/query", response_model=QueryResponse, tags=["query"])
async def query(request: Request, payload: QueryRequest) -> QueryResponse:
    """Answer a question with citations, verified figures, and a groundedness score."""
    engine = engine_of(request)
    REGISTRY.increment("secrag_queries_total")

    response = await engine.answer(payload)

    REGISTRY.observe("secrag_query_latency_ms", response.latency_ms)
    REGISTRY.increment("secrag_answers_total", status=response.answer.status.value)
    if response.cached:
        REGISTRY.increment("secrag_cache_hits_total")
    if response.route:
        REGISTRY.increment("secrag_routes_total", intent=response.route.intent.value)
    return response


@router.post("/v1/query/stream", tags=["query"])
async def query_stream(request: Request, payload: QueryRequest) -> StreamingResponse:
    """Stream an answer as server-sent events.

    Sources and the routing decision are emitted before the first token, so the
    client can render provenance while the answer is still being written.
    """
    engine = engine_of(request)
    REGISTRY.increment("secrag_queries_total")
    REGISTRY.increment("secrag_stream_requests_total")

    async def event_source() -> AsyncIterator[str]:
        try:
            async for event in engine.stream(payload):
                yield f"event: {event['event']}\ndata: {json.dumps(event['data'])}\n\n"
        except SecRagError as exc:
            yield f"event: error\ndata: {json.dumps(exc.to_dict())}\n\n"
        except Exception as exc:
            log.warning("stream_failed", error=str(exc))
            payload_error = {"code": "stream_error", "message": str(exc)}
            yield f"event: error\ndata: {json.dumps(payload_error)}\n\n"

    return StreamingResponse(
        event_source(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            # Without this, nginx-style proxies buffer the whole response and
            # the stream arrives as one lump, defeating the point.
            "X-Accel-Buffering": "no",
        },
    )


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------


@router.get("/", response_class=HTMLResponse, include_in_schema=False)
async def index() -> HTMLResponse:
    path = UI_DIR / "index.html"
    if not path.exists():
        return HTMLResponse(
            "<h1>sec-rag-platform</h1><p>API is running. See <a href='/docs'>/docs</a>.</p>"
        )
    return HTMLResponse(path.read_text(encoding="utf-8"))


# ---------------------------------------------------------------------------
# Module-level instance for `uvicorn secrag.api.app:app`.
#
# Constructed at the bottom of the module, after every handler it references
# has been defined. Building it immediately after the factory raises NameError
# on the exception handler declared further down.
# ---------------------------------------------------------------------------

app = create_app()