From 81eb2fdd30a1848db8f57197fe83acd8e99c9d0d Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 18 Jul 2025 18:25:26 -0700 Subject: [PATCH] [Feat] UI Vector Stores - Allow adding Vertex RAG Engine, OpenAI, Azure (#12752) * 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 * add get_vertex_ai_project * fixes to vertex transformation for RAG Engine * fix vectorStoreProviderFields * Add Vertex Rag engine * add oai, azure * fix validate_environment * fix provider name * fix tester * working vertex vector store --- .../base_llm/vector_store/transformation.py | 1 + .../bedrock/vector_stores/transformation.py | 1 + litellm/llms/custom_httpx/llm_http_handler.py | 2 + .../openai/vector_stores/transformation.py | 2 + .../pg_vector/vector_stores/transformation.py | 2 + .../vertex_ai/vector_stores/transformation.py | 35 ++++++---------- ...est_bedrock_vector_store_transformation.py | 1 + .../test_vertex_ai_vector_store.py | 2 +- .../VectorStoreForm.tsx | 29 ++++++++++++- .../src/components/vector_store_providers.tsx | 42 ++++++++++++++++++- 10 files changed, 91 insertions(+), 26 deletions(-) diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index e2edc75013f..b50fd957587 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -31,6 +31,7 @@ class BaseVectorStoreConfig: vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, api_base: str, litellm_logging_obj: LiteLLMLoggingObj, + litellm_params: dict, ) -> Tuple[str, Dict]: pass diff --git a/litellm/llms/bedrock/vector_stores/transformation.py b/litellm/llms/bedrock/vector_stores/transformation.py index e8a9131e769..c05b6ba3fb1 100644 --- a/litellm/llms/bedrock/vector_stores/transformation.py +++ b/litellm/llms/bedrock/vector_stores/transformation.py @@ -60,6 +60,7 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, api_base: str, litellm_logging_obj: LiteLLMLoggingObj, + litellm_params: dict, ) -> Tuple[str, Dict]: if isinstance(query, list): query = " ".join(query) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index c623c0dcde5..f77f87507ea 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2700,6 +2700,7 @@ class BaseLLMHTTPHandler: vector_store_search_optional_params=vector_store_search_optional_params, api_base=api_base, litellm_logging_obj=logging_obj, + litellm_params=dict(litellm_params), ) ) all_optional_params: Dict[str, Any] = dict(litellm_params) @@ -2796,6 +2797,7 @@ class BaseLLMHTTPHandler: vector_store_search_optional_params=vector_store_search_optional_params, api_base=api_base, litellm_logging_obj=logging_obj, + litellm_params=dict(litellm_params), ) ) diff --git a/litellm/llms/openai/vector_stores/transformation.py b/litellm/llms/openai/vector_stores/transformation.py index 8f89f53c639..0e890f0fd51 100644 --- a/litellm/llms/openai/vector_stores/transformation.py +++ b/litellm/llms/openai/vector_stores/transformation.py @@ -39,6 +39,7 @@ class OpenAIVectorStoreConfig(BaseVectorStoreConfig): headers.update( { "Authorization": f"Bearer {api_key}", + "Content-Type": "application/json", } ) @@ -83,6 +84,7 @@ class OpenAIVectorStoreConfig(BaseVectorStoreConfig): vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, api_base: str, litellm_logging_obj: LiteLLMLoggingObj, + litellm_params: dict, ) -> Tuple[str, Dict]: url = f"{api_base}/{vector_store_id}/search" typed_request_body = VectorStoreSearchRequest( diff --git a/litellm/llms/pg_vector/vector_stores/transformation.py b/litellm/llms/pg_vector/vector_stores/transformation.py index 6553a85ad9b..5d10faeba50 100644 --- a/litellm/llms/pg_vector/vector_stores/transformation.py +++ b/litellm/llms/pg_vector/vector_stores/transformation.py @@ -81,6 +81,7 @@ class PGVectorStoreConfig(OpenAIVectorStoreConfig): vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, api_base: str, litellm_logging_obj: LiteLLMLoggingObj, + litellm_params: dict, ) -> Tuple[str, Dict]: url = f"{api_base}/{vector_store_id}/search" _, request_body = super().transform_search_vector_store_request( @@ -89,5 +90,6 @@ class PGVectorStoreConfig(OpenAIVectorStoreConfig): vector_store_search_optional_params=vector_store_search_optional_params, api_base=api_base, litellm_logging_obj=litellm_logging_obj, + litellm_params=litellm_params, ) return url, request_body \ No newline at end of file diff --git a/litellm/llms/vertex_ai/vector_stores/transformation.py b/litellm/llms/vertex_ai/vector_stores/transformation.py index ca374728c03..80e76fd71a8 100644 --- a/litellm/llms/vertex_ai/vector_stores/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/transformation.py @@ -42,16 +42,8 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): litellm_params = litellm_params or GenericLiteLLMParams() # Get credentials and project info - vertex_credentials = ( - litellm_params.vertex_credentials - or get_secret_str("VERTEXAI_CREDENTIALS") - or get_secret_str("VERTEX_AI_CREDENTIALS") - ) - vertex_project = ( - litellm_params.vertex_project - or get_secret_str("VERTEXAI_PROJECT") - or get_secret_str("VERTEX_AI_PROJECT") - ) + vertex_credentials = self.get_vertex_ai_credentials(dict(litellm_params)) + vertex_project = self.get_vertex_ai_project(dict(litellm_params)) # Get access token using the base class method access_token, project_id = self._ensure_access_token( @@ -75,17 +67,8 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): """ Get the Base endpoint for Vertex AI RAG API """ - vertex_location = ( - litellm_params.get("vertex_location") - or get_secret_str("VERTEXAI_LOCATION") - or get_secret_str("VERTEX_AI_LOCATION") - or "us-central1" - ) - vertex_project = ( - litellm_params.get("vertex_project") - or get_secret_str("VERTEXAI_PROJECT") - or get_secret_str("VERTEX_AI_PROJECT") - ) + vertex_location = self.get_vertex_ai_location(litellm_params) + vertex_project = self.get_vertex_ai_project(litellm_params) if api_base: return api_base.rstrip("/") @@ -100,6 +83,7 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, api_base: str, litellm_logging_obj: LiteLLMLoggingObj, + litellm_params: dict, ) -> Tuple[str, Dict[str, Any]]: """ Transform search request for Vertex AI RAG API @@ -111,12 +95,19 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): # Vertex AI RAG API endpoint for retrieving contexts url = f"{api_base}:retrieveContexts" + # Use helper methods to get project and location, then construct full rag corpus path + vertex_project = self.get_vertex_ai_project(litellm_params) + vertex_location = self.get_vertex_ai_location(litellm_params) + + # Construct full rag corpus path + full_rag_corpus = f"projects/{vertex_project}/locations/{vertex_location}/ragCorpora/{vector_store_id}" + # Build the request body for Vertex AI RAG API request_body: Dict[str, Any] = { "vertex_rag_store": { "rag_resources": [ { - "rag_corpus": vector_store_id + "rag_corpus": full_rag_corpus } ] }, diff --git a/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py b/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py index e924fe83bb8..28b60e5e75f 100644 --- a/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py +++ b/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py @@ -20,6 +20,7 @@ def test_transform_search_request(): vector_store_search_optional_params={}, api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases", litellm_logging_obj=mock_log, + litellm_params={}, ) assert url.endswith("/kb123/retrieve") diff --git a/tests/vector_store_tests/test_vertex_ai_vector_store.py b/tests/vector_store_tests/test_vertex_ai_vector_store.py index e9dee407917..94371601075 100644 --- a/tests/vector_store_tests/test_vertex_ai_vector_store.py +++ b/tests/vector_store_tests/test_vertex_ai_vector_store.py @@ -19,7 +19,7 @@ class TestVertexAIVectorStore(BaseVectorStoreTest): def get_base_request_args(self): return { - "vector_store_id": "projects/reliablekeys/locations/us-central1/ragCorpora/6917529027641081856", + "vector_store_id": "6917529027641081856", "custom_llm_provider": "vertex_ai", "vertex_project": "reliablekeys", "vertex_location": "us-central1", diff --git a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreForm.tsx b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreForm.tsx index c36cc4ed6cd..99acccf3d3b 100644 --- a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreForm.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreForm.tsx @@ -164,6 +164,27 @@ const VectorStoreForm: React.FC = ({ /> )} + {/* Vertex RAG Engine Setup Instructions */} + {selectedProvider === "vertex_rag_engine" && ( + +

To use Vertex AI RAG Engine:

+
    +
  1. Set up your Vertex AI RAG Engine corpus following the guide: Vertex AI RAG Engine Overview
    2. +
  2. Create a corpus in your Google Cloud project
    3. +
  3. Note the corpus ID from the Vertex AI console
    4. +
  4. Enter the corpus ID in the Vector Store ID field below
    5. +
+ + } + type="info" + showIcon + style={{ marginBottom: '16px' }} + /> + )} + @@ -176,7 +197,13 @@ const VectorStoreForm: React.FC = ({ name="vector_store_id" rules={[{ required: true, message: "Please input the vector store ID from your api provider" }]} > - + {/* Provider-specific fields */} diff --git a/ui/litellm-dashboard/src/components/vector_store_providers.tsx b/ui/litellm-dashboard/src/components/vector_store_providers.tsx index 9cc97b99478..a188d6fba6f 100644 --- a/ui/litellm-dashboard/src/components/vector_store_providers.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_providers.tsx @@ -1,11 +1,17 @@ export enum VectorStoreProviders { Bedrock = "Amazon Bedrock", - PgVector = "PostgreSQL pgvector (LiteLLM Connector)" + PgVector = "PostgreSQL pgvector (LiteLLM Connector)", + VertexRagEngine = "Vertex AI RAG Engine", + OpenAI = "OpenAI", + Azure = "Azure OpenAI" } export const vectorStoreProviderMap: Record = { Bedrock: "bedrock", - PgVector: "pg_vector" + PgVector: "pg_vector", + VertexRagEngine: "vertex_ai", + OpenAI: "openai", + Azure: "azure" }; const asset_logos_folder = '/assets/logos/'; @@ -13,6 +19,9 @@ const asset_logos_folder = '/assets/logos/'; export const vectorStoreProviderLogoMap: Record = { [VectorStoreProviders.Bedrock]: `${asset_logos_folder}bedrock.svg`, [VectorStoreProviders.PgVector]: `${asset_logos_folder}postgresql.svg`, // Fallback to a generic database icon if needed + [VectorStoreProviders.VertexRagEngine]: `${asset_logos_folder}google.svg`, + [VectorStoreProviders.OpenAI]: `${asset_logos_folder}openai_small.svg`, + [VectorStoreProviders.Azure]: `${asset_logos_folder}microsoft_azure.svg`, }; // Define field types for provider-specific configurations @@ -45,6 +54,35 @@ export const vectorStoreProviderFields: Record required: true, type: "password" } + ], + vertex_rag_engine: [], + openai: [ + { + name: "api_key", + label: "API Key", + tooltip: "Enter your OpenAI API key", + placeholder: "sk-...", + required: true, + type: "password" + } + ], + azure: [ + { + name: "api_key", + label: "API Key", + tooltip: "Enter your Azure OpenAI API key", + placeholder: "your-azure-api-key", + required: true, + type: "password" + }, + { + name: "api_base", + label: "API Base", + tooltip: "Enter your Azure OpenAI endpoint (e.g., https://your-resource.openai.azure.com/)", + placeholder: "https://your-resource.openai.azure.com/", + required: true, + type: "text" + } ] };