| """HTTP service around an ``FGAORAG`` instance (FastAPI). |
| |
| Endpoints |
| |
| GET /health liveness |
| GET /ready readiness |
| GET /v1/info version, sources, costs, budget, controller configuration |
| POST /v1/answer {"question", "budget"?, "seed"?, "trace"?, "documents"?, "web_documents"?} |
| -> RAGResult as JSON; with "documents", FGAO answers from those passages |
| POST /v1/answer/stream same body; Server-Sent Events: "step" per decision, then "result" |
| POST /v1/batch {"questions": [...], ...} -> {"results": [...]} |
| GET /metrics Prometheus metrics |
| GET /docs, /openapi.json interactive API docs |
| |
| Operational behaviour |
| |
| * Admission control: at most ``max_inflight`` queries run at once; up to |
| ``max_queue`` more wait; beyond that the server answers 429 with Retry-After. |
| * Each query has a deadline (``request_timeout`` s): 504 on expiry, and the |
| in-flight model call is cancelled. |
| * A failing dependency (after retries) answers 502 with the error type, |
| unless the pipeline was built with ``on_error="record"``. |
| * Optional API keys: ``Authorization: Bearer <key>`` or ``X-API-Key: <key>``. |
| * Every response carries ``X-Request-ID`` (echoed if the client sent one), and |
| one JSON log line per query is written to the ``fgao.server`` logger. |
| """ |
| from __future__ import annotations |
|
|
| import asyncio |
| import dataclasses |
| import json |
| import logging |
| import time |
| import uuid |
| from typing import Any, Dict, Iterable, List, Optional, Union |
|
|
| from .. import __version__ |
| from ..rag.pipeline import FGAORAG, RAGResult |
| from ..resilience import _Resilient |
| from .metrics import Registry |
|
|
| from fastapi import Depends, FastAPI, Header, HTTPException, Request |
| from fastapi.responses import JSONResponse, PlainTextResponse, StreamingResponse |
| from pydantic import BaseModel, Field |
|
|
| log = logging.getLogger("fgao.server") |
|
|
|
|
| DocIn = Union[str, Dict[str, Any]] |
|
|
|
|
| class AnswerRequest(BaseModel): |
| question: str = Field(..., min_length=1, max_length=8000) |
| budget: Optional[float] = Field(None, gt=0) |
| seed: int = 0 |
| trace: bool = True |
| documents: Optional[List[DocIn]] = Field(None, description="answer from these passages instead of the configured KB") |
| web_documents: Optional[List[DocIn]] = Field(None, description="optional second source for the request documents") |
|
|
|
|
| class BatchRequest(BaseModel): |
| questions: List[str] = Field(..., min_length=1) |
| budget: Optional[float] = Field(None, gt=0) |
| seed: int = 0 |
| trace: bool = False |
| documents: Optional[List[DocIn]] = None |
| web_documents: Optional[List[DocIn]] = None |
|
|
|
|
| class Overloaded(Exception): |
| pass |
|
|
|
|
| class Admission: |
| """Bounded concurrency with a bounded wait queue (per event loop / worker process).""" |
|
|
| def __init__(self, max_inflight: int, max_queue: int) -> None: |
| self.max_inflight, self.max_queue = max_inflight, max_queue |
| self.inflight = 0 |
| self.waiting = 0 |
| self._sem: Optional[asyncio.Semaphore] = None |
|
|
| def _s(self) -> asyncio.Semaphore: |
| if self._sem is None: |
| self._sem = asyncio.Semaphore(self.max_inflight) |
| return self._sem |
|
|
| def check(self, n: int = 1) -> None: |
| free = self.max_inflight - self.inflight |
| if n > free and self.waiting + n - max(free, 0) > self.max_queue: |
| raise Overloaded() |
|
|
| async def acquire(self) -> None: |
| self.check() |
| self.waiting += 1 |
| try: |
| await self._s().acquire() |
| finally: |
| self.waiting -= 1 |
| self.inflight += 1 |
|
|
| def release(self) -> None: |
| self.inflight -= 1 |
| self._s().release() |
|
|
|
|
| def _resilient_parts(rag: FGAORAG) -> List[_Resilient]: |
| seen, out = set(), [] |
| cands = [rag.llm] + [s.retriever for s in rag.sources] |
| for comp in (rag.generator, rag.grader, rag.rewriter, rag.support_checker): |
| cands.append(getattr(comp, "llm", None)) |
| for c in cands: |
| if isinstance(c, _Resilient) and id(c) not in seen: |
| seen.add(id(c)) |
| out.append(c) |
| return out |
|
|
|
|
| def create_app( |
| rag: FGAORAG, |
| *, |
| api_keys: Optional[Iterable[str]] = None, |
| max_inflight: int = 256, |
| max_queue: int = 1024, |
| request_timeout: float = 60.0, |
| max_batch: int = 256, |
| max_budget: Optional[float] = None, |
| max_documents: int = 500, |
| title: str = "FGAO RAG", |
| ): |
| """Build the FastAPI application. ``max_budget`` caps per-request budgets (default: the pipeline budget x 4).""" |
| keys = {k for k in (api_keys or []) if k} |
| cap = float(max_budget if max_budget is not None else rag.budget * 4) |
| adm = Admission(max_inflight, max_queue) |
| reg = Registry() |
| m_http = reg.counter("fgao_http_requests_total", "HTTP requests by route and status code") |
| m_lat = reg.histogram("fgao_request_duration_seconds", "Query latency in seconds", |
| (0.05, 0.1, 0.25, 0.5, 1, 2, 4, 8, 16, 32, 64)) |
| m_out = reg.counter("fgao_queries_total", "Queries by outcome (verified, abstained, timeout, error)") |
| m_stop = reg.counter("fgao_stop_status_total", "Queries by FGAO stop status") |
| m_cost = reg.histogram("fgao_query_cost", "Budget spent per query", (1, 2, 4, 6, 8, 10, 12, 16, 24, 32, 64)) |
| m_ops = reg.counter("fgao_operation_calls_total", "Operation calls (generate, grade, rewrite, verify, retrieve:*)") |
| m_steps = reg.counter("fgao_steps_total", "Controller steps executed") |
| m_inflight = reg.gauge("fgao_inflight", "Queries running") |
| m_queue = reg.gauge("fgao_queue", "Queries waiting for a slot") |
| parts = _resilient_parts(rag) |
|
|
| def collect(): |
| m_inflight.set(adm.inflight) |
| m_queue.set(adm.waiting) |
| for p in parts: |
| for k, v in p.stats.as_dict().items(): |
| yield (f"fgao_dependency_{k}_total", "counter", f"Dependency {k} (resilient wrappers)", {"name": p.name}, v) |
|
|
| reg.collectors.append(lambda: list(collect())) |
|
|
| app = FastAPI(title=title, version=__version__, |
| description="Training-free FGAO orchestration of retrieval-augmented generation.") |
|
|
| def auth(authorization: Optional[str] = Header(None), x_api_key: Optional[str] = Header(None)) -> None: |
| if not keys: |
| return |
| token = x_api_key or (authorization[7:] if authorization and authorization.lower().startswith("bearer ") else None) |
| if token not in keys: |
| raise HTTPException(status_code=401, detail="invalid or missing API key") |
|
|
| @app.middleware("http") |
| async def request_id(request: Request, call_next): |
| rid = request.headers.get("x-request-id") or uuid.uuid4().hex |
| request.state.rid = rid |
| resp = await call_next(request) |
| resp.headers["X-Request-ID"] = rid |
| route = request.scope.get("route") |
| m_http.inc(route=getattr(route, "path", request.url.path), code=str(resp.status_code)) |
| return resp |
|
|
| def _budget(b: Optional[float]) -> Optional[float]: |
| if b is not None and b > cap: |
| raise HTTPException(status_code=422, detail=f"budget exceeds the server limit {cap:g}") |
| return b |
|
|
| def _payload(r: RAGResult, trace: bool) -> Dict[str, Any]: |
| d = r.to_dict() |
| if not trace: |
| d.pop("trace", None) |
| return d |
|
|
| def _record(r: RAGResult, seconds: float, rid: str) -> None: |
| m_out.inc(outcome="verified" if r.answer is not None else "abstained") |
| m_stop.inc(status=r.status) |
| m_lat.observe(seconds) |
| m_cost.observe(r.cost) |
| m_steps.inc(len(r.steps)) |
| for k, v in r.calls.items(): |
| m_ops.inc(v, op=k) |
| log.info(json.dumps({"event": "query", "request_id": rid, "status": r.status, "verified": r.answer is not None, |
| "cost": r.cost, "steps": len(r.steps), "seconds": round(seconds, 4)})) |
|
|
| def _pipeline(documents, web_documents) -> FGAORAG: |
| """The configured pipeline, or one over the request's own passages (same LLM, costs, priors, controller).""" |
| if not documents: |
| return rag |
| if len(documents) > max_documents or (web_documents and len(web_documents) > max_documents): |
| raise HTTPException(status_code=422, detail=f"at most {max_documents} documents per source") |
| if rag.llm is None: |
| raise HTTPException(status_code=422, detail="request documents need a pipeline with a default LLM") |
| return FGAORAG.from_retrievers(kb=list(documents), web=list(web_documents) if web_documents else None, |
| llm=rag.llm, costs=rag.costs, priors=rag.priors, weights=rag.weights, |
| budget=rag.budget, config=rag.config, on_error=rag.on_error) |
|
|
| async def _run(q: str, budget: Optional[float], seed: int, rid: str, pipe: Optional[FGAORAG] = None) -> RAGResult: |
| t0 = time.perf_counter() |
| pipe = pipe or rag |
| try: |
| r = await asyncio.wait_for(pipe.ainvoke(q, budget=budget, seed=seed), request_timeout) |
| except asyncio.TimeoutError: |
| m_out.inc(outcome="timeout") |
| m_lat.observe(time.perf_counter() - t0) |
| log.warning(json.dumps({"event": "timeout", "request_id": rid})) |
| raise |
| except Exception as e: |
| m_out.inc(outcome="error") |
| m_lat.observe(time.perf_counter() - t0) |
| log.error(json.dumps({"event": "error", "request_id": rid, "error": type(e).__name__, "detail": str(e)[:500]})) |
| raise |
| _record(r, time.perf_counter() - t0, rid) |
| return r |
|
|
| def _overloaded() -> JSONResponse: |
| return JSONResponse({"detail": "server overloaded, retry later"}, status_code=429, headers={"Retry-After": "1"}) |
|
|
| @app.get("/health") |
| async def health(): |
| return {"status": "ok"} |
|
|
| @app.get("/ready") |
| async def ready(): |
| return {"status": "ready", "inflight": adm.inflight, "queued": adm.waiting} |
|
|
| @app.get("/v1/info", dependencies=[Depends(auth)]) |
| async def info(): |
| return { |
| "version": __version__, |
| "budget": rag.budget, "max_budget": cap, |
| "sources": [{"name": s.name, "k": s.k, "cost": s.cost, "family": s.family, "rewrite": s.rewrite} |
| for s in rag.sources], |
| "costs": dataclasses.asdict(rag.costs), |
| "weights": dataclasses.asdict(rag.weights), |
| "priors": dataclasses.asdict(rag.priors), |
| "controller": dataclasses.asdict(rag.config), |
| "operations": [str(a) for a in rag.app.actions], |
| "on_error": rag.on_error, |
| "limits": {"max_inflight": max_inflight, "max_queue": max_queue, "request_timeout": request_timeout, |
| "max_batch": max_batch}, |
| } |
|
|
| @app.post("/v1/answer", dependencies=[Depends(auth)]) |
| async def answer(req: AnswerRequest, request: Request): |
| budget = _budget(req.budget) |
| pipe = _pipeline(req.documents, req.web_documents) |
| try: |
| await adm.acquire() |
| except Overloaded: |
| return _overloaded() |
| try: |
| r = await _run(req.question, budget, req.seed, request.state.rid, pipe) |
| except asyncio.TimeoutError: |
| raise HTTPException(status_code=504, detail=f"query exceeded {request_timeout:g} s") |
| except Exception as e: |
| raise HTTPException(status_code=502, detail=f"{type(e).__name__}: {e}"[:500]) |
| finally: |
| adm.release() |
| return _payload(r, req.trace) |
|
|
| @app.post("/v1/batch", dependencies=[Depends(auth)]) |
| async def batch(req: BatchRequest, request: Request): |
| if len(req.questions) > max_batch: |
| raise HTTPException(status_code=422, detail=f"at most {max_batch} questions per batch") |
| budget = _budget(req.budget) |
| pipe = _pipeline(req.documents, req.web_documents) |
| try: |
| adm.check(len(req.questions)) |
| except Overloaded: |
| return _overloaded() |
|
|
| async def one(i: int, q: str) -> Dict[str, Any]: |
| await adm.acquire() |
| try: |
| return _payload(await _run(q, budget, req.seed, f"{request.state.rid}:{i}", pipe), req.trace) |
| except asyncio.TimeoutError: |
| return {"question": q, "error": "timeout"} |
| except Exception as e: |
| return {"question": q, "error": f"{type(e).__name__}: {e}"[:500]} |
| finally: |
| adm.release() |
|
|
| results = await asyncio.gather(*(one(i, q) for i, q in enumerate(req.questions))) |
| return {"results": list(results)} |
|
|
| @app.post("/v1/answer/stream", dependencies=[Depends(auth)]) |
| async def stream(req: AnswerRequest, request: Request): |
| budget = _budget(req.budget) |
| pipe = _pipeline(req.documents, req.web_documents) |
| try: |
| await adm.acquire() |
| except Overloaded: |
| return _overloaded() |
| rid = request.state.rid |
| released = [False] |
|
|
| def release_once() -> None: |
| if not released[0]: |
| released[0] = True |
| adm.release() |
|
|
| async def events(): |
| t0 = time.perf_counter() |
| agen = pipe.astream(req.question, budget=budget, seed=req.seed).__aiter__() |
| try: |
| while True: |
| left = request_timeout - (time.perf_counter() - t0) |
| try: |
| item = await asyncio.wait_for(agen.__anext__(), max(left, 0.001)) |
| except StopAsyncIteration: |
| break |
| if isinstance(item, RAGResult): |
| _record(item, time.perf_counter() - t0, rid) |
| yield f"event: result\ndata: {json.dumps(_payload(item, req.trace))}\n\n" |
| else: |
| yield f"event: step\ndata: {json.dumps(item.to_dict())}\n\n" |
| except asyncio.TimeoutError: |
| m_out.inc(outcome="timeout") |
| yield f"event: error\ndata: {json.dumps({'error': 'timeout'})}\n\n" |
| except Exception as e: |
| m_out.inc(outcome="error") |
| yield f"event: error\ndata: {json.dumps({'error': f'{type(e).__name__}: {e}'[:500]})}\n\n" |
| finally: |
| release_once() |
|
|
| from starlette.background import BackgroundTask |
|
|
| return StreamingResponse(events(), media_type="text/event-stream", background=BackgroundTask(release_once), |
| headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}) |
|
|
| @app.get("/metrics") |
| async def metrics(): |
| return PlainTextResponse(reg.render(), media_type="text/plain; version=0.0.4") |
|
|
| app.state.rag = rag |
| app.state.admission = adm |
| app.state.registry = reg |
| return app |
|
|
