From db4e650f4ae93e8813d43d9e70adc630d730217b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 16:38:37 -0700 Subject: [PATCH] fix(vector-stores): gate RAG queries on the managed Milvus gRPC approval marker The approval check for managed Milvus gRPC connections ran in the search endpoint and the chat retrieval hook, but /v1/rag/query builds its search params from the same managed store without it. Move the check into build_request_data_from_managed_vector_store, the shared boundary both endpoints use, so an unmarked gRPC row is rejected everywhere and the search endpoint no longer needs its own post-merge copy of the check Also move MILVUS_ADMIN_CONFIGURED_CONNECTION into litellm/constants.py alongside the other server-owned sentinels --- litellm/constants.py | 1 + .../llms/milvus/vector_stores/connection.py | 2 +- .../proxy/vector_store_endpoints/endpoints.py | 18 +++---- .../management_endpoints.py | 3 +- litellm/types/vector_stores.py | 4 +- .../proxy/rag_endpoints/test_rag_endpoints.py | 54 +++++++++++++++++++ tests/test_litellm/proxy/test_proxy_server.py | 3 +- .../test_vector_store_endpoints.py | 4 +- 8 files changed, 71 insertions(+), 18 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index ce744e9c58a..1d4ada75a93 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -409,6 +409,7 @@ AZURE_COMPUTER_USE_OUTPUT_COST_PER_1K_TOKENS: Final = float( AZURE_VECTOR_STORE_COST_PER_GB_PER_DAY: Final = float( os.getenv("AZURE_VECTOR_STORE_COST_PER_GB_PER_DAY", 0.1) # $0.1 USD per 1 GB/Day (same as file search) ) +MILVUS_ADMIN_CONFIGURED_CONNECTION: Final = "_litellm_admin_configured_milvus_grpc" MIN_NON_ZERO_TEMPERATURE: Final = float(os.getenv("MIN_NON_ZERO_TEMPERATURE", 0.0001)) #### RELIABILITY #### REPEATED_STREAMING_CHUNK_LIMIT: Final = int( diff --git a/litellm/llms/milvus/vector_stores/connection.py b/litellm/llms/milvus/vector_stores/connection.py index f84e1651d37..f734f623450 100644 --- a/litellm/llms/milvus/vector_stores/connection.py +++ b/litellm/llms/milvus/vector_stores/connection.py @@ -4,7 +4,7 @@ from types import MappingProxyType from typing import Final import litellm -from litellm.types.vector_stores import MILVUS_ADMIN_CONFIGURED_CONNECTION +from litellm.constants import MILVUS_ADMIN_CONFIGURED_CONNECTION MILVUS_MANAGED_CONFIGURATION_FIELDS: Final = frozenset( { diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index f309c028aef..6643825128c 100644 --- a/litellm/proxy/vector_store_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_endpoints/endpoints.py @@ -9,6 +9,7 @@ from typing import ( from fastapi import APIRouter, Depends, HTTPException, Request, Response +from litellm.constants import MILVUS_ADMIN_CONFIGURED_CONNECTION from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( LiteLLM_ManagedVectorStore, ) @@ -24,7 +25,7 @@ from litellm.proxy.vector_store_endpoints.utils import ( get_litellm_managed_vector_store, ) from litellm.repositories.table_repositories import ManagedVectorStoreIndexRepository -from litellm.types.vector_stores import MILVUS_ADMIN_CONFIGURED_CONNECTION, IndexCreateRequest, IndexListResponse +from litellm.types.vector_stores import IndexCreateRequest, IndexListResponse from litellm.vector_stores.vector_store_registry import VectorStoreIndexRegistry router: Final = APIRouter() @@ -67,7 +68,13 @@ def build_request_data_from_managed_vector_store( } ) litellm_params: Final = vector_store.get("litellm_params") or MappingProxyType({}) - return MappingProxyType({**top_level, **litellm_params}) + request_data: Final = MappingProxyType({**top_level, **litellm_params}) + assert_proxy_admin_for_user_supplied_vector_store_connection( + custom_llm_provider=request_data.get("custom_llm_provider"), + litellm_params=request_data, + managed=True, + ) + return request_data async def _update_request_data_with_litellm_managed_vector_store_registry( @@ -111,13 +118,6 @@ async def _update_request_data_with_litellm_managed_vector_store_registry( **{key: value for key, value in data.items() if key not in blocked_fields}, **managed_data, } - if user_api_key_dict is not None: - assert_proxy_admin_for_user_supplied_vector_store_connection( - custom_llm_provider=request_data.get("custom_llm_provider"), - litellm_params=request_data, - user_api_key_dict=user_api_key_dict, - managed=True, - ) return request_data diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index ac10d9f0735..311850db678 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -20,7 +20,7 @@ if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient import litellm from litellm._logging import verbose_proxy_logger -from litellm.constants import REDACTED_BY_LITELM_STRING +from litellm.constants import MILVUS_ADMIN_CONFIGURED_CONNECTION, REDACTED_BY_LITELM_STRING from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker from litellm.proxy._types import ( @@ -38,7 +38,6 @@ from litellm.proxy.vector_store_endpoints.utils import ( from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import ManagedVectorStoresRepository from litellm.types.vector_stores import ( - MILVUS_ADMIN_CONFIGURED_CONNECTION, LiteLLM_ManagedVectorStore, LiteLLM_ManagedVectorStoreListResponse, VectorStoreDeleteRequest, diff --git a/litellm/types/vector_stores.py b/litellm/types/vector_stores.py index d48a3b3dea2..3ad09d27580 100644 --- a/litellm/types/vector_stores.py +++ b/litellm/types/vector_stores.py @@ -2,13 +2,11 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import datetime from enum import Enum -from typing import Any, Final, Literal +from typing import Any, 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/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py index 832435711c6..d5dc67a7e55 100644 --- a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py @@ -608,6 +608,60 @@ def test_rag_query_rejects_caller_embedding_selection_params(client_internal_use assert blocked_key in str(response.json()) +@pytest.mark.parametrize("approved", [False, True]) +def test_rag_query_gates_managed_milvus_grpc_store_on_admin_approval(client_internal_user, approved): + """ + Regression: /v1/rag/query resolves the managed store itself, so it must apply + the same admin-approval check as /v1/vector_stores/{id}/search. A Milvus gRPC + row without the server-issued approval marker (written before the gate + existed, or by an older proxy sharing the DB) must be rejected instead of + opening a gRPC channel to its api_base. + """ + import litellm + from litellm.constants import MILVUS_ADMIN_CONFIGURED_CONNECTION + from litellm.types.utils import ModelResponse + + connection = {"milvus_transport": "grpc", "api_base": "http://internal-milvus:19530"} + mock_vector_store = { + "vector_store_id": "legacy-milvus", + "custom_llm_provider": "milvus", + "litellm_params": {**connection, MILVUS_ADMIN_CONFIGURED_CONNECTION: True} if approved else connection, + } + mock_registry = MagicMock() + mock_registry.get_litellm_managed_vector_store_from_registry.return_value = mock_vector_store + + mock_response = ModelResponse( + id="chatcmpl-test", + choices=[{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + model="gpt-4o-mini", + ) + + with patch( # test-quality-ok: aquery is the endpoint's downstream boundary; whether it is reached is what the test asserts + "litellm.proxy.rag_endpoints.endpoints.litellm.aquery", + new_callable=AsyncMock, + return_value=mock_response, + ) as mock_aquery, patch.object(litellm, "vector_store_registry", mock_registry), patch( # test-quality-ok: seeds the managed-store registry and grants access so the connection gate is the only thing that can reject + "litellm.proxy.vector_store_endpoints.utils.can_user_access_vector_store", + new=AsyncMock(return_value=True), + ): + response = client_internal_user.post( + "/v1/rag/query", + json={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hello"}], + "retrieval_config": {"vector_store_id": "legacy-milvus"}, + }, + ) + + if approved: + assert response.status_code == 200, response.json() + assert mock_aquery.await_args.kwargs["vector_store_params"]["api_base"] == connection["api_base"] + return + assert response.status_code == 403, response.json() + assert "re-saved by a proxy admin" in str(response.json()) + mock_aquery.assert_not_awaited() + + EICAR = r"X5O!P%@AP[4\PZX54(P^)7CC)7}$EICAR-STANDARD-ANTIVIRUS-TEST-FILE!$H+H*" INGEST_REQUEST = '{"ingest_options":{"vector_store":{"custom_llm_provider":"openai"}}}' diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 23b5c91786a..3bd736b3c92 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -12785,8 +12785,9 @@ async def test_update_general_settings_keeps_yaml_openai_websocket_passthrough() @pytest.mark.asyncio async def test_init_vector_stores_in_db_refreshes_a_store_already_in_the_registry(monkeypatch): + from litellm.constants import MILVUS_ADMIN_CONFIGURED_CONNECTION from litellm.proxy.proxy_server import ProxyConfig - from litellm.types.vector_stores import MILVUS_ADMIN_CONFIGURED_CONNECTION, LiteLLM_ManagedVectorStore + from litellm.types.vector_stores import LiteLLM_ManagedVectorStore from litellm.vector_stores.vector_store_registry import VectorStoreRegistry stale: Final = LiteLLM_ManagedVectorStore( 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 1da2af2a79e..5143bb4155b 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 @@ -8,6 +8,7 @@ import pytest from fastapi import HTTPException, Request import litellm +from litellm.constants import MILVUS_ADMIN_CONFIGURED_CONNECTION from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( LiteLLM_ManagedVectorStore, VectorStorePreCallHook, @@ -39,10 +40,9 @@ from litellm.proxy.vector_store_files_endpoints.endpoints import ( _update_request_data_with_model_routing_hint, ) from litellm.types.utils import EmbeddingResponse, LlmProviders -from litellm.types.vector_stores import MILVUS_ADMIN_CONFIGURED_CONNECTION, IndexCreateRequest, IndexListResponse +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 - from tests.test_litellm.integrations.vector_store_integrations.test_vector_store_pre_call_hook import ( FakeLoggingObj, FakeProxyRuntime,