From 240ed0308a5b5db70935b16991bb2e120220044e Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Mon, 26 Jan 2026 15:17:03 -0800 Subject: [PATCH] UI features for RAG ingest --- .../src/components/networking.tsx | 43 +++ .../CreateVectorStore.test.tsx | 156 ++++++++++ .../CreateVectorStore.tsx | 273 ++++++++++++++++++ .../DocumentsTable.test.tsx | 102 +++++++ .../DocumentsTable.tsx | 98 +++++++ .../vector_store_management/index.tsx | 49 +++- .../vector_store_management/types.tsx | 29 ++ 7 files changed, 740 insertions(+), 10 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/vector_store_management/CreateVectorStore.test.tsx create mode 100644 ui/litellm-dashboard/src/components/vector_store_management/CreateVectorStore.tsx create mode 100644 ui/litellm-dashboard/src/components/vector_store_management/DocumentsTable.test.tsx create mode 100644 ui/litellm-dashboard/src/components/vector_store_management/DocumentsTable.tsx diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 6c8fb6f8340..2fef2abfad3 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -6948,6 +6948,49 @@ export const vectorStoreUpdateCall = async (accessToken: string, formValues: Rec } }; +export const ragIngestCall = async ( + accessToken: string, + file: File, + customLlmProvider: string, + vectorStoreId?: string +): Promise => { + try { + let url = proxyBaseUrl ? `${proxyBaseUrl}/rag/ingest` : `/rag/ingest`; + + const formData = new FormData(); + formData.append("file", file); + + const ingestOptions = { + ingest_options: { + vector_store: { + custom_llm_provider: customLlmProvider, + ...(vectorStoreId && { vector_store_id: vectorStoreId }), + }, + }, + }; + + formData.append("request", JSON.stringify(ingestOptions)); + + const response = await fetch(url, { + method: "POST", + headers: { + [globalLitellmHeaderName]: `Bearer ${accessToken}`, + }, + body: formData, + }); + + if (!response.ok) { + const error = await response.json(); + throw new Error(error.error?.message || error.detail || "Failed to ingest document"); + } + + return await response.json(); + } catch (error) { + console.error("Error ingesting document:", error); + throw error; + } +}; + export const getEmailEventSettings = async (accessToken: string): Promise => { try { const url = proxyBaseUrl ? `${proxyBaseUrl}/email/event_settings` : `/email/event_settings`; diff --git a/ui/litellm-dashboard/src/components/vector_store_management/CreateVectorStore.test.tsx b/ui/litellm-dashboard/src/components/vector_store_management/CreateVectorStore.test.tsx new file mode 100644 index 00000000000..d850afc1720 --- /dev/null +++ b/ui/litellm-dashboard/src/components/vector_store_management/CreateVectorStore.test.tsx @@ -0,0 +1,156 @@ +import { render, screen, fireEvent, waitFor, act } from "@testing-library/react"; +import { describe, it, expect, vi, beforeEach } from "vitest"; +import CreateVectorStore from "./CreateVectorStore"; +import * as networking from "../networking"; + +// Mock the networking module +vi.mock("../networking", () => ({ + ragIngestCall: vi.fn(), +})); + +// Mock NotificationsManager +vi.mock("../molecules/notifications_manager", () => ({ + default: { + success: vi.fn(), + fromBackend: vi.fn(), + }, +})); + +// Mock ProviderLogo +vi.mock("../molecules/models/ProviderLogo", () => ({ + ProviderLogo: ({ provider, className }: { provider: string; className?: string }) => ( +
+ {provider} +
+ ), +})); + +describe("CreateVectorStore", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("should render the component successfully", () => { + render(); + + expect(screen.getByText("Create Vector Store")).toBeInTheDocument(); + expect(screen.getByText("Step 1: Upload Documents")).toBeInTheDocument(); + expect(screen.getByText("Step 2: Select Provider")).toBeInTheDocument(); + }); + + it("should display upload area with correct text", () => { + render(); + + expect(screen.getByText("Click or drag files to this area to upload")).toBeInTheDocument(); + expect(screen.getByText(/Support for single or bulk upload/)).toBeInTheDocument(); + }); + + it("should have provider selection dropdown", () => { + render(); + + expect(screen.getByText("Provider")).toBeInTheDocument(); + }); + + it("should have create button disabled initially when no documents", () => { + render(); + + const createButton = screen.getByRole("button", { name: /Create Vector Store/i }); + expect(createButton).toBeDisabled(); + }); + + it("should show uploaded documents table when files are added", async () => { + render(); + + // Create a mock file + const file = new File(["test content"], "test.pdf", { type: "application/pdf" }); + + // Find the upload input (it's hidden but accessible) + const uploadInput = document.querySelector('input[type="file"]') as HTMLInputElement; + + await act(async () => { + if (uploadInput) { + fireEvent.change(uploadInput, { target: { files: [file] } }); + } + }); + + await waitFor(() => { + expect(screen.getByText("Uploaded Documents (1)")).toBeInTheDocument(); + }); + }); + + it("should call ragIngestCall when create button is clicked", async () => { + const mockRagIngestCall = vi.spyOn(networking, "ragIngestCall"); + mockRagIngestCall.mockResolvedValue({ + id: "test-id", + status: "completed", + vector_store_id: "vs_123", + file_id: "file_123", + }); + + const onSuccess = vi.fn(); + render(); + + // Create a mock file + const file = new File(["test content"], "test.pdf", { type: "application/pdf" }); + const uploadInput = document.querySelector('input[type="file"]') as HTMLInputElement; + + await act(async () => { + if (uploadInput) { + fireEvent.change(uploadInput, { target: { files: [file] } }); + } + }); + + // Wait for file to be added + await waitFor(() => { + expect(screen.getByText("Uploaded Documents (1)")).toBeInTheDocument(); + }); + + // Click create button + const createButton = screen.getByRole("button", { name: /Create Vector Store/i }); + + await act(async () => { + fireEvent.click(createButton); + }); + + await waitFor(() => { + expect(mockRagIngestCall).toHaveBeenCalledWith("test-token", expect.any(File), "openai", undefined); + }); + }); + + it("should display success message after successful creation", async () => { + const mockRagIngestCall = vi.spyOn(networking, "ragIngestCall"); + mockRagIngestCall.mockResolvedValue({ + id: "test-id", + status: "completed", + vector_store_id: "vs_123", + file_id: "file_123", + }); + + render(); + + // Create and upload a mock file + const file = new File(["test content"], "test.pdf", { type: "application/pdf" }); + const uploadInput = document.querySelector('input[type="file"]') as HTMLInputElement; + + await act(async () => { + if (uploadInput) { + fireEvent.change(uploadInput, { target: { files: [file] } }); + } + }); + + await waitFor(() => { + expect(screen.getByText("Uploaded Documents (1)")).toBeInTheDocument(); + }); + + // Click create button + const createButton = screen.getByRole("button", { name: /Create Vector Store/i }); + + await act(async () => { + fireEvent.click(createButton); + }); + + await waitFor(() => { + expect(screen.getByText("Vector Store Created Successfully")).toBeInTheDocument(); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/vector_store_management/CreateVectorStore.tsx b/ui/litellm-dashboard/src/components/vector_store_management/CreateVectorStore.tsx new file mode 100644 index 00000000000..7c61255f40b --- /dev/null +++ b/ui/litellm-dashboard/src/components/vector_store_management/CreateVectorStore.tsx @@ -0,0 +1,273 @@ +import React, { useState } from "react"; +import { Card, Title, Text } from "@tremor/react"; +import { Upload, Button, Select, Form, message, Alert } from "antd"; +import { InboxOutlined } from "@ant-design/icons"; +import type { UploadProps } from "antd"; +import { ragIngestCall } from "../networking"; +import { DocumentUpload, RAGIngestResponse } from "./types"; +import DocumentsTable from "./DocumentsTable"; +import { Providers, provider_map } from "../provider_info_helpers"; +import { ProviderLogo } from "../molecules/models/ProviderLogo"; +import NotificationsManager from "../molecules/notifications_manager"; + +const { Dragger } = Upload; + +interface CreateVectorStoreProps { + accessToken: string | null; + onSuccess?: (vectorStoreId: string) => void; +} + +const CreateVectorStore: React.FC = ({ accessToken, onSuccess }) => { + const [form] = Form.useForm(); + const [documents, setDocuments] = useState([]); + const [isCreating, setIsCreating] = useState(false); + const [selectedProvider, setSelectedProvider] = useState("openai"); + const [ingestResults, setIngestResults] = useState([]); + + const uploadProps: UploadProps = { + name: "file", + multiple: true, + accept: ".pdf,.txt,.docx,.md,.doc", + beforeUpload: (file) => { + const isValidType = [ + "application/pdf", + "text/plain", + "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + "application/msword", + "text/markdown", + ].includes(file.type); + + if (!isValidType) { + message.error(`${file.name} is not a supported file type. Please upload PDF, TXT, DOCX, or MD files.`); + return Upload.LIST_IGNORE; + } + + const isLt50M = file.size / 1024 / 1024 < 50; + if (!isLt50M) { + message.error(`${file.name} must be smaller than 50MB!`); + return Upload.LIST_IGNORE; + } + + const newDoc: DocumentUpload = { + uid: file.uid, + name: file.name, + status: "done", + size: file.size, + type: file.type, + originFileObj: file, + }; + + setDocuments((prev) => [...prev, newDoc]); + return false; // Prevent auto upload + }, + onRemove: (file) => { + setDocuments((prev) => prev.filter((doc) => doc.uid !== file.uid)); + }, + fileList: documents.map((doc) => ({ + uid: doc.uid, + name: doc.name, + status: doc.status, + size: doc.size, + })), + showUploadList: false, // We'll use our custom table + }; + + const handleRemoveDocument = (uid: string) => { + setDocuments((prev) => prev.filter((doc) => doc.uid !== uid)); + }; + + const handleCreateVectorStore = async () => { + if (documents.length === 0) { + message.warning("Please upload at least one document"); + return; + } + + if (!selectedProvider) { + message.warning("Please select a provider"); + return; + } + + if (!accessToken) { + message.error("No access token available"); + return; + } + + setIsCreating(true); + const results: RAGIngestResponse[] = []; + let vectorStoreId: string | undefined; + + try { + // Ingest each document + for (const doc of documents) { + if (!doc.originFileObj) continue; + + // Update document status to uploading + setDocuments((prev) => + prev.map((d) => (d.uid === doc.uid ? { ...d, status: "uploading" as const } : d)) + ); + + try { + const result = await ragIngestCall( + accessToken, + doc.originFileObj, + selectedProvider, + vectorStoreId // Use the same vector store ID for subsequent uploads + ); + + // Store the vector store ID from the first successful ingest + if (!vectorStoreId && result.vector_store_id) { + vectorStoreId = result.vector_store_id; + } + + results.push(result); + + // Update document status to done + setDocuments((prev) => + prev.map((d) => (d.uid === doc.uid ? { ...d, status: "done" as const } : d)) + ); + } catch (error) { + console.error(`Error ingesting ${doc.name}:`, error); + // Update document status to error + setDocuments((prev) => + prev.map((d) => (d.uid === doc.uid ? { ...d, status: "error" as const } : d)) + ); + throw error; // Stop processing on first error + } + } + + setIngestResults(results); + NotificationsManager.success( + `Successfully created vector store with ${results.length} document(s). Vector Store ID: ${vectorStoreId}` + ); + + if (onSuccess && vectorStoreId) { + onSuccess(vectorStoreId); + } + + // Clear documents after successful creation + setTimeout(() => { + setDocuments([]); + setIngestResults([]); + }, 3000); + } catch (error) { + console.error("Error creating vector store:", error); + NotificationsManager.fromBackend(`Failed to create vector store: ${error}`); + } finally { + setIsCreating(false); + } + }; + + const supportedProviders = ["openai", "bedrock"]; // Add more as needed + + return ( +
+
+ Create Vector Store + + Upload documents and select a provider to create a new vector store with embedded content. + +
+ + {/* Upload Area */} + +
+ Step 1: Upload Documents + + Upload one or more documents (PDF, TXT, DOCX, MD). Maximum file size: 50MB per file. + +
+ +

+ +

+

Click or drag files to this area to upload

+

+ Support for single or bulk upload. Supported formats: PDF, TXT, DOCX, MD +

+ + + + {/* Documents Table */} + {documents.length > 0 && ( + +
+ Uploaded Documents ({documents.length}) +
+ + + )} + + {/* Provider Selection and Create Button */} + +
+
+ Step 2: Select Provider + + Choose the LLM provider for embedding and vector store operations. + +
+ +
+ + + +
+ +
+ +
+
+ + + {/* Success Message */} + {ingestResults.length > 0 && ( + +

+ Vector Store ID: {ingestResults[0]?.vector_store_id} +

+

+ Documents Ingested: {ingestResults.length} +

+
+ } + type="success" + showIcon + closable + /> + )} + + ); +}; + +export default CreateVectorStore; diff --git a/ui/litellm-dashboard/src/components/vector_store_management/DocumentsTable.test.tsx b/ui/litellm-dashboard/src/components/vector_store_management/DocumentsTable.test.tsx new file mode 100644 index 00000000000..053c7efa623 --- /dev/null +++ b/ui/litellm-dashboard/src/components/vector_store_management/DocumentsTable.test.tsx @@ -0,0 +1,102 @@ +import { render, screen, fireEvent, act } from "@testing-library/react"; +import { describe, it, expect, vi } from "vitest"; +import DocumentsTable from "./DocumentsTable"; +import { DocumentUpload } from "./types"; + +// Mock antd message +vi.mock("antd", async () => { + const actual = await vi.importActual("antd"); + return { + ...actual, + message: { + success: vi.fn(), + }, + }; +}); + +describe("DocumentsTable", () => { + const mockDocuments: DocumentUpload[] = [ + { + uid: "1", + name: "test1.pdf", + status: "done", + size: 1024000, + type: "application/pdf", + }, + { + uid: "2", + name: "test2.txt", + status: "uploading", + size: 2048000, + type: "text/plain", + }, + { + uid: "3", + name: "test3.docx", + status: "error", + size: 512000, + type: "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + }, + ]; + + it("should render the table successfully", () => { + const onRemove = vi.fn(); + render(); + + expect(screen.getByText("test1.pdf")).toBeInTheDocument(); + expect(screen.getByText("test2.txt")).toBeInTheDocument(); + expect(screen.getByText("test3.docx")).toBeInTheDocument(); + }); + + it("should display correct status badges", () => { + const onRemove = vi.fn(); + render(); + + expect(screen.getByText("Ready")).toBeInTheDocument(); + expect(screen.getByText("Uploading")).toBeInTheDocument(); + expect(screen.getByText("Error")).toBeInTheDocument(); + }); + + it("should display file sizes", () => { + const onRemove = vi.fn(); + render(); + + expect(screen.getByText(/1000.00 KB/)).toBeInTheDocument(); + expect(screen.getByText(/2.00 MB/)).toBeInTheDocument(); + expect(screen.getByText(/500.00 KB/)).toBeInTheDocument(); + }); + + it("should call onRemove when delete button is clicked", () => { + const onRemove = vi.fn(); + render(); + + const deleteButtons = screen.getAllByLabelText(/remove/i); + + act(() => { + fireEvent.click(deleteButtons[0]); + }); + + expect(onRemove).toHaveBeenCalledWith("1"); + }); + + it("should show empty state when no documents", () => { + const onRemove = vi.fn(); + render(); + + expect(screen.getByText(/No documents uploaded yet/)).toBeInTheDocument(); + }); + + it("should have action buttons for each document", () => { + const onRemove = vi.fn(); + render(); + + // Each document should have 3 action buttons (view, copy, delete) + const viewButtons = screen.getAllByLabelText(/view details/i); + const copyButtons = screen.getAllByLabelText(/copy id/i); + const deleteButtons = screen.getAllByLabelText(/remove/i); + + expect(viewButtons).toHaveLength(3); + expect(copyButtons).toHaveLength(3); + expect(deleteButtons).toHaveLength(3); + }); +}); diff --git a/ui/litellm-dashboard/src/components/vector_store_management/DocumentsTable.tsx b/ui/litellm-dashboard/src/components/vector_store_management/DocumentsTable.tsx new file mode 100644 index 00000000000..aeb4240d366 --- /dev/null +++ b/ui/litellm-dashboard/src/components/vector_store_management/DocumentsTable.tsx @@ -0,0 +1,98 @@ +import React from "react"; +import { Table, Badge, Tooltip, message } from "antd"; +import { EyeOutlined, CopyOutlined, DeleteOutlined } from "@ant-design/icons"; +import { DocumentUpload } from "./types"; + +interface DocumentsTableProps { + documents: DocumentUpload[]; + onRemove: (uid: string) => void; +} + +const DocumentsTable: React.FC = ({ documents, onRemove }) => { + const handleCopyId = (uid: string) => { + navigator.clipboard.writeText(uid); + message.success("Document ID copied to clipboard"); + }; + + const getStatusBadge = (status: DocumentUpload["status"]) => { + const statusConfig = { + uploading: { color: "blue", text: "Uploading" }, + done: { color: "green", text: "Ready" }, + error: { color: "red", text: "Error" }, + removed: { color: "default", text: "Removed" }, + }; + + const config = statusConfig[status]; + return ; + }; + + const formatFileSize = (bytes?: number) => { + if (!bytes) return "-"; + const kb = bytes / 1024; + if (kb < 1024) return `${kb.toFixed(2)} KB`; + return `${(kb / 1024).toFixed(2)} MB`; + }; + + const columns = [ + { + title: "Name", + dataIndex: "name", + key: "name", + render: (name: string, record: DocumentUpload) => ( +
+ {name} + {record.size && ({formatFileSize(record.size)})} +
+ ), + }, + { + title: "Status", + dataIndex: "status", + key: "status", + width: 150, + render: (status: DocumentUpload["status"]) => getStatusBadge(status), + }, + { + title: "Actions", + key: "actions", + width: 120, + render: (_: any, record: DocumentUpload) => ( +
+ + console.log("View", record)} + /> + + + handleCopyId(record.uid)} + /> + + + onRemove(record.uid)} + /> + +
+ ), + }, + ]; + + return ( + + ); +}; + +export default DocumentsTable; diff --git a/ui/litellm-dashboard/src/components/vector_store_management/index.tsx b/ui/litellm-dashboard/src/components/vector_store_management/index.tsx index 6d21e861d4a..bcc5cd8059b 100644 --- a/ui/litellm-dashboard/src/components/vector_store_management/index.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_management/index.tsx @@ -1,5 +1,5 @@ import React, { useState, useEffect } from "react"; -import { Icon, Button as TremorButton, Col, Text, Grid } from "@tremor/react"; +import { Icon, Button as TremorButton, Col, Text, Grid, TabGroup, TabList, Tab, TabPanels, TabPanel } from "@tremor/react"; import { RefreshIcon } from "@heroicons/react/outline"; import { vectorStoreListCall, vectorStoreDeleteCall, credentialListCall, CredentialItem } from "../networking"; import { VectorStore } from "./types"; @@ -7,6 +7,7 @@ import VectorStoreTable from "./VectorStoreTable"; import VectorStoreForm from "./VectorStoreForm"; import DeleteResourceModal from "../common_components/DeleteResourceModal"; import VectorStoreInfoView from "./vector_store_info"; +import CreateVectorStore from "./CreateVectorStore"; import { isAdminRole } from "@/utils/roles"; import NotificationsManager from "../molecules/notifications_manager"; @@ -101,6 +102,12 @@ const VectorStoreManagement: React.FC = ({ accessToken, userID fetchVectorStores(); }; + const handleVectorStoreCreated = (vectorStoreId: string) => { + console.log("Vector store created:", vectorStoreId); + fetchVectorStores(); + // Optionally switch to the manage tab + }; + useEffect(() => { fetchVectorStores(); fetchCredentials(); @@ -134,18 +141,40 @@ const VectorStoreManagement: React.FC = ({ accessToken, userID -

You can use vector stores to store and retrieve LLM embeddings..

+

You can use vector stores to store and retrieve LLM embeddings.

 - setIsCreateModalVisible(true)}> - + Add Vector Store - + + + Create Vector Store + Manage Vector Stores + - - - - - + + {/* Tab 1: Create Vector Store */} + + + + + {/* Tab 2: Manage Vector Stores */} + + setIsCreateModalVisible(true)}> + + Add Vector Store + + + + + + + + + + {/* Create Vector Store Modal */}