diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 6d3eed9c371..85402ef5787 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -3,7 +3,7 @@ "limit": 13427 }, "reportArgumentType": { - "limit": 2187 + "limit": 2188 }, "reportAssignmentType": { "limit": 319 @@ -99,19 +99,19 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 44134 + "limit": 44131 }, "reportUnknownLambdaType": { "limit": 109 }, "reportUnknownMemberType": { - "limit": 38255 + "limit": 38252 }, "reportUnknownParameterType": { "limit": 19582 }, "reportUnknownVariableType": { - "limit": 29814 + "limit": 29806 }, "reportUnnecessaryCast": { "limit": 110 diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index d6a402e1860..5a85cccce23 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -44,10 +44,12 @@ from litellm.proxy.rag_endpoints.upload_security import ( validate_upload, ) from litellm.proxy.vector_store_endpoints.endpoints import ( + apply_managed_vector_store_connection_policy, build_request_data_from_managed_vector_store, reject_caller_embedding_selection_params, ) from litellm.proxy.vector_store_endpoints.utils import ( + assert_proxy_admin_for_user_supplied_vector_store_connection, assert_user_can_access_vector_store_id, ) from litellm.repositories.table_repositories import ManagedVectorStoresRepository @@ -134,7 +136,7 @@ async def _authorize_nested_vector_store_ids( user_api_key_dict: UserAPIKeyAuth, ) -> Mapping[str, LiteLLM_ManagedVectorStore]: """Authorize every nested vector store id and return the managed stores it resolved.""" - return MappingProxyType( + resolved_stores: Final = MappingProxyType( { vector_store_id: store for vector_store_id in sorted(_collect_vector_store_ids_from_payload(payload)) @@ -147,6 +149,14 @@ async def _authorize_nested_vector_store_ids( is not None } ) + for store in resolved_stores.values(): + assert_proxy_admin_for_user_supplied_vector_store_connection( + custom_llm_provider=store.get("custom_llm_provider"), + litellm_params=store.get("litellm_params"), + user_api_key_dict=user_api_key_dict, + managed=True, + ) + return resolved_stores def _build_file_metadata_entry( @@ -727,20 +737,17 @@ async def rag_query( user_api_key_dict=user_api_key_dict, ) - # Merge litellm-managed vector store params (provider, region, embedding - # model, credentials, ...) from the registry: the same source the direct - # /vector_stores/{id}/search endpoint uses. Store-managed keys win on - # conflict so callers cannot override the store's provider or credentials. managed_store: Final = resolved_stores.get(retrieval_config["vector_store_id"]) store_data: Final = ( build_request_data_from_managed_vector_store(managed_store) if managed_store is not None else MappingProxyType({}) ) - merged_retrieval_config: Final = { - **retrieval_config, - **store_data, - } # mutable-ok: litellm.aquery requires a plain dict payload + merged_retrieval_config: Final = apply_managed_vector_store_connection_policy( + data=retrieval_config, + vector_store=managed_store, + user_api_key_dict=user_api_key_dict, + ) # Add litellm data request_data: dict[str, object] = {} diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index 6643825128c..da925863d6e 100644 --- a/litellm/proxy/vector_store_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_endpoints/endpoints.py @@ -96,7 +96,22 @@ async def _update_request_data_with_litellm_managed_vector_store_registry( vector_store_to_run: Final[LiteLLM_ManagedVectorStore | None] = await get_litellm_managed_vector_store( vector_store_id=vector_store_id ) - if vector_store_to_run is None: + if vector_store_to_run is not None and user_api_key_dict is not None: + await assert_user_can_access_vector_store( + vector_store=vector_store_to_run, + user_api_key_dict=user_api_key_dict, + ) + return apply_managed_vector_store_connection_policy( + data=data, vector_store=vector_store_to_run, user_api_key_dict=user_api_key_dict + ) + + +def apply_managed_vector_store_connection_policy( + data: Mapping[str, object], + vector_store: LiteLLM_ManagedVectorStore | None, + user_api_key_dict: UserAPIKeyAuth | None, +) -> dict[str, object]: + if vector_store is None: caller_data: Final = {key: value for key, value in data.items() if key != MILVUS_ADMIN_CONFIGURED_CONNECTION} if user_api_key_dict is not None: assert_proxy_admin_for_user_supplied_vector_store_connection( @@ -105,18 +120,12 @@ async def _update_request_data_with_litellm_managed_vector_store_registry( user_api_key_dict=user_api_key_dict, ) return caller_data - if user_api_key_dict is not None: - await assert_user_can_access_vector_store( - vector_store=vector_store_to_run, - user_api_key_dict=user_api_key_dict, - ) blocked_fields: Final = managed_connection_fields( - vector_store_to_run.get("custom_llm_provider"), vector_store_to_run.get("litellm_params") + vector_store.get("custom_llm_provider"), vector_store.get("litellm_params") ) - managed_data: Final = build_request_data_from_managed_vector_store(vector_store_to_run) request_data: Final = { **{key: value for key, value in data.items() if key not in blocked_fields}, - **managed_data, + **build_request_data_from_managed_vector_store(vector_store), } return request_data 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 d5dc67a7e55..b3aad243c50 100644 --- a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py @@ -662,6 +662,137 @@ def test_rag_query_gates_managed_milvus_grpc_store_on_admin_approval(client_inte mock_aquery.assert_not_awaited() +def _managed_grpc_store(vector_store_id: str, *, approved: bool) -> dict[str, object]: + from litellm.constants import MILVUS_ADMIN_CONFIGURED_CONNECTION + + return { + "vector_store_id": vector_store_id, + "custom_llm_provider": "milvus", + "litellm_params": { + "custom_llm_provider": "milvus", + "api_base": "http://127.0.0.1:19530", + "api_key": "root:Milvus", + "milvus_transport": "grpc", + "milvus_text_field": "text", + "litellm_embedding_model": "team-embedding", + **({MILVUS_ADMIN_CONFIGURED_CONNECTION: True} if approved else {}), + }, + } + + +def _registry_with(store: dict[str, object]) -> MagicMock: + registry = MagicMock() + registry.get_litellm_managed_vector_store_from_registry.return_value = store + return registry + + +def test_rag_query_rejects_a_non_admin_inline_grpc_connection_on_an_unregistered_store(client_internal_user): + from litellm.llms.milvus.vector_stores.connection import MilvusConnectionRejection + + with ( + patch( # test-quality-ok: aquery is the downstream boundary; the test asserts it is never reached + "litellm.proxy.rag_endpoints.endpoints.litellm.aquery", new_callable=AsyncMock + ) as mock_aquery, + patch( # test-quality-ok: no registry so the id resolves as provider-native + "litellm.vector_store_registry", None + ), + patch( # test-quality-ok: proxy module global, no injection seam + "litellm.proxy.proxy_server.prisma_client", None + ), + ): + response = client_internal_user.post( + "/v1/rag/query", + json={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hello"}], + "retrieval_config": { + "vector_store_id": "vs_unregistered", + "custom_llm_provider": "milvus", + "milvus_transport": "grpc", + "api_base": "http://10.0.0.1:19530", + "api_key": "root:Milvus", + }, + }, + ) + + assert response.status_code == 403, response.text + assert MilvusConnectionRejection.ADMIN_REQUIRED.value in response.text + mock_aquery.assert_not_awaited() + + +def test_rag_query_drops_caller_connection_fields_for_an_approved_managed_grpc_store(client_internal_user): + import litellm + from litellm.constants import MILVUS_ADMIN_CONFIGURED_CONNECTION + from litellm.types.utils import ModelResponse + + 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 downstream boundary; the forwarded config is what the test asserts + "litellm.proxy.rag_endpoints.endpoints.litellm.aquery", new_callable=AsyncMock, return_value=mock_response + ) as mock_aquery, + patch.object( # test-quality-ok: seeds the approved managed store the policy under test reads + litellm, "vector_store_registry", _registry_with(_managed_grpc_store("proof_grpc", approved=True)) + ), + patch( # test-quality-ok: store access is not under test, so the request reaches the connection policy + "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": "proof_grpc", + "api_base": "http://10.0.0.1:19530", + "api_key": "attacker", + "milvus_db_name": "other_db", + "top_k": 3, + }, + }, + ) + + assert response.status_code == 200, response.text + forwarded_config = mock_aquery.await_args.kwargs["retrieval_config"] + assert forwarded_config["api_base"] == "http://127.0.0.1:19530" + assert forwarded_config["api_key"] == "root:Milvus" + assert "milvus_db_name" not in forwarded_config + assert forwarded_config["top_k"] == 3 + assert forwarded_config[MILVUS_ADMIN_CONFIGURED_CONNECTION] is True + + +def test_rag_ingest_rejects_an_unapproved_managed_grpc_store(client_internal_user): + import litellm + from litellm.llms.milvus.vector_stores.connection import MilvusConnectionRejection + + with ( + patch( # test-quality-ok: aingest is the downstream boundary; the test asserts it is never reached + "litellm.proxy.rag_endpoints.endpoints.litellm.aingest", new_callable=AsyncMock + ) as mock_aingest, + patch.object( # test-quality-ok: seeds the unapproved managed store the policy under test reads + litellm, "vector_store_registry", _registry_with(_managed_grpc_store("proof_legacy", approved=False)) + ), + patch( # test-quality-ok: store access is not under test, so the request reaches the connection policy + "litellm.proxy.vector_store_endpoints.utils.can_user_access_vector_store", new=AsyncMock(return_value=True) + ), + ): + response = client_internal_user.post( + "/v1/rag/ingest", + files={"file": ("sample.txt", io.BytesIO(b"test content"), "text/plain")}, + data={ + "request": '{"ingest_options":{"vector_store":{"custom_llm_provider":"milvus","vector_store_id":"proof_legacy"}}}' + }, + ) + + assert response.status_code == 403, response.text + assert MilvusConnectionRejection.ADMIN_SAVE_REQUIRED.value in response.text + mock_aingest.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/type-discipline-budget.json b/type-discipline-budget.json index 6690942bc89..303fd06fdf1 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22171 + "limit": 22172 }, "LIT002": { - "limit": 26728 + "limit": 26727 }, "LIT003": { "limit": 261