Coverage for src / quber / playground / app.py: 43%
602 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-09-23 22:14 -0400
1"""FastAPI backend for the answering playground.
3Serves the single-page UI and these endpoints:
4 GET /api/documents -> list ingested documents: filing metadata,
5 derived label, overlap and review flags
6 GET /api/document/{id}/pdf -> the source PDF (for PDF.js)
7 PATCH /api/documents/{id} -> edit metadata (title, folder, filing
8 fields); organization only, never identity
9 DELETE /api/documents/{id} -> remove the document outright (hard
10 delete; 409 while its ingest or a batch
11 run against it is still going)
12 POST /api/chat -> retrieve -> Opus answer -> grounded refs
13 POST /api/upload -> raw PDF body; runs the quber workflow
14 (fuse -> figure -> ingest) as a job
15 GET /api/jobs/{job_id} -> job stage (uploaded/waiting/extracting/
16 reading figures/ingesting/done/failed)
17 plus a tail of its log
18 POST /api/batch -> start a run of many questions -> job id
19 GET /api/batch/{id} -> the run's state and rows finished so far
20 GET /api/batch/{id}/events -> SSE, one event per question completing
21 GET /api/batch/{id}/export.xlsx and .csv -> the run as a workbook / text
23Document identity is the integer id plus the content hash. The storage key
24(`doc_key`, the hash's first 16 hex chars) names the staged files and is what
25ingest subprocesses take; it never reaches the UI. Labels are derived from the
26filing metadata at read time — see `quber.playground.metadata`.
28Uploads take the PDF as the raw request body (no multipart dependency). The
29bytes are hashed before anything runs: the same bytes as an existing document
30are rejected outright (409, kind "duplicate"), and a filing-tuple match is
31refused (409, kind "tuple") unless the request carries the replace intent —
32then the swap happens when the new ingest completes. An s3:// source cannot
33be hashed before its job stages it, so those checks run inside the job and a
34failure carries the same structured `error_info`. Every upload is admitted;
35a semaphore bounds how many run their pipelines at once
36(`settings.playground.upload_concurrency`), and a job holding no slot yet
37reports the `waiting` stage rather than appearing stalled on its first step.
39Run:
40 uv run quber playground
41"""
43from __future__ import annotations
45import hashlib
46import json
47import re
48import shutil
49import subprocess
50import threading
51import uuid
52from pathlib import Path
53from typing import Any, Dict, List, Literal, LiteralString, Optional
55from fastapi import FastAPI, HTTPException, Request
56from fastapi.responses import FileResponse, HTMLResponse, Response, StreamingResponse
57from loguru import logger
58from pydantic import BaseModel
60from quber.files.cache import fetch_file
61from quber.playground import auth, db, embedding, export, gpu, jobs, metadata, storage
62from quber.playground.answers.document import DocumentIdentity
63from quber.playground.answers.expectation import answer as agent_answer
64from quber.playground.answers.expectation import answer_stream as agent_answer_stream
65from quber.playground.batch import BatchRunner, RunState, parse_questions
66from quber.playground.retrieval import RetrievedChunk, retrieve
67from quber.playground.tracing import document_key
68from quber.settings import get_settings
70STATIC = Path(__file__).with_name("static")
71DATA_DIR = get_settings().playground.data_dir
72UPLOADS = DATA_DIR / "uploads"
73REPO_ROOT = Path(__file__).parents[3]
75app = FastAPI(title="Quber Playground")
77# On when the login settings are present, which is the hosted playground;
78# off on the developer host. A partial configuration is refused here.
79LOGIN_ENABLED = auth.install(app, get_settings().playground)
81# Set once the embedding model has loaded. The health route reports the app
82# ready only after this, so a request handed off by the load balancer never
83# lands on a task that would stall its first question on the model load.
84_EMBEDDING_READY = threading.Event()
87@app.on_event("startup")
88def warm_embedding_model() -> None:
89 """Load the embedding model in the background at startup. It otherwise
90 loads lazily on the first question, costing that question ~5 seconds."""
92 def warm() -> None:
93 embedding.embed_query("warm up")
94 _EMBEDDING_READY.set()
96 threading.Thread(target=warm, daemon=True).start()
99@app.on_event("startup")
100def open_job_records() -> None:
101 """The job tables exist before anything writes to them, and this process
102 starts beating for the jobs it runs."""
103 jobs.ensure_tables()
104 jobs.start_heartbeat(
105 upload_ids=lambda: [j["job_id"] for j in JOBS.values() if j["stage"] not in jobs.TERMINAL_STAGES],
106 run_ids=BATCH.running_ids,
107 store=BATCH_STORE,
108 )
111@app.get("/healthz", include_in_schema=False)
112def healthz() -> Response:
113 """Readiness, not liveness: 200 only once the database answers and the
114 embedding model is loaded, 503 before. The load balancer's target group
115 and the scale-to-zero startup page both poll this, so a 200 means the
116 app can take a question, not merely that the process is up."""
117 if not _EMBEDDING_READY.is_set():
118 return Response(
119 '{"status":"loading embedding model"}', status_code=503, media_type="application/json"
120 )
121 try:
122 with db.connect() as conn:
123 conn.execute("SELECT 1").fetchone()
124 except Exception as exc:
125 return Response(
126 json.dumps({"status": f"database unreachable: {type(exc).__name__}"}),
127 status_code=503,
128 media_type="application/json",
129 )
130 return Response('{"status":"ok"}', media_type="application/json")
133class ChatRequest(BaseModel):
134 doc_id: int
135 question: str
136 k: int = 10
137 # What the asker expects back: "auto" lets the model choose the shape,
138 # "value" forces a figure, "text" forces an explanation.
139 want: Literal["auto", "value", "text"] = "auto"
142class Reference(BaseModel):
143 ref_id: str
144 ref_type: Optional[str]
145 page: int # 1-based, to match the viewer
146 bbox: Optional[dict]
147 label: str
148 # Cell-level provenance, the GroundedCell attributes established at
149 # reconciliation (fusion documents only; None on ADE references):
150 status: Optional[str] = None # provenance code, e.g. 'reconciled'
151 note: Optional[str] = None # status inspector's evidence line, if any
152 text: Optional[str] = None # the served cell value
153 row: Optional[int] = None # 0-based position in the corrected table
154 col: Optional[int] = None
155 flagged: bool = False # True when the status is registered for review
156 # Figure-value provenance, from the grounding's position (None elsewhere):
157 chart: Optional[str] = None # the chart the value is printed in
158 segment: Optional[str] = None # what the value is labeled as on the chart
159 # The status in the reader's words, from the registry's label for the
160 # code. The transcript prints this, never the code. A misread carries
161 # the figure the page prints so the line can name both numbers.
162 reason: Optional[str] = None
163 printed: Optional[str] = None
166class RetrievedInfo(BaseModel):
167 chunk_id: str
168 chunk_type: str
169 page: int
170 score: float
173class ChatResponse(BaseModel):
174 answer: str
175 references: List[Reference]
176 retrieved: List[RetrievedInfo]
177 # The answer as a value rather than as text. `shape` names which payload
178 # came back — scalar, series, grid, prose, unanswerable — and `value` is
179 # that payload. `answer` stays populated with something readable so an
180 # existing consumer of this endpoint keeps working.
181 shape: str = "prose"
182 value: Optional[dict] = None
185# Statuses the provenance registry marks for review; cited cells carrying one
186# are served like any other but highlighted so the user knows they are flagged.
187def _inspect_codes() -> frozenset[str]:
188 from quber.core.extractors.base import CELL_STATUS_REFERENCE
190 return frozenset(s.code for s in CELL_STATUS_REFERENCE if s.inspect)
193def _status_reasons() -> Dict[str, str]:
194 from quber.core.extractors.base import CELL_STATUS_REFERENCE
196 return {s.code: s.label for s in CELL_STATUS_REFERENCE}
199INSPECT_CODES = _inspect_codes()
200STATUS_REASONS = _status_reasons()
202# The misread note is written by the figure-value reader in one of two
203# shapes, "the page prints '6%' where this value should appear" or "fragment
204# at the value's position prints '~ 87%'"; the figure inside the quotes is
205# what the transcript's line needs.
206_PRINTED_IN_NOTE = re.compile(r"prints '([^']*)'")
209def _reason(
210 status: Optional[str], note: Optional[str], served: Optional[str]
211) -> tuple[Optional[str], Optional[str]]:
212 """The plain-language line for a status, and the printed figure a misread names."""
213 if not status:
214 return None, None
215 printed = None
216 if status == "value_misread" and note:
217 m = _PRINTED_IN_NOTE.search(note)
218 if m:
219 printed = m.group(1)
220 if printed and served:
221 return f"The page prints {printed} here, not {served}. Check which is right.", printed
222 return STATUS_REASONS.get(status, status), printed
225def _label(ref_type: Optional[str], page1: int, status: Optional[str] = None) -> str:
226 """What the source is, in the reader's words. The status never rides
227 the label: it has its own field, and its plain-language line is `reason`."""
228 if ref_type == "tableCell":
229 return f"Page {page1}, table cell"
230 if ref_type in ("table", "chunkTable"):
231 return f"Page {page1}, table"
232 if ref_type == "figureValue":
233 return f"Page {page1}, chart value"
234 if ref_type == "line_item":
235 return f"Page {page1}, table row"
236 if ref_type == "picture":
237 return f"Page {page1}, figure"
238 return f"Page {page1}, text"
241def _stamped_index() -> str:
242 """The index shell with content-hashed asset links.
244 A browser that cached a previous deploy of app.js kept serving it on a
245 plain reload, because the asset URLs never change. Stamping each mutable
246 asset with a hash of its current bytes makes any change a new URL; the
247 pinned vendor files stay as they are. Hashing runs per request — the
248 files are small and the reload server would otherwise need cache
249 invalidation of its own.
250 """
251 html = (STATIC / "index.html").read_text()
252 for name in ("styles.css", "app.css", "app.js"):
253 digest = hashlib.sha256((STATIC / name).read_bytes()).hexdigest()[:8]
254 html = html.replace(f"/static/{name}", f"/static/{name}?v={digest}")
255 return html
258@app.get("/", response_class=HTMLResponse)
259def index() -> HTMLResponse:
260 # no-cache makes the shell revalidate on every load, so the stamped links
261 # inside it are always current; the assets themselves may cache freely.
262 return HTMLResponse(_stamped_index(), headers={"Cache-Control": "no-cache"})
265@app.get("/static/{name:path}")
266def static_asset(name: str) -> FileResponse:
267 # Nested paths are served (the design-system token stylesheets live in
268 # static/tokens/), but the resolved file must still sit inside static/ so
269 # that "../" segments cannot walk out of it.
270 path = STATIC / name
271 if not path.is_file() or not path.resolve().is_relative_to(STATIC.resolve()):
272 raise HTTPException(404, f"no asset {name}")
273 return FileResponse(path)
276_DOC_COLUMNS: LiteralString = (
277 "d.id, d.doc_key, d.filename, d.page_count, d.ade_version, d.title, d.folder, "
278 "d.filing_type, d.year, d.period, d.version"
279)
282# The processing track (documents.track) is deliberately NOT served here or
283# anywhere the browser can reach: naming vendor model internals in a client
284# payload makes them scrapeable. It is written at upload finalize and read in
285# the database only.
286def _doc_dict(r: "tuple[Any, ...]") -> Dict[str, Any]:
287 doc = {
288 "id": r[0],
289 "doc_key": r[1],
290 "filename": r[2],
291 "page_count": r[3],
292 "format": "fusion" if r[4] == "quber-fusion" else "ade",
293 "title": r[5],
294 "folder": r[6],
295 "filing_type": r[7],
296 "year": r[8],
297 "period": r[9],
298 "version": r[10],
299 }
300 doc["label"] = metadata.filing_label(r[7], r[8], r[9], r[10])
301 return doc
304def _doc_row(doc_id: int) -> Dict[str, Any]:
305 with db.connect() as conn:
306 r = conn.execute(
307 "SELECT " + _DOC_COLUMNS + " FROM ade_playground.documents d WHERE d.id = %s",
308 (doc_id,),
309 ).fetchone()
310 if r is None:
311 raise HTTPException(404, f"no document {doc_id}")
312 return _doc_dict(r)
315@app.get("/api/documents")
316def documents() -> list[dict]:
317 # The flag count rides the list because the alternative is one
318 # /api/document/{id}/flags call per document, each returning every flag
319 # object with its bounding box in order to take a length.
320 with db.connect() as conn:
321 rows = conn.execute(
322 "SELECT "
323 + _DOC_COLUMNS
324 + """, COUNT(g.id) FILTER (WHERE g.status = ANY(%s)) AS flags
325 FROM ade_playground.documents d
326 LEFT JOIN ade_playground.groundings g ON g.document_id = d.id
327 GROUP BY d.id ORDER BY lower(coalesce(d.title, d.filename))""",
328 (list(INSPECT_CODES),),
329 ).fetchall()
330 docs = []
331 for r in rows:
332 doc = _doc_dict(r)
333 # The storage key stays server-side; nothing in the UI needs it.
334 del doc["doc_key"]
335 doc["flags"] = r[11]
336 docs.append(doc)
337 metadata.annotate_overlaps(docs)
338 return docs
341class DocumentPatch(BaseModel):
342 """The editable metadata: presentation (title, folder) and the filing
343 fields. Edits change organization and labels only — never identity, and
344 never blocked by collisions; two documents landing on one tuple simply
345 both carry the possible-duplicate flag. Fields arrive as strings so an
346 emptied input clears its column; year and version parse to ints."""
348 title: Optional[str] = None
349 folder: Optional[str] = None
350 filing_type: Optional[str] = None
351 year: Optional[str] = None
352 period: Optional[str] = None
353 version: Optional[str] = None
356def _parse_int(name: str, raw: str, lo: int, hi: int) -> int:
357 try:
358 value = int(raw)
359 except ValueError:
360 raise HTTPException(400, f"{name} must be a number") from None
361 if not lo <= value <= hi:
362 raise HTTPException(400, f"{name} must be between {lo} and {hi}")
363 return value
366@app.patch("/api/documents/{doc_id}")
367def patch_document(doc_id: int, patch: DocumentPatch) -> dict:
368 """Update a document's metadata. The ingested artifacts are untouched;
369 an emptied field clears its column (folder falls back under the
370 ungrouped heading, filing fields drop out of the label)."""
371 if patch.title is not None and not patch.title.strip():
372 raise HTTPException(400, "title cannot be empty")
373 sets: List[LiteralString] = []
374 params: List[Any] = []
375 text_fields: List[tuple[LiteralString, Optional[str]]] = [
376 ("title", patch.title),
377 ("folder", patch.folder),
378 ("filing_type", patch.filing_type),
379 ("period", patch.period),
380 ]
381 for column, raw in text_fields:
382 if raw is not None:
383 sets.append(column + " = %s")
384 params.append(raw.strip() or None)
385 if patch.year is not None:
386 sets.append("year = %s")
387 params.append(_parse_int("year", patch.year, 1900, 2100) if patch.year.strip() else None)
388 if patch.version is not None:
389 sets.append("version = %s")
390 params.append(_parse_int("version", patch.version, 1, 999) if patch.version.strip() else None)
391 if not sets:
392 raise HTTPException(400, "nothing to update")
393 query: LiteralString = (
394 "UPDATE ade_playground.documents SET " # noqa: S608 - fragments are literals
395 + ", ".join(sets)
396 + " WHERE id = %s RETURNING id"
397 )
398 with db.connect() as conn:
399 row = conn.execute(query, (*params, doc_id)).fetchone()
400 if row is None:
401 raise HTTPException(404, f"no document {doc_id}")
402 return _doc_row(doc_id)
405def _ingest_busy(doc_id: int, doc_key: str) -> Optional[str]:
406 """Why this document cannot be pulled out from under running work, or
407 None. An ingest can touch it two ways: a replace targeting it, or an
408 s3-sourced job that will resolve to the same bytes and key."""
409 for job in jobs.live_uploads():
410 if job.get("replace_id") == doc_id:
411 return "a replacement for this document is still ingesting; wait for it to finish"
412 if job.get("doc_key") == doc_key:
413 return "an ingest of this document is still running; wait for it to finish"
414 if BATCH.running_for(doc_id):
415 return "a batch run against this document is still going; wait for it to finish"
416 return None
419def _doc_files(doc_key: str) -> List[Path]:
420 """Every staged file a document owns: the served PDF, the cached parse
421 and figure JSONs beside it, the uploaded original, and the pipeline
422 workdir."""
423 suffixes = (
424 ".pdf",
425 ".ade.json",
426 ".unified.json",
427 ".tables.json",
428 ".fusion.json",
429 ".scanned-tables.json",
430 ".figures.json",
431 ".figure-values.json",
432 ".charts.json",
433 )
434 paths = [DATA_DIR / f"{doc_key}{s}" for s in suffixes]
435 paths.append(UPLOADS / f"{doc_key}.pdf")
436 return paths
439def _remove_doc_files(doc_key: str) -> int:
440 removed = 0
441 for path in _doc_files(doc_key):
442 if path.exists():
443 path.unlink()
444 removed += 1
445 artifacts = UPLOADS / f"{doc_key}-artifacts"
446 if artifacts.exists():
447 shutil.rmtree(artifacts)
448 removed += 1
449 return removed
452@app.delete("/api/documents/{doc_id}")
453def delete_document(doc_id: int) -> dict:
454 """Remove a document outright: the row (chunks, groundings, vectors and
455 flags cascade with it), the staged files, and the upload artifacts.
457 A hard delete, documented as such: there is no tombstone, and the
458 cascade plus the file removal leave nothing for a re-ingest to
459 resurrect. Refused with 409 while an ingest touching the document or a
460 batch run against it is still going — the server owns no cancellation,
461 so the work finishes or fails before the ground moves.
462 """
463 doc = _doc_row(doc_id)
464 busy = _ingest_busy(doc_id, doc["doc_key"])
465 if busy:
466 raise HTTPException(409, busy)
467 with db.connect() as conn:
468 counts = conn.execute(
469 """SELECT COUNT(DISTINCT c.id), COUNT(DISTINCT g.id)
470 FROM ade_playground.documents d
471 LEFT JOIN ade_playground.chunks c ON c.document_id = d.id
472 LEFT JOIN ade_playground.groundings g ON g.document_id = d.id
473 WHERE d.id = %s""",
474 (doc_id,),
475 ).fetchone()
476 conn.execute("DELETE FROM ade_playground.documents WHERE id = %s", (doc_id,))
477 assert counts is not None
478 removed_files = _remove_doc_files(doc["doc_key"])
479 logger.info(
480 "document removed: {id} '{title}' ({chunks} chunks, {groundings} groundings, {files} files)",
481 id=doc_id,
482 title=doc["title"] or doc["filename"],
483 chunks=counts[0],
484 groundings=counts[1],
485 files=removed_files,
486 )
487 return {"removed": doc_id, "chunks": counts[0], "groundings": counts[1], "files": removed_files}
490@app.get("/api/document/{doc_id}/pdf")
491def pdf(doc_id: int) -> FileResponse:
492 doc = _doc_row(doc_id)
493 path = storage.local_pdf(doc["doc_key"])
494 if path is None:
495 raise HTTPException(404, f"No PDF staged for document {doc_id}")
496 return FileResponse(path, media_type="application/pdf")
499def _bbox_dict(raw: Any) -> dict:
500 """A groundings bbox as a dict, whether the driver returned it parsed."""
501 return raw if isinstance(raw, dict) else json.loads(raw)
504def _box_contains(outer: dict, inner: dict, eps: float = 0.005) -> bool:
505 """Whether `outer` wholly contains `inner`, with a hairline tolerance."""
506 return (
507 outer["left"] <= inner["left"] + eps
508 and outer["top"] <= inner["top"] + eps
509 and outer["right"] >= inner["right"] - eps
510 and outer["bottom"] >= inner["bottom"] - eps
511 )
514def _resolve_refs(doc_id: int, cited_ids: List[str]) -> List[Reference]:
515 """Map cited chunk/cell ids to page + bbox overlays via the groundings table."""
516 if not cited_ids:
517 return []
518 with db.connect() as conn:
519 rows = conn.execute(
520 """SELECT g.ref_id, g.ref_type, g.page, g.bbox, g.status, g.note,
521 g.cell_text, g.position
522 FROM ade_playground.groundings g
523 WHERE g.document_id = %s AND g.ref_id = ANY(%s)""",
524 (doc_id, cited_ids),
525 ).fetchall()
526 # Preserve the model's citation order.
527 by_id = {r[0]: r for r in rows}
528 # A cited figure value with its own box makes citing its containing
529 # picture redundant: the specific highlight is the answer's evidence,
530 # the whole-picture box only buries it.
531 covered_pictures = set()
532 for r in by_id.values():
533 if r[1] != "figureValue" or r[3] is None:
534 continue
535 position = r[7] if isinstance(r[7], dict) else (json.loads(r[7]) if r[7] else {})
536 picture_ref = (position or {}).get("picture_ref")
537 if picture_ref:
538 covered_pictures.add(picture_ref)
539 # A cited table cell with its own box makes citing its containing line
540 # record redundant on the same terms: the cell is the answer's evidence,
541 # the full-row band only buries it. The enclosing record's id is derived
542 # from the ingest's naming — cell `t<i>-<row>-<col>` sits in line record
543 # `t<i>-line-<row>`.
544 covered_lines = set()
545 for r in by_id.values():
546 if r[1] != "tableCell" or r[3] is None:
547 continue
548 position = r[7] if isinstance(r[7], dict) else (json.loads(r[7]) if r[7] else {})
549 row = (position or {}).get("row")
550 if row is not None:
551 covered_lines.add(f"{r[0].split('-', 1)[0]}-line-{row}")
552 # The finest citation wins the overlay outright. A cited figure value or
553 # table cell is the answer's evidence; any other cited region whose box
554 # wholly contains it on the same page — the stat panel's group, a
555 # wrapping paragraph — would only bury the tight highlight under a
556 # bigger rectangle, so it is not drawn as a reference.
557 fine_types = {"figureValue", "tableCell"}
558 fine_boxes = [
559 ((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
560 ]
562 def buries_fine(r: tuple) -> bool:
563 if r[1] in fine_types or r[3] is None:
564 return False
565 outer = _bbox_dict(r[3])
566 return any(page == (r[2] or 0) and _box_contains(outer, fine) for page, fine in fine_boxes)
568 refs: List[Reference] = []
569 for cid in cited_ids:
570 r = by_id.get(cid)
571 if not r:
572 continue
573 if r[1] == "picture" and r[0] in covered_pictures:
574 continue
575 if r[1] == "line_item" and r[0] in covered_lines:
576 continue
577 if buries_fine(r):
578 continue
579 page1 = (r[2] or 0) + 1
580 bbox = r[3] if isinstance(r[3], dict) else (json.loads(r[3]) if r[3] else None)
581 status, note, cell_text = r[4], r[5], r[6]
582 position = r[7] if isinstance(r[7], dict) else (json.loads(r[7]) if r[7] else None)
583 reason, printed = _reason(status, note, cell_text)
584 refs.append(
585 Reference(
586 ref_id=r[0],
587 ref_type=r[1],
588 page=page1,
589 bbox=bbox,
590 label=_label(r[1], page1, status),
591 status=status,
592 note=note,
593 text=cell_text,
594 row=(position or {}).get("row"),
595 col=(position or {}).get("col"),
596 flagged=bool(status and status in INSPECT_CODES),
597 chart=(position or {}).get("chart"),
598 segment=(position or {}).get("label"),
599 reason=reason,
600 printed=printed,
601 )
602 )
603 return refs
606class Flag(BaseModel):
607 """One item of the document's review queue: a cell whose provenance
608 status is registered for inspection, or a table-level dropped-text flag."""
610 ref_id: str
611 ref_type: Optional[str] # tableCell | tableFlag
612 page: int # 1-based
613 bbox: Optional[dict]
614 status: str
615 note: Optional[str] = None
616 text: str = "" # the flagged cell value, or the dropped fragment
617 reason: Optional[str] = None # the status in the reader's words
620@app.get("/api/document/{doc_id}/flags", response_model=List[Flag])
621def document_flags(doc_id: int) -> List[Flag]:
622 """Every flagged item for one document, in page order."""
623 with db.connect() as conn:
624 rows = conn.execute(
625 """SELECT g.ref_id, g.ref_type, g.page, g.bbox, g.status, g.note, g.cell_text
626 FROM ade_playground.groundings g
627 WHERE g.document_id = %s AND g.status = ANY(%s)
628 ORDER BY g.page, g.ref_id""",
629 (doc_id, list(INSPECT_CODES)),
630 ).fetchall()
631 return [
632 Flag(
633 ref_id=r[0],
634 ref_type=r[1],
635 page=(r[2] or 0) + 1,
636 bbox=r[3] if isinstance(r[3], dict) else (json.loads(r[3]) if r[3] else None),
637 status=r[4],
638 note=r[5],
639 text=r[6] or "",
640 reason=_reason(r[4], r[5], r[6])[0],
641 )
642 for r in rows
643 ]
646# ---------------------------------------------------------------- uploads
648# The jobs this process is running, in memory as the working copy. Every
649# stage transition is also written to the database (quber.playground.jobs),
650# which is what answers a status poll on any task and what outlives a task.
651JOBS: dict[str, dict] = {}
652_JOBS_LOCK = threading.Lock()
653_LOG_TAIL = 200 # lines kept per job
655# Every upload is admitted; this bounds how many run their pipelines at once.
656# Every stage is a subprocess, so two concurrent extractions are two
657# `quber fuse` processes on one CUDA device, each also making model calls.
658_UPLOAD_SLOTS = threading.Semaphore(get_settings().playground.upload_concurrency)
661# The ordered stages each pipeline moves through, ending in done or failed.
662# `waiting` is the admitted-but-queued state, ahead of the stages the pipeline
663# emits — without it a queued document looks stalled on its first step.
664PIPELINE_STAGES = {
665 "fusion": ["uploaded", "waiting", "extracting", "reading figures", "ingesting"],
666}
668#: The fusion pipeline's ingestion tracks, keyed by model line. A track selects
669#: the whole line at once — the figure scan's model and the ingest's chunking
670#: stream — so the two are never mixed within one document. The model is the
671#: CLI spelling: `quber figure` resolves "dpt-3" to its pinned dated version,
672#: so the pin stays single-sourced in the scan module. A future line (dpt-4,
673#: dpt-5) is one more entry here.
674TRACKS = {
675 "dpt-3": {"model": "dpt-3", "chunking": "grouped"},
676 "dpt-2": {"model": "dpt-2", "chunking": "flat"},
677}
678DEFAULT_TRACK = "dpt-3"
680#: Filing types whose figure scan stays on nominated pages only. These are the
681#: long prose-and-table filings where full coverage buys mostly empty scans at
682#: real cost; every other type — and that includes free-form types this set has
683#: never heard of — scans every page, because a deck's content routinely lives
684#: on pages the parse cannot nominate.
685NOMINATED_ONLY_TYPES = {"10-Q", "10-K"}
688def _scan_breadth(job: Dict[str, Any], pdf: Path, workdir: Path) -> tuple[str, str]:
689 """The figure scan's page breadth for one upload, and the reason for it.
691 Typed documents are decided by their filing type alone. An untyped document
692 is never allowed to default to full coverage — the type field is optional
693 and routinely left blank, and a blank must not buy an 80-page scan by
694 accident. A cover attribution agent reads the first page instead and
695 decides presentation (all pages) against prose filing (nominated only);
696 any failure in that read falls back to nominated-only, with the fallback
697 stated in the returned reason so the record shows breadth was reduced and
698 why. The verdict decides breadth only — it is never written to the
699 document's filing type, which stays exactly as the user entered it.
700 """
701 filing_type = (job.get("filing_type") or "").strip().upper()
702 if filing_type:
703 scope = "nominated" if filing_type in NOMINATED_ONLY_TYPES else "all"
704 return scope, f"filing type {filing_type}"
705 import asyncio
707 import fitz
709 from quber.agents.cover_attribution import get_cover_attributor
710 from quber.files.pdf import render_page
712 try:
713 with fitz.open(pdf) as doc:
714 page_count = doc.page_count
715 cover = workdir / f"{pdf.stem}.cover.png"
716 render_page(pdf, 1, 100, cover)
717 verdict = asyncio.run(get_cover_attributor().attribute(cover.read_bytes(), page_count))
718 cover.unlink(missing_ok=True)
719 if verdict.kind == "presentation":
720 return "all", f"attributed presentation: {verdict.reason}"
721 return "nominated", f"attributed filing: {verdict.reason}"
722 except Exception as exc:
723 return (
724 "nominated",
725 f"attribution failed ({type(exc).__name__}: {exc}); nominated-only fallback",
726 )
729class JobStatus(BaseModel):
730 job_id: str
731 filename: str
732 title: Optional[str] = None
733 folder: Optional[str] = None
734 # One of `stages`, or the terminal `done` / `failed`. The view reads the
735 # stage names off this route rather than keeping a parallel set, so these
736 # strings are the contract — the stepper draws `stages` and marks `stage`
737 # against it, which is how the ADE variant appears without the view
738 # learning anything new.
739 stage: str
740 stages: List[str] = PIPELINE_STAGES["fusion"]
741 error: Optional[str] = None
742 # Structured refusal for the states the modal renders specially: kind
743 # "duplicate" (with the existing document to link to), "tuple", or
744 # "busy". Set alongside `error` when an s3-sourced job fails admission —
745 # a body upload hits the same checks synchronously as a 409 whose detail
746 # carries this same shape.
747 error_info: Optional[dict] = None
748 # Set once the ingest completes: the new document's id.
749 doc_id: Optional[int] = None
750 log: List[str]
753def _run_stage(job: dict, stage: str, cmd: List[str]) -> None:
754 job["stage"] = stage
755 job["log"].append(f"--- {stage}: {' '.join(cmd)}")
756 jobs.save_upload(job)
757 proc = subprocess.Popen(cmd, cwd=REPO_ROOT, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
758 assert proc.stdout is not None
759 for line in proc.stdout:
760 job["log"].append(line.rstrip())
761 if len(job["log"]) > _LOG_TAIL:
762 del job["log"][:-_LOG_TAIL]
763 proc.wait()
764 # What the stage wrote is copied to the durable prefix whether it
765 # succeeded or not: a failed stage's partial outputs are the diagnosis.
766 storage.publish_workdir(job["doc_key"], UPLOADS / f"{job['doc_key']}-artifacts")
767 jobs.save_upload(job)
768 if proc.returncode != 0:
769 raise RuntimeError(f"{stage} exited with code {proc.returncode}")
772def _run_hosted_extraction(job: dict, pdf: str, workdir: Path) -> None:
773 """The extracting stage where this host has no GPU: the docling parse
774 goes to the worker while the table engine runs here, and fusion reads
775 the two sets of artifacts from the workdir once both are in it."""
776 doc_key = job["doc_key"]
778 def log(line: str) -> None:
779 job["log"].append(line)
780 if len(job["log"]) > _LOG_TAIL:
781 del job["log"][:-_LOG_TAIL]
783 job["stage"] = "extracting"
784 jobs.save_upload(job)
785 _run_stage(
786 job, "extracting", ["uv", "run", "quber", "table", pdf, "-o", str(workdir), "--llm-backend", "api"]
787 )
788 gpu.parse(doc_key, workdir, log)
789 jobs.save_upload(job)
790 _run_stage(
791 job,
792 "extracting",
793 ["uv", "run", "quber", "fuse", pdf, "--artifacts-dir", str(workdir), "-o", str(workdir)],
794 )
797def _await_slot(job: dict) -> None:
798 """Hold until a pipeline slot frees. Admitted-but-waiting is its own
799 stage: a queued document has something honest to show rather than
800 appearing stalled on its first step."""
801 job["stage"] = "waiting"
802 jobs.save_upload(job)
803 _UPLOAD_SLOTS.acquire()
806def _admission_error(job: dict) -> Optional[dict]:
807 """The reason this upload must not run, as the structured error_info the
808 modal renders — or None. Checks, in order: the same bytes already in the
809 library (or already being ingested), then the filing tuple against
810 existing documents and in-flight jobs. A tuple match with the replace
811 intent records the target on the job instead of refusing, unless the
812 target is under running work."""
813 with db.connect() as conn:
814 dup = conn.execute(
815 "SELECT id, title, folder, filename FROM ade_playground.documents WHERE content_hash = %s",
816 (job["content_hash"],),
817 ).fetchone()
818 if dup:
819 title = dup[1] or dup[3]
820 return {
821 "kind": "duplicate",
822 "message": f"This exact file is already in the library — '{title}' in {dup[2] or 'Ungrouped'}.",
823 "existing": {"id": dup[0], "title": dup[1], "folder": dup[2]},
824 }
825 others = [j for j in jobs.live_uploads() if j["job_id"] != job["job_id"]]
826 if any(j.get("doc_key") == job["doc_key"] for j in others):
827 return {"kind": "duplicate", "message": "This exact file is already being ingested."}
829 key = metadata.tuple_key(job["folder"], job["filing_type"], job["year"], job["period"], job["version"])
830 if key is None:
831 return None
832 for j in others:
833 other_key = metadata.tuple_key(
834 j.get("folder"), j.get("filing_type"), j.get("year"), j.get("period"), j.get("version")
835 )
836 if other_key == key:
837 return {
838 "kind": "tuple",
839 "message": "A document with this filing metadata is already being ingested.",
840 }
841 with db.connect() as conn:
842 rows = conn.execute("SELECT " + _DOC_COLUMNS + " FROM ade_playground.documents d").fetchall()
843 match = next(
844 (
845 d
846 for d in (_doc_dict(r) for r in rows)
847 if metadata.tuple_key(d["folder"], d["filing_type"], d["year"], d["period"], d["version"]) == key
848 ),
849 None,
850 )
851 if match is None:
852 return None
853 if not job["replace"]:
854 return {
855 "kind": "tuple",
856 "message": f"A {match['label']} already exists in {match['folder'] or 'Ungrouped'}.",
857 "existing": {"id": match["id"], "title": match["title"], "folder": match["folder"]},
858 }
859 busy = _ingest_busy(match["id"], match["doc_key"])
860 if busy:
861 return {"kind": "busy", "message": f"Cannot replace: {busy}."}
862 job["replace_id"] = match["id"]
863 job["replace_key"] = match["doc_key"]
864 return None
867def _finalize_job(job: dict) -> None:
868 """Stamp the upload's metadata onto the freshly ingested row and, for a
869 replace, swap the target out in the same transaction. Ingest recreates
870 the document row, so this runs after it."""
871 replaced_key: Optional[str] = None
872 with db.connect() as conn:
873 with conn.transaction():
874 conn.execute(
875 """UPDATE ade_playground.documents
876 SET title = %s, folder = %s, filing_type = %s, year = %s, period = %s,
877 version = %s, track = %s
878 WHERE doc_key = %s""",
879 (
880 job.get("title"),
881 job.get("folder"),
882 job.get("filing_type"),
883 job.get("year"),
884 job.get("period"),
885 job.get("version"),
886 job.get("track"),
887 job["doc_key"],
888 ),
889 )
890 if job.get("replace_id") is not None:
891 # A batch run started against the target while the new ingest
892 # ran: leave both documents standing (the overlap flag names
893 # them) rather than deleting under the run.
894 if BATCH.running_for(job["replace_id"]):
895 job["log"].append(
896 "replace target is under a running batch; both documents kept "
897 "(remove the old one from the library when the run finishes)"
898 )
899 else:
900 conn.execute("DELETE FROM ade_playground.documents WHERE id = %s", (job["replace_id"],))
901 replaced_key = job.get("replace_key")
902 row = conn.execute(
903 "SELECT id FROM ade_playground.documents WHERE doc_key = %s", (job["doc_key"],)
904 ).fetchone()
905 if replaced_key:
906 _remove_doc_files(replaced_key)
907 logger.info("replaced document {id} ({key})", id=job["replace_id"], key=replaced_key)
908 job["doc_id"] = row[0] if row else None
909 jobs.save_upload(job)
912def _admit_or_fail(job: dict, pdf_bytes: Optional[bytes], s3_uri: Optional[str]) -> Optional[str]:
913 """Resolve the upload to local bytes, hash them, and run admission.
914 Returns the staged PDF path to run the pipeline on, or None after
915 marking the job failed. For body uploads admission already ran in the
916 route; this stages the s3 source and repeats it (the authoritative
917 pass — the route could not hash bytes it never had)."""
918 if pdf_bytes is None:
919 assert s3_uri is not None
920 staged = UPLOADS / f"tmp-{job['job_id']}.pdf"
921 try:
922 fetch_file(s3_uri, staged)
923 pdf_bytes = staged.read_bytes()
924 except Exception as exc:
925 job["stage"] = "failed"
926 job["error"] = f"could not fetch {s3_uri}: {exc}"
927 return None
928 job["content_hash"] = hashlib.sha256(pdf_bytes).hexdigest()
929 job["doc_key"] = job["content_hash"][:16]
930 info = _admission_error(job)
931 if info is not None:
932 staged.unlink(missing_ok=True)
933 job["stage"] = "failed"
934 job["error"] = info["message"]
935 job["error_info"] = info
936 return None
937 staged.rename(UPLOADS / f"{job['doc_key']}.pdf")
938 return str(UPLOADS / f"{job['doc_key']}.pdf")
941def _run_job(job_id: str, pdf_bytes: Optional[bytes], s3_uri: Optional[str]) -> None:
942 """The quber workflow: fuse, then read the figures, then ingest.
944 The figure stage is part of the workflow rather than an option on it. Without
945 it a document arrives holding only what the text layer gave up: a chart's
946 plotted values are absent, a table printed as an image stands as the parse's
947 own reading of it, and the notes that qualify a figure are attached to
948 nothing. The standalone ADE path is a different pipeline and is unaffected.
950 The stage runs on every upload, and it spends: each nominated page is a scan.
951 A document the parse nominates no page in costs nothing and writes no figure
952 artifacts, which the ingest reads as "the run produced none of that kind"
953 rather than as an error.
954 """
955 job = JOBS[job_id]
956 try:
957 _await_slot(job)
958 pdf = _admit_or_fail(job, pdf_bytes, s3_uri)
959 if pdf is None:
960 return
961 doc_key = job["doc_key"]
962 workdir = UPLOADS / f"{doc_key}-artifacts"
963 workdir.mkdir(parents=True, exist_ok=True)
964 base = doc_key
965 # The PDF and its source manifest reach the durable prefix before any
966 # stage runs, and every stage's outputs follow it as the stage ends.
967 manifest = storage.write_source_manifest(
968 doc_key, filename=job["filename"], content_hash=job["content_hash"], source_uri=s3_uri
969 )
970 storage.publish(doc_key, [Path(pdf), manifest])
971 track = TRACKS[job.get("track") or DEFAULT_TRACK]
972 # Breadth is decided before anything runs and recorded on the job, so
973 # the run's record always says how many pages were in play and why —
974 # server-side only, never part of a client payload.
975 pages, breadth_reason = _scan_breadth(job, Path(pdf), workdir)
976 job["scan_pages"] = pages
977 job["scan_pages_reason"] = breadth_reason
978 job["source_uri"] = s3_uri
979 jobs.save_upload(job)
980 logger.info("Upload {}: figure scan breadth {} ({})", job_id, pages, breadth_reason)
981 if gpu.configured():
982 _run_hosted_extraction(job, pdf, workdir)
983 else:
984 _run_stage(job, "extracting", ["uv", "run", "quber", "fuse", pdf, "-o", str(workdir)])
985 _run_stage(
986 job,
987 "reading figures",
988 [
989 "uv",
990 "run",
991 "quber",
992 "figure",
993 pdf,
994 "--parse",
995 str(workdir / f"{base}.unified.json"),
996 "-o",
997 str(workdir),
998 "--model",
999 track["model"],
1000 "--pages",
1001 pages,
1002 ],
1003 )
1004 _run_stage(
1005 job,
1006 "ingesting",
1007 [
1008 "uv",
1009 "run",
1010 "python",
1011 "-m",
1012 "quber.playground.ingest_fusion",
1013 "--doc-key",
1014 doc_key,
1015 "--filename",
1016 job["filename"],
1017 "--unified-json",
1018 str(workdir / f"{base}.unified.json"),
1019 "--tables-json",
1020 str(workdir / f"{base}.tables.json"),
1021 "--fusion-json",
1022 str(workdir / f"{base}.fusion.json"),
1023 "--scanned-tables-json",
1024 str(workdir / f"{base}.scanned-tables.json"),
1025 "--figures-json",
1026 str(workdir / f"{base}.figures.json"),
1027 "--figure-values-json",
1028 str(workdir / f"{base}.figure-values.json"),
1029 "--pdf",
1030 pdf,
1031 "--chunking",
1032 track["chunking"],
1033 ],
1034 )
1035 _finalize_job(job)
1036 job["stage"] = "done"
1037 except Exception as exc:
1038 job["stage"] = "failed"
1039 job["error"] = str(exc)
1040 finally:
1041 jobs.save_upload(job)
1042 _UPLOAD_SLOTS.release()
1045@app.post("/api/upload", response_model=JobStatus)
1046async def upload(
1047 request: Request,
1048 filename: str,
1049 s3_uri: Optional[str] = None,
1050 title: Optional[str] = None,
1051 folder: Optional[str] = None,
1052 filing_type: Optional[str] = None,
1053 year: Optional[int] = None,
1054 period: Optional[str] = None,
1055 version: Optional[int] = None,
1056 replace: bool = False,
1057 track: str = DEFAULT_TRACK,
1058) -> JobStatus:
1059 """Run the fusion pipeline on a PDF sent as the raw request body, or on an
1060 s3:// object named by `s3_uri` (body then stays empty).
1062 `track` selects the model line the document is processed on — see TRACKS.
1064 `filename` is the original name, kept for presentation and never for
1065 identity. The metadata fields are how the library presents and keys the
1066 document; this route and the PATCH are the only places they are ever
1067 set. `replace` carries the modal's confirmed intent to swap out the
1068 document whose filing tuple this upload matches. Every admitted upload
1069 waits in the `waiting` stage until a pipeline slot frees.
1070 """
1071 if track not in TRACKS:
1072 raise HTTPException(400, f"track must be one of {sorted(TRACKS)}")
1073 body: Optional[bytes] = None
1074 if s3_uri is not None:
1075 if not s3_uri.startswith("s3://") or not s3_uri.lower().endswith(".pdf"):
1076 raise HTTPException(400, "s3_uri must be an s3:// URI ending in .pdf")
1077 else:
1078 body = await request.body()
1079 if not body.startswith(b"%PDF"):
1080 raise HTTPException(400, "request body is not a PDF")
1081 job_id = uuid.uuid4().hex[:12]
1082 job: dict = {
1083 "job_id": job_id,
1084 "filename": filename,
1085 "title": title,
1086 "folder": folder,
1087 "filing_type": filing_type,
1088 "year": year,
1089 "period": period,
1090 "version": version,
1091 "replace": replace,
1092 "track": track,
1093 "replace_id": None,
1094 "replace_key": None,
1095 "doc_key": None,
1096 "content_hash": None,
1097 "doc_id": None,
1098 "stage": "uploaded",
1099 "stages": PIPELINE_STAGES["fusion"],
1100 "error": None,
1101 "error_info": None,
1102 "log": [],
1103 }
1104 if body is not None:
1105 # The bytes are in hand: hash and run admission now, so a duplicate
1106 # or an unconfirmed tuple match is a synchronous refusal the modal
1107 # shows in place — not a job that fails a poll later.
1108 job["content_hash"] = hashlib.sha256(body).hexdigest()
1109 job["doc_key"] = job["content_hash"][:16]
1110 info = _admission_error(job)
1111 if info is not None:
1112 raise HTTPException(409, info)
1113 UPLOADS.mkdir(parents=True, exist_ok=True)
1114 (UPLOADS / f"{job['doc_key']}.pdf").write_bytes(body)
1115 else:
1116 UPLOADS.mkdir(parents=True, exist_ok=True)
1117 with _JOBS_LOCK:
1118 JOBS[job_id] = job
1119 jobs.save_upload(job)
1120 threading.Thread(target=_run_job, args=(job_id, body, s3_uri), daemon=True).start()
1121 return JobStatus(**{**job, "log": []})
1124@app.get("/api/jobs/{job_id}", response_model=JobStatus)
1125def job_status(job_id: str, tail: int = 12) -> JobStatus:
1126 """The job as this process holds it when it is running it, else as the
1127 database records it, which is how a poll lands on any task."""
1128 job = JOBS.get(job_id) or jobs.load_upload(job_id)
1129 if not job:
1130 raise HTTPException(404, f"no job {job_id}")
1131 return JobStatus(**{**job, "log": job["log"][-tail:]})
1134def _readable(payload) -> str:
1135 """What to put in `answer` for a consumer that only reads text.
1137 A figure renders as the figure, so the text field of a value-seeking
1138 answer is the value and not a sentence about it.
1139 """
1140 if payload.kind == "prose":
1141 return payload.text
1142 if payload.kind == "unanswerable":
1143 return payload.reason
1144 if payload.kind == "scalar":
1145 return payload.value
1146 return json.dumps(payload.model_dump(mode="json"), indent=2)
1149def cited_ids_source_first(payload, cited_ids: List[str]) -> List[str]:
1150 """The cited ids with the answer's own source leading.
1152 A value payload names the single cell or chunk it was read from in
1153 `source_id`. The UI activates the first reference that carries a box, so
1154 whatever id leads this list becomes the default highlight — and the
1155 value's own cell is the highlight the answer earned, not the wider line
1156 record the model happened to cite first. Payloads without a source
1157 (prose, unanswerable, series) keep the model's citation order.
1158 """
1159 source_id = getattr(payload, "source_id", None)
1160 if not source_id:
1161 return list(cited_ids)
1162 return [source_id] + [cid for cid in cited_ids if cid != source_id]
1165def _chat_response(doc_id: int, result, chunks) -> ChatResponse:
1166 payload = result.payload
1167 return ChatResponse(
1168 answer=_readable(payload),
1169 shape=payload.kind,
1170 value=payload.model_dump(mode="json"),
1171 references=_resolve_refs(doc_id, cited_ids_source_first(payload, result.cited_ids)),
1172 retrieved=[
1173 RetrievedInfo(
1174 chunk_id=c.chunk_id, chunk_type=c.chunk_type, page=c.page + 1, score=round(c.score, 4)
1175 )
1176 for c in chunks
1177 ],
1178 )
1181@app.post("/api/chat", response_model=ChatResponse)
1182async def chat(req: ChatRequest) -> ChatResponse:
1183 doc = _doc_row(req.doc_id)
1184 document_key.set(doc["doc_key"])
1185 chunks = await retrieve(doc["doc_key"], req.question, k=req.k)
1186 if not chunks:
1187 raise HTTPException(404, f"No chunks for document {req.doc_id} — is it ingested?")
1188 result = await agent_answer(req.question, chunks, req.want, DocumentIdentity.from_row(doc))
1189 return _chat_response(req.doc_id, result, chunks)
1192@app.post("/api/chat/stream")
1193async def chat_stream(req: ChatRequest) -> StreamingResponse:
1194 """Same contract as /api/chat, streamed: server-sent events carrying the
1195 answer text as it generates ('delta'), then the full response ('done')."""
1196 doc = _doc_row(req.doc_id)
1197 document_key.set(doc["doc_key"])
1198 chunks = await retrieve(doc["doc_key"], req.question, k=req.k)
1199 if not chunks:
1200 raise HTTPException(404, f"No chunks for document {req.doc_id} — is it ingested?")
1202 identity = DocumentIdentity.from_row(doc)
1204 async def events():
1205 try:
1206 async for kind, payload in agent_answer_stream(req.question, chunks, req.want, identity):
1207 if kind == "delta":
1208 yield f"data: {json.dumps({'type': 'delta', 'answer': payload})}\n\n"
1209 else:
1210 body = _chat_response(req.doc_id, payload, chunks).model_dump()
1211 yield f"data: {json.dumps({'type': 'done', 'payload': body})}\n\n"
1212 except Exception as exc:
1213 yield f"data: {json.dumps({'type': 'error', 'detail': str(exc)})}\n\n"
1215 return StreamingResponse(events(), media_type="text/event-stream")
1218# ---------------------------------------------------------------- batch
1220BATCH_STORE = jobs.BatchStore()
1221BATCH = BatchRunner(store=BATCH_STORE)
1224class BatchRequest(BaseModel):
1225 doc_id: int
1226 # The batch screen's prompt window, verbatim: one question per line. The
1227 # server parses it so there is exactly one reading of that format.
1228 questions: str
1229 k: int = 10
1230 # One expectation for the whole run; see ChatRequest.want.
1231 want: Literal["auto", "value", "text"] = "auto"
1234@app.post("/api/batch", response_model=RunState)
1235async def batch_start(req: BatchRequest) -> RunState:
1236 """Start a run. The work continues if the page is closed; the page keeps
1237 the returned job id in sessionStorage and re-attaches through the state
1238 and events routes."""
1239 questions = parse_questions(req.questions)
1240 if not questions:
1241 raise HTTPException(400, "no questions — the batch window takes one per line")
1242 doc = _doc_row(req.doc_id)
1243 doc_key = doc["doc_key"]
1244 # Set before the probe so its selection run carries the document too; the
1245 # run task copies this context when it is created.
1246 document_key.set(doc_key)
1247 probe = await retrieve(doc_key, questions[0], k=1)
1248 if not probe:
1249 raise HTTPException(404, f"No chunks for document {req.doc_id} — is it ingested?")
1251 identity = DocumentIdentity.from_row(doc)
1253 async def retrieve_one(question: str) -> List[RetrievedChunk]:
1254 return await retrieve(doc_key, question, k=req.k)
1256 async def answer_one(question: str, chunks: List[RetrievedChunk]) -> Dict[str, Any]:
1257 result = await agent_answer(question, chunks, req.want, identity)
1258 return _chat_response(req.doc_id, result, chunks).model_dump(mode="json")
1260 return BATCH.start(req.doc_id, questions, req.want, retrieve_one, answer_one)
1263def _run_state(job_id: str) -> RunState:
1264 state = BATCH.get(job_id)
1265 if state is None:
1266 # A restart and a cap eviction arrive here identically: an id the page
1267 # kept and a run this process cannot answer for. The page reports the
1268 # run as gone; it is never silence.
1269 raise HTTPException(404, f"no run {job_id} — the server no longer has it")
1270 return state
1273@app.get("/api/batch/{job_id}", response_model=RunState)
1274def batch_state(job_id: str) -> RunState:
1275 return _run_state(job_id)
1278@app.get("/api/batch/{job_id}/events")
1279async def batch_events(job_id: str) -> StreamingResponse:
1280 """One server-sent event per question completing — index, total, outcome —
1281 then a final `done` event. Watching a run and re-attaching to one are
1282 different needs; this is the watching route."""
1283 _run_state(job_id)
1285 async def stream():
1286 async for event in BATCH.events(job_id):
1287 yield f"data: {json.dumps(event)}\n\n"
1289 return StreamingResponse(stream(), media_type="text/event-stream")
1292def _export_stem(state: RunState) -> str:
1293 """The export filename's document part, down the label precedence. A run
1294 whose document was removed mid-session still exports; it just gets the
1295 generic stem."""
1296 try:
1297 doc = _doc_row(state.doc_id)
1298 except HTTPException:
1299 return "document"
1300 return metadata.export_stem(doc["label"], doc["title"], doc["filename"])
1303@app.get("/api/batch/{job_id}/export.xlsx")
1304def batch_export_xlsx(job_id: str) -> Response:
1305 state = _run_state(job_id)
1306 return Response(
1307 export.write_xlsx(state),
1308 media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
1309 headers={"Content-Disposition": f'attachment; filename="batch-{_export_stem(state)}-{job_id}.xlsx"'},
1310 )
1313@app.get("/api/batch/{job_id}/export.csv")
1314def batch_export_csv(job_id: str) -> Response:
1315 state = _run_state(job_id)
1316 return Response(
1317 export.write_csv(state),
1318 media_type="text/csv; charset=utf-8",
1319 headers={"Content-Disposition": f'attachment; filename="batch-{_export_stem(state)}-{job_id}.csv"'},
1320 )