From fc4a484cb187fb2a80f3892285552ae8d2a8fcab Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 27 Jan 2026 14:53:29 -0800 Subject: [PATCH] init ui for bedrock s3 vectors --- .../src/components/networking.tsx | 4 +- .../CreateVectorStore.test.tsx | 116 +++++++++- .../CreateVectorStore.tsx | 84 +++++++- .../S3VectorsConfig.test.tsx | 203 ++++++++++++++++++ .../S3VectorsConfig.tsx | 172 +++++++++++++++ .../src/components/vector_store_providers.tsx | 37 ++++ 6 files changed, 611 insertions(+), 5 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/vector_store_management/S3VectorsConfig.test.tsx create mode 100644 ui/litellm-dashboard/src/components/vector_store_management/S3VectorsConfig.tsx diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index a1625b6ffbe..e0dccbe6e80 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -6954,7 +6954,8 @@ export const ragIngestCall = async ( customLlmProvider: string, vectorStoreId?: string, vectorStoreName?: string, - vectorStoreDescription?: string + vectorStoreDescription?: string, + providerSpecificParams?: Record ): Promise => { try { let url = proxyBaseUrl ? `${proxyBaseUrl}/rag/ingest` : `/rag/ingest`; @@ -6967,6 +6968,7 @@ export const ragIngestCall = async ( vector_store: { custom_llm_provider: customLlmProvider, ...(vectorStoreId && { vector_store_id: vectorStoreId }), + ...(providerSpecificParams && providerSpecificParams), }, }, }; 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 index db2975781ba..a8ba55b91f6 100644 --- a/ui/litellm-dashboard/src/components/vector_store_management/CreateVectorStore.test.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_management/CreateVectorStore.test.tsx @@ -22,17 +22,51 @@ vi.mock("../vector_store_providers", () => ({ BEDROCK: "Amazon Bedrock", OPENAI: "OpenAI", AZURE_OPENAI: "Azure OpenAI", + S3Vectors: "AWS S3 Vectors", }, vectorStoreProviderMap: { BEDROCK: "bedrock", OPENAI: "openai", AZURE_OPENAI: "azure_openai", + S3Vectors: "s3_vectors", }, vectorStoreProviderLogoMap: { "Amazon Bedrock": "https://example.com/bedrock.png", "OpenAI": "https://example.com/openai.png", "Azure OpenAI": "https://example.com/azure.png", + "AWS S3 Vectors": "https://example.com/aws.png", }, + getProviderSpecificFields: vi.fn((provider: string) => { + if (provider === "s3_vectors") { + return [ + { + name: "vector_bucket_name", + label: "Vector Bucket Name", + tooltip: "S3 bucket name for vector storage", + placeholder: "my-vector-bucket", + required: true, + type: "text", + }, + { + name: "aws_region_name", + label: "AWS Region", + tooltip: "AWS region", + placeholder: "us-west-2", + required: true, + type: "text", + }, + { + name: "embedding_model", + label: "Embedding Model", + tooltip: "Embedding model to use", + placeholder: "text-embedding-3-small", + required: true, + type: "select", + }, + ]; + } + return []; + }), })); describe("CreateVectorStore", () => { @@ -43,9 +77,9 @@ describe("CreateVectorStore", () => { it("should render the component successfully", () => { render(); - expect(screen.getByText("Create Vector Store")).toBeInTheDocument(); + expect(screen.getAllByText("Create Vector Store").length).toBeGreaterThan(0); expect(screen.getByText("Step 1: Upload Documents")).toBeInTheDocument(); - expect(screen.getByText("Step 2: Select Provider")).toBeInTheDocument(); + expect(screen.getByText("Step 2: Configure Vector Store")).toBeInTheDocument(); }); it("should display upload area with correct text", () => { @@ -123,7 +157,15 @@ describe("CreateVectorStore", () => { }); await waitFor(() => { - expect(mockRagIngestCall).toHaveBeenCalledWith("test-token", expect.any(File), "bedrock", undefined); + expect(mockRagIngestCall).toHaveBeenCalledWith( + "test-token", + expect.any(File), + "bedrock", + undefined, + undefined, + undefined, + {} + ); }); }); @@ -163,4 +205,72 @@ describe("CreateVectorStore", () => { expect(screen.getByText("Vector Store Created Successfully")).toBeInTheDocument(); }); }); + + it("should display S3 Vectors provider-specific fields when selected", async () => { + render(); + + // Find and click the provider dropdown + const providerSelect = screen.getByRole("combobox"); + + await act(async () => { + fireEvent.mouseDown(providerSelect); + }); + + // Wait for dropdown options to appear + await waitFor(() => { + const s3Option = screen.queryByText("AWS S3 Vectors"); + if (s3Option) { + fireEvent.click(s3Option); + } + }); + + // Check if S3-specific fields are displayed + await waitFor(() => { + expect(screen.queryByText("Vector Bucket Name")).toBeInTheDocument(); + expect(screen.queryByText("AWS Region")).toBeInTheDocument(); + expect(screen.queryByText("Embedding Model")).toBeInTheDocument(); + }); + }); + + it("should validate S3 Vectors required fields before submission", async () => { + render(); + + // Upload a file first + 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(); + }); + + // Select S3 Vectors provider + const providerSelect = screen.getByRole("combobox"); + + await act(async () => { + fireEvent.mouseDown(providerSelect); + }); + + await waitFor(() => { + const s3Option = screen.queryByText("AWS S3 Vectors"); + if (s3Option) { + fireEvent.click(s3Option); + } + }); + + // Try to create without filling required fields + const createButton = screen.getByRole("button", { name: /Create Vector Store/i }); + + await act(async () => { + fireEvent.click(createButton); + }); + + // Should show validation warning (mocked message.warning would be called) + // The actual validation happens in the component + }); }); diff --git a/ui/litellm-dashboard/src/components/vector_store_management/CreateVectorStore.tsx b/ui/litellm-dashboard/src/components/vector_store_management/CreateVectorStore.tsx index 685e37ac739..ffacb2fa42d 100644 --- a/ui/litellm-dashboard/src/components/vector_store_management/CreateVectorStore.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_management/CreateVectorStore.tsx @@ -10,8 +10,11 @@ import { VectorStoreProviders, vectorStoreProviderLogoMap, vectorStoreProviderMap, + getProviderSpecificFields, + VectorStoreFieldConfig, } from "../vector_store_providers"; import NotificationsManager from "../molecules/notifications_manager"; +import S3VectorsConfig from "./S3VectorsConfig"; const { Dragger } = Upload; @@ -28,6 +31,7 @@ const CreateVectorStore: React.FC = ({ accessToken, onSu const [vectorStoreName, setVectorStoreName] = useState(""); const [vectorStoreDescription, setVectorStoreDescription] = useState(""); const [ingestResults, setIngestResults] = useState([]); + const [providerParams, setProviderParams] = useState>({}); const uploadProps: UploadProps = { name: "file", @@ -92,6 +96,15 @@ const CreateVectorStore: React.FC = ({ accessToken, onSu return; } + // Validate provider-specific required fields + const requiredFields = getProviderSpecificFields(selectedProvider).filter((field) => field.required); + for (const field of requiredFields) { + if (!providerParams[field.name]) { + message.warning(`Please provide ${field.label}`); + return; + } + } + if (!accessToken) { message.error("No access token available"); return; @@ -118,7 +131,8 @@ const CreateVectorStore: React.FC = ({ accessToken, onSu selectedProvider, vectorStoreId, // Use the same vector store ID for subsequent uploads vectorStoreName || undefined, - vectorStoreDescription || undefined + vectorStoreDescription || undefined, + providerParams ); // Store the vector store ID from the first successful ingest @@ -298,6 +312,74 @@ const CreateVectorStore: React.FC = ({ accessToken, onSu })} + + {/* S3 Vectors Configuration */} + {selectedProvider === "s3_vectors" && ( + + )} + + {/* Other Provider-specific fields */} + {selectedProvider !== "s3_vectors" && + getProviderSpecificFields(selectedProvider).map((field: VectorStoreFieldConfig) => { + if (field.type === "select") { + // For embedding model selection, we'd need to fetch available models + // For now, provide a text input as fallback + return ( + + {field.label}{" "} + + + + + } + required={field.required} + > + + setProviderParams((prev) => ({ ...prev, [field.name]: e.target.value })) + } + placeholder={field.placeholder} + size="large" + className="rounded-md" + /> + + ); + } + + return ( + + {field.label}{" "} + + + + + } + required={field.required} + > + + setProviderParams((prev) => ({ ...prev, [field.name]: e.target.value })) + } + placeholder={field.placeholder} + size="large" + className="rounded-md" + /> + + ); + })}
diff --git a/ui/litellm-dashboard/src/components/vector_store_management/S3VectorsConfig.test.tsx b/ui/litellm-dashboard/src/components/vector_store_management/S3VectorsConfig.test.tsx new file mode 100644 index 00000000000..beb13b5441e --- /dev/null +++ b/ui/litellm-dashboard/src/components/vector_store_management/S3VectorsConfig.test.tsx @@ -0,0 +1,203 @@ +import { render, screen, fireEvent, waitFor, act } from "@testing-library/react"; +import { describe, it, expect, vi, beforeEach } from "vitest"; +import S3VectorsConfig from "./S3VectorsConfig"; +import * as fetchModels from "../playground/llm_calls/fetch_models"; + +// Mock fetchAvailableModels +vi.mock("../playground/llm_calls/fetch_models", () => ({ + fetchAvailableModels: vi.fn(), +})); + +describe("S3VectorsConfig", () => { + const mockOnParamsChange = vi.fn(); + const defaultProps = { + accessToken: "test-token", + providerParams: {}, + onParamsChange: mockOnParamsChange, + }; + + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("should render the component successfully", () => { + vi.spyOn(fetchModels, "fetchAvailableModels").mockResolvedValue([]); + + render(); + + expect(screen.getByText("AWS S3 Vectors Setup")).toBeInTheDocument(); + expect(screen.getByText("Vector Bucket Name")).toBeInTheDocument(); + expect(screen.getByText("Index Name")).toBeInTheDocument(); + expect(screen.getByText("AWS Region")).toBeInTheDocument(); + expect(screen.getByText("Embedding Model")).toBeInTheDocument(); + }); + + it("should display setup instructions", () => { + vi.spyOn(fetchModels, "fetchAvailableModels").mockResolvedValue([]); + + render(); + + expect( + screen.getByText(/AWS S3 Vectors allows you to store and query vector embeddings directly in S3/) + ).toBeInTheDocument(); + expect(screen.getByText(/Vector buckets and indexes will be automatically created/)).toBeInTheDocument(); + expect(screen.getByText(/Vector dimensions are auto-detected/)).toBeInTheDocument(); + }); + + it("should fetch embedding models on mount", async () => { + const mockModels = [ + { model_group: "text-embedding-3-small", mode: "embedding" }, + { model_group: "text-embedding-3-large", mode: "embedding" }, + { model_group: "gpt-4", mode: "chat" }, + ]; + + const fetchSpy = vi.spyOn(fetchModels, "fetchAvailableModels").mockResolvedValue(mockModels); + + render(); + + await waitFor(() => { + expect(fetchSpy).toHaveBeenCalledWith("test-token"); + }); + }); + + it("should filter and display only embedding models", async () => { + const mockModels = [ + { model_group: "text-embedding-3-small", mode: "embedding" }, + { model_group: "text-embedding-3-large", mode: "embedding" }, + { model_group: "gpt-4", mode: "chat" }, + { model_group: "gpt-3.5-turbo", mode: "chat" }, + ]; + + vi.spyOn(fetchModels, "fetchAvailableModels").mockResolvedValue(mockModels); + + render(); + + // Wait for models to load + await waitFor(() => { + expect(fetchModels.fetchAvailableModels).toHaveBeenCalled(); + }); + + // The component should filter to only embedding models internally + // We can verify this by checking the component loaded successfully + expect(screen.getByText("Embedding Model")).toBeInTheDocument(); + }); + + it("should call onParamsChange when vector bucket name changes", async () => { + vi.spyOn(fetchModels, "fetchAvailableModels").mockResolvedValue([]); + + render(); + + const bucketInput = screen.getByPlaceholderText("my-vector-bucket"); + + await act(async () => { + fireEvent.change(bucketInput, { target: { value: "test-bucket" } }); + }); + + expect(mockOnParamsChange).toHaveBeenCalledWith({ + vector_bucket_name: "test-bucket", + }); + }); + + it("should call onParamsChange when AWS region changes", async () => { + vi.spyOn(fetchModels, "fetchAvailableModels").mockResolvedValue([]); + + render(); + + const regionInput = screen.getByPlaceholderText("us-west-2"); + + await act(async () => { + fireEvent.change(regionInput, { target: { value: "us-east-1" } }); + }); + + expect(mockOnParamsChange).toHaveBeenCalledWith({ + aws_region_name: "us-east-1", + }); + }); + + it("should call onParamsChange when embedding model is selected", async () => { + const mockModels = [ + { model_group: "text-embedding-3-small", mode: "embedding" }, + { model_group: "text-embedding-3-large", mode: "embedding" }, + ]; + + vi.spyOn(fetchModels, "fetchAvailableModels").mockResolvedValue(mockModels); + + render(); + + await waitFor(() => { + expect(fetchModels.fetchAvailableModels).toHaveBeenCalled(); + }); + + // Find the Select component and trigger change directly + const selectElement = screen.getByRole("combobox"); + + await act(async () => { + // Simulate selecting a value by firing the change event + fireEvent.change(selectElement, { target: { value: "text-embedding-3-small" } }); + }); + + // The component should handle the selection + expect(screen.getByText("Embedding Model")).toBeInTheDocument(); + }); + + it("should preserve existing params when updating a field", async () => { + vi.spyOn(fetchModels, "fetchAvailableModels").mockResolvedValue([]); + + const existingParams = { + vector_bucket_name: "existing-bucket", + aws_region_name: "us-west-2", + }; + + render(); + + const indexInput = screen.getByPlaceholderText("my-vector-index"); + + await act(async () => { + fireEvent.change(indexInput, { target: { value: "my-index" } }); + }); + + expect(mockOnParamsChange).toHaveBeenCalledWith({ + vector_bucket_name: "existing-bucket", + aws_region_name: "us-west-2", + index_name: "my-index", + }); + }); + + it("should display existing param values", () => { + vi.spyOn(fetchModels, "fetchAvailableModels").mockResolvedValue([]); + + const existingParams = { + vector_bucket_name: "my-bucket", + index_name: "my-index", + aws_region_name: "eu-west-1", + embedding_model: "text-embedding-3-small", + }; + + render(); + + expect(screen.getByDisplayValue("my-bucket")).toBeInTheDocument(); + expect(screen.getByDisplayValue("my-index")).toBeInTheDocument(); + expect(screen.getByDisplayValue("eu-west-1")).toBeInTheDocument(); + }); + + it("should handle model fetch error gracefully", async () => { + const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {}); + vi.spyOn(fetchModels, "fetchAvailableModels").mockRejectedValue(new Error("Failed to fetch models")); + + render(); + + await waitFor(() => { + expect(consoleErrorSpy).toHaveBeenCalledWith("Error fetching embedding models:", expect.any(Error)); + }); + + consoleErrorSpy.mockRestore(); + }); + + it("should not fetch models if accessToken is null", () => { + const fetchSpy = vi.spyOn(fetchModels, "fetchAvailableModels"); + + render(); + + expect(fetchSpy).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/vector_store_management/S3VectorsConfig.tsx b/ui/litellm-dashboard/src/components/vector_store_management/S3VectorsConfig.tsx new file mode 100644 index 00000000000..48685596211 --- /dev/null +++ b/ui/litellm-dashboard/src/components/vector_store_management/S3VectorsConfig.tsx @@ -0,0 +1,172 @@ +import React, { useState, useEffect } from "react"; +import { Alert, Form, Input, Select, Tooltip } from "antd"; +import { InfoCircleOutlined } from "@ant-design/icons"; +import { fetchAvailableModels, ModelGroup } from "../playground/llm_calls/fetch_models"; + +interface S3VectorsConfigProps { + accessToken: string | null; + providerParams: Record; + onParamsChange: (params: Record) => void; +} + +const S3VectorsConfig: React.FC = ({ + accessToken, + providerParams, + onParamsChange, +}) => { + const [embeddingModels, setEmbeddingModels] = useState([]); + const [isLoadingModels, setIsLoadingModels] = useState(false); + + useEffect(() => { + if (!accessToken) return; + + const loadModels = async () => { + setIsLoadingModels(true); + try { + const models = await fetchAvailableModels(accessToken); + // Filter for embedding models only + const embeddingOnly = models.filter((model) => model.mode === "embedding"); + setEmbeddingModels(embeddingOnly); + } catch (error) { + console.error("Error fetching embedding models:", error); + } finally { + setIsLoadingModels(false); + } + }; + + loadModels(); + }, [accessToken]); + + const handleFieldChange = (fieldName: string, value: string) => { + onParamsChange({ + ...providerParams, + [fieldName]: value, + }); + }; + + return ( + <> + {/* S3 Vectors Setup Instructions */} + +

AWS S3 Vectors allows you to store and query vector embeddings directly in S3:

+
    +
  • Vector buckets and indexes will be automatically created if they don't exist
    • +
  • Vector dimensions are auto-detected from your selected embedding model
    • +
  • Ensure your AWS credentials have permissions for S3 Vectors operations
    • +
  • + Learn more:{" "} + + AWS S3 Vectors Documentation + +
    • +
+
+ } + type="info" + showIcon + style={{ marginBottom: "16px" }} + /> + + {/* Vector Bucket Name */} + + Vector Bucket Name{" "} + + + + + } + required + > + handleFieldChange("vector_bucket_name", e.target.value)} + placeholder="my-vector-bucket" + size="large" + className="rounded-md" + /> + + + {/* Index Name (Optional) */} + + Index Name{" "} + + + + + } + > + handleFieldChange("index_name", e.target.value)} + placeholder="my-vector-index" + size="large" + className="rounded-md" + /> + + + {/* AWS Region */} + + AWS Region{" "} + + + + + } + required + > + handleFieldChange("aws_region_name", e.target.value)} + placeholder="us-west-2" + size="large" + className="rounded-md" + /> + + + {/* Embedding Model */} + + Embedding Model{" "} + + + + + } + required + > +