From ecd63e315ac05f2edb101b5144652fa224f73d04 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 1 Sep 2026 14:09:45 -0700 Subject: [PATCH] fix(vector-stores): trust config-defined Milvus gRPC --- litellm/proxy/vector_store_endpoints/utils.py | 6 ++-- litellm/types/vector_stores.py | 4 ++- .../vector_stores/vector_store_registry.py | 4 +++ .../test_vector_store_endpoints.py | 30 +++++++++++++++++++ 4 files changed, 41 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/vector_store_endpoints/utils.py b/litellm/proxy/vector_store_endpoints/utils.py index 997c6bae49c..9f527aec467 100644 --- a/litellm/proxy/vector_store_endpoints/utils.py +++ b/litellm/proxy/vector_store_endpoints/utils.py @@ -18,10 +18,12 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.types.utils import LlmProviders -from litellm.types.vector_stores import LiteLLM_ManagedVectorStore +from litellm.types.vector_stores import ( + MILVUS_ADMIN_CONFIGURED_CONNECTION, + LiteLLM_ManagedVectorStore, +) from litellm.utils import ProviderConfigManager -MILVUS_ADMIN_CONFIGURED_CONNECTION: Final = "_litellm_admin_configured_milvus_grpc" MILVUS_GRPC_CONNECTION_FIELDS: Final = frozenset( { "api_base", diff --git a/litellm/types/vector_stores.py b/litellm/types/vector_stores.py index 474c652ff3a..8c2101d91d7 100644 --- a/litellm/types/vector_stores.py +++ b/litellm/types/vector_stores.py @@ -1,11 +1,13 @@ from dataclasses import dataclass from datetime import datetime from enum import Enum -from typing import Any, Literal +from typing import Any, Final, Literal from pydantic import BaseModel from typing_extensions import TypedDict +MILVUS_ADMIN_CONFIGURED_CONNECTION: Final = "_litellm_admin_configured_milvus_grpc" + class SupportedVectorStoreIntegrations(str, Enum): """Supported vector store integrations.""" diff --git a/litellm/vector_stores/vector_store_registry.py b/litellm/vector_stores/vector_store_registry.py index b71d6784873..5580e867001 100644 --- a/litellm/vector_stores/vector_store_registry.py +++ b/litellm/vector_stores/vector_store_registry.py @@ -17,6 +17,7 @@ from litellm.repositories.table_repositories import ( ManagedVectorStoresRepository, ) from litellm.types.vector_stores import ( + MILVUS_ADMIN_CONFIGURED_CONNECTION, VECTOR_STORE_OPENAI_PARAMS, LiteLLM_ManagedVectorStore, LiteLLM_ManagedVectorStoreIndex, @@ -418,6 +419,9 @@ class VectorStoreRegistry: f"custom_llm_provider is required for initializing vector store, got custom_llm_provider={custom_llm_provider}" ) + if custom_llm_provider == "milvus" and vector_store_litellm_params.get("milvus_transport") == "grpc": + vector_store_litellm_params[MILVUS_ADMIN_CONFIGURED_CONNECTION] = True + litellm_managed_vector_store = LiteLLM_ManagedVectorStore( vector_store_id=vector_store_id, custom_llm_provider=custom_llm_provider, 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 index a5a74c174b3..04ec48447b7 100644 --- 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 @@ -37,6 +37,7 @@ from litellm.proxy.vector_store_files_endpoints.endpoints import ( from litellm.types.utils import EmbeddingResponse, LlmProviders from litellm.types.vector_stores import IndexCreateRequest, IndexListResponse from litellm.vector_stores.main import _direct_vector_store_embedding_executor +from litellm.vector_stores.vector_store_registry import VectorStoreRegistry def _serialize_litellm_params(litellm_params): @@ -769,6 +770,35 @@ async def test_unmarked_managed_milvus_connection_requires_admin_resave(): assert "re-saved" in str(exc_info.value.detail) +@pytest.mark.asyncio +async def test_config_loaded_milvus_grpc_connection_is_trusted(): + registry = VectorStoreRegistry() + registry.load_vector_stores_from_config( + [ + { + "vector_store_name": "configured", + "litellm_params": { + "vector_store_id": "configured", + "custom_llm_provider": "milvus", + "milvus_transport": "grpc", + "api_base": "https://configured-milvus:19530", + "litellm_embedding_model": "team-embedding-alias", + }, + } + ] + ) + + with patch.object(litellm, "vector_store_registry", registry): + result = await _update_request_data_with_litellm_managed_vector_store_registry( + data={"query": "allowed"}, + vector_store_id="configured", + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER), + ) + + assert result["api_base"] == "https://configured-milvus:19530" + assert result[MILVUS_ADMIN_CONFIGURED_CONNECTION] is True + + def test_admin_persistence_strips_forged_marker_and_adds_server_marker(): params = prepare_milvus_connection_for_persistence( custom_llm_provider="milvus/probe",