diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index 79b4fd6873d..79f182817f6 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -26,6 +26,184 @@ router = APIRouter() +def _build_file_metadata_entry( + response: Any, + file_data: Optional[Tuple[str, bytes, str]] = None, + file_url: Optional[str] = None, +) -> Dict[str, Any]: + """ + Build a file metadata entry for storing in vector_store_metadata. + + Args: + response: The response from litellm.aingest containing file_id + file_data: Optional tuple of (filename, content, content_type) + file_url: Optional URL if file was ingested from URL + + Returns: + Dictionary with file metadata (file_id, filename, file_url, ingested_at, etc.) + """ + from datetime import datetime, timezone + + # Extract file_id from response + file_id = None + if hasattr(response, "get"): + file_id = response.get("file_id") + elif hasattr(response, "file_id"): + file_id = response.file_id + + # Extract file information from file_data tuple + filename = None + file_size = None + content_type = None + + if file_data: + filename = file_data[0] + file_size = len(file_data[1]) if len(file_data) > 1 else None + content_type = file_data[2] if len(file_data) > 2 else None + + # Build file metadata entry + file_entry = { + "file_id": file_id, + "filename": filename, + "file_url": file_url, + "ingested_at": datetime.now(timezone.utc).isoformat(), + } + + # Add optional fields if available + if file_size is not None: + file_entry["file_size"] = file_size + if content_type is not None: + file_entry["content_type"] = content_type + + return file_entry + + +async def _save_vector_store_to_db_from_rag_ingest( + response: Any, + ingest_options: Dict[str, Any], + prisma_client, + user_api_key_dict: UserAPIKeyAuth, + file_data: Optional[Tuple[str, bytes, str]] = None, + file_url: Optional[str] = None, +) -> None: + """ + Helper function to save a newly created vector store from RAG ingest to the database. + + This function: + - Extracts vector store ID and config from the ingest response + - Checks if the vector store already exists in the database + - Creates a new database entry if it doesn't exist + - Adds the vector store to the registry + + Args: + response: The response from litellm.aingest() + ingest_options: The ingest options containing vector store config + prisma_client: The Prisma database client + user_api_key_dict: User API key authentication info + """ + from litellm.proxy.vector_store_endpoints.management_endpoints import ( + create_vector_store_in_db, + ) + + # Handle both dict and object responses + if hasattr(response, "get"): + vector_store_id = response.get("vector_store_id") + elif hasattr(response, "vector_store_id"): + vector_store_id = response.vector_store_id + else: + verbose_proxy_logger.warning( + f"Unable to extract vector_store_id from response type: {type(response)}" + ) + return + + if vector_store_id is None or not isinstance(vector_store_id, str): + verbose_proxy_logger.warning( + "Vector store ID is None or not a string, skipping database save" + ) + return + + vector_store_config = ingest_options.get("vector_store", {}) + custom_llm_provider = vector_store_config.get("custom_llm_provider") + + # Extract litellm_vector_store_params for custom name and description + litellm_vector_store_params = ingest_options.get("litellm_vector_store_params", {}) + custom_vector_store_name = litellm_vector_store_params.get("vector_store_name") + custom_vector_store_description = litellm_vector_store_params.get("vector_store_description") + + # Build file metadata entry using helper + file_entry = _build_file_metadata_entry( + response=response, + file_data=file_data, + file_url=file_url, + ) + + try: + # Check if vector store already exists in database + existing_vector_store = ( + await prisma_client.db.litellm_managedvectorstorestable.find_unique( + where={"vector_store_id": vector_store_id} + ) + ) + + # Only create if it doesn't exist + if existing_vector_store is None: + verbose_proxy_logger.info( + f"Saving newly created vector store {vector_store_id} to database" + ) + + # Initialize metadata with first file + initial_metadata = { + "ingested_files": [file_entry] + } + + # Use custom name if provided, otherwise default + vector_store_name = custom_vector_store_name or f"RAG Vector Store - {vector_store_id[:8]}" + vector_store_description = custom_vector_store_description or "Created via RAG ingest endpoint" + + await create_vector_store_in_db( + vector_store_id=vector_store_id, + custom_llm_provider=custom_llm_provider or "openai", + prisma_client=prisma_client, + vector_store_name=vector_store_name, + vector_store_description=vector_store_description, + vector_store_metadata=initial_metadata, + ) + + verbose_proxy_logger.info( + f"Vector store {vector_store_id} saved to database successfully" + ) + else: + verbose_proxy_logger.info( + f"Vector store {vector_store_id} already exists, appending file to metadata" + ) + + # Update existing vector store with new file + existing_metadata = existing_vector_store.vector_store_metadata or {} + if isinstance(existing_metadata, str): + import json + existing_metadata = json.loads(existing_metadata) + + ingested_files = existing_metadata.get("ingested_files", []) + ingested_files.append(file_entry) + existing_metadata["ingested_files"] = ingested_files + + # Update the vector store + from litellm.proxy.utils import safe_dumps + await prisma_client.db.litellm_managedvectorstorestable.update( + where={"vector_store_id": vector_store_id}, + data={"vector_store_metadata": safe_dumps(existing_metadata)} + ) + + verbose_proxy_logger.info( + f"Added file {file_entry.get('filename') or file_entry.get('file_url', 'Unknown')} to vector store {vector_store_id} metadata" + ) + except Exception as db_error: + # Log the error but don't fail the request since ingestion succeeded + verbose_proxy_logger.exception( + f"Failed to save vector store {vector_store_id} to database: {db_error}" + ) + + async def parse_rag_ingest_request( request: Request, ) -> Tuple[Dict[str, Any], Optional[Tuple[str, bytes, str]], Optional[str], Optional[str]]: @@ -158,6 +336,7 @@ async def rag_ingest( add_litellm_data_to_request, general_settings, llm_router, + prisma_client, proxy_config, version, ) @@ -189,6 +368,25 @@ async def rag_ingest( **request_data, ) + # Save vector store to database if it was newly created and prisma_client is available + verbose_proxy_logger.debug( + f"RAG Ingest - Checking database save conditions: prisma_client={prisma_client is not None}, response={response is not None}, response_type={type(response)}" + ) + + if prisma_client is not None and response is not None: + await _save_vector_store_to_db_from_rag_ingest( + response=response, + ingest_options=ingest_options, + prisma_client=prisma_client, + user_api_key_dict=user_api_key_dict, + file_data=file_data, + file_url=file_url, + ) + else: + verbose_proxy_logger.warning( + f"Skipping database save: prisma_client={prisma_client is not None}, response={response is not None}" + ) + return response except HTTPException: diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index bc61a60fe5a..f3787e62f4d 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -133,6 +133,112 @@ async def _resolve_embedding_config_from_db( return None +######################################################## +# Helper Functions +######################################################## +async def create_vector_store_in_db( + vector_store_id: str, + custom_llm_provider: str, + prisma_client, + vector_store_name: Optional[str] = None, + vector_store_description: Optional[str] = None, + vector_store_metadata: Optional[Dict] = None, + litellm_params: Optional[Dict] = None, + litellm_credential_name: Optional[str] = None, +) -> LiteLLM_ManagedVectorStore: + """ + Helper function to create a vector store in the database. + + This function handles: + - Checking if vector store already exists + - Creating the vector store in the database + - Adding it to the vector store registry + + Returns: + LiteLLM_ManagedVectorStore: The created vector store object + + Raises: + HTTPException: If vector store already exists or database error occurs + """ + from litellm.types.router import GenericLiteLLMParams + + if prisma_client is None: + raise HTTPException(status_code=500, detail="Database not connected") + + # Check if vector store already exists + existing_vector_store = ( + await prisma_client.db.litellm_managedvectorstorestable.find_unique( + where={"vector_store_id": vector_store_id} + ) + ) + if existing_vector_store is not None: + raise HTTPException( + status_code=400, + detail=f"Vector store with ID {vector_store_id} already exists", + ) + + # Prepare data for database + data_to_create: Dict[str, Any] = { + "vector_store_id": vector_store_id, + "custom_llm_provider": custom_llm_provider, + } + + if vector_store_name is not None: + data_to_create["vector_store_name"] = vector_store_name + if vector_store_description is not None: + data_to_create["vector_store_description"] = vector_store_description + if vector_store_metadata is not None: + data_to_create["vector_store_metadata"] = safe_dumps(vector_store_metadata) + if litellm_credential_name is not None: + data_to_create["litellm_credential_name"] = litellm_credential_name + + # Handle litellm_params - always provide at least an empty dict + if litellm_params: + # Auto-resolve embedding config if embedding model is provided but config is not + embedding_model = litellm_params.get("litellm_embedding_model") + if embedding_model and not litellm_params.get("litellm_embedding_config"): + resolved_config = await _resolve_embedding_config_from_db( + embedding_model=embedding_model, + prisma_client=prisma_client + ) + if resolved_config: + litellm_params["litellm_embedding_config"] = resolved_config + verbose_proxy_logger.info( + f"Auto-resolved embedding config for model {embedding_model}" + ) + + litellm_params_dict = GenericLiteLLMParams( + **litellm_params + ).model_dump(exclude_none=True) + data_to_create["litellm_params"] = safe_dumps(litellm_params_dict) + else: + # Provide empty dict if no litellm_params provided + data_to_create["litellm_params"] = safe_dumps({}) + + # Create in database + _new_vector_store = ( + await prisma_client.db.litellm_managedvectorstorestable.create( + data=data_to_create + ) + ) + + new_vector_store: LiteLLM_ManagedVectorStore = LiteLLM_ManagedVectorStore( + **_new_vector_store.model_dump() + ) + + # Add vector store to registry + if litellm.vector_store_registry is not None: + litellm.vector_store_registry.add_vector_store_to_registry( + vector_store=new_vector_store + ) + + verbose_proxy_logger.info( + f"Vector store {vector_store_id} created in database successfully" + ) + + return new_vector_store + + ######################################################## # Management Endpoints ######################################################## @@ -156,71 +262,34 @@ async def new_vector_store( - vector_store_metadata: Optional[Dict] - Additional metadata for the vector store """ from litellm.proxy.proxy_server import prisma_client - from litellm.types.router import GenericLiteLLMParams - - if prisma_client is None: - raise HTTPException(status_code=500, detail="Database not connected") try: - # Check if vector store already exists - existing_vector_store = ( - await prisma_client.db.litellm_managedvectorstorestable.find_unique( - where={"vector_store_id": vector_store.get("vector_store_id")} - ) - ) - if existing_vector_store is not None: + vector_store_id = vector_store.get("vector_store_id") + custom_llm_provider = vector_store.get("custom_llm_provider") + + if not vector_store_id or not custom_llm_provider: raise HTTPException( status_code=400, - detail=f"Vector store with ID {vector_store.get('vector_store_id')} already exists", - ) - - if vector_store.get("vector_store_metadata") is not None: - vector_store["vector_store_metadata"] = safe_dumps( - vector_store.get("vector_store_metadata") + detail="vector_store_id and custom_llm_provider are required" ) - - # Safely handle JSON serialization of litellm_params - litellm_params_json: Optional[str] = None - _input_litellm_params: dict = vector_store.get("litellm_params", {}) or {} - if _input_litellm_params is not None: - # Auto-resolve embedding config if embedding model is provided but config is not - embedding_model = _input_litellm_params.get("litellm_embedding_model") - if embedding_model and not _input_litellm_params.get("litellm_embedding_config"): - resolved_config = await _resolve_embedding_config_from_db( - embedding_model=embedding_model, - prisma_client=prisma_client - ) - if resolved_config: - _input_litellm_params["litellm_embedding_config"] = resolved_config - verbose_proxy_logger.info( - f"Auto-resolved embedding config for model {embedding_model}" - ) - - litellm_params_dict = GenericLiteLLMParams( - **_input_litellm_params - ).model_dump(exclude_none=True) - litellm_params_json = safe_dumps(litellm_params_dict) - del vector_store["litellm_params"] - - _new_vector_store = ( - await prisma_client.db.litellm_managedvectorstorestable.create( - data={ - **vector_store, - "litellm_params": litellm_params_json, - } - ) - ) - - new_vector_store: LiteLLM_ManagedVectorStore = LiteLLM_ManagedVectorStore( - **_new_vector_store.model_dump() + + # Extract and validate metadata + metadata = vector_store.get("vector_store_metadata") + validated_metadata: Optional[Dict] = None + if metadata is not None and isinstance(metadata, dict): + validated_metadata = metadata + + new_vector_store = await create_vector_store_in_db( + vector_store_id=vector_store_id, + custom_llm_provider=custom_llm_provider, + prisma_client=prisma_client, + vector_store_name=vector_store.get("vector_store_name"), + vector_store_description=vector_store.get("vector_store_description"), + vector_store_metadata=validated_metadata, + litellm_params=vector_store.get("litellm_params"), + litellm_credential_name=vector_store.get("litellm_credential_name"), ) - # Add vector store to registry - if litellm.vector_store_registry is not None: - litellm.vector_store_registry.add_vector_store_to_registry( - vector_store=new_vector_store - ) - return { "status": "success", "message": f"Vector store {vector_store.get('vector_store_id')} created successfully", diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 6c8fb6f8340..a1625b6ffbe 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -6948,6 +6948,62 @@ export const vectorStoreUpdateCall = async (accessToken: string, formValues: Rec } }; +export const ragIngestCall = async ( + accessToken: string, + file: File, + customLlmProvider: string, + vectorStoreId?: string, + vectorStoreName?: string, + vectorStoreDescription?: string +): Promise => { + try { + let url = proxyBaseUrl ? `${proxyBaseUrl}/rag/ingest` : `/rag/ingest`; + + const formData = new FormData(); + formData.append("file", file); + + const ingestOptions: any = { + ingest_options: { + vector_store: { + custom_llm_provider: customLlmProvider, + ...(vectorStoreId && { vector_store_id: vectorStoreId }), + }, + }, + }; + + // Add litellm_vector_store_params if name or description provided + if (vectorStoreName || vectorStoreDescription) { + ingestOptions.ingest_options.litellm_vector_store_params = {}; + if (vectorStoreName) { + ingestOptions.ingest_options.litellm_vector_store_params.vector_store_name = vectorStoreName; + } + if (vectorStoreDescription) { + ingestOptions.ingest_options.litellm_vector_store_params.vector_store_description = vectorStoreDescription; + } + } + + 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..db2975781ba --- /dev/null +++ b/ui/litellm-dashboard/src/components/vector_store_management/CreateVectorStore.test.tsx @@ -0,0 +1,166 @@ +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 vector_store_providers +vi.mock("../vector_store_providers", () => ({ + VectorStoreProviders: { + BEDROCK: "Amazon Bedrock", + OPENAI: "OpenAI", + AZURE_OPENAI: "Azure OpenAI", + }, + vectorStoreProviderMap: { + BEDROCK: "bedrock", + OPENAI: "openai", + AZURE_OPENAI: "azure_openai", + }, + vectorStoreProviderLogoMap: { + "Amazon Bedrock": "https://example.com/bedrock.png", + "OpenAI": "https://example.com/openai.png", + "Azure OpenAI": "https://example.com/azure.png", + }, +})); + +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), "bedrock", 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..685e37ac739 --- /dev/null +++ b/ui/litellm-dashboard/src/components/vector_store_management/CreateVectorStore.tsx @@ -0,0 +1,340 @@ +import React, { useState } from "react"; +import { Card, Title, Text } from "@tremor/react"; +import { Upload, Button, Select, Form, message, Alert, Tooltip, Input } from "antd"; +import { InboxOutlined, InfoCircleOutlined } from "@ant-design/icons"; +import type { UploadProps } from "antd"; +import { ragIngestCall } from "../networking"; +import { DocumentUpload, RAGIngestResponse } from "./types"; +import DocumentsTable from "./DocumentsTable"; +import { + VectorStoreProviders, + vectorStoreProviderLogoMap, + vectorStoreProviderMap, +} from "../vector_store_providers"; +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("bedrock"); + const [vectorStoreName, setVectorStoreName] = useState(""); + const [vectorStoreDescription, setVectorStoreDescription] = useState(""); + 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 + vectorStoreName || undefined, + vectorStoreDescription || undefined + ); + + // 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); + } + }; + + 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 Vector Store Details */} + +
+
+ Step 2: Configure Vector Store + + Choose the provider and optionally provide a name and description for your vector store. + +
+ +
+ + Vector Store Name{" "} + + + + + } + > + setVectorStoreName(e.target.value)} + placeholder="e.g., Product Documentation, Customer Support KB" + size="large" + className="rounded-md" + /> + + + + Description{" "} + + + + + } + > + setVectorStoreDescription(e.target.value)} + placeholder="e.g., Contains all product documentation and user guides" + rows={2} + size="large" + className="rounded-md" + /> + + + + Provider{" "} + + + + + } + required + > + + +
+ +
+ +
+
+ + + {/* 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..761bc2b7df7 --- /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(/delete/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(/eye/i); + const copyButtons = screen.getAllByLabelText(/copy/i); + const deleteButtons = screen.getAllByLabelText(/delete/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/TestVectorStoreTab.test.tsx b/ui/litellm-dashboard/src/components/vector_store_management/TestVectorStoreTab.test.tsx new file mode 100644 index 00000000000..ad6ef02e87b --- /dev/null +++ b/ui/litellm-dashboard/src/components/vector_store_management/TestVectorStoreTab.test.tsx @@ -0,0 +1,90 @@ +import { render, screen, fireEvent } from "@testing-library/react"; +import { describe, it, expect, vi } from "vitest"; +import TestVectorStoreTab from "./TestVectorStoreTab"; +import { VectorStore } from "./types"; + +// Mock VectorStoreTester component +vi.mock("./VectorStoreTester", () => ({ + VectorStoreTester: ({ vectorStoreId, accessToken }: { vectorStoreId: string; accessToken: string }) => ( +
+
{vectorStoreId}
+
{accessToken}
+
+ ), +})); + +const mockVectorStores: VectorStore[] = [ + { + vector_store_id: "vs_123", + custom_llm_provider: "openai", + vector_store_name: "Test Store 1", + vector_store_description: "Description 1", + created_at: "2024-01-01T00:00:00Z", + updated_at: "2024-01-01T00:00:00Z", + }, + { + vector_store_id: "vs_456", + custom_llm_provider: "bedrock", + vector_store_name: "Test Store 2", + vector_store_description: "Description 2", + created_at: "2024-01-02T00:00:00Z", + updated_at: "2024-01-02T00:00:00Z", + }, +]; + +describe("TestVectorStoreTab", () => { + it("should render the component successfully", () => { + render(); + + expect(screen.getByText("Select Vector Store")).toBeInTheDocument(); + expect(screen.getByText("Choose a vector store to test search queries against")).toBeInTheDocument(); + }); + + it("should show message when no access token", () => { + render(); + + expect(screen.getByText("Access token is required to test vector stores.")).toBeInTheDocument(); + }); + + it("should show message when no vector stores available", () => { + render(); + + expect(screen.getByText("No vector stores available. Create one first to test it.")).toBeInTheDocument(); + }); + + it("should render VectorStoreTester with first vector store by default", () => { + render(); + + expect(screen.getByTestId("vector-store-tester")).toBeInTheDocument(); + expect(screen.getByTestId("tester-vector-store-id")).toHaveTextContent("vs_123"); + expect(screen.getByTestId("tester-access-token")).toHaveTextContent("test-token"); + }); + + it("should update VectorStoreTester when selecting different vector store", () => { + render(); + + // Find the select component + const selectElement = screen.getByRole("combobox"); + + // Change selection + fireEvent.mouseDown(selectElement); + + // Wait for options to appear and click the second one + const option2 = screen.getByText("Test Store 2"); + fireEvent.click(option2); + + // Verify the tester component updated + expect(screen.getByTestId("tester-vector-store-id")).toHaveTextContent("vs_456"); + }); + + it("should display vector store names in select options", () => { + render(); + + const selectElement = screen.getByRole("combobox"); + fireEvent.mouseDown(selectElement); + + // Use getAllByText since the selected value also shows the name + expect(screen.getAllByText("Test Store 1").length).toBeGreaterThan(0); + expect(screen.getByText("Test Store 2")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/vector_store_management/TestVectorStoreTab.tsx b/ui/litellm-dashboard/src/components/vector_store_management/TestVectorStoreTab.tsx new file mode 100644 index 00000000000..da005491ca6 --- /dev/null +++ b/ui/litellm-dashboard/src/components/vector_store_management/TestVectorStoreTab.tsx @@ -0,0 +1,75 @@ +import React, { useState } from "react"; +import { Card, Select, Typography } from "antd"; +import { VectorStoreTester } from "./VectorStoreTester"; +import { VectorStore } from "./types"; + +const { Text, Title } = Typography; + +interface TestVectorStoreTabProps { + accessToken: string | null; + vectorStores: VectorStore[]; +} + +const TestVectorStoreTab: React.FC = ({ accessToken, vectorStores }) => { + const [selectedVectorStoreId, setSelectedVectorStoreId] = useState( + vectorStores.length > 0 ? vectorStores[0].vector_store_id : undefined + ); + + if (!accessToken) { + return ( + + Access token is required to test vector stores. + + ); + } + + if (vectorStores.length === 0) { + return ( + +
+ No vector stores available. Create one first to test it. +
+ + ); + } + + return ( +
+ +
+
+ Select Vector Store + Choose a vector store to test search queries against +
+ + +
+ + + {selectedVectorStoreId && ( + + )} +
+ ); +}; + +export default TestVectorStoreTab; diff --git a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTable.test.tsx b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTable.test.tsx index 65d15260c4c..0e2be7f62df 100644 --- a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTable.test.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTable.test.tsx @@ -130,9 +130,9 @@ describe("VectorStoreTable", () => { expect(screen.getByText("Provider")).toBeInTheDocument(); expect(screen.getByText("Created At")).toBeInTheDocument(); expect(screen.getByText("Updated At")).toBeInTheDocument(); - // Check that we have the expected number of header cells (6 data + 1 actions) + // Check that we have the expected number of header cells (7 data + 1 actions) const headers = screen.getAllByRole("columnheader"); - expect(headers).toHaveLength(7); + expect(headers).toHaveLength(8); }); it("should render all vector store rows", () => { @@ -183,7 +183,7 @@ describe("VectorStoreTable", () => { it("should render fallback for missing name", () => { renderComponent(); const fallbackElements = screen.getAllByText("-"); - expect(fallbackElements.length).toBe(2); // One for missing name, one for missing description + expect(fallbackElements.length).toBe(3); // One for missing name, one for missing description, one for missing files }); it("should wrap name in tooltip", () => { @@ -203,7 +203,7 @@ describe("VectorStoreTable", () => { it("should render fallback for missing description", () => { renderComponent(); const fallbackElements = screen.getAllByText("-"); - expect(fallbackElements.length).toBe(2); // One for missing name, one for missing description + expect(fallbackElements.length).toBe(3); // One for missing name, one for missing description, one for missing files }); it("should wrap description in tooltip", () => { @@ -386,7 +386,7 @@ describe("VectorStoreTable", () => { it("should span all columns in empty state", () => { renderComponent({ data: [] }); const emptyCell = screen.getByText("No vector stores found").closest("td"); - expect(emptyCell).toHaveAttribute("colSpan", "7"); // 6 data columns + 1 actions column + expect(emptyCell).toHaveAttribute("colSpan", "8"); // 7 data columns + 1 actions column }); }); @@ -403,7 +403,7 @@ describe("VectorStoreTable", () => { renderComponent({ data: minimalData }); expect(screen.getByText("minimal")).toBeInTheDocument(); - expect(screen.getAllByText("-")).toHaveLength(2); // Name and description fallbacks + expect(screen.getAllByText("-")).toHaveLength(3); // Name, description, and files fallbacks }); it("should handle single vector store", () => { diff --git a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTable.tsx b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTable.tsx index 52462c02e98..41e2b63112e 100644 --- a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTable.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTable.tsx @@ -66,6 +66,32 @@ const VectorStoreTable: React.FC = ({ data, onView, onEdi ); }, }, + { + header: "Files", + accessorKey: "vector_store_metadata", + cell: ({ row }) => { + const vectorStore = row.original; + const ingestedFiles = vectorStore.vector_store_metadata?.ingested_files || []; + + if (ingestedFiles.length === 0) { + return -; + } + + const filenames = ingestedFiles + .map((file) => file.filename || file.file_url || "Unknown") + .join(", "); + + const displayText = ingestedFiles.length === 1 + ? ingestedFiles[0].filename || ingestedFiles[0].file_url || "1 file" + : `${ingestedFiles.length} files`; + + return ( + + {displayText} + + ); + }, + }, { header: "Provider", accessorKey: "custom_llm_provider", 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..9cb57b8b9f6 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,8 @@ 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 TestVectorStoreTab from "./TestVectorStoreTab"; import { isAdminRole } from "@/utils/roles"; import NotificationsManager from "../molecules/notifications_manager"; @@ -101,6 +103,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 +142,46 @@ 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 + Test Vector Store + + + + {/* Tab 1: Create Vector Store */} + + + + + {/* Tab 2: Manage Vector Stores */} + + setIsCreateModalVisible(true)}> + + Add Vector Store + + + + + + + + + + {/* Tab 3: Test Vector Store */} + + + + + {/* Create Vector Store Modal */} ; + vector_store_metadata?: VectorStoreMetadata; created_at: string; updated_at: string; created_by?: string; @@ -41,3 +55,32 @@ export interface VectorStoreListResponse { current_page: number; total_pages: number; } + +// Document ingestion types +export interface DocumentUpload { + uid: string; + name: string; + status: "uploading" | "done" | "error" | "removed"; + size?: number; + type?: string; + originFileObj?: File; +} + +export interface RAGIngestRequest { + file_url?: string; + file_id?: string; + ingest_options: { + vector_store: { + custom_llm_provider: string; + vector_store_id?: string; + }; + }; +} + +export interface RAGIngestResponse { + id: string; + status: "completed" | "processing" | "failed"; + vector_store_id: string; + file_id: string; + error?: string; +}