diff --git a/studio/backend/routes/rag.py b/studio/backend/routes/rag.py index dc3afee76d..dddb6c66b5 100644 --- a/studio/backend/routes/rag.py +++ b/studio/backend/routes/rag.py @@ -407,6 +407,162 @@ def list_knowledge_bases( return KBListResponse(knowledge_bases = [_row_to_kb(r) for r in rows]) +class ReingestKBRequest(BaseModel): + """All fields optional — omitting one keeps the KB's current value.""" + chunking_strategy: ChunkingStrategy | None = None + mode: KBMode | None = None + embedding_model: str | None = None + + +class ReingestResponse(BaseModel): + job_ids: list[str] + document_ids: list[str] + + +def _reingest_scope( + *, + kb_id: str | None, + thread_id: str | None, + chunking_strategy: str, + mode: str, + embedding_model: str, +) -> ReingestResponse: + """Wipe scope artifacts and re-enqueue every stored document. + + The chat_settings / per-thread defaults aren't touched — caller is + responsible for updating any associated metadata before calling. + """ + scope = kb_scope(kb_id) if kb_id else thread_scope(thread_id) # type: ignore[arg-type] + with get_connection() as conn: + if kb_id: + rows = conn.execute( + "SELECT id, stored_path FROM rag_documents WHERE kb_id = ?", + (kb_id,), + ).fetchall() + else: + rows = conn.execute( + "SELECT id, stored_path FROM rag_documents WHERE thread_id = ?", + (thread_id,), + ).fetchall() + # Delete the rag_documents rows (cascade drops chunks); the + # uploaded file on disk is preserved so we can re-ingest from + # it. We re-INSERT a fresh row per stored_path below. + doc_ids = [r["id"] for r in rows] + if doc_ids: + placeholders = ",".join("?" for _ in doc_ids) + conn.execute( + f"DELETE FROM rag_documents WHERE id IN ({placeholders})", + doc_ids, + ) + conn.commit() + ingestion.delete_scope_artifacts(scope) + + job_ids: list[str] = [] + new_doc_ids: list[str] = [] + for row in rows: + stored_path = Path(row["stored_path"]) + if not stored_path.is_file(): + continue + filename = stored_path.name + # Strip the UUID prefix we attached at upload time so the + # re-inserted document carries the original name. + if "_" in filename: + _uuid_prefix, _, original = filename.partition("_") + if original: + filename = original + upload = _start_ingestion( + filename = filename, + stored_path = stored_path, + byte_size = stored_path.stat().st_size, + content_type = None, + kb_id = kb_id, + thread_id = thread_id, + embedding_model = embedding_model, + chunking_strategy = chunking_strategy, + mode = mode, + ) + job_ids.append(upload.job_id) + new_doc_ids.append(upload.document_id) + return ReingestResponse(job_ids = job_ids, document_ids = new_doc_ids) + + +@router.post( + "/knowledge-bases/{kb_id}/reingest", + response_model = ReingestResponse, +) +def reingest_knowledge_base( + kb_id: str, + payload: ReingestKBRequest, + current_subject: str = Depends(get_current_subject), +) -> ReingestResponse: + from utils.rag.config import resolve_embedder + + kb_row = _kb_or_404(kb_id) + keys = kb_row.keys() if hasattr(kb_row, "keys") else () + current_strategy = ( + kb_row["chunking_strategy"] + if "chunking_strategy" in keys + else "standard" + ) + current_mode = kb_row["mode"] if "mode" in keys else "text" + current_embedder = kb_row["embedding_model"] + + new_strategy = payload.chunking_strategy or current_strategy + new_mode = payload.mode or current_mode + _validate_mode_combo(new_mode, new_strategy) + + new_embedder = ( + payload.embedding_model + or ( + current_embedder + if (new_strategy == current_strategy and new_mode == current_mode) + else resolve_embedder(new_mode, new_strategy) + ) + ) + + with get_connection() as conn: + conn.execute( + """ + UPDATE rag_knowledge_bases + SET chunking_strategy = ?, mode = ?, embedding_model = ? + WHERE id = ? + """, + (new_strategy, new_mode, new_embedder, kb_id), + ) + conn.commit() + + return _reingest_scope( + kb_id = kb_id, + thread_id = None, + chunking_strategy = new_strategy, + mode = new_mode, + embedding_model = new_embedder, + ) + + +@router.post( + "/threads/{thread_id}/reingest", + response_model = ReingestResponse, +) +def reingest_thread_documents( + thread_id: str, + current_subject: str = Depends(get_current_subject), +) -> ReingestResponse: + """Rebuild a thread's RAG index using the current defaults. + + No body — per-thread strategy/mode overrides aren't exposed in v1. + """ + from utils.rag.config import RAG_EMBEDDING_MODEL + + return _reingest_scope( + kb_id = None, + thread_id = thread_id, + chunking_strategy = "standard", + mode = "text", + embedding_model = RAG_EMBEDDING_MODEL, + ) + + @router.delete("/knowledge-bases/{kb_id}") def delete_knowledge_base( kb_id: str, diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index 38ed0dff01..f6f0100d4d 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -447,6 +447,7 @@ export function ChatSettingsPanel({ ragSource.kind === "thread" ? activeThreadId : null, ); const clearThreadIndex = useRagStore((s) => s.clearThreadIndex); + const reingestThread = useRagStore((s) => s.reingestThread); const [kbCreateOpen, setKbCreateOpen] = useState(false); const ragEnabled = ragSource.kind !== "off"; const activeKbId = ragSource.kind === "kb" ? ragSource.kbId : null; @@ -1332,22 +1333,39 @@ export function ChatSettingsPanel({ /> ))} - +
+ + +
)} diff --git a/studio/frontend/src/features/rag/api/rag-api.ts b/studio/frontend/src/features/rag/api/rag-api.ts index affa0b4ffc..d6c18faf3b 100644 --- a/studio/frontend/src/features/rag/api/rag-api.ts +++ b/studio/frontend/src/features/rag/api/rag-api.ts @@ -206,6 +206,46 @@ export async function clearThreadDocuments(threadId: string): Promise { await throwOnError(response); } +export interface ReingestResponse { + job_ids: string[]; + document_ids: string[]; +} + +export interface ReingestKBOptions { + chunking_strategy?: ChunkingStrategy; + mode?: KBMode; + embedding_model?: string; +} + +export async function reingestKnowledgeBase( + kbId: string, + opts: ReingestKBOptions = {}, +): Promise { + const response = await authFetch( + `/api/rag/knowledge-bases/${encodeURIComponent(kbId)}/reingest`, + { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(opts), + }, + ); + return parseJsonOrThrow(response); +} + +export async function reingestThreadDocuments( + threadId: string, +): Promise { + const response = await authFetch( + `/api/rag/threads/${encodeURIComponent(threadId)}/reingest`, + { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: "{}", + }, + ); + return parseJsonOrThrow(response); +} + // ------------------------------------------------------------------ // Search // ------------------------------------------------------------------ diff --git a/studio/frontend/src/features/rag/components/kb-detail-panel.tsx b/studio/frontend/src/features/rag/components/kb-detail-panel.tsx index 033803eacc..d8005135aa 100644 --- a/studio/frontend/src/features/rag/components/kb-detail-panel.tsx +++ b/studio/frontend/src/features/rag/components/kb-detail-panel.tsx @@ -1,6 +1,7 @@ // 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 { Button } from "@/components/ui/button"; import { ScrollArea } from "@/components/ui/scroll-area"; import { Separator } from "@/components/ui/separator"; import { useState } from "react"; @@ -9,12 +10,14 @@ import { useKBDocuments } from "../hooks/use-kb-documents"; import { DocumentRow } from "./document-row"; import { DocumentUploadDropzone } from "./document-upload-dropzone"; import { IngestionProgress } from "./ingestion-progress"; +import { KBReconfigureDialog } from "./kb-reconfigure-dialog"; export function KBDetailPanel({ kb }: { kb: KnowledgeBase }) { const { documents, loading, error, upload, remove } = useKBDocuments(kb.id); const [activeJobsByDoc, setActiveJobsByDoc] = useState>( {}, ); + const [reconfigureOpen, setReconfigureOpen] = useState(false); const handleFiles = async (files: File[]) => { for (const file of files) { @@ -29,14 +32,31 @@ export function KBDetailPanel({ kb }: { kb: KnowledgeBase }) { return (
-
-

