From e8945cab46d0d5998c522b3faaff2ba19fbff7c9 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 30 Jun 2026 06:55:23 -0700 Subject: [PATCH] Whole-document context for RAG chat attachments (#6693) * Add whole-document context mode to RAG chat attachments Thread-attached files are injected in full when they fit a token budget, instead of only top-K retrieved chunks, so the model reads the entire file for summarize/reason-over-document requests. Oversized files fall back to top-K retrieval so the context window is never blown. KB and project corpora are unchanged (still retrieval). - core/rag/store.py: all_chunks_for_scope returns every completed-document chunk for a scope, ordered document-then-index, joined with filename. - core/rag/tool.py: whole_document_context renders the chunks as the same blocks + citation source-map retrieval produces, returns None when empty or over budget. - core/inference/tools.py: build_rag_autoinject tries whole-document first for thread scopes, falls through to search_for_autoinject otherwise. - core/rag/config.py: THREAD_WHOLE_DOC + WHOLE_DOC_MAX_TOKENS (env-tunable). - tests/test_rag_whole_document.py: store ordering, whole-doc render + budget cutoff, auto-inject whole-doc vs top-K fallback, KB never whole-doc. * Add scanned-PDF OCR fallback to RAG ingestion A PDF page with no extractable text layer (a scanned or image-only page) previously ingested as empty, so image PDFs were invisible to retrieval and whole-document context. Such pages are now rendered and transcribed by the loaded vision model during ingestion, so they become searchable and readable like any other page. This restores OCR for the RAG document flow without a separate extraction pipeline. - core/rag/parsers.py: render_pdf_pages renders whole pages (1-based) to PNG. - core/rag/captioner.py: factor the shared vision call into _vision_complete; add _ocr_one + ocr_pages (transcribe rendered pages, OCR_MAX_PAGES bound). - core/rag/ingestion.py: _ocr_scanned_pages runs right after parse, replacing text on near-empty PDF pages. No-op when OCR is off, no page is scanned, or no vision model is loaded (degrades like figure captioning). - core/rag/config.py: OCR_SCANNED, OCR_MIN_CHARS, OCR_MAX_PAGES, OCR_DPI, OCR_TIMEOUT_S, OCR_MAX_TOKENS (env-tunable). - tests/test_rag_ocr_fallback.py: page render, ocr_pages gating + cap, scanned PDF end-to-end OCR into chunks + whole-doc, born-digital skips OCR, disabled leaves the page empty. * Broaden OCR prompt to figures/tables and guard against repetition runaway The OCR prompt now asks the vision model to also transcribe text inside figures, diagrams, charts and tables, so labels and table cells on scanned pages are indexed rather than skipped. Verified on real documents that this does not regress plain-text transcription. Some vision models loop on sparse images (e.g. a title-only cover) and emit the same line hundreds of times. _collapse_runaway caps any run of identical consecutive lines so a pathological page cannot flood the index; legitimate short repeats (a label appearing a few times) survive. Applied in ocr_pages. * Restrict whole-document injection to thread attachments only whole_document_context resolved the combined project+thread scope, so a project chat (the frontend sends both thread_id and project_id) injected the entire project corpus in full, contradicting the design that project and KB corpora stay retrieval-only. A large project corpus could also push the total over budget and drop a small thread attachment back to top-K. Resolve the thread scope alone in whole_document_context, and in build_rag_autoinject only enter whole-doc mode when a thread attachment is present and no KB is selected (a KB pick is exclusive: search that corpus). Project sources and KBs keep top-K retrieval. Adds regression tests for the mixed project+thread payload, the budget isolation, and KB precedence. * Address review: keep project retrieval, harden budget + OCR guards Follow-up to the 8-reviewer pass on the whole-document + OCR work. - Preserve project grounding in project chats. The thread-scope-only fix made whole-doc exclusive of retrieval, so a thread attachment silently dropped the project corpus for that turn. build_rag_autoinject now whole-docs the thread attachment AND retrieves the project sources top-K, merged under one citation numbering via tool.render_sources. KB selection stays exclusive. - Budget: a NULL/zero token_count no longer bypasses the cap (length-based fallback in _row_token_count), so a malformed huge doc can't inject in full. - OCR runaway guard: _collapse_runaway now also caps each distinct line at a generous total across the page (not just consecutive), bounding the interleaved/alternating loops weak models emit; blank-line floods collapse too. - OCR: warn when a scanned PDF exceeds OCR_MAX_PAGES (pages past the cap stay untranscribed) instead of silently dropping them. - Document the known limits: OCR'd pages have no PDF highlight regions; vision models need a micro-batch >= image tokens (Gemma-family) or the server aborts. - Tests for project-retrieval composition, NULL-token budget, and interleaved runaway; drop the now-superseded exclude-project test. * Add OCR toggle to RAG retrieval settings Make scanned-PDF OCR user-controllable per upload instead of only via the RAG_OCR_SCANNED config default. The retrieval settings panel gains an OCR scanned pages switch (persisted in localStorage, on by default); the chosen value is read fresh at upload time and sent with each document upload. Backend: the three upload routes accept an optional ocr form field and pass it through start_ingestion to _ocr_scanned_pages, which now treats None as use the config default and an explicit bool as an override. The on/off policy lives only in _ocr_scanned_pages now, so ocr_pages no longer re-checks the config (that double gate would have blocked a per-upload ocr=True while the default was off). Tests cover both override directions (force on while config off, force off while config on). * Add "Describe figures & charts" toggle with chart-aware captions Surface RAG figure captioning as a user control and make it actually useful for graphs and plots. The figure detection already clustered vector drawings and raster images into regions and rendered them, but captioning was off by default, had no UI, and used a thin generic prompt. Accuracy: the caption prompt now asks for chart type, axis titles and units, legend or series, salient trends and readable values, and table columns, while forbidding invented numbers. The token budget is configurable (CAPTION_MAX_TOKENS) and captions pass through the same runaway guard as OCR so a looping vision model cannot flood the index. Control: a per-upload caption override threads from the three upload routes through start_ingestion and _run, with the on/off policy single-sourced in _run (caption self-gating removed from caption_images, mirroring the OCR change) so a force-on override works when the config default is off. The frontend adds a "Describe figures & charts" switch in the retrieval settings, persisted in localStorage and sent with each upload. Default on; it is a no-op without a vision model and bounded to CAPTION_MAX_IMAGES figures per document. Tests cover the new caption_images contract, the runaway guard on captions, the chart-aware prompt and token budget (and that OCR keeps its own prompt and budget), and both override directions end to end through ingestion. * Generalize figure understanding: transcribe-first prompt + high-DPI tiling Make figure/chart description work across any visual and any model strength, not just a strong VLM on simple figures. Two changes, validated by a recall benchmark on authoritative documents (ResNet/Attention papers, USDA, UN UDHR). 1. Transcribe-first caption prompt. The caption now asks the model to transcribe every visible label verbatim (titles, axis labels and units, legends, every box/node/arrow label, table cells, equations) and then add a one-line summary, instead of only describing the figure. Transcription is the most model-robust visual task, so weak models that cannot reason about a chart still recover its labels. 2. High-DPI tiling of figure pages. Figure-bearing pages are rendered as an overlapping grid of high-DPI tiles (plus a full-page pass for context); each tile is transcribed, then merged and de-duplicated. This keeps small diagram labels legible and covers every sub-figure without relying on exact region detection, which previously missed sub-figures and small labels. Supporting changes: figure render DPI 130 -> 200 with a clip margin so edge labels are not lost; vision calls are deterministic (temperature 0) so transcription does not randomly drop labels; the repetition guard now applies to captions too. New config knobs: FIGURE_DPI, FIGURE_MARGIN_FRAC, FIGURE_TILE_ROWS/COLS, FIGURE_TILE_ OVERLAP, FIGURE_FULLPAGE, CAPTION_MAX_PAGES, larger CAPTION_MAX_TOKENS, and CAPTION_MAX_IMAGES as a per-document tile budget. Measured figure context recall (per-label, dense academic figures): Qwen2.5-VL: 0.50 -> 0.83 (overall 0.81 -> 0.94) Gemma-4-E2B (weak): ~0 with loops -> 0.83 (overall 0.91) Born-digital text and scanned-page recall are unchanged (no regression). parsers gains _figure_boxes (shared detection), pages_with_figures, and render_pdf_figure_tiles; captioner gains merge_page_captions and a temperature parameter; ingestion routes figure captioning through the tiled path. * Fix RAG review issues: whole-doc budget pre-check, figure gating, empty re-ingest, vision auth Whole-document context now runs a cheap token-sum pre-check (store.scope_token_estimate) before hydrating every chunk's text, so an attachment that cannot fit the budget is rejected without loading the whole corpus into memory. The estimate mirrors all_chunks_for_scope's filter and the per-row token-count fallback exactly. Ingestion skips all figure work (PDF rasterization and detection, not just the caption call) unless a vision model is loaded, so a text-only deployment pays nothing. When OCR is enabled, scanned/image-only pages are excluded from figure tiling since OCR already transcribes them whole, avoiding double vision work and overlapping index entries; a scanned figure page is still tiled when OCR is off. start_ingestion no longer dedupes forever to a prior ingest that produced zero chunks (e.g. a scanned PDF uploaded before a vision model was loaded): the empty record is dropped and the content is re-ingested. Vision OCR and caption requests now send the backend Authorization header, so they match the chat endpoint and do not 401 under direct-stream (--api-key) mode. Adds tests for the budget estimate, scanned-page exclusion, the vision-model gate, the empty re-ingest path, and the auth-header passthrough. * Trim RAG vision-ingestion comments and docstrings Tighten the verbose multi-line docstrings and comments added across the RAG vision ingestion work (captioner, config, parsers, ingestion, store, tool, build_rag_autoinject, the RAG tests, and the chat-store/upload-hook frontend toggles) to one or two lines while keeping their intent. No code changed: verified comment/docstring-only against the prior commit, and the RAG test suite still passes. * Fix figure-tiling exclusion and client dedupe for re-ingestable docs Figure tiling now excludes only the pages OCR actually transcribed, not every text-less page. _ocr_scanned_pages returns the set of pages it OCR'd, and _run passes that to pages_with_figures as exclude_pages (replacing the ocr_on-keyed min_text_chars heuristic). A scanned page that OCR skipped (past OCR_MAX_PAGES, or whose OCR returned empty) is no longer dropped from captioning, so a chart on such a page still gets a caption. The document panel's upload dedupe no longer skips re-selecting a file whose only matching doc completed with zero chunks. Such a doc is re-ingestable (e.g. a scan attached before a vision model loaded), and the backend re-ingests on the same content hash, so the client must let it reach the backend; healthy or still-indexing docs are still skipped. The SSE complete frame's chunk count is recorded on the doc so the check is exact. Adds a regression test for the un-OCR'd scanned figure page and updates the pages_with_figures test to the exclude_pages interface. * Address review findings: whole-doc budget guard, job numChunks, dead code, upload cap whole_document_context now treats a non-positive max_tokens as "never inject" instead of injecting the whole corpus unbounded, so RAG_WHOLE_DOC_MAX_TOKENS=0 tightens rather than disables the budget (the real off switch stays RAG_THREAD_WHOLE_DOC=0). The job-status endpoint and get_job_status now expose num_chunks (joined from the document), and the upload hook threads it through the SSE-fallback completion paths (reconcile + poll). Previously a document that finished via the connection-cap fallback had no chunk count client-side, so the re-ingest dedupe wrongly treated it as empty and re-uploaded it. IndexJob/JobEvent gain the field and the untyped cast is dropped. Removes the dead render_pdf_figures function (superseded by the tiling path), its test, and the unused FIGURE_MARGIN_FRAC config knob. Adds an upload size cap (RAG_MAX_UPLOAD_BYTES, default 200 MB; 413 on exceed with the partial file cleaned up) so a pathological file can't drive unbounded parse + vision work. render_pdf_figure_tiles clamps rows/cols to >= 1 (no ZeroDivisionError on a misconfigured grid). Captioning progress is reported after OCR so the bar is monotonic. sqlite connections set busy_timeout=5000 so a long figure/scan ingest holding its connection doesn't make a concurrent ingest/read fail with "database is locked". Adds tests for the non-positive budget, the zero-grid clamp, job-status num_chunks, and the oversize-upload rejection. * Extract PDF text as layout-aware Markdown via pymupdf4llm parsers._pdf now extracts each PDF page as Markdown with pymupdf4llm.to_markdown (page_chunks=True) instead of flat page.get_text("text"), so tables, headings and lists keep their structure in the indexed chunks and retrieve far better (a table's cells stay associated with their row instead of flattening into a token stream). Gated by RAG_PDF_MARKDOWN (default on); falls back to plain PyMuPDF text when the toggle is off, pymupdf4llm is missing, extraction fails, or a page yields no Markdown. The scanned-page OCR and figure-tiling passes operate on rendered pixels and are unaffected; docx/html/txt keep their existing extractors. The preview-highlight locator already strips Markdown punctuation when building anchors; it now also splits anchor tokens on pipes so a Markdown table row still anchors to the raw PDF word stream. Declares pymupdf4llm as a studio/RAG dependency (was only transitively present via the data-designer plugin). Adds parser tests (Markdown table reaches the page text, the plain-text fallback, the missing-lib fallback) and a locator test for table-pipe anchoring. * Pin pymupdf4llm to 0.3.4 so the package scan does not pull onnxruntime The lockstep pymupdf4llm 1.27.x line makes pymupdf-layout a hard dependency, which in turn pulls onnxruntime (plus numpy/networkx/protobuf). The security-audit pip scan-packages job resolves requirements --with-deps, so adding pymupdf4llm to no-torch-runtime.txt and studio.txt surfaced onnxruntime's un-baselined CRITICAL finding and flipped the hf-stack shard from pass to fail. pymupdf4llm 0.3.x keeps pymupdf-layout behind an optional [layout] extra, so a plain install resolves to pymupdf + tabulate only and never touches onnxruntime. 0.3.4 requires pymupdf>=1.27.1, satisfied by our pinned pymupdf==1.27.2.3, and to_markdown (page_chunks=True) produces equivalent layout-aware Markdown on real PDFs (verified on the Attention, ResNet and USDA documents). Production already installs these files --no-deps, so onnxruntime was never shipped at runtime; this only fixes the scanner. The parser test now asserts Markdown markup (heading or table pipes) rather than table pipes specifically, since 0.3.4 emits a heading but not a pipe table on the tiny borderless synthetic fixture; both markers are absent from the plain-text fallback. * Fix RAG whole-doc review findings * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address RAG whole-doc review follow-ups * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address RAG review follow-up edge cases * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Reserve image budget for whole-document RAG * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: danielhanchen Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: wasimysaid Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/inference/tools.py | 143 ++++- studio/backend/core/rag/captioner.py | 173 +++++- studio/backend/core/rag/config.py | 49 +- studio/backend/core/rag/ingestion.py | 144 ++++- studio/backend/core/rag/locators.py | 6 +- studio/backend/core/rag/parsers.py | 212 +++++-- studio/backend/core/rag/store.py | 37 ++ studio/backend/core/rag/tool.py | 81 ++- .../backend/requirements/no-torch-runtime.txt | 5 + studio/backend/requirements/studio.txt | 3 + studio/backend/routes/rag.py | 32 +- studio/backend/storage/rag_db.py | 4 + studio/backend/tests/test_rag_captioning.py | 303 +++++++++- studio/backend/tests/test_rag_ingestion.py | 63 +++ studio/backend/tests/test_rag_ocr_fallback.py | 259 +++++++++ studio/backend/tests/test_rag_parsing.py | 88 +++ studio/backend/tests/test_rag_preview.py | 19 + .../backend/tests/test_rag_whole_document.py | 520 ++++++++++++++++++ .../src/features/chat/api/chat-adapter.ts | 6 + .../chat/stores/chat-runtime-store.ts | 26 + .../frontend/src/features/chat/types/api.ts | 3 + .../frontend/src/features/rag/api/rag-api.ts | 32 +- .../components/retrieval-settings-section.tsx | 52 ++ .../rag/components/use-rag-documents.ts | 66 ++- studio/frontend/src/features/rag/types/rag.ts | 2 + 25 files changed, 2171 insertions(+), 157 deletions(-) create mode 100644 studio/backend/tests/test_rag_ocr_fallback.py create mode 100644 studio/backend/tests/test_rag_parsing.py create mode 100644 studio/backend/tests/test_rag_whole_document.py diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index a5c193ff39..82c50933fc 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -1121,6 +1121,61 @@ def _autoinject_top_k() -> int: return _AUTOINJECT_DEFAULT_TOP_K +def _thread_whole_doc_enabled(scope: dict) -> bool: + """Whether a thread-attached file should be injected in full rather than + retrieved top-K. ``rag_scope.whole_doc=False`` disables it for this request.""" + override = scope.get("whole_doc") + if override is False: + return False + try: + from core.rag import config as _rag_config + except Exception: # noqa: BLE001 + return True + return _rag_config.THREAD_WHOLE_DOC + + +_IMAGE_PART_TOKEN_ESTIMATE = 1024 + + +def _message_token_estimate(conversation: list[dict]) -> int: + """Cheap prompt-size estimate for budget guards; exact tokenization happens later.""" + total = 0 + for msg in conversation: + content = msg.get("content") + if isinstance(content, str): + total += max(1, len(content) // 4) + elif isinstance(content, list): + for part in content: + if isinstance(part, dict): + if part.get("type") in ("image_url", "input_image"): + total += _IMAGE_PART_TOKEN_ESTIMATE + else: + total += max(1, len(str(part.get("text") or "")) // 4) + total += 4 # chat-template role / separator overhead estimate + return total + + +def _whole_doc_budget(scope: dict | None = None, conversation: list[dict] | None = None) -> int: + try: + from core.rag import config as _rag_config + except Exception: # noqa: BLE001 + budget = 6000 + else: + budget = _rag_config.WHOLE_DOC_MAX_TOKENS + if not scope: + return budget + context = _opt_int(scope.get("context_length") or scope.get("max_context_tokens")) + if context is None or context <= 0: + return budget + headroom = _opt_int(scope.get("response_headroom")) + if headroom is None: + headroom = max(1024, context // 4) + used = _message_token_estimate(conversation or []) + # Leave room for tool XML wrappers, citation metadata, and chat-template overhead. + available = context - headroom - used - 512 + return min(budget, max(0, available)) + + def _last_user_text(conversation: list[dict]) -> str: """Plain text of the most recent user turn (text parts only).""" for msg in reversed(conversation): @@ -1154,7 +1209,11 @@ def build_rag_autoinject(conversation: list[dict], rag_scope: dict | None) -> di enabled = rag_scope.get("autoinject") if enabled is None: enabled = _autoinject_enabled() - if not enabled: + thread_id = rag_scope.get("thread_id") + whole_doc_requested = ( + bool(thread_id) and not rag_scope.get("kb_id") and _thread_whole_doc_enabled(rag_scope) + ) + if not enabled and not whole_doc_requested: return None query = _last_user_text(conversation) if not query: @@ -1163,35 +1222,81 @@ def build_rag_autoinject(conversation: list[dict], rag_scope: dict | None) -> di from storage import rag_db if not rag_db.RAG_AVAILABLE: return None - from core.rag.tool import search_for_autoinject + from core.rag.tool import render_sources, search_for_autoinject, whole_document_context except Exception as exc: # noqa: BLE001 logger.warning("RAG auto-inject unavailable: %s", exc) return None + text: str | None = None + sources: list[dict] = [] + floor_override = rag_scope.get("autoinject_min_score") floor = float(floor_override) if floor_override is not None else _autoinject_floor() # Cap at the lean top_k, but honor a lower user setting. lean_k = _autoinject_top_k() sidebar_k = _opt_int(rag_scope.get("default_top_k")) top_k = min(sidebar_k, lean_k) if sidebar_k is not None else lean_k - try: - found = search_for_autoinject( - query = query, - scope_kb_id = rag_scope.get("kb_id"), - scope_thread_id = rag_scope.get("thread_id"), - scope_project_id = rag_scope.get("project_id"), - top_k = top_k, - min_dense_score = floor, - **_scope_retrieval_kwargs(rag_scope), - ) - except Exception as exc: # noqa: BLE001 - logger.warning("RAG auto-inject retrieval failed: %s", exc) - return None - if not found: - logger.info("RAG auto-inject: no passage >= %.2f; skipping", floor) + + # Whole-document mode: a thread-attached file under budget is injected in full so + # the model reads everything. A KB selection is exclusive, so whole-doc never + # preempts it; in a project chat the project sources are still retrieved top-K and + # appended under one citation numbering. Oversized files (or no thread doc) fall + # through to the combined top-K retrieval below. + if whole_doc_requested: + try: + budget = _whole_doc_budget(rag_scope, conversation) + + whole = whole_document_context( + scope_thread_id = thread_id, + max_tokens = budget, + ) + except Exception as exc: # noqa: BLE001 + logger.warning("RAG whole-document context failed: %s", exc) + whole = None + if whole is not None: + text, sources = whole + project_id = rag_scope.get("project_id") + if project_id: + try: + proj = search_for_autoinject( + query = query, + scope_project_id = project_id, + top_k = top_k, + min_dense_score = floor, + **_scope_retrieval_kwargs(rag_scope), + ) + except Exception as exc: # noqa: BLE001 + logger.warning("RAG project retrieval (whole-doc companion) failed: %s", exc) + proj = None + if proj is not None: + merged = sources + proj[1] + merged_text = render_sources(merged) + if max(1, len(merged_text) // 4) <= budget: + sources = merged + text = merged_text + logger.info("RAG auto-inject: whole-document context (%d chunk(s))", len(sources)) + + if text is None and enabled: + try: + found = search_for_autoinject( + query = query, + scope_kb_id = rag_scope.get("kb_id"), + scope_thread_id = rag_scope.get("thread_id"), + scope_project_id = rag_scope.get("project_id"), + top_k = top_k, + min_dense_score = floor, + **_scope_retrieval_kwargs(rag_scope), + ) + except Exception as exc: # noqa: BLE001 + logger.warning("RAG auto-inject retrieval failed: %s", exc) + return None + if not found: + logger.info("RAG auto-inject: no passage >= %.2f; skipping", floor) + return None + text, sources = found + if text is None: return None - text, sources = found import json as _json import uuid as _uuid @@ -1236,7 +1341,7 @@ def build_rag_autoinject(conversation: list[dict], rag_scope: dict | None) -> di "content": text, }, ] - logger.info("RAG auto-inject: %d passage(s) >= %.2f for %r", len(sources), floor, query[:80]) + logger.info("RAG auto-inject: %d passage(s) for %r", len(sources), query[:80]) return {"events": events, "messages": messages} diff --git a/studio/backend/core/rag/captioner.py b/studio/backend/core/rag/captioner.py index be8e341064..e29c9c9a7a 100644 --- a/studio/backend/core/rag/captioner.py +++ b/studio/backend/core/rag/captioner.py @@ -1,9 +1,12 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""Caption figures with the loaded vision model and splice the text into the page -so images are searchable via the normal FTS5 + dense path. No-op (never raises) -without a vision model or on failure; gated by ``config.CAPTION_IMAGES``.""" +"""Vision-model helpers for ingestion: figure captioning and scanned-page OCR. + +Both turn pixels into indexable text and are a no-op (never raise) without a loaded +vision model. They reuse the chat model's vision endpoint, so it must be served with +``--ubatch-size`` >= one image's tokens (some encoders, e.g. Gemma, attend +non-causally and abort otherwise); Studio's vision chat already requires this.""" from __future__ import annotations @@ -15,11 +18,54 @@ from . import config logger = logging.getLogger(__name__) _CAPTION_PROMPT = ( - "Describe this figure or image from a document in one or two concise " - "sentences, for search indexing. State what it depicts (e.g. a diagram, " - "chart, table or photo) and its key content. Do not add commentary." + "Read this figure or image from a document for search indexing.\n" + "First, on a line 'TEXT:', transcribe every piece of visible text exactly as " + "written, in reading order: the title, axis labels and units, legend and series " + "names, EVERY box / node / arrow label, table headers and cells, equations, and " + "footnotes. List each distinct label even if it is small.\n" + "Then, on a line 'SUMMARY:', add one or two sentences on what it shows (chart " + "type and trend, diagram subject, table topic, or photo content).\n" + "Report only what is visible. Transcribe exactly; do not invent or guess any " + "text, label, or number." ) +_OCR_PROMPT = ( + "Transcribe all text on this document page exactly as it appears, in reading " + "order, including any text inside figures, diagrams, charts, and tables (keep " + "table rows readable). Output only the transcribed text, with no commentary or " + "code fences. Preserve headings, lists, and line breaks. If the page has no " + "readable text, output nothing." +) + + +def _collapse_runaway( + text: str, + max_repeat: int = 3, + max_total: int = 8, +) -> str: + """Cap runaway repetition: vision models sometimes loop a line many times. Keep + each distinct line to ``max_repeat`` in a row and ``max_total`` total, and collapse + blank-line floods, so a degenerate page cannot flood the index.""" + out: list[str] = [] + seen: dict[str, int] = {} + prev: str | None = None + run = 0 + for line in text.splitlines(): + key = line.strip() + if not key: + if prev == "": # collapse runs of blank lines to a single separator + continue + prev = "" + out.append("") + continue + run = run + 1 if key == prev else 1 + prev = key + seen[key] = seen.get(key, 0) + 1 + if run > max_repeat or seen[key] > max_total: + continue + out.append(line) + return "\n".join(out) + def vision_endpoint() -> tuple[str, str] | None: """``(base_url, model)`` for a loaded vision GGUF model, else None.""" @@ -33,7 +79,28 @@ def vision_endpoint() -> tuple[str, str] | None: return None -def _caption_one(base_url: str, model: str, image_bytes: bytes, timeout: float) -> str | None: +def _vision_auth_headers() -> dict | None: + """Bearer header for the backend's API, or None. Vision calls share the chat + endpoint, so they need the same key under direct-stream (``--api-key``) mode.""" + try: + from routes.inference import get_llama_cpp_backend + return get_llama_cpp_backend()._auth_headers or None + except Exception: # noqa: BLE001 - auth discovery must never break ingestion + return None + + +def _vision_complete( + base_url: str, + model: str, + image_bytes: bytes, + *, + prompt: str, + timeout: float, + max_tokens: int, + temperature: float = 0.0, +) -> str | None: + """One image-in / text-out call to the loaded vision model's OpenAI-compatible + endpoint. Returns the stripped text or ``None`` on empty/failure (non-fatal).""" import httpx data_url = "data:image/png;base64," + base64.b64encode(image_bytes).decode("ascii") @@ -43,33 +110,62 @@ def _caption_one(base_url: str, model: str, image_bytes: bytes, timeout: float) { "role": "user", "content": [ - {"type": "text", "text": _CAPTION_PROMPT}, + {"type": "text", "text": prompt}, {"type": "image_url", "image_url": {"url": data_url}}, ], } ], - "max_tokens": 200, - "temperature": 0.2, + "max_tokens": max_tokens, + # Deterministic by default: transcription must not randomly drop labels. + "temperature": temperature, "stream": False, # Off: thinking models would spend the budget reasoning, returning "". "chat_template_kwargs": {"enable_thinking": False}, } try: - r = httpx.post(f"{base_url}/v1/chat/completions", json = payload, timeout = timeout) + r = httpx.post( + f"{base_url}/v1/chat/completions", + json = payload, + timeout = timeout, + headers = _vision_auth_headers(), + ) r.raise_for_status() text = r.json()["choices"][0]["message"]["content"] return text.strip() or None - except Exception: # noqa: BLE001 - a failed caption is non-fatal - logger.debug("caption request failed", exc_info = True) + except Exception: # noqa: BLE001 - a failed vision call is non-fatal + logger.debug("vision request failed", exc_info = True) return None +def _caption_one(base_url: str, model: str, image_bytes: bytes, timeout: float) -> str | None: + return _vision_complete( + base_url, + model, + image_bytes, + prompt = _CAPTION_PROMPT, + timeout = timeout, + max_tokens = config.CAPTION_MAX_TOKENS, + ) + + +def _ocr_one(base_url: str, model: str, image_bytes: bytes, timeout: float) -> str | None: + return _vision_complete( + base_url, + model, + image_bytes, + prompt = _OCR_PROMPT, + timeout = timeout, + max_tokens = config.OCR_MAX_TOKENS, + ) + + def caption_images( images: list, *, endpoint: tuple[str, str] | None = None ) -> dict[int, list[str]]: - """Caption ``ParsedImage`` objects, keyed by 1-based page number; ``{}`` when - disabled, no vision model, or no images. Bounded by ``CAPTION_MAX_IMAGES``.""" - if not config.CAPTION_IMAGES or not images: + """Caption ``ParsedImage`` objects, keyed by 1-based page number; ``{}`` when there + are no images or no vision model. The caller (`ingestion._run`) owns the on/off + policy. Bounded by ``CAPTION_MAX_IMAGES``; each caption passes ``_collapse_runaway``.""" + if not images: return {} ep = endpoint or vision_endpoint() if ep is None: @@ -84,7 +180,50 @@ def caption_images( caption = _caption_one(base_url, model, image_bytes, config.CAPTION_TIMEOUT_S) if caption: page = getattr(img, "page_number", None) or 0 - out.setdefault(int(page), []).append(caption) + out.setdefault(int(page), []).append(_collapse_runaway(caption)) + return out + + +def ocr_pages( + page_pngs: dict[int, bytes], *, endpoint: tuple[str, str] | None = None +) -> dict[int, str]: + """OCR rendered page PNGs (keyed by 1-based page number) to text; ``{}`` when there + is no vision model or no pages. The caller (`ingestion._ocr_scanned_pages`) owns the + on/off policy. Bounded by ``OCR_MAX_PAGES``.""" + if not page_pngs: + return {} + ep = endpoint or vision_endpoint() + if ep is None: + return {} + base_url, model = ep + + out: dict[int, str] = {} + for page_num in sorted(page_pngs)[: config.OCR_MAX_PAGES]: + text = _ocr_one(base_url, model, page_pngs[page_num], config.OCR_TIMEOUT_S) + if text: + out[int(page_num)] = _collapse_runaway(text) + return out + + +def merge_page_captions(captions: dict[int, list[str]]) -> dict[int, list[str]]: + """Merge a page's per-tile captions into one deduped block: drop lines repeated + across overlapping tiles (first kept, order preserved), then ``_collapse_runaway``, + so ``splice_captions`` adds a single figure block per page.""" + out: dict[int, list[str]] = {} + for page, caps in captions.items(): + seen: set[str] = set() + lines: list[str] = [] + for cap in caps: + for line in (cap or "").splitlines(): + stripped = line.strip() + key = stripped.lower() + if not stripped or key in seen: + continue + seen.add(key) + lines.append(stripped) + merged = _collapse_runaway("\n".join(lines)) + if merged.strip(): + out[page] = [merged] return out diff --git a/studio/backend/core/rag/config.py b/studio/backend/core/rag/config.py index 993423683c..54a224d081 100644 --- a/studio/backend/core/rag/config.py +++ b/studio/backend/core/rag/config.py @@ -17,13 +17,50 @@ TOP_K_DENSE = int(os.environ.get("RAG_TOP_K_DENSE", "30")) TOP_K_HYBRID = int(os.environ.get("RAG_TOP_K_HYBRID", "10")) RRF_K = int(os.environ.get("RAG_RRF_K", "60")) -UPLOAD_EXTS = {".pdf", ".txt", ".md", ".markdown", ".docx", ".html", ".htm"} +# Whole-document context: a thread-attached file under the token budget is injected +# in full (every chunk, in order) instead of top-K retrieval; above it, use retrieval. +THREAD_WHOLE_DOC = os.environ.get("RAG_THREAD_WHOLE_DOC", "1") == "1" +WHOLE_DOC_MAX_TOKENS = int(os.environ.get("RAG_WHOLE_DOC_MAX_TOKENS", "6000")) -# Figure captioning via the loaded vision model; off by default since each caption -# is a model call. MAX_IMAGES bounds per-doc cost. -CAPTION_IMAGES = os.environ.get("RAG_CAPTION_IMAGES", "0") == "1" -CAPTION_MAX_IMAGES = int(os.environ.get("RAG_CAPTION_MAX_IMAGES", "8")) -CAPTION_TIMEOUT_S = float(os.environ.get("RAG_CAPTION_TIMEOUT_S", "30")) +UPLOAD_EXTS = {".pdf", ".txt", ".md", ".markdown", ".docx", ".html", ".htm"} +# Reject uploads larger than this, so one pathological file can't drive unbounded parse +# + vision work at ingest. 0 disables the cap. Default 200 MB. +MAX_UPLOAD_BYTES = int(os.environ.get("RAG_MAX_UPLOAD_BYTES", str(200 * 1024 * 1024))) + +# Extract PDF text as layout-aware Markdown (pymupdf4llm) instead of flat text, so +# tables, headings and lists survive into chunks and retrieval. Falls back to plain +# PyMuPDF text when off, when pymupdf4llm is missing, or when extraction fails. +PDF_MARKDOWN = os.environ.get("RAG_PDF_MARKDOWN", "1") == "1" + +# Figure captioning via the loaded vision model: detected figures are transcribed + +# described so they become searchable. On by default, a no-op without a vision model; +# the chat's "Describe figures & charts" toggle overrides it per upload. +CAPTION_IMAGES = os.environ.get("RAG_CAPTION_IMAGES", "1") == "1" +# Total per-document tile budget (figure-bearing pages are tiled, see below). +CAPTION_MAX_IMAGES = int(os.environ.get("RAG_CAPTION_MAX_IMAGES", "24")) +CAPTION_TIMEOUT_S = float(os.environ.get("RAG_CAPTION_TIMEOUT_S", "60")) +# Larger than a one-line caption since captions transcribe every label. FIGURE_DPI is +# high enough to keep small box/axis labels legible when tiles are rendered. +CAPTION_MAX_TOKENS = int(os.environ.get("RAG_CAPTION_MAX_TOKENS", "768")) +FIGURE_DPI = int(os.environ.get("RAG_FIGURE_DPI", "200")) +# Figure pages are tiled into an overlapping ROWS x COLS grid of high-DPI tiles (plus +# an optional full page), so small labels and every sub-figure are covered without +# exact region detection. MAX_PAGES bounds figure pages; MAX_IMAGES bounds total tiles. +FIGURE_TILE_ROWS = int(os.environ.get("RAG_FIGURE_TILE_ROWS", "2")) +FIGURE_TILE_COLS = int(os.environ.get("RAG_FIGURE_TILE_COLS", "2")) +FIGURE_TILE_OVERLAP = float(os.environ.get("RAG_FIGURE_TILE_OVERLAP", "0.12")) +FIGURE_FULLPAGE = os.environ.get("RAG_FIGURE_FULLPAGE", "1") == "1" +CAPTION_MAX_PAGES = int(os.environ.get("RAG_CAPTION_MAX_PAGES", "4")) + +# Scanned-PDF OCR: a page with little extractable text is rendered and transcribed by +# the vision model so it becomes searchable. Needs a vision model, else skipped (page +# stays empty). MIN_CHARS is the text length below which a page is treated as scanned. +OCR_SCANNED = os.environ.get("RAG_OCR_SCANNED", "1") == "1" +OCR_MIN_CHARS = int(os.environ.get("RAG_OCR_MIN_CHARS", "16")) +OCR_MAX_PAGES = int(os.environ.get("RAG_OCR_MAX_PAGES", "20")) +OCR_DPI = int(os.environ.get("RAG_OCR_DPI", "150")) +OCR_TIMEOUT_S = float(os.environ.get("RAG_OCR_TIMEOUT_S", "60")) +OCR_MAX_TOKENS = int(os.environ.get("RAG_OCR_MAX_TOKENS", "2048")) # Embedder backend. "auto": sentence-transformers on a CUDA/ROCm GPU (torch fp16 # wins bulk indexing), else torch-free GGUF llama-server. Switching backends changes diff --git a/studio/backend/core/rag/ingestion.py b/studio/backend/core/rag/ingestion.py index 77bbdc92f5..04365ab76b 100644 --- a/studio/backend/core/rag/ingestion.py +++ b/studio/backend/core/rag/ingestion.py @@ -99,25 +99,108 @@ def _embed_all(texts: list[str], model_name: str | None): return vectors +def _ocr_scanned_pages( + pages: list, + stored_path: str, + conn, + job_id: str, + ocr: bool | None = None, +) -> tuple[list, set[int]]: + """Replace text on near-empty (scanned/image-only) PDF pages with vision-model OCR + so image PDFs become searchable. ``ocr`` overrides ``config.OCR_SCANNED`` per upload + (``None`` = config default); no-op without scanned pages or a vision model. OCR'd + pages have no text layer, so no preview highlight regions, but stay searchable. + Returns ``(pages, ocred)``: new ``Page`` objects for OCR'd pages (originals + otherwise) and the set of page numbers actually transcribed.""" + if not (config.OCR_SCANNED if ocr is None else ocr): + return pages, set() + scanned = [ + p.page_number + for p in pages + if p.page_number is not None and len((p.text or "").strip()) < config.OCR_MIN_CHARS + ] + if not scanned or captioner.vision_endpoint() is None: + return pages, set() + if len(scanned) > config.OCR_MAX_PAGES: + logger.warning( + "OCR: %d scanned pages exceed OCR_MAX_PAGES=%d; pages past the cap stay " + "untranscribed (raise RAG_OCR_MAX_PAGES to cover them)", + len(scanned), + config.OCR_MAX_PAGES, + ) + scanned = scanned[: config.OCR_MAX_PAGES] + _progress(conn, job_id, "ocr", 0.25) + page_pngs = parsers.render_pdf_pages(stored_path, scanned, dpi = config.OCR_DPI) + texts = captioner.ocr_pages(page_pngs) + if not texts: + return pages, set() + + from .parsers import Page + + out: list = [] + ocred: set[int] = set() + for page in pages: + text = texts.get(page.page_number) + if text: + original = (page.text or "").strip() + merged = text if not original or original in text else f"{original}\n\n{text}" + out.append(Page(text = merged, page_number = page.page_number, char_count = len(merged))) + ocred.add(page.page_number) + else: + out.append(page) + return out, ocred + + def _run( - job_id: str, document_id: str, scope: str, stored_path: str, model_name: str | None + job_id: str, + document_id: str, + scope: str, + stored_path: str, + model_name: str | None, + ocr: bool | None = None, + caption: bool | None = None, ) -> None: conn = rag_db.get_connection() try: _progress(conn, job_id, "parsing", 0.1) pages = parsers.parse(stored_path) - if config.CAPTION_IMAGES and stored_path.lower().endswith(".pdf"): - # Caption figures, splice into page text (no-op without a vision model). + is_pdf = stored_path.lower().endswith(".pdf") + ocred: set[int] = set() + if is_pdf: + pages, ocred = _ocr_scanned_pages(pages, stored_path, conn, job_id, ocr = ocr) + caption_on = config.CAPTION_IMAGES if caption is None else caption + # Skip all figure work (PDF rasterization included) without a vision model. + if caption_on and is_pdf and captioner.vision_endpoint() is not None: + # Tile figure pages, transcribe+describe each tile, then merge/dedup/splice + # into the page text so small labels and every sub-figure are captured. try: - figures = parsers.render_pdf_figures( - stored_path, max_figures = config.CAPTION_MAX_IMAGES + fig_pages = parsers.pages_with_figures( + stored_path, + max_pages = config.CAPTION_MAX_PAGES, + # Skip only pages OCR actually transcribed (it covers them whole); a + # scanned figure page past the OCR cap or with empty OCR still tiles. + exclude_pages = ocred, + ) + tiles = ( + parsers.render_pdf_figure_tiles( + stored_path, + fig_pages, + dpi = config.FIGURE_DPI, + rows = config.FIGURE_TILE_ROWS, + cols = config.FIGURE_TILE_COLS, + overlap = config.FIGURE_TILE_OVERLAP, + fullpage = config.FIGURE_FULLPAGE, + max_tiles = config.CAPTION_MAX_IMAGES, + ) + if fig_pages + else [] ) except Exception: - logger.warning("figure rendering failed for job %s", job_id, exc_info = True) - figures = [] - if figures: - _progress(conn, job_id, "captioning", 0.2) - captions = captioner.caption_images(figures) + logger.warning("figure tiling failed for job %s", job_id, exc_info = True) + tiles = [] + if tiles: + _progress(conn, job_id, "captioning", 0.28) + captions = captioner.merge_page_captions(captioner.caption_images(tiles)) pages = captioner.splice_captions(pages, captions) _progress(conn, job_id, "chunking", 0.3) @@ -175,6 +258,8 @@ def start_ingestion( *, project_id: str | None = None, model_name: str | None = None, + ocr: bool | None = None, + caption: bool | None = None, ) -> tuple[str, str]: """Create the document + job rows and spawn the worker, returning ``(document_id, job_id)``. A duplicate content hash in this scope returns the @@ -191,13 +276,26 @@ def start_ingestion( try: existing = store.document_by_hash(conn, scope, sha) if existing is not None: - job_id = _new_job(conn, existing, scope, status = "completed", progress = 1.0) - _remove_upload(stored_path) - with _jobs_lock: - _jobs[job_id] = queue.Queue() - _emit(job_id, {"type": "complete", "num_chunks": 0, "deduped": True}) - _emit(job_id, None) - return existing, job_id + doc = store.get_document(conn, existing) + empty_completed = ( + doc is not None and doc.get("status") == "completed" and not doc.get("num_chunks") + ) + if empty_completed: + # A prior ingest of identical bytes yielded zero chunks (e.g. a scanned + # PDF uploaded before a vision model loaded). Re-ingest, don't dedupe. + store.delete_document(conn, existing) + _remove_upload(doc.get("stored_path"), keep_path = stored_path) + else: + job_id = _new_job(conn, existing, scope, status = "completed", progress = 1.0) + _remove_upload(stored_path) + with _jobs_lock: + _jobs[job_id] = queue.Queue() + _emit( + job_id, + {"type": "complete", "num_chunks": doc.get("num_chunks") or 0, "deduped": True}, + ) + _emit(job_id, None) + return existing, job_id for failed in store.failed_documents_by_hash(conn, scope, sha): store.delete_document(conn, failed["id"]) _remove_upload(failed.get("stored_path"), keep_path = stored_path) @@ -221,7 +319,7 @@ def start_ingestion( _jobs[job_id] = queue.Queue() threading.Thread( target = _run, - args = (job_id, document_id, scope, stored_path, model_name), + args = (job_id, document_id, scope, stored_path, model_name, ocr, caption), daemon = True, ).start() return document_id, job_id @@ -339,10 +437,16 @@ def job_events(job_id: str): def get_job_status(job_id: str) -> dict | None: - """Read the persisted ingestion job row (status / stage / progress / error).""" + """Read the persisted ingestion job row (status / stage / progress / error), plus + the document's ``num_chunks`` so a client polling to completion learns the chunk + count (the SSE ``complete`` frame carries it, but the poll/reconcile path does not).""" conn = rag_db.get_connection() try: - row = conn.execute("SELECT * FROM ingestion_jobs WHERE id=?", (job_id,)).fetchone() + row = conn.execute( + "SELECT j.*, d.num_chunks AS num_chunks FROM ingestion_jobs j " + "LEFT JOIN documents d ON d.id = j.document_id WHERE j.id=?", + (job_id,), + ).fetchone() return dict(row) if row else None finally: conn.close() diff --git a/studio/backend/core/rag/locators.py b/studio/backend/core/rag/locators.py index 57c0487486..9331bb15ac 100644 --- a/studio/backend/core/rag/locators.py +++ b/studio/backend/core/rag/locators.py @@ -39,9 +39,11 @@ def _norm_token(token: str) -> str: def _anchor_tokens(page_text: str, match: LocatorMatch) -> list[str]: """Normalized anchor tokens from the chunk's leading span. Drops first and last - token (boundaries often slice mid-word) when long enough.""" + token (boundaries often slice mid-word) when long enough. Pipes are split out so + Markdown table cells (``|Q1|$1.2M|``) become individual words that match the PDF + word stream.""" segment = page_text[match.start : match.end] - raw = segment.split() + raw = segment.replace("|", " ").split() if len(raw) >= MIN_ANCHOR_WORDS + 2: raw = raw[1:-1] tokens = [t for t in (_norm_token(w) for w in raw) if t] diff --git a/studio/backend/core/rag/parsers.py b/studio/backend/core/rag/parsers.py index 84da941762..ba248cf9a6 100644 --- a/studio/backend/core/rag/parsers.py +++ b/studio/backend/core/rag/parsers.py @@ -15,6 +15,8 @@ import os from dataclasses import dataclass from html.parser import HTMLParser +from . import config + logger = logging.getLogger(__name__) @@ -67,6 +69,28 @@ def _html(raw: str) -> list[Page]: return [_page("\n".join(parser.out), 1)] +def _pdf_markdown(doc) -> list[str] | None: + """Per-page layout-aware Markdown (tables, headings, lists) via pymupdf4llm; index + i maps to page i+1. Returns None when the lib is missing, extraction fails, or the + page count does not line up, so the caller falls back to plain PyMuPDF text.""" + try: + import pymupdf4llm + except Exception: + return None + try: + chunks = pymupdf4llm.to_markdown( + doc, + page_chunks = True, + show_progress = False, + ) + except Exception: # noqa: BLE001 - never let Markdown extraction break ingestion + logger.warning("pymupdf4llm extraction failed; using plain text", exc_info = True) + return None + if not isinstance(chunks, list) or len(chunks) != doc.page_count: + return None + return [str(c.get("text") or "") for c in chunks] + + def _pdf(path: str, want_images: bool) -> tuple[list[Page], list[ParsedImage]]: import fitz # PyMuPDF @@ -74,8 +98,11 @@ def _pdf(path: str, want_images: bool) -> tuple[list[Page], list[ParsedImage]]: images: list[ParsedImage] = [] doc = fitz.open(path) try: + md = _pdf_markdown(doc) if config.PDF_MARKDOWN else None for i, page in enumerate(doc): - text = page.get_text("text") or "" + # Prefer layout-aware Markdown (keeps tables/headings legible for retrieval); + # fall back to plain text when Markdown is off, unavailable, or empty here. + text = (md[i] if md else "") or page.get_text("text") or "" pages.append(_page(text, i + 1)) if want_images: for img in page.get_images(full = True): @@ -118,63 +145,164 @@ def _merge_rects(boxes: list) -> list: return merged -def render_pdf_figures( - path: str, +def _figure_boxes( + page, *, - dpi: int = 130, min_area_frac: float = 0.04, min_side: float = 40.0, - max_figures: int = 8, -) -> list[ParsedImage]: - """Detect figure regions and render each to a PNG for captioning. +) -> list: + """Qualifying figure-region rectangles on a page: cluster vector drawings + raster + placements, merge overlaps, keep the page-spanning ones (area/side filtered).""" + boxes: list = [] + try: + boxes.extend(info["bbox"] for info in page.get_image_info()) + except Exception: + pass + try: + boxes.extend(page.cluster_drawings()) + except Exception: + pass + if not boxes: + return [] + page_area = page.rect.width * page.rect.height + keep: list = [] + for box in _merge_rects(boxes): + if ( + box.get_area() >= min_area_frac * page_area + and box.width >= min_side + and box.height >= min_side + ): + keep.append(box) + return keep - Academic figures are vector, so raster extraction yields fragments; instead - cluster vector drawings + raster placements into boxes, keep the page-spanning - ones, and render them. Any failure yields [], never an exception. - """ + +def pages_with_figures( + path: str, + *, + max_pages: int = 4, + min_area_frac: float = 0.04, + min_side: float = 40.0, + exclude_pages: set[int] | None = None, +) -> list[int]: + """1-based page numbers with a qualifying figure region, capped at ``max_pages``; + drives figure tiling. ``exclude_pages`` (1-based) are skipped: those are the pages + OCR already transcribed whole, so tiling them would duplicate the vision work. Any + failure yields [].""" + exclude = exclude_pages or set() try: import pymupdf except Exception: return [] - - out: list[ParsedImage] = [] try: doc = pymupdf.open(path) except Exception: return [] + pages: list[int] = [] try: for i, page in enumerate(doc): - boxes: list = [] - try: - boxes.extend(info["bbox"] for info in page.get_image_info()) - except Exception: - pass - try: - boxes.extend(page.cluster_drawings()) - except Exception: - pass - if not boxes: + if (i + 1) in exclude: continue - page_area = page.rect.width * page.rect.height - for box in _merge_rects(boxes): - if ( - box.get_area() >= min_area_frac * page_area - and box.width >= min_side - and box.height >= min_side - ): - try: - pix = page.get_pixmap(dpi = dpi, clip = box) - out.append( - ParsedImage( - image_bytes = pix.tobytes("png"), - page_number = i + 1, - xref = 0, - ) + if _figure_boxes(page, min_area_frac = min_area_frac, min_side = min_side): + pages.append(i + 1) + if len(pages) >= max_pages: + break + return pages + finally: + doc.close() + + +def render_pdf_figure_tiles( + path: str, + page_numbers, + *, + dpi: int = 200, + rows: int = 2, + cols: int = 2, + overlap: float = 0.12, + fullpage: bool = True, + max_tiles: int = 24, +) -> list[ParsedImage]: + """Render figure-bearing pages as overlapping high-DPI tiles (plus an optional full + page), each a ``ParsedImage`` keyed by page number. Tiling keeps small labels legible + and covers every sub-figure without exact region detection. Any failure yields [].""" + wanted = [int(n) for n in page_numbers] + if not wanted: + return [] + rows, cols = max(1, int(rows)), max(1, int(cols)) # never divide by zero + try: + import pymupdf + except Exception: + return [] + try: + doc = pymupdf.open(path) + except Exception: + return [] + out: list[ParsedImage] = [] + try: + for num in wanted: + if num < 1 or num > doc.page_count: + continue + page = doc[num - 1] + rect = page.rect + clips: list = [rect] if fullpage else [] + cw, ch = rect.width / cols, rect.height / rows + ox, oy = cw * overlap, ch * overlap + for r in range(rows): + for c in range(cols): + clips.append( + pymupdf.Rect( + rect.x0 + c * cw - ox, + rect.y0 + r * ch - oy, + rect.x0 + (c + 1) * cw + ox, + rect.y0 + (r + 1) * ch + oy, ) - except Exception: - continue - if len(out) >= max_figures: - return out + & rect + ) + for clip in clips: + try: + pix = page.get_pixmap(dpi = dpi, clip = clip) + out.append(ParsedImage(image_bytes = pix.tobytes("png"), page_number = num, xref = 0)) + except Exception: + continue + if len(out) >= max_tiles: + return out + return out + finally: + doc.close() + + +def render_pdf_pages( + path: str, + page_numbers, + *, + dpi: int = 150, +) -> dict[int, bytes]: + """Render whole PDF pages (given as 1-based numbers) to PNG bytes, keyed by + page number. Backs scanned-page OCR. Any failure yields ``{}`` (or skips that + page), never an exception. + """ + wanted = {int(n) for n in page_numbers} + if not wanted: + return {} + try: + import pymupdf + except Exception: + return {} + try: + doc = pymupdf.open(path) + except Exception: + return {} + out: dict[int, bytes] = {} + try: + for i, page in enumerate(doc): + num = i + 1 + if num not in wanted: + continue + try: + pix = page.get_pixmap(dpi = dpi) + out[num] = pix.tobytes("png") + except Exception: + continue return out finally: doc.close() diff --git a/studio/backend/core/rag/store.py b/studio/backend/core/rag/store.py index 7d58931e53..8e59c5fbf6 100644 --- a/studio/backend/core/rag/store.py +++ b/studio/backend/core/rag/store.py @@ -292,3 +292,40 @@ def chunks_by_id(conn: sqlite3.Connection, ids) -> dict: list(ids), ).fetchall() return {r["id"]: r for r in rows} + + +def all_chunks_for_scope(conn: sqlite3.Connection, scope) -> list[dict]: + """Every completed-document chunk for a scope, ordered document-then-index and + joined with the document filename. Backs whole-document context injection, so + it does no retrieval or embedding.""" + scopes = _scopes(scope) + if not scopes: + return [] + placeholders = ",".join("?" * len(scopes)) + rows = conn.execute( + f"SELECT c.id, c.text, c.document_id, c.chunk_index, c.page_number, " + f"c.token_count, d.filename, d.created_at " + f"FROM chunks c JOIN documents d ON d.id=c.document_id " + f"WHERE c.scope IN ({placeholders}) AND d.status='completed' " + f"ORDER BY d.created_at, c.document_id, c.chunk_index", + list(scopes), + ).fetchall() + return [dict(r) for r in rows] + + +def scope_token_estimate(conn: sqlite3.Connection, scope) -> int: + """Upper-bound token total for a scope's completed chunks without hydrating text. + Mirrors ``all_chunks_for_scope`` + the ``tool._row_token_count`` fallback (stored + count, else length/4), so the whole-doc budget can be checked before loading text.""" + scopes = _scopes(scope) + if not scopes: + return 0 + placeholders = ",".join("?" * len(scopes)) + row = conn.execute( + f"SELECT COALESCE(SUM(CASE WHEN c.token_count > 0 THEN c.token_count " + f"ELSE MAX(1, length(COALESCE(c.text, '')) / 4) END), 0) AS total " + f"FROM chunks c JOIN documents d ON d.id=c.document_id " + f"WHERE c.scope IN ({placeholders}) AND d.status='completed'", + list(scopes), + ).fetchone() + return int(row["total"] or 0) diff --git a/studio/backend/core/rag/tool.py b/studio/backend/core/rag/tool.py index ccb1b47e63..b05f8dd3a3 100644 --- a/studio/backend/core/rag/tool.py +++ b/studio/backend/core/rag/tool.py @@ -16,7 +16,13 @@ from xml.sax.saxutils import quoteattr from storage import rag_db from . import config, retrieval -from .store import kb_scope, project_scope, thread_scope +from .store import ( + all_chunks_for_scope, + kb_scope, + project_scope, + scope_token_estimate, + thread_scope, +) SEARCH_KNOWLEDGE_BASE_TOOL = { "type": "function", @@ -90,6 +96,30 @@ def _format(rows, hits) -> tuple[str, list[dict]]: return "\n\n".join(blocks), sources +def render_sources(sources: list[dict]) -> str: + """Render a citation-source list to sequentially-numbered ```` blocks, + rewriting each source's ``citationId`` to match its 1-based position. Lets + independently-built source lists (a whole-document thread attachment plus + retrieved project passages) be merged under one citation numbering.""" + blocks: list[str] = [] + for i, s in enumerate(sources, 1): + s["citationId"] = i + src = quoteattr(s.get("filename") or "unknown") + page = s.get("page") + page_attr = f" page={quoteattr(str(page))}" if page else "" + blocks.append(f'\n{s.get("text") or ""}\n') + return "\n\n".join(blocks) + + +def _row_token_count(row) -> int: + """Chunk token count for budgeting, falling back to a length estimate when the + stored count is missing or zero, so a malformed chunk cannot bypass the budget.""" + tc = row["token_count"] + if tc: + return int(tc) + return max(1, len(row["text"] or "") // 4) + + def search_knowledge_base_with_sources( *, query: str, @@ -186,6 +216,55 @@ def search_for_autoinject( return (text, sources) if sources else None +def whole_document_context( + *, scope_thread_id: str | None = None, max_tokens: int +) -> tuple[str, list[dict]] | None: + """Render EVERY chunk of the THREAD's attached documents (in order) as the same + ```` blocks + citation source-map as retrieval, so the model reads the whole + file rather than top-K passages. Thread-attached files only: KB and project corpora + are search corpora, never whole-document, so this resolves the thread scope alone. + ``None`` (caller falls back to retrieval) when there is no thread scope, no completed + chunks, or the total exceeds ``max_tokens``.""" + if not scope_thread_id: + return None + # A non-positive budget means "never inject" (disable whole-doc via + # RAG_THREAD_WHOLE_DOC=0), not "inject the whole corpus unbounded". + if max_tokens <= 0: + return None + scope = thread_scope(scope_thread_id) + conn = rag_db.get_connection() + try: + # Cheap budget pre-check (SUM, no text hydration): reject an oversized attachment + # before loading the whole corpus; all_chunks_for_scope runs only once it fits. + if scope_token_estimate(conn, scope) > max_tokens: + return None + rows = all_chunks_for_scope(conn, scope) + finally: + conn.close() + if not rows: + return None + total = sum(_row_token_count(r) for r in rows) + if total > max_tokens: + return None + + sources: list[dict] = [ + { + "citationId": i, + "chunkId": r["id"], + "documentId": r["document_id"], + "filename": r["filename"] or "unknown", + "page": r["page_number"], + "text": r["text"] or "", + "score": None, + } + for i, r in enumerate(rows, 1) + ] + rendered = render_sources(sources) + if max(1, len(rendered) // 4) > max_tokens: + return None + return rendered, sources + + def search_knowledge_base( *, query: str, diff --git a/studio/backend/requirements/no-torch-runtime.txt b/studio/backend/requirements/no-torch-runtime.txt index 9796fc3a50..de321f80ed 100644 --- a/studio/backend/requirements/no-torch-runtime.txt +++ b/studio/backend/requirements/no-torch-runtime.txt @@ -73,4 +73,9 @@ pillow # this file installs --no-deps; without them Studio runs with RAG disabled. sqlite-vec==0.1.9 pymupdf==1.27.2.3 +# 0.3.x keeps pymupdf-layout (which pulls onnxruntime) an optional extra; the +# lockstep 1.27.x line makes it a hard dep we do not need for to_markdown(). +pymupdf4llm==0.3.4 python-docx==1.2.0 + +lxml==6.0.2 diff --git a/studio/backend/requirements/studio.txt b/studio/backend/requirements/studio.txt index 1b7f7a668c..6f4a5c3292 100644 --- a/studio/backend/requirements/studio.txt +++ b/studio/backend/requirements/studio.txt @@ -22,4 +22,7 @@ fastmcp>=3.0.2 # extras-no-deps.txt; these add the lexical+dense store and document parsing. sqlite-vec==0.1.9 pymupdf==1.27.2.3 +# 0.3.x keeps pymupdf-layout (which pulls onnxruntime) an optional extra; the +# lockstep 1.27.x line makes it a hard dep we do not need for to_markdown(). +pymupdf4llm==0.3.4 python-docx==1.2.0 diff --git a/studio/backend/routes/rag.py b/studio/backend/routes/rag.py index 8d23240fd5..4e35fce3c2 100644 --- a/studio/backend/routes/rag.py +++ b/studio/backend/routes/rag.py @@ -19,7 +19,7 @@ import secrets import time import uuid -from fastapi import APIRouter, Depends, File, HTTPException, Query, UploadFile +from fastapi import APIRouter, Depends, File, Form, HTTPException, Query, UploadFile from fastapi.responses import FileResponse, StreamingResponse from pydantic import BaseModel, Field @@ -62,13 +62,24 @@ def _save_upload(file: UploadFile) -> tuple[str, str]: uploads = ensure_dir(rag_uploads_root()) stored_path = str(uploads / f"{uuid.uuid4().hex}{ext}") size = 0 + cap = config.MAX_UPLOAD_BYTES + too_big = False with open(stored_path, "wb") as out: while True: block = file.file.read(1 << 20) if not block: break size += len(block) + if cap and size > cap: + too_big = True + break out.write(block) + if too_big: + os.remove(stored_path) + raise HTTPException( + status_code = 413, + detail = f"File exceeds the {cap // (1024 * 1024)} MB upload limit.", + ) if size == 0: os.remove(stored_path) raise HTTPException(status_code = 400, detail = "Uploaded file is empty.") @@ -207,6 +218,8 @@ def delete_knowledge_base(kb_id: str, subject: str = Depends(get_current_subject async def upload_kb_document( kb_id: str, file: UploadFile = File(...), + ocr: bool | None = Form(None), + caption: bool | None = Form(None), subject: str = Depends(get_current_subject), ) -> dict: _require_rag() @@ -218,7 +231,7 @@ async def upload_kb_document( conn.close() stored_path, filename = _save_upload(file) document_id, job_id = ingestion.start_ingestion( - store.kb_scope(kb_id), kb_id, None, filename, stored_path + store.kb_scope(kb_id), kb_id, None, filename, stored_path, ocr = ocr, caption = caption ) return {"documentId": document_id, "jobId": job_id, "filename": filename} @@ -238,12 +251,20 @@ def list_kb_documents(kb_id: str, subject: str = Depends(get_current_subject)) - async def upload_thread_document( thread_id: str, file: UploadFile = File(...), + ocr: bool | None = Form(None), + caption: bool | None = Form(None), subject: str = Depends(get_current_subject), ) -> dict: _require_rag() stored_path, filename = _save_upload(file) document_id, job_id = ingestion.start_ingestion( - store.thread_scope(thread_id), None, thread_id, filename, stored_path + store.thread_scope(thread_id), + None, + thread_id, + filename, + stored_path, + ocr = ocr, + caption = caption, ) return {"documentId": document_id, "jobId": job_id, "filename": filename} @@ -263,6 +284,8 @@ def list_thread_documents(thread_id: str, subject: str = Depends(get_current_sub async def upload_project_document( project_id: str, file: UploadFile = File(...), + ocr: bool | None = Form(None), + caption: bool | None = Form(None), subject: str = Depends(get_current_subject), ) -> dict: _require_rag() @@ -278,6 +301,8 @@ async def upload_project_document( filename, stored_path, project_id = project_id, + ocr = ocr, + caption = caption, ) return {"documentId": document_id, "jobId": job_id, "filename": filename} @@ -321,6 +346,7 @@ def job_status(job_id: str, subject: str = Depends(get_current_subject)) -> dict "stage": row.get("stage"), "progress": row.get("progress") or 0.0, "error": row.get("error"), + "numChunks": row.get("num_chunks") or 0, } diff --git a/studio/backend/storage/rag_db.py b/studio/backend/storage/rag_db.py index 4da600d768..ce27326562 100644 --- a/studio/backend/storage/rag_db.py +++ b/studio/backend/storage/rag_db.py @@ -119,6 +119,10 @@ def get_connection() -> sqlite3.Connection: ensure_dir(db_path.parent) conn = sqlite3.connect(str(db_path)) conn.row_factory = sqlite3.Row + # Wait for a lock instead of erroring immediately: a figure/scan-heavy ingest can + # hold its connection across many seconds of vision calls, and a concurrent ingest + # or autoinject read would otherwise hit "database is locked". + conn.execute("PRAGMA busy_timeout = 5000") try: conn.enable_load_extension(True) sqlite_vec.load(conn) diff --git a/studio/backend/tests/test_rag_captioning.py b/studio/backend/tests/test_rag_captioning.py index 5d83a7d38d..f475c9e374 100644 --- a/studio/backend/tests/test_rag_captioning.py +++ b/studio/backend/tests/test_rag_captioning.py @@ -13,13 +13,15 @@ def _img(page): return ParsedImage(image_bytes = b"\x89PNG fake", page_number = page, xref = page) -def test_caption_images_disabled_by_default(monkeypatch): +def test_caption_images_runs_when_images_present(monkeypatch): + # Policy lives in ingestion (_run); caption_images captions given images + endpoint. monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", False) - assert captioner.caption_images([_img(1)], endpoint = ("http://x", "local")) == {} + monkeypatch.setattr(captioner, "_caption_one", lambda *a: "a chart") + out = captioner.caption_images([_img(1)], endpoint = ("http://x", "local")) + assert out == {1: ["a chart"]} def test_caption_images_groups_by_page(monkeypatch): - monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", True) monkeypatch.setattr(captioner.config, "CAPTION_MAX_IMAGES", 8) monkeypatch.setattr(captioner, "_caption_one", lambda base, model, b, t: "a chart of results") out = captioner.caption_images([_img(1), _img(1), _img(3)], endpoint = ("http://x", "local")) @@ -27,7 +29,6 @@ def test_caption_images_groups_by_page(monkeypatch): def test_caption_images_respects_cap(monkeypatch): - monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", True) monkeypatch.setattr(captioner.config, "CAPTION_MAX_IMAGES", 2) calls = [] monkeypatch.setattr(captioner, "_caption_one", lambda *a: (calls.append(1) or "cap")) @@ -36,11 +37,183 @@ def test_caption_images_respects_cap(monkeypatch): def test_caption_images_no_endpoint(monkeypatch): - monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", True) monkeypatch.setattr(captioner, "vision_endpoint", lambda: None) assert captioner.caption_images([_img(1)]) == {} +def test_caption_runaway_guard_applied(monkeypatch): + # A looping vision model must not flood the index; captions pass _collapse_runaway. + monkeypatch.setattr(captioner, "_caption_one", lambda *a: "\n".join(["LOOP"] * 40)) + out = captioner.caption_images([_img(1)], endpoint = ("http://x", "local")) + assert out[1][0].splitlines().count("LOOP") == 3 # 40 -> 3 + + +def test_caption_prompt_and_token_budget(monkeypatch): + # Caption and OCR keep separate prompts + token caps over the shared _vision_complete. + captured: dict = {} + + def fake_vision_complete(base_url, model, image_bytes, *, prompt, timeout, max_tokens): + captured.update(prompt = prompt, timeout = timeout, max_tokens = max_tokens) + return "ok" + + monkeypatch.setattr(captioner, "_vision_complete", fake_vision_complete) + monkeypatch.setattr(captioner.config, "CAPTION_MAX_TOKENS", 277) + + captioner._caption_one("http://x", "local", b"img", 12.0) + prompt = captured["prompt"].lower() + # Unified prompt: transcribe every label (recall) + axis/legend coverage + describe. + assert "transcribe" in prompt + assert ("axis" in prompt or "axes" in prompt) and "legend" in prompt + assert "do not invent" in prompt + assert captured["max_tokens"] == 277 + assert captured["timeout"] == 12.0 + + captured.clear() + monkeypatch.setattr(captioner.config, "OCR_MAX_TOKENS", 999) + captioner._ocr_one("http://x", "local", b"img", 5.0) + assert captured["max_tokens"] == 999 + assert "transcribe" in captured["prompt"].lower() + + +def test_pages_with_figures_and_tiles(tmp_path): + from core.rag import parsers + + pdf = tmp_path / "fig.pdf" + _figure_pdf(pdf) + pgs = parsers.pages_with_figures(str(pdf), max_pages = 4) + assert pgs == [1] + tiles = parsers.render_pdf_figure_tiles(str(pdf), pgs, rows = 2, cols = 2, fullpage = True) + assert len(tiles) == 5 # full page + 2x2 grid + assert all(t.image_bytes[:8] == b"\x89PNG\r\n\x1a\n" and t.page_number == 1 for t in tiles) + capped = parsers.render_pdf_figure_tiles( + str(pdf), pgs, rows = 2, cols = 2, fullpage = True, max_tiles = 3 + ) + assert len(capped) == 3 # max_tiles budget honored + + +def test_render_pdf_figure_tiles_zero_grid_no_crash(tmp_path): + # A misconfigured rows/cols=0 must clamp to 1, not raise ZeroDivisionError. + import pymupdf + + from core.rag import parsers + + pdf = tmp_path / "blank.pdf" + doc = pymupdf.open() + doc.new_page() + doc.save(str(pdf)) + doc.close() + + out = parsers.render_pdf_figure_tiles(str(pdf), [1], rows = 0, cols = 0, fullpage = True) + assert len(out) == 2 # full page + a single 1x1 tile, no crash + + +def test_pages_with_figures_excludes_given_pages(tmp_path): + # Pages OCR already transcribed (passed as exclude_pages) are skipped; every other + # figure page is still returned for tiling. + import pymupdf + + from core.rag import parsers + + def _draw_chart(page): + shape = page.new_shape() + shape.draw_rect(pymupdf.Rect(60, 140, 540, 520)) + for i in range(8): + shape.draw_line((80, 160 + i * 40), (520, 160 + i * 40)) + shape.finish(color = (0, 0, 0), fill = (0.8, 0.8, 0.9)) + shape.commit() + + pdf = tmp_path / "charts.pdf" + doc = pymupdf.open() + _draw_chart(doc.new_page()) + _draw_chart(doc.new_page()) + doc.save(str(pdf)) + doc.close() + + assert parsers.pages_with_figures(str(pdf), max_pages = 4) == [1, 2] + assert parsers.pages_with_figures(str(pdf), max_pages = 4, exclude_pages = {1}) == [2] + assert parsers.pages_with_figures(str(pdf), max_pages = 4, exclude_pages = {2}) == [1] + + +def test_run_skips_figure_work_without_vision_model( + rag_conn, stub_embeddings, monkeypatch, tmp_path +): + # No vision model -> the whole figure pass (detection + rasterization) is skipped. + from core.rag import parsers + + monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", True) + monkeypatch.setattr(captioner, "vision_endpoint", lambda: None) + touched: list[str] = [] + monkeypatch.setattr( + parsers, "pages_with_figures", lambda *a, **k: touched.append("detect") or [] + ) + monkeypatch.setattr( + parsers, "render_pdf_figure_tiles", lambda *a, **k: touched.append("render") or [] + ) + + pdf = tmp_path / "fig.pdf" + _figure_pdf(pdf) + _ingest_with_caption(rag_conn, "t1", pdf, None) # follow config (ON), but no model + assert touched == [] # neither figure detection nor tiling ran + + +def test_vision_complete_sends_auth_header(monkeypatch): + # Direct-stream serves llama-server with --api-key; vision calls must send the bearer. + import httpx + + monkeypatch.setattr( + captioner, "_vision_auth_headers", lambda: {"Authorization": "Bearer secret"} + ) + captured: dict = {} + + class _Resp: + def raise_for_status(self): + pass + + def json(self): + return {"choices": [{"message": {"content": "ok"}}]} + + def fake_post(url, *, json, timeout, headers): + captured.update(url = url, headers = headers) + return _Resp() + + monkeypatch.setattr(httpx, "post", fake_post) + out = captioner._vision_complete( + "http://x", "local", b"img", prompt = "p", timeout = 5.0, max_tokens = 8 + ) + assert out == "ok" + assert captured["headers"] == {"Authorization": "Bearer secret"} + + +def test_vision_complete_omits_header_when_unauthenticated(monkeypatch): + # No api-key configured -> no spurious Authorization header on plain llama-server. + import httpx + + monkeypatch.setattr(captioner, "_vision_auth_headers", lambda: None) + captured: dict = {} + + class _Resp: + def raise_for_status(self): + pass + + def json(self): + return {"choices": [{"message": {"content": "ok"}}]} + + def fake_post(url, *, json, timeout, headers): + captured["headers"] = headers + return _Resp() + + monkeypatch.setattr(httpx, "post", fake_post) + captioner._vision_complete("http://x", "local", b"i", prompt = "p", timeout = 5.0, max_tokens = 8) + assert captured["headers"] is None + + +def test_merge_page_captions_dedups(): + out = captioner.merge_page_captions({1: ["MatMul\nScale", "Scale\nSoftMax"]}) + text = out[1][0] + assert text.lower().count("scale") == 1 # repeated label from overlapping tiles dropped + assert "MatMul" in text and "SoftMax" in text + + def test_splice_captions_appends_to_right_page(): pages = [Page("body one", 1, 8), Page("body two", 2, 8)] out = captioner.splice_captions(pages, {2: ["a diagram of X"]}) @@ -55,29 +228,6 @@ def test_splice_captions_noop_when_empty(): assert captioner.splice_captions(pages, {}) is pages -def test_render_pdf_figures_detects_drawing(tmp_path): - import pymupdf - - from core.rag.parsers import render_pdf_figures - - pdf = tmp_path / "fig.pdf" - doc = pymupdf.open() - page = doc.new_page() - shape = page.new_shape() - shape.draw_rect(pymupdf.Rect(60, 60, 540, 460)) - for i in range(8): - shape.draw_line((80, 80 + i * 40), (520, 80 + i * 40)) - shape.finish(color = (0, 0, 0), fill = (0.8, 0.8, 0.9)) - shape.commit() - doc.save(str(pdf)) - doc.close() - - figs = render_pdf_figures(str(pdf)) - assert figs, "expected at least one rendered figure region" - assert figs[0].image_bytes[:8] == b"\x89PNG\r\n\x1a\n" - assert figs[0].page_number == 1 - - def test_captioned_text_is_searchable(rag_home, stub_embeddings, monkeypatch): from core.rag import retrieval, store from storage import rag_db @@ -103,3 +253,100 @@ def test_captioned_text_is_searchable(rag_home, stub_embeddings, monkeypatch): finally: conn.close() assert hits, "spliced caption text should be retrievable via lexical search" + + +# ── per-upload caption override (parallels test_rag_ocr_fallback.py) ── + + +def _figure_pdf(path): + """A born-digital PDF: a page with real text (so it is not treated as scanned) + plus a vector drawing region that figure detection picks up as a figure.""" + import pymupdf + + doc = pymupdf.open() + page = doc.new_page() + page.insert_textbox( + pymupdf.Rect(40, 40, 550, 120), + "Quarterly revenue report. The chart below shows the trend.", + fontsize = 11, + ) + shape = page.new_shape() + shape.draw_rect(pymupdf.Rect(60, 140, 540, 520)) + for i in range(8): + shape.draw_line((80, 160 + i * 40), (520, 160 + i * 40)) + shape.finish(color = (0, 0, 0), fill = (0.8, 0.8, 0.9)) + shape.commit() + doc.save(str(path)) + doc.close() + + +def _ingest_with_caption(rag_conn, thread_id, path, caption): + from core.rag import ingestion, store + + scope = store.thread_scope(thread_id) + document_id = store.create_document( + rag_conn, + scope = scope, + filename = "fig.pdf", + sha256 = str(path) + str(caption), + thread_id = thread_id, + status = "pending", + stored_path = str(path), + ) + job_id = ingestion._new_job(rag_conn, document_id, scope) + # _run(job_id, document_id, scope, stored_path, model_name, ocr, caption) + ingestion._run(job_id, document_id, scope, str(path), None, None, caption) + return store.get_document(rag_conn, document_id) + + +def test_caption_override_true_runs_when_config_off( + rag_conn, stub_embeddings, monkeypatch, tmp_path +): + # Config default OFF, but the per-upload toggle (caption=True) forces captioning. + from core.rag import tool + + monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", False) + monkeypatch.setattr(captioner, "vision_endpoint", lambda: ("http://x", "local")) + monkeypatch.setattr(captioner, "_caption_one", lambda *a: "bar chart of revenue wombat-7") + + pdf = tmp_path / "fig.pdf" + _figure_pdf(pdf) + _ingest_with_caption(rag_conn, "t1", pdf, True) + + text, _ = tool.whole_document_context(scope_thread_id = "t1", max_tokens = 6000) + assert "wombat-7" in text # the spliced figure caption reached the index + + +def test_caption_override_false_skips_when_config_on( + rag_conn, stub_embeddings, monkeypatch, tmp_path +): + # Config default ON, but the per-upload toggle (caption=False) skips captioning. + monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", True) + monkeypatch.setattr(captioner, "vision_endpoint", lambda: ("http://x", "local")) + called = [] + monkeypatch.setattr(captioner, "_caption_one", lambda *a: called.append(1) or "should not run") + + pdf = tmp_path / "fig.pdf" + _figure_pdf(pdf) + _ingest_with_caption(rag_conn, "t1", pdf, False) + + assert called == [] # no vision caption calls despite config ON + + +def test_caption_none_follows_config(rag_conn, stub_embeddings, monkeypatch, tmp_path): + # Omitted override (None) falls back to config.CAPTION_IMAGES. + monkeypatch.setattr(captioner, "vision_endpoint", lambda: ("http://x", "local")) + seen = [] + monkeypatch.setattr(captioner, "_caption_one", lambda *a: seen.append(1) or "chart caption") + + monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", False) + pdf_off = tmp_path / "off.pdf" + _figure_pdf(pdf_off) + _ingest_with_caption(rag_conn, "t1", pdf_off, None) + assert seen == [] # config OFF + no override -> no captioning + + monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", True) + pdf_on = tmp_path / "on.pdf" + _figure_pdf(pdf_on) + _ingest_with_caption(rag_conn, "t2", pdf_on, None) + assert seen # config ON + no override -> captioning runs diff --git a/studio/backend/tests/test_rag_ingestion.py b/studio/backend/tests/test_rag_ingestion.py index f0b71bc23b..7e9e803687 100644 --- a/studio/backend/tests/test_rag_ingestion.py +++ b/studio/backend/tests/test_rag_ingestion.py @@ -83,6 +83,34 @@ def test_ingestion_dedupe_by_hash(rag_home, stub_embeddings, tmp_path): conn.close() +def test_ingestion_reingests_when_existing_has_zero_chunks(rag_home, stub_embeddings, tmp_path): + # A prior ingest of identical bytes that yielded no chunks (e.g. a scanned PDF + # before a vision model loaded) must re-ingest, not dedupe to the empty record. + path = _write(tmp_path, "doc.txt", "alpha bravo charlie " * 50) + sha = ingestion._sha256_file(path) + scope = store.kb_scope("K1") + conn = rag_db.get_connection() + try: + empty_id = store.create_document(conn, scope = scope, filename = "old.txt", sha256 = sha) + store.set_document_status(conn, empty_id, "completed", num_chunks = 0) + finally: + conn.close() + + doc_id, job_id = ingestion.start_ingestion(scope, "K1", None, "doc.txt", path) + events = _drain(job_id) + _wait_completed(job_id) + + assert not any(e.get("deduped") for e in events) # not a dedupe -> real ingest + assert doc_id != empty_id + conn = rag_db.get_connection() + try: + docs = store.list_documents(conn, scope) + assert len(docs) == 1 # the empty record was removed, replaced by the new one + assert docs[0]["num_chunks"] > 0 + finally: + conn.close() + + def test_ingestion_dedupe_removes_duplicate_upload(rag_home, stub_embeddings): from utils.paths import ensure_dir, rag_uploads_root @@ -210,6 +238,41 @@ def test_delete_document_route_removes_stored_upload(rag_home): conn.close() +def test_get_job_status_includes_num_chunks(rag_home, stub_embeddings, tmp_path): + # The poll/reconcile path reads num_chunks from get_job_status (the SSE complete + # frame carries it, but a client that falls back to polling needs it here too). + path = _write(tmp_path, "doc.txt", "alpha bravo charlie " * 50) + scope = store.kb_scope("K1") + _doc_id, job_id = ingestion.start_ingestion(scope, "K1", None, "doc.txt", path) + _drain(job_id) + _wait_completed(job_id) + status = ingestion.get_job_status(job_id) + assert status["status"] == "completed" + assert status["num_chunks"] and status["num_chunks"] > 0 + + +def test_save_upload_rejects_oversize_file(rag_home, monkeypatch): + # A file over the cap is rejected (413) and its partial bytes are cleaned up. + import io + + from fastapi import HTTPException + + from core.rag import config + from routes import rag as rag_routes + from utils.paths import rag_uploads_root + + monkeypatch.setattr(config, "MAX_UPLOAD_BYTES", 1024) + + class _Up: + filename = "big.txt" + file = io.BytesIO(b"x" * 4096) + + with pytest.raises(HTTPException) as ei: + rag_routes._save_upload(_Up()) + assert ei.value.status_code == 413 + assert list(rag_uploads_root().glob("*.txt")) == [] # partial upload removed + + def test_ingestion_delete_removes_all_rows(rag_home, stub_embeddings, tmp_path): path = _write(tmp_path, "doc.txt", "alpha bravo charlie delta") scope = store.kb_scope("K1") diff --git a/studio/backend/tests/test_rag_ocr_fallback.py b/studio/backend/tests/test_rag_ocr_fallback.py new file mode 100644 index 0000000000..c7be1fe60b --- /dev/null +++ b/studio/backend/tests/test_rag_ocr_fallback.py @@ -0,0 +1,259 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Scanned-PDF OCR fallback: a PDF page with no text layer is rendered and transcribed +by the vision model during ingestion, so image-only PDFs become searchable. The vision +call is stubbed, so no model is needed.""" + +import pymupdf + +from core.rag import captioner, ingestion, parsers, store, tool + + +def _image_only_pdf(path, *, pages = 1): + """A PDF whose pages carry only a raster image, so get_text returns ''.""" + doc = pymupdf.open() + pix = pymupdf.Pixmap(pymupdf.csRGB, pymupdf.IRect(0, 0, 120, 120)) + pix.clear_with(220) + for _ in range(pages): + page = doc.new_page() + page.insert_image(page.rect, pixmap = pix) + doc.save(str(path)) + doc.close() + + +def _text_pdf(path, body): + doc = pymupdf.open() + page = doc.new_page() + page.insert_textbox(pymupdf.Rect(40, 40, 550, 800), body, fontsize = 11) + doc.save(str(path)) + doc.close() + + +def _ingest(rag_conn, thread_id, filename, path): + """Drive the real ingestion worker synchronously and return the document row.""" + scope = store.thread_scope(thread_id) + document_id = store.create_document( + rag_conn, + scope = scope, + filename = filename, + sha256 = filename, + thread_id = thread_id, + status = "pending", + stored_path = str(path), + ) + job_id = ingestion._new_job(rag_conn, document_id, scope) + ingestion._run(job_id, document_id, scope, str(path), None) + return store.get_document(rag_conn, document_id) + + +# ── parsers.render_pdf_pages ───────────────────────────────────────── + + +def test_render_pdf_pages_returns_png_per_page(tmp_path): + pdf = tmp_path / "two.pdf" + _image_only_pdf(pdf, pages = 2) + out = parsers.render_pdf_pages(str(pdf), [1, 2], dpi = 72) + assert set(out) == {1, 2} + assert all(b.startswith(b"\x89PNG") for b in out.values()) + + +def test_render_pdf_pages_excludes_unwanted(tmp_path): + pdf = tmp_path / "three.pdf" + _image_only_pdf(pdf, pages = 3) + out = parsers.render_pdf_pages(str(pdf), [2], dpi = 72) + assert set(out) == {2} + + +def test_render_pdf_pages_empty_request(tmp_path): + pdf = tmp_path / "one.pdf" + _image_only_pdf(pdf, pages = 1) + assert parsers.render_pdf_pages(str(pdf), [], dpi = 72) == {} + + +# ── captioner.ocr_pages gating ─────────────────────────────────────── + + +def test_ocr_pages_no_endpoint(monkeypatch): + monkeypatch.setattr(captioner, "vision_endpoint", lambda: None) + assert captioner.ocr_pages({1: b"x"}) == {} + + +def test_collapse_runaway_caps_repeated_lines(): + # A looping model repeats a line hundreds of times; the guard caps it, keeps repeats. + text = "\n".join(["TITLE"] * 200 + ["body"] + ["Add & Norm"] * 3) + out = captioner._collapse_runaway(text) + lines = out.splitlines() + assert lines.count("TITLE") == 3 # 200 -> 3 + assert lines.count("Add & Norm") == 3 # legitimate triple survives + assert "body" in lines + + +def test_collapse_runaway_caps_interleaved_repeats(): + # Models also loop non-consecutively; the global per-line cap bounds those too. + text = "\n".join(["Llion Vaswani Google", "Niki Parmar Google"] * 40) + out = captioner._collapse_runaway(text) + lines = [ln for ln in out.splitlines() if ln.strip()] + assert lines.count("Llion Vaswani Google") <= 8 + assert lines.count("Niki Parmar Google") <= 8 + + +def test_collapse_runaway_noop_on_normal_text(): + text = "Heading\n\nFirst paragraph.\nSecond paragraph.\n\nFooter" + assert captioner._collapse_runaway(text) == text + + +def test_ocr_pages_applies_runaway_guard(monkeypatch): + monkeypatch.setattr(captioner.config, "OCR_SCANNED", True) + monkeypatch.setattr(captioner, "_ocr_one", lambda *a: "\n".join(["X"] * 50)) + out = captioner.ocr_pages({1: b"img"}, endpoint = ("http://x", "local")) + assert out[1].splitlines().count("X") == 3 # guard applied to stored text + + +def test_ocr_pages_transcribes_and_caps(monkeypatch): + monkeypatch.setattr(captioner.config, "OCR_SCANNED", True) + monkeypatch.setattr(captioner.config, "OCR_MAX_PAGES", 1) + calls = [] + monkeypatch.setattr( + captioner, + "_ocr_one", + lambda base, model, b, t: (calls.append(1) or "transcribed text"), + ) + out = captioner.ocr_pages({1: b"a", 2: b"b"}, endpoint = ("http://x", "local")) + assert out == {1: "transcribed text"} # page 2 dropped by the cap + assert len(calls) == 1 + + +def test_ocr_scanned_pages_merges_short_text_layer(rag_conn, monkeypatch): + # Near-empty pages can still have meaningful extractable text; OCR augments it + # rather than replacing it with a fallible vision transcription. + scope = store.thread_scope("t1") + document_id = store.create_document(rag_conn, scope = scope, filename = "scan.pdf", sha256 = "h") + job_id = ingestion._new_job(rag_conn, document_id, scope) + pages = [parsers.Page("ID-42", 1, 5)] + + monkeypatch.setattr(captioner.config, "OCR_SCANNED", True) + monkeypatch.setattr(captioner.config, "OCR_MIN_CHARS", 16) + monkeypatch.setattr(captioner, "vision_endpoint", lambda: ("http://x", "local")) + monkeypatch.setattr(parsers, "render_pdf_pages", lambda *a, **k: {1: b"png"}) + monkeypatch.setattr(captioner, "ocr_pages", lambda page_pngs: {1: "OCR body text"}) + + out, ocred = ingestion._ocr_scanned_pages(pages, "scan.pdf", rag_conn, job_id) + assert ocred == {1} + assert out[0].text == "ID-42\n\nOCR body text" + + +# ── end-to-end ingestion ───────────────────────────────────────────── + + +def test_scanned_pdf_is_ocred_into_chunks(rag_conn, stub_embeddings, monkeypatch, tmp_path): + monkeypatch.setattr(captioner.config, "OCR_SCANNED", True) + monkeypatch.setattr(captioner, "vision_endpoint", lambda: ("http://x", "local")) + monkeypatch.setattr( + captioner, "_ocr_one", lambda base, model, b, t: "Invoice total is zebra-42 due Friday" + ) + + pdf = tmp_path / "scan.pdf" + _image_only_pdf(pdf, pages = 1) + doc = _ingest(rag_conn, "t1", "scan.pdf", pdf) + + assert doc["status"] == "completed" + assert doc["num_chunks"] >= 1 + # The OCR'd text is now indexed and reaches whole-document injection. + text, _sources = tool.whole_document_context(scope_thread_id = "t1", max_tokens = 6000) + assert "zebra-42" in text + + +def test_scanned_page_past_ocr_cap_is_still_captioned( + rag_conn, stub_embeddings, monkeypatch, tmp_path +): + # OCR is capped to one page, so page 2 is scanned but never transcribed. Figure + # captioning must still cover it (we exclude only the pages OCR actually handled), + # so a chart on an un-OCR'd scanned page is not silently dropped. + monkeypatch.setattr(captioner.config, "OCR_SCANNED", True) + monkeypatch.setattr(captioner.config, "OCR_MAX_PAGES", 1) + monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", True) + monkeypatch.setattr(captioner, "vision_endpoint", lambda: ("http://x", "local")) + monkeypatch.setattr(captioner, "_ocr_one", lambda *a: "scanned page alpha") + monkeypatch.setattr(captioner, "_caption_one", lambda *a: "figure caption bravo") + + pdf = tmp_path / "scan2.pdf" + _image_only_pdf(pdf, pages = 2) + doc = _ingest(rag_conn, "t1", "scan2.pdf", pdf) + + assert doc["status"] == "completed" + text, _ = tool.whole_document_context(scope_thread_id = "t1", max_tokens = 6000) + assert "scanned page alpha" in text # page 1 OCR'd, within the cap + assert "figure caption bravo" in text # page 2 past the cap -> captioned, not dropped + + +def test_born_digital_pdf_skips_ocr(rag_conn, stub_embeddings, monkeypatch, tmp_path): + called = [] + monkeypatch.setattr(captioner.config, "OCR_SCANNED", True) + monkeypatch.setattr(captioner, "_ocr_one", lambda *a: called.append(1) or "should not run") + + pdf = tmp_path / "digital.pdf" + _text_pdf(pdf, "Real born digital body text. " * 30 + "marker-quokka") + doc = _ingest(rag_conn, "t1", "digital.pdf", pdf) + + assert doc["status"] == "completed" + assert called == [] # page had real text -> never considered scanned + text, _sources = tool.whole_document_context(scope_thread_id = "t1", max_tokens = 6000) + assert "marker-quokka" in text + + +def _ingest_with_ocr(rag_conn, thread_id, path, ocr): + scope = store.thread_scope(thread_id) + document_id = store.create_document( + rag_conn, + scope = scope, + filename = "scan.pdf", + sha256 = str(path) + str(ocr), + thread_id = thread_id, + status = "pending", + stored_path = str(path), + ) + job_id = ingestion._new_job(rag_conn, document_id, scope) + ingestion._run(job_id, document_id, scope, str(path), None, ocr = ocr) + return store.get_document(rag_conn, document_id) + + +def test_ocr_override_false_skips_ocr_when_config_on( + rag_conn, stub_embeddings, monkeypatch, tmp_path +): + # Config default ON, but the per-upload toggle (ocr=False) skips OCR. + monkeypatch.setattr(captioner.config, "OCR_SCANNED", True) + monkeypatch.setattr(captioner, "vision_endpoint", lambda: ("http://x", "local")) + monkeypatch.setattr(captioner, "_ocr_one", lambda *a: "should not run") + pdf = tmp_path / "scan.pdf" + _image_only_pdf(pdf, pages = 1) + doc = _ingest_with_ocr(rag_conn, "t1", pdf, ocr = False) + assert doc["num_chunks"] == 0 # scanned page left empty + + +def test_ocr_override_true_runs_ocr_when_config_off( + rag_conn, stub_embeddings, monkeypatch, tmp_path +): + # Config default OFF, but the per-upload toggle (ocr=True) forces OCR on. + monkeypatch.setattr(captioner.config, "OCR_SCANNED", False) + monkeypatch.setattr(captioner, "vision_endpoint", lambda: ("http://x", "local")) + monkeypatch.setattr(captioner, "_ocr_one", lambda *a: "forced ocr text quokka") + pdf = tmp_path / "scan.pdf" + _image_only_pdf(pdf, pages = 1) + doc = _ingest_with_ocr(rag_conn, "t1", pdf, ocr = True) + assert doc["num_chunks"] >= 1 + text, _ = tool.whole_document_context(scope_thread_id = "t1", max_tokens = 6000) + assert "quokka" in text + + +def test_ocr_disabled_leaves_scanned_pdf_empty(rag_conn, stub_embeddings, monkeypatch, tmp_path): + monkeypatch.setattr(captioner.config, "OCR_SCANNED", False) + + pdf = tmp_path / "scan.pdf" + _image_only_pdf(pdf, pages = 1) + doc = _ingest(rag_conn, "t1", "scan.pdf", pdf) + + # With OCR off, a text-less scanned page yields no chunks (prior behavior). + assert doc["status"] == "completed" + assert doc["num_chunks"] == 0 + assert tool.whole_document_context(scope_thread_id = "t1", max_tokens = 6000) is None diff --git a/studio/backend/tests/test_rag_parsing.py b/studio/backend/tests/test_rag_parsing.py new file mode 100644 index 0000000000..4c46f49495 --- /dev/null +++ b/studio/backend/tests/test_rag_parsing.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""PDF text extraction: layout-aware Markdown (pymupdf4llm) with plain-text fallback.""" + +from __future__ import annotations + +import pytest + +pytest.importorskip("pymupdf") + + +def _table_pdf(path): + import pymupdf + + doc = pymupdf.open() + page = doc.new_page() + page.insert_textbox(pymupdf.Rect(40, 40, 550, 70), "Quarterly Results", fontsize = 16) + rows = [("Quarter", "Revenue", "Growth"), ("Q1", "$1.2M", "12%"), ("Q2", "$1.5M", "25%")] + y = 90 + for r in rows: + page.insert_textbox(pymupdf.Rect(40, y, 250, y + 20), r[0], fontsize = 11) + page.insert_textbox(pymupdf.Rect(250, y, 400, y + 20), r[1], fontsize = 11) + page.insert_textbox(pymupdf.Rect(400, y, 540, y + 20), r[2], fontsize = 11) + y += 24 + doc.save(str(path)) + doc.close() + + +def test_pdf_extracts_markdown_table(tmp_path, monkeypatch): + # With Markdown on, the layout is emitted as Markdown markup (heading, and a pipe table + # where the extractor detects one) that flat get_text never produces. + pytest.importorskip("pymupdf4llm") + from core.rag import config, parsers + + monkeypatch.setattr(config, "PDF_MARKDOWN", True) + pdf = tmp_path / "table.pdf" + _table_pdf(pdf) + text = "\n".join(p.text for p in parsers.parse(str(pdf))) + assert "Q2" in text and "$1.5M" in text # cell values preserved + assert "#" in text or "|" in text # Markdown markup (heading or table pipes) + + +def test_pdf_markdown_off_uses_plain_text(tmp_path, monkeypatch): + # The toggle (RAG_PDF_MARKDOWN=0) falls back to flat PyMuPDF text: content is still + # there, but with no Markdown markup. + from core.rag import config, parsers + + monkeypatch.setattr(config, "PDF_MARKDOWN", False) + pdf = tmp_path / "table.pdf" + _table_pdf(pdf) + text = "\n".join(p.text for p in parsers.parse(str(pdf))) + assert "Q2" in text and "$1.5M" in text + assert "#" not in text and "|" not in text # plain text path emits no Markdown markup + + +def test_pdf_markdown_passes_only_supported_legacy_kwargs(monkeypatch): + # The pinned PyMuPDF4LLM legacy path ignores unknown kwargs; do not pass the + # newer layout-only OCR knobs or Markdown extraction silently loses policy control. + from core.rag import parsers + + captured = {} + + class _FakePymupdf4llm: + @staticmethod + def to_markdown(doc, **kwargs): + captured.update(kwargs) + return [{"text": "plain markdown"}] + + class _Doc: + page_count = 1 + + monkeypatch.setitem(__import__("sys").modules, "pymupdf4llm", _FakePymupdf4llm) + assert parsers._pdf_markdown(_Doc()) == ["plain markdown"] + assert captured == {"page_chunks": True, "show_progress": False} + + +def test_pdf_markdown_falls_back_when_lib_missing(tmp_path, monkeypatch): + # If pymupdf4llm extraction returns None (missing/failed), parsing still yields the + # plain-text pages rather than raising. + from core.rag import config, parsers + + monkeypatch.setattr(config, "PDF_MARKDOWN", True) + monkeypatch.setattr(parsers, "_pdf_markdown", lambda doc: None) + pdf = tmp_path / "table.pdf" + _table_pdf(pdf) + pages = parsers.parse(str(pdf)) + assert pages and "Quarter" in pages[0].text diff --git a/studio/backend/tests/test_rag_preview.py b/studio/backend/tests/test_rag_preview.py index e7f2a39792..0ff27897bd 100644 --- a/studio/backend/tests/test_rag_preview.py +++ b/studio/backend/tests/test_rag_preview.py @@ -165,6 +165,25 @@ def test_locator_handles_midword_anchor_and_locates_line(): assert r["width"] > 0 and r["height"] > 0 +def test_locator_anchors_through_markdown_table_pipes(): + # Markdown table cells are pipe-joined with no spaces; the locator splits on pipes + # so a table-row chunk still anchors to the raw PDF word stream. + import pymupdf + + from core.rag.locators import LocatorMatch, _regions_for_match + + doc = pymupdf.open() + page = doc.new_page() + page.insert_text((72, 200), "Quarter Revenue Growth Q1 sales strong here", fontsize = 12) + # What the Markdown parser stores for the row (cells joined by pipes, no spaces). + page_text = "|Quarter|Revenue|Growth|Q1|sales|strong|here|" + match = LocatorMatch(page_index = 0, page_number = 1, start = 0, end = len(page_text)) + rects = _regions_for_match(doc, page_text, match) + doc.close() + + assert rects, "a Markdown table row should still anchor to the page words" + + def test_sign_verify_roundtrip(rag_home): from routes import rag as rag_routes diff --git a/studio/backend/tests/test_rag_whole_document.py b/studio/backend/tests/test_rag_whole_document.py new file mode 100644 index 0000000000..545d731fd2 --- /dev/null +++ b/studio/backend/tests/test_rag_whole_document.py @@ -0,0 +1,520 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Whole-document context mode: a thread-attached file small enough to fit is +injected in full (every chunk, in order) instead of top-K retrieval. Covers the +new store query, the tool-level renderer, and the auto-inject wiring + fallback. +No embedder is needed - the whole-doc path does no query embedding.""" + +import json + +from core.rag import store, tool +from core.rag.chunking import Chunk +from core.inference import tools as inf_tools + +# A vector per chunk just to satisfy add_chunks (the whole-doc path never reads +# vectors); dimension is arbitrary but must be consistent within a connection. +_VEC = [0.1, 0.2, 0.3, 0.4] + + +def _chunk( + text, + index = 0, + page = None, + tokens = None, +): + return Chunk( + text = text, + token_count = tokens if tokens is not None else len(text.split()), + page_number = page, + source_page_index = 0, + chunk_index = index, + page_char_start = 0, + page_char_end = len(text), + ) + + +def _add_doc( + conn, + scope, + doc_id, + filename, + sha, + texts, + *, + status = "completed", + tokens = None, + pages = None, +): + chunks = [ + _chunk( + t, + i, + page = (pages[i] if pages else None), + tokens = (tokens[i] if tokens else None), + ) + for i, t in enumerate(texts) + ] + vectors = [list(_VEC) for _ in texts] + store.create_document(conn, scope = scope, filename = filename, sha256 = sha, document_id = doc_id) + store.add_chunks(conn, scope, doc_id, chunks, vectors) + store.set_document_status(conn, doc_id, status, num_chunks = len(texts)) + + +def _injected_text(result) -> str: + """The text spliced into the conversation as the synthetic tool result.""" + tool_msg = next(m for m in result["messages"] if m.get("role") == "tool") + return tool_msg["content"] + + +# ── store.all_chunks_for_scope ─────────────────────────────────────── + + +def test_all_chunks_for_scope_orders_by_document_then_index(rag_conn): + scope = store.thread_scope("t1") + _add_doc(rag_conn, scope, "d1", "first.pdf", "h1", ["a", "b", "c"]) + _add_doc(rag_conn, scope, "d2", "second.pdf", "h2", ["x", "y"]) + rows = store.all_chunks_for_scope(rag_conn, scope) + assert [r["id"] for r in rows] == ["d1:0", "d1:1", "d1:2", "d2:0", "d2:1"] + assert rows[0]["filename"] == "first.pdf" + assert rows[-1]["filename"] == "second.pdf" + assert rows[0]["text"] == "a" + + +def test_all_chunks_for_scope_excludes_non_completed(rag_conn): + scope = store.thread_scope("t1") + _add_doc(rag_conn, scope, "done", "done.pdf", "h1", ["ready"]) + _add_doc(rag_conn, scope, "pend", "pend.pdf", "h2", ["indexing"], status = "pending") + rows = store.all_chunks_for_scope(rag_conn, scope) + assert [r["id"] for r in rows] == ["done:0"] + + +def test_all_chunks_for_scope_empty_scope(rag_conn): + assert store.all_chunks_for_scope(rag_conn, store.thread_scope("nope")) == [] + + +def test_all_chunks_for_scope_isolates_scopes(rag_conn): + _add_doc(rag_conn, store.thread_scope("t1"), "d1", "f", "h1", ["mine"]) + _add_doc(rag_conn, store.thread_scope("t2"), "d2", "f", "h2", ["theirs"]) + rows = store.all_chunks_for_scope(rag_conn, store.thread_scope("t1")) + assert [r["text"] for r in rows] == ["mine"] + + +# ── store.scope_token_estimate (cheap whole-doc budget pre-check) ───── + + +def test_scope_token_estimate_sums_without_hydrating(rag_conn): + # Stored counts sum directly; zero/missing falls back to length/4; non-completed out. + scope = store.thread_scope("t1") + _add_doc(rag_conn, scope, "d1", "a.pdf", "h1", ["alpha", "bravo"], tokens = [10, 20]) + # token_count 0 -> length/4 fallback: a 40-char chunk estimates to 10 tokens. + _add_doc(rag_conn, scope, "d2", "b.pdf", "h2", ["x" * 40], tokens = [0]) + _add_doc(rag_conn, scope, "d3", "c.pdf", "h3", ["pending"], status = "pending", tokens = [99]) + assert store.scope_token_estimate(rag_conn, scope) == 10 + 20 + 10 + assert store.scope_token_estimate(rag_conn, store.thread_scope("none")) == 0 + + +def test_scope_token_estimate_matches_row_sum(rag_conn): + # Must agree with the exact per-row sum it short-circuits (one stored count, one + # length/4 fallback), so the pre-check never disagrees with the full path. + from core.rag.tool import _row_token_count + + scope = store.thread_scope("t1") + _add_doc( + rag_conn, scope, "d1", "a.pdf", "h1", ["a long-ish chunk body here", "tail"], tokens = [0, 5] + ) + rows = store.all_chunks_for_scope(rag_conn, scope) + assert store.scope_token_estimate(rag_conn, scope) == sum(_row_token_count(r) for r in rows) + + +# ── tool.whole_document_context ────────────────────────────────────── + + +def test_whole_document_context_returns_full_text_and_sources(rag_conn): + scope = store.thread_scope("t1") + _add_doc( + rag_conn, + scope, + "d1", + "report.pdf", + "h1", + ["chapter one body", "chapter two body"], + pages = [1, 2], + ) + result = tool.whole_document_context(scope_thread_id = "t1", max_tokens = 6000) + assert result is not None + text, sources = result + # Every chunk is present, in order, as blocks. + assert "chapter one body" in text + assert "chapter two body" in text + assert ' None (whole-doc is thread-attachment only). + assert tool.whole_document_context(max_tokens = 6000) is None + + +def test_whole_document_context_null_token_count_enforces_budget(rag_conn): + # A missing token_count must not bypass the budget; fall back to a length estimate. + big = "word " * 20_000 # ~20k tokens by length estimate + _add_doc(rag_conn, store.thread_scope("t1"), "d1", "big.pdf", "h1", [big], tokens = [None]) + assert tool.whole_document_context(scope_thread_id = "t1", max_tokens = 6000) is None + assert tool.whole_document_context(scope_thread_id = "t1", max_tokens = 1_000_000) is not None + + +def test_whole_document_context_spans_multiple_docs(rag_conn): + scope = store.thread_scope("t1") + _add_doc(rag_conn, scope, "d1", "a.pdf", "h1", ["alpha text"]) + _add_doc(rag_conn, scope, "d2", "b.pdf", "h2", ["bravo text"]) + text, sources = tool.whole_document_context(scope_thread_id = "t1", max_tokens = 6000) + assert "alpha text" in text and "bravo text" in text + assert {s["filename"] for s in sources} == {"a.pdf", "b.pdf"} + + +# ── build_rag_autoinject wiring ────────────────────────────────────── + + +def _convo(text = "summarize the whole document"): + return [{"role": "user", "content": text}] + + +def test_build_rag_autoinject_uses_whole_doc(rag_conn): + scope = store.thread_scope("t1") + _add_doc(rag_conn, scope, "d1", "doc.pdf", "h1", ["whole alpha part", "whole bravo part"]) + result = inf_tools.build_rag_autoinject(_convo(), {"thread_id": "t1"}) + assert result is not None + injected = _injected_text(result) + # Both chunks present -> the model receives the entire file, not top-K. + assert "whole alpha part" in injected + assert "whole bravo part" in injected + # Tool-message content is chunk text only; the citation JSON tail is internal. + assert inf_tools.RAG_SOURCES_SENTINEL not in injected + + +def test_build_rag_autoinject_whole_doc_runs_when_autoinject_false(rag_conn, monkeypatch): + # Large-model Auto sets autoinject=False, but whole-doc is a separate thread-doc + # context mode and should still inject a fitting attachment. + _add_doc(rag_conn, store.thread_scope("t1"), "d1", "doc.pdf", "h1", ["entire file body"]) + monkeypatch.setattr( + tool, + "search_for_autoinject", + lambda **kw: (_ for _ in ()).throw(AssertionError("retrieval should not run")), + ) + result = inf_tools.build_rag_autoinject(_convo(), {"thread_id": "t1", "autoinject": False}) + assert result is not None + assert "entire file body" in _injected_text(result) + + +def test_build_rag_autoinject_explicit_off_disables_whole_doc(rag_conn, monkeypatch): + # The UI Off switch sends both autoinject=False and whole_doc=False. + _add_doc(rag_conn, store.thread_scope("t1"), "d1", "doc.pdf", "h1", ["small body"]) + monkeypatch.setattr( + tool, + "search_for_autoinject", + lambda **kw: (_ for _ in ()).throw(AssertionError("retrieval should not run")), + ) + assert ( + inf_tools.build_rag_autoinject( + _convo(), {"thread_id": "t1", "autoinject": False, "whole_doc": False} + ) + is None + ) + + +def test_build_rag_autoinject_falls_back_over_budget(rag_conn, monkeypatch): + scope = store.thread_scope("t1") + _add_doc(rag_conn, scope, "d1", "big.pdf", "h1", ["overflow"], tokens = [50_000]) + + sentinel = ("TOPK_FALLBACK_TEXT", [{"citationId": 1, "filename": "big.pdf", "text": "x"}]) + monkeypatch.setattr(tool, "search_for_autoinject", lambda **kw: sentinel) + + result = inf_tools.build_rag_autoinject(_convo(), {"thread_id": "t1"}) + assert result is not None + assert _injected_text(result) == "TOPK_FALLBACK_TEXT" + + +def test_build_rag_autoinject_context_budget_falls_back(rag_conn, monkeypatch): + # Runtime context can be smaller than RAG_WHOLE_DOC_MAX_TOKENS; cap whole-doc to + # the active context and fall back to retrieval when it would overflow. + _add_doc( + rag_conn, store.thread_scope("t1"), "d1", "small.pdf", "h1", ["fits global"], tokens = [900] + ) + sentinel = ("TOPK_CONTEXT_FALLBACK", [{"citationId": 1, "filename": "small.pdf", "text": "x"}]) + monkeypatch.setattr(tool, "search_for_autoinject", lambda **kw: sentinel) + result = inf_tools.build_rag_autoinject( + _convo(), {"thread_id": "t1", "context_length": 1200, "whole_doc": True} + ) + assert result is not None + assert _injected_text(result) == "TOPK_CONTEXT_FALLBACK" + + +def test_whole_doc_budget_reserves_image_parts(monkeypatch): + from core.rag import config + + monkeypatch.setattr(config, "WHOLE_DOC_MAX_TOKENS", 10_000) + scope = {"context_length": 7000, "response_headroom": 1000} + text_only = [{"role": "user", "content": [{"type": "text", "text": "summarize"}]}] + with_image = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "summarize"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,abc"}}, + ], + } + ] + + assert ( + inf_tools._whole_doc_budget(scope, text_only) + - inf_tools._whole_doc_budget(scope, with_image) + == inf_tools._IMAGE_PART_TOKEN_ESTIMATE + ) + + +def test_build_rag_autoinject_server_kill_switch_blocks_whole_doc(rag_conn, monkeypatch): + # RAG_THREAD_WHOLE_DOC=0 stays authoritative; browser requests should not + # turn it back on by default. + from core.rag import config + + monkeypatch.setattr(config, "THREAD_WHOLE_DOC", False) + _add_doc(rag_conn, store.thread_scope("t1"), "d1", "doc.pdf", "h1", ["small body"]) + monkeypatch.setattr( + tool, + "search_for_autoinject", + lambda **kw: (_ for _ in ()).throw(AssertionError("retrieval should not run")), + ) + assert ( + inf_tools.build_rag_autoinject(_convo(), {"thread_id": "t1", "autoinject": False}) is None + ) + + +def test_whole_document_context_budgets_rendered_wrappers(rag_conn): + # Many tiny chunks add wrapper overhead beyond raw chunk token counts; budget + # the rendered prompt, not just stored text. + texts = ["x" for _ in range(120)] + _add_doc( + rag_conn, + store.thread_scope("t1"), + "d1", + "many-pages.pdf", + "h1", + texts, + tokens = [1 for _ in texts], + ) + assert tool.whole_document_context(scope_thread_id = "t1", max_tokens = 500) is None + + +def test_build_rag_autoinject_whole_doc_disabled_via_override(rag_conn, monkeypatch): + scope = store.thread_scope("t1") + _add_doc(rag_conn, scope, "d1", "doc.pdf", "h1", ["small body"]) + + sentinel = ("TOPK_TEXT", [{"citationId": 1, "filename": "doc.pdf", "text": "x"}]) + monkeypatch.setattr(tool, "search_for_autoinject", lambda **kw: sentinel) + + # whole_doc=False forces retrieval even though the doc fits. + result = inf_tools.build_rag_autoinject(_convo(), {"thread_id": "t1", "whole_doc": False}) + assert result is not None + assert _injected_text(result) == "TOPK_TEXT" + + +def test_build_rag_autoinject_kb_scope_never_whole_doc(rag_conn, monkeypatch): + # A KB-only scope (no thread) goes through retrieval, never whole-doc. + kb_scope = store.kb_scope("K1") + _add_doc(rag_conn, kb_scope, "d1", "kb.pdf", "h1", ["kb body one", "kb body two"]) + + sentinel = ("KB_RETRIEVAL_TEXT", [{"citationId": 1, "filename": "kb.pdf", "text": "x"}]) + monkeypatch.setattr(tool, "search_for_autoinject", lambda **kw: sentinel) + + result = inf_tools.build_rag_autoinject(_convo(), {"kb_id": "K1"}) + assert result is not None + assert _injected_text(result) == "KB_RETRIEVAL_TEXT" + + +def test_whole_document_context_thread_scope_only(rag_conn): + # A project corpus chunk is never whole-doc injected, even with a thread attachment. + _add_doc(rag_conn, store.thread_scope("t1"), "td", "thread.txt", "h1", ["thread attachment"]) + _add_doc(rag_conn, store.project_scope("p1"), "pd", "project.txt", "h2", ["project corpus"]) + text, sources = tool.whole_document_context(scope_thread_id = "t1", max_tokens = 6000) + assert "thread attachment" in text + assert "project corpus" not in text + assert {s["filename"] for s in sources} == {"thread.txt"} + + +def test_build_rag_autoinject_appends_project_retrieval(rag_conn, monkeypatch): + # Project chat: thread attachment whole-doc'd AND project sources retrieved, merged. + _add_doc( + rag_conn, + store.thread_scope("t1"), + "td", + "thread.txt", + "h1", + ["thread chunk one", "thread chunk two"], + ) + proj = ( + "PROJ", + [ + { + "citationId": 1, + "chunkId": "pj:0", + "documentId": "pj", + "filename": "project.txt", + "page": None, + "text": "project passage zeta", + "score": 0.91, + } + ], + ) + captured = {} + + def fake_search(**kw): + captured.update(kw) + return proj + + monkeypatch.setattr(tool, "search_for_autoinject", fake_search) + result = inf_tools.build_rag_autoinject(_convo(), {"thread_id": "t1", "project_id": "p1"}) + injected = _injected_text(result) + # Whole thread attachment AND the project passage are both injected. + assert "thread chunk one" in injected + assert "thread chunk two" in injected + assert "project passage zeta" in injected + # The companion retrieval was scoped to the project only (not thread or KB). + assert captured.get("scope_project_id") == "p1" + assert captured.get("scope_thread_id") is None + assert captured.get("scope_kb_id") is None + # Citation ids are sequential across the merged set: thread 1,2 then project 3. + assert ' whole-doc injection ──────── + + +def test_real_ingestion_feeds_whole_document(rag_conn, stub_embeddings, tmp_path): + """Drive the real ingestion worker on a multi-paragraph file, then confirm whole-doc + injection splices the entire document, not just retrieved chunks.""" + from core.rag import ingestion + + scope = store.thread_scope("t1") + body = ( + "# Quarterly Report\n\n" + + ("Revenue rose across every region this period. " * 40) + + "\n\nThe unique closing marker is xyzzy-sentinel for the final page. " * 40 + ) + src = tmp_path / "report.md" + src.write_text(body, encoding = "utf-8") + + document_id = store.create_document( + rag_conn, + scope = scope, + filename = "report.md", + sha256 = "sha-e2e", + thread_id = "t1", + status = "pending", + stored_path = str(src), + ) + job_id = ingestion._new_job(rag_conn, document_id, scope) + ingestion._run(job_id, document_id, scope, str(src), None) + + doc = store.get_document(rag_conn, document_id) + assert doc["status"] == "completed" + assert doc["num_chunks"] >= 2 # the doc chunked into multiple pieces + + result = inf_tools.build_rag_autoinject(_convo(), {"thread_id": "t1"}) + assert result is not None + injected = _injected_text(result) + # Opening and ending both present -> the whole file reached the model. + assert "Revenue rose" in injected + assert "xyzzy-sentinel" in injected + # Every stored chunk is represented as a numbered block. + assert injected.count(" void; setRagAutoInject: (value: RagAutoInject) => void; setRagAutoInjectMinScore: (score: number) => void; + setRagOcrScanned: (enabled: boolean) => void; + setRagCaptionFigures: (enabled: boolean) => void; setToolStatus: (status: string | null) => void; setGeneratingStatus: (status: string | null) => void; setActiveDiffusionCanvas: (canvas: DiffusionCanvasFrame | null) => void; @@ -1077,6 +1091,8 @@ export const useChatRuntimeStore = create((set, get) => ({ DEFAULT_RAG_AUTOINJECT_MIN_SCORE, { min: 0, max: 1 }, ), + ragOcrScanned: loadBool(CHAT_RAG_OCR_KEY, DEFAULT_RAG_OCR), + ragCaptionFigures: loadBool(CHAT_RAG_CAPTION_KEY, DEFAULT_RAG_CAPTION), toolStatus: null, generatingStatus: null, activeDiffusionCanvas: null, @@ -1498,6 +1514,16 @@ export const useChatRuntimeStore = create((set, get) => ({ ); return { ragAutoInjectMinScore }; }), + setRagOcrScanned: (ragOcrScanned) => + set(() => { + saveBool(CHAT_RAG_OCR_KEY, ragOcrScanned); + return { ragOcrScanned }; + }), + setRagCaptionFigures: (ragCaptionFigures) => + set(() => { + saveBool(CHAT_RAG_CAPTION_KEY, ragCaptionFigures); + return { ragCaptionFigures }; + }), setToolStatus: (toolStatus) => set({ toolStatus }), setActiveDiffusionCanvas: (activeDiffusionCanvas) => set({ activeDiffusionCanvas }), diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index b8391e31e7..412f35d1a0 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -350,6 +350,9 @@ export interface OpenAIChatCompletionsRequest { mode: "hybrid" | "lexical" | "dense"; autoinject?: boolean; autoinject_min_score?: number; + + whole_doc?: boolean; + context_length?: number; }; auto_heal_tool_calls?: boolean; max_tool_calls_per_message?: number; diff --git a/studio/frontend/src/features/rag/api/rag-api.ts b/studio/frontend/src/features/rag/api/rag-api.ts index b1215d2d46..20800230b5 100644 --- a/studio/frontend/src/features/rag/api/rag-api.ts +++ b/studio/frontend/src/features/rag/api/rag-api.ts @@ -39,9 +39,17 @@ async function ragRequest( return json as T; } -async function ragUpload(path: string, file: File): Promise { +async function ragUpload( + path: string, + file: File, + ocr?: boolean, + caption?: boolean, +): Promise { const form = new FormData(); form.append("file", file); + // Per-upload overrides for the vision passes; omitted -> backend config default. + if (ocr !== undefined) form.append("ocr", String(ocr)); + if (caption !== undefined) form.append("caption", String(caption)); // No Content-Type: let the browser set the multipart boundary. const response = await authFetch(`${RAG_BASE}${path}`, { method: "POST", @@ -103,10 +111,14 @@ export async function listKnowledgeBaseDocuments( export function uploadKnowledgeBaseDocument( kbId: string, file: File, + ocr?: boolean, + caption?: boolean, ): Promise { return ragUpload( `/knowledge-bases/${encodeURIComponent(kbId)}/documents`, file, + ocr, + caption, ); } @@ -122,8 +134,15 @@ export async function listThreadDocuments( export function uploadThreadDocument( threadId: string, file: File, + ocr?: boolean, + caption?: boolean, ): Promise { - return ragUpload(`/threads/${encodeURIComponent(threadId)}/documents`, file); + return ragUpload( + `/threads/${encodeURIComponent(threadId)}/documents`, + file, + ocr, + caption, + ); } export async function listProjectDocuments( @@ -138,8 +157,15 @@ export async function listProjectDocuments( export function uploadProjectDocument( projectId: string, file: File, + ocr?: boolean, + caption?: boolean, ): Promise { - return ragUpload(`/projects/${encodeURIComponent(projectId)}/documents`, file); + return ragUpload( + `/projects/${encodeURIComponent(projectId)}/documents`, + file, + ocr, + caption, + ); } // Cached "does this project have indexed sources?" probe so the chat adapter can diff --git a/studio/frontend/src/features/rag/components/retrieval-settings-section.tsx b/studio/frontend/src/features/rag/components/retrieval-settings-section.tsx index 103ed5c8e8..be791b8e6d 100644 --- a/studio/frontend/src/features/rag/components/retrieval-settings-section.tsx +++ b/studio/frontend/src/features/rag/components/retrieval-settings-section.tsx @@ -9,6 +9,7 @@ import { SelectValue, } from "@/components/ui/select"; import { Slider } from "@/components/ui/slider"; +import { Switch } from "@/components/ui/switch"; import { ToggleGroup, ToggleGroupItem } from "@/components/ui/toggle-group"; import { Tooltip, @@ -109,6 +110,12 @@ export function RetrievalSettingsSection() { const setRagAutoInjectMinScore = useChatRuntimeStore( (s) => s.setRagAutoInjectMinScore, ); + const ragOcrScanned = useChatRuntimeStore((s) => s.ragOcrScanned); + const setRagOcrScanned = useChatRuntimeStore((s) => s.setRagOcrScanned); + const ragCaptionFigures = useChatRuntimeStore((s) => s.ragCaptionFigures); + const setRagCaptionFigures = useChatRuntimeStore( + (s) => s.setRagCaptionFigures, + ); return (
@@ -202,6 +209,51 @@ export function RetrievalSettingsSection() { format={(v) => v.toFixed(2)} />
+ +
+
+ + OCR scanned pages + + Read text off scanned or image-only PDF pages with the loaded + model's vision, at upload time, so picture-only documents become + searchable. Needs a vision model; pages with a text layer are + unaffected. + + + + Transcribe image-only PDF pages when attaching. + +
+ +
+ +
+
+ + Describe figures & charts + + Caption PDF figures, charts, tables and diagrams at upload with the + loaded model's vision, so their content becomes searchable. Needs a + vision model; adds vision calls for detected figures. + + + + Read charts and diagrams when attaching. + +
+ +
); } diff --git a/studio/frontend/src/features/rag/components/use-rag-documents.ts b/studio/frontend/src/features/rag/components/use-rag-documents.ts index 488d9ed6d0..8e756b2782 100644 --- a/studio/frontend/src/features/rag/components/use-rag-documents.ts +++ b/studio/frontend/src/features/rag/components/use-rag-documents.ts @@ -2,6 +2,12 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import { useCallback, useEffect, useRef, useState } from "react"; +import { useChatRuntimeStore } from "@/features/chat"; + +import { + CHAT_RAG_CAPTION_KEY, + CHAT_RAG_OCR_KEY, +} from "@/features/chat/stores/chat-runtime-store"; import { toast } from "@/lib/toast"; import { deleteDocument, @@ -45,13 +51,19 @@ export function useRagDocuments( }, [documents]); // documentId -> signature; forgotten on delete, cleared on scope change. const sigByDocId = useRef>(new Map()); - const sigAttached = useCallback( - (sig: string) => { - for (const s of sigByDocId.current.values()) if (s === sig) return true; - return false; - }, - [], - ); + // Skip a re-selected file only if a matching doc is healthy or still indexing. A doc + // that completed with 0 chunks is re-ingestable (e.g. a scan attached before a vision + // model loaded); the backend re-ingests on the same hash, so let it through. + const sigBlocksReupload = useCallback((sig: string) => { + const ids = new Set(); + for (const [id, s] of sigByDocId.current) if (s === sig) ids.add(id); + if (ids.size === 0) return false; + const docs = documentsRef.current.filter((d) => ids.has(d.id)); + if (docs.length === 0) return false; // sig tracked but doc gone -> allow re-upload + return docs.some( + (d) => d.status !== "completed" || (d.numChunks ?? 0) > 0, + ); + }, []); // True while upload() runs, so the scope-change effect can tell a real switch // from lazy thread materialization mid-upload (which must not reset). const uploadInFlightRef = useRef(false); @@ -82,7 +94,11 @@ export function useRagDocuments( const controller = new AbortController(); trackedJobs.current.set(jobId, controller); - const finish = (status: DocumentStatus, error?: string | null) => { + const finish = ( + status: DocumentStatus, + error?: string | null, + numChunks?: number | null, + ) => { if (status === "failed") { // Drop the chip rather than show "Failed"; warn via toast. sigByDocId.current.delete(documentId); @@ -91,7 +107,14 @@ export function useRagDocuments( description: error ?? "Indexing failed", }); } else { - patchDoc(documentId, { status, error: null, progress: 1 }); + // Record numChunks so re-selecting this file dedups (vs a 0-chunk doc, which + // stays re-ingestable); the SSE "complete" frame carries it. + patchDoc(documentId, { + status, + error: null, + progress: 1, + ...(numChunks != null ? { numChunks } : {}), + }); } trackedJobs.current.delete(jobId); }; @@ -105,7 +128,7 @@ export function useRagDocuments( progress: ev.progress ?? null, }); } else if (ev.type === "complete") { - finish("completed"); + finish("completed", null, ev.num_chunks); return; } else if (ev.type === "error") { finish("failed", ev.error ?? "Indexing failed"); @@ -121,6 +144,7 @@ export function useRagDocuments( ? "failed" : "completed", job.error, + job.numChunks, ); } catch { if (controller.signal.aborted) { @@ -132,7 +156,8 @@ export function useRagDocuments( for (let i = 0; i < 600; i++) { if (controller.signal.aborted) break; const job = await getJob(jobId); - if (job.status === "completed") return finish("completed"); + if (job.status === "completed") + return finish("completed", null, job.numChunks); if (job.status === "failed") { return finish("failed", job.error ?? "Indexing failed"); } @@ -231,12 +256,21 @@ export function useRagDocuments( tempId: string, ) => { try { + // Send vision-pass overrides only after the user has explicitly set them; + // otherwise backend env defaults own the ingest policy. + const state = useChatRuntimeStore.getState(); + const hasLocal = (key: string) => + typeof window !== "undefined" && window.localStorage.getItem(key) !== null; + const ocr = hasLocal(CHAT_RAG_OCR_KEY) ? state.ragOcrScanned : undefined; + const caption = hasLocal(CHAT_RAG_CAPTION_KEY) + ? state.ragCaptionFigures + : undefined; const result = activeScope.type === "kb" - ? await uploadKnowledgeBaseDocument(activeScope.kbId, file) + ? await uploadKnowledgeBaseDocument(activeScope.kbId, file, ocr, caption) : activeScope.type === "project" - ? await uploadProjectDocument(activeScope.projectId, file) - : await uploadThreadDocument(activeScope.threadId, file); + ? await uploadProjectDocument(activeScope.projectId, file, ocr, caption) + : await uploadThreadDocument(activeScope.threadId, file, ocr, caption); sigByDocId.current.set(result.documentId, fileSignature(file)); if (seenIds.has(result.documentId)) { setDocuments((rows) => rows.filter((row) => row.id !== tempId)); @@ -286,7 +320,7 @@ export function useRagDocuments( // one look like nothing happened. Dedup re-selections up front. const fresh: Array<{ tempId: string; file: File }> = []; for (const file of Array.from(files)) { - if (sigAttached(fileSignature(file))) { + if (sigBlocksReupload(fileSignature(file))) { toast.info(`${file.name} is already indexed - skipping`); continue; } @@ -332,7 +366,7 @@ export function useRagDocuments( uploadInFlightRef.current = false; } }, - [scope, uploadOne, sigAttached], + [scope, uploadOne, sigBlocksReupload], ); const remove = useCallback( diff --git a/studio/frontend/src/features/rag/types/rag.ts b/studio/frontend/src/features/rag/types/rag.ts index d01854f5a7..1277500ae6 100644 --- a/studio/frontend/src/features/rag/types/rag.ts +++ b/studio/frontend/src/features/rag/types/rag.ts @@ -39,6 +39,7 @@ export interface IndexJob { stage?: string | null; progress?: number | null; error?: string | null; + numChunks?: number | null; } /** One SSE frame from /jobs/{jobId}/events. */ @@ -47,6 +48,7 @@ export interface JobEvent { stage?: string | null; progress?: number | null; error?: string | null; + num_chunks?: number | null; } /** Coords 0..1, top-left origin. */