From 8a9ff04b4d369c227645c5b361ec1f0b65c94e27 Mon Sep 17 00:00:00 2001 From: milan Date: Thu, 23 Jul 2026 20:48:01 +0000 Subject: [PATCH 1/5] 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 From 232b9e8e78ff5b7d14376bf7d72173cb469a05f6 Mon Sep 17 00:00:00 2001 From: milan Date: Thu, 23 Jul 2026 21:04:42 +0000 Subject: [PATCH 2/5] fix(rag): consume top-level filters kwarg to avoid duplicate keyword in search Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/rag/main.py | 3 ++- tests/test_litellm/rag/test_main.py | 35 +++++++++++++++++++++++++++++ 2 files changed, 37 insertions(+), 1 deletion(-) diff --git a/litellm/rag/main.py b/litellm/rag/main.py index a1a5bbd38e9..628680e8a62 100644 --- a/litellm/rag/main.py +++ b/litellm/rag/main.py @@ -233,7 +233,8 @@ 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") + kwargs_filters = kwargs.pop("filters", None) + filters = retrieval_config.get("retrieval_filter") or retrieval_config.get("filters") or kwargs_filters with _suppressed_sub_call_billing(): search_response = 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 30ff0b76cd8..796f8c38750 100644 --- a/tests/test_litellm/rag/test_main.py +++ b/tests/test_litellm/rag/test_main.py @@ -301,6 +301,41 @@ async def test_aquery_forwards_retrieval_filter_to_vector_store_search(filter_ke assert fake_search.await_args.kwargs["max_num_results"] == 50 +@pytest.mark.asyncio +async def test_aquery_top_level_filters_kwarg_does_not_collide(): + """ + An SDK caller may pass a top-level `filters` kwarg (it used to flow to the + search via **kwargs). Now that the pipeline passes `filters` explicitly, the + top-level kwarg must be consumed rather than forwarded twice, otherwise + asearch raises TypeError for a duplicate keyword before any search runs. + """ + from litellm.types.vector_stores import VectorStoreSearchResponse + + top_level_filter = {"equals": {"key": "tenant", "value": "a"}} + + 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": "hello"}], + retrieval_config={"vector_store_id": "vs_test_123", "custom_llm_provider": "openai"}, + filters=top_level_filter, + mock_response="hi", + ) + + assert isinstance(response, ModelResponse) + fake_search.assert_awaited_once() + assert fake_search.await_args.kwargs["filters"] == top_level_filter + assert "filters" not in fake_search.await_args.kwargs.get("kwargs", {}) + + @pytest.mark.asyncio async def test_aquery_without_filter_forwards_none(): """ From c6a5aab74c9a2c11b2d55ba289659a4ff3942ad5 Mon Sep 17 00:00:00 2001 From: milan Date: Thu, 23 Jul 2026 21:27:25 +0000 Subject: [PATCH 3/5] ci: run orphaned tests/test_litellm/rag suite in misc shard Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/test-unit-misc.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/test-unit-misc.yml b/.github/workflows/test-unit-misc.yml index 7c3b195f0ad..34d5eef6caa 100644 --- a/.github/workflows/test-unit-misc.yml +++ b/.github/workflows/test-unit-misc.yml @@ -35,6 +35,7 @@ jobs: tests/test_litellm/interactions tests/test_litellm/ocr tests/test_litellm/passthrough + tests/test_litellm/rag tests/test_litellm/sandbox tests/test_litellm/vector_stores tests/test_litellm/videos From d2c574608ad430bcc5f2f9bd00fd603a0c3aacbd Mon Sep 17 00:00:00 2001 From: milan Date: Tue, 25 Aug 2026 14:56:22 +0000 Subject: [PATCH 4/5] test(rag): validate retrieval filters at HTTP boundary Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/types/rag.py | 13 +- tests/test_litellm/rag/test_main.py | 176 ++++++++++++---------------- 2 files changed, 82 insertions(+), 107 deletions(-) diff --git a/litellm/types/rag.py b/litellm/types/rag.py index bed3fdfa7e9..629979afde9 100644 --- a/litellm/types/rag.py +++ b/litellm/types/rag.py @@ -2,10 +2,11 @@ Type definitions for RAG (Retrieval Augmented Generation) Ingest API. """ +from collections.abc import Mapping from typing import Any, Literal from pydantic import BaseModel, ConfigDict -from typing_extensions import TypedDict +from typing_extensions import ReadOnly, TypedDict from litellm.types.utils import ModelResponse @@ -237,11 +238,11 @@ class RAGIngestRequest(BaseModel): class RAGRetrievalConfig(TypedDict, total=False): """Configuration for vector store retrieval.""" - vector_store_id: str - custom_llm_provider: str - top_k: int # max results from vector store - filters: dict[str, Any] | None # optional - vector store filters - retrieval_filter: dict[str, Any] | None # optional - alias forwarded as vector store filters + vector_store_id: ReadOnly[str] + custom_llm_provider: ReadOnly[str] + top_k: ReadOnly[int] + filters: ReadOnly[Mapping[str, object] | None] + retrieval_filter: ReadOnly[Mapping[str, object] | None] class RAGRerankConfig(TypedDict, total=False): diff --git a/tests/test_litellm/rag/test_main.py b/tests/test_litellm/rag/test_main.py index 796f8c38750..ab0a9c0f002 100644 --- a/tests/test_litellm/rag/test_main.py +++ b/tests/test_litellm/rag/test_main.py @@ -11,9 +11,13 @@ aquery carries the completion response with real usage and cost. """ import asyncio -from unittest.mock import AsyncMock, patch +import json +from typing import Final +from unittest.mock import patch +import httpx import pytest +import respx import litellm from litellm._internal_context import is_internal_call @@ -255,113 +259,83 @@ async def test_aquery_streaming_bills_sub_call_costs_into_final_event(): @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 +@pytest.mark.parametrize( + ("retrieval_config_json", "top_level_filter_json", "expected_filter_json"), + ( + ( + '{"vector_store_id":"vs_test_123","custom_llm_provider":"openai","top_k":50,' + '"retrieval_filter":{"equals":{"key":"tenant","value":"retrieval"}}}', + None, + '{"equals":{"key":"tenant","value":"retrieval"}}', + ), + ( + '{"vector_store_id":"vs_test_123","custom_llm_provider":"openai","top_k":50,' + '"filters":{"equals":{"key":"tenant","value":"alias"}}}', + None, + '{"equals":{"key":"tenant","value":"alias"}}', + ), + ( + '{"vector_store_id":"vs_test_123","custom_llm_provider":"openai","top_k":50}', + '{"equals":{"key":"tenant","value":"top-level"}}', + '{"equals":{"key":"tenant","value":"top-level"}}', + ), + ( + '{"vector_store_id":"vs_test_123","custom_llm_provider":"openai","top_k":50,' + '"retrieval_filter":{"equals":{"key":"tenant","value":"retrieval"}},' + '"filters":{"equals":{"key":"tenant","value":"alias"}}}', + '{"equals":{"key":"tenant","value":"top-level"}}', + '{"equals":{"key":"tenant","value":"retrieval"}}', + ), + ( + '{"vector_store_id":"vs_test_123","custom_llm_provider":"openai","top_k":50}', + None, + None, + ), + ), +) +async def test_aquery_forwards_filters_to_vector_store_search( + retrieval_config_json: str, + top_level_filter_json: str | None, + expected_filter_json: str | None, + monkeypatch, +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + retrieval_config: Final = json.loads(retrieval_config_json) + top_level_filter: Final = json.loads(top_level_filter_json) if top_level_filter_json is not None else None + expected_filter: Final = json.loads(expected_filter_json) if expected_filter_json is not None else None - 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 respx.mock(assert_all_called=True) as respx_mock: + search_route: Final = respx_mock.post("https://example.com/v1/vector_stores/vs_test_123/search").mock( + return_value=httpx.Response( + 200, + content='{"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", + respx_mock.post("https://example.com/v1/chat/completions").mock( + return_value=httpx.Response( + 200, + content=( + '{"id":"chatcmpl-test","object":"chat.completion","created":1,"model":"gpt-4o-mini",' + '"choices":[{"index":0,"message":{"role":"assistant","content":"answer"},"finish_reason":"stop"}],' + '"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}' + ), + ) ) - - 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_top_level_filters_kwarg_does_not_collide(): - """ - An SDK caller may pass a top-level `filters` kwarg (it used to flow to the - search via **kwargs). Now that the pipeline passes `filters` explicitly, the - top-level kwarg must be consumed rather than forwarded twice, otherwise - asearch raises TypeError for a duplicate keyword before any search runs. - """ - from litellm.types.vector_stores import VectorStoreSearchResponse - - top_level_filter = {"equals": {"key": "tenant", "value": "a"}} - - 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": "hello"}], - retrieval_config={"vector_store_id": "vs_test_123", "custom_llm_provider": "openai"}, + response: Final = await litellm.aquery( + model="openai/gpt-4o-mini", + messages=json.loads('[{"role":"user","content":"most frequent causes of low nicotine"}]'), + retrieval_config=retrieval_config, filters=top_level_filter, - mock_response="hi", + api_key="sk-test", + api_base="https://example.com/v1", ) + request_body: Final = json.loads(search_route.calls.last.request.content) assert isinstance(response, ModelResponse) - fake_search.assert_awaited_once() - assert fake_search.await_args.kwargs["filters"] == top_level_filter - assert "filters" not in fake_search.await_args.kwargs.get("kwargs", {}) - - -@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 + assert response.choices[0].message.content == "answer" + assert request_body["query"] == "most frequent causes of low nicotine" + assert request_body["filters"] == expected_filter + assert request_body["max_num_results"] == 50 def test_rag_call_types_are_registered(): From b4f9e319fc2a0e55c775db371ce5d09e806e6c4f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 15 Sep 2026 17:53:57 -0700 Subject: [PATCH 5/5] fix(proxy): surface the upstream status code when a RAG query fails --- litellm/proxy/rag_endpoints/endpoints.py | 7 +++- .../proxy/rag_endpoints/test_rag_endpoints.py | 36 +++++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index d6a402e1860..c09f9c755ed 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -69,6 +69,11 @@ def _response_attr(source: object, name: str) -> object: return getattr(source, name, None) +def _upstream_status_code(error: Exception) -> int: + code: Final = getattr(error, "status_code", None) + return code if isinstance(code, int) else 500 + + def _raise_vector_store_scan_depth_exceeded() -> None: raise HTTPException( status_code=400, @@ -814,6 +819,6 @@ async def rag_query( except Exception as e: verbose_proxy_logger.exception("RAG Query failed: %s", e) raise HTTPException( - status_code=500, + status_code=_upstream_status_code(e), detail={"error": str(e)}, ) diff --git a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py index 832435711c6..1cceaf95b09 100644 --- a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py @@ -11,6 +11,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi.testclient import TestClient +import litellm from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.proxy_server import app @@ -282,6 +283,41 @@ def test_rag_query_returns_response_cost_header(client_internal_user): assert response.headers.get("x-litellm-response-cost") == "3.45e-06" +@pytest.mark.parametrize( + ("upstream_error", "expected_status"), + [ + (litellm.BadRequestError(message="filter andAll needs two clauses", model="kb", llm_provider="bedrock"), 400), + (litellm.NotFoundError(message="Knowledge Base does not exist", model="kb", llm_provider="bedrock"), 404), + (RuntimeError("pipeline blew up"), 500), + ], +) +def test_rag_query_surfaces_upstream_status_code(client_internal_user, upstream_error, expected_status): + """A vector store rejection must reach the caller with its own status code, never a blanket 500.""" + with ( + patch( # test-quality-ok: the handler calls the module-level litellm.aquery directly; no injection seam + "litellm.proxy.rag_endpoints.endpoints.litellm.aquery", + new=AsyncMock(side_effect=upstream_error), + ), + patch("litellm.vector_store_registry", None), # test-quality-ok: proxy module global, no injection seam + patch("litellm.proxy.proxy_server.prisma_client", None), # test-quality-ok: proxy module global, no injection seam + ): + response = client_internal_user.post( + "/v1/rag/query", + json={ + "model": "bedrock/us.anthropic.claude-sonnet-5", + "messages": [{"role": "user", "content": "How was this document ingested?"}], + "retrieval_config": { + "vector_store_id": "L7INRFMVQT", + "custom_llm_provider": "bedrock", + "retrieval_filter": {"andAll": [{"equals": {"key": "department", "value": "billing"}}]}, + }, + }, + ) + + assert response.status_code == expected_status, response.text + assert str(upstream_error) in response.json()["detail"]["error"] + + def test_rag_query_stream_returns_event_stream(client_internal_user): """ A stream=true /v1/rag/query must return an SSE response. Returning the raw