mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-05 02:41:56 +00:00
ragIngestCall
This commit is contained in:
parent
d476e36c34
commit
1f0d54d2fc
8 changed files with 337 additions and 36 deletions
|
|
@ -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, {
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
});
|
||||
});
|
||||
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
});
|
||||
});
|
||||
|
|
@ -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;
|
||||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue