ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio: tighten RAG code comments

Browse source 
Shorten and condense comments across the RAG backend, frontend, and
tests for readability. Comment text only; no code, strings, identifiers,
or logic changed. License headers and lint/type pragmas are preserved.
This commit is contained in:
Daniel Han 2026-05-31 08:31:08 +00:00
commit d1348cac3f
60 changed files with 541 additions and 631 deletions
Download patch file Download diff file Expand all files Collapse all files

6
studio/backend/core/inference/llama_cpp.py
View file

@ -4285,9 +4285,9 @@ class LlamaCppBackend:
# without triggering a retry storm. Cancel during both
# prefill and streaming is handled by the watcher thread
# which closes the response, unblocking any httpx read.
# 300 s headroom for large models (30B+) re-prefilling after
# a tool call that returned a long result (e.g. RAG chunks
# with images) — prior 120 s was tripping on Gemma-4-31B.
# 300 s headroom for large models (30B+) re-prefilling after a tool
# call with a long result (e.g. RAG chunks with images) — prior 120 s
# tripped on Gemma-4-31B.
prefill_timeout = httpx.Timeout(
connect = 30,
read = 300.0,

15
studio/backend/core/rag/authorization.py
View file

@ -43,9 +43,8 @@ def document_for_subject_or_404(
referenced thread actually exists in `chat_threads`. A missing
thread row collapses to 404 so a non-existent thread cannot
silently grant access through a dangling `thread_id`.
# TODO(thread-owner): once `chat_threads.owner_user_id` exists,
# join through it the same way KB documents do and drop the
# single-user invariant. Update the test
# TODO(thread-owner): once `chat_threads.owner_user_id` exists, join
# through it like KB docs and drop the single-user invariant. Update
# `tests/test_rag_authorization.py::test_thread_doc_other_user_404`
# to assert per-user isolation rather than thread existence.
@ -83,9 +82,8 @@ def document_for_subject_or_404(
return row
if thread_id is not None:
# Single-user invariant (see TODO above). We require the
# thread row to exist; an unknown thread_id is treated as
# not-found, not as silent grant.
# Single-user invariant (see TODO above): require the thread row to
# exist; an unknown thread_id is not-found, not a silent grant.
thread_row = conn.execute(
"SELECT id FROM chat_threads WHERE id = ?",
(thread_id,),
@ -94,9 +92,8 @@ def document_for_subject_or_404(
raise HTTPException(status_code = 404, detail = _NOT_FOUND_DETAIL)
return row
# Documents must belong to either a KB or a thread (DB CHECK
# constraint enforces XOR on insert); a row that satisfies
# neither is corrupt — treat as 404.
# Docs must belong to a KB or a thread (DB CHECK enforces XOR on insert);
# a row satisfying neither is corrupt — treat as 404.
raise HTTPException(status_code = 404, detail = _NOT_FOUND_DETAIL)

31
studio/backend/core/rag/captioner.py
View file

@ -43,16 +43,14 @@ _PROMPT = (
"body paragraphs."
)
_MAX_NEW_TOKENS = 200
# Downscale large images so the base64 payload stays manageable; the chat
# model's prefill cost scales with image-tile count, not pixel count, but
# very large inputs still bloat the JSON body. 1600 px on the long side
# matches PR #5351's chat-composer extractor.
# Downscale large images to keep the base64 payload manageable; prefill cost
# scales with tile count not pixels, but huge inputs still bloat the JSON body.
# 1600 px on the long side matches PR #5351's chat-composer extractor.
_MAX_IMAGE_SIZE = 1600
_REQUEST_TIMEOUT_SECONDS = 120.0
# Helper VLM (used when no vision-capable chat model is loaded).
# Matches the model pre-cached by precache_helper_gguf() at studio
# startup so the captioner doesn't have to wait on a fresh download.
# Helper VLM (used when no vision-capable chat model is loaded). Matches the
# model pre-cached by precache_helper_gguf() at startup to avoid a fresh download.
_HELPER_REPO = "unsloth/gemma-4-E2B-it-GGUF"
_HELPER_VARIANT = "UD-Q4_K_XL"
_HELPER_MODEL_NAME = "helper"
@ -80,11 +78,9 @@ def _load_helper_vlm() -> Optional[tuple[Any, str, str]]:
try:
from core.inference.llama_cpp import LlamaCppBackend
# kill_orphans=False is critical: the global singleton is
# already running the user's chat-model llama-server. Killing
# "orphans" here would reap that healthy chat process because
# the orphan-killer can't tell two LlamaCppBackend instances
# apart by PID ownership.
# kill_orphans=False is critical: the global singleton already runs the
# user's chat-model llama-server. Killing "orphans" here would reap that
# healthy process — the orphan-killer can't tell two backends apart by PID.
backend = LlamaCppBackend(kill_orphans = False)
logger.info(
"RAG captioner: loading helper VLM as fallback",
@ -127,10 +123,9 @@ def _post_one(client: Any, endpoint: str, model: str, blob: bytes) -> str:
],
"max_tokens": _MAX_NEW_TOKENS,
"temperature": 0.0,
# Reasoning models (gemma-4, qwen3-thinking, etc.) burn the whole
# token budget on <thinking> output and emit empty visible content
# — useless for a short image caption. Disable thinking for this
# request only; the user's chat sessions stay unaffected.
# Reasoning models (gemma-4, qwen3-thinking, etc.) spend the whole budget on
# <thinking> and emit empty visible content — useless for a short caption.
# Disable thinking for this request only; the user's chat sessions stay unaffected.
"chat_template_kwargs": {"enable_thinking": False},
}
response = client.post(endpoint, json = payload)
@ -202,8 +197,8 @@ def caption_images(
)
return out
finally:
# Always tear down the helper if we spawned one. Chat VLM (when
# provided by the parent) is left alone — it's not ours to manage.
# Tear down the helper if we spawned one. The parent-provided chat VLM is
# left alone — it's not ours to manage.
if helper_backend is not None:
try:
helper_backend.unload_model()

10
studio/backend/core/rag/chunking.py
View file

@ -10,13 +10,11 @@ from typing import Callable
from .parsers import ParsedPage
# Match "Figure 1:", "Figure 1.2:", "Fig. 3.", "Table 4:" etc. at line-start,
# tolerating leading bold markers. Used to break chunks BEFORE such captions
# so the caption ends up at the start of its own chunk — dense embeddings
# pool over the whole chunk, so figure references buried at the end get
# diluted by surrounding body text.
# tolerating leading bold markers. Breaks chunks BEFORE such captions so the
# caption starts its own chunk — dense embeddings pool over the whole chunk, so
# figure refs buried at the end get diluted by surrounding body text.
_FIGURE_BOUNDARY_RE = re.compile(
# Number forms covered: "1", "12", "1.2", "B.1" (appendix-style),
# tolerating bold wrappers around either the label or the number.
# Number forms: "1", "12", "1.2", "B.1" (appendix); bold wrappers tolerated.
r"^\**(?:Figure|Fig\.|Table|Tab\.)\s+[A-Z]?\.?\d+(?:\.\d+)?\**[\.:]",
re.MULTILINE | re.IGNORECASE,
)

8
studio/backend/core/rag/embeddings.py
View file

@ -92,11 +92,9 @@ class _BGEVLAdapter:
for start in range(0, len(inputs), batch_size):
batch = list(inputs[start : start + batch_size])
if is_image:
# BGE-VL's internal data_process re-opens each item with
# Image.open(...), which needs a file-like (has .read())
# or a path — NOT a pre-opened PIL Image. Pass BytesIO so
# the model's own opener works. PIL Images get rebuffered
# via an in-memory PNG round-trip.
# BGE-VL's data_process re-opens each item via Image.open(...),
# which needs a file-like (.read()) or path — NOT a pre-opened PIL
# Image. Pass BytesIO; PIL Images get rebuffered via an in-memory PNG.
file_likes: list[Any] = []
for b in batch:
if isinstance(b, (bytes, bytearray)):

49
studio/backend/core/rag/ingestion.py
View file

@ -65,10 +65,9 @@ def _subprocess_worker(
vlm_model: str | None = None,
enable_captions: bool = True,
) -> None:
# Spawned subprocess: structlog isn't configured here (the parent's
# setup runs in the FastAPI process only), so configure it the same
# way so captioner / parser logs render as JSON like the rest, not
# structlog's default dev ConsoleRenderer.
# Spawned subprocess: structlog setup only ran in the parent's FastAPI
# process. Configure it here too so captioner/parser logs render as JSON,
# not structlog's default dev ConsoleRenderer.
try:
import os as _os
@ -85,9 +84,8 @@ def _subprocess_worker(
from core.rag.parsers import inline_image_captions, parse
out_queue.put({"type": "progress", "stage": "parse", "progress": 0.05})
# Always extract images so we can caption + splice for both
# modes. Text mode uses the captions inline in markdown; multimodal
# additionally embeds the raw images as image-kind chunks.
# Always extract images to caption + splice for both modes. Text mode uses
# captions inline in markdown; multimodal also embeds raw images as image-kind chunks.
parsed = parse(Path(stored_path), want_images = True)
pages = parsed.pages
if not pages and not parsed.images:
@ -96,11 +94,9 @@ def _subprocess_worker(
)
return
# Caption figures once (chat VLM if available, else helper VLM
# fallback), then splice captions into the page markdown so the
# chunker indexes them like any other text. Multimodal mode also
# passes these same captions through to _stream_image_chunks
# below — no duplicate VLM calls per image.
# Caption figures once (chat VLM if available, else helper VLM), then splice
# captions into the page markdown so the chunker indexes them like any text.
# Multimodal reuses these captions in _stream_image_chunks below — no duplicate VLM calls.
captions: list[str] = []
if parsed.images and enable_captions:
out_queue.put(
@ -718,9 +714,8 @@ def _pump(
msg["vectors"],
)
except sqlite3.IntegrityError as exc:
# rag_documents row was deleted mid-ingest (user removed
# the chip / cleared the index). Fail the job cleanly
# rather than crashing the pump thread.
# rag_documents row deleted mid-ingest (chip removed / index
# cleared). Fail the job cleanly rather than crashing the pump thread.
final_error = (
f"document was removed before ingestion finished ({exc})"
)
@ -743,9 +738,8 @@ def _pump(
finished_at = int(time.time())
if state.cancelled:
# User cancelled mid-flight. The route-side deleteDocument removes the
# row, file, and chunk artifacts; here we just mark terminal and notify
# subscribers so the SSE stream closes cleanly.
# User cancelled mid-flight. Route-side deleteDocument removes the row, file,
# and chunk artifacts; here we just mark terminal and notify subscribers so the SSE closes.
_update_document_row(state.document_id, status = "cancelled")
_update_job_row(
state.job_id,
@ -806,10 +800,9 @@ def _probe_loaded_vlm() -> tuple[str | None, str | None]:
unsloth in-process VLMs would need a different bridge.
"""
try:
# The singleton getter lives in routes.inference, not the
# llama_cpp module. Importing from the wrong place silently
# returned None for every probe — captioner always fell back
# to the helper VLM even when the chat model was vision-capable.
# The singleton getter lives in routes.inference, not the llama_cpp module.
# The wrong import silently returned None for every probe, so the captioner
# always fell back to the helper VLM even when the chat model was vision-capable.
from routes.inference import get_llama_cpp_backend
except Exception as exc:
logger.warning("RAG probe: get_llama_cpp_backend import failed", error = str(exc))
@ -850,13 +843,11 @@ def enqueue_ingestion(
or resolve_embedder(mode, chunking_strategy)
or RAG_EMBEDDING_MODEL
)
# Probe the loaded chat backend so the subprocess can route figure
# captioning to the user's own vision model (no extra VRAM). Runs
# for both modes — text mode splices captions into markdown, and
# multimodal mode additionally feeds them to the image-vector
# encoder. If no vision chat model is loaded, the subprocess falls
# back to the helper VLM (pre-cached at studio startup). Skipped
# entirely when captioning is disabled for this upload.
# Probe the loaded chat backend so the subprocess can caption figures with the
# user's own vision model (no extra VRAM). Runs for both modes — text splices
# captions into markdown, multimodal also feeds them to the image-vector encoder.
# No vision chat model loaded → falls back to the helper VLM (pre-cached at startup).
# Skipped when captioning is disabled for this upload.
vlm_url: str | None = None
vlm_model: str | None = None
if enable_captions:

7
studio/backend/core/rag/parsers/__init__.py
View file

@ -7,10 +7,9 @@ import re
from dataclasses import dataclass, field
from pathlib import Path
# Same shape as the chunker's figure-boundary regex but captures the
# figure label (Figure / Fig. / Table / Tab.) AND the number so we can
# attribute a VLM caption back to a specific figure on a multi-figure
# page.
# Same shape as the chunker's figure-boundary regex but also captures the label
# (Figure / Fig. / Table / Tab.) AND number, so a VLM caption maps to a specific
# figure on a multi-figure page.
_FIGURE_LINE_RE = re.compile(
r"^(?P<lead>\**)(?P<label>Figure|Fig\.|Table|Tab\.)\s+"
r"(?P<num>[A-Z]?\.?\d+(?:\.\d+)?)(?P<trail>\**[\.:])",

5
studio/backend/core/rag/parsers/pdf.py
View file

@ -14,9 +14,8 @@ from . import ParsedImage, ParsedPage, ParseResult
logger = logging.getLogger(__name__)
# pymupdf4llm wraps OCR'd vector-graphics text with these markers even when
# `ignore_images=True`. Strip the whole block — the VLM captioner produces
# a proper description for the figure, and the marker text just pollutes
# the chunked body / shows up verbatim in citations.
# `ignore_images=True`. Strip the whole block — the VLM captioner describes the
# figure, and the marker text just pollutes the chunked body / citations.
_PICTURE_TEXT_BLOCK_RE = re.compile(
r"-{3,}\s*Start of picture text\s*-{3,}.*?-{3,}\s*End of picture text\s*-{3,}",
re.DOTALL | re.IGNORECASE,

13
studio/backend/core/rag/reranker.py
View file

@ -18,9 +18,8 @@ from .retrieval import Hit
logger = get_logger(__name__)
# Reentrant: get_reranker() holds the lock while calling unload(), which
# also enters `with _lock`. A plain Lock would deadlock the same thread on
# the second acquisition; RLock allows re-entry from the holding thread.
# Reentrant: get_reranker() holds the lock while calling unload(), which also
# enters `with _lock`. A plain Lock would self-deadlock; RLock allows re-entry.
_lock = threading.RLock()
_model: Any | None = None
_model_name: str | None = None
@ -41,9 +40,8 @@ def _resolve_device() -> str:
def _load(model_name: str) -> Any:
# Stderr print is unconditional so we can see this line even when
# structlog routing is misbehaving — diagnostics for a previously
# invisible hang.
# Unconditional stderr print so this shows even when structlog routing
# misbehaves — diagnostics for a previously invisible hang.
print(
f"[rag.reranker] _load entered: model={model_name}",
file = sys.stderr,
@ -107,8 +105,7 @@ def precache_reranker(model_name: str | None = None) -> None:
elapsed_seconds = round(time.perf_counter() - started, 2),
)
except Exception as exc: # noqa: BLE001
# Non-critical: the lazy loader will retry the download on first
# use. We log so the user can see what happened.
# Non-critical: the lazy loader retries the download on first use; log it.
logger.warning(
"RAG reranker precache failed; will download lazily",
model = target,

10
studio/backend/core/rag/retrieval.py
View file

@ -17,12 +17,10 @@ from utils.rag.config import (
from . import bm25, embeddings, vector_store
# Match "Figure 1", "Figure 1.2", "Figure B.1", "Table 4", "Fig. 5" anywhere
# in the query. Used to inject a third retrieval source that directly looks
# up chunks anchored by these references — dense vectors don't preserve
# figure numbers, so without this an exact-numbered query gets out-ranked
# by chunks describing other figures that share more vocabulary with the
# question.
# Match "Figure 1", "Figure 1.2", "Figure B.1", "Table 4", "Fig. 5" anywhere in
# the query. Feeds a third retrieval source that looks up chunks anchored by these
# refs — dense vectors don't preserve figure numbers, so without this an exact-numbered
# query gets out-ranked by chunks describing other figures with more shared vocabulary.
_FIGURE_REF_RE = re.compile(
r"\b(Figure|Fig\.|Table|Tab\.)\s+([A-Z]?\.?\d+(?:\.\d+)?)\b",
re.IGNORECASE,

26
studio/backend/core/rag/tool.py
View file

@ -17,10 +17,9 @@ from loggers import get_logger
logger = get_logger(__name__)
# Per-request chunk-id counter. Task-local (FastAPI runs each request in
# its own asyncio task → its own Context). Lets the model cite chunks
# unambiguously when multiple search_knowledge_base calls run in the
# same chat turn: call 1 returns ids 1..N, call 2 returns N+1..N+M, etc.
# Per-request chunk-id counter. Task-local (FastAPI runs each request in its own
# asyncio task → its own Context). Keeps citation ids unique across multiple
# search_knowledge_base calls in a turn: call 1 → 1..N, call 2 → N+1..N+M, etc.
_chunk_id_counter: ContextVar[int] = ContextVar("rag_chunk_id_counter", default = 0)
@ -93,11 +92,9 @@ def _format_hits_for_llm(hits: list[dict], start_id: int = 0) -> str:
f'id="{index}"',
f'source="{_xml_attr(hit.get("filename") or "unknown")}"',
]
# Durable backend ids — additive per contracts.md §3.1 (T3).
# ``id`` above stays as the visible citation id (used by the model
# as `[N]`); ``document_id`` + ``chunk_id`` are what the preview
# route consumes. Old XML without these attrs still parses on
# the frontend (hover-only), per contracts §3.2.
# Durable backend ids — additive per contracts.md §3.1 (T3). ``id`` stays
# the visible citation (model's `[N]`); ``document_id`` + ``chunk_id`` feed
# the preview route. Old XML without these still parses (hover-only) per §3.2.
document_id = hit.get("document_id")
if document_id:
attrs.append(f'document_id="{_xml_attr(document_id)}"')
@ -128,8 +125,8 @@ def _format_hits_for_llm(hits: list[dict], start_id: int = 0) -> str:
attrs.append(f'kind="{_xml_attr(kind)}"')
image_path = hit.get("image_path")
if kind == "image" and image_path and document_id:
# Mirror routes/rag.py search-response shape so the frontend
# tool card can render the image inline via the same route.
# Mirror routes/rag.py search-response shape so the frontend tool card
# can render the image inline via the same route.
image_url = f"/api/rag/images/{document_id}/{Path(image_path).name}"
attrs.append(f'image_url="{_xml_attr(image_url)}"')
text = (hit.get("text") or "").strip()
@ -266,10 +263,9 @@ def search_knowledge_base(
else:
hits = hits[:k]
# Merge Hit-side metadata (score, dense_score, chunk_index) into the
# sqlite-side row so the formatter sees one flat dict per chunk.
# Image-kind hits flow through so the multimodal embedder's match
# can reach the LLM; their image_url lets the UI render the picture.
# Merge Hit metadata (score, dense_score, chunk_index) into the sqlite row so
# the formatter sees one flat dict per chunk. Image-kind hits flow through so
# the multimodal match reaches the LLM; their image_url lets the UI render it.
formatted: list[dict] = []
for hit in hits:
row = lookup.get(hit.chunk_id)

4
studio/backend/routes/chat_history.py
View file

@ -230,8 +230,8 @@ async def delete_threads(
payload: ChatDeleteRequest,
current_subject: str = Depends(get_current_subject),
):
# rag_documents has no FK cascade to chat_threads, so purge their
# files + vectors + bm25 explicitly before deleting the threads.
# No FK cascade from rag_documents to chat_threads, so purge their files +
# vectors + bm25 explicitly before deleting the threads.
purge_thread_documents(payload.ids)
delete_chat_threads(payload.ids)
return {"status": "deleted"}

4
studio/backend/routes/inference.py
View file

@ -2712,8 +2712,8 @@ async def openai_chat_completions(
# MCP-only request: skip built-ins, leave room for MCP tools.
tools_to_use = []
elif payload.enabled_tools is not None:
# Preserve client-supplied order so prioritised tools
# (e.g. search_knowledge_base when RAG is on) appear first.
# Preserve client order so prioritised tools (e.g. search_knowledge_base
# when RAG is on) appear first.
_by_name = {t["function"]["name"]: t for t in ALL_TOOLS}
tools_to_use = [
_by_name[name] for name in payload.enabled_tools if name in _by_name

65
studio/backend/routes/rag.py
View file

@ -122,8 +122,7 @@ class UploadResponse(BaseModel):
document_id: str
job_id: str
filename: str
# True when an identical file (same content hash) was already indexed
# in this scope, so no new ingestion job was started. job_id is "".
# Identical content hash already indexed in this scope; no job started, job_id "".
already_indexed: bool = False
@ -268,12 +267,10 @@ async def _save_upload(file: UploadFile) -> tuple[Path, str, int, str]:
stored_path = upload_dir / stored_name
max_bytes = RAG_MAX_UPLOAD_MB * 1024 * 1024
written = 0
# Hash the bytes as they stream so we can dedup identical re-uploads
# within a scope without re-reading the file.
# Hash bytes while streaming to dedup identical re-uploads within a scope.
hasher = hashlib.sha256()
# Route writes through anyio worker thread so the event loop stays free.
# Outer try/except cleans up partial files after async-with closes the fd
# (Windows refuses unlink on an open fd).
# anyio worker thread keeps the event loop free. Outer try/except cleans up
# partial files after the async-with closes the fd (Windows refuses unlink on an open fd).
try:
async with await anyio.open_file(stored_path, "wb") as f:
while True:
@ -313,11 +310,9 @@ def _start_ingestion(
) -> UploadResponse:
document_id = str(uuid4())
with get_connection() as conn:
# Dedup: if an identical file (same content hash) is already
# indexed in this scope, skip re-ingestion. Only a 'completed'
# row counts — a failed/in-flight prior attempt should be allowed
# to retry. Scope is the same kb_id or thread_id the upload
# targets (a file shared across two KBs is indexed in each).
# Dedup: skip re-ingestion if the same content hash is already indexed
# in this scope. Only 'completed' counts — failed/in-flight may retry.
# Scope is the target kb_id or thread_id (a file in two KBs indexes in each).
if content_hash:
if kb_id is not None:
existing = conn.execute(
@ -334,8 +329,7 @@ def _start_ingestion(
(thread_id, content_hash),
).fetchone()
if existing is not None:
# Drop the redundant upload we just wrote to disk; the
# already-indexed copy stays the source of truth.
# Drop the redundant upload; the already-indexed copy is source of truth.
_unlink_if_under_uploads(stored_path)
return UploadResponse(
document_id = existing["id"],
@ -508,8 +502,7 @@ def warmup_rag_embedder(
try:
embeddings.get_embedder(model_name)
except Exception as exc: # noqa: BLE001
# Log the detailed exception server-side; return a generic message
# so internal paths / stack info aren't exposed to the client.
# Log details server-side; return a generic message so paths/stack stay hidden.
logger.warning("RAG warmup failed for %s: %s", model_name, exc)
return {"ok": False, "model": model_name, "error": "Failed to load embedder"}
return {"ok": True, "model": model_name}
@ -532,8 +525,7 @@ def precache_rag_reranker(
try:
precache_reranker()
except Exception as exc: # noqa: BLE001
# Detailed exception logged server-side; client gets a generic
# message so internal paths / stack info aren't exposed.
# Log details server-side; client gets a generic message so paths/stack stay hidden.
logger.warning(
"RAG reranker precache failed",
model = RAG_RERANKER_MODEL,
@ -590,8 +582,7 @@ class UpdateThreadRagSettingsRequest(BaseModel):
chunking_strategy: ChunkingStrategy | None = None
mode: KBMode | None = None
embedding_model: str | None = None
# Only consulted by reingest (not persisted as a thread setting); omit or
# None keeps captioning on.
# Reingest-only (not persisted); omit or None keeps captioning on.
caption_images: bool | None = None
@ -697,7 +688,7 @@ def _reingest_scope(
"SELECT id, stored_path FROM rag_documents WHERE thread_id = ?",
(thread_id,),
).fetchall()
# Drop rag_documents (chunks cascade); files on disk are reused below.
# Drop rag_documents (chunks cascade); disk files reused below.
doc_ids = [r["id"] for r in rows]
if doc_ids:
placeholders = ",".join("?" for _ in doc_ids)
@ -1168,10 +1159,9 @@ class LocatorBackfillResponse(BaseModel):
pagesRefreshed: int
# Extension allowlist for inline rendering / disposition. Anything not in
# this map collapses to ("application/octet-stream", attachment, "unknown").
# .html / .htm intentionally serve as text/plain attachment (decisions Q7 +
# Risk #3) so an uploaded HTML cannot execute in the app origin.
# Extension allowlist for inline rendering / disposition. Unlisted ext collapses
# to ("application/octet-stream", attachment, "unknown"). .html / .htm serve as
# text/plain attachment (decisions Q7 + Risk #3) so uploaded HTML can't execute in the app origin.
_PREVIEW_EXT_MAP: dict[str, tuple[str, str, PreviewMediaKind]] = {
".pdf": ("application/pdf", "inline", "pdf"),
".txt": ("text/plain; charset=utf-8", "inline", "text"),
@ -1318,8 +1308,8 @@ def _resolve_document_file_or_404(doc_row: Any, document_id: str) -> Path:
root = rag_uploads_root(),
)
except ValueError as exc:
# Symlink escape / ``..`` / absolute outside root — collapse to
# "file not found" (the auth row exists, the bytes do not).
# Escape (symlink / ``..`` / absolute outside root) collapses to
# "file not found" — the auth row exists, the bytes do not.
logger.warning(
"RAG preview: stored_path escaped uploads root for doc %s: %s",
document_id,
@ -1410,9 +1400,8 @@ def _serve_document_file_row(
headers = range_headers,
)
# Starlette's FileResponse handles ordinary downloads efficiently. We
# still advertise Accept-Ranges so PDF.js can switch to explicit range
# requests via the signed URL path.
# FileResponse handles ordinary downloads; still advertise Accept-Ranges so
# PDF.js can switch to explicit range requests via the signed URL path.
return FileResponse(
path = str(resolved),
media_type = content_type,
@ -1451,8 +1440,7 @@ def get_document_preview_target(
}
if not chunk_id:
# Q2: document-row preview returns metadata only — frontend MUST
# NOT fall back to "first chunk".
# Q2: document-row preview is metadata-only — frontend MUST NOT fall back to "first chunk".
return PreviewTargetResponse(
**base,
chunkId = None,
@ -1469,14 +1457,11 @@ def get_document_preview_target(
pdfRegions = [],
)
# Single connection enforces membership AND fetches the row in one
# query. Splitting this into a separate `chunk_belongs_to_document`
# call would open a second SQLite connection and create a TOCTOU
# window — if the chunk is deleted between the two calls, the data
# fetch returns None and the route 500s on the next attribute access
# (devils-advocate D1.1). The cross-document case still collapses to
# the same 404 the auth helper emits — never 400 (would leak doc
# existence).
# One connection enforces membership AND fetches the row in a single query.
# A separate `chunk_belongs_to_document` call would open a second SQLite
# connection and open a TOCTOU window — if the chunk is deleted between the
# two calls, the fetch returns None and the route 500s (D1.1). Cross-document
# collapses to the same 404 — never 400 (would leak doc existence).
with get_connection() as conn:
chunk_row = conn.execute(
"""

4
studio/backend/storage/studio_db.py
View file

@ -262,8 +262,8 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_rag_documents_thread_id ON rag_documents(thread_id)"
)
# content_hash: sha256 of the uploaded bytes, used to skip re-indexing a
# file that already exists in the same scope (kb_id / thread_id).
# content_hash: sha256 of uploaded bytes; skips re-indexing a file already
# in the same scope (kb_id / thread_id).
rag_documents_columns = {
row[1] for row in conn.execute("PRAGMA table_info(rag_documents)").fetchall()
}

8
studio/backend/tests/test_rag_authorization.py
View file

@ -151,9 +151,9 @@ def test_kb_doc_missing_kb_raises_404(tmp_path, monkeypatch):
with studio_db.get_connection() as conn:
_insert_kb(conn, kb_id, owner = "alice")
_insert_kb_doc(conn, doc_id, kb_id)
# Delete the KB — ON DELETE CASCADE should also drop the doc.
# Delete the KB; cascade drops the doc too.
conn.execute("DELETE FROM rag_knowledge_bases WHERE id = ?", (kb_id,))
# After cascade deletion the doc_id no longer exists → 404.
# doc_id is gone post-cascade → 404.
with pytest.raises(HTTPException) as exc_info:
document_for_subject_or_404(doc_id, "alice")
assert exc_info.value.status_code == 404
@ -177,7 +177,7 @@ def test_thread_doc_nonexistent_thread_raises_404(tmp_path, monkeypatch):
"""A missing thread_id does NOT silently grant access — it must be 404."""
_reset_db(tmp_path, monkeypatch)
doc_id, thread_id = _uid(), _uid()
# Insert doc with a thread_id that has no matching chat_threads row.
# Doc's thread_id has no matching chat_threads row.
with studio_db.get_connection() as conn:
conn.execute(
"""
@ -250,7 +250,7 @@ def test_chunk_belongs_returns_false_for_wrong_doc(tmp_path, monkeypatch):
_insert_kb_doc(conn, doc_a, kb_id, "a.pdf")
_insert_kb_doc(conn, doc_b, kb_id, "b.pdf")
_insert_chunk(conn, chunk_id, doc_a)
# chunk belongs to doc_a — probing with doc_b must return False
# chunk is doc_a's; probing doc_b → False.
assert chunk_belongs_to_document(chunk_id, doc_b) is False

18
studio/backend/tests/test_rag_preview_routes.py
View file

@ -229,7 +229,7 @@ class TestPreviewTarget:
assert resp.status_code == 200
body = resp.json()
# All chunk fields MUST be null — UI must not guess a first chunk.
# All chunk fields null — no first-chunk guess.
assert body["chunkId"] is None
assert body["chunkIndex"] is None
assert body["targetPage"] is None
@ -288,19 +288,19 @@ class TestPreviewTarget:
client = _make_client(app, "alice")
try:
# Probe doc_b with chunk_a (which belongs to doc_a)
# Probe doc_b with chunk_a (belongs to doc_a)
resp = client.get(
f"/api/rag/documents/{doc_b}/preview-target?chunk_id={chunk_a}"
)
finally:
_clear_overrides(app)
# Must be 404, NOT 200 with doc_a's chunk data
# 404, not 200 with doc_a's chunk data
assert resp.status_code == 404
def test_unauthenticated_returns_401(self, app, db_env, monkeypatch):
"""No bearer token → 401."""
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env))
# No override — let the real dependency raise
# No override: let the real dependency raise.
client = TestClient(app, raise_server_exceptions = False)
resp = client.get(f"/api/rag/documents/{_uid()}/preview-target")
assert resp.status_code == 401
@ -419,7 +419,7 @@ class TestFileRoute:
assert resp.status_code == 200
ct = resp.headers.get("content-type", "").lower()
# MUST NOT be text/html — must be text/plain
# Must be text/plain, not text/html.
assert "text/html" not in ct, f"HTML executed inline! content-type={ct}"
assert "text/plain" in ct
disp = resp.headers.get("content-disposition", "").lower()
@ -501,7 +501,7 @@ class TestFileRoute:
_insert_kb(conn, kb_id)
_insert_doc(conn, doc_id, kb_id, str(stored))
# Delete the file after inserting the row
# Delete the file after the row exists.
stored.unlink()
client = _make_client(app, "alice")
@ -521,7 +521,7 @@ class TestFileRoute:
doc_id, kb_id = _uid(), _uid()
uploads = db_env / "rag" / "uploads"
uploads.mkdir(parents = True, exist_ok = True)
# A legitimate-looking path that is outside the RAG uploads root
# A plausible path outside the RAG uploads root.
outside = tmp_path / "etc" / "passwd"
outside.parent.mkdir(parents = True, exist_ok = True)
outside.write_bytes(b"root:x:0:0")
@ -529,7 +529,7 @@ class TestFileRoute:
with studio_db.get_connection() as conn:
_insert_kb(conn, kb_id)
# Insert with stored_path pointing outside root
# stored_path points outside root.
conn.execute(
"INSERT INTO rag_documents "
"(id, kb_id, thread_id, filename, content_type, stored_path, status, "
@ -544,7 +544,7 @@ class TestFileRoute:
finally:
_clear_overrides(app)
# Must NOT serve the file — containment violation must return 404
# Containment violation: 404, never serve the file.
assert resp.status_code == 404
def test_nosniff_and_cache_headers_on_txt_file(self, app, db_env, monkeypatch):

14
studio/backend/tests/test_rag_source_identity.py
View file

@ -60,10 +60,10 @@ def _hit(
def _parse_chunks(xml_output: str) -> list[dict]:
"""Parse <chunk ...> elements from the multi-block tool output."""
chunks = []
# Each block is wrapped in <chunk ...>\n...\n</chunk>; parse directly.
# Blocks are <chunk ...>...</chunk>; parse directly.
for match in re.finditer(r"<chunk\s([^>]*)>", xml_output):
attrs_raw = match.group(1)
# Quick attribute parser for "key="value"" pairs.
# Parse "key="value"" pairs.
attrs: dict = {}
for m in re.finditer(r'(\w+)="([^"]*)"', attrs_raw):
attrs[m.group(1)] = m.group(2)
@ -92,7 +92,7 @@ def test_citation_id_is_sequential_counter_not_uuid():
output = _format_hits_for_llm(hits, start_id = 0)
chunks = _parse_chunks(output)
visible_id = chunks[0]["id"]
# Must be a small integer string, NOT the UUID
# Small integer string, not the UUID.
assert visible_id == "1", f"expected '1' got {visible_id!r}"
assert visible_id != chunk_id
@ -114,7 +114,7 @@ def test_citation_ids_are_globally_sequential_across_calls():
assert chunks2[0]["id"] == "2"
assert chunks2[1]["id"] == "3"
# No id overlap
# No id overlap.
all_ids = {c["id"] for c in chunks1 + chunks2}
assert len(all_ids) == 3
@ -131,7 +131,7 @@ def test_same_filename_docs_have_distinct_document_ids():
output = _format_hits_for_llm(hits)
chunks = _parse_chunks(output)
assert len(chunks) == 2
# Both use the same filename but MUST have distinct document_id values
# Same filename, distinct document_id values.
assert chunks[0]["document_id"] != chunks[1]["document_id"]
assert chunks[0]["document_id"] == doc_a
assert chunks[1]["document_id"] == doc_b
@ -213,10 +213,10 @@ def test_xml_special_chars_in_filename_escaped():
)
]
output = _format_hits_for_llm(hits)
# The output must parse cleanly (no unescaped < or " in attrs)
# Output parses cleanly (no unescaped < or " in attrs).
chunks = _parse_chunks(output)
assert len(chunks) == 1
# source attribute should have the filename escaped
# source attribute has the filename escaped.
source_attr = chunks[0].get("source", "")
assert "<" not in source_attr and '"' not in source_attr

12
studio/backend/utils/datasets/llm_assist.py
View file

@ -101,11 +101,9 @@ def precache_helper_gguf():
else:
logger.warning(f"No GGUF matching variant '{variant}' in {repo}")
# If the repo also ships an mmproj (vision projection), grab it
# so the helper can be used as a vision-language model by the
# RAG captioner path. Preference order: F16 → BF16 → F32 → any.
# Best-effort — the LLM-assist path doesn't need vision, so a
# missing mmproj is fine and only logged.
# Grab an mmproj (vision projection) if the repo ships one, so the helper
# can serve as a VLM for the RAG captioner. Preference: F16 → BF16 → F32 → any.
# Best-effort — LLM-assist doesn't need vision, so a missing mmproj is fine and logged.
mmproj_files = [
f for f in files if "mmproj" in f.lower() and f.endswith(".gguf")
]
@ -155,8 +153,8 @@ def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]:
try:
from core.inference.llama_cpp import LlamaCppBackend
# kill_orphans=False so the helper backend doesn't reap the
# parent's chat-model llama-server while loading itself.
# kill_orphans=False so the helper doesn't reap the parent's chat-model
# llama-server while loading itself.
backend = LlamaCppBackend(kill_orphans = False)
logger.info(f"Loading helper model: {repo} ({variant})")

22
studio/backend/utils/rag/config.py
View file

@ -31,20 +31,18 @@ RAG_EMBEDDING_MODEL: str = (
or "BAAI/bge-small-en-v1.5"
)
# Default embedder per (mode, chunking). (multimodal, late) is unsupported
# and rejected at KB-create time in routes/rag.py.
# Default embedder per (mode, chunking). (multimodal, late) is rejected at
# KB-create time in routes/rag.py.
#
# Text mode is the default; figures from PDFs are captioned at ingest by
# the loaded chat VLM (or a helper gemma-3n fallback) and spliced into
# the page markdown before chunking, so a single 384-d text embedder
# handles all retrieval. Multimodal mode adds image-vector rows on top,
# embedded by Qwen3-VL-Embedding-2B (2 B params, 2048-d, no CLIP text
# cap — full 512-token chunks embed losslessly).
# Text mode is the default: PDF figures are captioned at ingest (chat VLM or
# helper gemma-3n fallback) and spliced into the page markdown before chunking,
# so a single 384-d text embedder handles retrieval. Multimodal adds image-vector
# rows on top via Qwen3-VL-Embedding-2B (2 B, 2048-d, no CLIP text cap — full
# 512-token chunks embed losslessly).
#
# Alternative multimodal embedders left in tree for manual override:
# - "BAAI/BGE-VL-large" — smaller (~400 M / 768-d) but CLIP-family
# with a 77-token text cap; routed via `_BGEVLAdapter` in
# core/rag/embeddings.py.
# Alternative multimodal embedders kept for manual override:
# - "BAAI/BGE-VL-large" — smaller (~400 M / 768-d) but CLIP-family with a
# 77-token text cap; routed via `_BGEVLAdapter` in core/rag/embeddings.py.
RAG_EMBEDDER_MATRIX: dict[tuple[str, str], str] = {
("text", "standard"): "BAAI/bge-small-en-v1.5",
("text", "late"): "nomic-ai/nomic-embed-text-v1.5",

10
studio/frontend/src/__tests__/chat-adapter.test.ts
View file

@ -2,10 +2,10 @@
* Tests for chat-adapter XML parsing — contracts §3 / §4, T3.
*
* Coverage:
* - New XML with document_id + chunk_id attributes → citationId, documentId, backendChunkId populated.
* - Legacy XML without durable IDs → citationId populated, documentId/backendChunkId absent.
* - New XML (document_id + chunk_id) → citationId, documentId, backendChunkId set.
* - Legacy XML (no durable IDs) → citationId set, documentId/backendChunkId absent.
* - Same visible [N] across turns does NOT imply same backendChunkId.
* - Same filename in two chunks → distinct documentId values preserved.
* - Same filename in two chunks → distinct documentId values kept.
* - Missing attributes degrade gracefully — no throw.
* - citationId is always the visible "N" counter, never the UUID.
*/
@ -36,7 +36,7 @@ The margin rose to 18%.
});
it("legacy XML without durable IDs leaves documentId and backendChunkId absent", () => {
// Old XML: no document_id, no chunk_id — hover-only, NOT preview-clickable (Q3).
// Old XML: no document_id/chunk_id — hover-only, not preview-clickable (Q3).
const xml = `
<chunk id="2" source="old-doc.pdf" page="3">
Legacy chunk text.
@ -115,7 +115,7 @@ Legacy chunk text.
});
it("XML entity encoding in source attribute is decoded", () => {
// &amp; should decode to & in the source attribute (decodeXml in parseChunks)
// &amp; should decode to & in source (decodeXml in parseChunks)
const xml = `<chunk id="1" source="report &amp; summary.pdf" document_id="doc-1" chunk_id="ck-1">text</chunk>`;
const parts: ParsedChunk[] = parseChunks(xml);
expect(parts).toHaveLength(1);

36
studio/frontend/src/__tests__/preview-panel.test.tsx
View file

@ -1,13 +1,13 @@
/**
* Tests for preview-panel — HTML/DOCX/unknown must NEVER render inline (T5 / Risk #3).
*
* Acceptance criteria (contracts §5.4, PLAN.md T5, decisions Q7):
* - mediaKind === "pdf" → react-pdf view is mounted (or loading indicator shown).
* - mediaKind === "html" → text-view fallback shown, NO object/embed/iframe with blob URL.
* - mediaKind === "docx" → text-view fallback shown, NO inline rendering.
* - mediaKind === "unknown" → unavailable/download state, NOT inline.
* - mediaKind === "text" → text/snippet view shown.
* - Panel without a target renders nothing or unavailable state.
* Acceptance (contracts §5.4, PLAN.md T5, decisions Q7):
* - "pdf" → react-pdf view mounted (or loading indicator).
* - "html" → text-view fallback, NO object/embed/iframe with blob URL.
* - "docx" → text-view fallback, NO inline rendering.
* - "unknown" → unavailable/download state, NOT inline.
* - "text" → text/snippet view.
* - No target → nothing or unavailable state.
*/
import {
@ -31,9 +31,9 @@ const DOWNLOAD_BUTTON_NAME = /download/i;
const LONG_CONTENT_TEXT = /Some long content/;
// ── Mock preview store ────────────────────────────────────────────────
// The real component uses per-field selectors: usePreviewStore((s) => s.target)
// so the mock must handle the selector pattern.
// vi.hoisted ensures the mock fn is initialised before vi.mock factory runs.
// The component uses per-field selectors (usePreviewStore((s) => s.target)),
// so the mock must handle the selector pattern. vi.hoisted ensures the
// mock fn is initialised before the vi.mock factory runs.
interface MockStoreState {
target: PreviewTarget | null;
@ -61,8 +61,8 @@ let mockState: MockStoreState = {
const { mockAuthFetch, mockUsePreviewStore } = vi.hoisted(() => {
// usePreviewStore is called two ways:
// usePreviewStore((s) => s.field) — selector form (React hook)
// usePreviewStore.getState().close() — outside React (cleanup effect)
// usePreviewStore((s) => s.field) — selector (React hook)
// usePreviewStore.getState().close() — outside React (cleanup)
const fn = vi.fn((selector?: (s: MockStoreState) => unknown) => {
if (typeof selector === "function") {
return selector(mockState);
@ -125,7 +125,7 @@ beforeEach(() => {
return mockState;
},
);
// Restore getState after mockImplementation replaces the fn internals
// Restore getState; mockImplementation replaces the fn internals
mockUsePreviewStore.getState = () => mockState;
// Mock window.matchMedia globally for tests
@ -282,7 +282,7 @@ describe("preview-panel inline rendering safety (contracts §5.4 / Risk #3)", ()
);
rerender(React.createElement(PreviewPanel, { open: false }));
// The useEffect for open=false should have called close()
// open=false useEffect should have called close()
expect(closeFn).toHaveBeenCalled();
});
@ -333,7 +333,7 @@ describe("preview-panel inline rendering safety (contracts §5.4 / Risk #3)", ()
render(React.createElement(PreviewPanel, { open: true }));
// Radix UI Sheet component should render dialog role in mobile viewports
// Radix Sheet renders dialog role in mobile viewports
const dialog = screen.getByRole("dialog");
expect(dialog).toBeInTheDocument();
expect(dialog).toHaveClass("preview-sheet-content");
@ -342,7 +342,7 @@ describe("preview-panel inline rendering safety (contracts §5.4 / Risk #3)", ()
});
// ── Pure-logic: inline allowlist (always green) ───────────────────────
// Uses the real production isInlineBlobAllowed (D1.5 fix: no local copy).
// Uses the real isInlineBlobAllowed (D1.5 fix: no local copy).
describe("inline object URL allowlist (contracts §5.4, pure logic)", () => {
const inlineSafe: PreviewMediaKind[] = ["pdf", "text", "image"];
@ -384,7 +384,7 @@ describe("preview-panel stable scrollbars, sheets, layouts, and downloads", () =
render(React.createElement(PreviewPanel, { open: true }));
// The snippet is rendered in a <pre> element. Check if it has overflow-auto
// Snippet renders in a <pre>; check it has overflow-auto
const preElement = screen.getByText(LONG_CONTENT_TEXT);
expect(preElement).toHaveClass("overflow-auto");
expect(preElement).toHaveClass("flex-1");
@ -510,7 +510,7 @@ describe("PreviewTextView precise highlights matching", () => {
mediaKind: "text",
filename: "notes.txt",
snippet: "...\nAlphanumericDensity123456\n...",
lineStart: 999, // Trigger hasLocator without matching any specific line range
lineStart: 999, // hasLocator true, but matches no line range
});
render(React.createElement(PreviewPanel, { open: true }));

42
studio/frontend/src/__tests__/preview-pdf-smoke.test.tsx
View file

@ -14,9 +14,9 @@ import { beforeEach, describe, expect, it, vi } from "vitest";
const SOURCE_EXCERPT_TEXT = /source excerpt/i;
/** The thumbnail rail also renders mocked `<Page>` elements, so every
* test that targets the main render must scope through the
* `pdf-main-page` wrapper instead of taking the first `pdf-page`. */
/** The thumbnail rail also renders mocked `<Page>` elements, so tests
* targeting the main render scope through the `pdf-main-page` wrapper
* instead of taking the first `pdf-page`. */
async function findMainPdfPage(): Promise<HTMLElement> {
const wrapper = await screen.findByTestId("pdf-main-page");
return within(wrapper).getByTestId("pdf-page");
@ -100,13 +100,13 @@ function target(overrides: Partial<PreviewTarget> = {}): PreviewTarget {
beforeEach(() => {
class ResizeObserverMock implements ResizeObserver {
observe(_target: Element, _options?: ResizeObserverOptions) {
// jsdom has no layout observer; the component only needs the API shape.
// jsdom has no layout observer; only the API shape is needed.
}
unobserve(_target: Element) {
// jsdom has no layout observer; the component only needs the API shape.
// jsdom has no layout observer; only the API shape is needed.
}
disconnect() {
// jsdom has no layout observer; the component only needs the API shape.
// jsdom has no layout observer; only the API shape is needed.
}
}
vi.stubGlobal("ResizeObserver", ResizeObserverMock);
@ -155,9 +155,9 @@ describe("PreviewPdfView smoke", () => {
expect(regionHighlight).toHaveClass("bg-primary/20");
expect(regionHighlight).toHaveClass("ring-primary/60");
// Verify Tailwind v4 light-mode isolation reset wrapper. After the
// thumbnail-rail refactor, the light wrapper lives INSIDE the
// Document and directly wraps the main-page block.
// Verify Tailwind v4 light-mode isolation reset wrapper. Post
// thumbnail-rail refactor it lives INSIDE the Document, directly
// wrapping the main-page block.
const wrapper = screen.getByTestId("pdf-main-page").parentElement;
expect(wrapper).toHaveClass("light");
expect(wrapper).toHaveClass("bg-white");
@ -173,7 +173,7 @@ describe("PreviewPdfView smoke", () => {
);
// Source-excerpt card uses a neutral muted surface (no brand-coloured
// left rail) so it sits inside the panel without visually competing.
// left rail) so it doesn't visually compete in the panel.
const excerptCard = screen.getByText(SOURCE_EXCERPT_TEXT).parentElement;
expect(excerptCard).toHaveClass("border-border/60");
expect(excerptCard).toHaveClass("bg-muted/30");
@ -204,13 +204,13 @@ describe("PreviewPdfView smoke", () => {
resizeCallbacks.push(callback);
}
observe(_target: Element, _options?: ResizeObserverOptions) {
// jsdom has no layout observer; the component only needs the API shape.
// jsdom has no layout observer; only the API shape is needed.
}
unobserve(_target: Element) {
// jsdom has no layout observer; the component only needs the API shape.
// jsdom has no layout observer; only the API shape is needed.
}
disconnect() {
// jsdom has no layout observer; the component only needs the API shape.
// jsdom has no layout observer; only the API shape is needed.
}
}
vi.stubGlobal("ResizeObserver", FakeResizeObserver);
@ -222,11 +222,11 @@ describe("PreviewPdfView smoke", () => {
}),
);
// Initial render sets width synchronously on mount. Let's capture the initial width.
// Initial render sets width synchronously on mount; capture it.
const page = await findMainPdfPage();
const initialWidth = Number(page.getAttribute("data-width"));
// Activate fake timers AFTER finding the elements to avoid findByTestId timeout
// Fake timers AFTER finding elements, to avoid findByTestId timeout
vi.useFakeTimers();
// Set up HTMLDivElement.prototype.clientWidth mock
@ -242,7 +242,7 @@ describe("PreviewPdfView smoke", () => {
configurable: true,
});
// Now trigger resize callback after changing clientWidth
// Trigger resize callback after changing clientWidth
clientWidthValue = 600;
const resizeCallback = resizeCallbacks[0];
if (!resizeCallback) {
@ -250,23 +250,23 @@ describe("PreviewPdfView smoke", () => {
}
const resizeObserver: ResizeObserver = {
observe() {
// The callback under test ignores the observer instance.
// The callback under test ignores the observer arg.
},
unobserve() {
// The callback under test ignores the observer instance.
// The callback under test ignores the observer arg.
},
disconnect() {
// The callback under test ignores the observer instance.
// The callback under test ignores the observer arg.
},
};
resizeCallback([], resizeObserver);
// Width should NOT be updated immediately because of the 100ms debounce
// Width must NOT update immediately (100ms debounce)
expect(Number(getMainPdfPage().getAttribute("data-width"))).toBe(
initialWidth,
);
// Fast-forward time by 100ms to trigger the debounced callback and flush updates
// Advance 100ms to fire the debounced callback and flush updates
act(() => {
vi.advanceTimersByTime(100);
vi.runAllTimers();

16
studio/frontend/src/__tests__/preview-store.test.ts
View file

@ -2,13 +2,13 @@
* Tests for preview-store object URL lifecycle (contracts §5, T4).
*
* Coverage:
* - open() revokes previous object URL before assigning a new one.
* - open() revokes the previous object URL before assigning a new one.
* - close() revokes any live object URL.
* - Opening doc B while doc A is loaded revokes doc A's URL.
* - PDFs use a signed range URL instead of a full blob download.
* - Inline object URLs are created ONLY for safe non-PDF mediaKind (text/image).
* - For unsafe mediaKind (html/docx/unknown) blob fetch is skipped; previewBlobUrl = null.
* - isInlineBlobAllowed pure predicate matches contracts §5.4 allowlist.
* - PDFs use a signed range URL, not a full blob download.
* - Inline object URLs created ONLY for safe non-PDF kinds (text/image).
* - Unsafe kinds (html/docx/unknown) skip blob fetch; previewBlobUrl = null.
* - isInlineBlobAllowed matches the contracts §5.4 allowlist.
* - __previewStoreInternals() verifies module-scoped cleanup.
*/
@ -19,8 +19,8 @@ import type {
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
// ── Mock rag-api BEFORE importing the store ───────────────────────────
// vi.hoisted ensures the mock refs are initialised before vi.mock factory
// runs (vi.mock is hoisted to the top of the file by Vitest's transformer).
// vi.hoisted initialises the mock refs before the vi.mock factory runs
// (Vitest hoists vi.mock to the top of the file).
const {
mockFetchPreviewTarget,
@ -75,7 +75,7 @@ beforeEach(() => {
mockFetchPreviewTarget.mockReset();
mockFetchPreviewFileBlob.mockReset();
mockFetchPreviewFileUrl.mockReset();
// Reset store to idle between tests
// Reset store to idle per test
usePreviewStore.getState().close();
});

4
studio/frontend/src/app/routes/knowledge-bases.tsx
View file

@ -7,8 +7,8 @@ import { useSettingsDialogStore } from "@/features/settings";
import { requireAuth } from "../auth-guards";
import { Route as rootRoute } from "./__root";
// /knowledge-bases is a deep link to the settings modal's Knowledge
// Bases tab. Open it, then redirect home. Mirrors /settings.
// /knowledge-bases deep-links to the settings modal's Knowledge Bases
// tab: open it, then redirect home. Mirrors /settings.
export const Route = createRoute({
getParentRoute: () => rootRoute,
path: "/knowledge-bases",

7
studio/frontend/src/components/assistant-ui/sources.tsx
View file

@ -208,9 +208,8 @@ const DocumentSourceBadge: FC<{ source: DocSourceData }> = ({ source }) => {
const metaParts: string[] = [];
if (source.page) metaParts.push(`page ${source.page}`);
// Preview is clickable IFF both durable IDs are present (contracts
// §4.1 routing rule + Q3). Legacy sources fall through to a
// non-interactive badge with hover-only behavior.
// Clickable IFF both durable IDs are present (contracts §4.1 + Q3);
// legacy sources fall through to a non-interactive hover-only badge.
const isClickable =
source.documentId !== null && source.backendChunkId !== null;
const openPreview = usePreviewStore((s) => s.open);
@ -304,7 +303,7 @@ const SourcesGroup: FC = () => {
const [visibleCount, setVisibleCount] = useState<number | null>(null);
const [expanded, setExpanded] = useState(false);
// Extract source parts (both URL and document) from the message
// Extract URL and document source parts from the message
const sources: SourceData[] = [];
if (message.content) {
for (const part of message.content) {

12
studio/frontend/src/components/assistant-ui/thread.tsx
View file

@ -1153,8 +1153,8 @@ const RagToggle: FC = () => {
const setRagSource = useChatRuntimeStore((s) => s.setRagSource);
const supportsTools = useChatRuntimeStore((s) => s.supportsTools);
// RAG runs through the local search_knowledge_base tool, so it's hidden
// entirely for external providers (whose request path doesn't execute
// local tools) and gated on tool-calling support otherwise.
// for external providers (no local-tool execution) and gated on
// tool-calling support otherwise.
if (parseExternalModelId(checkpoint) !== null) return null;
const disabled = !modelLoaded || !supportsTools;
return (
@ -1243,9 +1243,9 @@ const ToolStatusDisplay: FC = () => {
);
};
// RAG-aware + button: picks doc formats and routes to ingest pipeline.
// RAG-aware + button: picks doc formats and routes to the ingest pipeline.
// A second button picks a whole folder (webkitdirectory) and routes every
// compatible file through the same pipeline (which drains at the configured
// compatible file through the same pipeline (drains at the configured
// parallel-indexing rate).
const RagDocAttachment: FC<{ onSelect: (file: File) => void }> = ({
onSelect,
@ -1274,8 +1274,8 @@ const RagDocAttachment: FC<{ onSelect: (file: File) => void }> = ({
/>
<input
// webkitdirectory isn't in React's input prop types; set it on the
// element directly so the picker selects a folder (returns every
// file recursively, which selectCompatible then filters).
// element directly so the picker selects a folder (returns every file
// recursively, which selectCompatible then filters).
ref={(el) => {
folderInputRef.current = el;
if (el) el.setAttribute("webkitdirectory", "");

18
studio/frontend/src/components/assistant-ui/tool-ui-search-knowledge-base.tsx
View file

@ -19,8 +19,7 @@ import {
} from "./tool-fallback";
export interface ParsedChunk {
/** Visible citation id the model uses inside `[N]` references. Display
* only; never sent to the backend as a chunk_id. */
/** Visible `[N]` citation id. Display only; never sent as a chunk_id. */
id: string;
source: string;
page?: string;
@ -34,13 +33,12 @@ export interface ParsedChunk {
kind?: string;
imageUrl?: string;
text: string;
/** Durable `rag_documents.id`. Carries through when the tool XML
* includes `document_id="..."`. Absent on legacy tool output. */
/** Durable `rag_documents.id` from tool XML `document_id=`. Absent on
* legacy tool output. */
documentId?: string;
/** Durable `rag_chunks.id`. Carries through when the tool XML
* includes `chunk_id="..."`. Absent on legacy tool output. The
* preview routing value sent as `?chunk_id=` to `/preview-target`;
* never the same as the visible `id`. */
/** Durable `rag_chunks.id` from tool XML `chunk_id=`. Absent on legacy
* output. Sent as `?chunk_id=` to `/preview-target`; never the
* visible `id`. */
backendChunkId?: string;
}
@ -95,8 +93,8 @@ export function parseChunks(raw: string): ParsedChunk[] {
return out;
}
/** Fetch a backend image via the bearer-authed `authFetch`, expose it
* as a blob URL for `<img src>`. Cleans up the object URL on unmount. */
/** Fetch a backend image via bearer-authed `authFetch`, expose it as a
* blob URL for `<img src>`. Revokes the object URL on unmount. */
function useAuthedImageUrl(path: string | undefined): string | undefined {
const [url, setUrl] = useState<string | undefined>(undefined);
useEffect(() => {

93
studio/frontend/src/features/chat/api/chat-adapter.ts
View file

@ -382,18 +382,17 @@ interface DocumentSourcePart {
type: "source";
sourceType: "document";
id: string;
/** Display alias of `citationId`. Kept so the existing sources.tsx
* renderer keeps working; new code SHOULD use `citationId`. NEVER
* sent to backend as the durable chunk_id. */
/** Display alias of `citationId`, kept for the existing sources.tsx
* renderer; new code SHOULD use `citationId`. NEVER sent to backend
* as the durable chunk_id. */
chunkId: string;
/** Visible model-citation id (the `[N]` reference). Display only. */
citationId: string;
/** Durable `rag_documents.id` from tool XML `document_id=`. Null
* when the source came from legacy XML lacking the attribute;
* preview routing is gated off in that case. */
/** Durable `rag_documents.id` from tool XML `document_id=`. Null on
* legacy XML lacking it; preview routing is gated off then. */
documentId: string | null;
/** Durable `rag_chunks.id` from tool XML `chunk_id=`. Null on
* legacy XML. Sent as `?chunk_id=` to `/preview-target`. */
/** Durable `rag_chunks.id` from tool XML `chunk_id=`. Null on legacy
* XML. Sent as `?chunk_id=` to `/preview-target`. */
backendChunkId: string | null;
filename: string;
page?: string;
@ -405,10 +404,9 @@ interface DocumentSourcePart {
text: string;
}
/** Pull every `[N]` token the model wrote in its reply.
* Naive: regex over the whole text. False positives (e.g. `[1]` inside a
* code fence or list marker) are tolerated — the worst case is a stray
* badge for an id that exists in the retrieval set. */
/** Pull every `[N]` token the model wrote. Naive whole-text regex;
* false positives (e.g. `[1]` in a code fence) are tolerated — worst
* case is a stray badge for an id already in the retrieval set. */
const CITATION_RE = /\[(\d+)\]/g;
function extractCitedIds(text: string): Set<string> {
const ids = new Set<string>();
@ -482,11 +480,11 @@ function toDocumentSourcePart(
}
/** Build doc-shaped source parts for chunks the model cited.
* `allChunks` is the flat union of every search_knowledge_base tool
* result in this turn (deduped by id). If the model forgets literal
* `[N]` ids, fall back to retrieved chunks so source chips remain
* visible and previewable. Hallucinated `[99]` refs without a matching
* chunk are silently dropped. */
* `allChunks` is the flat, id-deduped union of every
* search_knowledge_base result this turn. If the model omits literal
* `[N]` ids, fall back to retrieved chunks so chips stay visible and
* previewable. Hallucinated `[99]` refs with no matching chunk are
* silently dropped. */
function buildDocumentSourceParts(
allChunks: ParsedChunk[],
citedIds: Set<string>,
@ -1665,19 +1663,18 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
);
}
// Temporary debug toggle: when false, the pre-fetch path is skipped
// entirely so retrieval only happens via the LLM-invoked
// search_knowledge_base tool. Flip back to true to restore the
// always-on grounding for external providers / non-tool models.
// Temporary debug toggle: when false, pre-fetch is skipped so
// retrieval only runs via the LLM-invoked search_knowledge_base
// tool. Flip back to true to restore always-on grounding for
// external providers / non-tool models.
const ragPrefetchEnabled = false;
const ragSource = runtime.ragSource;
const ragToolEnabled = runtime.ragToolEnabled;
// Even when RAG is toggled on, the tool + system-prompt nudge are
// useless if the active scope has no indexed documents — the model
// would call the tool, get back "no chunks", and waste a turn. Do
// a lightweight scope-has-docs check up front and treat the empty
// scope as effectively "off" for this turn.
// Even with RAG on, the tool + prompt nudge are useless if the
// scope has no indexed docs — the model would call the tool, get
// "no chunks", and waste a turn. Do a lightweight scope-has-docs
// check up front and treat an empty scope as "off" for this turn.
let ragScopeHasDocs = false;
if (ragToolEnabled && ragSource.kind !== "off") {
try {
@ -1689,9 +1686,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
ragScopeHasDocs = docs.length > 0;
}
} catch (err) {
// If the doc-list endpoint is unreachable we err on the side
// of letting the tool through — better to attempt retrieval
// and surface an error than to silently skip RAG.
// If the doc-list endpoint is unreachable, let the tool
// through — better to attempt retrieval and surface an error
// than to silently skip RAG.
console.warn("RAG scope-has-docs check failed:", err);
ragScopeHasDocs = true;
}
@ -2346,12 +2343,12 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
externalProvider.enablePromptCaching ?? true,
}
: {}),
// Anthropic-only: pass the cache TTL the user picked in
// Configuration → Provider. Omitted = inherit the default
// 5-minute pool. The backend's `_stream_anthropic` only
// attaches `cache_control.ttl` when the value is one of
// "5m" / "1h" (see external_provider.py near line 1375),
// so unknown values are a no-op end-to-end.
// Anthropic-only: pass the cache TTL picked in
// Configuration → Provider. Omitted = default 5-minute
// pool. The backend's `_stream_anthropic` only attaches
// `cache_control.ttl` for "5m" / "1h" (see
// external_provider.py near line 1375), so unknown values
// are a no-op end-to-end.
...(supportsProviderPromptCacheTtl(
externalProvider.providerType,
) &&
@ -2417,7 +2414,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
? {
enable_tools: true,
enabled_tools: [
// RAG goes first so the model sees it before any other
// RAG first so the model sees it before any other
// tool when scanning the spec list.
...(ragToolPathTaken ? ["search_knowledge_base"] : []),
...(toolsEnabled ? ["web_search"] : []),
@ -2443,10 +2440,10 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
mcp_enabled: mcpEnabledForChat,
auto_heal_tool_calls:
useChatRuntimeStore.getState().autoHealToolCalls,
// With RAG active the model only needs up to 3 retrieval
// calls; cap low so it can't spiral into web search / fetch
// after already answering from the documents. Off-RAG turns
// keep the user's full budget.
// With RAG active the model needs at most 3 retrieval
// calls; cap low so it can't spiral into web search /
// fetch after answering from the documents. Off-RAG
// turns keep the user's full budget.
max_tool_calls_per_message: ragToolPathTaken
? Math.min(
useChatRuntimeStore.getState().maxToolCallsPerMessage,
@ -3046,10 +3043,10 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
);
});
// RAG: flatten chunks across every search_knowledge_base call this
// turn, then emit previewable doc-source chips for cited chunks.
// If the model omits literal [N] ids, show the retrieved chunks so
// the answer still has a visible citation/preview affordance.
// RAG: flatten chunks across every search_knowledge_base call
// this turn, then emit previewable doc-source chips for cited
// chunks. If the model omits literal [N] ids, show the retrieved
// chunks so the answer still has a citation/preview affordance.
const ragChunks = toolCallParts.flatMap((tc) => {
if (tc.toolName !== "search_knowledge_base" || !tc.result) {
return [];
@ -3064,10 +3061,10 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
)
: [];
// SDK's SourceMessagePart only types `sourceType: "url"` with a
// required `url` field. SourcesGroup branches on `sourceType` at
// runtime, so cast the doc-shaped parts through `unknown` rather
// than weakening the helper's strict typing.
// SDK's SourceMessagePart only types `sourceType: "url"` (with a
// required `url`). SourcesGroup branches on `sourceType` at
// runtime, so cast doc-shaped parts through `unknown` rather than
// weakening the helper's strict typing.
const sourceParts = [
...urlSourceParts,
...(documentSourceParts as unknown as typeof urlSourceParts),

25
studio/frontend/src/features/chat/chat-settings-sheet.tsx
View file

@ -331,7 +331,7 @@ function getRightSlotSheetSnapshot(): boolean {
}
function noopRightSlotSheetSubscription(): void {
// No browser media-query subscription is needed during SSR/tests.
// No media-query subscription during SSR/tests.
}
function subscribeRightSlotSheet(callback: () => void): () => void {
@ -519,7 +519,7 @@ export function ChatSettingsPanel({
const ragDefaults = useRagStore((s) => s.defaults);
// Load thread RAG settings once per threadId. Ref-guarded; keep
// `threadSettings` out of deps to avoid post-load update loops.
// `threadSettings` out of deps to avoid update loops.
const threadSettingsLoadedRef = useRef<string | null>(null);
useEffect(() => {
if (
@ -542,7 +542,7 @@ export function ChatSettingsPanel({
const aui = useAui();
// Brand-new chat has no backend thread yet — initialize the local
// assistant-ui thread to mint a remoteId so per-thread RAG settings
// can be saved before the user sends a message or attaches a file.
// save before the user sends a message or attaches a file.
const ensureThreadId = async (): Promise<string | null> => {
const stored = useChatRuntimeStore.getState().activeThreadId;
if (stored) return stored;
@ -1703,9 +1703,9 @@ export function ChatSettingsPanel({
onCheckedChange={(next) => {
setEnableRerank(next);
if (!next) return;
// First flip-on may have to download ~1.1 GB; the
// toast covers the latency so the user doesn't think
// the next query is hung waiting on the reranker.
// First flip-on may download ~1.1 GB; the toast
// covers the latency so the next query doesn't look
// hung waiting on the reranker.
const toastId = toast.loading(
"Preparing reranker (one-time download)…",
);
@ -1970,11 +1970,10 @@ export function ChatSettingsPanel({
</>
);
// Right-slot routing (decision Q6 / Risk #11): the same slot hosts
// EITHER the inference settings panel OR the document-preview panel,
// never both. Preview takes precedence when a target is open. The
// slot widens for preview because PDFs need more real estate than
// a 17rem settings strip.
// Right-slot routing (decision Q6 / Risk #11): the slot hosts EITHER
// the inference settings panel OR the document-preview panel, never
// both. Preview wins when a target is open, and the slot widens for
// it since PDFs need more room than a 17rem settings strip.
const previewActive =
previewTarget !== null ||
previewStatus === "loading" ||
@ -1995,8 +1994,8 @@ export function ChatSettingsPanel({
if (next) {
onOpenChange?.(true);
} else {
// Closing the sheet from outside closes BOTH the preview
// and the settings — single right slot.
// Closing the sheet from outside closes BOTH preview and
// settings — single right slot.
usePreviewStore.getState().close();
onOpenChange?.(false);
}

6
studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts
View file

@ -719,9 +719,9 @@ export function useChatModelRuntime() {
activeNativePathToken: nativePathToken ?? null,
});
// Reset RAG to off on every successful model load so the
// user always opts in explicitly per session. Goes through
// the setter so the persisted toggle in localStorage is
// kept in sync (setState alone would skip the saveBool).
// user opts in explicitly per session. Via the setter so the
// persisted localStorage toggle stays in sync (setState alone
// would skip the saveBool).
useChatRuntimeStore.getState().setRagToolEnabled(false);
// Qwen3/3.5/3.6: apply thinking-mode-specific params after load
if (

52
studio/frontend/src/features/chat/hooks/use-thread-doc-uploads.ts
View file

@ -48,23 +48,23 @@ export interface UseThreadDocUploadsResult {
isIndexing: boolean;
}
/** RAG upload from the composer "+" button. Routes to whichever source is
* currently selected in the Retrieval dropdown: KB → uploads to that KB;
* thread or off → uploads to (and lazy-initializes) the current chat
* thread, then flips source to "thread" on first ingest if it was "off". */
/** RAG upload from the composer "+" button. Routes by the Retrieval
* dropdown source: KB → that KB; thread/off → the current chat thread
* (lazy-initialized), flipping source to "thread" on first ingest if
* it was "off". */
export function useThreadDocUploads(): UseThreadDocUploadsResult {
const aui = useAui();
const activeThreadId = useChatRuntimeStore((s) => s.activeThreadId);
const [pendingDocs, setPendingDocs] = useState<PendingDoc[]>([]);
// Track the scope key per chip so removeDoc can dispatch the right
// delete (KB docs vs thread docs use different scope keys in the store).
// Track scope key per chip so removeDoc dispatches the right delete
// (KB vs thread docs use different scope keys in the store).
const [chipScopeKeys, setChipScopeKeys] = useState<Record<string, string>>(
{},
);
// Brand-new chats have no backend thread until the first message is sent.
// Initialize the current local thread to mint a remoteId so RAG uploads
// can attach before-send; the saved thread row gets created on first send.
// Brand-new chats have no backend thread until the first message.
// Initialize the local thread to mint a remoteId so RAG uploads can
// attach before-send; the saved row is created on first send.
const ensureThreadId = useCallback(async (): Promise<string | null> => {
const stored = useChatRuntimeStore.getState().activeThreadId;
if (stored) return stored;
@ -86,9 +86,9 @@ export function useThreadDocUploads(): UseThreadDocUploadsResult {
const addDoc = useCallback(
(file: File) => {
const localChipId = crypto.randomUUID();
// Lifecycle state shared between the upload flow and the cancel thunk.
// The cancel thunk closes over these `let`s by reference, so it always
// sees the latest job/document ids no matter when the user cancels.
// Lifecycle state shared by the upload flow and cancel thunk. The
// thunk closes over these `let`s by reference, so it always sees
// the latest job/document ids whenever the user cancels.
const abort = new AbortController();
let jobId: string | undefined;
let documentId: string | undefined;
@ -110,9 +110,9 @@ export function useThreadDocUploads(): UseThreadDocUploadsResult {
return rest;
});
};
// Stop the backend job (if started) and delete its document so the
// index resets. Idempotent: both a late in-flight abort and the toast
// cancel can reach here.
// Stop the backend job (if started) and delete its document to
// reset the index. Idempotent: both a late in-flight abort and the
// toast cancel can reach here.
const cleanupBackend = async () => {
if (cleaned) return;
cleaned = true;
@ -128,8 +128,8 @@ export function useThreadDocUploads(): UseThreadDocUploadsResult {
...prev,
{ id: localChipId, file, status: "uploading" },
]);
// Register in the aggregate-progress store now (synchronously, for the
// whole batch) so the single toast counts queued files too.
// Register in the aggregate-progress store now (synchronously, for
// the whole batch) so the single toast counts queued files too.
const indexProgress = useIndexProgressStore.getState();
indexProgress.add(localChipId, file.name);
indexProgress.setCancel(localChipId, async () => {
@ -141,7 +141,7 @@ export function useThreadDocUploads(): UseThreadDocUploadsResult {
});
void (async () => {
// Hold an indexing slot for this document's whole lifecycle so bulk /
// Hold an indexing slot for this doc's whole lifecycle so bulk /
// folder uploads drain at the configured concurrency instead of
// spawning every ingestion at once. Released on every terminal path.
await acquireIndexSlot();
@ -160,7 +160,7 @@ export function useThreadDocUploads(): UseThreadDocUploadsResult {
scope = { kind: "kb", kbId: ragSource.kbId };
scopeKey = `kb:${ragSource.kbId}`;
} else {
// ragSource is "thread" or "off" — fall back to thread.
// ragSource is "thread" or "off" — use thread.
const threadId = await ensureThreadId();
if (abort.signal.aborted) {
releaseSlot();
@ -199,8 +199,8 @@ export function useThreadDocUploads(): UseThreadDocUploadsResult {
documentId = did;
jobId = jid;
if (abort.signal.aborted) {
// Cancelled while the upload was in flight: the document now
// exists on the backend, so tear it down here.
// Cancelled mid-upload: the doc now exists on the backend,
// so tear it down here.
releaseSlot();
await cleanupBackend();
removeChip();
@ -208,9 +208,9 @@ export function useThreadDocUploads(): UseThreadDocUploadsResult {
}
if (alreadyIndexed) {
// Identical file already in this scope — no re-index. If a
// chip for this document already exists, drop the one we just
// added so the composer doesn't show the same doc twice;
// otherwise mark this chip ready.
// chip for this doc already exists, drop the one we just
// added so the composer doesn't show it twice; otherwise mark
// this chip ready.
setPendingDocs((prev) => {
const dupExists = prev.some(
(d) => d.id !== localChipId && d.documentId === did,
@ -253,8 +253,8 @@ export function useThreadDocUploads(): UseThreadDocUploadsResult {
),
);
// First ingest in an off-source chat: flip to thread so
// the model has somewhere to search. KB uploads don't
// need this — source is already a KB.
// the model has somewhere to search. KB uploads skip this
// — source is already a KB.
if (
scope?.kind === "thread" &&
useChatRuntimeStore.getState().ragSource.kind === "off"

33
studio/frontend/src/features/chat/shared-composer.tsx
View file

@ -520,10 +520,9 @@ export function SharedComposer({
// Images pill is only ever lit on OpenAI cloud's Responses-API models
// and Gemini Nano Banana family. No local tool runtime fallback.
const showImagePill = supportsBuiltinImageGeneration;
// RAG retrieval runs entirely through the local search_knowledge_base
// tool, so it needs the tool-calling loop. No external-builtin RAG
// equivalent — gate purely on supportsTools (mirrors web/code when not
// backed by a provider builtin).
// RAG runs entirely through the local search_knowledge_base tool, so
// it needs the tool-calling loop. No external-builtin equivalent —
// gate purely on supportsTools (mirrors web/code without a builtin).
const ragDisabled = !modelLoaded || !supportsTools;
// Fetch pill: Anthropic-only (web_fetch_20250910 / web_fetch_20260209).
const webFetchDisabled = !modelLoaded || !supportsBuiltinWebFetch;
@ -581,13 +580,13 @@ export function SharedComposer({
}, [aui]);
// Composer "+" upload — routes to whichever scope the Retrieval
// dropdown currently points at (KB or thread). Keeps "what you see is
// what you upload to" so users don't get silent thread-vs-KB mismatches.
// dropdown points at (KB or thread). "What you see is what you upload
// to" avoids silent thread-vs-KB mismatches.
const addDoc = useCallback(
(file: File) => {
const localChipId = crypto.randomUUID();
// Lifecycle state shared between the upload flow and the cancel thunk;
// the thunk closes over these `let`s so it sees the latest ids whenever
// Lifecycle state shared by the upload flow and cancel thunk; the
// thunk closes over these `let`s so it sees the latest ids whenever
// the user cancels.
const abort = new AbortController();
let jobId: string | undefined;
@ -621,8 +620,8 @@ export function SharedComposer({
...prev,
{ id: localChipId, file, status: "uploading" },
]);
// Register in the aggregate-progress store now (whole batch) so the
// single toast counts queued files too.
// Register in the aggregate-progress store now (whole batch) so
// the single toast counts queued files too.
const indexProgress = useIndexProgressStore.getState();
indexProgress.add(localChipId, file.name);
indexProgress.setCancel(localChipId, async () => {
@ -633,9 +632,9 @@ export function SharedComposer({
removeChip();
});
void (async () => {
// Hold an indexing slot for the document's whole lifecycle so bulk /
// folder uploads drain at the configured concurrency. Released on
// every terminal path below.
// Hold an indexing slot for the doc's whole lifecycle so bulk /
// folder uploads drain at the configured concurrency. Released
// on every terminal path below.
await acquireIndexSlot();
slotAcquired = true;
if (abort.signal.aborted) {
@ -689,16 +688,16 @@ export function SharedComposer({
documentId = did;
jobId = jid;
if (abort.signal.aborted) {
// Cancelled while uploading: the document now exists on the
// backend, so tear it down here.
// Cancelled mid-upload: the doc now exists on the backend,
// so tear it down here.
releaseSlot();
await cleanupBackend();
removeChip();
return;
}
if (alreadyIndexed) {
// Drop the just-added chip if this doc is already represented
// so the composer never shows the same document twice.
// Drop the just-added chip if this doc is already shown, so
// the composer never lists the same document twice.
setPendingDocs((prev) => {
const dupExists = prev.some(
(d) => d.id !== localChipId && d.documentId === did,

20
studio/frontend/src/features/chat/stores/chat-runtime-store.ts
View file

@ -348,11 +348,11 @@ type ChatRuntimeStore = {
// Cosine floor; 0 disables. Set > 0 to drop off-topic hits.
ragMinScore: number;
// Max documents indexed in parallel (bulk/folder uploads drain at this
// rate). 1 = sequential. Keeps many concurrent ingestion subprocesses
// from thrashing the GPU/CPU.
// rate). 1 = sequential. Keeps concurrent ingestion subprocesses from
// thrashing the GPU/CPU.
ragIndexConcurrency: number;
// Caption figures/images during ingestion (default on). Off skips the VLM
// captioning pass for faster, text-only indexing.
// Caption figures/images during ingestion (default on). Off skips the
// VLM captioning pass for faster, text-only indexing.
ragCaptionImages: boolean;
hydratePersistedSettings: () => Promise<void>;
setModelLoading: (loading: boolean) => void;
@ -716,9 +716,9 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
),
...getHydratedSettingsState(settings, state, hydrationVersions),
};
// After hydration, if RAG is explicitly on (persisted), warm
// the embedder so the first message doesn't pay the cold load
// inline. RAG is opt-in by default — no auto-enable migration.
// After hydration, if RAG is persisted on, warm the embedder
// so the first message doesn't pay the cold load inline. RAG
// is opt-in by default — no auto-enable migration.
if (
nextState.ragToolEnabled === true ||
(nextState.ragToolEnabled === undefined && state.ragToolEnabled)
@ -995,9 +995,9 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
setRagToolEnabled: (ragToolEnabled) =>
set((state) => {
saveBool(CHAT_RAG_TOOL_ENABLED_KEY, ragToolEnabled);
// Warmup on off→on transitions: kick the backend to preload the
// embedder so the user's first RAG-using message doesn't pay the
// cold-start (~30s for Qwen3-VL-Embedding-2B) inline. Fire-and-forget.
// Warmup on off→on: preload the embedder so the first RAG-using
// message doesn't pay the cold-start (~30s for
// Qwen3-VL-Embedding-2B) inline. Fire-and-forget.
if (ragToolEnabled && !state.ragToolEnabled) {
void import("@/features/rag/api/rag-api")
.then((m) => m.warmupRagEmbedder())

12
studio/frontend/src/features/chat/utils/rag-index-queue.ts
View file

@ -3,12 +3,12 @@
import { useChatRuntimeStore } from "../stores/chat-runtime-store";
/** Bounds how many documents index in parallel. Each RAG upload acquires a
* slot before it starts and releases it once its ingestion job finishes
* (complete/error/already-indexed), so bulk/folder uploads drain at the
* user-configured `ragIndexConcurrency` rate instead of spawning a
* subprocess per file all at once. Module-scoped singleton — shared across
* both composer surfaces. */
/** Bounds how many documents index in parallel. Each RAG upload
* acquires a slot before it starts and releases it once its ingestion
* job finishes (complete/error/already-indexed), so bulk/folder uploads
* drain at the user-configured `ragIndexConcurrency` rate instead of
* spawning a subprocess per file at once. Module-scoped singleton,
* shared across both composer surfaces. */
let active = 0;
const waiters: Array<() => void> = [];

47
studio/frontend/src/features/rag/api/rag-api.ts
View file

@ -36,8 +36,7 @@ export interface UploadResponse {
document_id: string;
job_id: string;
filename: string;
/** True when an identical file was already indexed in this scope; no
* new ingestion job was started and job_id is "". */
/** Identical file already indexed in this scope; no job started, job_id is "". */
already_indexed?: boolean;
}
@ -396,16 +395,15 @@ export async function setRagDefaults(
return parseJsonOrThrow<RagDefaults>(response);
}
/** Preload the configured embedder on the backend. Long-running (cold
* load can take 30s+). Fire-and-forget: failure is non-fatal because
* the first real query will lazy-load again. */
/** Preload the configured embedder on the backend. Long-running (cold load
* 30s+). Fire-and-forget: failure is fine, the first query lazy-loads. */
export async function warmupRagEmbedder(): Promise<void> {
await authFetch("/api/rag/warmup", { method: "POST" });
}
/** Download reranker weights into the HF cache. ~1.1 GB on first call;
* no-op when cached. Called when the user flips the reranker toggle so
* the download lands on an explicit action, not the first chat turn. */
/** Download reranker weights into the HF cache. ~1.1 GB on first call,
* no-op when cached. Triggered by the reranker toggle so the download
* happens on an explicit action, not the first chat turn. */
export async function precacheRagReranker(): Promise<{
ok: boolean;
model: string;
@ -433,17 +431,17 @@ export async function search(req: SearchRequest): Promise<SearchHit[]> {
// --- Ingestion SSE ---
/** Cancel an in-flight ingestion job. Best-effort: a 404/already-terminal job
* resolves without error so batch cancellation never throws on stale ids. */
/** Cancel an in-flight ingestion job. Best-effort: a 404/terminal job
* resolves without error so batch cancel never throws on stale ids. */
export async function cancelJob(jobId: string): Promise<void> {
await authFetch(`/api/rag/jobs/${encodeURIComponent(jobId)}/cancel`, {
method: "POST",
}).catch(() => {});
}
/** Subscribe to a job's SSE stream; returns an unsubscribe fn.
* Use the EventSource polyfill so the bearer token rides in an
* Authorization header instead of leaking through URL query params. */
/** Subscribe to a job's SSE stream; returns an unsubscribe fn. Uses the
* EventSource polyfill so the bearer token rides in an Authorization
* header instead of leaking through URL query params. */
export function subscribeToJobEvents(
jobId: string,
handlers: {
@ -483,7 +481,7 @@ export function subscribeToJobEvents(
};
source.onerror = () => {
// Browser auto-reconnects unless we close; let the consumer decide instead.
// Browser auto-reconnects unless we close; let the consumer decide.
source.close();
handlers.onError?.(new Error("SSE connection lost"));
handlers.onClose?.();
@ -497,9 +495,9 @@ export function subscribeToJobEvents(
// --- Preview target + blob ---
/** Fetch the preview-target metadata for a document. When `chunkId` is
* provided it must belong to `documentId`; mismatch collapses to 404
* with the same "Document not found" body as missing/unauthorized. */
/** Fetch preview-target metadata for a document. A provided `chunkId` must
* belong to `documentId`; mismatch collapses to the same 404 / "Document
* not found" body as missing/unauthorized. */
export async function fetchPreviewTarget(
documentId: string,
chunkId?: string | null,
@ -511,8 +509,8 @@ export async function fetchPreviewTarget(
return parseJsonOrThrow<PreviewTarget>(response);
}
/** Mint a short-lived URL that PDF.js can range-load without putting
* the user's bearer token in a query string. */
/** Mint a short-lived URL PDF.js can range-load without putting the
* user's bearer token in a query string. */
export async function fetchPreviewFileUrl(
documentId: string,
signal?: AbortSignal,
@ -538,13 +536,10 @@ export async function backfillDocumentLocators(
return parseJsonOrThrow<LocatorBackfillResult>(response);
}
/** Download the original uploaded file as a Blob via `authFetch` (so
* the bearer token rides in the Authorization header — never a query
* string). The caller (preview-store) creates and revokes the object
* URL so blob lifecycle stays in one place.
*
* `signal` lets the caller abort the fetch when the user switches
* documents mid-load. */
/** Download the original uploaded file as a Blob via `authFetch` (bearer
* token in the Authorization header, never a query string). The caller
* (preview-store) creates and revokes the object URL so blob lifecycle
* stays in one place. `signal` aborts the fetch on document switch. */
export async function fetchPreviewFileBlob(
documentId: string,
signal?: AbortSignal,

16
studio/frontend/src/features/rag/components/document-row.tsx
View file

@ -34,9 +34,8 @@ export function DocumentRow({
}: {
doc: RagDocument;
onDelete?: () => void;
/** Fired when the row body (not the delete button) is clicked.
* Per decision Q9, callers should only pass this for completed
* documents. */
/** Fired when the row body (not the delete button) is clicked. Per
* decision Q9, callers pass this only for completed documents. */
onPreview?: () => void;
rightSlot?: React.ReactNode;
className?: string;
@ -45,10 +44,9 @@ export function DocumentRow({
const handleRowClick = (e: MouseEvent<HTMLDivElement>) => {
if (!onPreview) return;
// If a button/anchor/control was clicked (e.g. the delete icon),
// skip preview — let that handler win. Buttons inside this row
// additionally call stopPropagation, but this is defense in depth
// for any descendant Button that forgets to.
// Skip preview if a button/anchor/control was clicked (e.g. delete) —
// let that handler win. Those buttons also call stopPropagation; this
// is defense in depth for any descendant Button that forgets to.
const target = e.target as HTMLElement | null;
const interactive = target?.closest("button, a, [role=button]");
if (interactive && interactive !== e.currentTarget) return;
@ -108,8 +106,8 @@ export function DocumentRow({
size="icon"
aria-label="Delete document"
onClick={(e) => {
// Stop propagation so the row's onClick (preview open)
// does not fire when the user is asking to delete.
// Stop propagation so the row's preview-open onClick
// doesn't fire on delete.
e.stopPropagation();
onDelete();
}}

18
studio/frontend/src/features/rag/components/ingestion-toast-stack.tsx
View file

@ -9,9 +9,9 @@ import { AnimatePresence, motion, useReducedMotion } from "motion/react";
import { useEffect, useRef, useState } from "react";
import { useIndexProgressStore } from "../stores/index-progress-store";
/** Single aggregate indexing toast (top-right). One entry per upload batch:
* "Indexing documents · 2/5 · 40%" while in flight, "RAG index ready" when
* the last document finishes. Replaces the old one-toast-per-file stack. */
/** Single aggregate indexing toast (top-right), one per upload batch:
* "Indexing documents · 2/5 · 40%" in flight, "RAG index ready" when the
* last document finishes. Replaces the old one-toast-per-file stack. */
const DISMISS_DELAY_MS = 4000;
@ -40,8 +40,8 @@ export function IngestionToastStack() {
const errored = items.filter((e) => e.status === "error").length;
const totalChunks = items.reduce((sum, e) => sum + (e.chunks || 0), 0);
const allDone = total > 0 && done === total;
// Overall progress: completed/errored files count as 1, in-flight files
// contribute their fractional progress. Smooth even for a handful of files.
// Overall progress: ready/errored files count as 1, in-flight files add
// their fraction. Smooth even for a handful of files.
const overall =
total === 0
? 0
@ -52,8 +52,8 @@ export function IngestionToastStack() {
) / total;
const pct = Math.round(overall * 100);
// Auto-dismiss once the whole batch is terminal; cancel if a new upload
// re-opens the batch (entries change back to not-all-done).
// Auto-dismiss once the batch is terminal; cancel the timer if a new
// upload re-opens it (entries change back to not-all-done).
useEffect(() => {
if (allDone) {
if (dismissTimerRef.current === null) {
@ -90,8 +90,8 @@ export function IngestionToastStack() {
`${totalChunks} chunk${totalChunks === 1 ? "" : "s"} indexed` +
(errored > 0 ? ` · ${errored} failed` : "");
} else {
// Show the file currently being worked on (1-based), not the completed
// count — so a fresh batch reads "1/8" rather than "0/8".
// Show the in-progress file (1-based), not the completed count, so a
// fresh batch reads "1/8" rather than "0/8".
const current = Math.min(total, done + 1);
subtitle = `${current}/${total} · ${pct}%`;
}

2
studio/frontend/src/features/rag/components/kb-create-dialog.tsx
View file

@ -81,7 +81,7 @@ export function KBCreateDialog({
setSubmitting(false);
};
// Forbid (multimodal, late); disable the other side when one is picked.
// Forbid (multimodal, late): disable each side when the other is picked.
const lateDisabled = mode === "multimodal";
const multimodalDisabled = chunkingStrategy === "late";

14
studio/frontend/src/features/rag/components/kb-detail-panel.tsx
View file

@ -45,10 +45,10 @@ export function KBDetailPanel({
const [reconfigureOpen, setReconfigureOpen] = useState(false);
const openPreview = usePreviewStore((s) => s.open);
// Route uploads through the shared aggregate indexing toast (same as the
// chat composer) so a multi-file upload shows ONE progress bar instead of a
// per-row spinner. The doc list refreshes via uploadDocument's own job
// subscription; here we only drive the toast + concurrency semaphore.
// Route uploads through the shared aggregate indexing toast (like the chat
// composer) so a multi-file upload shows ONE progress bar, not a per-row
// spinner. uploadDocument's own job subscription refreshes the doc list;
// here we only drive the toast + concurrency semaphore.
const handleFiles = (files: File[]) => {
const indexProgress = useIndexProgressStore.getState();
const captionImages = useChatRuntimeStore.getState().ragCaptionImages;
@ -170,9 +170,9 @@ export function KBDetailPanel({
<div
key={doc.id}
className={cn(
// grid minmax(0,1fr)/auto (same as the Connections
// rows) so the name column shrinks and the delete
// button stays put — never widening the panel.
// grid minmax(0,1fr)/auto (like the Connections rows)
// so the name column shrinks and the delete button
// stays put — never widening the panel.
"grid w-full grid-cols-[minmax(0,1fr)_auto] items-center gap-2 rounded-lg border border-foreground/20 bg-muted px-3 py-1.5 text-xs",
previewable && "cursor-pointer hover:bg-muted/70",
)}

29
studio/frontend/src/features/rag/components/preview-panel.tsx
View file

@ -27,10 +27,9 @@ import { PreviewTextView } from "./preview-text-view";
import { PreviewUnavailable } from "./preview-unavailable";
interface PreviewPanelProps {
/** Whether the panel is currently being shown in its host slot.
* When the host hides the slot (e.g. the user closes both
* settings and preview from the chat header), we run the close
* side-effect so the blob URL is revoked. */
/** Whether the panel is shown in its host slot. When the host hides
* the slot (e.g. closing both settings and preview from the chat
* header), we run the close side-effect so the blob URL is revoked. */
open: boolean;
disableDrawer?: boolean;
}
@ -89,8 +88,8 @@ function renderPreviewBody({
}
if (status === "error") {
// Treat 404s as "document missing" so a stale citation reads
// like "no longer available" rather than a generic error.
// Treat 404s as "document missing" so a stale citation reads as
// "no longer available" rather than a generic error.
const isMissing = (error ?? "").toLowerCase().includes("not found");
return (
<PreviewUnavailable
@ -111,18 +110,18 @@ function renderPreviewBody({
return <PreviewPdfView target={target} file={pdfFile} />;
}
// text / image / docx / html / unknown — all routed through
// text-view. text gets the snippet rendered inline; docx/html
// /unknown skip inline-render entirely (contracts §5.4 + Risk #3).
// text/image/docx/html/unknown all route through text-view: text
// renders the snippet inline; docx/html/unknown skip inline-render
// entirely (contracts §5.4 + Risk #3).
return <PreviewTextView target={target} />;
}
return null;
}
/** Body-only renderer. The host slot (desktop aside or mobile sheet)
* is owned by the host (chat-settings panel slot, kb-detail panel,
* etc.) — this component is purely the content of the right slot. */
/** Body-only renderer. The host owns the slot (desktop aside or mobile
* sheet — chat-settings panel slot, kb-detail panel, etc.); this is
* purely the content of the right slot. */
export const PreviewPanel: FC<PreviewPanelProps> = ({
open,
disableDrawer = false,
@ -135,8 +134,8 @@ export const PreviewPanel: FC<PreviewPanelProps> = ({
const close = usePreviewStore((s) => s.close);
const isSqueezed = useIsViewportSqueezed() && !disableDrawer;
// Unmount + visibility cleanup: when the panel is hidden or
// unmounted, revoke the live object URL (contracts §5.5).
// Unmount + visibility cleanup: revoke the live object URL when the
// panel is hidden or unmounted (contracts §5.5).
useEffect(() => {
if (!open) {
close();
@ -144,7 +143,7 @@ export const PreviewPanel: FC<PreviewPanelProps> = ({
}, [open, close]);
useEffect(() => {
return () => {
// Component truly unmounting (e.g. navigation away). Cleanup
// Component truly unmounting (e.g. navigation away); clean up
// anything still live.
usePreviewStore.getState().close();
};

33
studio/frontend/src/features/rag/components/preview-pdf-view.tsx
View file

@ -31,10 +31,10 @@ import "react-pdf/dist/Page/TextLayer.css";
import type { PreviewPdfRegion, PreviewTarget } from "../api/rag-api";
import { PreviewUnavailable } from "./preview-unavailable";
// Configure pdfjs worker in the same module where react-pdf is used,
// per the react-pdf README. `import.meta.url` resolves to the JS bundle
// containing this module, and Vite (+ Tauri) rewrites the URL during
// build so the worker is co-located with the chunk.
// Configure the pdfjs worker in the same module as react-pdf, per its
// README. `import.meta.url` resolves to this module's JS bundle, and
// Vite (+ Tauri) rewrites the URL at build so the worker sits with the
// chunk.
pdfjs.GlobalWorkerOptions.workerSrc = new URL(
"pdfjs-dist/build/pdf.worker.min.mjs",
import.meta.url,
@ -52,8 +52,8 @@ type PdfLightThemeStyle = CSSProperties & Record<`--${string}`, string>;
const RESIZE_DEBOUNCE_MS = 100;
const MIN_PDF_WIDTH = 280;
// Body has p-2 (8px each side) + stable scrollbar gutter (~10px) + a tiny
// breathing margin so the page render doesn't kiss the scrollbar.
// Body p-2 (8px each side) + stable scrollbar gutter (~10px) + a small
// margin so the page render doesn't kiss the scrollbar.
const PDF_BODY_GUTTER_PX = 28;
const PDF_THUMBNAIL_WIDTH = 64;
@ -102,8 +102,8 @@ function markFirstMatch(text: string, needle: string): string | null {
)}</mark>${escapeHtml(text.slice(end))}`;
}
// Keep text-layer highlighting opt-in. Citation snippets render in the card
// below; using them here would mark common words across unrelated PDF text.
// Keep text-layer highlighting opt-in. Citation snippets render in the
// card below; using them here would mark common words across the PDF.
function highlightPdfText(text: string, searchTerm: string): string {
const trimmed = searchTerm.trim();
if (trimmed.length < 2) {
@ -132,9 +132,9 @@ interface PdfThumbnailProps {
onSelect: (pageNumber: number) => void;
}
/** Lazy thumbnail rendered via IntersectionObserver — only mounts the
* inner <Page> when scrolled into view (or close to it), so large PDFs
* stay responsive even when the rail caps at 80 buttons. */
/** Lazy thumbnail via IntersectionObserver — mounts the inner <Page>
* only when scrolled near view, so large PDFs stay responsive even
* with the rail capped at 80 buttons. */
const PdfThumbnail: FC<PdfThumbnailProps> = ({
pageNumber,
active,
@ -295,12 +295,11 @@ export const PreviewPdfView: FC<PreviewPdfViewProps> = ({ target, file }) => {
setWidth(next);
}, []);
// Callback ref instead of useRef + mount effect: the scroll container
// lives INSIDE <Document>, so it only enters the DOM after the PDF
// loads. Attaching the ResizeObserver the instant the node mounts
// (rather than on the component's mount effect, when the node is still
// absent) is what keeps the main page from rendering at width 0 — the
// thin white strip regression.
// Callback ref, not useRef + mount effect: the scroll container lives
// INSIDE <Document> and only enters the DOM after the PDF loads.
// Attaching the ResizeObserver the instant the node mounts (vs the
// component mount effect, when the node is still absent) keeps the main
// page from rendering at width 0 — the thin white strip regression.
const attachContainer = useCallback(
(node: HTMLDivElement | null) => {
if (observerRef.current) {

75
studio/frontend/src/features/rag/components/preview-text-view.tsx
View file

@ -12,11 +12,10 @@ interface PreviewTextViewProps {
target: PreviewTarget;
}
/** Fetch the original document bytes via authFetch so the bearer
* token rides in the Authorization header. `window.open(url)` and
* `<a download href=url>` cannot set custom headers, so handing
* them the raw `/file` URL gets a 401 (HTTPBearer-only backend —
* see D1.3). */
/** Fetch the original document bytes via authFetch so the bearer token
* rides in the Authorization header. `window.open(url)` / `<a download>`
* can't set custom headers, so the raw `/file` URL gets a 401 on the
* HTTPBearer-only backend (see D1.3). */
async function fetchOriginalBlob(target: PreviewTarget): Promise<Blob> {
const response = await authFetch(
`/api/rag/documents/${encodeURIComponent(target.documentId)}/file`,
@ -41,16 +40,14 @@ async function downloadOriginal(target: PreviewTarget): Promise<void> {
const blob = await fetchOriginalBlob(target);
const url = URL.createObjectURL(blob);
clickDownloadUrl(url, target.filename);
// Defer revocation so the browser's download pipeline gets the bytes
// before the URL goes away.
// Defer revocation so the download pipeline gets the bytes first.
setTimeout(() => URL.revokeObjectURL(url), 0);
}
async function openOriginalInNewTab(target: PreviewTarget): Promise<void> {
// Defense in depth: refuse to create an inline blob URL for the
// unsafe types even if a future caller forgets the gate. The
// browser would render an html-blob as live HTML in the new tab,
// which is the Risk #3 / contracts §2.3 trip.
// Defense in depth: refuse an inline blob URL for the unsafe types
// even if a caller forgets the gate. An html-blob would render as
// live HTML in the new tab — the Risk #3 / contracts §2.3 trip.
if (!isInlineBlobAllowed(target.mediaKind)) {
throw new Error(
`Inline open not allowed for mediaKind "${target.mediaKind}" — use Download instead.`,
@ -64,15 +61,15 @@ async function openOriginalInNewTab(target: PreviewTarget): Promise<void> {
}
/** Extracted-text / snippet preview used for:
* - `text` mediaKind (txt/md) — the snippet is the only inline
* rendering we trust, and the original is one click away.
* - `docx`, `html`, `unknown` — the original is NEVER rendered
* inline from an object URL (Risk #3); we show the cited chunk
* text plus a safe download/open action.
* - `text` (txt/md) — snippet is the only inline rendering we trust;
* the original is one click away.
* - `docx` / `html` / `unknown` — original is NEVER rendered inline
* from an object URL (Risk #3); show the cited chunk text plus a
* safe download/open action.
*
* When `chunk_id` was not supplied (document-row preview per
* contracts §1.3 + decision Q2), `snippet` is `null` and we show a
* metadata-only state instead of guessing a first chunk. */
* When no `chunk_id` was supplied (document-row preview, contracts
* §1.3 + Q2), `snippet` is null and we show a metadata-only state
* instead of guessing a first chunk. */
interface MatchRange {
start: number;
end: number;
@ -186,15 +183,15 @@ const renderHighlightedSnippet = (
};
/** Extracted-text / snippet preview used for:
* - `text` mediaKind (txt/md) — the snippet is the only inline
* rendering we trust, and the original is one click away.
* - `docx`, `html`, `unknown` — the original is NEVER rendered
* inline from an object URL (Risk #3); we show the cited chunk
* text plus a safe download/open action.
* - `text` (txt/md) — snippet is the only inline rendering we trust;
* the original is one click away.
* - `docx` / `html` / `unknown` — original is NEVER rendered inline
* from an object URL (Risk #3); show the cited chunk text plus a
* safe download/open action.
*
* When `chunk_id` was not supplied (document-row preview per
* contracts §1.3 + decision Q2), `snippet` is `null` and we show a
* metadata-only state instead of guessing a first chunk. */
* When no `chunk_id` was supplied (document-row preview, contracts
* §1.3 + Q2), `snippet` is null and we show a metadata-only state
* instead of guessing a first chunk. */
export const PreviewTextView: FC<PreviewTextViewProps> = ({ target }) => {
const snippet = target.snippet;
const hasSnippet = snippet !== null && snippet.trim().length > 0;
@ -203,15 +200,12 @@ export const PreviewTextView: FC<PreviewTextViewProps> = ({ target }) => {
target.lineEnd !== null ||
target.pageCharStart !== null ||
target.pageCharEnd !== null;
// "Open original" creates a blob: URL of the original bytes and
// passes it to `window.open`. For `html` the new tab would render
// it as live HTML — exactly the Risk #3 / contracts §2.3 trip
// ("MUST refuse to create an inline object URL for mediaKind ==
// 'html' | 'docx' | 'unknown'"). For those types the only safe
// action is Download (backend already sets
// Content-Disposition: attachment for those Content-Types). The
// pdf/text/image allowlist is the same one the preview-store
// uses to decide whether to fetch the blob at all (§5.4). */
// "Open original" blobs the bytes and hands them to `window.open`. For
// html/docx/unknown the new tab would render live HTML — the Risk #3 /
// contracts §2.3 trip ("MUST refuse an inline object URL" for those
// kinds). Download is the only safe action there (backend sends
// Content-Disposition: attachment). The pdf/text/image allowlist is the
// same one preview-store uses to gate the blob fetch (§5.4).
const canOpenInline = isInlineBlobAllowed(target.mediaKind);
const handleDownload = useCallback(() => {
@ -221,11 +215,10 @@ export const PreviewTextView: FC<PreviewTextViewProps> = ({ target }) => {
}, [target]);
const handleOpenExternal = useCallback(() => {
// Re-fetch through authFetch and hand the new tab a blob URL.
// `window.open(rawApiUrl)` would send the request WITHOUT the
// Authorization header (window.open can't set custom headers)
// and the HTTPBearer-protected /file route would respond 401.
// See D1.3 finding.
// Re-fetch via authFetch and hand the new tab a blob URL.
// `window.open(rawApiUrl)` would omit the Authorization header
// (window.open can't set custom headers), so the HTTPBearer-protected
// /file route would 401. See D1.3 finding.
openOriginalInNewTab(target).catch(() => {
// best-effort; the user can retry.
});

10
studio/frontend/src/features/rag/components/preview-unavailable.tsx
View file

@ -7,12 +7,12 @@ import type { FC } from "react";
interface PreviewUnavailableProps {
/** Filename if known; "Document" otherwise. */
filename?: string;
/** One-line reason — pulled from the backend's error body when
* available, otherwise a generic copy. */
/** One-line reason — from the backend's error body if available,
* else a generic copy. */
reason: string;
/** "missing" → deleted/404 case; "error" → other failures. The icon
* + tone change so a stale citation reads as "no longer available"
* rather than a transient blip. */
/** "missing" → deleted/404; "error" → other failures. Icon + tone
* change so a stale citation reads as "no longer available" rather
* than a transient blip. */
variant?: "missing" | "error";
}

2
studio/frontend/src/features/rag/hooks/use-kb-documents.ts
View file

@ -5,7 +5,7 @@ import { useEffect } from "react";
import type { RagDocument } from "../api/rag-api";
import { kbScopeKey, threadScopeKey, useRagStore } from "../stores/rag-store";
// Stable sentinel: inline `[]` in the selector causes React error #185
// Stable sentinel: inline `[]` in the selector triggers React error #185
// (new ref each call → Zustand re-renders → infinite loop).
const EMPTY_DOCS: RagDocument[] = [];

11
studio/frontend/src/features/rag/hooks/use-resizable-width.ts
View file

@ -43,10 +43,10 @@ function readStored(key: string, fallback: number): number {
}
}
/** Drag-to-resize hook for a right-anchored panel. The handle sits on
* the panel's LEFT edge — width grows as the pointer moves toward
* viewport x=0. Persists to localStorage and re-clamps on viewport
* changes so a wide panel cannot eclipse the host content. */
/** Drag-to-resize hook for a right-anchored panel. The handle is on the
* panel's LEFT edge, so width grows as the pointer moves toward x=0.
* Persists to localStorage and re-clamps on viewport resize so a wide
* panel cannot eclipse the host content. */
export function useResizablePanelWidth({
storageKey,
defaultWidth,
@ -78,8 +78,7 @@ export function useResizablePanelWidth({
try {
window.localStorage.setItem(storageKey, String(width));
} catch {
// localStorage may be unavailable (private mode, quota); persist
// is best-effort.
// localStorage may be unavailable (private mode, quota); best-effort.
}
}, [width, storageKey, enabled]);

17
studio/frontend/src/features/rag/stores/index-progress-store.ts
View file

@ -3,11 +3,10 @@
import { create } from "zustand";
/** Tracks every document in the current upload batch so the toast can show
* ONE aggregate "Indexing documents" entry (with overall %) instead of a
* separate toast per file. Unlike the per-job rag-store map, entries are
* registered at addDoc time, so queued-but-not-yet-started files (held by
* the concurrency semaphore) are counted in the denominator too. */
/** Tracks every document in the current upload batch so the toast shows ONE
* aggregate "Indexing documents" entry (overall %) instead of one per file.
* Unlike the per-job rag-store map, entries register at addDoc time, so
* queued files held by the concurrency semaphore count in the denominator. */
export type IndexEntryStatus = "queued" | "indexing" | "ready" | "error";
@ -19,8 +18,8 @@ export interface IndexEntry {
/** Chunks this file produced (from the job's complete event); 0 until done. */
chunks: number;
/** Tear down this upload and remove its document from the index. Registered
* by the upload surface so the aggregate toast can cancel the whole batch
* without owning the per-file job/SSE/semaphore handles. */
* by the upload surface so the toast can cancel the whole batch without
* owning the per-file job/SSE/semaphore handles. */
cancel?: () => Promise<void> | void;
}
@ -64,8 +63,8 @@ export const useIndexProgressStore = create<IndexProgressState>((set, get) => ({
patch(set, id, { status: "ready", progress: 1, chunks }),
setError: (id) => patch(set, id, { status: "error" }),
setCancel: (id, cancel) => patch(set, id, { cancel }),
// Cancel every file in the batch (running, queued, and already-finished) so
// the index returns to its pre-batch state, then drop all toast entries.
// Cancel every file in the batch (running, queued, finished) to restore the
// pre-batch index state, then drop all toast entries.
cancelAll: async () => {
const handles = Object.values(get().entries)
.map((e) => e.cancel)

63
studio/frontend/src/features/rag/stores/preview-store.ts
View file

@ -10,10 +10,10 @@ import {
fetchPreviewTarget,
} from "../api/rag-api";
/** mediaKinds that may safely back an inline object URL (e.g. PDF.js
* worker, plain text, raster image). HTML, DOCX, and unknown are
* forced through the extracted-text fallback per contracts §5.4 +
* Risk #3 (unsafe HTML inline rendering). */
/** mediaKinds that may safely back an inline object URL (PDF.js worker,
* plain text, raster image). HTML/DOCX/unknown go through the
* extracted-text fallback per contracts §5.4 + Risk #3 (unsafe HTML
* inline rendering). */
const INLINE_BLOB_ALLOWLIST: ReadonlySet<PreviewMediaKind> = new Set([
"pdf",
"text",
@ -24,28 +24,28 @@ export function isInlineBlobAllowed(mediaKind: PreviewMediaKind): boolean {
return INLINE_BLOB_ALLOWLIST.has(mediaKind);
}
/** What the panel should mount for the current target. Computed from
* `target.mediaKind` so the panel never has to re-derive it. */
/** What the panel mounts for the current target, derived from
* `target.mediaKind` so the panel never re-derives it. */
export type PreviewLoadStatus = "idle" | "loading" | "ready" | "error";
export interface PreviewRequest {
/** Durable `rag_documents.id`. The only field required to open. */
documentId: string;
/** Durable `rag_chunks.id`. Optional — absence means document-row
* preview (contracts §1.3 + decision Q2: snippet/targetPage stay
* null, no first-chunk fallback). */
/** Durable `rag_chunks.id`. Optional; absence means document-row
* preview (contracts §1.3 + Q2: snippet/targetPage stay null, no
* first-chunk fallback). */
backendChunkId?: string | null;
}
interface PreviewState {
/** Currently-open preview, or null when closed. */
target: PreviewTarget | null;
/** Object URL for the original file blob (PDF / text / image only).
* Null for docx / html / unknown (extracted-text fallback) and
* while the fetch is still in flight. */
/** Object URL for the original file blob (PDF/text/image only). Null
* for docx/html/unknown (extracted-text fallback) and while the
* fetch is in flight. */
previewBlobUrl: string | null;
/** Original fetched file blob for text/image fallback previews. PDFs
* prefer `previewFileUrl` so PDF.js can issue range requests. */
/** Fetched file blob for text/image fallback previews. PDFs prefer
* `previewFileUrl` so PDF.js can issue range requests. */
previewBlob: Blob | null;
/** Short-lived signed URL for PDF.js range requests. */
previewFileUrl: string | null;
@ -54,23 +54,22 @@ interface PreviewState {
status: PreviewLoadStatus;
/** Last error message, if `status === "error"`. */
error: string | null;
/** Open key uniquely identifying the current request — used by tests
* and by consumers that need to react to "the open call changed
* underneath me" (e.g. re-fetch after stale closure). */
/** Key uniquely identifying the current request — used by tests and
* consumers reacting to "the open call changed underneath me" (e.g.
* re-fetch after stale closure). */
openKey: number;
/** Open or replace the current preview. If a previous preview is
* open, its object URL is revoked and its in-flight fetch is
* aborted before the new request begins. */
/** Open or replace the current preview. A previously-open preview's
* object URL is revoked and its in-flight fetch aborted before the
* new request begins. */
open: (req: PreviewRequest) => Promise<void>;
/** Close the current preview. Revokes the object URL and aborts any
/** Close the current preview: revoke the object URL and abort any
* in-flight fetch. Safe to call when nothing is open. */
close: () => void;
}
// State that can't live inside the zustand object without being
// part of the React render cycle. Kept module-scoped because the
// preview store is a singleton.
// State kept out of the zustand object to stay off the React render
// cycle. Module-scoped because the preview store is a singleton.
let activeAbortController: AbortController | null = null;
let activeBlobUrl: string | null = null;
let activeOpenKey = 0;
@ -101,8 +100,8 @@ export const usePreviewStore = create<PreviewState>((set) => ({
openKey: 0,
async open(req) {
// Single-slot invariant (contracts §5.1): tear down whatever was
// there before assigning the new target. revoke → abort → reset.
// Single-slot invariant (contracts §5.1): tear down the previous
// target before assigning the new one. revoke → abort → reset.
revokeActiveBlobUrl();
abortActive();
@ -147,9 +146,9 @@ export const usePreviewStore = create<PreviewState>((set) => ({
return;
}
// For mediaKinds outside the allowlist (docx / html / unknown),
// skip the blob fetch entirely — the panel mounts the
// extracted-text fallback (contracts §5.4 + Risk #3).
// For mediaKinds outside the allowlist (docx/html/unknown), skip the
// blob fetch — the panel mounts the extracted-text fallback
// (contracts §5.4 + Risk #3).
if (!isInlineBlobAllowed(target.mediaKind)) {
if (activeAbortController === controller) activeAbortController = null;
set({
@ -261,9 +260,9 @@ export const usePreviewStore = create<PreviewState>((set) => ({
},
}));
/** Test-only inspector: returns whether the module-scoped blob URL is
* still live. Used by `preview-store.test.ts` to assert
* URL.revokeObjectURL was paired with URL.createObjectURL. */
/** Test-only inspector: whether the module-scoped blob URL is still
* live. Used by `preview-store.test.ts` to assert revokeObjectURL was
* paired with createObjectURL. */
export function __previewStoreInternals(): {
activeBlobUrl: string | null;
hasInflightController: boolean;

22
studio/frontend/src/features/settings/tabs/knowledge-bases-tab.tsx
View file

@ -17,9 +17,8 @@ import { Add01Icon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { type CSSProperties, useState } from "react";
// Height of the master-detail workspace, only rendered once a KB's upload or
// files panel is open (or a preview is active) so the section stays compact
// when just browsing the list.
// Master-detail workspace height, applied only once a KB's upload/files panel
// is open (or a preview is active) so the section stays compact while browsing.
const KB_WORKSPACE_HEIGHT = "h-[360px]";
export function KnowledgeBasesTab() {
@ -41,10 +40,10 @@ export function KnowledgeBasesTab() {
};
const handlePanel = (kb: KnowledgeBase, panel: KBPanel) => {
// A stale preview from a previous selection would otherwise linger and
// force the workspace wider; clear it whenever the panel changes.
// Clear any stale preview from a prior selection; it would otherwise
// linger and force the workspace wider.
closePreview();
// Re-clicking the active KB's active button collapses back to full width.
// Re-clicking the active button collapses back to full width.
if (activeKb?.id === kb.id && activePanel === panel) {
closePanel();
} else {
@ -71,9 +70,8 @@ export function KnowledgeBasesTab() {
});
return (
// pr-2.5 insets right-aligned counts (e.g. "3 total", "60 threads") from
// the scroll container's scrollbar, which overlays the content edge in
// webviews that don't honor scrollbar-gutter.
// pr-2.5 insets right-aligned counts (e.g. "3 total") from the scrollbar,
// which overlays the content edge in webviews ignoring scrollbar-gutter.
<div className="flex min-w-0 flex-col gap-4 pr-2.5">
<div>
<h2 className="text-lg font-semibold">Knowledge bases</h2>
@ -169,9 +167,9 @@ export function KnowledgeBasesTab() {
</button>
<div
className={cn(
// shrink + min-w-0 + a width cap so the preview can never
// push the workspace wider than the dialog (which clipped
// the close button and right-aligned content).
// shrink + min-w-0 + width cap so the preview can't push the
// workspace wider than the dialog (clipping the close button
// and right-aligned content).
"w-0 min-w-0 shrink overflow-hidden max-lg:hidden lg:w-[var(--preview-w)] lg:max-w-[60%]",
!previewResizing && "transition-[width] duration-200 ease-out",
)}

8
studio/install_python_stack.py
View file

@ -1775,11 +1775,9 @@ def install_python_stack() -> int:
[sys.executable, str(SINGLE_ENV / "patch_metadata.py")],
)
# 14. AMD ROCm: final torch repair. Multiple install steps above can
# pull in CUDA torch from PyPI (base packages, extras, overrides,
# studio deps, etc.). Running the repair as the very last step
# ensures ROCm torch is in place at runtime, regardless of which
# intermediate step clobbered it.
# 14. AMD ROCm: final torch repair. Earlier steps can pull in CUDA torch
# from PyPI; running last ensures ROCm torch wins regardless of which
# step clobbered it.
if not IS_WINDOWS and not IS_MACOS and not NO_TORCH:
_progress("ROCm torch (final)")
_ensure_rocm_torch()

4
tests/fixtures/rag-preview/make_fixture_pdf.py vendored
View file

@ -19,8 +19,8 @@ def _compress(data: bytes) -> bytes:
def _pdf() -> bytes:
# Minimal PDF 1.4 with one page, one text stream.
# Structure: header, catalog, pages, page, content stream, xref, trailer.
# Minimal one-page PDF 1.4: header, catalog, pages, page, content
# stream, xref, trailer.
page_text = b"BT /F1 12 Tf 72 720 Td (RAG preview fixture - page 1) Tj ET"
compressed = _compress(page_text)
stream_len = len(compressed)

2
tests/python/test_rag_bm25.py
View file

@ -18,7 +18,7 @@ def isolated_bm25_root(tmp_path, monkeypatch):
from utils.paths import storage_roots
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
# Reset module-level cache between tests.
# Reset module cache between tests.
from core.rag import bm25
bm25._cache.clear()

13
tests/python/test_rag_chunking.py
View file

@ -28,7 +28,7 @@ def test_chunk_pages_splits_long_text():
for chunk in chunks:
assert (
_wc_counter(chunk.text) <= 55
) # max + small slack from atomic split granularity
) # max + slack from atomic split granularity
def test_chunk_pages_short_text_is_one_chunk():
@ -70,7 +70,7 @@ def test_chunk_pages_no_empty_chunks():
def test_chunk_pages_overlap_produces_repeated_tokens():
# Build a list of unique numbered sentences so we can detect overlap.
# Unique numbered sentences let us detect overlap.
sentences = [f"sentence-{i}" for i in range(40)]
text = " ".join(sentences)
chunks = chunk_pages(
@ -82,14 +82,13 @@ def test_chunk_pages_overlap_produces_repeated_tokens():
if len(chunks) >= 2:
first_tail_words = set(chunks[0].text.split()[-4:])
second_head_words = set(chunks[1].text.split()[:4])
# At least one word should appear in both
# At least one shared word.
assert first_tail_words & second_head_words
def test_chunk_pages_splits_on_markdown_headings():
# Phase 3A: heading separators take priority over paragraph breaks
# so chunks start at section boundaries when the parser emits
# Markdown.
# Phase 3A: heading separators outrank paragraph breaks, so chunks
# start at section boundaries for Markdown.
md = (
"# First Section\n\n"
+ "alpha " * 30
@ -104,7 +103,7 @@ def test_chunk_pages_splits_on_markdown_headings():
overlap_tokens = 0,
token_counter = _wc_counter,
)
# We expect multiple chunks and at least one to begin at a heading.
# Expect multiple chunks, at least one starting at a heading.
assert len(chunks) >= 2
starts_at_heading = sum(
1 for c in chunks if c.text.lstrip().startswith(("# ", "## "))

10
tests/python/test_rag_late_chunking.py
View file

@ -37,7 +37,7 @@ def test_spans_index_back_to_full_doc_text():
assert chunks
assert len(chunks) == len(char_spans)
for chunk, (start, end) in zip(chunks, char_spans):
# The chunk text must be exactly the slice of full_doc it claims.
# Chunk text must equal the full_doc slice it claims.
assert full_doc[start:end] == chunk.text
@ -54,7 +54,7 @@ def test_chunks_inherit_page_number_by_overlap():
)
pages_seen = {c.page_number for c in chunks}
assert pages_seen <= {1, 2}
# Both pages should contribute at least one chunk.
# Each page contributes at least one chunk.
assert 1 in pages_seen
assert 2 in pages_seen
@ -72,7 +72,7 @@ def test_full_doc_joins_pages_with_blank_line_separator():
)
assert "first" in full_doc
assert "second" in full_doc
# The two pages must be separated by exactly one blank line.
# Pages separated by exactly one blank line.
assert "first\n\nsecond" in full_doc
@ -80,7 +80,7 @@ def test_full_doc_joins_pages_with_blank_line_separator():
def test_late_chunk_encode_returns_one_vector_per_span():
pytest.importorskip("sentence_transformers")
pytest.importorskip("torch")
# all-MiniLM-L6-v2 is ~80MB and embeds at 384 dims.
# all-MiniLM-L6-v2: ~80MB, 384 dims.
import os
os.environ.setdefault(
@ -100,7 +100,7 @@ def test_late_chunk_encode_returns_one_vector_per_span():
"# Results\n\n"
"Accuracy improved by 12% over the baseline."
)
# char_spans for three chunks — one per section, picked manually.
# char_spans: one per section, picked manually.
char_spans = [
(doc_text.index("The quick"), doc_text.index("\n\n# Methods")),
(doc_text.index("We trained"), doc_text.index("\n\n# Results")),

13
tests/python/test_rag_multimodal.py
View file

@ -24,7 +24,7 @@ def test_html_parser_returns_images_when_requested(tmp_path):
pytest.importorskip("markdownify")
from core.rag.parsers import parse
# A tiny 1x1 transparent PNG.
# 1x1 transparent PNG.
png_bytes = (
b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01"
b"\x08\x06\x00\x00\x00\x1f\x15\xc4\x89\x00\x00\x00\rIDATx\x9cc\xfc\xff"
@ -57,12 +57,12 @@ def test_multimodal_late_combo_validator():
from routes.rag import _validate_mode_combo
# Allowed combos return None.
# Allowed combos → None.
assert _validate_mode_combo("text", "standard") is None
assert _validate_mode_combo("text", "late") is None
assert _validate_mode_combo("multimodal", "standard") is None
# Forbidden combo raises 400.
# Forbidden combo → 400.
with pytest.raises(HTTPException) as excinfo:
_validate_mode_combo("multimodal", "late")
assert excinfo.value.status_code == 400
@ -76,7 +76,7 @@ def test_rag_embedder_matrix_excludes_multimodal_late():
assert ("text", "late") in RAG_EMBEDDER_MATRIX
assert ("multimodal", "standard") in RAG_EMBEDDER_MATRIX
# Unknown combos fall back to the legacy default rather than KeyError.
# Unknown combos fall back to the legacy default, not KeyError.
fallback = resolve_embedder("multimodal", "late")
assert isinstance(fallback, str) and fallback
@ -101,7 +101,7 @@ def test_multimodal_encode_image_returns_vector(tmp_path, monkeypatch):
pytest.importorskip("sentence_transformers")
pytest.importorskip("PIL")
monkeypatch.setenv("UNSLOTH_RAG_EMBEDDING_MODEL", "BAAI/BGE-VL-base")
# Reset the embedder singleton so the env var takes effect.
# Reset the embedder singleton so the env var applies.
from core.rag import embeddings as embeddings_module
embeddings_module._model = None
@ -121,7 +121,6 @@ def test_multimodal_encode_image_returns_vector(tmp_path, monkeypatch):
dim = vectors[0].shape[0]
assert dim > 0
# Text from the same model should also be `dim`-d — shared space is
# the whole point of multimodal embedders.
# Text shares the same dim — the point of a multimodal embedder.
text_vec = embeddings_module.encode(["a red square"])[0]
assert text_vec.shape[0] == dim

22
tests/python/test_rag_multimodal_integration.py
View file

@ -39,14 +39,13 @@ def test_multimodal_subprocess_emits_image_and_caption_chunks(
pytest.importorskip("PIL")
pytest.importorskip("torch")
# Use a tmp studio root so the ingest subprocess writes images
# somewhere isolated.
# tmp studio root isolates the subprocess's image writes.
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
monkeypatch.setenv("UNSLOTH_RAG_EMBEDDING_MODEL", "BAAI/BGE-VL-base")
monkeypatch.setenv("UNSLOTH_RAG_CHUNK_SIZE", "200")
monkeypatch.setenv("UNSLOTH_RAG_CHUNK_OVERLAP", "20")
# Reset module-level caches so the new env vars take effect.
# Reset module caches so the new env vars apply.
import importlib
import utils.rag.config as rag_config
@ -57,7 +56,7 @@ def test_multimodal_subprocess_emits_image_and_caption_chunks(
embeddings_module._model = None
embeddings_module._model_name = None
# Generate a small PDF with text + one embedded image.
# Small PDF: text + one embedded image.
from PIL import Image
img = Image.new("RGB", (96, 64), (200, 100, 50))
@ -83,7 +82,7 @@ def test_multimodal_subprocess_emits_image_and_caption_chunks(
doc.save(str(pdf_path))
doc.close()
# Drive the subprocess worker in-process with a regular queue.
# Drive the worker in-process via a regular queue.
from core.rag.ingestion import _subprocess_worker
out_queue: "queue_module.Queue[dict]" = queue_module.Queue()
@ -99,19 +98,18 @@ def test_multimodal_subprocess_emits_image_and_caption_chunks(
document_id = "test-doc-1",
)
# Drain everything (the queue is in-process so order is stable).
# Drain all events (in-process queue, stable order).
events: list[dict] = []
while not out_queue.empty():
events.append(out_queue.get_nowait())
# The worker must emit at least one chunks_batch and exactly one
# terminal complete/error event.
# Expect >=1 chunks_batch and exactly one terminal complete/error.
assert any(e["type"] == "chunks_batch" for e in events)
terminals = [e for e in events if e["type"] in ("complete", "error")]
assert len(terminals) == 1, terminals
assert terminals[0]["type"] == "complete"
# Collect all chunks across batches.
# Collect chunks across batches.
all_chunks: list[dict] = []
for e in events:
if e["type"] == "chunks_batch":
@ -120,12 +118,10 @@ def test_multimodal_subprocess_emits_image_and_caption_chunks(
kinds = [c.get("kind") for c in all_chunks]
assert "text" in kinds, "expected at least one text chunk"
assert "image" in kinds, "expected at least one image chunk"
# The PDF has a paragraph immediately after the image, so caption
# pairing should fire.
# Paragraph right after the image triggers caption pairing.
assert "caption" in kinds, "expected at least one caption chunk"
# Image chunks must carry a file path that exists on disk under
# the tmp studio root.
# Image chunks carry a path on disk under the tmp studio root.
image_chunks = [c for c in all_chunks if c.get("kind") == "image"]
for chunk in image_chunks:
assert chunk.get("image_path"), chunk

9
tests/python/test_rag_parsers.py
View file

@ -34,7 +34,7 @@ def test_markdown_parser_preserves_headings(tmp_path):
file.write_text("# Title\n\nBody text with **emphasis**.", encoding = "utf-8")
result = parse(file)
assert result.pages
# Markdown should pass through unchanged — heading marker preserved.
# Markdown passes through; heading marker preserved.
assert "# Title" in result.pages[0].text
@ -67,7 +67,7 @@ def test_html_parser_emits_markdown_headings(tmp_path):
result = parse(file)
assert result.pages
md = result.pages[0].text
# markdownify converts <h1> → '# ', <h2> → '## '
# markdownify: <h1> → '# ', <h2> → '## '
assert "# Main Title" in md
assert "## Sub Section" in md
assert "visible text" in md
@ -91,8 +91,7 @@ def test_pdf_parser_extracts_pages(tmp_path):
with open(file, "wb") as f:
writer.write(f)
# Blank page yields no extractable text — should return empty pages
# without error.
# Blank page: no text, returns empty pages without error.
result = parse(file)
assert isinstance(result.pages, list)
assert isinstance(result.images, list)
@ -116,7 +115,7 @@ def test_docx_parser_emits_markdown_headings(tmp_path):
result = parse(file)
assert result.pages
md = result.pages[0].text
# mammoth via _STYLE_MAP maps Heading 1/2 → h1/h2 → '# '/'## '.
# mammoth _STYLE_MAP: Heading 1/2 → h1/h2 → '# '/'## '.
assert "# Top Level Heading" in md
assert "## Sub Heading" in md
assert "First paragraph" in md

6
tests/python/test_rag_retrieval.py
View file

@ -16,7 +16,7 @@ def test_rrf_fuses_two_rankings():
dense = [Hit("c", 0.9), Hit("b", 0.8), Hit("d", 0.5)]
fused = _rrf_fuse([bm25, dense], rrf_k = 60, top_k = 3)
ids = [h.chunk_id for h in fused]
# b appears at rank 2 in both -> highest fused score
# b ranks 2 in both -> highest fused score.
assert ids[0] == "b"
assert set(ids) == {"a", "b", "c"} or set(ids) == {"b", "c", "a"}
@ -31,7 +31,7 @@ def test_rrf_top_k_limits_output():
def test_rrf_unique_ranking():
# Single ranking — fused order matches input order.
# Single ranking: fused order matches input.
ranking = [Hit("x", 0.0), Hit("y", 0.0), Hit("z", 0.0)]
fused = _rrf_fuse([ranking], rrf_k = 60, top_k = 3)
assert [h.chunk_id for h in fused] == ["x", "y", "z"]
@ -41,6 +41,6 @@ def test_rrf_preserves_payload_from_first_ranking():
a = Hit("a", 1.0, document_id = "doc1", chunk_index = 5)
b = Hit("a", 2.0, document_id = "doc2", chunk_index = 7)
fused = _rrf_fuse([[a], [b]], rrf_k = 60, top_k = 1)
# First sighting wins for payload (deterministic)
# First sighting wins for payload (deterministic).
assert fused[0].document_id == "doc1"
assert fused[0].chunk_index == 5

7
tests/python/test_rag_tool_handler.py
View file

@ -127,7 +127,7 @@ def test_format_hits_produces_fenced_chunks():
assert 'tokens="42"' in result
assert "first body\n</chunk>" in result
assert '<chunk id="2" source="beta.md" score="0.610">' in result
# Blocks separated by a blank line so the model can scan the list.
# Blank line between blocks so the model can scan them.
assert "</chunk>\n\n<chunk" in result
@ -188,8 +188,7 @@ def test_tool_spec_shape_is_openai_compatible():
assert fn["name"] == "search_knowledge_base"
assert "query" in fn["parameters"]["required"]
assert "top_k" in fn["parameters"]["properties"]
# Description should hint at when to call so the LLM picks it up
# appropriately. Don't lock the exact wording.
# Description hints when to call; don't lock the exact wording.
assert "documents" in fn["description"].lower()
@ -263,5 +262,5 @@ def test_all_tools_includes_rag():
names = [t["function"]["name"] for t in ALL_TOOLS]
assert "search_knowledge_base" in names
assert "web_search" in names # regression — we shouldn't have removed the others
assert "web_search" in names # regression: others must stay
assert "python" in names

2
tests/python/test_rag_vector_store.py
View file

@ -49,7 +49,7 @@ def test_upsert_and_search_returns_nearest_first(isolated_rag_db):
results = vector_store.search(scope, [1.0, 0.0, 0.0, 0.0], top_k = 2)
assert len(results) == 2
assert results[0]["chunk_id"] == "p1"
# Cosine similarity converted to [0, 1]; closer = higher.
# Cosine mapped to [0, 1]; closer = higher.
assert results[0]["score"] > results[1]["score"]