diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index 4d62f1d6d71..3f4055ebe3b 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -8,6 +8,7 @@ Provides: import base64 from collections.abc import Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final import orjson @@ -44,6 +45,43 @@ if TYPE_CHECKING: router: Final = APIRouter() +# Fields that must never be logged verbatim from a client-supplied +# ``retrieval_config`` (provider-native stores legitimately carry connection +# credentials in the dict; the debug log below redacts them). +_SENSITIVE_RETRIEVAL_CONFIG_FIELDS: Final = frozenset( + { + "api_key", + "aws_access_key_id", + "aws_secret_access_key", + "aws_session_token", + "vertex_credentials", + "vertex_ai_credentials", + "litellm_embedding_config", + } +) + +# Connection credential/endpoint fields that an authenticated client must not +# be able to inject into the vector store search through ``retrieval_config``. +# For a managed store the operator-resolved registry params are authoritative +# (they travel on a dedicated channel); for a provider-native store connection +# params must come from the proxy's server-side configuration (env / litellm +# credentials), never from the request body - a client-supplied ``api_base`` +# would let a tenant pivot the proxy's egress to an internal host (SSRF). +_RETRIEVAL_CONFIG_CLIENT_BLOCKED_FIELDS: Final = frozenset( + { + "api_key", + "api_base", + "api_version", + "aws_region_name", + "region_name", + "aws_access_key_id", + "aws_secret_access_key", + "aws_session_token", + "litellm_embedding_config", + "litellm_embedding_model", + } +) + def _raise_vector_store_scan_depth_exceeded() -> None: raise HTTPException( @@ -524,6 +562,7 @@ async def rag_ingest( # Add litellm data request_data: dict[str, Any] = {} + # rebind-ok: request_data is rebuilt from the litellm-enriched request. request_data = await add_litellm_data_to_request( data=request_data, request=request, @@ -693,10 +732,17 @@ async def rag_query( user_api_key_dict=user_api_key_dict, ) + # Deferred import: keeps vector_store_endpoints out of sys.modules at + # proxy boot so lazy router registration and the OpenAPI snapshot stub + # injection are not skipped for vector store routes. + from litellm.proxy.vector_store_endpoints.endpoints import ( + _update_request_data_with_litellm_managed_vector_store_registry, # pyright: ignore[reportPrivateUsage] # one canonical managed-store registry resolver, shared with the vector store endpoints + ) + # Add litellm data request_data: dict[str, object] = {} request_data = await add_litellm_data_to_request( - data=request_data, + data={}, # mutable-ok: resolver populates this dict in place request=request, general_settings=general_settings, user_api_key_dict=user_api_key_dict, @@ -704,15 +750,51 @@ async def rag_query( proxy_config=proxy_config, ) - verbose_proxy_logger.debug("RAG Query - model: %s, retrieval_config: %s", model, retrieval_config) + resolved_registry = await _update_request_data_with_litellm_managed_vector_store_registry( + data={}, # mutable-ok: resolver populates this dict in place + vector_store_id=retrieval_config["vector_store_id"], + user_api_key_dict=user_api_key_dict, + ) + resolved_registry.pop("litellm_credential_name", None) + # An authenticated client must not control the search connection + # params: strip credential/endpoint fields from the client-supplied + # retrieval_config. Connection params reach the search only through + # the trusted managed-store registry channel below (or the proxy's + # own server-side configuration for provider-native stores). + search_retrieval_config: Final[Mapping[str, Any]] = MappingProxyType( + { + key: value + for key, value in retrieval_config.items() + if key not in _RETRIEVAL_CONFIG_CLIENT_BLOCKED_FIELDS + } + ) + # Trusted registry params travel on a separate channel (never merged + # into the client-controlled ``retrieval_config``), so client-supplied + # credential/endpoint fields can never be forwarded to the search for + # a managed store, and resolved credentials are never logged. The + # search call site forwards only an allowlisted subset of this dict + # (see ``_VECTOR_STORE_SEARCH_PARAMS`` in ``litellm/rag/main.py``), + # so stored metadata / guardrail fields never reach the LLM call. + + verbose_proxy_logger.debug( + "RAG Query - model: %s, retrieval_config: %s", + model, + MappingProxyType( + { + key: ("***" if key in _SENSITIVE_RETRIEVAL_CONFIG_FIELDS and value else value) + for key, value in search_retrieval_config.items() + } + ), + ) # Call query response: Final = await litellm.aquery( model=model, messages=messages, - retrieval_config=retrieval_config, + retrieval_config=search_retrieval_config, rerank=rerank, stream=stream, + litellm_managed_vector_store_registry=resolved_registry, router=llm_router, **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 abbf6892a98..14afb932245 100644 --- a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py @@ -11,7 +11,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi.testclient import TestClient - +import litellm from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.proxy_server import app @@ -57,19 +57,13 @@ def test_internal_user_viewer_rag_ingest_without_vector_store_id_rejected( response = client_internal_user_viewer.post( "/v1/rag/ingest", files={"file": ("sample.txt", io.BytesIO(b"test content"), "text/plain")}, - data={ - "request": '{"ingest_options":{"vector_store":{"custom_llm_provider":"openai"}}}' - }, + data={"request": '{"ingest_options":{"vector_store":{"custom_llm_provider":"openai"}}}'}, ) assert response.status_code == 403 detail = response.json() assert "detail" in detail - error_msg = ( - detail["detail"]["error"] - if isinstance(detail["detail"], dict) - else str(detail["detail"]) - ) + error_msg = detail["detail"]["error"] if isinstance(detail["detail"], dict) else str(detail["detail"]) assert "internal_user_viewer" in error_msg assert "vector_store_id" in error_msg @@ -96,8 +90,7 @@ def test_internal_user_viewer_rag_ingest_with_vector_store_id_passes_check( # Should not be 403 (role check passed) assert response.status_code != 403, ( - f"internal_user_viewer with vector_store_id should pass role check. " - f"Response: {response.json()}" + f"internal_user_viewer with vector_store_id should pass role check. Response: {response.json()}" ) @@ -113,15 +106,12 @@ def test_internal_user_rag_ingest_without_vector_store_id_allowed(client_interna 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":"openai"}}}' - }, + data={"request": '{"ingest_options":{"vector_store":{"custom_llm_provider":"openai"}}}'}, ) # Should not be 403 assert response.status_code != 403, ( - f"internal_user should be allowed to create new vector stores. " - f"Response: {response.json()}" + f"internal_user should be allowed to create new vector stores. Response: {response.json()}" ) @@ -169,13 +159,13 @@ def test_rag_ingest_blocks_clientside_credentials(client_internal_user, blocked_ }, }, ) - assert ( - response.status_code == 400 - ), f"Expected 400 when '{blocked_field}' is set clientside, got {response.status_code}: {response.json()}" + assert response.status_code == 400, ( + f"Expected 400 when '{blocked_field}' is set clientside, got {response.status_code}: {response.json()}" + ) body = response.json() - assert blocked_field in str( - body - ), f"Response should mention '{blocked_field}': {body}" + assert blocked_field in str(body), f"Response should mention '{blocked_field}': {body}" + + class TestRagIngestSSRFBlocked: """ aws_sts_endpoint and related credential-redirect fields must be rejected @@ -192,9 +182,7 @@ class TestRagIngestSSRFBlocked: ("aws_bedrock_runtime_endpoint", "https://attacker.example/bedrock"), ], ) - def test_ssrf_field_in_vector_store_config_rejected( - self, field, value, client_internal_user - ): + def test_ssrf_field_in_vector_store_config_rejected(self, field, value, client_internal_user): payload = { "file_url": "https://example.com/doc.pdf", "ingest_options": { @@ -214,9 +202,7 @@ class TestRagIngestSSRFBlocked: ) body = response.json() detail = body.get("detail", {}) - error_text = ( - detail.get("error", "") if isinstance(detail, dict) else str(detail) - ) + error_text = detail.get("error", "") if isinstance(detail, dict) else str(detail) assert field in error_text, f"Error should name the offending field: {error_text}" def test_clean_bedrock_ingest_options_not_rejected(self, client_internal_user): @@ -229,14 +215,10 @@ class TestRagIngestSSRFBlocked: "/v1/rag/ingest", json={ "file_url": "https://example.com/doc.pdf", - "ingest_options": { - "vector_store": {"custom_llm_provider": "bedrock"} - }, + "ingest_options": {"vector_store": {"custom_llm_provider": "bedrock"}}, }, ) - assert response.status_code != 400, ( - f"Clean Bedrock ingest_options should not be rejected: {response.json()}" - ) + assert response.status_code != 400, f"Clean Bedrock ingest_options should not be rejected: {response.json()}" def test_rag_query_returns_response_cost_header(client_internal_user): @@ -260,12 +242,14 @@ def test_rag_query_returns_response_cost_header(client_internal_user): ) mock_response._hidden_params["response_cost"] = 3.45e-06 - with patch( - "litellm.proxy.rag_endpoints.endpoints.litellm.aquery", - new_callable=AsyncMock, - return_value=mock_response, - ), patch("litellm.vector_store_registry", None), patch( - "litellm.proxy.proxy_server.prisma_client", None + with ( + patch( + "litellm.proxy.rag_endpoints.endpoints.litellm.aquery", + new_callable=AsyncMock, + return_value=mock_response, + ), + patch("litellm.vector_store_registry", None), + patch("litellm.proxy.proxy_server.prisma_client", None), ): response = client_internal_user.post( "/v1/rag/query", @@ -283,6 +267,167 @@ def test_rag_query_returns_response_cost_header(client_internal_user): assert response.headers.get("x-litellm-response-cost") == "3.45e-06" +@pytest.mark.asyncio +async def test_rag_query_resolves_managed_vector_store_registry(client_internal_user): + """ + /v1/rag/query must resolve the LiteLLM-managed vector store registry so the + configured provider and litellm_params (api_key, api_base) reach the + internal vector store search instead of defaulting to OpenAI. + + Before the fix, retrieval_config["custom_llm_provider"] stayed whatever the + client sent and the registry credentials were never forwarded to + litellm.aquery, so a registered non-OpenAI vector store was always queried + via api.openai.com. The fix also keeps the registry params on a dedicated + channel (``litellm_managed_vector_store_registry``), separate from the + client-controlled ``retrieval_config``: stored fields such as metadata + must never surface as top-level kwargs, which flow to the generation call. + """ + from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( + LiteLLM_ManagedVectorStore, + ) + from litellm.types.utils import ModelResponse + + mock_vector_store: LiteLLM_ManagedVectorStore = { + "vector_store_id": "vs_registered", + "custom_llm_provider": "bedrock", + "litellm_credential_name": "my_bedrock_creds", + "litellm_params": { + "aws_region_name": "us-east-1", + # Security probe: a stored field that must never reach the + # generation-call kwargs / sanitization boundary. + "metadata": {"disable_global_guardrails": True}, + }, + } + + mock_registry = MagicMock() + mock_registry.get_litellm_managed_vector_store_from_registry.return_value = mock_vector_store + + captured = {} + + async def fake_aquery(**kwargs): + captured.update(kwargs) + return ModelResponse( + id="chatcmpl-test", + choices=[ + { + "index": 0, + "message": { + "role": "assistant", + "content": "The codename is AZURE-FALCON-42.", + }, + "finish_reason": "stop", + } + ], + model="gpt-4o-mini", + usage={"prompt_tokens": 35, "completion_tokens": 14, "total_tokens": 49}, + ) + + with ( + patch( + "litellm.proxy.rag_endpoints.endpoints.litellm.aquery", + new=AsyncMock(side_effect=fake_aquery), + ), + patch.object(litellm, "vector_store_registry", mock_registry), + patch("litellm.proxy.proxy_server.prisma_client", None), + ): + response = client_internal_user.post( + "/v1/rag/query", + json={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "What is the codename?"}], + "retrieval_config": { + "vector_store_id": "vs_registered", + "custom_llm_provider": "openai", + }, + }, + ) + + assert response.status_code == 200, response.text + # Registry params travel on the dedicated managed-store channel, separate + # from the client-supplied retrieval_config (which is left untouched)... + assert captured["retrieval_config"]["custom_llm_provider"] == "openai" + assert captured["retrieval_config"]["vector_store_id"] == "vs_registered" + assert "aws_region_name" not in captured["retrieval_config"] + resolved_registry = captured["litellm_managed_vector_store_registry"] + assert resolved_registry["custom_llm_provider"] == "bedrock" + assert resolved_registry["aws_region_name"] == "us-east-1" + assert resolved_registry["metadata"] == {"disable_global_guardrails": True} + # ...and never surface as top-level kwargs, which flow to the generation + # call. The credential name is popped (the search resolves it by id). + assert "aws_region_name" not in captured + # `metadata` at the top level is the request's own metadata (agent_id, + # endpoint, headers...); the registry's guardrail override must not leak + # into it, or the generation call would skip global guardrails. + assert "disable_global_guardrails" not in captured.get("metadata", {}) + assert "litellm_credential_name" not in captured + assert "litellm_credential_name" not in captured["retrieval_config"] + assert "litellm_credential_name" not in resolved_registry + + +@pytest.mark.asyncio +async def test_rag_query_strips_client_connection_params(client_internal_user): + """ + Credential/endpoint fields (api_base, api_key, aws creds, ...) supplied by + the client inside retrieval_config must never reach the vector store + search: for a provider-native store they would let a tenant point the + proxy's egress at an arbitrary host (SSRF), and for a managed store the + operator-resolved registry params are authoritative. The proxy strips + these fields before calling litellm.aquery. + """ + from litellm.types.utils import ModelResponse + + captured = {} + + async def fake_aquery(**kwargs): + captured.update(kwargs) + return ModelResponse( + id="chatcmpl-test", + choices=[ + { + "index": 0, + "message": { + "role": "assistant", + "content": "The codename is AZURE-FALCON-42.", + }, + "finish_reason": "stop", + } + ], + model="gpt-4o-mini", + usage={"prompt_tokens": 35, "completion_tokens": 14, "total_tokens": 49}, + ) + + with ( + patch( + "litellm.proxy.rag_endpoints.endpoints.litellm.aquery", + new=AsyncMock(side_effect=fake_aquery), + ), + patch("litellm.vector_store_registry", None), + patch("litellm.proxy.proxy_server.prisma_client", None), + ): + response = client_internal_user.post( + "/v1/rag/query", + json={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "What is the codename?"}], + "retrieval_config": { + "vector_store_id": "vs_test_123", + "custom_llm_provider": "openai", + "api_base": "http://attacker.example/internal", + "api_key": "sk-attacker", + "aws_secret_access_key": "AKIATTACKER", + }, + }, + ) + + assert response.status_code == 200, response.text + retrieval_config = captured["retrieval_config"] + assert retrieval_config["vector_store_id"] == "vs_test_123" + assert retrieval_config["custom_llm_provider"] == "openai" + assert "api_base" not in retrieval_config + assert "api_key" not in retrieval_config + assert "aws_secret_access_key" not in retrieval_config + + def test_rag_query_stream_returns_event_stream(client_internal_user): """ A stream=true /v1/rag/query must return an SSE response. Returning the raw @@ -301,10 +446,14 @@ def test_rag_query_stream_returns_event_stream(client_internal_user): api_key="test-key", ) - with patch( - "litellm.proxy.rag_endpoints.endpoints.litellm.aquery", - new=AsyncMock(side_effect=fake_aquery), - ), patch("litellm.vector_store_registry", None), patch("litellm.proxy.proxy_server.prisma_client", None): + with ( + patch( + "litellm.proxy.rag_endpoints.endpoints.litellm.aquery", + new=AsyncMock(side_effect=fake_aquery), + ), + patch("litellm.vector_store_registry", None), + patch("litellm.proxy.proxy_server.prisma_client", None), + ): response = client_internal_user.post( "/v1/rag/query", json={