| """FastAPI entrypoint for MBF Insight."""
|
| from __future__ import annotations
|
|
|
| import asyncio
|
| import io
|
| import json
|
| import re
|
| from typing import AsyncIterator
|
|
|
| from fastapi import FastAPI, File, HTTPException, UploadFile
|
| from fastapi.middleware.cors import CORSMiddleware
|
| from pydantic import BaseModel
|
| from sse_starlette.sse import EventSourceResponse
|
|
|
| from .config import get_settings
|
| from .graph.agent import run_agent
|
| from .graph.live_dashboard import build_live_dashboard
|
| from .graph.suggest import get_followups, get_suggestions
|
| from .retrieval.duck import get_con, schema_version, with_writer
|
|
|
| app = FastAPI(title="MBF Insight API", version="0.1.0")
|
|
|
| app.add_middleware(
|
| CORSMiddleware,
|
| allow_origins=["*"],
|
| allow_credentials=False,
|
| allow_methods=["*"],
|
| allow_headers=["*"],
|
| )
|
|
|
|
|
| class AskBody(BaseModel):
|
| question: str
|
| thread_id: str | None = None
|
|
|
|
|
| class FollowupBody(BaseModel):
|
| question: str
|
| narrative: str | None = ""
|
| component_titles: list[str] | None = None
|
| n: int = 4
|
|
|
|
|
| @app.get("/health")
|
| def health() -> dict:
|
| from .keys.pool import get_pool
|
| from .retrieval.duck import get_con
|
| s = get_settings()
|
| try:
|
| n = get_con().execute(
|
| "SELECT COUNT(*) FROM information_schema.tables WHERE table_schema='main'"
|
| ).fetchone()[0]
|
| db_ok = True
|
| except Exception as e:
|
| n = 0
|
| db_ok = False
|
| pool = get_pool()
|
| from .keys.pool import ROLE_RANGES, Role
|
| role_map = {
|
| r.value: {"keys": list(range(lo, hi + 1)),
|
| "active": sum(1 for st in pool._states if lo <= st.idx <= hi and st.idx not in pool._dead)}
|
| for r, (lo, hi) in ROLE_RANGES.items()
|
| }
|
| return {
|
| "status": "ok",
|
| "duckdb": s.duckdb_path,
|
| "duckdb_ok": db_ok,
|
| "tables": n,
|
| "llm_keys": len(pool._states),
|
| "dead_keys": sorted(pool._dead),
|
| "mode": "groq" if pool.available() else "offline",
|
| "roles": role_map,
|
| }
|
|
|
|
|
| @app.post("/ask")
|
| async def ask(body: AskBody) -> dict:
|
| state = await run_agent(body.question)
|
| return {
|
| "question": body.question,
|
| "mode": state.mode,
|
| "intent": state.intent,
|
| "sql": state.sql,
|
| "row_count": len(state.rows),
|
| "plan": state.plan.model_dump() if state.plan else None,
|
| "trace": state.trace,
|
| "keys_used": state.keys_used,
|
| }
|
|
|
|
|
| @app.post("/ask/stream")
|
| async def ask_stream(body: AskBody) -> EventSourceResponse:
|
| """Real-time pipeline streaming. Emits one SSE event per stage transition.
|
|
|
| Event types:
|
| - `stage`: incremental updates (status: running|done|error)
|
| - `final`: terminal payload with full plan + sql + trace + keys_used
|
| """
|
| queue: asyncio.Queue[dict] = asyncio.Queue()
|
| STOP = object()
|
|
|
| async def emit(evt: dict) -> None:
|
| await queue.put(evt)
|
|
|
| async def runner() -> None:
|
| try:
|
| state = await run_agent(body.question, emit=emit)
|
| await queue.put({
|
| "__final__": True,
|
| "question": body.question,
|
| "mode": state.mode,
|
| "intent": state.intent,
|
| "sql": state.sql,
|
| "row_count": len(state.rows),
|
| "plan": state.plan.model_dump() if state.plan else None,
|
| "trace": state.trace,
|
| "keys_used": state.keys_used,
|
| "confidence": state.confidence,
|
| "followups": state.followups,
|
| })
|
| except Exception as e:
|
| await queue.put({"__error__": True, "error": str(e)[:300]})
|
| finally:
|
| await queue.put(STOP)
|
|
|
| async def gen() -> AsyncIterator[dict]:
|
|
|
| yield {"event": "open", "data": json.dumps({"ok": True})}
|
| task = asyncio.create_task(runner())
|
| try:
|
| while True:
|
| item = await queue.get()
|
| if item is STOP:
|
| break
|
| if isinstance(item, dict) and item.get("__final__"):
|
| payload = {k: v for k, v in item.items() if k != "__final__"}
|
| yield {"event": "final", "data": json.dumps(payload, default=str)}
|
| elif isinstance(item, dict) and item.get("__error__"):
|
| yield {"event": "error", "data": json.dumps({"error": item["error"]})}
|
| else:
|
| yield {"event": "stage", "data": json.dumps(item, default=str)}
|
| finally:
|
| if not task.done():
|
| task.cancel()
|
|
|
| return EventSourceResponse(gen())
|
|
|
|
|
| @app.get("/suggest")
|
| async def suggest(n: int = 8, seed: str = "default") -> dict:
|
| items = await get_suggestions(seed=seed, n=n)
|
| return {"items": items}
|
|
|
|
|
| @app.post("/followups")
|
| async def followups(body: FollowupBody) -> dict:
|
| items = await get_followups(
|
| question=body.question,
|
| narrative=body.narrative or "",
|
| component_titles=body.component_titles or [],
|
| n=body.n,
|
| )
|
| return {"items": items}
|
|
|
|
|
| @app.get("/dashboard/live")
|
| async def dashboard_live(seed: int | None = None) -> dict:
|
| """Live, ever-changing snapshot of the warehouse — no question required.
|
|
|
| The plan adapts to which tables hold data, their shape, and cardinality.
|
| Pass `seed` to get a deterministic layout; omit it for a fresh shuffle.
|
| """
|
|
|
| return await asyncio.to_thread(build_live_dashboard, seed)
|
|
|
|
|
|
|
|
|
| _TABLE_NAME_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]{0,62}$")
|
|
|
|
|
| def _ident(name: str) -> str:
|
| """Quote a SQL identifier safely (DuckDB)."""
|
| return '"' + name.replace('"', '""') + '"'
|
|
|
|
|
| def _list_kb_tables() -> list[dict]:
|
| con = get_con()
|
| rows = con.execute(
|
| "SELECT table_name FROM information_schema.tables "
|
| "WHERE table_schema='main' ORDER BY table_name"
|
| ).fetchall()
|
| out: list[dict] = []
|
| for (t,) in rows:
|
| try:
|
| n = con.execute(f"SELECT COUNT(*) FROM {_ident(t)}").fetchone()[0]
|
| except Exception:
|
| n = None
|
| try:
|
| cols = [r[1] for r in con.execute(f"PRAGMA table_info({_ident(t)})").fetchall()]
|
| except Exception:
|
| cols = []
|
| out.append({"name": t, "rows": n, "columns": cols})
|
| return out
|
|
|
|
|
| @app.get("/kb/version")
|
| def kb_version() -> dict:
|
| return {"version": schema_version()}
|
|
|
|
|
| @app.get("/kb/tables")
|
| async def kb_tables() -> dict:
|
| items = await asyncio.to_thread(_list_kb_tables)
|
| return {"version": schema_version(), "tables": items}
|
|
|
|
|
| class RenameBody(BaseModel):
|
| old_name: str
|
| new_name: str
|
|
|
|
|
| @app.post("/kb/rename")
|
| async def kb_rename(body: RenameBody) -> dict:
|
| if not _TABLE_NAME_RE.match(body.new_name):
|
| raise HTTPException(status_code=400, detail="Invalid new table name")
|
| if not _TABLE_NAME_RE.match(body.old_name):
|
| raise HTTPException(status_code=400, detail="Invalid current table name")
|
|
|
| def _do(w):
|
| w.execute(f"ALTER TABLE {_ident(body.old_name)} RENAME TO {_ident(body.new_name)}")
|
|
|
| try:
|
| await asyncio.to_thread(with_writer, _do)
|
| except Exception as e:
|
| raise HTTPException(status_code=400, detail=str(e)[:300]) from e
|
| return {"ok": True, "version": schema_version()}
|
|
|
|
|
| @app.delete("/kb/tables/{name}")
|
| async def kb_delete(name: str) -> dict:
|
| if not _TABLE_NAME_RE.match(name):
|
| raise HTTPException(status_code=400, detail="Invalid table name")
|
|
|
| def _do(w):
|
| w.execute(f"DROP TABLE IF EXISTS {_ident(name)}")
|
|
|
| try:
|
| await asyncio.to_thread(with_writer, _do)
|
| except Exception as e:
|
| raise HTTPException(status_code=400, detail=str(e)[:300]) from e
|
| return {"ok": True, "version": schema_version()}
|
|
|
|
|
| @app.post("/kb/upload")
|
| async def kb_upload(
|
| file: UploadFile = File(...),
|
| table: str | None = None,
|
| mode: str = "replace",
|
| ) -> dict:
|
| """Upload a CSV file and load it as a DuckDB table.
|
|
|
| The table name is derived from `table` (preferred) or the filename stem.
|
| Modes: `replace` (drop+create) or `append` (insert into existing).
|
| """
|
| raw = await file.read()
|
| if not raw:
|
| raise HTTPException(status_code=400, detail="Empty file")
|
| if len(raw) > 50 * 1024 * 1024:
|
| raise HTTPException(status_code=413, detail="File exceeds 50MB limit")
|
|
|
|
|
| candidate = (table or file.filename or "uploaded").rsplit("/", 1)[-1]
|
| candidate = candidate.rsplit(".", 1)[0]
|
| candidate = re.sub(r"[^A-Za-z0-9_]+", "_", candidate).strip("_")
|
| if not candidate or candidate[0].isdigit():
|
| candidate = f"t_{candidate}" if candidate else "uploaded"
|
| if not _TABLE_NAME_RE.match(candidate):
|
| raise HTTPException(status_code=400, detail="Could not derive a valid table name")
|
|
|
| if mode not in {"replace", "append"}:
|
| raise HTTPException(status_code=400, detail="mode must be replace|append")
|
|
|
|
|
| head = raw[:4].lower()
|
| if head.startswith(b"pk\x03\x04") or head.startswith(b"\xd0\xcf\x11\xe0"):
|
| raise HTTPException(status_code=400, detail="Only CSV/TSV files are supported")
|
|
|
| def _do(w):
|
|
|
|
|
| import tempfile
|
| import os
|
| tmp = tempfile.NamedTemporaryFile(suffix=".csv", delete=False)
|
| try:
|
| tmp.write(raw)
|
| tmp.close()
|
| tmp_path = tmp.name.replace("\\", "/")
|
| if mode == "replace":
|
| w.execute(f"DROP TABLE IF EXISTS {_ident(candidate)}")
|
| w.execute(
|
| f"CREATE TABLE {_ident(candidate)} AS "
|
| f"SELECT * FROM read_csv_auto('{tmp_path}', SAMPLE_SIZE=-1, IGNORE_ERRORS=TRUE)"
|
| )
|
| else:
|
| w.execute(
|
| f"INSERT INTO {_ident(candidate)} "
|
| f"SELECT * FROM read_csv_auto('{tmp_path}', SAMPLE_SIZE=-1, IGNORE_ERRORS=TRUE)"
|
| )
|
| n = w.execute(f"SELECT COUNT(*) FROM {_ident(candidate)}").fetchone()[0]
|
| return n
|
| finally:
|
| try:
|
| os.unlink(tmp.name)
|
| except Exception:
|
| pass
|
|
|
| try:
|
| rows = await asyncio.to_thread(with_writer, _do)
|
| except HTTPException:
|
| raise
|
| except Exception as e:
|
| raise HTTPException(status_code=400, detail=f"Load failed: {str(e)[:300]}") from e
|
|
|
| return {
|
| "ok": True,
|
| "table": candidate,
|
| "rows": rows,
|
| "mode": mode,
|
| "version": schema_version(),
|
| }
|
|
|
