diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 7e26882e95..b1aceb42cd 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -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, diff --git a/studio/backend/core/rag/authorization.py b/studio/backend/core/rag/authorization.py index dbd1787785..77f59c4435 100644 --- a/studio/backend/core/rag/authorization.py +++ b/studio/backend/core/rag/authorization.py @@ -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) diff --git a/studio/backend/core/rag/captioner.py b/studio/backend/core/rag/captioner.py index b488bd7dcd..3aea7238f5 100644 --- a/studio/backend/core/rag/captioner.py +++ b/studio/backend/core/rag/captioner.py @@ -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 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 + # 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() diff --git a/studio/backend/core/rag/chunking.py b/studio/backend/core/rag/chunking.py index 8ea5ca57fd..21151e158f 100644 --- a/studio/backend/core/rag/chunking.py +++ b/studio/backend/core/rag/chunking.py @@ -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, ) diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index aa0052c5cc..7cf3141ad1 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -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)): diff --git a/studio/backend/core/rag/ingestion.py b/studio/backend/core/rag/ingestion.py index f3030c22ba..27e5d54840 100644 --- a/studio/backend/core/rag/ingestion.py +++ b/studio/backend/core/rag/ingestion.py @@ -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: diff --git a/studio/backend/core/rag/parsers/__init__.py b/studio/backend/core/rag/parsers/__init__.py index 99e565a123..74c7d605e4 100644 --- a/studio/backend/core/rag/parsers/__init__.py +++ b/studio/backend/core/rag/parsers/__init__.py @@ -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\**)(?P