From c22fb88b4f45c9048da6d29a354f4df7f554da4b Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Tue, 26 May 2026 20:51:26 +0400 Subject: [PATCH] Studio: composer + button uploads to the currently-selected RAG source (KB or thread) --- .../chat/hooks/use-thread-doc-uploads.ts | 78 +++++++++++++------ .../src/features/chat/shared-composer.tsx | 53 ++++++++----- 2 files changed, 88 insertions(+), 43 deletions(-) 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 fa1eab9501..a9b184e801 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 @@ -46,11 +46,19 @@ export interface UseThreadDocUploadsResult { isIndexing: boolean; } -/** Per-thread RAG upload: file → POST → SSE → chip status. */ +/** RAG upload from the composer "+" button. Routes to whichever source is + * currently selected in the Retrieval dropdown: KB → uploads to that KB; + * thread or off → uploads to (and lazy-initializes) the current chat + * thread, then flips source to "thread" on first ingest if it was "off". */ export function useThreadDocUploads(): UseThreadDocUploadsResult { const aui = useAui(); const activeThreadId = useChatRuntimeStore((s) => s.activeThreadId); const [pendingDocs, setPendingDocs] = useState([]); + // Track the scope key per chip so removeDoc can dispatch the right + // delete (KB docs vs thread docs use different scope keys in the store). + const [chipScopeKeys, setChipScopeKeys] = useState>( + {}, + ); // Brand-new chats have no backend thread until the first message is sent. // Initialize the current local thread to mint a remoteId so RAG uploads @@ -82,28 +90,40 @@ export function useThreadDocUploads(): UseThreadDocUploadsResult { ]); void (async () => { - const threadId = await ensureThreadId(); - if (!threadId) { - setPendingDocs((prev) => - prev.map((d) => - d.id === localChipId - ? { - ...d, - status: "error", - errorMessage: "Could not create thread for upload", - } - : d, - ), - ); - toast.error("Could not create thread for upload"); - return; + const ragSource = useChatRuntimeStore.getState().ragSource; + let scope: + | { kind: "kb"; kbId: string } + | { kind: "thread"; threadId: string } + | null = null; + let scopeKey: string | null = null; + if (ragSource.kind === "kb") { + scope = { kind: "kb", kbId: ragSource.kbId }; + scopeKey = `kb:${ragSource.kbId}`; + } else { + // ragSource is "thread" or "off" — fall back to thread. + const threadId = await ensureThreadId(); + if (!threadId) { + setPendingDocs((prev) => + prev.map((d) => + d.id === localChipId + ? { + ...d, + status: "error", + errorMessage: "Could not create thread for upload", + } + : d, + ), + ); + toast.error("Could not create thread for upload"); + return; + } + scope = { kind: "thread", threadId }; + scopeKey = `thread:${threadId}`; } + setChipScopeKeys((m) => ({ ...m, [localChipId]: scopeKey })); const uploadDocument = useRagStore.getState().uploadDocument; try { - const { documentId, jobId } = await uploadDocument( - { kind: "thread", threadId }, - file, - ); + const { documentId, jobId } = await uploadDocument(scope, file); setPendingDocs((prev) => prev.map((d) => d.id === localChipId @@ -119,7 +139,13 @@ export function useThreadDocUploads(): UseThreadDocUploadsResult { d.id === localChipId ? { ...d, status: "ready" } : d, ), ); - if (useChatRuntimeStore.getState().ragSource.kind === "off") { + // First ingest in an off-source chat: flip to thread so + // the model has somewhere to search. KB uploads don't + // need this — source is already a KB. + if ( + scope?.kind === "thread" && + useChatRuntimeStore.getState().ragSource.kind === "off" + ) { useChatRuntimeStore .getState() .setRagSource({ kind: "thread" }); @@ -156,15 +182,21 @@ export function useThreadDocUploads(): UseThreadDocUploadsResult { setPendingDocs((prev) => { const doc = prev.find((d) => d.id === id); if (doc?.documentId) { + const scopeKey = + chipScopeKeys[id] ?? `thread:${activeThreadId ?? ""}`; void useRagStore .getState() - .deleteDocument(doc.documentId, `thread:${activeThreadId ?? ""}`) + .deleteDocument(doc.documentId, scopeKey) .catch(() => {}); } return prev.filter((d) => d.id !== id); }); + setChipScopeKeys((m) => { + const { [id]: _gone, ...rest } = m; + return rest; + }); }, - [activeThreadId], + [activeThreadId, chipScopeKeys], ); const clearDocs = useCallback(() => setPendingDocs([]), []); diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx index 121bb028e4..8292b1e5d2 100644 --- a/studio/frontend/src/features/chat/shared-composer.tsx +++ b/studio/frontend/src/features/chat/shared-composer.tsx @@ -524,6 +524,9 @@ export function SharedComposer({ } }, [aui]); + // Composer "+" upload — routes to whichever scope the Retrieval + // dropdown currently points at (KB or thread). Keeps "what you see is + // what you upload to" so users don't get silent thread-vs-KB mismatches. const addDoc = useCallback( (file: File) => { const localChipId = crypto.randomUUID(); @@ -532,28 +535,35 @@ export function SharedComposer({ { id: localChipId, file, status: "uploading" }, ]); void (async () => { - const threadId = await ensureThreadId(); - if (!threadId) { - setPendingDocs((prev) => - prev.map((d) => - d.id === localChipId - ? { - ...d, - status: "error", - errorMessage: "Could not create thread for upload", - } - : d, - ), - ); - toast.error("Could not create thread for upload"); - return; + const ragSource = useChatRuntimeStore.getState().ragSource; + let scope: + | { kind: "kb"; kbId: string } + | { kind: "thread"; threadId: string } + | null = null; + if (ragSource.kind === "kb") { + scope = { kind: "kb", kbId: ragSource.kbId }; + } else { + const threadId = await ensureThreadId(); + if (!threadId) { + setPendingDocs((prev) => + prev.map((d) => + d.id === localChipId + ? { + ...d, + status: "error", + errorMessage: "Could not create thread for upload", + } + : d, + ), + ); + toast.error("Could not create thread for upload"); + return; + } + scope = { kind: "thread", threadId }; } const uploadDocument = useRagStore.getState().uploadDocument; try { - const { documentId, jobId } = await uploadDocument( - { kind: "thread", threadId }, - file, - ); + const { documentId, jobId } = await uploadDocument(scope, file); setPendingDocs((prev) => prev.map((d) => d.id === localChipId @@ -569,7 +579,10 @@ export function SharedComposer({ d.id === localChipId ? { ...d, status: "ready" } : d, ), ); - if (useChatRuntimeStore.getState().ragSource.kind === "off") { + if ( + scope?.kind === "thread" && + useChatRuntimeStore.getState().ragSource.kind === "off" + ) { useChatRuntimeStore .getState() .setRagSource({ kind: "thread" });