From 7ffa8eb1d08b83eca5f56bac462ac02b9322de1e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 13:07:46 -0700 Subject: [PATCH] fix(vector-stores): approve config gRPC stores with prefixed providers approve_configured_connection only stamped the admin marker when the provider was exactly "milvus", while gRPC detection and the managed rejection both normalize prefixed values such as "milvus/team", so a config-defined gRPC store with a prefixed provider was rejected with 403 on every search. Normalize the provider the same way before stamping Also give VectorStoreRegistry and VectorStoreIndexRegistry their own list per instance instead of a shared mutable default, which made every registry built without an explicit list see the stores loaded into any other one. The budget files are ratcheted from the merge-base values, so the earlier merge commit's double-lowered limits are corrected here --- basedpyright-code-budget.json | 18 +++++++++--------- .../llms/milvus/vector_stores/connection.py | 2 +- litellm/vector_stores/vector_store_registry.py | 10 ++++++---- ruff-strict-budget.json | 4 ++-- test-quality-budget.json | 2 +- .../test_vector_store_endpoints.py | 5 +++-- .../test_vector_store_registry.py | 13 +++++++++++++ 7 files changed, 35 insertions(+), 19 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index ba0a796fbe8..733526e72f7 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,9 +1,9 @@ { "reportAny": { - "limit": 13422 + "limit": 13427 }, "reportArgumentType": { - "limit": 2137 + "limit": 2187 }, "reportAssignmentType": { "limit": 319 @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 3363 + "limit": 3368 }, "reportFunctionMemberAccess": { "limit": 7 @@ -57,7 +57,7 @@ "limit": 5570 }, "reportMissingTypeArgument": { - "limit": 15246 + "limit": 15275 }, "reportMissingTypeStubs": { "limit": 40 @@ -99,19 +99,19 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 42551 + "limit": 44133 }, "reportUnknownLambdaType": { "limit": 109 }, "reportUnknownMemberType": { - "limit": 38191 + "limit": 38267 }, "reportUnknownParameterType": { - "limit": 19576 + "limit": 19582 }, "reportUnknownVariableType": { - "limit": 29814 + "limit": 29826 }, "reportUnnecessaryCast": { "limit": 110 @@ -123,7 +123,7 @@ "limit": 4 }, "reportUnnecessaryIsInstance": { - "limit": 810 + "limit": 815 }, "reportUntypedBaseClass": { "limit": 0 diff --git a/litellm/llms/milvus/vector_stores/connection.py b/litellm/llms/milvus/vector_stores/connection.py index 920576c78ac..73d0dbab7fb 100644 --- a/litellm/llms/milvus/vector_stores/connection.py +++ b/litellm/llms/milvus/vector_stores/connection.py @@ -125,6 +125,6 @@ def managed_connection_fields(custom_llm_provider: object) -> frozenset[str]: def approve_configured_connection( custom_llm_provider: object, litellm_params: Mapping[str, object] ) -> Mapping[str, object]: - if custom_llm_provider == "milvus" and litellm_params.get("milvus_transport") == "grpc": + if _normalize_provider(custom_llm_provider) == "milvus" and litellm_params.get("milvus_transport") == "grpc": return MappingProxyType({**litellm_params, MILVUS_ADMIN_CONFIGURED_CONNECTION: True}) return litellm_params diff --git a/litellm/vector_stores/vector_store_registry.py b/litellm/vector_stores/vector_store_registry.py index 542d06ed772..1ea53cd95fe 100644 --- a/litellm/vector_stores/vector_store_registry.py +++ b/litellm/vector_stores/vector_store_registry.py @@ -58,8 +58,10 @@ def _normalized_vector_store(vector_store: LiteLLM_ManagedVectorStore) -> LiteLL class VectorStoreIndexRegistry: - def __init__(self, vector_store_indexes: list[LiteLLM_ManagedVectorStoreIndex] = []): - self.vector_store_indexes: list[LiteLLM_ManagedVectorStoreIndex] = vector_store_indexes + def __init__(self, vector_store_indexes: list[LiteLLM_ManagedVectorStoreIndex] | None = None): + self.vector_store_indexes: list[LiteLLM_ManagedVectorStoreIndex] = ( + vector_store_indexes if vector_store_indexes is not None else [] + ) def get_vector_store_indexes(self) -> list[LiteLLM_ManagedVectorStoreIndex]: """ @@ -127,8 +129,8 @@ class VectorStoreIndexRegistry: class VectorStoreRegistry: - def __init__(self, vector_stores: list[LiteLLM_ManagedVectorStore] = []): - self.vector_stores: list[LiteLLM_ManagedVectorStore] = vector_stores + def __init__(self, vector_stores: list[LiteLLM_ManagedVectorStore] | None = None): + self.vector_stores: list[LiteLLM_ManagedVectorStore] = vector_stores if vector_stores is not None else [] self.vector_store_ids_to_vector_store_map: dict[str, LiteLLM_ManagedVectorStore] = {} self.config_vector_store_ids: frozenset[str] = frozenset() diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 18b6622eec3..f5323a0fd05 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -33,7 +33,7 @@ "limit": 2 }, "B006": { - "limit": 176 + "limit": 174 }, "B008": { "limit": 503 @@ -246,7 +246,7 @@ "limit": 109 }, "TRY300": { - "limit": 848 + "limit": 851 }, "UP028": { "limit": 2 diff --git a/test-quality-budget.json b/test-quality-budget.json index 3a600a19537..1c6282a9251 100644 --- a/test-quality-budget.json +++ b/test-quality-budget.json @@ -21,6 +21,6 @@ "limit": 117 }, "TQ008": { - "limit": 10972 + "limit": 10990 } } 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 39bd83beb98..4b5057ca47d 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 @@ -794,13 +794,14 @@ async def test_unmarked_managed_milvus_connection_requires_admin_resave(): @pytest.mark.asyncio -async def test_config_loaded_milvus_grpc_connection_is_trusted(): +@pytest.mark.parametrize("custom_llm_provider", ["milvus", "milvus/probe"]) +async def test_config_loaded_milvus_grpc_connection_is_trusted(custom_llm_provider: str): registry = VectorStoreRegistry() source = { "vector_store_name": "configured", "litellm_params": { "vector_store_id": "configured", - "custom_llm_provider": "milvus", + "custom_llm_provider": custom_llm_provider, "milvus_transport": "grpc", "api_base": "https://configured-milvus:19530", "litellm_embedding_model": "team-embedding-alias", diff --git a/tests/test_litellm/vector_stores/test_vector_store_registry.py b/tests/test_litellm/vector_stores/test_vector_store_registry.py index ece38287a35..1f78a707104 100644 --- a/tests/test_litellm/vector_stores/test_vector_store_registry.py +++ b/tests/test_litellm/vector_stores/test_vector_store_registry.py @@ -105,6 +105,19 @@ def test_get_credentials_for_vector_store(): assert result == {} +def test_fresh_registries_do_not_share_config_loaded_stores(): + VectorStoreRegistry().load_vector_stores_from_config( + [ + { + "vector_store_name": "configured", + "litellm_params": {"vector_store_id": "configured", "custom_llm_provider": "openai"}, + } + ] + ) + + assert VectorStoreRegistry().get_litellm_managed_vector_store_from_registry("configured") is None + + def test_add_vector_store_to_registry(): """Test that add_vector_store_to_registry adds vector store correctly when there are pre-existing stores""" # Create pre-existing vector stores