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>
This commit is contained in:
milan 2026-07-23 20:48:01 +00:00
parent df3050f538
commit 8a9ff04b4d
3 changed files with 79 additions and 1 deletions

View file

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

View file

@ -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):

View file

@ -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