From 08cee7cc0a384c56048cefab22d1e4b52ef49782 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Thu, 28 May 2026 19:52:46 +0400 Subject: [PATCH] Studio: bounded parallel RAG indexing + folder upload Uploading many docs (or a folder) previously spawned an ingestion subprocess per file all at once, thrashing the GPU/CPU. Add a configurable concurrency limit and a folder picker. - ragIndexConcurrency setting (default 1) in the chat runtime store, persisted like the other RAG scalar settings; exposed as a 'Parallel indexing' slider (1-8) at the bottom of the sidebar Retrieval section. - New rag-index-queue.ts semaphore: each document upload acquires a slot before it starts and releases it once its ingestion job finishes (complete / error / already-indexed), so bulk uploads drain at the configured rate. Wired into both composer upload paths (use-thread-doc-uploads + shared-composer). - Folder upload: a second 'Attach a folder' button on the RAG attach control uses a webkitdirectory input; every compatible file is routed through the same queue. Multi-file select already worked (the input has 'multiple' and loops addDoc). - Content-hash dedup (shipped earlier) means re-scanning a folder skips already-indexed files. Not build/UI verified here (no bun); needs bun typecheck + a browser check of bulk/folder upload draining at the set concurrency. --- .../src/components/assistant-ui/thread.tsx | 44 ++++++++++++++++--- .../features/chat/api/chat-settings-api.ts | 1 + .../src/features/chat/chat-settings-sheet.tsx | 30 +++++++++++++ .../chat/hooks/use-thread-doc-uploads.ts | 17 +++++++ .../src/features/chat/shared-composer.tsx | 17 +++++++ .../chat/stores/chat-runtime-store.ts | 20 ++++++++- .../features/chat/utils/rag-index-queue.ts | 41 +++++++++++++++++ 7 files changed, 163 insertions(+), 7 deletions(-) create mode 100644 studio/frontend/src/features/chat/utils/rag-index-queue.ts diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 8050bcd86f..acdab53539 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -81,6 +81,7 @@ import { ChevronLeftIcon, ChevronRightIcon, DownloadIcon, + FolderIcon, GlobeIcon, HeadphonesIcon, ImageIcon, @@ -1243,10 +1244,21 @@ const ToolStatusDisplay: FC = () => { }; // RAG-aware + button: picks doc formats and routes to ingest pipeline. +// A second button picks a whole folder (webkitdirectory) and routes every +// compatible file through the same pipeline (which drains at the configured +// parallel-indexing rate). const RagDocAttachment: FC<{ onSelect: (file: File) => void }> = ({ onSelect, }) => { const inputRef = useRef(null); + const folderInputRef = useRef(null); + const selectCompatible = (files: FileList | null) => { + if (!files) return; + for (let i = 0; i < files.length; i++) { + const f = files[i]; + if (f && isDocumentFile(f)) onSelect(f); + } + }; return ( <> void }> = ({ multiple className="hidden" onChange={(e) => { - const files = e.target.files; - if (!files) return; - for (let i = 0; i < files.length; i++) { - const f = files[i]; - if (f && isDocumentFile(f)) onSelect(f); - } + selectCompatible(e.target.files); + e.target.value = ""; + }} + /> + { + folderInputRef.current = el; + if (el) el.setAttribute("webkitdirectory", ""); + }} + type="file" + multiple + className="hidden" + onChange={(e) => { + selectCompatible(e.target.files); e.target.value = ""; }} /> @@ -1274,6 +1297,15 @@ const RagDocAttachment: FC<{ onSelect: (file: File) => void }> = ({ > + folderInputRef.current?.click()} + > + + ); }; 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 2b7b0f8aa3..ee47b6a0fe 100644 --- a/studio/frontend/src/features/chat/api/chat-settings-api.ts +++ b/studio/frontend/src/features/chat/api/chat-settings-api.ts @@ -38,6 +38,7 @@ export interface PersistedChatSettings { enableRerank?: boolean; ragTopK?: number; ragMinScore?: number; + ragIndexConcurrency?: number; } 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 de3450119c..fb9386ad0c 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -494,6 +494,12 @@ export function ChatSettingsPanel({ const setEnableRerank = useChatRuntimeStore((s) => s.setEnableRerank); const ragTopK = useChatRuntimeStore((s) => s.ragTopK); const setRagTopK = useChatRuntimeStore((s) => s.setRagTopK); + const ragIndexConcurrency = useChatRuntimeStore( + (s) => s.ragIndexConcurrency, + ); + const setRagIndexConcurrency = useChatRuntimeStore( + (s) => s.setRagIndexConcurrency, + ); const activeThreadId = useChatRuntimeStore((s) => s.activeThreadId); const { knowledgeBases, deleteKB } = useKnowledgeBases(); const { documents: threadDocs, remove: removeThreadDoc } = useThreadDocuments( @@ -1703,6 +1709,30 @@ export function ChatSettingsPanel({ disabled={!ragEnabled} /> +
+
+ + + {ragIndexConcurrency} + +
+ + v != null && setRagIndexConcurrency(v) + } + /> +

