mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-15 14:21:58 +00:00
[Feat] Add UI for /rag/ingest API - upload docs, pdfs etc to create vector stores (#19822)
* feat: _save_vector_store_to_db_from_rag_ingest * UI features for RAG ingest * fix: Endpoints * ragIngestCall * _save_vector_store_to_db_from_rag_ingest * fix: rag_ingest Code QA CHECK * UI fixes unit tests
This commit is contained in:
Ishaan Jaff
@@ -26,6 +26,184 @@ from litellm.proxy.common_utils.http_parsing_utils import (
|
||||
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:
|
||||
|
||||
@@ -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")
|
||||
)
|
||||
|
||||
# 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,
|
||||
}
|
||||
detail="vector_store_id and custom_llm_provider are required"
|
||||
)
|
||||
|
||||
# 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"),
|
||||
)
|
||||
|
||||
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
|
||||
)
|
||||
|
||||
return {
|
||||
"status": "success",
|
||||
"message": f"Vector store {vector_store.get('vector_store_id')} created successfully",
|
||||
|
||||
@@ -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<any> => {
|
||||
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<EmailEventSettingsResponse> => {
|
||||
try {
|
||||
const url = proxyBaseUrl ? `${proxyBaseUrl}/email/event_settings` : `/email/event_settings`;
|
||||
|
||||
+166
@@ -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(<CreateVectorStore accessToken="test-token" />);
|
||||
|
||||
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(<CreateVectorStore accessToken="test-token" />);
|
||||
|
||||
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(<CreateVectorStore accessToken="test-token" />);
|
||||
|
||||
expect(screen.getByText("Provider")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should have create button disabled initially when no documents", () => {
|
||||
render(<CreateVectorStore accessToken="test-token" />);
|
||||
|
||||
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(<CreateVectorStore accessToken="test-token" />);
|
||||
|
||||
// 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(<CreateVectorStore accessToken="test-token" onSuccess={onSuccess} />);
|
||||
|
||||
// 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(<CreateVectorStore accessToken="test-token" />);
|
||||
|
||||
// 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();
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -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<CreateVectorStoreProps> = ({ accessToken, onSuccess }) => {
|
||||
const [form] = Form.useForm();
|
||||
const [documents, setDocuments] = useState<DocumentUpload[]>([]);
|
||||
const [isCreating, setIsCreating] = useState(false);
|
||||
const [selectedProvider, setSelectedProvider] = useState<string>("bedrock");
|
||||
const [vectorStoreName, setVectorStoreName] = useState<string>("");
|
||||
const [vectorStoreDescription, setVectorStoreDescription] = useState<string>("");
|
||||
const [ingestResults, setIngestResults] = useState<RAGIngestResponse[]>([]);
|
||||
|
||||
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 (
|
||||
<div className="space-y-6">
|
||||
<div>
|
||||
<Title>Create Vector Store</Title>
|
||||
<Text className="text-gray-500">
|
||||
Upload documents and select a provider to create a new vector store with embedded content.
|
||||
</Text>
|
||||
</div>
|
||||
|
||||
{/* Upload Area */}
|
||||
<Card>
|
||||
<div className="mb-4">
|
||||
<Text className="font-medium">Step 1: Upload Documents</Text>
|
||||
<Text className="text-sm text-gray-500 block mt-1">
|
||||
Upload one or more documents (PDF, TXT, DOCX, MD). Maximum file size: 50MB per file.
|
||||
</Text>
|
||||
</div>
|
||||
<Dragger {...uploadProps}>
|
||||
<p className="ant-upload-drag-icon">
|
||||
<InboxOutlined style={{ fontSize: "48px", color: "#1890ff" }} />
|
||||
</p>
|
||||
<p className="ant-upload-text">Click or drag files to this area to upload</p>
|
||||
<p className="ant-upload-hint">
|
||||
Support for single or bulk upload. Supported formats: PDF, TXT, DOCX, MD
|
||||
</p>
|
||||
</Dragger>
|
||||
</Card>
|
||||
|
||||
{/* Documents Table */}
|
||||
{documents.length > 0 && (
|
||||
<Card>
|
||||
<div className="mb-4">
|
||||
<Text className="font-medium">Uploaded Documents ({documents.length})</Text>
|
||||
</div>
|
||||
<DocumentsTable documents={documents} onRemove={handleRemoveDocument} />
|
||||
</Card>
|
||||
)}
|
||||
|
||||
{/* Provider Selection and Vector Store Details */}
|
||||
<Card>
|
||||
<div className="space-y-4">
|
||||
<div>
|
||||
<Text className="font-medium">Step 2: Configure Vector Store</Text>
|
||||
<Text className="text-sm text-gray-500 block mt-1">
|
||||
Choose the provider and optionally provide a name and description for your vector store.
|
||||
</Text>
|
||||
</div>
|
||||
|
||||
<Form form={form} layout="vertical">
|
||||
<Form.Item
|
||||
label={
|
||||
<span>
|
||||
Vector Store Name{" "}
|
||||
<Tooltip title="Optional: Give your vector store a meaningful name">
|
||||
<InfoCircleOutlined style={{ marginLeft: "4px" }} />
|
||||
</Tooltip>
|
||||
</span>
|
||||
}
|
||||
>
|
||||
<Input
|
||||
value={vectorStoreName}
|
||||
onChange={(e) => setVectorStoreName(e.target.value)}
|
||||
placeholder="e.g., Product Documentation, Customer Support KB"
|
||||
size="large"
|
||||
className="rounded-md"
|
||||
/>
|
||||
</Form.Item>
|
||||
|
||||
<Form.Item
|
||||
label={
|
||||
<span>
|
||||
Description{" "}
|
||||
<Tooltip title="Optional: Describe what this vector store contains">
|
||||
<InfoCircleOutlined style={{ marginLeft: "4px" }} />
|
||||
</Tooltip>
|
||||
</span>
|
||||
}
|
||||
>
|
||||
<Input.TextArea
|
||||
value={vectorStoreDescription}
|
||||
onChange={(e) => setVectorStoreDescription(e.target.value)}
|
||||
placeholder="e.g., Contains all product documentation and user guides"
|
||||
rows={2}
|
||||
size="large"
|
||||
className="rounded-md"
|
||||
/>
|
||||
</Form.Item>
|
||||
|
||||
<Form.Item
|
||||
label={
|
||||
<span>
|
||||
Provider{" "}
|
||||
<Tooltip title="Select the provider for embedding and vector store operations">
|
||||
<InfoCircleOutlined style={{ marginLeft: "4px" }} />
|
||||
</Tooltip>
|
||||
</span>
|
||||
}
|
||||
required
|
||||
>
|
||||
<Select
|
||||
value={selectedProvider}
|
||||
onChange={setSelectedProvider}
|
||||
placeholder="Select a provider"
|
||||
size="large"
|
||||
style={{ width: "100%" }}
|
||||
>
|
||||
{Object.entries(VectorStoreProviders).map(([providerEnum, providerDisplayName]) => {
|
||||
return (
|
||||
<Select.Option key={providerEnum} value={vectorStoreProviderMap[providerEnum]}>
|
||||
<div className="flex items-center space-x-2">
|
||||
<img
|
||||
src={vectorStoreProviderLogoMap[providerDisplayName]}
|
||||
alt={`${providerEnum} logo`}
|
||||
className="w-5 h-5"
|
||||
onError={(e) => {
|
||||
// Create a div with provider initial as fallback
|
||||
const target = e.target as HTMLImageElement;
|
||||
const parent = target.parentElement;
|
||||
if (parent) {
|
||||
const fallbackDiv = document.createElement("div");
|
||||
fallbackDiv.className =
|
||||
"w-5 h-5 rounded-full bg-gray-200 flex items-center justify-center text-xs";
|
||||
fallbackDiv.textContent = providerDisplayName.charAt(0);
|
||||
parent.replaceChild(fallbackDiv, target);
|
||||
}
|
||||
}}
|
||||
/>
|
||||
<span>{providerDisplayName}</span>
|
||||
</div>
|
||||
</Select.Option>
|
||||
);
|
||||
})}
|
||||
</Select>
|
||||
</Form.Item>
|
||||
</Form>
|
||||
|
||||
<div className="flex justify-end">
|
||||
<Button
|
||||
type="primary"
|
||||
size="large"
|
||||
onClick={handleCreateVectorStore}
|
||||
loading={isCreating}
|
||||
disabled={documents.length === 0 || !selectedProvider}
|
||||
>
|
||||
{isCreating ? "Creating Vector Store..." : "Create Vector Store"}
|
||||
</Button>
|
||||
</div>
|
||||
</div>
|
||||
</Card>
|
||||
|
||||
{/* Success Message */}
|
||||
{ingestResults.length > 0 && (
|
||||
<Alert
|
||||
message="Vector Store Created Successfully"
|
||||
description={
|
||||
<div>
|
||||
<p>
|
||||
<strong>Vector Store ID:</strong> {ingestResults[0]?.vector_store_id}
|
||||
</p>
|
||||
<p>
|
||||
<strong>Documents Ingested:</strong> {ingestResults.length}
|
||||
</p>
|
||||
</div>
|
||||
}
|
||||
type="success"
|
||||
showIcon
|
||||
closable
|
||||
/>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
||||
export default CreateVectorStore;
|
||||
@@ -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(<DocumentsTable documents={mockDocuments} onRemove={onRemove} />);
|
||||
|
||||
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(<DocumentsTable documents={mockDocuments} onRemove={onRemove} />);
|
||||
|
||||
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(<DocumentsTable documents={mockDocuments} onRemove={onRemove} />);
|
||||
|
||||
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(<DocumentsTable documents={mockDocuments} onRemove={onRemove} />);
|
||||
|
||||
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(<DocumentsTable documents={[]} onRemove={onRemove} />);
|
||||
|
||||
expect(screen.getByText(/No documents uploaded yet/)).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should have action buttons for each document", () => {
|
||||
const onRemove = vi.fn();
|
||||
render(<DocumentsTable documents={mockDocuments} onRemove={onRemove} />);
|
||||
|
||||
// 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);
|
||||
});
|
||||
});
|
||||
@@ -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<DocumentsTableProps> = ({ 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 <Badge color={config.color} text={config.text} />;
|
||||
};
|
||||
|
||||
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) => (
|
||||
<div className="flex items-center space-x-2">
|
||||
<span className="text-sm">{name}</span>
|
||||
{record.size && <span className="text-xs text-gray-400">({formatFileSize(record.size)})</span>}
|
||||
</div>
|
||||
),
|
||||
},
|
||||
{
|
||||
title: "Status",
|
||||
dataIndex: "status",
|
||||
key: "status",
|
||||
width: 150,
|
||||
render: (status: DocumentUpload["status"]) => getStatusBadge(status),
|
||||
},
|
||||
{
|
||||
title: "Actions",
|
||||
key: "actions",
|
||||
width: 120,
|
||||
render: (_: any, record: DocumentUpload) => (
|
||||
<div className="flex items-center space-x-2">
|
||||
<Tooltip title="View details">
|
||||
<EyeOutlined
|
||||
className="cursor-pointer text-gray-600 hover:text-blue-500"
|
||||
onClick={() => console.log("View", record)}
|
||||
/>
|
||||
</Tooltip>
|
||||
<Tooltip title="Copy ID">
|
||||
<CopyOutlined
|
||||
className="cursor-pointer text-gray-600 hover:text-blue-500"
|
||||
onClick={() => handleCopyId(record.uid)}
|
||||
/>
|
||||
</Tooltip>
|
||||
<Tooltip title="Remove">
|
||||
<DeleteOutlined
|
||||
className="cursor-pointer text-gray-600 hover:text-red-500"
|
||||
onClick={() => onRemove(record.uid)}
|
||||
/>
|
||||
</Tooltip>
|
||||
</div>
|
||||
),
|
||||
},
|
||||
];
|
||||
|
||||
return (
|
||||
<Table
|
||||
dataSource={documents}
|
||||
columns={columns}
|
||||
rowKey="uid"
|
||||
pagination={false}
|
||||
locale={{
|
||||
emptyText: "No documents uploaded yet. Upload documents above to get started.",
|
||||
}}
|
||||
size="small"
|
||||
/>
|
||||
);
|
||||
};
|
||||
|
||||
export default DocumentsTable;
|
||||
+90
@@ -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 }) => (
|
||||
<div data-testid="vector-store-tester">
|
||||
<div data-testid="tester-vector-store-id">{vectorStoreId}</div>
|
||||
<div data-testid="tester-access-token">{accessToken}</div>
|
||||
</div>
|
||||
),
|
||||
}));
|
||||
|
||||
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(<TestVectorStoreTab accessToken="test-token" vectorStores={mockVectorStores} />);
|
||||
|
||||
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(<TestVectorStoreTab accessToken={null} vectorStores={mockVectorStores} />);
|
||||
|
||||
expect(screen.getByText("Access token is required to test vector stores.")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should show message when no vector stores available", () => {
|
||||
render(<TestVectorStoreTab accessToken="test-token" vectorStores={[]} />);
|
||||
|
||||
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(<TestVectorStoreTab accessToken="test-token" vectorStores={mockVectorStores} />);
|
||||
|
||||
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(<TestVectorStoreTab accessToken="test-token" vectorStores={mockVectorStores} />);
|
||||
|
||||
// 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(<TestVectorStoreTab accessToken="test-token" vectorStores={mockVectorStores} />);
|
||||
|
||||
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();
|
||||
});
|
||||
});
|
||||
@@ -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<TestVectorStoreTabProps> = ({ accessToken, vectorStores }) => {
|
||||
const [selectedVectorStoreId, setSelectedVectorStoreId] = useState<string | undefined>(
|
||||
vectorStores.length > 0 ? vectorStores[0].vector_store_id : undefined
|
||||
);
|
||||
|
||||
if (!accessToken) {
|
||||
return (
|
||||
<Card>
|
||||
<Text type="secondary">Access token is required to test vector stores.</Text>
|
||||
</Card>
|
||||
);
|
||||
}
|
||||
|
||||
if (vectorStores.length === 0) {
|
||||
return (
|
||||
<Card>
|
||||
<div className="text-center py-8">
|
||||
<Text type="secondary">No vector stores available. Create one first to test it.</Text>
|
||||
</div>
|
||||
</Card>
|
||||
);
|
||||
}
|
||||
|
||||
return (
|
||||
<div className="space-y-4">
|
||||
<Card>
|
||||
<div className="space-y-4">
|
||||
<div>
|
||||
<Title level={5}>Select Vector Store</Title>
|
||||
<Text type="secondary">Choose a vector store to test search queries against</Text>
|
||||
</div>
|
||||
|
||||
<Select
|
||||
value={selectedVectorStoreId}
|
||||
onChange={setSelectedVectorStoreId}
|
||||
placeholder="Select a vector store"
|
||||
size="large"
|
||||
style={{ width: "100%" }}
|
||||
showSearch
|
||||
optionFilterProp="children"
|
||||
>
|
||||
{vectorStores.map((vs) => (
|
||||
<Select.Option key={vs.vector_store_id} value={vs.vector_store_id}>
|
||||
<div className="flex flex-col">
|
||||
<span className="font-medium">{vs.vector_store_name || vs.vector_store_id}</span>
|
||||
{vs.vector_store_name && (
|
||||
<span className="text-xs text-gray-500 font-mono">{vs.vector_store_id}</span>
|
||||
)}
|
||||
</div>
|
||||
</Select.Option>
|
||||
))}
|
||||
</Select>
|
||||
</div>
|
||||
</Card>
|
||||
|
||||
{selectedVectorStoreId && (
|
||||
<VectorStoreTester vectorStoreId={selectedVectorStoreId} accessToken={accessToken} />
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
||||
export default TestVectorStoreTab;
|
||||
+6
-6
@@ -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", () => {
|
||||
|
||||
@@ -66,6 +66,32 @@ const VectorStoreTable: React.FC<VectorStoreTableProps> = ({ 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 <span className="text-xs text-gray-400">-</span>;
|
||||
}
|
||||
|
||||
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 (
|
||||
<Tooltip title={filenames}>
|
||||
<span className="text-xs text-blue-600">{displayText}</span>
|
||||
</Tooltip>
|
||||
);
|
||||
},
|
||||
},
|
||||
{
|
||||
header: "Provider",
|
||||
accessorKey: "custom_llm_provider",
|
||||
|
||||
@@ -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<VectorStoreProps> = ({ 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<VectorStoreProps> = ({ accessToken, userID
|
||||
</div>
|
||||
|
||||
<Text className="mb-4">
|
||||
<p>You can use vector stores to store and retrieve LLM embeddings..</p>
|
||||
<p>You can use vector stores to store and retrieve LLM embeddings.</p>
|
||||
</Text>
|
||||
|
||||
<TremorButton className="mb-4" onClick={() => setIsCreateModalVisible(true)}>
|
||||
+ Add Vector Store
|
||||
</TremorButton>
|
||||
<TabGroup>
|
||||
<TabList className="mb-6">
|
||||
<Tab>Create Vector Store</Tab>
|
||||
<Tab>Manage Vector Stores</Tab>
|
||||
<Tab>Test Vector Store</Tab>
|
||||
</TabList>
|
||||
|
||||
<Grid numItems={1} className="gap-2 pt-2 pb-2 h-[75vh] w-full mt-2">
|
||||
<Col numColSpan={1}>
|
||||
<VectorStoreTable data={vectorStores} onView={handleView} onEdit={handleEdit} onDelete={handleDelete} />
|
||||
</Col>
|
||||
</Grid>
|
||||
<TabPanels>
|
||||
{/* Tab 1: Create Vector Store */}
|
||||
<TabPanel>
|
||||
<CreateVectorStore accessToken={accessToken} onSuccess={handleVectorStoreCreated} />
|
||||
</TabPanel>
|
||||
|
||||
{/* Tab 2: Manage Vector Stores */}
|
||||
<TabPanel>
|
||||
<TremorButton className="mb-4" onClick={() => setIsCreateModalVisible(true)}>
|
||||
+ Add Vector Store
|
||||
</TremorButton>
|
||||
|
||||
<Grid numItems={1} className="gap-2 pt-2 pb-2 w-full mt-2">
|
||||
<Col numColSpan={1}>
|
||||
<VectorStoreTable
|
||||
data={vectorStores}
|
||||
onView={handleView}
|
||||
onEdit={handleEdit}
|
||||
onDelete={handleDelete}
|
||||
/>
|
||||
</Col>
|
||||
</Grid>
|
||||
</TabPanel>
|
||||
|
||||
{/* Tab 3: Test Vector Store */}
|
||||
<TabPanel>
|
||||
<TestVectorStoreTab accessToken={accessToken} vectorStores={vectorStores} />
|
||||
</TabPanel>
|
||||
</TabPanels>
|
||||
</TabGroup>
|
||||
|
||||
{/* Create Vector Store Modal */}
|
||||
<VectorStoreForm
|
||||
|
||||
@@ -1,9 +1,23 @@
|
||||
export interface IngestedFile {
|
||||
file_id?: string;
|
||||
filename?: string;
|
||||
file_url?: string;
|
||||
ingested_at: string;
|
||||
file_size?: number;
|
||||
content_type?: string;
|
||||
}
|
||||
|
||||
export interface VectorStoreMetadata {
|
||||
ingested_files?: IngestedFile[];
|
||||
[key: string]: any;
|
||||
}
|
||||
|
||||
export interface VectorStore {
|
||||
vector_store_id: string;
|
||||
custom_llm_provider: string;
|
||||
vector_store_name?: string;
|
||||
vector_store_description?: string;
|
||||
vector_store_metadata?: Record<string, any>;
|
||||
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;
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user