From 5802a5bbe3b10fc4b0a8bffbd482f17c75ca8b0b Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 18 Jul 2025 18:18:53 -0700 Subject: [PATCH] [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 --- litellm/proxy/_types.py | 6 + .../proxy/vector_store_endpoints/endpoints.py | 36 ++ litellm/router.py | 31 +- .../vector_stores/vector_store_registry.py | 11 + .../test_router_endpoints.py | 48 ++ .../test_vector_store_endpoints.py | 102 ++++ .../test_bedrock_vector_store.py | 15 +- .../src/components/networking.tsx | 32 ++ .../VectorStoreTester.tsx | 280 +++++++++++ .../vector_store_info.tsx | 462 ++++++++++-------- 10 files changed, 805 insertions(+), 218 deletions(-) create mode 100644 tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py create mode 100644 ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTester.tsx diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 44f9da39f1..540b3e6b74 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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 = [ diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index 22ea7fd546..7859c1c4b8 100644 --- a/litellm/proxy/vector_store_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_endpoints/endpoints.py @@ -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( diff --git a/litellm/router.py b/litellm/router.py index 0b7c36325d..644422cda3 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -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, ): diff --git a/litellm/vector_stores/vector_store_registry.py b/litellm/vector_stores/vector_store_registry.py index c5bb809d26..680200ca17 100644 --- a/litellm/vector_stores/vector_store_registry.py +++ b/litellm/vector_stores/vector_store_registry.py @@ -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]: diff --git a/tests/router_unit_tests/test_router_endpoints.py b/tests/router_unit_tests/test_router_endpoints.py index cc813c5f78..de1d22d5a8 100644 --- a/tests/router_unit_tests/test_router_endpoints.py +++ b/tests/router_unit_tests/test_router_endpoints.py @@ -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" + ) + diff --git a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py new file mode 100644 index 0000000000..384c37ec1e --- /dev/null +++ b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py @@ -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 \ No newline at end of file diff --git a/tests/vector_store_tests/test_bedrock_vector_store.py b/tests/vector_store_tests/test_bedrock_vector_store.py index 98ea070667..5f473fac00 100644 --- a/tests/vector_store_tests/test_bedrock_vector_store.py +++ b/tests/vector_store_tests/test_bedrock_vector_store.py @@ -98,4 +98,17 @@ class TestBedrockVectorStore(BaseVectorStoreTest): # Test with None attributes = config._get_attributes_from_metadata(None) - assert attributes == {} \ No newline at end of file + 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) \ No newline at end of file diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 98010ee9a8..9054d70bd4 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -6226,3 +6226,35 @@ export const testMCPToolsListRequest = async ( throw error; } }; + +export const vectorStoreSearchCall = async ( + accessToken: string, + vectorStoreId: string, + query: string +): Promise => { + 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; + } +}; diff --git a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTester.tsx b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTester.tsx new file mode 100644 index 0000000000..a70c0d6400 --- /dev/null +++ b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTester.tsx @@ -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; +} + +interface VectorStoreSearchResponse { + object: string; + search_query: string; + data: VectorStoreResult[]; +} + +interface VectorStoreTesterProps { + vectorStoreId: string; + accessToken: string; + className?: string; +} + +export const VectorStoreTester: React.FC = ({ + 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>({}); + + 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) => { + 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 ( + +
+ {/* Header */} +
+
+ + Test Vector Store +
+ {searchHistory.length > 0 && ( + + )} +
+ + {/* Results Area */} +
+ {searchHistory.length === 0 ? ( +
+ + Test your vector store by entering a search query below +
+ ) : ( +
+ {searchHistory.map((entry, index) => ( +
+ {/* User Query */} +
+
+
+ Query + + {formatTimestamp(entry.timestamp)} + +
+
+ {entry.query} +
+
+
+ + {/* Vector Store Response */} +
+
+
+ + Vector Store Results + {entry.response && ( + + {entry.response.data?.length || 0} results + + )} +
+ + {entry.response && entry.response.data && entry.response.data.length > 0 ? ( +
+ {entry.response.data.map((result, resultIndex) => { + const isExpanded = expandedResults[`${index}-${resultIndex}`] || false; + + return ( +
+ {/* Clickable Header */} +
toggleResultExpansion(index, resultIndex)} + > +
+ {isExpanded ? ( + + ) : ( + + )} + Result {resultIndex + 1} + {/* Show preview of content when collapsed */} + {!isExpanded && result.content && result.content[0] && ( + + - {result.content[0].text.substring(0, 100)}... + + )} +
+ + Score: {result.score.toFixed(4)} + +
+ + {/* Expandable Content */} + {isExpanded && ( +
+ {/* Content */} + {result.content && result.content.map((content, contentIndex) => ( +
+
Content ({content.type})
+
+ {content.text} +
+
+ ))} + + {/* Metadata */} + {(result.file_id || result.filename || result.attributes) && ( +
+
Metadata
+
+ {result.file_id && ( +
+ File ID: {result.file_id} +
+ )} + {result.filename && ( +
+ Filename: {result.filename} +
+ )} + {result.attributes && Object.keys(result.attributes).length > 0 && ( +
+ Attributes: + 
+                                                  {JSON.stringify(result.attributes, null, 2)}
+
+
+ )} +
+
+ )} +
+ )} +
+ ); + })} +
+ ) : ( +
No results found
+ )} +
+
+ + {index < searchHistory.length - 1 && } +
+ ))} +
+ )} + + {isLoading && ( +
+ } /> +
+ )} +
+ + {/* Input Area */} +
+
+
+