From 8a9ff04b4d369c227645c5b361ec1f0b65c94e27 Mon Sep 17 00:00:00 2001
From: milan <milan@berri.ai>
Date: Thu, 23 Jul 2026 20:48:01 +0000
Subject: [PATCH] fix(rag): forward retrieval_filter from retrieval_config to
 Bedrock KB search

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 litellm/rag/main.py                 |  2 +
 litellm/types/rag.py                |  1 +
 tests/test_litellm/rag/test_main.py | 77 ++++++++++++++++++++++++++++-
 3 files changed, 79 insertions(+), 1 deletion(-)

diff --git a/litellm/rag/main.py b/litellm/rag/main.py
index 29891ccfd24..a1a5bbd38e9 100644
--- a/litellm/rag/main.py
+++ b/litellm/rag/main.py
@@ -233,10 +233,12 @@ async def _execute_query_pipeline(
         raise ValueError("No query found in messages for RAG query")
 
     # 2. Search vector store
+    filters = retrieval_config.get("retrieval_filter") or retrieval_config.get("filters")
     with _suppressed_sub_call_billing():
         search_response = await litellm.vector_stores.asearch(
             vector_store_id=retrieval_config["vector_store_id"],
             query=query_text,
+            filters=filters,
             max_num_results=retrieval_config.get("top_k", 10),
             custom_llm_provider=retrieval_config.get("custom_llm_provider", "openai"),
             **kwargs,
diff --git a/litellm/types/rag.py b/litellm/types/rag.py
index 802bf7a6e9e..216c5a611ca 100644
--- a/litellm/types/rag.py
+++ b/litellm/types/rag.py
@@ -244,6 +244,7 @@ class RAGRetrievalConfig(TypedDict, total=False):
     custom_llm_provider: str
     top_k: int  # max results from vector store
     filters: Optional[Dict[str, Any]]  # optional - vector store filters
+    retrieval_filter: Optional[Dict[str, Any]]  # optional - alias forwarded as vector store filters
 
 
 class RAGRerankConfig(TypedDict, total=False):
diff --git a/tests/test_litellm/rag/test_main.py b/tests/test_litellm/rag/test_main.py
index 584124ba06a..30ff0b76cd8 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
 
@@ -254,6 +254,81 @@ async def test_aquery_streaming_bills_sub_call_costs_into_final_event():
     assert standard_logging_object["response_cost"] >= 0.003
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("filter_key", ["retrieval_filter", "filters"])
+async def test_aquery_forwards_retrieval_filter_to_vector_store_search(filter_key):
+    """
+    The retrieval_config filter (AWS Bedrock KB metadata filter) must reach the
+    vector store search call. Before the fix it was dropped, so Bedrock ran an
+    unfiltered Retrieve and returned documents from the wrong metadata partition.
+    Both the customer-facing `retrieval_filter` key and the typed `filters` alias
+    must be forwarded as the search `filters` argument.
+    """
+    from litellm.types.vector_stores import VectorStoreSearchResponse
+
+    retrieval_filter = {
+        "andAll": [
+            {"equals": {"key": "Technology", "value": "Blade"}},
+            {"equals": {"key": "Parameter", "value": "Nicotine"}},
+        ]
+    }
+
+    fake_search = AsyncMock(
+        return_value=VectorStoreSearchResponse(
+            object="vector_store.search_results.page",
+            search_query="q",
+            data=[],
+        )
+    )
+
+    with patch("litellm.vector_stores.asearch", new=fake_search):
+        response = await litellm.aquery(
+            model="gpt-4o-mini",
+            messages=[{"role": "user", "content": "most frequent causes of low nicotine"}],
+            retrieval_config={
+                "vector_store_id": "CBVFYF3MYF",
+                "custom_llm_provider": "bedrock",
+                "top_k": 50,
+                filter_key: retrieval_filter,
+            },
+            mock_response="answer",
+        )
+
+    assert isinstance(response, ModelResponse)
+    fake_search.assert_awaited_once()
+    assert fake_search.await_args.kwargs["filters"] == retrieval_filter
+    assert fake_search.await_args.kwargs["vector_store_id"] == "CBVFYF3MYF"
+    assert fake_search.await_args.kwargs["max_num_results"] == 50
+
+
+@pytest.mark.asyncio
+async def test_aquery_without_filter_forwards_none():
+    """
+    When no filter is provided, the search call must receive filters=None rather
+    than a truthy default that would silently constrain an unfiltered query.
+    """
+    from litellm.types.vector_stores import VectorStoreSearchResponse
+
+    fake_search = AsyncMock(
+        return_value=VectorStoreSearchResponse(
+            object="vector_store.search_results.page",
+            search_query="q",
+            data=[],
+        )
+    )
+
+    with patch("litellm.vector_stores.asearch", new=fake_search):
+        await litellm.aquery(
+            model="gpt-4o-mini",
+            messages=[{"role": "user", "content": "hello"}],
+            retrieval_config={"vector_store_id": "vs_test_123", "custom_llm_provider": "openai"},
+            mock_response="hi",
+        )
+
+    fake_search.assert_awaited_once()
+    assert fake_search.await_args.kwargs["filters"] is None
+
+
 def test_rag_call_types_are_registered():
     """
     query/aquery/ingest/aingest are @client-decorated entry points, so their