119 lines
3.4 KiB
Python
119 lines
3.4 KiB
Python
"""FastAPI entry point for nvi.
|
|
|
|
Endpoints:
|
|
GET /healthz — liveness probe
|
|
POST /ask — submit a question, returns run_id
|
|
GET /runs/{run_id}/stream — SSE stream of run events
|
|
GET /runs/{run_id} — final state snapshot
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import uuid
|
|
from contextlib import asynccontextmanager
|
|
from datetime import datetime, timezone
|
|
|
|
from fastapi import FastAPI, HTTPException
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from fastapi.responses import StreamingResponse
|
|
from pydantic import BaseModel
|
|
|
|
from api.runtime import events
|
|
from api.runtime.runner import run as run_graph
|
|
|
|
logger = logging.getLogger("nvi")
|
|
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s %(message)s")
|
|
|
|
|
|
# In-memory run store. Each entry holds final-state + lifecycle metadata; the
|
|
# event queue is owned by api.runtime.events.
|
|
runs: dict[str, dict] = {}
|
|
RUN_TTL_SECONDS = 3600
|
|
RUN_CLEANUP_INTERVAL = 300
|
|
|
|
|
|
async def _cleanup_runs():
|
|
while True:
|
|
await asyncio.sleep(RUN_CLEANUP_INTERVAL)
|
|
now = datetime.now(timezone.utc).timestamp()
|
|
expired = [
|
|
rid for rid, r in runs.items()
|
|
if r["status"] in ("completed", "error")
|
|
and now - r.get("created_at", now) > RUN_TTL_SECONDS
|
|
]
|
|
for rid in expired:
|
|
del runs[rid]
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(_app: FastAPI):
|
|
task = asyncio.create_task(_cleanup_runs())
|
|
yield
|
|
task.cancel()
|
|
|
|
|
|
app = FastAPI(title="nvi", lifespan=lifespan)
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"],
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
|
|
@app.get("/healthz")
|
|
async def healthz():
|
|
return {"status": "ok", "runs_in_memory": len(runs)}
|
|
|
|
|
|
class AskRequest(BaseModel):
|
|
question: str
|
|
|
|
|
|
@app.post("/ask")
|
|
async def ask(req: AskRequest):
|
|
run_id = uuid.uuid4().hex[:8]
|
|
now = datetime.now(timezone.utc)
|
|
runs[run_id] = {
|
|
"status": "running",
|
|
"question": req.question,
|
|
"created_at": now.timestamp(),
|
|
}
|
|
events.open_run(run_id)
|
|
|
|
async def _drive():
|
|
try:
|
|
final = await run_graph(run_id, req.question)
|
|
runs[run_id] = {
|
|
**runs[run_id],
|
|
"status": "error" if final.get("error") else "completed",
|
|
"plan_rationale": final.get("plan_rationale"),
|
|
"plan_steps": final.get("plan_steps", []),
|
|
"findings": final.get("findings", []),
|
|
"answer": final.get("answer", ""),
|
|
"error": final.get("error"),
|
|
}
|
|
except Exception as e:
|
|
logger.exception("run %s failed at top level", run_id)
|
|
runs[run_id] = {**runs[run_id], "status": "error", "error": str(e)}
|
|
|
|
asyncio.create_task(_drive())
|
|
logger.info("ask_received run_id=%s q=%r", run_id, req.question)
|
|
return {"run_id": run_id, "status": "running"}
|
|
|
|
|
|
@app.get("/runs/{run_id}")
|
|
async def get_run(run_id: str):
|
|
if run_id not in runs:
|
|
raise HTTPException(404, detail=f"Run {run_id} not found")
|
|
return runs[run_id]
|
|
|
|
|
|
@app.get("/runs/{run_id}/stream")
|
|
async def stream_run(run_id: str):
|
|
if run_id not in runs:
|
|
raise HTTPException(404, detail=f"Run {run_id} not found")
|
|
return StreamingResponse(events.stream(run_id), media_type="text/event-stream")
|