diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index a4ef30fb3c4..050a59c9fe5 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -26,6 +26,9 @@ from litellm.proxy.common_utils.http_parsing_utils import ( _safe_get_request_headers, get_form_data, ) +from litellm.proxy.vector_store_endpoints.endpoints import ( + _update_request_data_with_litellm_managed_vector_store_registry, +) from litellm.proxy.vector_store_endpoints.utils import ( assert_user_can_access_vector_store_id, ) @@ -647,6 +650,11 @@ async def rag_query( status_code=400, detail={"error": "retrieval_config must contain 'vector_store_id'"}, ) + retrieval_config = await _update_request_data_with_litellm_managed_vector_store_registry( + data=retrieval_config, + vector_store_id=retrieval_config["vector_store_id"], + user_api_key_dict=user_api_key_dict, + ) await _authorize_nested_vector_store_ids( payload=retrieval_config, user_api_key_dict=user_api_key_dict, diff --git a/litellm/rag/main.py b/litellm/rag/main.py index 64b55b4320f..027c199f0d5 100644 --- a/litellm/rag/main.py +++ b/litellm/rag/main.py @@ -226,13 +226,18 @@ async def _execute_query_pipeline( raise ValueError("No query found in messages for RAG query") # 2. Search vector store + _reserved_retrieval_keys = frozenset( + ("vector_store_id", "top_k", "custom_llm_provider", "query", "max_num_results") + ) + search_params = {k: v for k, v in {**retrieval_config, **kwargs}.items() if k not in _reserved_retrieval_keys} + with _suppressed_sub_call_billing(): search_response = await litellm.vector_stores.asearch( vector_store_id=retrieval_config["vector_store_id"], query=query_text, max_num_results=retrieval_config.get("top_k", 10), custom_llm_provider=retrieval_config.get("custom_llm_provider", "openai"), - **kwargs, + **search_params, ) search_provider = retrieval_config.get("custom_llm_provider", "openai") 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 15a117bd6fc..dba3ed4c462 100644 --- a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py @@ -13,9 +13,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi.testclient import TestClient -sys.path.insert( - 0, os.path.abspath("../../../../..") -) # Adds the parent directory to the system path +sys.path.insert(0, os.path.abspath("../../../../..")) # Adds the parent directory to the system path from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -62,19 +60,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 @@ -101,8 +93,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()}" ) @@ -118,15 +109,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()}" ) @@ -174,13 +162,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 @@ -197,9 +185,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": { @@ -219,9 +205,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): @@ -234,14 +218,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): @@ -265,12 +245,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", @@ -306,10 +288,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={ @@ -327,3 +313,73 @@ def test_rag_query_stream_returns_event_stream(client_internal_user): assert response.headers.get("content-type", "").startswith("text/event-stream") assert '"object":"chat.completion.chunk"' in response.text assert "data: [DONE]" in response.text + + +def test_rag_query_resolves_vector_store_registry_credentials(client_internal_user): + """ + /v1/rag/query must resolve vector store registry entries and pass the resolved + custom_llm_provider, api_key, api_base, etc. into litellm.vector_stores.asearch. + """ + import litellm + + mock_vector_store = { + "vector_store_id": "vs_azure_123", + "custom_llm_provider": "azure_ai_search", + "litellm_credential_name": "azure_search_cred", + "litellm_params": { + "api_key": "fake-azure-search-key", + "api_base": "https://fake-azure-search.search.windows.net", + "api_version": "2023-11-01", + }, + } + mock_registry = MagicMock() + mock_registry.get_litellm_managed_vector_store_from_registry.return_value = mock_vector_store + mock_registry.pop_vector_stores_to_run_with_db_fallback = AsyncMock(return_value=None) + + captured_asearch_kwargs = {} + + async def fake_asearch(**kwargs): + captured_asearch_kwargs.update(kwargs) + mock_search_resp = MagicMock() + mock_search_resp.data = [] + return mock_search_resp + + mock_model_response = MagicMock( + id="chatcmpl-test", + choices=[ + MagicMock( + message=MagicMock(role="assistant", content="Answer"), + finish_reason="stop", + ) + ], + _hidden_params={}, + model="gpt-4o-mini", + ) + + with ( + patch.object(litellm, "vector_store_registry", mock_registry), + patch( + "litellm.vector_stores.asearch", + new=AsyncMock(side_effect=fake_asearch), + ), + patch( + "litellm.acompletion", + new=AsyncMock(return_value=mock_model_response), + ), + ): + response = client_internal_user.post( + "/v1/rag/query", + json={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "search test query"}], + "retrieval_config": { + "vector_store_id": "vs_azure_123", + }, + }, + ) + + assert response.status_code == 200, response.json() + assert captured_asearch_kwargs.get("custom_llm_provider") == "azure_ai_search" + assert captured_asearch_kwargs.get("api_key") == "fake-azure-search-key" + assert captured_asearch_kwargs.get("api_base") == "https://fake-azure-search.search.windows.net" + assert captured_asearch_kwargs.get("api_version") == "2023-11-01" diff --git a/tests/test_litellm/rag/test_main.py b/tests/test_litellm/rag/test_main.py index 584124ba06a..70559629535 100644 --- a/tests/test_litellm/rag/test_main.py +++ b/tests/test_litellm/rag/test_main.py @@ -11,7 +11,7 @@ aquery carries the completion response with real usage and cost. """ import asyncio -from unittest.mock import patch +from unittest.mock import AsyncMock, patch import pytest @@ -89,6 +89,35 @@ async def test_aquery_single_billing_event_carries_completion_usage_and_cost(use assert standard_logging_object["response_cost"] > 0 +@pytest.mark.asyncio +async def test_aquery_forwards_provider_credentials_to_vector_store_search(): + """Retrieval config credentials must be forwarded into the vector store search call.""" + captured_kwargs = {} + + async def fake_asearch(**kwargs): + captured_kwargs.update(kwargs) + return {"data": []} + + with patch("litellm.vector_stores.asearch", new=AsyncMock(side_effect=fake_asearch)): + await litellm.aquery( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hello"}], + retrieval_config={ + "vector_store_id": "vs_test_123", + "custom_llm_provider": "azure_ai_search", + "api_key": "search-key", + "api_base": "https://example.com", + "api_version": "2024-01-01", + }, + mock_response="hi there", + ) + + assert captured_kwargs["custom_llm_provider"] == "azure_ai_search" + assert captured_kwargs["api_key"] == "search-key" + assert captured_kwargs["api_base"] == "https://example.com" + assert captured_kwargs["api_version"] == "2024-01-01" + + @pytest.mark.asyncio async def test_aquery_response_hidden_params_carry_completion_cost(): """