{kb.name}

- {kb.description ? ( -

{kb.description}

- ) : null} -

- Embedding model: {kb.embedding_model} -

+
+
+

{kb.name}

+ {kb.description ? ( +

{kb.description}

+ ) : null} +

+ {kb.mode === "multimodal" ? "🖼️ Multimodal · " : ""} + {kb.chunking_strategy === "late" ? "⚡ Late · " : ""} + Embedder: {kb.embedding_model} +

+
+
@@ -82,6 +102,13 @@ export function KBDetailPanel({ kb }: { kb: KnowledgeBase }) { })}
+ +
); } diff --git a/studio/frontend/src/features/rag/components/kb-reconfigure-dialog.tsx b/studio/frontend/src/features/rag/components/kb-reconfigure-dialog.tsx new file mode 100644 index 0000000000..ac6d95c6a4 --- /dev/null +++ b/studio/frontend/src/features/rag/components/kb-reconfigure-dialog.tsx @@ -0,0 +1,207 @@ +// 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 { Button } from "@/components/ui/button"; +import { + Dialog, + DialogContent, + DialogDescription, + DialogFooter, + DialogHeader, + DialogTitle, +} from "@/components/ui/dialog"; +import { Input } from "@/components/ui/input"; +import { Label } from "@/components/ui/label"; +import { + Select, + SelectContent, + SelectItem, + SelectTrigger, + SelectValue, +} from "@/components/ui/select"; +import { useEffect, useState } from "react"; +import type { + ChunkingStrategy, + KBMode, + KnowledgeBase, +} from "../api/rag-api"; +import { useRagStore } from "../stores/rag-store"; + +export function KBReconfigureDialog({ + open, + onOpenChange, + kb, + documentCount, +}: { + open: boolean; + onOpenChange: (open: boolean) => void; + kb: KnowledgeBase; + documentCount: number; +}) { + const reingestKB = useRagStore((s) => s.reingestKB); + const [chunkingStrategy, setChunkingStrategy] = useState( + kb.chunking_strategy, + ); + const [mode, setMode] = useState(kb.mode); + const [embeddingModel, setEmbeddingModel] = useState(""); + const [submitting, setSubmitting] = useState(false); + const [error, setError] = useState(null); + + // Re-sync when the dialog opens against a different KB. + useEffect(() => { + if (open) { + setChunkingStrategy(kb.chunking_strategy); + setMode(kb.mode); + setEmbeddingModel(""); + setError(null); + setSubmitting(false); + } + }, [open, kb.id, kb.chunking_strategy, kb.mode]); + + const lateDisabled = mode === "multimodal"; + const multimodalDisabled = chunkingStrategy === "late"; + + const placeholderEmbedder = + mode === "multimodal" + ? `Current: ${kb.embedding_model}` + : chunkingStrategy === "late" + ? `Current: ${kb.embedding_model}` + : `Current: ${kb.embedding_model}`; + + const changedSettings = + chunkingStrategy !== kb.chunking_strategy || + mode !== kb.mode || + embeddingModel.trim() !== ""; + + const handleSubmit = async (e: React.FormEvent) => { + e.preventDefault(); + if (submitting) return; + const verb = changedSettings ? "Reconfigure and re-index" : "Re-index"; + if ( + !window.confirm( + `${verb} ${documentCount} document${documentCount === 1 ? "" : "s"}? ` + + `Existing chunks will be deleted and rebuilt from the original files. ` + + `Search will be unavailable until ingestion finishes.`, + ) + ) { + return; + } + setSubmitting(true); + setError(null); + try { + await reingestKB(kb.id, { + chunking_strategy: chunkingStrategy, + mode, + embedding_model: embeddingModel.trim() || undefined, + }); + onOpenChange(false); + } catch (err) { + setError(err instanceof Error ? err.message : String(err)); + setSubmitting(false); + } + }; + + return ( + + +
+ + Reconfigure “{kb.name}” + + Change the chunking strategy, mode, or embedder for this KB. + All {documentCount} document{documentCount === 1 ? "" : "s"}{" "} + will be re-ingested from the originals on disk. + + +
+
+ + +
+
+ + +
+
+ + setEmbeddingModel(e.target.value)} + placeholder={placeholderEmbedder} + /> +

