Whole-document context for RAG chat attachments (#6693)
* Add whole-document context mode to RAG chat attachments
Thread-attached files are injected in full when they fit a token budget,
instead of only top-K retrieved chunks, so the model reads the entire file
for summarize/reason-over-document requests. Oversized files fall back to
top-K retrieval so the context window is never blown. KB and project
corpora are unchanged (still retrieval).
- core/rag/store.py: all_chunks_for_scope returns every completed-document
chunk for a scope, ordered document-then-index, joined with filename.
- core/rag/tool.py: whole_document_context renders the chunks as the same
<chunk> blocks + citation source-map retrieval produces, returns None
when empty or over budget.
- core/inference/tools.py: build_rag_autoinject tries whole-document first
for thread scopes, falls through to search_for_autoinject otherwise.
- core/rag/config.py: THREAD_WHOLE_DOC + WHOLE_DOC_MAX_TOKENS (env-tunable).
- tests/test_rag_whole_document.py: store ordering, whole-doc render +
budget cutoff, auto-inject whole-doc vs top-K fallback, KB never whole-doc.
* Add scanned-PDF OCR fallback to RAG ingestion
A PDF page with no extractable text layer (a scanned or image-only page)
previously ingested as empty, so image PDFs were invisible to retrieval and
whole-document context. Such pages are now rendered and transcribed by the
loaded vision model during ingestion, so they become searchable and readable
like any other page. This restores OCR for the RAG document flow without a
separate extraction pipeline.
- core/rag/parsers.py: render_pdf_pages renders whole pages (1-based) to PNG.
- core/rag/captioner.py: factor the shared vision call into _vision_complete;
add _ocr_one + ocr_pages (transcribe rendered pages, OCR_MAX_PAGES bound).
- core/rag/ingestion.py: _ocr_scanned_pages runs right after parse, replacing
text on near-empty PDF pages. No-op when OCR is off, no page is scanned, or
no vision model is loaded (degrades like figure captioning).
- core/rag/config.py: OCR_SCANNED, OCR_MIN_CHARS, OCR_MAX_PAGES, OCR_DPI,
OCR_TIMEOUT_S, OCR_MAX_TOKENS (env-tunable).
- tests/test_rag_ocr_fallback.py: page render, ocr_pages gating + cap, scanned
PDF end-to-end OCR into chunks + whole-doc, born-digital skips OCR, disabled
leaves the page empty.
* Broaden OCR prompt to figures/tables and guard against repetition runaway
The OCR prompt now asks the vision model to also transcribe text inside figures,
diagrams, charts and tables, so labels and table cells on scanned pages are
indexed rather than skipped. Verified on real documents that this does not
regress plain-text transcription.
Some vision models loop on sparse images (e.g. a title-only cover) and emit the
same line hundreds of times. _collapse_runaway caps any run of identical
consecutive lines so a pathological page cannot flood the index; legitimate
short repeats (a label appearing a few times) survive. Applied in ocr_pages.
* Restrict whole-document injection to thread attachments only
whole_document_context resolved the combined project+thread scope, so a project
chat (the frontend sends both thread_id and project_id) injected the entire
project corpus in full, contradicting the design that project and KB corpora stay
retrieval-only. A large project corpus could also push the total over budget and
drop a small thread attachment back to top-K.
Resolve the thread scope alone in whole_document_context, and in
build_rag_autoinject only enter whole-doc mode when a thread attachment is present
and no KB is selected (a KB pick is exclusive: search that corpus). Project
sources and KBs keep top-K retrieval. Adds regression tests for the mixed
project+thread payload, the budget isolation, and KB precedence.
* Address review: keep project retrieval, harden budget + OCR guards
Follow-up to the 8-reviewer pass on the whole-document + OCR work.
- Preserve project grounding in project chats. The thread-scope-only fix made
whole-doc exclusive of retrieval, so a thread attachment silently dropped the
project corpus for that turn. build_rag_autoinject now whole-docs the thread
attachment AND retrieves the project sources top-K, merged under one citation
numbering via tool.render_sources. KB selection stays exclusive.
- Budget: a NULL/zero token_count no longer bypasses the cap (length-based
fallback in _row_token_count), so a malformed huge doc can't inject in full.
- OCR runaway guard: _collapse_runaway now also caps each distinct line at a
generous total across the page (not just consecutive), bounding the
interleaved/alternating loops weak models emit; blank-line floods collapse too.
- OCR: warn when a scanned PDF exceeds OCR_MAX_PAGES (pages past the cap stay
untranscribed) instead of silently dropping them.
- Document the known limits: OCR'd pages have no PDF highlight regions; vision
models need a micro-batch >= image tokens (Gemma-family) or the server aborts.
- Tests for project-retrieval composition, NULL-token budget, and interleaved
runaway; drop the now-superseded exclude-project test.
* Add OCR toggle to RAG retrieval settings
Make scanned-PDF OCR user-controllable per upload instead of only via the
RAG_OCR_SCANNED config default. The retrieval settings panel gains an OCR
scanned pages switch (persisted in localStorage, on by default); the chosen
value is read fresh at upload time and sent with each document upload.
Backend: the three upload routes accept an optional ocr form field and pass it
through start_ingestion to _ocr_scanned_pages, which now treats None as use the
config default and an explicit bool as an override. The on/off policy lives only
in _ocr_scanned_pages now, so ocr_pages no longer re-checks the config (that
double gate would have blocked a per-upload ocr=True while the default was off).
Tests cover both override directions (force on while config off, force off while
config on).
* Add "Describe figures & charts" toggle with chart-aware captions
Surface RAG figure captioning as a user control and make it actually useful for
graphs and plots. The figure detection already clustered vector drawings and
raster images into regions and rendered them, but captioning was off by default,
had no UI, and used a thin generic prompt.
Accuracy: the caption prompt now asks for chart type, axis titles and units,
legend or series, salient trends and readable values, and table columns, while
forbidding invented numbers. The token budget is configurable (CAPTION_MAX_TOKENS)
and captions pass through the same runaway guard as OCR so a looping vision model
cannot flood the index.
Control: a per-upload caption override threads from the three upload routes through
start_ingestion and _run, with the on/off policy single-sourced in _run (caption
self-gating removed from caption_images, mirroring the OCR change) so a force-on
override works when the config default is off. The frontend adds a "Describe
figures & charts" switch in the retrieval settings, persisted in localStorage and
sent with each upload. Default on; it is a no-op without a vision model and bounded
to CAPTION_MAX_IMAGES figures per document.
Tests cover the new caption_images contract, the runaway guard on captions, the
chart-aware prompt and token budget (and that OCR keeps its own prompt and budget),
and both override directions end to end through ingestion.
* Generalize figure understanding: transcribe-first prompt + high-DPI tiling
Make figure/chart description work across any visual and any model strength, not
just a strong VLM on simple figures. Two changes, validated by a recall benchmark
on authoritative documents (ResNet/Attention papers, USDA, UN UDHR).
1. Transcribe-first caption prompt. The caption now asks the model to transcribe
every visible label verbatim (titles, axis labels and units, legends, every
box/node/arrow label, table cells, equations) and then add a one-line summary,
instead of only describing the figure. Transcription is the most model-robust
visual task, so weak models that cannot reason about a chart still recover its
labels.
2. High-DPI tiling of figure pages. Figure-bearing pages are rendered as an
overlapping grid of high-DPI tiles (plus a full-page pass for context); each
tile is transcribed, then merged and de-duplicated. This keeps small diagram
labels legible and covers every sub-figure without relying on exact region
detection, which previously missed sub-figures and small labels.
Supporting changes: figure render DPI 130 -> 200 with a clip margin so edge labels
are not lost; vision calls are deterministic (temperature 0) so transcription does
not randomly drop labels; the repetition guard now applies to captions too. New
config knobs: FIGURE_DPI, FIGURE_MARGIN_FRAC, FIGURE_TILE_ROWS/COLS, FIGURE_TILE_
OVERLAP, FIGURE_FULLPAGE, CAPTION_MAX_PAGES, larger CAPTION_MAX_TOKENS, and
CAPTION_MAX_IMAGES as a per-document tile budget.
Measured figure context recall (per-label, dense academic figures):
Qwen2.5-VL: 0.50 -> 0.83 (overall 0.81 -> 0.94)
Gemma-4-E2B (weak): ~0 with loops -> 0.83 (overall 0.91)
Born-digital text and scanned-page recall are unchanged (no regression).
parsers gains _figure_boxes (shared detection), pages_with_figures, and
render_pdf_figure_tiles; captioner gains merge_page_captions and a temperature
parameter; ingestion routes figure captioning through the tiled path.
* Fix RAG review issues: whole-doc budget pre-check, figure gating, empty re-ingest, vision auth
Whole-document context now runs a cheap token-sum pre-check (store.scope_token_estimate)
before hydrating every chunk's text, so an attachment that cannot fit the budget is
rejected without loading the whole corpus into memory. The estimate mirrors
all_chunks_for_scope's filter and the per-row token-count fallback exactly.
Ingestion skips all figure work (PDF rasterization and detection, not just the caption
call) unless a vision model is loaded, so a text-only deployment pays nothing. When OCR
is enabled, scanned/image-only pages are excluded from figure tiling since OCR already
transcribes them whole, avoiding double vision work and overlapping index entries; a
scanned figure page is still tiled when OCR is off.
start_ingestion no longer dedupes forever to a prior ingest that produced zero chunks
(e.g. a scanned PDF uploaded before a vision model was loaded): the empty record is
dropped and the content is re-ingested.
Vision OCR and caption requests now send the backend Authorization header, so they
match the chat endpoint and do not 401 under direct-stream (--api-key) mode.
Adds tests for the budget estimate, scanned-page exclusion, the vision-model gate, the
empty re-ingest path, and the auth-header passthrough.
* Trim RAG vision-ingestion comments and docstrings
Tighten the verbose multi-line docstrings and comments added across the RAG vision
ingestion work (captioner, config, parsers, ingestion, store, tool, build_rag_autoinject,
the RAG tests, and the chat-store/upload-hook frontend toggles) to one or two lines while
keeping their intent. No code changed: verified comment/docstring-only against the prior
commit, and the RAG test suite still passes.
* Fix figure-tiling exclusion and client dedupe for re-ingestable docs
Figure tiling now excludes only the pages OCR actually transcribed, not every
text-less page. _ocr_scanned_pages returns the set of pages it OCR'd, and _run passes
that to pages_with_figures as exclude_pages (replacing the ocr_on-keyed min_text_chars
heuristic). A scanned page that OCR skipped (past OCR_MAX_PAGES, or whose OCR returned
empty) is no longer dropped from captioning, so a chart on such a page still gets a
caption.
The document panel's upload dedupe no longer skips re-selecting a file whose only
matching doc completed with zero chunks. Such a doc is re-ingestable (e.g. a scan
attached before a vision model loaded), and the backend re-ingests on the same content
hash, so the client must let it reach the backend; healthy or still-indexing docs are
still skipped. The SSE complete frame's chunk count is recorded on the doc so the
check is exact.
Adds a regression test for the un-OCR'd scanned figure page and updates the
pages_with_figures test to the exclude_pages interface.
* Address review findings: whole-doc budget guard, job numChunks, dead code, upload cap
whole_document_context now treats a non-positive max_tokens as "never inject" instead
of injecting the whole corpus unbounded, so RAG_WHOLE_DOC_MAX_TOKENS=0 tightens rather
than disables the budget (the real off switch stays RAG_THREAD_WHOLE_DOC=0).
The job-status endpoint and get_job_status now expose num_chunks (joined from the
document), and the upload hook threads it through the SSE-fallback completion paths
(reconcile + poll). Previously a document that finished via the connection-cap fallback
had no chunk count client-side, so the re-ingest dedupe wrongly treated it as empty and
re-uploaded it. IndexJob/JobEvent gain the field and the untyped cast is dropped.
Removes the dead render_pdf_figures function (superseded by the tiling path), its test,
and the unused FIGURE_MARGIN_FRAC config knob.
Adds an upload size cap (RAG_MAX_UPLOAD_BYTES, default 200 MB; 413 on exceed with the
partial file cleaned up) so a pathological file can't drive unbounded parse + vision
work. render_pdf_figure_tiles clamps rows/cols to >= 1 (no ZeroDivisionError on a
misconfigured grid). Captioning progress is reported after OCR so the bar is monotonic.
sqlite connections set busy_timeout=5000 so a long figure/scan ingest holding its
connection doesn't make a concurrent ingest/read fail with "database is locked".
Adds tests for the non-positive budget, the zero-grid clamp, job-status num_chunks, and
the oversize-upload rejection.
* Extract PDF text as layout-aware Markdown via pymupdf4llm
parsers._pdf now extracts each PDF page as Markdown with pymupdf4llm.to_markdown
(page_chunks=True) instead of flat page.get_text("text"), so tables, headings and lists
keep their structure in the indexed chunks and retrieve far better (a table's cells stay
associated with their row instead of flattening into a token stream). Gated by
RAG_PDF_MARKDOWN (default on); falls back to plain PyMuPDF text when the toggle is off,
pymupdf4llm is missing, extraction fails, or a page yields no Markdown. The scanned-page
OCR and figure-tiling passes operate on rendered pixels and are unaffected; docx/html/txt
keep their existing extractors.
The preview-highlight locator already strips Markdown punctuation when building anchors;
it now also splits anchor tokens on pipes so a Markdown table row still anchors to the
raw PDF word stream.
Declares pymupdf4llm as a studio/RAG dependency (was only transitively present via the
data-designer plugin). Adds parser tests (Markdown table reaches the page text, the
plain-text fallback, the missing-lib fallback) and a locator test for table-pipe anchoring.
* Pin pymupdf4llm to 0.3.4 so the package scan does not pull onnxruntime
The lockstep pymupdf4llm 1.27.x line makes pymupdf-layout a hard dependency,
which in turn pulls onnxruntime (plus numpy/networkx/protobuf). The security-audit
pip scan-packages job resolves requirements --with-deps, so adding pymupdf4llm to
no-torch-runtime.txt and studio.txt surfaced onnxruntime's un-baselined CRITICAL
finding and flipped the hf-stack shard from pass to fail.
pymupdf4llm 0.3.x keeps pymupdf-layout behind an optional [layout] extra, so a plain
install resolves to pymupdf + tabulate only and never touches onnxruntime. 0.3.4
requires pymupdf>=1.27.1, satisfied by our pinned pymupdf==1.27.2.3, and to_markdown
(page_chunks=True) produces equivalent layout-aware Markdown on real PDFs (verified on
the Attention, ResNet and USDA documents). Production already installs these files
--no-deps, so onnxruntime was never shipped at runtime; this only fixes the scanner.
The parser test now asserts Markdown markup (heading or table pipes) rather than table
pipes specifically, since 0.3.4 emits a heading but not a pipe table on the tiny
borderless synthetic fixture; both markers are absent from the plain-text fallback.
* Fix RAG whole-doc review findings
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Address RAG whole-doc review follow-ups
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Address RAG review follow-up edge cases
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Reserve image budget for whole-document RAG
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Co-authored-by: danielhanchen <michaelhan2050@gmail.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: wasimysaid <wasimysdev@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han
GitHub
parent
43d3caf38b
commit
e8945cab46
25 changed files with 2168 additions and 154 deletions
|
|
@ -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}
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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 ``<chunk>`` 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'<chunk id="{i}" source={src}{page_attr}>\n{s.get("text") or ""}\n</chunk>')
|
||||
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
|
||||
``<chunk>`` 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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
259
studio/backend/tests/test_rag_ocr_fallback.py
Normal file
259
studio/backend/tests/test_rag_ocr_fallback.py
Normal file
|
|
@ -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
|
||||
88
studio/backend/tests/test_rag_parsing.py
Normal file
88
studio/backend/tests/test_rag_parsing.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
520
studio/backend/tests/test_rag_whole_document.py
Normal file
520
studio/backend/tests/test_rag_whole_document.py
Normal file
|
|
@ -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 <chunk> blocks.
|
||||
assert "chapter one body" in text
|
||||
assert "chapter two body" in text
|
||||
assert '<chunk id="1"' in text
|
||||
assert '<chunk id="2"' in text
|
||||
assert text.index("chapter one") < text.index("chapter two")
|
||||
# Source-map mirrors retrieval's shape, with no score on the whole-doc path.
|
||||
assert [s["citationId"] for s in sources] == [1, 2]
|
||||
assert all(s["filename"] == "report.pdf" for s in sources)
|
||||
assert all(s["score"] is None for s in sources)
|
||||
assert [s["page"] for s in sources] == [1, 2]
|
||||
assert [s["chunkId"] for s in sources] == ["d1:0", "d1:1"]
|
||||
|
||||
|
||||
def test_whole_document_context_none_over_budget(rag_conn):
|
||||
scope = store.thread_scope("t1")
|
||||
_add_doc(rag_conn, scope, "d1", "big.pdf", "h1", ["huge"], tokens = [50_000])
|
||||
assert tool.whole_document_context(scope_thread_id = "t1", max_tokens = 6000) is None
|
||||
# Same doc fits under a larger budget.
|
||||
assert tool.whole_document_context(scope_thread_id = "t1", max_tokens = 100_000) is not None
|
||||
|
||||
|
||||
def test_whole_document_context_none_when_empty(rag_conn):
|
||||
assert tool.whole_document_context(scope_thread_id = "t1", max_tokens = 6000) is None
|
||||
|
||||
|
||||
def test_whole_document_context_non_positive_budget_returns_none(rag_conn):
|
||||
# A non-positive budget disables whole-doc (RAG_WHOLE_DOC_MAX_TOKENS=0 footgun)
|
||||
# rather than injecting the whole corpus unbounded.
|
||||
scope = store.thread_scope("t1")
|
||||
_add_doc(rag_conn, scope, "d1", "a.pdf", "h1", ["tiny body"])
|
||||
assert tool.whole_document_context(scope_thread_id = "t1", max_tokens = 0) is None
|
||||
assert tool.whole_document_context(scope_thread_id = "t1", max_tokens = -5) is None
|
||||
|
||||
|
||||
def test_whole_document_context_none_without_scope(rag_conn):
|
||||
# No thread scope -> 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 '<chunk id="1"' in injected
|
||||
assert '<chunk id="2"' in injected
|
||||
assert '<chunk id="3"' in injected
|
||||
|
||||
|
||||
def test_build_rag_autoinject_skips_project_companion_over_budget(rag_conn, monkeypatch):
|
||||
_add_doc(rag_conn, store.thread_scope("t1"), "td", "thread.txt", "h1", ["thread body"])
|
||||
project_text = "project overflow " * 2000
|
||||
proj = (
|
||||
"PROJ",
|
||||
[
|
||||
{
|
||||
"citationId": 1,
|
||||
"chunkId": "pj:0",
|
||||
"documentId": "pj",
|
||||
"filename": "project.txt",
|
||||
"page": None,
|
||||
"text": project_text,
|
||||
"score": 0.91,
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
monkeypatch.setattr(tool, "search_for_autoinject", lambda **kw: proj)
|
||||
result = inf_tools.build_rag_autoinject(_convo(), {"thread_id": "t1", "project_id": "p1"})
|
||||
injected = _injected_text(result)
|
||||
assert "thread body" in injected
|
||||
assert "project overflow" not in injected
|
||||
|
||||
|
||||
def test_build_rag_autoinject_thread_whole_doc_ignores_project_size(rag_conn, monkeypatch):
|
||||
# A large project corpus must not push a small thread attachment over budget;
|
||||
# whole-doc resolves the thread scope alone (companion retrieval stubbed out).
|
||||
monkeypatch.setattr(tool, "search_for_autoinject", lambda **kw: None)
|
||||
_add_doc(rag_conn, store.thread_scope("t1"), "td", "thread.txt", "h1", ["small thread file"])
|
||||
_add_doc(
|
||||
rag_conn, store.project_scope("p1"), "pd", "project.txt", "h2", ["big"], tokens = [50_000]
|
||||
)
|
||||
result = inf_tools.build_rag_autoinject(_convo(), {"thread_id": "t1", "project_id": "p1"})
|
||||
assert "small thread file" in _injected_text(result)
|
||||
|
||||
|
||||
def test_build_rag_autoinject_kb_defers_to_retrieval(rag_conn, monkeypatch):
|
||||
# A KB selection is exclusive: a thread attachment can't preempt it; KB uses retrieval.
|
||||
_add_doc(rag_conn, store.thread_scope("t1"), "td", "thread.txt", "h1", ["thread attachment"])
|
||||
sentinel = ("KB_RETRIEVAL", [{"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", "thread_id": "t1"})
|
||||
assert _injected_text(result) == "KB_RETRIEVAL"
|
||||
|
||||
|
||||
def test_build_rag_autoinject_no_scope_returns_none(rag_conn):
|
||||
assert inf_tools.build_rag_autoinject(_convo(), None) is None
|
||||
assert inf_tools.build_rag_autoinject(_convo(), {}) is None
|
||||
|
||||
|
||||
def test_build_rag_autoinject_args_carry_user_query(rag_conn):
|
||||
scope = store.thread_scope("t1")
|
||||
_add_doc(rag_conn, scope, "d1", "doc.pdf", "h1", ["small body"])
|
||||
result = inf_tools.build_rag_autoinject(_convo("what is in here"), {"thread_id": "t1"})
|
||||
assistant_msg = next(m for m in result["messages"] if m.get("role") == "assistant")
|
||||
args = json.loads(assistant_msg["tool_calls"][0]["function"]["arguments"])
|
||||
assert args["query"] == "what is in here"
|
||||
|
||||
|
||||
# ── end-to-end: real ingestion pipeline -> 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("<chunk id=") == doc["num_chunks"]
|
||||
|
|
@ -2728,6 +2728,12 @@ export function createOpenAIStreamAdapter(
|
|||
params.checkpoint,
|
||||
),
|
||||
autoinject_min_score: ragAutoInjectMinScore,
|
||||
|
||||
...(ragAutoInject === "off"
|
||||
? { whole_doc: false }
|
||||
: {}),
|
||||
context_length:
|
||||
runtime.ggufContextLength ?? params.maxSeqLength ?? undefined,
|
||||
},
|
||||
}
|
||||
: {}),
|
||||
|
|
|
|||
|
|
@ -51,6 +51,8 @@ export const CHAT_RAG_TOP_K_KEY = "unsloth_chat_rag_top_k";
|
|||
export const CHAT_RAG_AUTOINJECT_KEY = "unsloth_chat_rag_autoinject";
|
||||
export const CHAT_RAG_AUTOINJECT_MIN_SCORE_KEY =
|
||||
"unsloth_chat_rag_autoinject_min_score";
|
||||
export const CHAT_RAG_OCR_KEY = "unsloth_chat_rag_ocr_scanned";
|
||||
export const CHAT_RAG_CAPTION_KEY = "unsloth_chat_rag_caption_figures";
|
||||
export const CHAT_SPECULATIVE_TYPE_KEY = "unsloth_chat_speculative_type";
|
||||
|
||||
// Persist only the model-agnostic intents (auto/ngram/off). MTP modes
|
||||
|
|
@ -69,6 +71,12 @@ export const DEFAULT_RAG_TOP_K = 5;
|
|||
export type RagAutoInject = "auto" | "on" | "off";
|
||||
export const DEFAULT_RAG_AUTOINJECT: RagAutoInject = "auto";
|
||||
export const DEFAULT_RAG_AUTOINJECT_MIN_SCORE = 0.7;
|
||||
// OCR scanned/image-only PDF pages at ingest time. On by default; off skips the
|
||||
// extra vision pass (only matters when the loaded chat model has vision).
|
||||
export const DEFAULT_RAG_OCR = true;
|
||||
// Describe figures/charts in PDFs at ingest time so they become searchable. On by
|
||||
// default (no-op without a vision model); off skips the per-figure vision calls.
|
||||
export const DEFAULT_RAG_CAPTION = true;
|
||||
|
||||
function loadRagSource(): RagSource {
|
||||
if (typeof window === "undefined") return DEFAULT_RAG_SOURCE;
|
||||
|
|
@ -592,6 +600,10 @@ type ChatRuntimeStore = {
|
|||
// autoInject = forced first-pass retrieval before answering.
|
||||
ragAutoInject: RagAutoInject;
|
||||
ragAutoInjectMinScore: number;
|
||||
// OCR scanned/image-only PDF pages at ingest time (vision model required).
|
||||
ragOcrScanned: boolean;
|
||||
// Describe figures/charts at ingest time (vision model required).
|
||||
ragCaptionFigures: boolean;
|
||||
/**
|
||||
* When on, local Studio tool calls pause for an explicit allow/deny in the
|
||||
* chat before they run.
|
||||
|
|
@ -757,6 +769,8 @@ type ChatRuntimeStore = {
|
|||
setRagTopK: (topK: number) => 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<ChatRuntimeStore>((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<ChatRuntimeStore>((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 }),
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -39,9 +39,17 @@ async function ragRequest<T>(
|
|||
return json as T;
|
||||
}
|
||||
|
||||
async function ragUpload(path: string, file: File): Promise<DocumentUploadResult> {
|
||||
async function ragUpload(
|
||||
path: string,
|
||||
file: File,
|
||||
ocr?: boolean,
|
||||
caption?: boolean,
|
||||
): Promise<DocumentUploadResult> {
|
||||
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<DocumentUploadResult> {
|
||||
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<DocumentUploadResult> {
|
||||
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<DocumentUploadResult> {
|
||||
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
|
||||
|
|
|
|||
|
|
@ -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 (
|
||||
<div className="flex flex-col gap-5 pt-1">
|
||||
|
|
@ -202,6 +209,51 @@ export function RetrievalSettingsSection() {
|
|||
format={(v) => v.toFixed(2)}
|
||||
/>
|
||||
</div>
|
||||
|
||||
<div className="flex items-start justify-between gap-3">
|
||||
<div className="flex flex-col">
|
||||
<span className="flex items-center gap-1.5 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg">
|
||||
OCR scanned pages
|
||||
<InfoHint>
|
||||
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.
|
||||
</InfoHint>
|
||||
</span>
|
||||
<span className="text-[12px] leading-[1.3] text-muted-foreground">
|
||||
Transcribe image-only PDF pages when attaching.
|
||||
</span>
|
||||
</div>
|
||||
<Switch
|
||||
checked={ragOcrScanned}
|
||||
onCheckedChange={setRagOcrScanned}
|
||||
aria-label="OCR scanned pages"
|
||||
className="mt-0.5"
|
||||
/>
|
||||
</div>
|
||||
|
||||
<div className="flex items-start justify-between gap-3">
|
||||
<div className="flex flex-col">
|
||||
<span className="flex items-center gap-1.5 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg">
|
||||
Describe figures & charts
|
||||
<InfoHint>
|
||||
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.
|
||||
</InfoHint>
|
||||
</span>
|
||||
<span className="text-[12px] leading-[1.3] text-muted-foreground">
|
||||
Read charts and diagrams when attaching.
|
||||
</span>
|
||||
</div>
|
||||
<Switch
|
||||
checked={ragCaptionFigures}
|
||||
onCheckedChange={setRagCaptionFigures}
|
||||
aria-label="Describe figures and charts"
|
||||
className="mt-0.5"
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<Map<string, string>>(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<string>();
|
||||
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(
|
||||
|
|
|
|||
|
|
@ -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. */
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue