mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-14 02:22:54 +00:00
[Feat] LLM API Endpoint - Expose OpenAI Compatible /vector_stores/{vector_store_id}/search endpoint (#12749)
* fix _pass_through_endpoint_without_required_model * add get_litellm_managed_vector_store_from_registry * undo router change * fix for using router + vector search methods * add simple helper for _update_request_data_with_litellm_managed_vector_store_registry * add vector_stores routes * test_router_avector_store_search_passes_correct_args * [Feat] UI - Allow clicking into Vector Stores (#12741) * Add View Vector Store * add /info for vector store * fix updated_at * allow easily testing the KB on litellm * fix * rename test * test_init_vector_store_api_endpoints * test_update_request_data_with_litellm_managed_vector_store_registry
This commit is contained in:
Ishaan Jaff
@@ -307,6 +307,12 @@ class LiteLLMRoutes(enum.Enum):
|
||||
"/v1/responses/{response_id}",
|
||||
"/responses/{response_id}/input_items",
|
||||
"/v1/responses/{response_id}/input_items",
|
||||
|
||||
# vector stores
|
||||
"/vector_stores",
|
||||
"/v1/vector_stores",
|
||||
"/vector_stores/{vector_store_id}/search",
|
||||
"/v1/vector_stores/{vector_store_id}/search",
|
||||
]
|
||||
|
||||
mapped_pass_through_routes = [
|
||||
|
||||
@@ -1,5 +1,11 @@
|
||||
from typing import Dict, Optional
|
||||
|
||||
from fastapi import APIRouter, Depends, Request, Response
|
||||
|
||||
import litellm
|
||||
from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import (
|
||||
LiteLLM_ManagedVectorStore,
|
||||
)
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
@@ -9,6 +15,30 @@ router = APIRouter()
|
||||
# OpenAI Compatible Endpoints
|
||||
########################################################
|
||||
|
||||
def _update_request_data_with_litellm_managed_vector_store_registry(
|
||||
data: Dict,
|
||||
vector_store_id: str,
|
||||
) -> Dict:
|
||||
"""
|
||||
Update the request data with the litellm managed vector store registry.
|
||||
|
||||
"""
|
||||
if litellm.vector_store_registry is not None:
|
||||
vector_store_to_run: Optional[LiteLLM_ManagedVectorStore] = litellm.vector_store_registry.get_litellm_managed_vector_store_from_registry(
|
||||
vector_store_id=vector_store_id
|
||||
)
|
||||
if vector_store_to_run is not None:
|
||||
if "custom_llm_provider" in vector_store_to_run:
|
||||
data["custom_llm_provider"] = vector_store_to_run.get("custom_llm_provider")
|
||||
|
||||
if "litellm_credential_name" in vector_store_to_run:
|
||||
data["litellm_credential_name"] = vector_store_to_run.get("litellm_credential_name")
|
||||
|
||||
if "litellm_params" in vector_store_to_run:
|
||||
litellm_params = vector_store_to_run.get("litellm_params", {}) or {}
|
||||
data.update(litellm_params)
|
||||
return data
|
||||
|
||||
@router.post("/v1/vector_stores/{vector_store_id}/search", dependencies=[Depends(user_api_key_auth)])
|
||||
@router.post("/vector_stores/{vector_store_id}/search", dependencies=[Depends(user_api_key_auth)])
|
||||
async def vector_store_search(
|
||||
@@ -41,6 +71,12 @@ async def vector_store_search(
|
||||
data = await _read_request_body(request=request)
|
||||
if "vector_store_id" not in data:
|
||||
data["vector_store_id"] = vector_store_id
|
||||
|
||||
data = _update_request_data_with_litellm_managed_vector_store_registry(
|
||||
data=data,
|
||||
vector_store_id=vector_store_id
|
||||
)
|
||||
|
||||
processor = ProxyBaseLLMRequestProcessing(data=data)
|
||||
try:
|
||||
return await processor.base_process_llm_request(
|
||||
|
||||
+26
-5
@@ -3262,6 +3262,7 @@ class Router:
|
||||
async def _pass_through_moderation_endpoint_factory(
|
||||
self,
|
||||
original_function: Callable,
|
||||
custom_llm_provider: Optional[str] = None,
|
||||
**kwargs,
|
||||
):
|
||||
# update kwargs with model_group
|
||||
@@ -3328,7 +3329,7 @@ class Router:
|
||||
|
||||
def sync_wrapper(
|
||||
custom_llm_provider: Optional[
|
||||
Literal["openai", "azure", "anthropic"]
|
||||
str
|
||||
] = None,
|
||||
client: Optional[Any] = None,
|
||||
**kwargs,
|
||||
@@ -3342,7 +3343,7 @@ class Router:
|
||||
# Handle asynchronous call types
|
||||
async def async_wrapper(
|
||||
custom_llm_provider: Optional[
|
||||
Literal["openai", "azure", "anthropic"]
|
||||
str
|
||||
] = None,
|
||||
client: Optional[Any] = None,
|
||||
**kwargs,
|
||||
@@ -3370,8 +3371,6 @@ class Router:
|
||||
"aimage_edit",
|
||||
"agenerate_content",
|
||||
"agenerate_content_stream",
|
||||
"avector_store_search",
|
||||
"avector_store_create",
|
||||
):
|
||||
return await self._ageneric_api_call_with_fallbacks(
|
||||
original_function=original_function,
|
||||
@@ -3392,6 +3391,15 @@ class Router:
|
||||
original_function=original_function,
|
||||
**kwargs,
|
||||
)
|
||||
elif call_type in (
|
||||
"avector_store_search",
|
||||
"avector_store_create",
|
||||
):
|
||||
return await self._init_vector_store_api_endpoints(
|
||||
original_function=original_function,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
**kwargs,
|
||||
)
|
||||
elif call_type in ("afile_delete", "afile_content"):
|
||||
return await self._ageneric_api_call_with_fallbacks(
|
||||
original_function=original_function,
|
||||
@@ -3401,6 +3409,19 @@ class Router:
|
||||
)
|
||||
|
||||
return async_wrapper
|
||||
|
||||
async def _init_vector_store_api_endpoints(
|
||||
self,
|
||||
original_function: Callable,
|
||||
custom_llm_provider: Optional[str] = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
Initialize the Vector Store API endpoints on the router.
|
||||
"""
|
||||
if custom_llm_provider and "custom_llm_provider" not in kwargs:
|
||||
kwargs["custom_llm_provider"] = custom_llm_provider
|
||||
return await original_function(**kwargs)
|
||||
|
||||
async def _init_responses_api_endpoints(
|
||||
self,
|
||||
@@ -3427,7 +3448,7 @@ class Router:
|
||||
async def _pass_through_assistants_endpoint_factory(
|
||||
self,
|
||||
original_function: Callable,
|
||||
custom_llm_provider: Optional[Literal["openai", "azure", "anthropic"]] = None,
|
||||
custom_llm_provider: Optional[str] = None,
|
||||
client: Optional[AsyncOpenAI] = None,
|
||||
**kwargs,
|
||||
):
|
||||
|
||||
@@ -128,6 +128,17 @@ class VectorStoreRegistry:
|
||||
return vector_store
|
||||
return None
|
||||
|
||||
def get_litellm_managed_vector_store_from_registry(
|
||||
self, vector_store_id: str
|
||||
) -> Optional[LiteLLM_ManagedVectorStore]:
|
||||
"""
|
||||
Returns the vector store from the registry
|
||||
"""
|
||||
for vector_store in self.vector_stores:
|
||||
if vector_store.get("vector_store_id") == vector_store_id:
|
||||
return vector_store
|
||||
return None
|
||||
|
||||
def pop_vector_stores_to_run(
|
||||
self, non_default_params: Dict, tools: Optional[List[Dict]] = None
|
||||
) -> List[LiteLLM_ManagedVectorStore]:
|
||||
|
||||
@@ -617,3 +617,51 @@ async def test_init_responses_api_endpoints():
|
||||
assert second_call_kwargs["model"] == "claude-3-sonnet"
|
||||
assert second_call_kwargs["response_id"] == "resp_claude_123"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_init_vector_store_api_endpoints():
|
||||
"""
|
||||
Test that _init_vector_store_api_endpoints correctly passes custom_llm_provider to kwargs
|
||||
"""
|
||||
# Create a router with a basic model
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "test-model",
|
||||
"litellm_params": {
|
||||
"model": "openai/test-model",
|
||||
"api_key": "fake-api-key",
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
# Mock the original function
|
||||
mock_original_function = AsyncMock(return_value={"status": "success"})
|
||||
|
||||
# Call without custom_llm_provider
|
||||
result = await router._init_vector_store_api_endpoints(
|
||||
original_function=mock_original_function,
|
||||
vector_store_id="test-store"
|
||||
)
|
||||
|
||||
# Verify original function was called with correct kwargs
|
||||
mock_original_function.assert_called_once_with(vector_store_id="test-store")
|
||||
assert result == {"status": "success"}
|
||||
|
||||
# Reset the mock
|
||||
mock_original_function.reset_mock()
|
||||
|
||||
# Call with custom_llm_provider
|
||||
await router._init_vector_store_api_endpoints(
|
||||
original_function=mock_original_function,
|
||||
custom_llm_provider="openai",
|
||||
vector_store_id="test-store"
|
||||
)
|
||||
|
||||
# Verify custom_llm_provider was added to kwargs
|
||||
mock_original_function.assert_called_once_with(
|
||||
vector_store_id="test-store",
|
||||
custom_llm_provider="openai"
|
||||
)
|
||||
|
||||
|
||||
@@ -0,0 +1,102 @@
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(
|
||||
0, os.path.abspath("../../../../..")
|
||||
) # Adds the parent directory to the system path
|
||||
|
||||
import litellm
|
||||
from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import (
|
||||
LiteLLM_ManagedVectorStore,
|
||||
)
|
||||
from litellm.proxy.vector_store_endpoints.endpoints import (
|
||||
_update_request_data_with_litellm_managed_vector_store_registry,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_router_avector_store_search_passes_correct_args():
|
||||
"""
|
||||
Test that router.avector_store_search() passes the correct arguments
|
||||
to downstream litellm.vector_stores.asearch() with custom_llm_provider and query.
|
||||
"""
|
||||
# Create a router
|
||||
router = litellm.Router(model_list=[])
|
||||
|
||||
# Mock the router's _init_vector_store_api_endpoints method to avoid real API calls
|
||||
with patch.object(router, '_init_vector_store_api_endpoints') as mock_init:
|
||||
mock_init.return_value = {
|
||||
"object": "vector_store.search_results.page",
|
||||
"search_query": "test query",
|
||||
"data": []
|
||||
}
|
||||
|
||||
# Call router's avector_store_search
|
||||
result = await router.avector_store_search(
|
||||
vector_store_id="test_store_id",
|
||||
query="test query",
|
||||
custom_llm_provider="bedrock"
|
||||
)
|
||||
|
||||
# Verify the internal method was called with correct args
|
||||
mock_init.assert_called_once()
|
||||
call_args = mock_init.call_args
|
||||
|
||||
# Check that the original function is passed correctly
|
||||
assert call_args[1]["vector_store_id"] == "test_store_id"
|
||||
assert call_args[1]["query"] == "test query"
|
||||
assert call_args[1]["custom_llm_provider"] == "bedrock"
|
||||
|
||||
|
||||
def test_update_request_data_with_litellm_managed_vector_store_registry():
|
||||
"""
|
||||
Test that _update_request_data_with_litellm_managed_vector_store_registry
|
||||
correctly updates request data with vector store registry information.
|
||||
"""
|
||||
# Setup test data
|
||||
data = {"existing_key": "existing_value"}
|
||||
vector_store_id = "test_store_id"
|
||||
|
||||
# Mock vector store registry
|
||||
mock_vector_store: LiteLLM_ManagedVectorStore = {
|
||||
"vector_store_id": "test_store_id",
|
||||
"custom_llm_provider": "bedrock",
|
||||
"litellm_credential_name": "test_credential",
|
||||
"litellm_params": {"api_key": "test_key", "aws_region_name": "us-east-1"}
|
||||
}
|
||||
|
||||
mock_registry = MagicMock()
|
||||
mock_registry.get_litellm_managed_vector_store_from_registry.return_value = mock_vector_store
|
||||
|
||||
# Test with vector store registry
|
||||
with patch.object(litellm, 'vector_store_registry', mock_registry):
|
||||
result = _update_request_data_with_litellm_managed_vector_store_registry(
|
||||
data=data,
|
||||
vector_store_id=vector_store_id
|
||||
)
|
||||
|
||||
# Verify the data was updated correctly
|
||||
assert result["existing_key"] == "existing_value" # Original data preserved
|
||||
assert result["custom_llm_provider"] == "bedrock"
|
||||
assert result["litellm_credential_name"] == "test_credential"
|
||||
assert result["api_key"] == "test_key"
|
||||
assert result["aws_region_name"] == "us-east-1"
|
||||
|
||||
# Verify registry was called correctly
|
||||
mock_registry.get_litellm_managed_vector_store_from_registry.assert_called_once_with(
|
||||
vector_store_id="test_store_id"
|
||||
)
|
||||
|
||||
# Test with no vector store registry
|
||||
with patch.object(litellm, 'vector_store_registry', None):
|
||||
original_data = {"existing_key": "existing_value"}
|
||||
result = _update_request_data_with_litellm_managed_vector_store_registry(
|
||||
data=original_data,
|
||||
vector_store_id=vector_store_id
|
||||
)
|
||||
|
||||
# Verify data remains unchanged when no registry
|
||||
assert result == original_data
|
||||
@@ -98,4 +98,17 @@ class TestBedrockVectorStore(BaseVectorStoreTest):
|
||||
|
||||
# Test with None
|
||||
attributes = config._get_attributes_from_metadata(None)
|
||||
assert attributes == {}
|
||||
assert attributes == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bedrock_search_with_router():
|
||||
from litellm.router import Router
|
||||
# init router
|
||||
_router = Router(model_list=[])
|
||||
search_response = await _router.avector_store_search(
|
||||
query="what happens after we add a model",
|
||||
vector_store_id="T37J8R4WTM",
|
||||
custom_llm_provider="bedrock",
|
||||
)
|
||||
print(search_response)
|
||||
@@ -6226,3 +6226,35 @@ export const testMCPToolsListRequest = async (
|
||||
throw error;
|
||||
}
|
||||
};
|
||||
|
||||
export const vectorStoreSearchCall = async (
|
||||
accessToken: string,
|
||||
vectorStoreId: string,
|
||||
query: string
|
||||
): Promise<any> => {
|
||||
try {
|
||||
const url = `${getProxyBaseUrl()}/v1/vector_stores/${vectorStoreId}/search`;
|
||||
const response = await fetch(url, {
|
||||
method: "POST",
|
||||
headers: {
|
||||
"Authorization": `Bearer ${accessToken}`,
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
body: JSON.stringify({
|
||||
query: query
|
||||
}),
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
const errorData = await response.text();
|
||||
await handleError(errorData);
|
||||
return null;
|
||||
}
|
||||
|
||||
const data = await response.json();
|
||||
return data;
|
||||
} catch (error) {
|
||||
console.error("Error testing vector store search:", error);
|
||||
throw error;
|
||||
}
|
||||
};
|
||||
|
||||
@@ -0,0 +1,280 @@
|
||||
import React, { useState } from "react";
|
||||
import { Button, Input, Card, Typography, Spin, message, Divider } from "antd";
|
||||
import { SendOutlined, DatabaseOutlined, LoadingOutlined, DownOutlined, RightOutlined } from "@ant-design/icons";
|
||||
import { vectorStoreSearchCall } from "../networking";
|
||||
|
||||
const { TextArea } = Input;
|
||||
const { Text, Title } = Typography;
|
||||
|
||||
interface VectorStoreContent {
|
||||
text: string;
|
||||
type: string;
|
||||
}
|
||||
|
||||
interface VectorStoreResult {
|
||||
score: number;
|
||||
content: VectorStoreContent[];
|
||||
file_id?: string;
|
||||
filename?: string;
|
||||
attributes?: Record<string, any>;
|
||||
}
|
||||
|
||||
interface VectorStoreSearchResponse {
|
||||
object: string;
|
||||
search_query: string;
|
||||
data: VectorStoreResult[];
|
||||
}
|
||||
|
||||
interface VectorStoreTesterProps {
|
||||
vectorStoreId: string;
|
||||
accessToken: string;
|
||||
className?: string;
|
||||
}
|
||||
|
||||
export const VectorStoreTester: React.FC<VectorStoreTesterProps> = ({
|
||||
vectorStoreId,
|
||||
accessToken,
|
||||
className = "",
|
||||
}) => {
|
||||
const [query, setQuery] = useState("");
|
||||
const [isLoading, setIsLoading] = useState(false);
|
||||
const [searchHistory, setSearchHistory] = useState<{
|
||||
query: string;
|
||||
response: VectorStoreSearchResponse | null;
|
||||
timestamp: number;
|
||||
}[]>([]);
|
||||
const [expandedResults, setExpandedResults] = useState<Record<string, boolean>>({});
|
||||
|
||||
const handleSearch = async () => {
|
||||
if (!query.trim()) {
|
||||
message.warning("Please enter a search query");
|
||||
return;
|
||||
}
|
||||
|
||||
setIsLoading(true);
|
||||
|
||||
try {
|
||||
const response = await vectorStoreSearchCall(accessToken, vectorStoreId, query);
|
||||
|
||||
const historyEntry = {
|
||||
query,
|
||||
response,
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
|
||||
setSearchHistory(prev => [historyEntry, ...prev]);
|
||||
setQuery("");
|
||||
} catch (error) {
|
||||
console.error("Error searching vector store:", error);
|
||||
message.error("Failed to search vector store");
|
||||
} finally {
|
||||
setIsLoading(false);
|
||||
}
|
||||
};
|
||||
|
||||
const handleKeyDown = (event: React.KeyboardEvent<HTMLTextAreaElement>) => {
|
||||
if (event.key === 'Enter' && !event.shiftKey) {
|
||||
event.preventDefault();
|
||||
handleSearch();
|
||||
}
|
||||
};
|
||||
|
||||
const formatTimestamp = (timestamp: number): string => {
|
||||
return new Date(timestamp).toLocaleString();
|
||||
};
|
||||
|
||||
const clearHistory = () => {
|
||||
setSearchHistory([]);
|
||||
setExpandedResults({});
|
||||
message.success("Search history cleared");
|
||||
};
|
||||
|
||||
const toggleResultExpansion = (historyIndex: number, resultIndex: number) => {
|
||||
const key = `${historyIndex}-${resultIndex}`;
|
||||
setExpandedResults(prev => ({
|
||||
...prev,
|
||||
[key]: !prev[key]
|
||||
}));
|
||||
};
|
||||
|
||||
return (
|
||||
<Card className="w-full rounded-xl shadow-md">
|
||||
<div className="flex flex-col h-[600px]">
|
||||
{/* Header */}
|
||||
<div className="p-4 border-b border-gray-200 flex justify-between items-center">
|
||||
<div className="flex items-center">
|
||||
<DatabaseOutlined className="mr-2 text-blue-500" />
|
||||
<Title level={4} className="mb-0">Test Vector Store</Title>
|
||||
</div>
|
||||
{searchHistory.length > 0 && (
|
||||
<Button onClick={clearHistory} size="small">
|
||||
Clear History
|
||||
</Button>
|
||||
)}
|
||||
</div>
|
||||
|
||||
{/* Results Area */}
|
||||
<div className="flex-1 overflow-auto p-4 pb-0">
|
||||
{searchHistory.length === 0 ? (
|
||||
<div className="h-full flex flex-col items-center justify-center text-gray-400">
|
||||
<DatabaseOutlined style={{ fontSize: '48px', marginBottom: '16px' }} />
|
||||
<Text>Test your vector store by entering a search query below</Text>
|
||||
</div>
|
||||
) : (
|
||||
<div className="space-y-4">
|
||||
{searchHistory.map((entry, index) => (
|
||||
<div key={index} className="space-y-2">
|
||||
{/* User Query */}
|
||||
<div className="text-right">
|
||||
<div className="inline-block max-w-[80%] rounded-lg shadow-sm p-3 bg-blue-50 border border-blue-200">
|
||||
<div className="flex items-center gap-2 mb-1">
|
||||
<strong className="text-sm">Query</strong>
|
||||
<span className="text-xs text-gray-500">
|
||||
{formatTimestamp(entry.timestamp)}
|
||||
</span>
|
||||
</div>
|
||||
<div className="text-left">
|
||||
{entry.query}
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{/* Vector Store Response */}
|
||||
<div className="text-left">
|
||||
<div className="inline-block max-w-[80%] rounded-lg shadow-sm p-3 bg-white border border-gray-200">
|
||||
<div className="flex items-center gap-2 mb-2">
|
||||
<DatabaseOutlined className="text-green-500" />
|
||||
<strong className="text-sm">Vector Store Results</strong>
|
||||
{entry.response && (
|
||||
<span className="text-xs px-2 py-0.5 rounded bg-gray-100 text-gray-600">
|
||||
{entry.response.data?.length || 0} results
|
||||
</span>
|
||||
)}
|
||||
</div>
|
||||
|
||||
{entry.response && entry.response.data && entry.response.data.length > 0 ? (
|
||||
<div className="space-y-3">
|
||||
{entry.response.data.map((result, resultIndex) => {
|
||||
const isExpanded = expandedResults[`${index}-${resultIndex}`] || false;
|
||||
|
||||
return (
|
||||
<div key={resultIndex} className="border rounded-lg overflow-hidden bg-gray-50">
|
||||
{/* Clickable Header */}
|
||||
<div
|
||||
className="flex justify-between items-center p-3 cursor-pointer hover:bg-gray-100 transition-colors"
|
||||
onClick={() => toggleResultExpansion(index, resultIndex)}
|
||||
>
|
||||
<div className="flex items-center">
|
||||
{isExpanded ? (
|
||||
<DownOutlined className="text-gray-500 mr-2" />
|
||||
) : (
|
||||
<RightOutlined className="text-gray-500 mr-2" />
|
||||
)}
|
||||
<span className="font-medium text-sm">Result {resultIndex + 1}</span>
|
||||
{/* Show preview of content when collapsed */}
|
||||
{!isExpanded && result.content && result.content[0] && (
|
||||
<span className="ml-2 text-xs text-gray-500 truncate max-w-md">
|
||||
- {result.content[0].text.substring(0, 100)}...
|
||||
</span>
|
||||
)}
|
||||
</div>
|
||||
<span className="text-xs bg-blue-100 text-blue-800 px-2 py-1 rounded">
|
||||
Score: {result.score.toFixed(4)}
|
||||
</span>
|
||||
</div>
|
||||
|
||||
{/* Expandable Content */}
|
||||
{isExpanded && (
|
||||
<div className="border-t bg-white p-3">
|
||||
{/* Content */}
|
||||
{result.content && result.content.map((content, contentIndex) => (
|
||||
<div key={contentIndex} className="mb-3">
|
||||
<div className="text-xs text-gray-500 mb-1">Content ({content.type})</div>
|
||||
<div className="text-sm bg-gray-50 p-3 rounded border text-gray-800 max-h-40 overflow-y-auto">
|
||||
{content.text}
|
||||
</div>
|
||||
</div>
|
||||
))}
|
||||
|
||||
{/* Metadata */}
|
||||
{(result.file_id || result.filename || result.attributes) && (
|
||||
<div className="mt-3 pt-3 border-t border-gray-200">
|
||||
<div className="text-xs text-gray-500 mb-2 font-medium">Metadata</div>
|
||||
<div className="space-y-2 text-xs">
|
||||
{result.file_id && (
|
||||
<div className="bg-gray-50 p-2 rounded">
|
||||
<span className="font-medium">File ID:</span> {result.file_id}
|
||||
</div>
|
||||
)}
|
||||
{result.filename && (
|
||||
<div className="bg-gray-50 p-2 rounded">
|
||||
<span className="font-medium">Filename:</span> {result.filename}
|
||||
</div>
|
||||
)}
|
||||
{result.attributes && Object.keys(result.attributes).length > 0 && (
|
||||
<div className="bg-gray-50 p-2 rounded">
|
||||
<span className="font-medium block mb-1">Attributes:</span>
|
||||
<pre className="text-xs bg-white p-2 rounded border overflow-x-auto">
|
||||
{JSON.stringify(result.attributes, null, 2)}
|
||||
</pre>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
})}
|
||||
</div>
|
||||
) : (
|
||||
<div className="text-gray-500 text-sm">No results found</div>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{index < searchHistory.length - 1 && <Divider />}
|
||||
</div>
|
||||
))}
|
||||
</div>
|
||||
)}
|
||||
|
||||
{isLoading && (
|
||||
<div className="flex justify-center items-center my-4">
|
||||
<Spin indicator={<LoadingOutlined style={{ fontSize: 24 }} spin />} />
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
|
||||
{/* Input Area */}
|
||||
<div className="p-4 border-t border-gray-200 bg-white">
|
||||
<div className="flex items-end space-x-2">
|
||||
<div className="flex-1">
|
||||
<TextArea
|
||||
value={query}
|
||||
onChange={(e) => setQuery(e.target.value)}
|
||||
onKeyDown={handleKeyDown}
|
||||
placeholder="Enter your search query... (Shift+Enter for new line)"
|
||||
disabled={isLoading}
|
||||
autoSize={{ minRows: 1, maxRows: 4 }}
|
||||
style={{ resize: 'none' }}
|
||||
/>
|
||||
</div>
|
||||
<Button
|
||||
type="primary"
|
||||
onClick={handleSearch}
|
||||
disabled={isLoading || !query.trim()}
|
||||
icon={<SendOutlined />}
|
||||
loading={isLoading}
|
||||
>
|
||||
Search
|
||||
</Button>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</Card>
|
||||
);
|
||||
};
|
||||
|
||||
export default VectorStoreTester;
|
||||
+250
-212
@@ -5,6 +5,11 @@ import {
|
||||
Title,
|
||||
Button,
|
||||
Badge,
|
||||
TabGroup,
|
||||
TabList,
|
||||
Tab,
|
||||
TabPanels,
|
||||
TabPanel,
|
||||
} from "@tremor/react";
|
||||
import {
|
||||
Form,
|
||||
@@ -19,6 +24,7 @@ import { ArrowLeftIcon } from "@heroicons/react/outline";
|
||||
import { vectorStoreInfoCall, vectorStoreUpdateCall, credentialListCall, CredentialItem } from "../networking";
|
||||
import { VectorStore } from "./types";
|
||||
import { Providers, providerLogoMap, provider_map } from "../provider_info_helpers";
|
||||
import VectorStoreTester from "./VectorStoreTester";
|
||||
|
||||
interface VectorStoreInfoViewProps {
|
||||
vectorStoreId: string;
|
||||
@@ -40,6 +46,7 @@ const VectorStoreInfoView: React.FC<VectorStoreInfoViewProps> = ({
|
||||
const [isEditing, setIsEditing] = useState<boolean>(editVectorStore);
|
||||
const [metadataString, setMetadataString] = useState<string>("{}");
|
||||
const [credentials, setCredentials] = useState<CredentialItem[]>([]);
|
||||
const [activeTab, setActiveTab] = useState<string>(editVectorStore ? "details" : "details");
|
||||
|
||||
const fetchVectorStoreDetails = async () => {
|
||||
if (!accessToken) return;
|
||||
@@ -134,224 +141,255 @@ const VectorStoreInfoView: React.FC<VectorStoreInfoViewProps> = ({
|
||||
{is_admin && !isEditing && (
|
||||
<Button onClick={() => setIsEditing(true)}>Edit Vector Store</Button>
|
||||
)}
|
||||
|
||||
</div>
|
||||
|
||||
{isEditing ? (
|
||||
<Card>
|
||||
<Form
|
||||
form={form}
|
||||
onFinish={handleSave}
|
||||
layout="vertical"
|
||||
initialValues={vectorStoreDetails}
|
||||
>
|
||||
<Form.Item
|
||||
label="Vector Store ID"
|
||||
name="vector_store_id"
|
||||
rules={[{ required: true, message: "Please input a vector store ID" }]}
|
||||
>
|
||||
<Input disabled />
|
||||
</Form.Item>
|
||||
|
||||
<Form.Item
|
||||
label="Vector Store Name"
|
||||
name="vector_store_name"
|
||||
>
|
||||
<Input />
|
||||
</Form.Item>
|
||||
|
||||
<Form.Item
|
||||
label="Description"
|
||||
name="vector_store_description"
|
||||
>
|
||||
<Input.TextArea rows={4} />
|
||||
</Form.Item>
|
||||
|
||||
<Form.Item
|
||||
label={
|
||||
<span>
|
||||
Provider{' '}
|
||||
<Tooltip title="Select the provider for this vector store">
|
||||
<InfoCircleOutlined style={{ marginLeft: '4px' }} />
|
||||
</Tooltip>
|
||||
</span>
|
||||
}
|
||||
name="custom_llm_provider"
|
||||
rules={[{ required: true, message: "Please select a provider" }]}
|
||||
>
|
||||
<Select2>
|
||||
{Object.entries(Providers).map(([providerEnum, providerDisplayName]) => {
|
||||
// Currently only showing Bedrock since it's the only supported provider
|
||||
if (providerEnum === 'Bedrock') {
|
||||
return (
|
||||
<Select2.Option key={providerEnum} value={provider_map[providerEnum]}>
|
||||
<div className="flex items-center space-x-2">
|
||||
<img
|
||||
src={providerLogoMap[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>
|
||||
</Select2.Option>
|
||||
);
|
||||
}
|
||||
return null;
|
||||
})}
|
||||
</Select2>
|
||||
</Form.Item>
|
||||
|
||||
{/* Credentials */}
|
||||
<div className="mb-4">
|
||||
<Text className="text-sm text-gray-500 mb-2">
|
||||
Either select existing credentials OR enter provider credentials below
|
||||
</Text>
|
||||
</div>
|
||||
|
||||
<Form.Item
|
||||
label="Existing Credentials"
|
||||
name="litellm_credential_name"
|
||||
>
|
||||
<Select2
|
||||
showSearch
|
||||
placeholder="Select or search for existing credentials"
|
||||
optionFilterProp="children"
|
||||
filterOption={(input, option) =>
|
||||
(option?.label ?? '').toLowerCase().includes(input.toLowerCase())
|
||||
}
|
||||
options={[
|
||||
{ value: null, label: 'None' },
|
||||
...credentials.map((credential) => ({
|
||||
value: credential.credential_name,
|
||||
label: credential.credential_name
|
||||
}))
|
||||
]}
|
||||
allowClear
|
||||
/>
|
||||
</Form.Item>
|
||||
|
||||
<div className="flex items-center my-4">
|
||||
<div className="flex-grow border-t border-gray-200"></div>
|
||||
<span className="px-4 text-gray-500 text-sm">OR</span>
|
||||
<div className="flex-grow border-t border-gray-200"></div>
|
||||
</div>
|
||||
|
||||
<Form.Item
|
||||
label={
|
||||
<span>
|
||||
Metadata{' '}
|
||||
<Tooltip title="JSON metadata for the vector store">
|
||||
<InfoCircleOutlined style={{ marginLeft: '4px' }} />
|
||||
</Tooltip>
|
||||
</span>
|
||||
}
|
||||
>
|
||||
<Input.TextArea
|
||||
rows={4}
|
||||
value={metadataString}
|
||||
onChange={(e) => setMetadataString(e.target.value)}
|
||||
placeholder='{"key": "value"}'
|
||||
/>
|
||||
</Form.Item>
|
||||
|
||||
<div className="flex justify-end space-x-2">
|
||||
<AntButton onClick={() => setIsEditing(false)}>Cancel</AntButton>
|
||||
<AntButton type="primary" htmlType="submit">Save Changes</AntButton>
|
||||
</div>
|
||||
</Form>
|
||||
</Card>
|
||||
) : (
|
||||
<div className="space-y-6">
|
||||
<Card>
|
||||
<Title>Vector Store Details</Title>
|
||||
<div className="space-y-4 mt-4">
|
||||
<TabGroup>
|
||||
<TabList className="mb-6">
|
||||
<Tab>Details</Tab>
|
||||
<Tab>Test Vector Store</Tab>
|
||||
</TabList>
|
||||
|
||||
<TabPanels>
|
||||
{/* Details Tab */}
|
||||
<TabPanel>
|
||||
{isEditing ? (
|
||||
<div>
|
||||
<Text className="font-medium">ID</Text>
|
||||
<Text>{vectorStoreDetails.vector_store_id}</Text>
|
||||
</div>
|
||||
<div>
|
||||
<Text className="font-medium">Name</Text>
|
||||
<Text>{vectorStoreDetails.vector_store_name || "-"}</Text>
|
||||
</div>
|
||||
<div>
|
||||
<Text className="font-medium">Description</Text>
|
||||
<Text>{vectorStoreDetails.vector_store_description || "-"}</Text>
|
||||
</div>
|
||||
<div>
|
||||
<Text className="font-medium">Provider</Text>
|
||||
<div className="flex items-center space-x-2 mt-1">
|
||||
{(() => {
|
||||
const provider = vectorStoreDetails.custom_llm_provider || "bedrock";
|
||||
const { displayName, logo } = (() => {
|
||||
// Find the enum key by matching provider_map values
|
||||
const enumKey = Object.keys(provider_map).find(
|
||||
key => provider_map[key].toLowerCase() === provider.toLowerCase()
|
||||
);
|
||||
|
||||
if (!enumKey) {
|
||||
return { displayName: provider, logo: "" };
|
||||
<div className="flex justify-between items-center mb-4">
|
||||
<Title>Edit Vector Store</Title>
|
||||
</div>
|
||||
<Card>
|
||||
<Form
|
||||
form={form}
|
||||
onFinish={handleSave}
|
||||
layout="vertical"
|
||||
initialValues={vectorStoreDetails}
|
||||
>
|
||||
<Form.Item
|
||||
label="Vector Store ID"
|
||||
name="vector_store_id"
|
||||
rules={[{ required: true, message: "Please input a vector store ID" }]}
|
||||
>
|
||||
<Input disabled />
|
||||
</Form.Item>
|
||||
|
||||
<Form.Item
|
||||
label="Vector Store Name"
|
||||
name="vector_store_name"
|
||||
>
|
||||
<Input />
|
||||
</Form.Item>
|
||||
|
||||
<Form.Item
|
||||
label="Description"
|
||||
name="vector_store_description"
|
||||
>
|
||||
<Input.TextArea rows={4} />
|
||||
</Form.Item>
|
||||
|
||||
<Form.Item
|
||||
label={
|
||||
<span>
|
||||
Provider{' '}
|
||||
<Tooltip title="Select the provider for this vector store">
|
||||
<InfoCircleOutlined style={{ marginLeft: '4px' }} />
|
||||
</Tooltip>
|
||||
</span>
|
||||
}
|
||||
|
||||
// Get the display name from Providers enum and logo from map
|
||||
const displayName = Providers[enumKey as keyof typeof Providers];
|
||||
const logo = providerLogoMap[displayName];
|
||||
|
||||
return { displayName, logo };
|
||||
})();
|
||||
|
||||
return (
|
||||
<>
|
||||
{logo && (
|
||||
<img
|
||||
src={logo}
|
||||
alt={`${displayName} logo`}
|
||||
className="w-5 h-5"
|
||||
onError={(e) => {
|
||||
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 = displayName.charAt(0);
|
||||
parent.replaceChild(fallbackDiv, target);
|
||||
}
|
||||
}}
|
||||
/>
|
||||
)}
|
||||
<Badge color="blue">{displayName}</Badge>
|
||||
</>
|
||||
);
|
||||
})()}
|
||||
name="custom_llm_provider"
|
||||
rules={[{ required: true, message: "Please select a provider" }]}
|
||||
>
|
||||
<Select2>
|
||||
{Object.entries(Providers).map(([providerEnum, providerDisplayName]) => {
|
||||
// Currently only showing Bedrock since it's the only supported provider
|
||||
if (providerEnum === 'Bedrock') {
|
||||
return (
|
||||
<Select2.Option key={providerEnum} value={provider_map[providerEnum]}>
|
||||
<div className="flex items-center space-x-2">
|
||||
<img
|
||||
src={providerLogoMap[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>
|
||||
</Select2.Option>
|
||||
);
|
||||
}
|
||||
return null;
|
||||
})}
|
||||
</Select2>
|
||||
</Form.Item>
|
||||
|
||||
{/* Credentials */}
|
||||
<div className="mb-4">
|
||||
<Text className="text-sm text-gray-500 mb-2">
|
||||
Either select existing credentials OR enter provider credentials below
|
||||
</Text>
|
||||
</div>
|
||||
|
||||
<Form.Item
|
||||
label="Existing Credentials"
|
||||
name="litellm_credential_name"
|
||||
>
|
||||
<Select2
|
||||
showSearch
|
||||
placeholder="Select or search for existing credentials"
|
||||
optionFilterProp="children"
|
||||
filterOption={(input, option) =>
|
||||
(option?.label ?? '').toLowerCase().includes(input.toLowerCase())
|
||||
}
|
||||
options={[
|
||||
{ value: null, label: 'None' },
|
||||
...credentials.map((credential) => ({
|
||||
value: credential.credential_name,
|
||||
label: credential.credential_name
|
||||
}))
|
||||
]}
|
||||
allowClear
|
||||
/>
|
||||
</Form.Item>
|
||||
|
||||
<div className="flex items-center my-4">
|
||||
<div className="flex-grow border-t border-gray-200"></div>
|
||||
<span className="px-4 text-gray-500 text-sm">OR</span>
|
||||
<div className="flex-grow border-t border-gray-200"></div>
|
||||
</div>
|
||||
|
||||
<Form.Item
|
||||
label={
|
||||
<span>
|
||||
Metadata{' '}
|
||||
<Tooltip title="JSON metadata for the vector store">
|
||||
<InfoCircleOutlined style={{ marginLeft: '4px' }} />
|
||||
</Tooltip>
|
||||
</span>
|
||||
}
|
||||
>
|
||||
<Input.TextArea
|
||||
rows={4}
|
||||
value={metadataString}
|
||||
onChange={(e) => setMetadataString(e.target.value)}
|
||||
placeholder='{"key": "value"}'
|
||||
/>
|
||||
</Form.Item>
|
||||
|
||||
<div className="flex justify-end space-x-2">
|
||||
<AntButton onClick={() => setIsEditing(false)}>Cancel</AntButton>
|
||||
<AntButton type="primary" htmlType="submit">Save Changes</AntButton>
|
||||
</div>
|
||||
</Form>
|
||||
</Card>
|
||||
</div>
|
||||
) : (
|
||||
<div>
|
||||
<div className="flex justify-between items-center mb-4">
|
||||
<Title>Vector Store Details</Title>
|
||||
{is_admin && (
|
||||
<Button onClick={() => setIsEditing(true)}>Edit Vector Store</Button>
|
||||
)}
|
||||
</div>
|
||||
<Card>
|
||||
<div className="space-y-4">
|
||||
<div>
|
||||
<Text className="font-medium">ID</Text>
|
||||
<Text>{vectorStoreDetails.vector_store_id}</Text>
|
||||
</div>
|
||||
<div>
|
||||
<Text className="font-medium">Name</Text>
|
||||
<Text>{vectorStoreDetails.vector_store_name || "-"}</Text>
|
||||
</div>
|
||||
<div>
|
||||
<Text className="font-medium">Description</Text>
|
||||
<Text>{vectorStoreDetails.vector_store_description || "-"}</Text>
|
||||
</div>
|
||||
<div>
|
||||
<Text className="font-medium">Provider</Text>
|
||||
<div className="flex items-center space-x-2 mt-1">
|
||||
{(() => {
|
||||
const provider = vectorStoreDetails.custom_llm_provider || "bedrock";
|
||||
const { displayName, logo } = (() => {
|
||||
// Find the enum key by matching provider_map values
|
||||
const enumKey = Object.keys(provider_map).find(
|
||||
key => provider_map[key].toLowerCase() === provider.toLowerCase()
|
||||
);
|
||||
|
||||
if (!enumKey) {
|
||||
return { displayName: provider, logo: "" };
|
||||
}
|
||||
|
||||
// Get the display name from Providers enum and logo from map
|
||||
const displayName = Providers[enumKey as keyof typeof Providers];
|
||||
const logo = providerLogoMap[displayName];
|
||||
|
||||
return { displayName, logo };
|
||||
})();
|
||||
|
||||
return (
|
||||
<>
|
||||
{logo && (
|
||||
<img
|
||||
src={logo}
|
||||
alt={`${displayName} logo`}
|
||||
className="w-5 h-5"
|
||||
onError={(e) => {
|
||||
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 = displayName.charAt(0);
|
||||
parent.replaceChild(fallbackDiv, target);
|
||||
}
|
||||
}}
|
||||
/>
|
||||
)}
|
||||
<Badge color="blue">{displayName}</Badge>
|
||||
</>
|
||||
);
|
||||
})()}
|
||||
</div>
|
||||
</div>
|
||||
<div>
|
||||
<Text className="font-medium">Metadata</Text>
|
||||
<div className="bg-gray-50 p-3 rounded mt-2 font-mono text-xs overflow-auto max-h-48">
|
||||
<pre>{metadataString}</pre>
|
||||
</div>
|
||||
</div>
|
||||
<div>
|
||||
<Text className="font-medium">Created</Text>
|
||||
<Text>{vectorStoreDetails.created_at ? new Date(vectorStoreDetails.created_at).toLocaleString() : "-"}</Text>
|
||||
</div>
|
||||
<div>
|
||||
<Text className="font-medium">Last Updated</Text>
|
||||
<Text>{vectorStoreDetails.updated_at ? new Date(vectorStoreDetails.updated_at).toLocaleString() : "-"}</Text>
|
||||
</div>
|
||||
</div>
|
||||
</Card>
|
||||
</div>
|
||||
<div>
|
||||
<Text className="font-medium">Metadata</Text>
|
||||
<div className="bg-gray-50 p-3 rounded mt-2 font-mono text-xs overflow-auto max-h-48">
|
||||
<pre>{metadataString}</pre>
|
||||
</div>
|
||||
</div>
|
||||
<div>
|
||||
<Text className="font-medium">Created</Text>
|
||||
<Text>{vectorStoreDetails.created_at ? new Date(vectorStoreDetails.created_at).toLocaleString() : "-"}</Text>
|
||||
</div>
|
||||
<div>
|
||||
<Text className="font-medium">Last Updated</Text>
|
||||
<Text>{vectorStoreDetails.updated_at ? new Date(vectorStoreDetails.updated_at).toLocaleString() : "-"}</Text>
|
||||
</div>
|
||||
</div>
|
||||
</Card>
|
||||
</div>
|
||||
)}
|
||||
)}
|
||||
</TabPanel>
|
||||
|
||||
{/* Test Tab */}
|
||||
<TabPanel>
|
||||
<VectorStoreTester
|
||||
vectorStoreId={vectorStoreDetails.vector_store_id}
|
||||
accessToken={accessToken || ""}
|
||||
/>
|
||||
</TabPanel>
|
||||
</TabPanels>
|
||||
</TabGroup>
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
||||
Reference in New Issue
Block a user