diff --git a/litellm/rag/main.py b/litellm/rag/main.py index 8ddc4c231dd..1f63152632e 100644 --- a/litellm/rag/main.py +++ b/litellm/rag/main.py @@ -246,7 +246,8 @@ 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} ) @@ -257,7 +258,7 @@ async def _execute_query_pipeline( if k not in _SEARCH_ARGS_SET_BY_PIPELINE } ) - forwarded_search_params: Final = MappingProxyType({**provider_search_params, **store_search_params, **kwargs}) + 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"], diff --git a/tests/test_litellm/rag/test_main.py b/tests/test_litellm/rag/test_main.py index f119088b0e9..264bcd6fb75 100644 --- a/tests/test_litellm/rag/test_main.py +++ b/tests/test_litellm/rag/test_main.py @@ -394,8 +394,9 @@ 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 and the completion never inherits - the store's connection params. + 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 @@ -420,6 +421,8 @@ async def test_aquery_forwards_vector_store_params_to_search_but_not_completion( 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", @@ -445,8 +448,10 @@ async def test_aquery_forwards_vector_store_params_to_search_but_not_completion( assert search_kwargs["milvus_text_field"] == "book_intro_text" assert search_kwargs["outputFields"] == ["book_intro_text"] fake_completion.assert_awaited_once() - store_only_keys = {"api_base", "api_key", "milvus_text_field", "outputFields"} - assert not (store_only_keys & set(fake_completion.await_args.kwargs)) + 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():