From 7335dc07a9f6f347dfd4d7b0a50ddbd888a07d1e Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Fri, 29 May 2026 15:05:29 +0400 Subject: [PATCH] Studio: make figure captioning optional with a Retrieval toggle --- studio/backend/core/rag/ingestion.py | 35 ++++++++++++------- studio/backend/routes/rag.py | 15 ++++++++ .../features/chat/api/chat-settings-api.ts | 1 + .../src/features/chat/chat-settings-sheet.tsx | 29 +++++++++++++-- .../chat/hooks/use-thread-doc-uploads.ts | 4 ++- .../src/features/chat/shared-composer.tsx | 4 ++- .../chat/stores/chat-runtime-store.ts | 18 +++++++++- .../frontend/src/features/rag/api/rag-api.ts | 9 +++-- .../rag/components/kb-reconfigure-dialog.tsx | 2 ++ .../src/features/rag/stores/rag-store.ts | 7 ++-- 10 files changed, 101 insertions(+), 23 deletions(-) diff --git a/studio/backend/core/rag/ingestion.py b/studio/backend/core/rag/ingestion.py index 4e8915fb07..f3030c22ba 100644 --- a/studio/backend/core/rag/ingestion.py +++ b/studio/backend/core/rag/ingestion.py @@ -63,6 +63,7 @@ def _subprocess_worker( document_id: str = "", vlm_url: str | None = None, 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 @@ -101,7 +102,7 @@ def _subprocess_worker( # passes these same captions through to _stream_image_chunks # below — no duplicate VLM calls per image. captions: list[str] = [] - if parsed.images: + if parsed.images and enable_captions: out_queue.put( {"type": "progress", "stage": "caption_images", "progress": 0.08} ) @@ -838,6 +839,7 @@ def enqueue_ingestion( embedding_model: str | None = None, chunking_strategy: str = "standard", mode: str = "text", + enable_captions: bool = True, ) -> str: """Create the job row, spawn the subprocess, start the pump; return job_id.""" from utils.rag.config import resolve_embedder @@ -853,19 +855,25 @@ def enqueue_ingestion( # 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). - vlm_url, vlm_model = _probe_loaded_vlm() - if vlm_url: - logger.info( - "RAG ingest: will caption figures via loaded chat VLM", - vlm_model = vlm_model, - vlm_url = vlm_url, - ) + # back to the helper VLM (pre-cached at studio startup). Skipped + # entirely when captioning is disabled for this upload. + vlm_url: str | None = None + vlm_model: str | None = None + if enable_captions: + vlm_url, vlm_model = _probe_loaded_vlm() + if vlm_url: + logger.info( + "RAG ingest: will caption figures via loaded chat VLM", + vlm_model = vlm_model, + vlm_url = vlm_url, + ) + else: + logger.info( + "RAG ingest: no vision-capable chat model loaded; " + "subprocess will use the helper gemma-3n VLM fallback." + ) else: - logger.info( - "RAG ingest: no vision-capable chat model loaded; " - "subprocess will use the helper gemma-3n VLM fallback." - ) + logger.info("RAG ingest: figure captioning disabled for this upload") job_id = str(uuid4()) with get_connection() as conn: conn.execute( @@ -898,6 +906,7 @@ def enqueue_ingestion( document_id, vlm_url, vlm_model, + enable_captions, ), daemon = True, ) diff --git a/studio/backend/routes/rag.py b/studio/backend/routes/rag.py index bbf657b851..0d8aca6ce3 100644 --- a/studio/backend/routes/rag.py +++ b/studio/backend/routes/rag.py @@ -328,6 +328,7 @@ def _start_ingestion( embedding_model: str, chunking_strategy: str = "standard", mode: str = "text", + caption_images: bool = True, content_hash: str | None = None, ) -> UploadResponse: document_id = str(uuid4()) @@ -390,6 +391,7 @@ def _start_ingestion( embedding_model = embedding_model, chunking_strategy = chunking_strategy, mode = mode, + enable_captions = caption_images, ) return UploadResponse(document_id = document_id, job_id = job_id, filename = filename) @@ -608,6 +610,9 @@ 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. + caption_images: bool | None = None def _thread_settings_key(thread_id: str) -> str: @@ -681,6 +686,8 @@ class ReingestKBRequest(BaseModel): chunking_strategy: ChunkingStrategy | None = None mode: KBMode | None = None embedding_model: str | None = None + # Not persisted on the KB; omit or None keeps captioning on for the rebuild. + caption_images: bool | None = None class ReingestResponse(BaseModel): @@ -695,6 +702,7 @@ def _reingest_scope( chunking_strategy: str, mode: str, embedding_model: str, + caption_images: bool = True, ) -> ReingestResponse: """Wipe scope artifacts and re-enqueue every document; metadata untouched.""" scope = kb_scope(kb_id) if kb_id else thread_scope(thread_id) # type: ignore[arg-type] @@ -742,6 +750,7 @@ def _reingest_scope( embedding_model = embedding_model, chunking_strategy = chunking_strategy, mode = mode, + caption_images = caption_images, ) job_ids.append(upload.job_id) new_doc_ids.append(upload.document_id) @@ -794,6 +803,7 @@ def reingest_knowledge_base( chunking_strategy = new_strategy, mode = new_mode, embedding_model = new_embedder, + caption_images = payload.caption_images is not False, ) @@ -834,6 +844,7 @@ def reingest_thread_documents( chunking_strategy = settings.chunking_strategy, mode = settings.mode, embedding_model = embedder, + caption_images = payload.caption_images is not False, ) @@ -863,6 +874,7 @@ def delete_knowledge_base( async def upload_kb_document( kb_id: str, file: UploadFile, + caption_images: bool = True, current_subject: str = Depends(get_current_subject), ) -> UploadResponse: kb_row = _kb_or_404(kb_id) @@ -883,6 +895,7 @@ async def upload_kb_document( embedding_model = kb_row["embedding_model"], chunking_strategy = chunking_strategy, mode = mode, + caption_images = caption_images, content_hash = content_hash, ) @@ -891,6 +904,7 @@ async def upload_kb_document( async def upload_thread_document( thread_id: str, file: UploadFile, + caption_images: bool = True, current_subject: str = Depends(get_current_subject), ) -> UploadResponse: from utils.rag.config import resolve_embedder @@ -912,6 +926,7 @@ async def upload_thread_document( embedding_model = embedder, chunking_strategy = settings.chunking_strategy, mode = settings.mode, + caption_images = caption_images, content_hash = content_hash, ) diff --git a/studio/frontend/src/features/chat/api/chat-settings-api.ts b/studio/frontend/src/features/chat/api/chat-settings-api.ts index ee47b6a0fe..05898f9597 100644 --- a/studio/frontend/src/features/chat/api/chat-settings-api.ts +++ b/studio/frontend/src/features/chat/api/chat-settings-api.ts @@ -39,6 +39,7 @@ export interface PersistedChatSettings { ragTopK?: number; ragMinScore?: number; ragIndexConcurrency?: number; + ragCaptionImages?: boolean; } interface ChatSettingsResponse { diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index fb9386ad0c..d2c9ddb1aa 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -500,6 +500,10 @@ export function ChatSettingsPanel({ const setRagIndexConcurrency = useChatRuntimeStore( (s) => s.setRagIndexConcurrency, ); + const ragCaptionImages = useChatRuntimeStore((s) => s.ragCaptionImages); + const setRagCaptionImages = useChatRuntimeStore( + (s) => s.setRagCaptionImages, + ); const activeThreadId = useChatRuntimeStore((s) => s.activeThreadId); const { knowledgeBases, deleteKB } = useKnowledgeBases(); const { documents: threadDocs, remove: removeThreadDoc } = useThreadDocuments( @@ -573,7 +577,10 @@ export function ChatSettingsPanel({ `with the new settings? Existing chunks will be deleted and rebuilt.`, ); if (ok) { - void reingestThread(threadId, patch); + void reingestThread(threadId, { + ...patch, + caption_images: ragCaptionImages, + }); } else { // User declined: refresh so the select snaps back. void loadThreadSettings(threadId); @@ -1466,6 +1473,22 @@ export function ChatSettingsPanel({ source before sending.

+
+
+ + Caption images + + + Describe figures with a vision model during indexing so + they're searchable. Off = faster, text-only indexing. + +
+ setRagCaptionImages(next)} + disabled={!ragEnabled} + /> +