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

Studio: bounded parallel RAG indexing + folder upload

Browse source 
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.
This commit is contained in:
Roland Tannous 2026-05-28 19:52:46 +04:00
commit 08cee7cc0a
7 changed files with 163 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

44
studio/frontend/src/components/assistant-ui/thread.tsx
View file

@ -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<HTMLInputElement | null>(null);
const folderInputRef = useRef<HTMLInputElement | null>(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 (
<>
<input
@ -1256,12 +1268,23 @@ const RagDocAttachment: FC<{ onSelect: (file: File) => 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 = "";
}}
/>
<input
// webkitdirectory isn't in React's input prop types; set it on the
// element directly so the picker selects a folder (returns every
// file recursively, which selectCompatible then filters).
ref={(el) => {
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 }> = ({
>
<PaperclipIcon className="size-4" />
</TooltipIconButton>
<TooltipIconButton
tooltip="Attach a folder for RAG"
aria-label="Attach a folder for RAG"
variant="ghost"
className="size-8 rounded-full text-muted-foreground"
onClick={() => folderInputRef.current?.click()}
>
<FolderIcon className="size-4" />
</TooltipIconButton>
</>
);
};

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

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

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

@ -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}
/>
</div>
<div className="flex flex-col gap-1.5">
<div className="flex items-center justify-between">
<label className="text-[12px] font-medium text-muted-foreground">
Parallel indexing
</label>
<span className="text-[12px] tabular-nums text-muted-foreground">
{ragIndexConcurrency}
</span>
</div>
<Slider
value={[ragIndexConcurrency]}
min={1}
max={8}
step={1}
onValueChange={([v]) =>
v != null && setRagIndexConcurrency(v)
}
/>
<p className="text-[11px] text-muted-foreground">
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.
</p>
</div>
</div>
</CollapsibleSection>
) : null}

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

@ -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();
}
})();
},

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

@ -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();
}
})();
},

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

@ -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<void>;
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<ChatRuntimeStore>((set, get) => ({
enableRerank: false,
ragTopK: 5,
ragMinScore: 0,
ragIndexConcurrency: 1,
hydratePersistedSettings: async () => {
if (get().settingsHydrated) {
return;
@ -934,6 +942,16 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((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) {

41
studio/frontend/src/features/chat/utils/rag-index-queue.ts Normal file
View file

@ -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<void> {
return new Promise<void>((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();
}