Merge pull request #39452 from BerriAI/litellm_fix_rag_query_store_credentials

fix(rag): forward the managed vector store's params to the search call

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,

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,
                 )
             )

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"],

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