"""FastAPI backend for the answering playground. Serves the single-page UI and these endpoints: GET /api/documents -> list ingested documents: filing metadata, derived label, overlap and review flags GET /api/document/{id}/pdf -> the source PDF (for PDF.js) GET /api/document/{id}/flags -> the document's review queue, in page order PATCH /api/documents/{id} -> edit metadata (title, folder, filing fields); organization only, never identity DELETE /api/documents/{id} -> remove the document outright (hard delete; 409 while its ingest or a batch run against it is still going) POST /api/chat -> retrieve -> Opus answer -> grounded refs POST /api/chat/stream -> the same answer as server-sent events POST /api/upload -> raw PDF body; runs the quber workflow (fuse -> figure -> ingest) as a job GET /api/jobs/{job_id} -> job stage (uploaded/waiting/extracting/ reading figures/ingesting/done/failed) plus a tail of its log POST /api/batch -> start a run of many questions -> job id GET /api/batch/{id} -> the run's state and rows finished so far GET /api/batch/{id}/events -> SSE, one event per question completing GET /api/batch/{id}/export.xlsx and .csv -> the run as a workbook / text GET /healthz -> readiness: 200 once the database answers and the embedding model is loaded GET /static/{path} -> the UI's assets `quber.playground.auth.install` adds `/api/me`, and when sign-in is on also `/login`, `/callback` and `/logout`. Document identity is the integer id plus the content hash. The storage key (`doc_key`, the hash's first 16 hex chars) names the staged files and is what ingest subprocesses take. The document list leaves it out, but the PATCH response returns the full document row and so carries it. Labels are derived from the filing metadata at read time — see `quber.playground.metadata`. Uploads take the PDF as the raw request body (no multipart dependency). The bytes are hashed before anything runs: the same bytes as an existing document are rejected outright (409, kind "duplicate"), and a filing-tuple match is refused (409, kind "tuple") unless the request carries the replace intent — then the swap happens when the new ingest completes. An s3:// source cannot be hashed before its job stages it, so those checks run inside the job and a failure carries the same structured `error_info`. Every upload is admitted; a semaphore bounds how many run their pipelines at once (`settings.playground.upload_concurrency`), and a job holding no slot yet reports the `waiting` stage rather than appearing stalled on its first step. Run: uv run quber playground """ from __future__ import annotations import hashlib import json import re import shutil import subprocess import threading import uuid from pathlib import Path from typing import Any, Dict, List, Literal, LiteralString, Optional from fastapi import FastAPI, HTTPException, Request from fastapi.responses import FileResponse, HTMLResponse, Response, StreamingResponse from loguru import logger from pydantic import BaseModel from quber.files.cache import fetch_file from quber.playground import auth, db, embedding, export, gpu, jobs, metadata, storage from quber.playground.answers.document import DocumentIdentity from quber.playground.answers.expectation import answer as agent_answer from quber.playground.answers.expectation import answer_stream as agent_answer_stream from quber.playground.batch import BatchRunner, RunState, parse_questions from quber.playground.retrieval import RetrievedChunk, retrieve from quber.playground.tracing import document_key from quber.settings import get_settings STATIC = Path(__file__).with_name("static") DATA_DIR = get_settings().playground.data_dir UPLOADS = DATA_DIR / "uploads" REPO_ROOT = Path(__file__).parents[3] app = FastAPI(title="Quber Playground") # On when the login settings are present, which is the hosted playground; # off on the developer host. A partial configuration is refused here. LOGIN_ENABLED = auth.install(app, get_settings().playground) # Set once the embedding model has loaded. The health route reports the app # ready only after this, so a request handed off by the load balancer never # lands on a task that would stall its first question on the model load. _EMBEDDING_READY = threading.Event() @app.on_event("startup") def warm_embedding_model() -> None: """Load the embedding model in the background at startup. It otherwise loads lazily on the first question, costing that question ~5 seconds.""" def warm() -> None: embedding.embed_query("warm up") _EMBEDDING_READY.set() threading.Thread(target=warm, daemon=True).start() @app.on_event("startup") def open_job_records() -> None: """The job tables exist before anything writes to them, and this process starts beating for the jobs it runs.""" jobs.ensure_tables() jobs.start_heartbeat( upload_ids=lambda: [j["job_id"] for j in JOBS.values() if j["stage"] not in jobs.TERMINAL_STAGES], run_ids=BATCH.running_ids, store=BATCH_STORE, ) @app.get("/healthz", include_in_schema=False) def healthz() -> Response: """Readiness, not liveness: 200 only once the database answers and the embedding model is loaded, 503 before. The load balancer's target group and the scale-to-zero startup page both poll this, so a 200 means the app can take a question, not merely that the process is up.""" if not _EMBEDDING_READY.is_set(): return Response( '{"status":"loading embedding model"}', status_code=503, media_type="application/json" ) try: with db.connect() as conn: conn.execute("SELECT 1").fetchone() except Exception as exc: return Response( json.dumps({"status": f"database unreachable: {type(exc).__name__}"}), status_code=503, media_type="application/json", ) return Response('{"status":"ok"}', media_type="application/json") class ChatRequest(BaseModel): doc_id: int question: str k: int = 10 # What the asker expects back: "auto" lets the model choose the shape, # "value" forces a figure, "text" forces an explanation. want: Literal["auto", "value", "text"] = "auto" class Reference(BaseModel): ref_id: str ref_type: Optional[str] page: int # 1-based, to match the viewer bbox: Optional[dict] label: str # Cell-level provenance, the GroundedCell attributes established at # reconciliation (fusion documents only; None on ADE references): status: Optional[str] = None # provenance code, e.g. 'reconciled' note: Optional[str] = None # status inspector's evidence line, if any text: Optional[str] = None # the served cell value row: Optional[int] = None # 0-based position in the corrected table col: Optional[int] = None flagged: bool = False # True when the status is registered for review # Figure-value provenance, from the grounding's position (None elsewhere): chart: Optional[str] = None # the chart the value is printed in segment: Optional[str] = None # what the value is labeled as on the chart # The status in the reader's words, from the registry's label for the # code. The transcript prints this, never the code. A misread carries # the figure the page prints so the line can name both numbers. reason: Optional[str] = None printed: Optional[str] = None class RetrievedInfo(BaseModel): chunk_id: str chunk_type: str page: int score: float class ChatResponse(BaseModel): answer: str references: List[Reference] retrieved: List[RetrievedInfo] # The answer as a value rather than as text. `shape` names which payload # came back — scalar, series, grid, prose, unanswerable — and `value` is # that payload. `answer` stays populated with something readable so an # existing consumer of this endpoint keeps working. shape: str = "prose" value: Optional[dict] = None # Statuses the provenance registry marks for review; cited cells carrying one # are served like any other but highlighted so the user knows they are flagged. def _inspect_codes() -> frozenset[str]: from quber.core.extractors.base import CELL_STATUS_REFERENCE return frozenset(s.code for s in CELL_STATUS_REFERENCE if s.inspect) def _status_reasons() -> Dict[str, str]: from quber.core.extractors.base import CELL_STATUS_REFERENCE return {s.code: s.label for s in CELL_STATUS_REFERENCE} INSPECT_CODES = _inspect_codes() STATUS_REASONS = _status_reasons() # The misread note is written by the figure-value reconciliation in # `quber.core.figures.values` in one shape: "the page prints '6%' where this # value should appear - possible misread". The figure inside the quotes is what # the transcript's line needs. The figure is formatted with repr, so a figure # that itself contains a single quote is not quoted the way this pattern # expects and is not read correctly. _PRINTED_IN_NOTE = re.compile(r"prints '([^']*)'") def _reason( status: Optional[str], note: Optional[str], served: Optional[str] ) -> tuple[Optional[str], Optional[str]]: """The plain-language line for a status, and the printed figure a misread names.""" if not status: return None, None printed = None if status == "value_misread" and note: m = _PRINTED_IN_NOTE.search(note) if m: printed = m.group(1) if printed and served: return f"The page prints {printed} here, not {served}. Check which is right.", printed return STATUS_REASONS.get(status, status), printed def _label(ref_type: Optional[str], page1: int, status: Optional[str] = None) -> str: """What the source is, in the reader's words. The status never rides the label: it has its own field, and its plain-language line is `reason`.""" if ref_type == "tableCell": return f"Page {page1}, table cell" if ref_type in ("table", "chunkTable"): return f"Page {page1}, table" if ref_type == "figureValue": return f"Page {page1}, chart value" if ref_type == "line_item": return f"Page {page1}, table row" if ref_type == "picture": return f"Page {page1}, figure" return f"Page {page1}, text" def _stamped_index() -> str: """The index shell with content-hashed asset links. A browser that cached a previous deploy of app.js kept serving it on a plain reload, because the asset URLs never change. Stamping each mutable asset with a hash of its current bytes makes any change a new URL; the pinned vendor files stay as they are. Hashing runs per request — the files are small and the reload server would otherwise need cache invalidation of its own. """ html = (STATIC / "index.html").read_text() for name in ("styles.css", "app.css", "app.js"): digest = hashlib.sha256((STATIC / name).read_bytes()).hexdigest()[:8] html = html.replace(f"/static/{name}", f"/static/{name}?v={digest}") return html @app.get("/", response_class=HTMLResponse) def index() -> HTMLResponse: # no-cache makes the shell revalidate on every load, so the stamped links # inside it are always current; the assets themselves may cache freely. return HTMLResponse(_stamped_index(), headers={"Cache-Control": "no-cache"}) @app.get("/static/{name:path}") def static_asset(name: str) -> FileResponse: # Nested paths are served (the design-system token stylesheets live in # static/tokens/), but the resolved file must still sit inside static/ so # that "../" segments cannot walk out of it. path = STATIC / name if not path.is_file() or not path.resolve().is_relative_to(STATIC.resolve()): raise HTTPException(404, f"no asset {name}") return FileResponse(path) _DOC_COLUMNS: LiteralString = ( "d.id, d.doc_key, d.filename, d.page_count, d.ade_version, d.title, d.folder, " "d.filing_type, d.year, d.period, d.version" ) # The processing track (documents.track) is deliberately NOT served here or # anywhere the browser can reach: naming vendor model internals in a client # payload makes them scrapeable. It is written at upload finalize and read in # the database only. def _doc_dict(r: "tuple[Any, ...]") -> Dict[str, Any]: doc = { "id": r[0], "doc_key": r[1], "filename": r[2], "page_count": r[3], "format": "fusion" if r[4] == "quber-fusion" else "ade", "title": r[5], "folder": r[6], "filing_type": r[7], "year": r[8], "period": r[9], "version": r[10], } doc["label"] = metadata.filing_label(r[7], r[8], r[9], r[10]) return doc def _doc_row(doc_id: int) -> Dict[str, Any]: with db.connect() as conn: r = conn.execute( "SELECT " + _DOC_COLUMNS + " FROM ade_playground.documents d WHERE d.id = %s", (doc_id,), ).fetchone() if r is None: raise HTTPException(404, f"no document {doc_id}") return _doc_dict(r) @app.get("/api/documents") def documents() -> list[dict]: # The flag count rides the list because the alternative is one # /api/document/{id}/flags call per document, each returning every flag # object with its bounding box in order to take a length. with db.connect() as conn: rows = conn.execute( "SELECT " + _DOC_COLUMNS + """, COUNT(g.id) FILTER (WHERE g.status = ANY(%s)) AS flags FROM ade_playground.documents d LEFT JOIN ade_playground.groundings g ON g.document_id = d.id GROUP BY d.id ORDER BY lower(coalesce(d.title, d.filename))""", (list(INSPECT_CODES),), ).fetchall() docs = [] for r in rows: doc = _doc_dict(r) # The list leaves the storage key out; nothing in the UI needs it. del doc["doc_key"] doc["flags"] = r[11] docs.append(doc) metadata.annotate_overlaps(docs) return docs class DocumentPatch(BaseModel): """The editable metadata: presentation (title, folder) and the filing fields. Edits change organization and labels only — never identity, and never blocked by collisions; two documents landing on one tuple simply both carry the possible-duplicate flag. Fields arrive as strings so an emptied input clears its column; year and version parse to ints.""" title: Optional[str] = None folder: Optional[str] = None filing_type: Optional[str] = None year: Optional[str] = None period: Optional[str] = None version: Optional[str] = None def _parse_int(name: str, raw: str, lo: int, hi: int) -> int: try: value = int(raw) except ValueError: raise HTTPException(400, f"{name} must be a number") from None if not lo <= value <= hi: raise HTTPException(400, f"{name} must be between {lo} and {hi}") return value @app.patch("/api/documents/{doc_id}") def patch_document(doc_id: int, patch: DocumentPatch) -> dict: """Update a document's metadata. The ingested artifacts are untouched; an emptied field clears its column (folder falls back under the ungrouped heading, filing fields drop out of the label).""" if patch.title is not None and not patch.title.strip(): raise HTTPException(400, "title cannot be empty") sets: List[LiteralString] = [] params: List[Any] = [] text_fields: List[tuple[LiteralString, Optional[str]]] = [ ("title", patch.title), ("folder", patch.folder), ("filing_type", patch.filing_type), ("period", patch.period), ] for column, raw in text_fields: if raw is not None: sets.append(column + " = %s") params.append(raw.strip() or None) if patch.year is not None: sets.append("year = %s") params.append(_parse_int("year", patch.year, 1900, 2100) if patch.year.strip() else None) if patch.version is not None: sets.append("version = %s") params.append(_parse_int("version", patch.version, 1, 999) if patch.version.strip() else None) if not sets: raise HTTPException(400, "nothing to update") query: LiteralString = ( "UPDATE ade_playground.documents SET " # noqa: S608 - fragments are literals + ", ".join(sets) + " WHERE id = %s RETURNING id" ) with db.connect() as conn: row = conn.execute(query, (*params, doc_id)).fetchone() if row is None: raise HTTPException(404, f"no document {doc_id}") return _doc_row(doc_id) def _ingest_busy(doc_id: int, doc_key: str) -> Optional[str]: """Why this document cannot be pulled out from under running work, or None. An ingest can touch it two ways: a replace targeting it, or an s3-sourced job that will resolve to the same bytes and key.""" for job in jobs.live_uploads(): if job.get("replace_id") == doc_id: return "a replacement for this document is still ingesting; wait for it to finish" if job.get("doc_key") == doc_key: return "an ingest of this document is still running; wait for it to finish" if BATCH.running_for(doc_id): return "a batch run against this document is still going; wait for it to finish" return None def _doc_files(doc_key: str) -> List[Path]: """The staged files a document's storage key names: the served PDF and the cached parse and figure JSONs in the data directory, and the uploaded original. The pipeline workdir is not in the list; `_remove_doc_files` removes it separately. The `uploads/.source.json` manifest is in neither.""" suffixes = ( ".pdf", ".ade.json", ".unified.json", ".tables.json", ".fusion.json", ".scanned-tables.json", ".figures.json", ".figure-values.json", ".charts.json", ) paths = [DATA_DIR / f"{doc_key}{s}" for s in suffixes] paths.append(UPLOADS / f"{doc_key}.pdf") return paths def _remove_doc_files(doc_key: str) -> int: removed = 0 for path in _doc_files(doc_key): if path.exists(): path.unlink() removed += 1 artifacts = UPLOADS / f"{doc_key}-artifacts" if artifacts.exists(): shutil.rmtree(artifacts) removed += 1 return removed @app.delete("/api/documents/{doc_id}") def delete_document(doc_id: int) -> dict: """Remove a document outright: the row (chunks, groundings, vectors and flags cascade with it), the staged files, and the upload artifacts. A hard delete, documented as such: there is no tombstone, and the cascade plus the file removal leave nothing for a re-ingest to resurrect. Refused with 409 while an ingest touching the document or a batch run against it is still going — the server owns no cancellation, so the work finishes or fails before the ground moves. """ doc = _doc_row(doc_id) busy = _ingest_busy(doc_id, doc["doc_key"]) if busy: raise HTTPException(409, busy) with db.connect() as conn: counts = conn.execute( """SELECT COUNT(DISTINCT c.id), COUNT(DISTINCT g.id) FROM ade_playground.documents d LEFT JOIN ade_playground.chunks c ON c.document_id = d.id LEFT JOIN ade_playground.groundings g ON g.document_id = d.id WHERE d.id = %s""", (doc_id,), ).fetchone() conn.execute("DELETE FROM ade_playground.documents WHERE id = %s", (doc_id,)) assert counts is not None removed_files = _remove_doc_files(doc["doc_key"]) logger.info( "document removed: {id} '{title}' ({chunks} chunks, {groundings} groundings, {files} files)", id=doc_id, title=doc["title"] or doc["filename"], chunks=counts[0], groundings=counts[1], files=removed_files, ) return {"removed": doc_id, "chunks": counts[0], "groundings": counts[1], "files": removed_files} @app.get("/api/document/{doc_id}/pdf") def pdf(doc_id: int) -> FileResponse: doc = _doc_row(doc_id) path = storage.local_pdf(doc["doc_key"]) if path is None: raise HTTPException(404, f"No PDF staged for document {doc_id}") return FileResponse(path, media_type="application/pdf") def _bbox_dict(raw: Any) -> dict: """A groundings bbox as a dict, whether the driver returned it parsed.""" return raw if isinstance(raw, dict) else json.loads(raw) def _box_contains(outer: dict, inner: dict, eps: float = 0.005) -> bool: """Whether `outer` wholly contains `inner`, with a hairline tolerance.""" return ( outer["left"] <= inner["left"] + eps and outer["top"] <= inner["top"] + eps and outer["right"] >= inner["right"] - eps and outer["bottom"] >= inner["bottom"] - eps ) def _resolve_refs(doc_id: int, cited_ids: List[str]) -> List[Reference]: """Map cited chunk/cell ids to page + bbox overlays via the groundings table.""" if not cited_ids: return [] with db.connect() as conn: rows = conn.execute( """SELECT g.ref_id, g.ref_type, g.page, g.bbox, g.status, g.note, g.cell_text, g.position FROM ade_playground.groundings g WHERE g.document_id = %s AND g.ref_id = ANY(%s)""", (doc_id, cited_ids), ).fetchall() # Preserve the model's citation order. by_id = {r[0]: r for r in rows} # A cited figure value with its own box makes citing its containing # picture redundant: the specific highlight is the answer's evidence, # the whole-picture box only buries it. covered_pictures = set() for r in by_id.values(): if r[1] != "figureValue" or r[3] is None: continue position = r[7] if isinstance(r[7], dict) else (json.loads(r[7]) if r[7] else {}) picture_ref = (position or {}).get("picture_ref") if picture_ref: covered_pictures.add(picture_ref) # A cited table cell with its own box makes citing its containing line # record redundant on the same terms: the cell is the answer's evidence, # the full-row band only buries it. The enclosing record's id is derived # from the ingest's naming — cell `t--` sits in line record # `t-line-`. covered_lines = set() for r in by_id.values(): if r[1] != "tableCell" or r[3] is None: continue position = r[7] if isinstance(r[7], dict) else (json.loads(r[7]) if r[7] else {}) row = (position or {}).get("row") if row is not None: covered_lines.add(f"{r[0].split('-', 1)[0]}-line-{row}") # The finest citation wins the overlay outright. A cited figure value or # table cell is the answer's evidence; any other cited region whose box # wholly contains it on the same page — the stat panel's group, a # wrapping paragraph — would only bury the tight highlight under a # bigger rectangle, so it is not drawn as a reference. fine_types = {"figureValue", "tableCell"} fine_boxes = [ ((r[2] or 0), _bbox_dict(r[3])) for r in by_id.values() if r[1] in fine_types and r[3] is not None ] def buries_fine(r: tuple) -> bool: if r[1] in fine_types or r[3] is None: return False outer = _bbox_dict(r[3]) return any(page == (r[2] or 0) and _box_contains(outer, fine) for page, fine in fine_boxes) refs: List[Reference] = [] for cid in cited_ids: r = by_id.get(cid) if not r: continue if r[1] == "picture" and r[0] in covered_pictures: continue if r[1] == "line_item" and r[0] in covered_lines: continue if buries_fine(r): continue page1 = (r[2] or 0) + 1 bbox = r[3] if isinstance(r[3], dict) else (json.loads(r[3]) if r[3] else None) status, note, cell_text = r[4], r[5], r[6] position = r[7] if isinstance(r[7], dict) else (json.loads(r[7]) if r[7] else None) reason, printed = _reason(status, note, cell_text) refs.append( Reference( ref_id=r[0], ref_type=r[1], page=page1, bbox=bbox, label=_label(r[1], page1, status), status=status, note=note, text=cell_text, row=(position or {}).get("row"), col=(position or {}).get("col"), flagged=bool(status and status in INSPECT_CODES), chart=(position or {}).get("chart"), segment=(position or {}).get("label"), reason=reason, printed=printed, ) ) return refs class Flag(BaseModel): """One item of the document's review queue: a cell whose provenance status is registered for inspection, or a table-level dropped-text flag.""" ref_id: str ref_type: Optional[str] # tableCell | figureValue | tableFlag page: int # 1-based bbox: Optional[dict] status: str note: Optional[str] = None text: str = "" # the flagged value, or the flag's quoted text reason: Optional[str] = None # the status in the reader's words @app.get("/api/document/{doc_id}/flags", response_model=List[Flag]) def document_flags(doc_id: int) -> List[Flag]: """Every flagged item for one document, in page order.""" with db.connect() as conn: rows = conn.execute( """SELECT g.ref_id, g.ref_type, g.page, g.bbox, g.status, g.note, g.cell_text FROM ade_playground.groundings g WHERE g.document_id = %s AND g.status = ANY(%s) ORDER BY g.page, g.ref_id""", (doc_id, list(INSPECT_CODES)), ).fetchall() return [ Flag( ref_id=r[0], ref_type=r[1], page=(r[2] or 0) + 1, bbox=r[3] if isinstance(r[3], dict) else (json.loads(r[3]) if r[3] else None), status=r[4], note=r[5], text=r[6] or "", reason=_reason(r[4], r[5], r[6])[0], ) for r in rows ] # ---------------------------------------------------------------- uploads # The jobs this process is running, in memory as the working copy. Every # stage transition is also written to the database (quber.playground.jobs), # which is what answers a status poll on any task and what outlives a task. JOBS: dict[str, dict] = {} _JOBS_LOCK = threading.Lock() _LOG_TAIL = 200 # lines kept per job # Every upload is admitted; this bounds how many run their pipelines at once. # On the developer host two concurrent extractions are two `quber fuse` # subprocesses on one CUDA device, each also making model calls. On the hosted # app the docling parse goes to the RunPod worker instead (see # `_run_hosted_extraction`). _UPLOAD_SLOTS = threading.Semaphore(get_settings().playground.upload_concurrency) # The ordered stages each pipeline moves through, ending in done or failed. # `waiting` is the admitted-but-queued state, ahead of the stages the pipeline # emits — without it a queued document looks stalled on its first step. PIPELINE_STAGES = { "fusion": ["uploaded", "waiting", "extracting", "reading figures", "ingesting"], } #: The fusion pipeline's ingestion tracks, keyed by model line. A track selects #: the whole line at once — the figure scan's model and the ingest's chunking #: stream — so the two are never mixed within one document. The model is the #: CLI spelling: `quber figure` resolves "dpt-3" to its pinned dated version, #: so the pin stays single-sourced in the scan module. A future line (dpt-4, #: dpt-5) is one more entry here. TRACKS = { "dpt-3": {"model": "dpt-3", "chunking": "grouped"}, "dpt-2": {"model": "dpt-2", "chunking": "flat"}, } DEFAULT_TRACK = "dpt-3" #: Filing types whose figure scan stays on nominated pages only. These are the #: long prose-and-table filings where full coverage buys mostly empty scans at #: real cost; every other type — and that includes free-form types this set has #: never heard of — scans every page, because a deck's content routinely lives #: on pages the parse cannot nominate. NOMINATED_ONLY_TYPES = {"10-Q", "10-K"} def _scan_breadth(job: Dict[str, Any], pdf: Path, workdir: Path) -> tuple[str, str]: """The figure scan's page breadth for one upload, and the reason for it. Typed documents are decided by their filing type alone. An untyped document is never allowed to default to full coverage — the type field is optional and routinely left blank, and a blank must not buy an 80-page scan by accident. A cover attribution agent reads the first page instead and decides presentation (all pages) against prose filing (nominated only); any failure in that read falls back to nominated-only, with the fallback stated in the returned reason so the record shows breadth was reduced and why. The verdict decides breadth only — it is never written to the document's filing type, which stays exactly as the user entered it. """ filing_type = (job.get("filing_type") or "").strip().upper() if filing_type: scope = "nominated" if filing_type in NOMINATED_ONLY_TYPES else "all" return scope, f"filing type {filing_type}" import asyncio import fitz from quber.agents.cover_attribution import get_cover_attributor from quber.files.pdf import render_page try: with fitz.open(pdf) as doc: page_count = doc.page_count cover = workdir / f"{pdf.stem}.cover.png" render_page(pdf, 1, 100, cover) verdict = asyncio.run(get_cover_attributor().attribute(cover.read_bytes(), page_count)) cover.unlink(missing_ok=True) if verdict.kind == "presentation": return "all", f"attributed presentation: {verdict.reason}" return "nominated", f"attributed filing: {verdict.reason}" except Exception as exc: return ( "nominated", f"attribution failed ({type(exc).__name__}: {exc}); nominated-only fallback", ) class JobStatus(BaseModel): job_id: str filename: str title: Optional[str] = None folder: Optional[str] = None # One of `stages`, or the terminal `done` / `failed`. The view reads the # stage names off this route rather than keeping a parallel set, so these # strings are the contract — the stepper draws `stages` and marks `stage` # against it, so the view learns a pipeline's stages from the job itself. stage: str stages: List[str] = PIPELINE_STAGES["fusion"] error: Optional[str] = None # Structured refusal for the states the modal renders specially: kind # "duplicate" (with the existing document to link to), "tuple", or # "busy". Set alongside `error` when an s3-sourced job fails admission — # a body upload hits the same checks synchronously as a 409 whose detail # carries this same shape. error_info: Optional[dict] = None # Set once the ingest completes: the new document's id. doc_id: Optional[int] = None log: List[str] def _run_stage(job: dict, stage: str, cmd: List[str]) -> None: job["stage"] = stage job["log"].append(f"--- {stage}: {' '.join(cmd)}") jobs.save_upload(job) proc = subprocess.Popen(cmd, cwd=REPO_ROOT, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True) assert proc.stdout is not None for line in proc.stdout: job["log"].append(line.rstrip()) if len(job["log"]) > _LOG_TAIL: del job["log"][:-_LOG_TAIL] proc.wait() # What the stage wrote is copied to the durable prefix whether it # succeeded or not: a failed stage's partial outputs are the diagnosis. storage.publish_workdir(job["doc_key"], UPLOADS / f"{job['doc_key']}-artifacts") jobs.save_upload(job) if proc.returncode != 0: raise RuntimeError(f"{stage} exited with code {proc.returncode}") def _run_hosted_extraction(job: dict, pdf: str, workdir: Path) -> None: """The extracting stage where this host has no GPU, in three steps run one after another: the table engine runs here as `quber table`, then the docling parse goes to the worker and is waited on, then `quber fuse` reads the two sets of artifacts from the workdir.""" doc_key = job["doc_key"] def log(line: str) -> None: job["log"].append(line) if len(job["log"]) > _LOG_TAIL: del job["log"][:-_LOG_TAIL] job["stage"] = "extracting" jobs.save_upload(job) _run_stage( job, "extracting", ["uv", "run", "quber", "table", pdf, "-o", str(workdir), "--llm-backend", "api"] ) gpu.parse(doc_key, workdir, log) jobs.save_upload(job) _run_stage( job, "extracting", ["uv", "run", "quber", "fuse", pdf, "--artifacts-dir", str(workdir), "-o", str(workdir)], ) def _await_slot(job: dict) -> None: """Hold until a pipeline slot frees. Admitted-but-waiting is its own stage: a queued document has something honest to show rather than appearing stalled on its first step.""" job["stage"] = "waiting" jobs.save_upload(job) _UPLOAD_SLOTS.acquire() def _admission_error(job: dict) -> Optional[dict]: """The reason this upload must not run, as the structured error_info the modal renders — or None. Checks, in order: the same bytes already in the library (or already being ingested), then the filing tuple against existing documents and in-flight jobs. A tuple match with the replace intent records the target on the job instead of refusing, unless the target is under running work.""" with db.connect() as conn: dup = conn.execute( "SELECT id, title, folder, filename FROM ade_playground.documents WHERE content_hash = %s", (job["content_hash"],), ).fetchone() if dup: title = dup[1] or dup[3] return { "kind": "duplicate", "message": f"This exact file is already in the library — '{title}' in {dup[2] or 'Ungrouped'}.", "existing": {"id": dup[0], "title": dup[1], "folder": dup[2]}, } others = [j for j in jobs.live_uploads() if j["job_id"] != job["job_id"]] if any(j.get("doc_key") == job["doc_key"] for j in others): return {"kind": "duplicate", "message": "This exact file is already being ingested."} key = metadata.tuple_key(job["folder"], job["filing_type"], job["year"], job["period"], job["version"]) if key is None: return None for j in others: other_key = metadata.tuple_key( j.get("folder"), j.get("filing_type"), j.get("year"), j.get("period"), j.get("version") ) if other_key == key: return { "kind": "tuple", "message": "A document with this filing metadata is already being ingested.", } with db.connect() as conn: rows = conn.execute("SELECT " + _DOC_COLUMNS + " FROM ade_playground.documents d").fetchall() match = next( ( d for d in (_doc_dict(r) for r in rows) if metadata.tuple_key(d["folder"], d["filing_type"], d["year"], d["period"], d["version"]) == key ), None, ) if match is None: return None if not job["replace"]: return { "kind": "tuple", "message": f"A {match['label']} already exists in {match['folder'] or 'Ungrouped'}.", "existing": {"id": match["id"], "title": match["title"], "folder": match["folder"]}, } busy = _ingest_busy(match["id"], match["doc_key"]) if busy: return {"kind": "busy", "message": f"Cannot replace: {busy}."} job["replace_id"] = match["id"] job["replace_key"] = match["doc_key"] return None def _finalize_job(job: dict) -> None: """Stamp the upload's metadata onto the freshly ingested row and, for a replace, swap the target out in the same transaction. Ingest recreates the document row, so this runs after it.""" replaced_key: Optional[str] = None with db.connect() as conn: with conn.transaction(): conn.execute( """UPDATE ade_playground.documents SET title = %s, folder = %s, filing_type = %s, year = %s, period = %s, version = %s, track = %s WHERE doc_key = %s""", ( job.get("title"), job.get("folder"), job.get("filing_type"), job.get("year"), job.get("period"), job.get("version"), job.get("track"), job["doc_key"], ), ) if job.get("replace_id") is not None: # A batch run started against the target while the new ingest # ran: leave both documents standing (the overlap flag names # them) rather than deleting under the run. if BATCH.running_for(job["replace_id"]): job["log"].append( "replace target is under a running batch; both documents kept " "(remove the old one from the library when the run finishes)" ) else: conn.execute("DELETE FROM ade_playground.documents WHERE id = %s", (job["replace_id"],)) replaced_key = job.get("replace_key") row = conn.execute( "SELECT id FROM ade_playground.documents WHERE doc_key = %s", (job["doc_key"],) ).fetchone() if replaced_key: _remove_doc_files(replaced_key) logger.info("replaced document {id} ({key})", id=job["replace_id"], key=replaced_key) job["doc_id"] = row[0] if row else None jobs.save_upload(job) def _admit_or_fail(job: dict, pdf_bytes: Optional[bytes], s3_uri: Optional[str]) -> Optional[str]: """Resolve the upload to local bytes, hash them, and run admission. Returns the staged PDF path to run the pipeline on, or None after marking the job failed. Each upload is admitted once. A body upload was hashed and admitted in the route, so it passes straight through. An s3 source is fetched, hashed and admitted here, because the route never had its bytes.""" if pdf_bytes is None: assert s3_uri is not None staged = UPLOADS / f"tmp-{job['job_id']}.pdf" try: fetch_file(s3_uri, staged) pdf_bytes = staged.read_bytes() except Exception as exc: job["stage"] = "failed" job["error"] = f"could not fetch {s3_uri}: {exc}" return None job["content_hash"] = hashlib.sha256(pdf_bytes).hexdigest() job["doc_key"] = job["content_hash"][:16] info = _admission_error(job) if info is not None: staged.unlink(missing_ok=True) job["stage"] = "failed" job["error"] = info["message"] job["error_info"] = info return None staged.rename(UPLOADS / f"{job['doc_key']}.pdf") return str(UPLOADS / f"{job['doc_key']}.pdf") def _run_job(job_id: str, pdf_bytes: Optional[bytes], s3_uri: Optional[str]) -> None: """The quber workflow: fuse, then read the figures, then ingest. The figure stage is part of the workflow rather than an option on it. Without it a document arrives holding only what the text layer gave up: a chart's plotted values are absent, a table printed as an image stands as the parse's own reading of it, and the notes that qualify a figure are attached to nothing. The standalone ADE path is a different pipeline and is unaffected. The stage runs on every upload, and it spends: each page it scans is paid for. `_scan_breadth` decides which pages. At nominated breadth only the pages the parse nominates are scanned, so a document with no nominated page costs nothing. At all-pages breadth every page is scanned. An artifact file the run did not write, such as the scanned tables of a document with no table captured, the ingest reads as "the run produced none of that kind" rather than as an error. """ job = JOBS[job_id] try: _await_slot(job) pdf = _admit_or_fail(job, pdf_bytes, s3_uri) if pdf is None: return doc_key = job["doc_key"] workdir = UPLOADS / f"{doc_key}-artifacts" workdir.mkdir(parents=True, exist_ok=True) base = doc_key # The PDF and its source manifest reach the durable prefix before any # stage runs, and every stage's outputs follow it as the stage ends. manifest = storage.write_source_manifest( doc_key, filename=job["filename"], content_hash=job["content_hash"], source_uri=s3_uri ) storage.publish(doc_key, [Path(pdf), manifest]) track = TRACKS[job.get("track") or DEFAULT_TRACK] # Breadth is decided before anything runs and recorded on the job, so # the run's record always says how many pages were in play and why — # server-side only, never part of a client payload. pages, breadth_reason = _scan_breadth(job, Path(pdf), workdir) job["scan_pages"] = pages job["scan_pages_reason"] = breadth_reason job["source_uri"] = s3_uri jobs.save_upload(job) logger.info("Upload {}: figure scan breadth {} ({})", job_id, pages, breadth_reason) if gpu.configured(): _run_hosted_extraction(job, pdf, workdir) else: _run_stage(job, "extracting", ["uv", "run", "quber", "fuse", pdf, "-o", str(workdir)]) _run_stage( job, "reading figures", [ "uv", "run", "quber", "figure", pdf, "--parse", str(workdir / f"{base}.unified.json"), "-o", str(workdir), "--model", track["model"], "--pages", pages, ], ) _run_stage( job, "ingesting", [ "uv", "run", "python", "-m", "quber.playground.ingest_fusion", "--doc-key", doc_key, "--filename", job["filename"], "--unified-json", str(workdir / f"{base}.unified.json"), "--tables-json", str(workdir / f"{base}.tables.json"), "--fusion-json", str(workdir / f"{base}.fusion.json"), "--scanned-tables-json", str(workdir / f"{base}.scanned-tables.json"), "--figures-json", str(workdir / f"{base}.figures.json"), "--figure-values-json", str(workdir / f"{base}.figure-values.json"), "--pdf", pdf, "--chunking", track["chunking"], ], ) _finalize_job(job) job["stage"] = "done" except Exception as exc: job["stage"] = "failed" job["error"] = str(exc) finally: jobs.save_upload(job) _UPLOAD_SLOTS.release() @app.post("/api/upload", response_model=JobStatus) async def upload( request: Request, filename: str, s3_uri: Optional[str] = None, title: Optional[str] = None, folder: Optional[str] = None, filing_type: Optional[str] = None, year: Optional[int] = None, period: Optional[str] = None, version: Optional[int] = None, replace: bool = False, track: str = DEFAULT_TRACK, ) -> JobStatus: """Run the fusion pipeline on a PDF sent as the raw request body, or on an s3:// object named by `s3_uri` (body then stays empty). `track` selects the model line the document is processed on — see TRACKS. `filename` is the original name, kept for presentation and never for identity. The metadata fields are how the library presents and keys the document; this route and the PATCH are the only places they are ever set. `replace` carries the modal's confirmed intent to swap out the document whose filing tuple this upload matches. Every admitted upload waits in the `waiting` stage until a pipeline slot frees. """ if track not in TRACKS: raise HTTPException(400, f"track must be one of {sorted(TRACKS)}") body: Optional[bytes] = None if s3_uri is not None: if not s3_uri.startswith("s3://") or not s3_uri.lower().endswith(".pdf"): raise HTTPException(400, "s3_uri must be an s3:// URI ending in .pdf") else: body = await request.body() if not body.startswith(b"%PDF"): raise HTTPException(400, "request body is not a PDF") job_id = uuid.uuid4().hex[:12] job: dict = { "job_id": job_id, "filename": filename, "title": title, "folder": folder, "filing_type": filing_type, "year": year, "period": period, "version": version, "replace": replace, "track": track, "replace_id": None, "replace_key": None, "doc_key": None, "content_hash": None, "doc_id": None, "stage": "uploaded", "stages": PIPELINE_STAGES["fusion"], "error": None, "error_info": None, "log": [], } if body is not None: # The bytes are in hand: hash and run admission now, so a duplicate # or an unconfirmed tuple match is a synchronous refusal the modal # shows in place — not a job that fails a poll later. job["content_hash"] = hashlib.sha256(body).hexdigest() job["doc_key"] = job["content_hash"][:16] info = _admission_error(job) if info is not None: raise HTTPException(409, info) UPLOADS.mkdir(parents=True, exist_ok=True) (UPLOADS / f"{job['doc_key']}.pdf").write_bytes(body) else: UPLOADS.mkdir(parents=True, exist_ok=True) with _JOBS_LOCK: JOBS[job_id] = job jobs.save_upload(job) threading.Thread(target=_run_job, args=(job_id, body, s3_uri), daemon=True).start() return JobStatus(**{**job, "log": []}) @app.get("/api/jobs/{job_id}", response_model=JobStatus) def job_status(job_id: str, tail: int = 12) -> JobStatus: """The job as this process holds it when it is running it, else as the database records it, which is how a poll lands on any task.""" job = JOBS.get(job_id) or jobs.load_upload(job_id) if not job: raise HTTPException(404, f"no job {job_id}") return JobStatus(**{**job, "log": job["log"][-tail:]}) def _readable(payload) -> str: """What to put in `answer` for a consumer that only reads text. A figure renders as the figure, so the text field of a value-seeking answer is the value and not a sentence about it. """ if payload.kind == "prose": return payload.text if payload.kind == "unanswerable": return payload.reason if payload.kind == "scalar": return payload.value return json.dumps(payload.model_dump(mode="json"), indent=2) def cited_ids_source_first(payload, cited_ids: List[str]) -> List[str]: """The cited ids with the answer's own source leading. A value payload names the single cell or chunk it was read from in `source_id`. The UI activates the first reference that carries a box, so whatever id leads this list becomes the default highlight — and the value's own cell is the highlight the answer earned, not the wider line record the model happened to cite first. Payloads without a source (prose, unanswerable, series) keep the model's citation order. """ source_id = getattr(payload, "source_id", None) if not source_id: return list(cited_ids) return [source_id] + [cid for cid in cited_ids if cid != source_id] def _chat_response(doc_id: int, result, chunks) -> ChatResponse: payload = result.payload return ChatResponse( answer=_readable(payload), shape=payload.kind, value=payload.model_dump(mode="json"), references=_resolve_refs(doc_id, cited_ids_source_first(payload, result.cited_ids)), retrieved=[ RetrievedInfo( chunk_id=c.chunk_id, chunk_type=c.chunk_type, page=c.page + 1, score=round(c.score, 4) ) for c in chunks ], ) @app.post("/api/chat", response_model=ChatResponse) async def chat(req: ChatRequest) -> ChatResponse: doc = _doc_row(req.doc_id) document_key.set(doc["doc_key"]) chunks = await retrieve(doc["doc_key"], req.question, k=req.k) if not chunks: raise HTTPException(404, f"No chunks for document {req.doc_id} — is it ingested?") result = await agent_answer(req.question, chunks, req.want, DocumentIdentity.from_row(doc)) return _chat_response(req.doc_id, result, chunks) @app.post("/api/chat/stream") async def chat_stream(req: ChatRequest) -> StreamingResponse: """Same contract as /api/chat, streamed: server-sent events carrying the answer text as it generates ('delta'), then the full response ('done').""" doc = _doc_row(req.doc_id) document_key.set(doc["doc_key"]) chunks = await retrieve(doc["doc_key"], req.question, k=req.k) if not chunks: raise HTTPException(404, f"No chunks for document {req.doc_id} — is it ingested?") identity = DocumentIdentity.from_row(doc) async def events(): try: async for kind, payload in agent_answer_stream(req.question, chunks, req.want, identity): if kind == "delta": yield f"data: {json.dumps({'type': 'delta', 'answer': payload})}\n\n" else: body = _chat_response(req.doc_id, payload, chunks).model_dump() yield f"data: {json.dumps({'type': 'done', 'payload': body})}\n\n" except Exception as exc: yield f"data: {json.dumps({'type': 'error', 'detail': str(exc)})}\n\n" return StreamingResponse(events(), media_type="text/event-stream") # ---------------------------------------------------------------- batch BATCH_STORE = jobs.BatchStore() BATCH = BatchRunner(store=BATCH_STORE) class BatchRequest(BaseModel): doc_id: int # The batch screen's prompt window, verbatim: one question per line. The # server parses it so there is exactly one reading of that format. questions: str k: int = 10 # One expectation for the whole run; see ChatRequest.want. want: Literal["auto", "value", "text"] = "auto" @app.post("/api/batch", response_model=RunState) async def batch_start(req: BatchRequest) -> RunState: """Start a run. The work continues if the page is closed; the page keeps the returned job id in sessionStorage and re-attaches through the state and events routes.""" questions = parse_questions(req.questions) if not questions: raise HTTPException(400, "no questions — the batch window takes one per line") doc = _doc_row(req.doc_id) doc_key = doc["doc_key"] # Set before the probe so its selection run carries the document too; the # run task copies this context when it is created. document_key.set(doc_key) probe = await retrieve(doc_key, questions[0], k=1) if not probe: raise HTTPException(404, f"No chunks for document {req.doc_id} — is it ingested?") identity = DocumentIdentity.from_row(doc) async def retrieve_one(question: str) -> List[RetrievedChunk]: return await retrieve(doc_key, question, k=req.k) async def answer_one(question: str, chunks: List[RetrievedChunk]) -> Dict[str, Any]: result = await agent_answer(question, chunks, req.want, identity) return _chat_response(req.doc_id, result, chunks).model_dump(mode="json") return BATCH.start(req.doc_id, questions, req.want, retrieve_one, answer_one) def _run_state(job_id: str) -> RunState: state = BATCH.get(job_id) if state is None: # Neither this process's memory nor the database holds the id, so the # page kept an id nobody can answer for. The page reports the run as # gone; it is never silence. raise HTTPException(404, f"no run {job_id} — the server no longer has it") return state @app.get("/api/batch/{job_id}", response_model=RunState) def batch_state(job_id: str) -> RunState: return _run_state(job_id) @app.get("/api/batch/{job_id}/events") async def batch_events(job_id: str) -> StreamingResponse: """One server-sent event per question completing — index, total, outcome — then a final `done` event. Watching a run and re-attaching to one are different needs; this is the watching route.""" _run_state(job_id) async def stream(): async for event in BATCH.events(job_id): yield f"data: {json.dumps(event)}\n\n" return StreamingResponse(stream(), media_type="text/event-stream") def _export_stem(state: RunState) -> str: """The export filename's document part, down the label precedence. A run whose document was removed mid-session still exports; it just gets the generic stem.""" try: doc = _doc_row(state.doc_id) except HTTPException: return "document" return metadata.export_stem(doc["label"], doc["title"], doc["filename"]) @app.get("/api/batch/{job_id}/export.xlsx") def batch_export_xlsx(job_id: str) -> Response: state = _run_state(job_id) return Response( export.write_xlsx(state), media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", headers={"Content-Disposition": f'attachment; filename="batch-{_export_stem(state)}-{job_id}.xlsx"'}, ) @app.get("/api/batch/{job_id}/export.csv") def batch_export_csv(job_id: str) -> Response: state = _run_state(job_id) return Response( export.write_csv(state), media_type="text/csv; charset=utf-8", headers={"Content-Disposition": f'attachment; filename="batch-{_export_stem(state)}-{job_id}.csv"'}, )