| """FastAPI surface for the agent.
|
|
|
| `AGENT_MODE=mock` (default) replays canned traces so the UI can be developed
|
| without a GPU. `AGENT_MODE=live` runs the LangGraph supervisor. The event
|
| contract is the same either way; only the generator behind StreamingResponse
|
| changes.
|
| """
|
|
|
| from __future__ import annotations
|
|
|
| import asyncio
|
| from contextlib import asynccontextmanager
|
| from pathlib import Path
|
| from typing import AsyncIterator
|
|
|
| from fastapi import FastAPI, Query
|
| from fastapi.middleware.cors import CORSMiddleware
|
| from fastapi.responses import StreamingResponse
|
| from fastapi.staticfiles import StaticFiles
|
|
|
| from .config import get_settings
|
| from .events import ChatRequest
|
| from .mock_agent import replay
|
| from .stream import live_stream
|
|
|
| SSE_HEADERS = {
|
| "Cache-Control": "no-cache, no-transform",
|
| "Connection": "keep-alive",
|
| "X-Accel-Buffering": "no",
|
| }
|
|
|
|
|
| def _warmup() -> None:
|
| """Load embedding + reranker weights and open stores before the first query.
|
|
|
| The first Librarian call otherwise spends ~15s downloading/loading models,
|
| which on Fargate is paid by whoever happens to ask first.
|
| """
|
| from .agent.embeddings import get_embeddings
|
| from .agent.resources import get_qdrant, get_reranker, get_sql_database
|
|
|
| get_embeddings()
|
| get_reranker()
|
| get_qdrant()
|
| get_sql_database()
|
|
|
|
|
| @asynccontextmanager
|
| async def lifespan(_app: FastAPI):
|
| settings = get_settings()
|
| if settings.agent_mode == "live" and settings.warmup_on_start:
|
| await asyncio.to_thread(_warmup)
|
| yield
|
|
|
|
|
| app = FastAPI(title="FinRAG API", version="0.2.0", lifespan=lifespan)
|
|
|
| app.add_middleware(
|
| CORSMiddleware,
|
| allow_origins=["http://localhost:5173", "http://127.0.0.1:5173"],
|
| allow_credentials=True,
|
| allow_methods=["*"],
|
| allow_headers=["*"],
|
| )
|
|
|
|
|
| @app.get("/api/health")
|
| async def health() -> dict[str, object]:
|
| settings = get_settings()
|
| return {
|
| "status": "ok",
|
| "agent": settings.agent_mode,
|
| "vllm": settings.vllm_base_url,
|
| "embeddings": settings.embedding_provider,
|
| "qdrant": str(settings.qdrant_url or settings.qdrant_path),
|
| }
|
|
|
|
|
| @app.post("/api/chat/stream")
|
| async def chat_stream(
|
| body: ChatRequest,
|
| speed: float = Query(1.0, gt=0, le=50, description="Mock replay speed. Ignored in live mode."),
|
| ) -> StreamingResponse:
|
| """Stream agent state and answer tokens as Server-Sent Events.
|
|
|
| POST rather than GET because the query goes in the body and, once auth
|
| lands, the JWT goes in an Authorization header. That rules out `EventSource`
|
| on the client — it is GET-only and cannot set headers.
|
| """
|
| settings = get_settings()
|
|
|
| async def stream() -> AsyncIterator[bytes]:
|
| frames = (
|
| live_stream(body.query, body.conversation_id)
|
| if settings.agent_mode == "live"
|
| else replay(body.query, body.conversation_id, speed)
|
| )
|
| async for frame in frames:
|
| yield frame.encode("utf-8")
|
|
|
| return StreamingResponse(stream(), media_type="text/event-stream", headers=SSE_HEADERS)
|
|
|
|
|
| _dist = Path(__file__).resolve().parents[2] / "frontend" / "dist"
|
| if _dist.is_dir():
|
| app.mount("/", StaticFiles(directory=str(_dist), html=True), name="spa")
|
|
|
