mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-17 23:51:30 +00:00
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:
parent
df3050f538
commit
8a9ff04b4d
3 changed files with 79 additions and 1 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue