diff --git a/studio/backend/core/rag/captioner.py b/studio/backend/core/rag/captioner.py new file mode 100644 index 0000000000..b6fdf5df3b --- /dev/null +++ b/studio/backend/core/rag/captioner.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Figure captioning via a small generative VLM (PaddleOCR-VL, 4-bit). + +Defensive by design: any load or per-image failure returns an empty +string so the ingestion pipeline can fall back to its prior page-text +caption behaviour rather than crashing. + +Lifecycle: lives inside the ingestion subprocess (`_subprocess_worker` +in `ingestion.py`). The model loads on first `caption_images` call and +is released when the subprocess exits — no impact on the parent chat +model. +""" + +from __future__ import annotations + +import logging +import threading +from typing import Any + +logger = logging.getLogger(__name__) + +_CAPTION_MODEL_NAME = "unsloth/PaddleOCR-VL" +_PROMPT = ( + "Describe this figure in <=60 words. Focus on factual content " + "(axes, labels, captions, visible text, main objects). " + "Do not speculate beyond what is visible." +) +_MAX_NEW_TOKENS = 120 + +_lock = threading.Lock() +_model: Any | None = None +_processor: Any | None = None +_load_failed: bool = False + + +def _load() -> tuple[Any, Any] | None: + """Lazy-load PaddleOCR-VL. Returns (model, processor) or None on failure. + + Sentinel-cached: once a load failure occurs in this subprocess we + don't keep retrying for every batch of images. + """ + global _model, _processor, _load_failed + if _load_failed: + return None + with _lock: + if _model is not None and _processor is not None: + return _model, _processor + try: + import torch + from transformers import AutoModelForVision2Seq, AutoProcessor + + logger.info("Loading RAG captioner: %s", _CAPTION_MODEL_NAME) + processor = AutoProcessor.from_pretrained( + _CAPTION_MODEL_NAME, + trust_remote_code = True, + ) + model = AutoModelForVision2Seq.from_pretrained( + _CAPTION_MODEL_NAME, + trust_remote_code = True, + load_in_4bit = True, + device_map = "auto", + torch_dtype = torch.float16, + ) + model.eval() + _model = model + _processor = processor + return _model, _processor + except Exception as exc: + logger.warning( + "RAG captioner %s failed to load (%s). Falling back to " + "page-text captions for this ingestion.", + _CAPTION_MODEL_NAME, + exc, + ) + _load_failed = True + return None + + +def caption_images(image_bytes_list: list[bytes]) -> list[str]: + """Generate one short caption per image. Same-length output. + + Returns ``""`` for any image whose captioning failed (or all images + if the model couldn't load). Never raises — ingestion must stay + resilient to VLM unavailability. + """ + if not image_bytes_list: + return [] + loaded = _load() + if loaded is None: + return ["" for _ in image_bytes_list] + + model, processor = loaded + from io import BytesIO + + import torch + from PIL import Image + + out: list[str] = [] + for blob in image_bytes_list: + try: + image = Image.open(BytesIO(blob)).convert("RGB") + messages = [ + { + "role": "user", + "content": [ + {"type": "image"}, + {"type": "text", "text": _PROMPT}, + ], + } + ] + prompt = processor.apply_chat_template( + messages, + add_generation_prompt = True, + ) + inputs = processor( + images = image, + text = prompt, + return_tensors = "pt", + ) + input_ids_len = ( + int(inputs["input_ids"].shape[1]) if "input_ids" in inputs else 0 + ) + inputs = { + k: (v.to(model.device) if hasattr(v, "to") else v) + for k, v in inputs.items() + } + with torch.no_grad(): + output_ids = model.generate( + **inputs, + max_new_tokens = _MAX_NEW_TOKENS, + do_sample = False, + ) + # Decode only the newly-generated tokens, not the prompt echo. + new_tokens = output_ids[:, input_ids_len:] + caption = processor.batch_decode( + new_tokens, + skip_special_tokens = True, + )[0] + out.append(caption.strip()) + except Exception as exc: # noqa: BLE001 + logger.warning("caption_images: skipping one image: %s", exc) + out.append("") + return out diff --git a/studio/backend/core/rag/ingestion.py b/studio/backend/core/rag/ingestion.py index 8520c2c96d..3965b40e0f 100644 --- a/studio/backend/core/rag/ingestion.py +++ b/studio/backend/core/rag/ingestion.py @@ -191,6 +191,7 @@ def _stream_image_chunks( first_index: int, ) -> int: """Persist images, emit image+caption chunks; pairs share pair_group.""" + from core.rag.captioner import caption_images from core.rag.embeddings import encode, encode_images from utils.paths.storage_roots import ensure_dir, rag_uploads_root @@ -203,7 +204,7 @@ def _stream_image_chunks( paths: list[str] = [] bytes_for_encoding: list[bytes] = [] - captions: list[str] = [] + fallback_captions: list[str] = [] pages: list[int | None] = [] for idx, img in enumerate(images): ext = _MIME_TO_EXT.get(img.mime_type, ".bin") @@ -215,12 +216,25 @@ def _stream_image_chunks( continue paths.append(str(path)) bytes_for_encoding.append(img.image_bytes) - captions.append(img.nearest_caption or "") + fallback_captions.append(img.nearest_caption or "") pages.append(img.page_number) if not paths: return 0 + # VLM-generated captions; falls back to the parser's nearest_caption + # (page-text blob) when the VLM is unavailable or fails for an image. + out_queue.put({"type": "progress", "stage": "caption_images", "progress": 0.87}) + vlm_captions = caption_images(bytes_for_encoding) + captions: list[str] = [ + ( + vlm_captions[i].strip() + if i < len(vlm_captions) and vlm_captions[i].strip() + else fallback_captions[i] + ) + for i in range(len(bytes_for_encoding)) + ] + image_vectors = encode_images(bytes_for_encoding, model_name = model_name) caption_to_image: list[int] = [i for i, cap in enumerate(captions) if cap.strip()] diff --git a/studio/backend/core/rag/tool.py b/studio/backend/core/rag/tool.py index 836750df78..b0800aee27 100644 --- a/studio/backend/core/rag/tool.py +++ b/studio/backend/core/rag/tool.py @@ -10,6 +10,7 @@ so the model never sees KB UUIDs. from __future__ import annotations from contextvars import ContextVar +from pathlib import Path from typing import Any, Literal from loggers import get_logger @@ -110,6 +111,13 @@ def _format_hits_for_llm(hits: list[dict], start_id: int = 0) -> str: kind = hit.get("kind") if kind and kind != "text": attrs.append(f'kind="{_xml_attr(kind)}"') + image_path = hit.get("image_path") + document_id = hit.get("document_id") + 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. + 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() blocks.append(f"\n{text}\n") return "\n\n".join(blocks) @@ -210,7 +218,8 @@ def search_knowledge_base( rows = conn.execute( f""" SELECT c.id AS chunk_id, c.text, c.page_number, - c.token_count, c.kind, d.filename + c.token_count, c.kind, c.image_path, + c.document_id, d.filename FROM rag_chunks c JOIN rag_documents d ON d.id = c.document_id WHERE c.id IN ({placeholders}) @@ -241,13 +250,14 @@ def search_knowledge_base( else: hits = hits[:k] - # Skip image-kind hits; the paired caption surfaces separately. # 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. formatted: list[dict] = [] for hit in hits: row = lookup.get(hit.chunk_id) - if row is None or row.get("kind") == "image": + if row is None: continue formatted.append( { diff --git a/studio/frontend/src/components/assistant-ui/tool-ui-search-knowledge-base.tsx b/studio/frontend/src/components/assistant-ui/tool-ui-search-knowledge-base.tsx index 02cf157f5d..fe74c43fb2 100644 --- a/studio/frontend/src/components/assistant-ui/tool-ui-search-knowledge-base.tsx +++ b/studio/frontend/src/components/assistant-ui/tool-ui-search-knowledge-base.tsx @@ -7,7 +7,8 @@ import { type ToolCallMessagePartComponent, useAuiState, } from "@assistant-ui/react"; -import { FileTextIcon, LoaderIcon } from "lucide-react"; +import { authFetch } from "@/features/auth"; +import { FileTextIcon, ImageIcon, LoaderIcon } from "lucide-react"; import { memo, useEffect, useState } from "react"; import { cn } from "@/lib/utils"; import { @@ -25,6 +26,7 @@ export interface ParsedChunk { chunkIndex?: string; tokens?: string; kind?: string; + imageUrl?: string; text: string; } @@ -62,6 +64,7 @@ export function parseChunks(raw: string): ParsedChunk[] { chunkIndex: attrs.chunk_index, tokens: attrs.tokens, kind: attrs.kind, + imageUrl: attrs.image_url, text, }); } @@ -72,6 +75,57 @@ export function parseChunks(raw: string): ParsedChunk[] { return out; } +/** Fetch a backend image via the bearer-authed `authFetch`, expose it + * as a blob URL for ``. Cleans up the object URL on unmount. */ +function useAuthedImageUrl(path: string | undefined): string | undefined { + const [url, setUrl] = useState(undefined); + useEffect(() => { + if (!path) { + setUrl(undefined); + return; + } + let cancelled = false; + let objectUrl: string | undefined; + void authFetch(path) + .then(async (response) => { + if (!response.ok) throw new Error(`image fetch ${response.status}`); + return response.blob(); + }) + .then((blob) => { + if (cancelled) return; + objectUrl = URL.createObjectURL(blob); + setUrl(objectUrl); + }) + .catch(() => { + if (!cancelled) setUrl(undefined); + }); + return () => { + cancelled = true; + if (objectUrl) URL.revokeObjectURL(objectUrl); + }; + }, [path]); + return url; +} + +function ChunkImage({ url, alt }: { url: string; alt: string }) { + const blobUrl = useAuthedImageUrl(url); + if (!blobUrl) { + return ( +
+ + Loading image… +
+ ); + } + return ( + {alt} + ); +} + function ChunkCard({ chunk }: { chunk: ParsedChunk }) { const meta: string[] = []; if (chunk.page) meta.push(`page ${chunk.page}`); @@ -93,7 +147,11 @@ function ChunkCard({ chunk }: { chunk: ParsedChunk }) { [{chunk.id}] - + {chunk.kind === "image" ? ( + + ) : ( + + )} {chunk.source} @@ -104,9 +162,14 @@ function ChunkCard({ chunk }: { chunk: ParsedChunk }) { ) : null} -
-        {chunk.text}
-      
+ {chunk.kind === "image" && chunk.imageUrl ? ( + + ) : null} + {chunk.text ? ( +
+          {chunk.text}
+        
+ ) : null} ); } diff --git a/tests/python/test_rag_tool_handler.py b/tests/python/test_rag_tool_handler.py index f6f12016ee..8d31b45fa2 100644 --- a/tests/python/test_rag_tool_handler.py +++ b/tests/python/test_rag_tool_handler.py @@ -153,6 +153,25 @@ def test_format_hits_offsets_ids_by_start_id(): assert '