From abde56391b87a24e69db8889c8f1c7ffd7d77fd4 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 12 Nov 2025 12:38:50 -0800 Subject: [PATCH] [Fix] - Bedrock Knowledge Bases - add support for filtering kb queries (#16543) * test_e2e_bedrock_knowledgebase_retrieval_with_llm_api_call_with_tools_and_filters * fix vs registry * fix merging params * test_bedrock_kb_request_body_has_transformed_filters * fix typing / linting --- .../bedrock/chat/agentcore/transformation.py | 4 +- ...odel_prices_and_context_window_backup.json | 92 +++++++---- litellm/types/vector_stores.py | 22 +++ .../vector_stores/vector_store_registry.py | 149 +++++++++++------- .../test_bedrock_knowledgebase_hook.py | 148 ++++++++++++++++- 5 files changed, 322 insertions(+), 93 deletions(-) diff --git a/litellm/llms/bedrock/chat/agentcore/transformation.py b/litellm/llms/bedrock/chat/agentcore/transformation.py index 730efec2ee..677bd91f98 100644 --- a/litellm/llms/bedrock/chat/agentcore/transformation.py +++ b/litellm/llms/bedrock/chat/agentcore/transformation.py @@ -5,7 +5,7 @@ https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agentcore_InvokeAgen """ import json -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast from urllib.parse import quote import httpx @@ -517,7 +517,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): from litellm.utils import CustomStreamWrapper if client is None or not isinstance(client, AsyncHTTPHandler): - client = get_async_httpx_client(llm_provider="bedrock", params={}) + client = get_async_httpx_client(llm_provider=cast(Any, "bedrock"), params={}) # Make async streaming request response = await client.post( diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 1333f2f182..aa7774c43e 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -18309,6 +18309,21 @@ "supports_reasoning": true, "supports_tool_choice": true }, + "openrouter/deepseek/deepseek-v3.2-exp": { + "input_cost_per_token": 2e-07, + "input_cost_per_token_cache_hit": 2e-08, + "litellm_provider": "openrouter", + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 4e-07, + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": false, + "supports_tool_choice": true + }, "openrouter/deepseek/deepseek-coder": { "input_cost_per_token": 1.4e-07, "litellm_provider": "openrouter", @@ -18552,6 +18567,19 @@ "output_cost_per_token": 1e-06, "supports_tool_choice": true }, + "openrouter/minimax/minimax-m2": { + "input_cost_per_token": 2.55e-7, + "litellm_provider": "openrouter", + "max_input_tokens": 204800, + "max_output_tokens": 204800, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 1.02e-6, + "supports_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": true, + "supports_tool_choice": true + }, "openrouter/mistralai/mistral-7b-instruct": { "input_cost_per_token": 1.3e-07, "litellm_provider": "openrouter", @@ -19023,15 +19051,16 @@ "supports_vision": true }, "openrouter/qwen/qwen3-coder": { - "input_cost_per_token": 1e-06, + "input_cost_per_token": 2.2e-7, "litellm_provider": "openrouter", - "max_input_tokens": 1000000, - "max_output_tokens": 1000000, - "max_tokens": 1000000, + "max_input_tokens": 262100, + "max_output_tokens": 262100, + "max_tokens": 262100, "mode": "chat", - "output_cost_per_token": 5e-06, + "output_cost_per_token": 9.5e-7, "source": "https://openrouter.ai/qwen/qwen3-coder", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "openrouter/switchpoint/router": { "input_cost_per_token": 8.5e-07, @@ -19080,6 +19109,32 @@ "supports_tool_choice": true, "supports_web_search": false }, + "openrouter/z-ai/glm-4.6": { + "input_cost_per_token": 4.0e-7, + "litellm_provider": "openrouter", + "max_input_tokens": 202800, + "max_output_tokens": 131000, + "max_tokens": 202800, + "mode": "chat", + "output_cost_per_token": 1.75e-6, + "source": "https://openrouter.ai/z-ai/glm-4.6", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "openrouter/z-ai/glm-4.6:exacto": { + "input_cost_per_token": 4.5e-7, + "litellm_provider": "openrouter", + "max_input_tokens": 202800, + "max_output_tokens": 131000, + "max_tokens": 202800, + "mode": "chat", + "output_cost_per_token": 1.9e-6, + "source": "https://openrouter.ai/z-ai/glm-4.6:exacto", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, "ovhcloud/DeepSeek-R1-Distill-Llama-70B": { "input_cost_per_token": 6.7e-07, "litellm_provider": "ovhcloud", @@ -24059,7 +24114,6 @@ "supports_parallel_function_calling": false, "supports_vision": false }, - "whisper-1": { "input_cost_per_second": 0.0001, "litellm_provider": "openai", @@ -24069,30 +24123,6 @@ "/v1/audio/transcriptions" ] }, - "vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas": { - "input_cost_per_token": 1.5e-07, - "litellm_provider": "vertex_ai-qwen_models", - "max_input_tokens": 262144, - "max_output_tokens": 262144, - "max_tokens": 262144, - "mode": "chat", - "output_cost_per_token": 1.2e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", - "supports_function_calling": true, - "supports_tool_choice": true - }, - "vertex_ai/qwen/qwen3-next-80b-a3b-thinking-maas": { - "input_cost_per_token": 1.5e-07, - "litellm_provider": "vertex_ai-qwen_models", - "max_input_tokens": 262144, - "max_output_tokens": 262144, - "max_tokens": 262144, - "mode": "chat", - "output_cost_per_token": 1.2e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", - "supports_function_calling": true, - "supports_tool_choice": true - }, "xai/grok-2": { "input_cost_per_token": 2e-06, "litellm_provider": "xai", diff --git a/litellm/types/vector_stores.py b/litellm/types/vector_stores.py index 5456eb90e3..6ae0b4bd2f 100644 --- a/litellm/types/vector_stores.py +++ b/litellm/types/vector_stores.py @@ -1,3 +1,4 @@ +from dataclasses import dataclass from datetime import datetime from enum import Enum from typing import Any, Dict, List, Literal, Optional, Tuple, Union @@ -246,3 +247,24 @@ VECTOR_STORE_OPENAI_PARAMS = Literal[ "ranking_options", "rewrite_query", ] + + + +@dataclass +class VectorStoreToolParams: + """Parameters extracted from a file_search tool definition""" + filters: Optional[Dict] = None + max_num_results: Optional[int] = None + ranking_options: Optional[Dict] = None + + def to_dict(self) -> Dict: + """Convert to dict, excluding None values""" + return { + k: v + for k, v in { + "filters": self.filters, + "max_num_results": self.max_num_results, + "ranking_options": self.ranking_options, + }.items() + if v is not None + } diff --git a/litellm/vector_stores/vector_store_registry.py b/litellm/vector_stores/vector_store_registry.py index f7a2ddaec8..9578c8e349 100644 --- a/litellm/vector_stores/vector_store_registry.py +++ b/litellm/vector_stores/vector_store_registry.py @@ -1,15 +1,17 @@ # litellm/proxy/vector_stores/vector_store_registry.py import json from datetime import datetime, timezone -from typing import TYPE_CHECKING, Any, Dict, List, Optional +from typing import TYPE_CHECKING, Any, Dict, List, Optional, get_args from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import remove_items_at_indices from litellm.types.vector_stores import ( + VECTOR_STORE_OPENAI_PARAMS, LiteLLM_ManagedVectorStore, LiteLLM_ManagedVectorStoreIndex, LiteLLM_ManagedVectorStoreListResponse, LiteLLM_VectorStoreConfig, + VectorStoreToolParams, ) if TYPE_CHECKING: @@ -17,7 +19,6 @@ if TYPE_CHECKING: else: PrismaClient = Any - class VectorStoreIndexRegistry: def __init__( self, vector_store_indexes: List[LiteLLM_ManagedVectorStoreIndex] = [] @@ -110,6 +111,20 @@ class VectorStoreRegistry: str, LiteLLM_ManagedVectorStore ] = {} + def _extract_tool_params(self, tool: Dict) -> VectorStoreToolParams: + """ + Extract supported parameters from a tool definition. + + Dynamically extracts all parameters defined in VECTOR_STORE_OPENAI_PARAMS. + """ + # Get the list of supported param names from the Literal type + supported_params = get_args(VECTOR_STORE_OPENAI_PARAMS) + + # Extract only the params that exist in the tool + kwargs = {param: tool.get(param) for param in supported_params if param in tool} + + return VectorStoreToolParams(**kwargs) + def get_vector_store_ids_to_run( self, non_default_params: Dict, tools: Optional[List[Dict]] = None ) -> List[str]: @@ -130,62 +145,53 @@ class VectorStoreRegistry: return vector_store_ids - def pop_vector_store_ids_to_run( - self, non_default_params: Dict, tools: Optional[List[Dict]] = None - ) -> List[str]: - """ - Pops the vector store ids from the non_default_params and tools - """ - vector_store_ids: List[str] = [] - - # 1. check if vector_store_ids is provided in the non_default_params - vector_store_ids = non_default_params.pop("vector_store_ids", None) or [] - - # 2. check if vector_store_ids is provided as a tool in the request - vector_store_ids = self.get_and_pop_recognised_vector_store_tools( - tools=tools, - vector_store_ids=vector_store_ids, - ) - - return vector_store_ids - def get_and_pop_recognised_vector_store_tools( - self, tools: Optional[List[Dict]] = None, vector_store_ids: List[str] = [] - ) -> List[str]: + self, tools: Optional[List[Dict]] = None, vector_store_ids: Optional[List[str]] = None + ) -> Dict[str, VectorStoreToolParams]: """ - Returns and pops the vector store ids from the tool calls - - It only pops the recognised vector store tools from the tools list. - + Returns and pops recognized vector store tools from the tools list. + Args: - tools: The tools to pop the vector store ids from - vector_store_ids: The list of vector store IDs the user provided - + tools: The tools to extract and remove vector store IDs from + vector_store_ids: Mutable list to append found vector_store_ids to + Returns: - The vector store ids that were popped + Dict mapping vector_store_id to its extracted tool parameters """ - if tools: - tools_to_remove: List[int] = [] - for i, tool in enumerate(tools): - tool_vector_store_ids: List[str] = tool.get("vector_store_ids", []) - if len(tool_vector_store_ids) == 0: - continue - # remove the tool if all vector_store_ids are recognised in the registry - recognised = all( - any(vs.get("vector_store_id") == vs_id for vs in self.vector_stores) - for vs_id in tool_vector_store_ids - ) - if recognised: - tools_to_remove.append(i) - vector_store_ids.extend(tool_vector_store_ids) - - # remove recognised tools from the original list - remove_items_at_indices( - items=tools, - indices=tools_to_remove, + params_by_id: Dict[str, VectorStoreToolParams] = {} + + if not tools: + return params_by_id + + if vector_store_ids is None: + vector_store_ids = [] + + tools_to_remove: List[int] = [] + + for i, tool in enumerate(tools): + tool_vector_store_ids = tool.get("vector_store_ids", []) + if not tool_vector_store_ids: + continue + + # Check if all vector_store_ids are recognized in the registry + recognised = all( + any(vs.get("vector_store_id") == vs_id for vs in self.vector_stores) + for vs_id in tool_vector_store_ids ) - - return vector_store_ids + + if recognised: + tools_to_remove.append(i) + vector_store_ids.extend(tool_vector_store_ids) + + # Extract and store params for each vector store + tool_params = self._extract_tool_params(tool) + for vs_id in tool_vector_store_ids: + params_by_id[vs_id] = tool_params + + # Remove recognized tools from the original list + remove_items_at_indices(items=tools, indices=tools_to_remove) + + return params_by_id def get_vector_store_to_run( self, non_default_params: Dict, tools: Optional[List[Dict]] = None @@ -240,18 +246,45 @@ class VectorStoreRegistry: self, non_default_params: Dict, tools: Optional[List[Dict]] = None ) -> List[LiteLLM_ManagedVectorStore]: """ - Pops the vector stores to run - - Primary function to use for vector store pre call hook + Pops the vector stores to run with their tool parameters merged. + + Primary function to use for vector store pre call hook. + + Args: + non_default_params: Parameters dict to pop vector_store_ids from + tools: Optional list of tools to extract vector store params from + + Returns: + List of vector stores with tool parameters merged into litellm_params """ - vector_store_ids = self.pop_vector_store_ids_to_run( - non_default_params=non_default_params, tools=tools + # Pop vector_store_ids from params + vector_store_ids: List[str] = non_default_params.pop("vector_store_ids", None) or [] + + # Extract params from tools and collect IDs + params_by_id = self.get_and_pop_recognised_vector_store_tools( + tools=tools, + vector_store_ids=vector_store_ids ) + vector_stores_to_run: List[LiteLLM_ManagedVectorStore] = [] + for vector_store_id in vector_store_ids: for vector_store in self.vector_stores: if vector_store.get("vector_store_id") == vector_store_id: - vector_stores_to_run.append(vector_store) + # Create a copy to avoid modifying the registry + vector_store_copy = vector_store.copy() + + # Merge tool params if they exist + if vector_store_id in params_by_id: + existing_params = vector_store_copy.get("litellm_params", {}) or {} + tool_params_dict = params_by_id[vector_store_id].to_dict() + # Tool params take precedence over existing params + tool_params_dict.update(existing_params) + vector_store_copy["litellm_params"] = tool_params_dict + + vector_stores_to_run.append(vector_store_copy) + break + return vector_stores_to_run def _get_vector_store_ids_from_tool_calls( diff --git a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py index e8c696e2b9..f87351edb0 100644 --- a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py +++ b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py @@ -7,6 +7,7 @@ sys.path.insert(0, os.path.abspath("../..")) import asyncio import litellm +import litellm.vector_stores.main import gzip import json import logging @@ -23,15 +24,24 @@ from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook i from litellm.llms.custom_httpx.http_handler import HTTPHandler, AsyncHTTPHandler from litellm.integrations.custom_logger import CustomLogger from litellm.types.utils import StandardLoggingPayload, StandardLoggingVectorStoreRequest -from litellm.types.vector_stores import VectorStoreSearchResponse +from litellm.types.vector_stores import ( + VectorStoreSearchResponse, + VectorStoreResultContent, + VectorStoreSearchResult, +) class MockCustomLogger(CustomLogger): def __init__(self): self.standard_logging_payload: Optional[StandardLoggingPayload] = None + self.completion_logging_payload: Optional[StandardLoggingPayload] = None super().__init__() async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - self.standard_logging_payload = kwargs.get("standard_logging_object") + payload = kwargs.get("standard_logging_object") + # Store the payload - completion calls have call_type='acompletion' + if payload and payload.get("call_type") == "acompletion": + self.completion_logging_payload = payload + self.standard_logging_payload = payload pass @pytest.fixture(autouse=True) @@ -234,6 +244,140 @@ async def test_e2e_bedrock_knowledgebase_retrieval_with_llm_api_call_with_tools( ) assert response is not None +@pytest.mark.asyncio +async def test_e2e_bedrock_knowledgebase_retrieval_with_llm_api_call_with_tools_and_filters(setup_vector_store_registry): + """ + Test that filters from file_search tools are properly passed through to vector store search. + This test verifies the entire flow: tool parsing -> filter extraction -> vector store API call. + + In this case we filter for a non-existent user_id, which should return no results. + """ + litellm._turn_on_debug() + + response = await litellm.acompletion( + model="anthropic/claude-3-5-haiku-latest", + messages=[{"role": "user", "content": "what is litellm?"}], + max_tokens=10, + tools=[ + { + "type": "file_search", + "vector_store_ids": ["T37J8R4WTM"], + "filters": { + "key": "user_id", + "value": "fake-user-id", + "operator": "eq" + } + } + ], + ) + + # Verify response is not None + assert response is not None + + # Verify search results were added to the response (this proves the search was called) + assert hasattr(response.choices[0].message, "provider_specific_fields") + provider_fields = response.choices[0].message.provider_specific_fields + assert provider_fields is not None + assert "search_results" in provider_fields, "search_results not in provider_specific_fields" + + search_results = provider_fields["search_results"] + assert search_results is not None and len(search_results) > 0, "No search results found" + + # The search was performed - this confirms filters were passed through + # The logs above show: litellm.asearch(... filters={'key': 'user_id', 'value': 'fake-user-id', 'operator': 'eq'}) + # And the Bedrock API request contains: {'filter': {'equals': {'key': 'user_id', 'value': 'fake-user-id'}}} + + print("✅ Filters were successfully passed through to vector store search") + print(f" Search was performed and {len(search_results)} result(s) returned") + + +@pytest.mark.asyncio +async def test_bedrock_kb_request_body_has_transformed_filters(setup_vector_store_registry): + """ + Validate that the Bedrock Knowledge Base request body contains the transformed filters. + """ + captured_request_body: dict = {} + + async def fake_async_vector_store_search_handler( + vector_store_id, + query, + vector_store_search_optional_params, + vector_store_provider_config, + custom_llm_provider, + litellm_params, + logging_obj, + extra_headers=None, + extra_body=None, + timeout=None, + client=None, + _is_async=False, + ): + litellm_params_dict = ( + litellm_params.model_dump(exclude_none=False) + if hasattr(litellm_params, "model_dump") + else dict(litellm_params) + ) + api_base = vector_store_provider_config.get_complete_url( + api_base=litellm_params_dict.get("api_base"), + litellm_params=litellm_params_dict, + ) + + url, request_body = vector_store_provider_config.transform_search_vector_store_request( + vector_store_id=vector_store_id, + query=query, + vector_store_search_optional_params=vector_store_search_optional_params, + api_base=api_base, + litellm_logging_obj=logging_obj, + litellm_params=litellm_params_dict, + ) + captured_request_body["url"] = url + captured_request_body["body"] = request_body + + return VectorStoreSearchResponse( + object="vector_store.search_results.page", + search_query=query if isinstance(query, str) else " ".join(query), + data=[ + VectorStoreSearchResult( + score=0.9, + content=[VectorStoreResultContent(text="LiteLLM is a library", type="text")], + ) + ], + ) + + with patch.object( + litellm.vector_stores.main.base_llm_http_handler, + "async_vector_store_search_handler", + new=AsyncMock(side_effect=fake_async_vector_store_search_handler), + ): + response = await litellm.acompletion( + model="anthropic/claude-3-5-haiku-latest", + messages=[{"role": "user", "content": "what is litellm?"}], + max_tokens=10, + tools=[ + { + "type": "file_search", + "vector_store_ids": ["T37J8R4WTM"], + "filters": { + "key": "user_id", + "value": "fake-user-id", + "operator": "eq", + }, + } + ], + ) + + assert response is not None + print("captured_request_body:", json.dumps(captured_request_body, indent=4, default=str)) + assert "body" in captured_request_body, "Bedrock KB request body was not captured" + + vector_search = captured_request_body["body"]["retrievalConfiguration"]["vectorSearchConfiguration"] + aws_filter = vector_search["filter"] + assert "equals" in aws_filter, f"Expected 'equals' in AWS format, got: {aws_filter}" + assert aws_filter["equals"]["key"] == "user_id" + assert aws_filter["equals"]["value"] == "fake-user-id" + + print("✅ Filters transformed correctly: OpenAI format -> AWS Bedrock format") + @pytest.mark.asyncio async def test_openai_with_knowledge_base_mock_openai(setup_vector_store_registry): """