diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index e144ff965ae..c8c6c505375 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -761,6 +761,7 @@ async def rag_query( model=model, messages=messages, retrieval_config=merged_retrieval_config, + vector_store_params=store_data, rerank=rerank, stream=stream, router=llm_router, diff --git a/litellm/rag/main.py b/litellm/rag/main.py index 94bfc305a6a..1f63152632e 100644 --- a/litellm/rag/main.py +++ b/litellm/rag/main.py @@ -11,7 +11,7 @@ __all__ = ["aingest", "aquery", "ingest", "query"] import asyncio import contextvars -from collections.abc import Coroutine, Iterator +from collections.abc import Coroutine, Iterator, Mapping from contextlib import contextmanager from functools import partial from types import MappingProxyType @@ -66,6 +66,10 @@ _FORWARDABLE_RETRIEVAL_CONFIG_KEYS: Final = frozenset( } ) +_SEARCH_ARGS_SET_BY_PIPELINE: Final = frozenset( + {"vector_store_id", "query", "max_num_results", "custom_llm_provider", "router"} +) + def get_ingestion_class(provider: str) -> type[BaseRAGIngestion]: """ @@ -225,6 +229,7 @@ async def _execute_query_pipeline( retrieval_config: dict[str, Any], rerank: dict[str, Any] | None = None, stream: bool = False, + vector_store_params: Mapping[str, object] | None = None, **kwargs, ) -> ModelResponse: """ @@ -241,11 +246,19 @@ async def _execute_query_pipeline( # 2. Search vector store # Forward allowlisted provider retrieval_config extras (region, embedding - # model, bucket, credential refs) to the search call; kwargs win on conflict. + # model, bucket, credential refs) to the search call; the managed store's + # params win on conflict. provider_search_params: Final = MappingProxyType( {k: v for k, v in retrieval_config.items() if k in _FORWARDABLE_RETRIEVAL_CONFIG_KEYS} ) - forwarded_search_params: Final = MappingProxyType({**provider_search_params, **kwargs}) + store_search_params: Final = MappingProxyType( + { + k: v + for k, v in (vector_store_params.items() if vector_store_params else ()) + if k not in _SEARCH_ARGS_SET_BY_PIPELINE + } + ) + forwarded_search_params: Final = MappingProxyType({**provider_search_params, **kwargs, **store_search_params}) with _suppressed_sub_call_billing(): search_response: Final = await litellm.vector_stores.asearch( vector_store_id=retrieval_config["vector_store_id"], @@ -339,6 +352,7 @@ async def aquery( retrieval_config: dict[str, Any], rerank: dict[str, Any] | None = None, stream: bool = False, + vector_store_params: Mapping[str, object] | None = None, **kwargs, ) -> ModelResponse: """ @@ -356,6 +370,7 @@ async def aquery( retrieval_config=retrieval_config, rerank=rerank, stream=stream, + vector_store_params=vector_store_params, **kwargs, ) @@ -386,6 +401,7 @@ def query( retrieval_config: dict[str, Any], rerank: dict[str, Any] | None = None, stream: bool = False, + vector_store_params: Mapping[str, object] | None = None, **kwargs, ) -> ModelResponse | Coroutine[None, None, ModelResponse]: """ @@ -402,6 +418,7 @@ def query( retrieval_config=retrieval_config, rerank=rerank, stream=stream, + vector_store_params=vector_store_params, **kwargs, ) else: @@ -412,6 +429,7 @@ def query( retrieval_config=retrieval_config, rerank=rerank, stream=stream, + vector_store_params=vector_store_params, **kwargs, ) ) 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 0085b6ebd36..a176e91eaa4 100644 --- a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py @@ -421,6 +421,76 @@ def test_rag_query_store_params_win_over_user_retrieval_config(client_internal_u assert forwarded_config["aws_region_name"] == "eu-west-1" +def test_rag_query_forwards_managed_store_credentials_to_search(client_internal_user): + """ + Regression for LIT-6773: the registry store's api_key / api_base and its + provider extras (Milvus outputFields, milvus_text_field) must reach the + vector store search the way the direct /v1/vector_stores/{id}/search + endpoint forwards them. Pre-fix the RAG path allowlisted them away and a + managed Milvus store 500'd with "MILVUS_API_KEY is not set". + """ + import litellm + from litellm import Router + from litellm.types.vector_stores import VectorStoreSearchResponse + + mock_vector_store = { + "vector_store_id": "customer_kb", + "custom_llm_provider": "milvus", + "litellm_params": { + "vector_store_id": "customer_kb", + "custom_llm_provider": "milvus", + "api_base": "http://127.0.0.1:19530", + "api_key": "root:Milvus", + "litellm_embedding_model": "multilingual-e5-large", + "milvus_text_field": "book_intro_text", + "outputFields": ["book_intro_text"], + }, + } + mock_registry = MagicMock() + mock_registry.get_litellm_managed_vector_store_from_registry.return_value = mock_vector_store + fake_search = AsyncMock( + return_value=VectorStoreSearchResponse(object="vector_store.search_results.page", search_query="q", data=[]) + ) + router = Router( + model_list=[ + { + "model_name": "gpt-4o-mini", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test", "mock_response": "hi"}, + } + ] + ) + + with ( + patch("litellm.vector_stores.asearch", new=fake_search), # test-quality-ok: the search boundary under test + patch.object(litellm, "vector_store_registry", mock_registry), # test-quality-ok: seeds the store under test + patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: mock-response router for completion + patch( # test-quality-ok: store access is not under test, so the request reaches the search boundary + "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": "which database is built for similarity search?"}], + "retrieval_config": {"vector_store_id": "customer_kb", "custom_llm_provider": "milvus", "top_k": 2}, + }, + ) + + assert response.status_code == 200, response.json() + fake_search.assert_awaited_once() + search_kwargs = fake_search.await_args.kwargs + assert search_kwargs["vector_store_id"] == "customer_kb" + assert search_kwargs["custom_llm_provider"] == "milvus" + assert search_kwargs["max_num_results"] == 2 + assert search_kwargs["api_base"] == "http://127.0.0.1:19530" + assert search_kwargs["api_key"] == "root:Milvus" + assert search_kwargs["litellm_embedding_model"] == "multilingual-e5-large" + assert search_kwargs["milvus_text_field"] == "book_intro_text" + assert search_kwargs["outputFields"] == ["book_intro_text"] + + @pytest.mark.parametrize( "blocked_key", ["embedding_model", "litellm_embedding_model", "litellm_embedding_config", "litellm_credential_name"], diff --git a/tests/test_litellm/rag/test_main.py b/tests/test_litellm/rag/test_main.py index 51d03544910..264bcd6fb75 100644 --- a/tests/test_litellm/rag/test_main.py +++ b/tests/test_litellm/rag/test_main.py @@ -388,6 +388,72 @@ async def test_aquery_does_not_forward_connection_override_keys_to_search(): assert not (blocked & set(search_kwargs.keys())) +@pytest.mark.asyncio +async def test_aquery_forwards_vector_store_params_to_search_but_not_completion(): + """ + Regression for LIT-6773: the server-trusted vector_store_params (a managed + store's litellm_params) must reach the search call wholesale, including the + connection keys the caller allowlist blocks, while the caller's own + retrieval_config overrides stay blocked, the caller's top-level api_key and + api_base stay on the completion only, and the completion never inherits the + store's connection params. + """ + from unittest.mock import AsyncMock + + from litellm.types.vector_stores import VectorStoreSearchResponse + + fake_search = AsyncMock( + return_value=VectorStoreSearchResponse( + object="vector_store.search_results.page", search_query="q", data=[] + ) + ) + fake_completion = AsyncMock( + return_value=ModelResponse( + id="chatcmpl-test", + choices=[{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + model="gpt-4o-mini", + ) + ) + with ( + patch("litellm.vector_stores.asearch", new=fake_search), # test-quality-ok: the search boundary under test + patch("litellm.acompletion", new=fake_completion), # test-quality-ok: the completion boundary under test + ): + await litellm.aquery( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hello"}], + api_key="sk-llm-key", + api_base="https://llm.example.com", + retrieval_config={ + "vector_store_id": "customer_kb", + "custom_llm_provider": "milvus", + "api_base": "https://attacker.example.com", + "api_key": "attacker-key", + }, + vector_store_params={ + "vector_store_id": "customer_kb", + "custom_llm_provider": "milvus", + "api_base": "http://127.0.0.1:19530", + "api_key": "root:Milvus", + "milvus_text_field": "book_intro_text", + "outputFields": ["book_intro_text"], + }, + ) + + fake_search.assert_awaited_once() + search_kwargs = fake_search.await_args.kwargs + assert search_kwargs["vector_store_id"] == "customer_kb" + assert search_kwargs["custom_llm_provider"] == "milvus" + assert search_kwargs["api_base"] == "http://127.0.0.1:19530" + assert search_kwargs["api_key"] == "root:Milvus" + assert search_kwargs["milvus_text_field"] == "book_intro_text" + assert search_kwargs["outputFields"] == ["book_intro_text"] + fake_completion.assert_awaited_once() + completion_kwargs = fake_completion.await_args.kwargs + assert completion_kwargs["api_key"] == "sk-llm-key" + assert completion_kwargs["api_base"] == "https://llm.example.com" + assert not ({"milvus_text_field", "outputFields"} & set(completion_kwargs)) + + def test_rag_call_types_are_registered(): """ query/aquery/ingest/aingest are @client-decorated entry points, so their