Merge d2c574608a into 3746ba58d7

litellm/rag/main.py

@@ -223,10 +223,13 @@ async def _execute_query_pipeline(
         raise ValueError("No query found in messages for RAG query")
 
     # 2. Search vector store
+    kwargs_filters: Final = kwargs.pop("filters", None)
+    filters: Final = retrieval_config.get("retrieval_filter") or retrieval_config.get("filters") or kwargs_filters
     with _suppressed_sub_call_billing():
         search_response: Final = 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,

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,10 +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
+    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):

tests/test_litellm/rag/test_main.py

@@ -11,9 +11,13 @@ aquery carries the completion response with real usage and cost.
 """
 
 import asyncio
+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
@@ -254,6 +258,86 @@ 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(
+    ("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
+
+    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":[]}',
+            )
+        )
+        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}}'
+                ),
+            )
+        )
+        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,
+            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)
+    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():
     """
     query/aquery/ingest/aingest are @client-decorated entry points, so their