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

Studio: make figure captioning optional with a Retrieval toggle

Browse source
This commit is contained in:
Roland Tannous 2026-05-29 15:05:29 +04:00
commit 7335dc07a9
10 changed files with 101 additions and 23 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -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,
)

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

@ -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,
)

1
studio/frontend/src/features/chat/api/chat-settings-api.ts
View file

@ -39,6 +39,7 @@ export interface PersistedChatSettings {
ragTopK?: number;
ragMinScore?: number;
ragIndexConcurrency?: number;
ragCaptionImages?: boolean;
}
interface ChatSettingsResponse {

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

@ -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.
</p>
</div>
<div className="flex items-center justify-between">
<div className="flex flex-col">
<span className="text-[13px] font-medium">
Caption images
</span>
<span className="text-[11px] text-muted-foreground">
Describe figures with a vision model during indexing so
they're searchable. Off = faster, text-only indexing.
</span>
</div>
<Switch
checked={ragCaptionImages}
onCheckedChange={(next) => setRagCaptionImages(next)}
disabled={!ragEnabled}
/>
</div>
<div className="flex flex-col gap-1.5">
<label className="text-[12px] font-medium text-muted-foreground">
Search mode
@ -1611,7 +1634,9 @@ export function ChatSettingsPanel({
`Re-index all ${threadDocs.length} document${threadDocs.length === 1 ? "" : "s"}? Existing chunks will be deleted and rebuilt; search will be unavailable until ingestion finishes.`,
)
) {
void reingestThread(activeThreadId);
void reingestThread(activeThreadId, {
caption_images: ragCaptionImages,
});
}
}}
>

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

@ -188,12 +188,14 @@ export function useThreadDocUploads(): UseThreadDocUploadsResult {
}
setChipScopeKeys((m) => ({ ...m, [localChipId]: scopeKey as string }));
const uploadDocument = useRagStore.getState().uploadDocument;
const captionImages =
useChatRuntimeStore.getState().ragCaptionImages;
try {
const {
documentId: did,
jobId: jid,
alreadyIndexed,
} = await uploadDocument(scope, file);
} = await uploadDocument(scope, file, captionImages);
documentId = did;
jobId = jid;
if (abort.signal.aborted) {

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

@ -679,12 +679,14 @@ export function SharedComposer({
scopeKey = `thread:${threadId}`;
}
const uploadDocument = useRagStore.getState().uploadDocument;
const captionImages =
useChatRuntimeStore.getState().ragCaptionImages;
try {
const {
documentId: did,
jobId: jid,
alreadyIndexed,
} = await uploadDocument(scope, file);
} = await uploadDocument(scope, file, captionImages);
documentId = did;
jobId = jid;
if (abort.signal.aborted) {

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

@ -334,6 +334,9 @@ type ChatRuntimeStore = {
// rate). 1 = sequential. Keeps many 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.
ragCaptionImages: boolean;
hydratePersistedSettings: () => Promise<void>;
setModelLoading: (loading: boolean) => void;
setModelRequiresTrustRemoteCode: (required: boolean) => void;
@ -389,6 +392,7 @@ type ChatRuntimeStore = {
setRagTopK: (value: number) => void;
setRagMinScore: (value: number) => void;
setRagIndexConcurrency: (value: number) => void;
setRagCaptionImages: (value: boolean) => void;
setRagToolEnabled: (value: boolean) => void;
};
@ -411,7 +415,8 @@ type ScalarSettingKey =
| "enableRerank"
| "ragTopK"
| "ragMinScore"
| "ragIndexConcurrency";
| "ragIndexConcurrency"
| "ragCaptionImages";
type PresetHydrationVersions = {
customPresets: number;
@ -452,6 +457,7 @@ const SCALAR_SETTING_KEYS = [
"ragTopK",
"ragMinScore",
"ragIndexConcurrency",
"ragCaptionImages",
] as const satisfies readonly ScalarSettingKey[];
const inferenceParamMutationVersions = Object.fromEntries(
@ -668,6 +674,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
ragTopK: 5,
ragMinScore: 0,
ragIndexConcurrency: 1,
ragCaptionImages: true,
hydratePersistedSettings: async () => {
if (get().settingsHydrated) {
return;
@ -952,6 +959,15 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
);
return { ragIndexConcurrency: clamped };
}),
setRagCaptionImages: (ragCaptionImages) =>
set((state) => {
setScalarSettingVersion(
"ragCaptionImages",
ragCaptionImages,
state.ragCaptionImages,
);
return { ragCaptionImages };
}),
setToolsEnabled: (toolsEnabled, options) =>
set(() => {
if (options?.persist !== false) {

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

@ -232,11 +232,12 @@ export async function listThreadDocuments(
export async function uploadKBDocument(
kbId: string,
file: File,
captionImages = true,
): Promise<UploadResponse> {
const form = new FormData();
form.append("file", file);
const response = await authFetch(
`/api/rag/knowledge-bases/${encodeURIComponent(kbId)}/documents`,
`/api/rag/knowledge-bases/${encodeURIComponent(kbId)}/documents?caption_images=${captionImages}`,
{ method: "POST", body: form },
);
return parseJsonOrThrow<UploadResponse>(response);
@ -245,11 +246,12 @@ export async function uploadKBDocument(
export async function uploadThreadDocument(
threadId: string,
file: File,
captionImages = true,
): Promise<UploadResponse> {
const form = new FormData();
form.append("file", file);
const response = await authFetch(
`/api/rag/threads/${encodeURIComponent(threadId)}/documents`,
`/api/rag/threads/${encodeURIComponent(threadId)}/documents?caption_images=${captionImages}`,
{ method: "POST", body: form },
);
return parseJsonOrThrow<UploadResponse>(response);
@ -295,6 +297,7 @@ export interface ReingestKBOptions {
chunking_strategy?: ChunkingStrategy;
mode?: KBMode;
embedding_model?: string;
caption_images?: boolean;
}
export async function reingestKnowledgeBase(
@ -322,6 +325,8 @@ export interface UpdateThreadRagSettingsRequest {
chunking_strategy?: ChunkingStrategy;
mode?: KBMode;
embedding_model?: string | null;
// Only consulted by reingest (not persisted as a thread setting).
caption_images?: boolean;
}
export async function getThreadRagSettings(

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

@ -25,6 +25,7 @@ import type {
KBMode,
KnowledgeBase,
} from "../api/rag-api";
import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store";
import { useRagStore } from "../stores/rag-store";
export function KBReconfigureDialog({
@ -93,6 +94,7 @@ export function KBReconfigureDialog({
chunking_strategy: chunkingStrategy,
mode,
embedding_model: embeddingModel.trim() || undefined,
caption_images: useChatRuntimeStore.getState().ragCaptionImages,
});
onOpenChange(false);
} catch (err) {

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

@ -56,6 +56,7 @@ interface RagStoreState {
uploadDocument: (
scope: { kind: "kb"; kbId: string } | { kind: "thread"; threadId: string },
file: File,
captionImages?: boolean,
) => Promise<{
documentId: string;
jobId: string;
@ -189,11 +190,11 @@ export const useRagStore = create<RagStoreState>((set, get) => ({
}
},
async uploadDocument(scope, file) {
async uploadDocument(scope, file, captionImages = true) {
const result =
scope.kind === "kb"
? await uploadKBDocument(scope.kbId, file)
: await uploadThreadDocument(scope.threadId, file);
? await uploadKBDocument(scope.kbId, file, captionImages)
: await uploadThreadDocument(scope.threadId, file, captionImages);
const scopeKey =
scope.kind === "kb"
? kbScopeKey(scope.kbId)