tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-08 22:25:25 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Fix] - Bedrock Knowledge Bases - add support for filtering kb queries (#16543)

Browse Source 
* 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
This commit is contained in:
Ishaan Jaff
2025-11-12 12:38:50 -08:00
committed by GitHub
parent e6bbf3ce81
commit abde56391b
5 changed files with 322 additions and 93 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/llms/bedrock/chat/agentcore/transformation.py
+2 -2
View File
@@ -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(
litellm/model_prices_and_context_window_backup.json
+61 -31
View File
@@ -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",
litellm/types/vector_stores.py
+22
View File
@@ -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
}
litellm/vector_stores/vector_store_registry.py
+91 -58
View File
@@ -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(
tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py
+146 -2
View File
@@ -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):
"""