+ Leave blank to keep the current model (or pick the matrix + default when mode/strategy changes). +

+
+ {error ? ( +
{error}
+ ) : null} +
+ + + + +
+ + + ); +} diff --git a/studio/frontend/src/features/rag/stores/rag-store.ts b/studio/frontend/src/features/rag/stores/rag-store.ts index 1d777ad582..3cf9d98d66 100644 --- a/studio/frontend/src/features/rag/stores/rag-store.ts +++ b/studio/frontend/src/features/rag/stores/rag-store.ts @@ -15,6 +15,9 @@ import { listThreadDocuments, listThreadIndexes, type RagDocument, + type ReingestKBOptions, + reingestKnowledgeBase as apiReingestKB, + reingestThreadDocuments as apiReingestThread, subscribeToJobEvents, type ThreadIndexSummary, uploadKBDocument, @@ -51,6 +54,9 @@ interface RagStoreState { loadThreadIndexes: () => Promise; clearThreadIndex: (threadId: string) => Promise; + reingestKB: (kbId: string, opts?: ReingestKBOptions) => Promise; + reingestThread: (threadId: string) => Promise; + subscribeJob: (jobId: string, onComplete?: () => void) => void; } @@ -220,6 +226,35 @@ export const useRagStore = create((set, get) => ({ }); }, + async reingestKB(kbId, opts) { + const response = await apiReingestKB(kbId, opts ?? {}); + // Refresh the KB list so updated chunking_strategy / mode / embedder + // values flow back into the UI, and the doc list so old chunk + // counts reset to 0 until each job completes. + void get().loadKnowledgeBases(); + void get().loadKBDocuments(kbId); + // Subscribe to every new job so progress chips render and the doc + // list refreshes on completion (mirrors uploadDocument's pattern). + for (const jobId of response.job_ids) { + get().subscribeJob(jobId, () => { + void get().loadKBDocuments(kbId); + }); + } + return response.job_ids; + }, + + async reingestThread(threadId) { + const response = await apiReingestThread(threadId); + void get().loadThreadDocuments(threadId); + void get().loadThreadIndexes(); + for (const jobId of response.job_ids) { + get().subscribeJob(jobId, () => { + void get().loadThreadDocuments(threadId); + }); + } + return response.job_ids; + }, + subscribeJob(jobId, onComplete) { const existing = get().jobUnsubscribers[jobId]; if (existing) return; diff --git a/tests/python/test_rag_reingest.py b/tests/python/test_rag_reingest.py new file mode 100644 index 0000000000..70508a2397 --- /dev/null +++ b/tests/python/test_rag_reingest.py @@ -0,0 +1,57 @@ +"""Reingest endpoint tests (Backfill UX). + +Full end-to-end reingest needs a running studio + a real embedder; that's +covered manually via the curl smoke flow in the plan. Here we cover the +parts that are testable without external models: payload validation and +the (multimodal, late) constraint propagation. +""" + +import sys +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] +STUDIO_BACKEND = REPO_ROOT / "studio" / "backend" +if str(STUDIO_BACKEND) not in sys.path: + sys.path.insert(0, str(STUDIO_BACKEND)) + + +def test_reingest_request_accepts_all_optional_fields(): + from routes.rag import ReingestKBRequest + + empty = ReingestKBRequest() + assert empty.chunking_strategy is None + assert empty.mode is None + assert empty.embedding_model is None + + partial = ReingestKBRequest(chunking_strategy = "late") + assert partial.chunking_strategy == "late" + assert partial.mode is None + + +def test_reingest_request_rejects_unknown_strategy(): + from routes.rag import ReingestKBRequest + from pydantic import ValidationError + + with pytest.raises(ValidationError): + ReingestKBRequest(chunking_strategy = "telekinetic") + + +def test_reingest_request_rejects_unknown_mode(): + from routes.rag import ReingestKBRequest + from pydantic import ValidationError + + with pytest.raises(ValidationError): + ReingestKBRequest(mode = "augmented") + + +def test_constraint_still_enforced_for_reingest_combos(): + """The combination guard is shared with create — verify it still bites.""" + from fastapi import HTTPException + + from routes.rag import _validate_mode_combo + + with pytest.raises(HTTPException) as excinfo: + _validate_mode_combo("multimodal", "late") + assert excinfo.value.status_code == 400