+ How many documents index at once when you upload several + (or a folder). 1 = one at a time. Higher is faster but uses + more GPU/CPU. +

+
) : null} diff --git a/studio/frontend/src/features/chat/hooks/use-thread-doc-uploads.ts b/studio/frontend/src/features/chat/hooks/use-thread-doc-uploads.ts index 0b83662459..11bf3ffb86 100644 --- a/studio/frontend/src/features/chat/hooks/use-thread-doc-uploads.ts +++ b/studio/frontend/src/features/chat/hooks/use-thread-doc-uploads.ts @@ -8,6 +8,7 @@ import { toast } from "sonner"; import { subscribeToJobEvents } from "@/features/rag/api/rag-api"; import { useRagStore } from "@/features/rag/stores/rag-store"; import { useChatRuntimeStore } from "../stores/chat-runtime-store"; +import { acquireIndexSlot, releaseIndexSlot } from "../utils/rag-index-queue"; export type PendingDoc = { id: string; @@ -90,6 +91,17 @@ export function useThreadDocUploads(): UseThreadDocUploadsResult { ]); void (async () => { + // Hold an indexing slot for this document'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(); + let slotReleased = false; + const releaseSlot = () => { + if (!slotReleased) { + slotReleased = true; + releaseIndexSlot(); + } + }; const ragSource = useChatRuntimeStore.getState().ragSource; let scope: | { kind: "kb"; kbId: string } @@ -115,6 +127,7 @@ export function useThreadDocUploads(): UseThreadDocUploadsResult { ), ); toast.error("Could not create thread for upload"); + releaseSlot(); return; } scope = { kind: "thread", threadId }; @@ -152,6 +165,7 @@ export function useThreadDocUploads(): UseThreadDocUploadsResult { ) { useChatRuntimeStore.getState().setRagSource({ kind: "thread" }); } + releaseSlot(); return; } setPendingDocs((prev) => @@ -180,6 +194,7 @@ export function useThreadDocUploads(): UseThreadDocUploadsResult { .getState() .setRagSource({ kind: "thread" }); } + releaseSlot(); } else if (event.type === "error") { setPendingDocs((prev) => prev.map((d) => @@ -188,6 +203,7 @@ export function useThreadDocUploads(): UseThreadDocUploadsResult { : d, ), ); + releaseSlot(); } }, }); @@ -201,6 +217,7 @@ export function useThreadDocUploads(): UseThreadDocUploadsResult { ), ); toast.error(`Document upload failed: ${message}`); + releaseSlot(); } })(); }, diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx index b6a718c532..d3b0aacaef 100644 --- a/studio/frontend/src/features/chat/shared-composer.tsx +++ b/studio/frontend/src/features/chat/shared-composer.tsx @@ -39,6 +39,7 @@ import { Image03Icon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { subscribeToJobEvents } from "@/features/rag/api/rag-api"; import { useRagStore } from "@/features/rag/stores/rag-store"; +import { acquireIndexSlot, releaseIndexSlot } from "./utils/rag-index-queue"; import { toast } from "@/lib/toast"; import { loadModel, validateModel } from "./api/chat-api"; import { parseExternalModelId, providerTypeSupportsVision } from "./external-providers"; @@ -590,6 +591,17 @@ export function SharedComposer({ { id: localChipId, file, status: "uploading" }, ]); 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. + await acquireIndexSlot(); + let slotReleased = false; + const releaseSlot = () => { + if (!slotReleased) { + slotReleased = true; + releaseIndexSlot(); + } + }; const ragSource = useChatRuntimeStore.getState().ragSource; let scope: | { kind: "kb"; kbId: string } @@ -612,6 +624,7 @@ export function SharedComposer({ ), ); toast.error("Could not create thread for upload"); + releaseSlot(); return; } scope = { kind: "thread", threadId }; @@ -645,6 +658,7 @@ export function SharedComposer({ ) { useChatRuntimeStore.getState().setRagSource({ kind: "thread" }); } + releaseSlot(); return; } setPendingDocs((prev) => @@ -670,6 +684,7 @@ export function SharedComposer({ .getState() .setRagSource({ kind: "thread" }); } + releaseSlot(); } else if (event.type === "error") { setPendingDocs((prev) => prev.map((d) => @@ -678,6 +693,7 @@ export function SharedComposer({ : d, ), ); + releaseSlot(); } }, }); @@ -692,6 +708,7 @@ export function SharedComposer({ ), ); toast.error(`Document upload failed: ${message}`); + releaseSlot(); } })(); }, diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index 85407201a4..a1116e287e 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -330,6 +330,10 @@ type ChatRuntimeStore = { ragTopK: number; // 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. + ragIndexConcurrency: number; hydratePersistedSettings: () => Promise; setModelLoading: (loading: boolean) => void; setModelRequiresTrustRemoteCode: (required: boolean) => void; @@ -384,6 +388,7 @@ type ChatRuntimeStore = { setEnableRerank: (value: boolean) => void; setRagTopK: (value: number) => void; setRagMinScore: (value: number) => void; + setRagIndexConcurrency: (value: number) => void; setRagToolEnabled: (value: boolean) => void; }; @@ -405,7 +410,8 @@ type ScalarSettingKey = | "ragMode" | "enableRerank" | "ragTopK" - | "ragMinScore"; + | "ragMinScore" + | "ragIndexConcurrency"; type PresetHydrationVersions = { customPresets: number; @@ -445,6 +451,7 @@ const SCALAR_SETTING_KEYS = [ "enableRerank", "ragTopK", "ragMinScore", + "ragIndexConcurrency", ] as const satisfies readonly ScalarSettingKey[]; const inferenceParamMutationVersions = Object.fromEntries( @@ -660,6 +667,7 @@ export const useChatRuntimeStore = create((set, get) => ({ enableRerank: false, ragTopK: 5, ragMinScore: 0, + ragIndexConcurrency: 1, hydratePersistedSettings: async () => { if (get().settingsHydrated) { return; @@ -934,6 +942,16 @@ export const useChatRuntimeStore = create((set, get) => ({ setScalarSettingVersion("ragMinScore", ragMinScore, state.ragMinScore); return { ragMinScore }; }), + setRagIndexConcurrency: (ragIndexConcurrency) => + set((state) => { + const clamped = Math.max(1, Math.min(8, Math.round(ragIndexConcurrency))); + setScalarSettingVersion( + "ragIndexConcurrency", + clamped, + state.ragIndexConcurrency, + ); + return { ragIndexConcurrency: clamped }; + }), setToolsEnabled: (toolsEnabled, options) => set(() => { if (options?.persist !== false) { diff --git a/studio/frontend/src/features/chat/utils/rag-index-queue.ts b/studio/frontend/src/features/chat/utils/rag-index-queue.ts new file mode 100644 index 0000000000..43595d4603 --- /dev/null +++ b/studio/frontend/src/features/chat/utils/rag-index-queue.ts @@ -0,0 +1,41 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +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. */ + +let active = 0; +const waiters: Array<() => void> = []; + +function limit(): number { + const n = useChatRuntimeStore.getState().ragIndexConcurrency; + return Math.max(1, Number.isFinite(n) ? Math.round(n) : 1); +} + +function admitWaiters(): void { + while (waiters.length > 0 && active < limit()) { + active += 1; + const next = waiters.shift(); + next?.(); + } +} + +/** Resolves once a slot is free (immediately if under the limit). */ +export function acquireIndexSlot(): Promise { + return new Promise((resolve) => { + waiters.push(resolve); + admitWaiters(); + }); +} + +/** Release a previously-acquired slot and admit the next waiter. */ +export function releaseIndexSlot(): void { + active = Math.max(0, active - 1); + admitWaiters(); +}