ragIngestCall

This commit is contained in:
Ishaan Jaffer 2026-01-26 15:58:00 -08:00
parent d476e36c34
commit 1f0d54d2fc
8 changed files with 337 additions and 36 deletions

View file

@ -6952,7 +6952,9 @@ export const ragIngestCall = async (
accessToken: string,
file: File,
customLlmProvider: string,
vectorStoreId?: string
vectorStoreId?: string,
vectorStoreName?: string,
vectorStoreDescription?: string
): Promise<any> => {
try {
let url = proxyBaseUrl ? `${proxyBaseUrl}/rag/ingest` : `/rag/ingest`;
@ -6960,7 +6962,7 @@ export const ragIngestCall = async (
const formData = new FormData();
formData.append("file", file);
const ingestOptions = {
const ingestOptions: any = {
ingest_options: {
vector_store: {
custom_llm_provider: customLlmProvider,
@ -6969,6 +6971,17 @@ export const ragIngestCall = async (
},
};
// 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, {

View file

@ -16,13 +16,23 @@ vi.mock("../molecules/notifications_manager", () => ({
},
}));
// Mock ProviderLogo
vi.mock("../molecules/models/ProviderLogo", () => ({
ProviderLogo: ({ provider, className }: { provider: string; className?: string }) => (
<div className={className} data-testid={`provider-logo-${provider}`}>
{provider}
</div>
),
// 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", () => {
@ -113,7 +123,7 @@ describe("CreateVectorStore", () => {
});
await waitFor(() => {
expect(mockRagIngestCall).toHaveBeenCalledWith("test-token", expect.any(File), "openai", undefined);
expect(mockRagIngestCall).toHaveBeenCalledWith("test-token", expect.any(File), "bedrock", undefined);
});
});

View file

@ -1,13 +1,16 @@
import React, { useState } from "react";
import { Card, Title, Text } from "@tremor/react";
import { Upload, Button, Select, Form, message, Alert } from "antd";
import { InboxOutlined } from "@ant-design/icons";
import { 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 { Providers, provider_map } from "../provider_info_helpers";
import { ProviderLogo } from "../molecules/models/ProviderLogo";
import {
VectorStoreProviders,
vectorStoreProviderLogoMap,
vectorStoreProviderMap,
} from "../vector_store_providers";
import NotificationsManager from "../molecules/notifications_manager";
const { Dragger } = Upload;
@ -21,7 +24,9 @@ const CreateVectorStore: React.FC<CreateVectorStoreProps> = ({ accessToken, onSu
const [form] = Form.useForm();
const [documents, setDocuments] = useState<DocumentUpload[]>([]);
const [isCreating, setIsCreating] = useState(false);
const [selectedProvider, setSelectedProvider] = useState<string>("openai");
const [selectedProvider, setSelectedProvider] = useState<string>("bedrock");
const [vectorStoreName, setVectorStoreName] = useState<string>("");
const [vectorStoreDescription, setVectorStoreDescription] = useState<string>("");
const [ingestResults, setIngestResults] = useState<RAGIngestResponse[]>([]);
const uploadProps: UploadProps = {
@ -111,7 +116,9 @@ const CreateVectorStore: React.FC<CreateVectorStoreProps> = ({ accessToken, onSu
accessToken,
doc.originFileObj,
selectedProvider,
vectorStoreId // Use the same vector store ID for subsequent uploads
vectorStoreId, // Use the same vector store ID for subsequent uploads
vectorStoreName || undefined,
vectorStoreDescription || undefined
);
// Store the vector store ID from the first successful ingest
@ -157,8 +164,6 @@ const CreateVectorStore: React.FC<CreateVectorStoreProps> = ({ accessToken, onSu
}
};
const supportedProviders = ["openai", "bedrock"]; // Add more as needed
return (
<div className="space-y-6">
<div>
@ -197,18 +202,67 @@ const CreateVectorStore: React.FC<CreateVectorStoreProps> = ({ accessToken, onSu
</Card>
)}
{/* Provider Selection and Create Button */}
{/* Provider Selection and Vector Store Details */}
<Card>
<div className="space-y-4">
<div>
<Text className="font-medium">Step 2: Select Provider</Text>
<Text className="font-medium">Step 2: Configure Vector Store</Text>
<Text className="text-sm text-gray-500 block mt-1">
Choose the LLM provider for embedding and vector store operations.
Choose the provider and optionally provide a name and description for your vector store.
</Text>
</div>
<Form form={form} layout="vertical">
<Form.Item label="Provider" required>
<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}
@ -216,19 +270,32 @@ const CreateVectorStore: React.FC<CreateVectorStoreProps> = ({ accessToken, onSu
size="large"
style={{ width: "100%" }}
>
{Object.entries(Providers)
.filter(([providerEnum]) => supportedProviders.includes(provider_map[providerEnum]?.toLowerCase()))
.map(([providerEnum, providerDisplayName]) => {
const providerValue = provider_map[providerEnum]?.toLowerCase();
return (
<Select.Option key={providerEnum} value={providerValue}>
<div className="flex items-center space-x-2">
<ProviderLogo provider={providerEnum} className="w-5 h-5" />
<span>{providerDisplayName}</span>
</div>
</Select.Option>
);
})}
{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>

View file

@ -0,0 +1,89 @@
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);
expect(screen.getByText("Test Store 1")).toBeInTheDocument();
expect(screen.getByText("Test Store 2")).toBeInTheDocument();
});
});

View file

@ -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;

View file

@ -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",

View file

@ -8,6 +8,7 @@ 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";
@ -148,6 +149,7 @@ const VectorStoreManagement: React.FC<VectorStoreProps> = ({ accessToken, userID
<TabList className="mb-6">
<Tab>Create Vector Store</Tab>
<Tab>Manage Vector Stores</Tab>
<Tab>Test Vector Store</Tab>
</TabList>
<TabPanels>
@ -173,6 +175,11 @@ const VectorStoreManagement: React.FC<VectorStoreProps> = ({ accessToken, userID
</Col>
</Grid>
</TabPanel>
{/* Tab 3: Test Vector Store */}
<TabPanel>
<TestVectorStoreTab accessToken={accessToken} vectorStores={vectorStores} />
</TabPanel>
</TabPanels>
</TabGroup>

View file

@ -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;