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. */