scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

fix(bedrock): map OpenAI vector store search filters to AWS RetrievalFilter

Browse source 
Bedrock Knowledge Base search forwarded OpenAI's type-keyed filters
({"type": "eq", ...}, {"type": "and", "filters": [...]}) to AWS unchanged,
and AWS answered 400 for every one. Comparison and compound filters now
map to AWS's equals/notEquals/greaterThan/.../andAll/orAll keys, a
one-member group is unwrapped, and AWS-native filters pass through.

A malformed OpenAI filter answers 400 before any AWS call, or is dropped
under drop_params, which vector store search now reads from the request
or the global setting instead of hard-coding it off.
This commit is contained in:
mateo-berri 2026-10-30 04:41:00 -07:00
parent 4c8b90945f
commit ea0153a416
4 changed files with 309 additions and 49 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

114
litellm/llms/bedrock/vector_stores/transformation.py
View file

@ -1,11 +1,14 @@
from collections.abc import Mapping
from copy import deepcopy
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, cast
from urllib.parse import urlparse
import httpx
from pydantic import TypeAdapter, ValidationError
from litellm._logging import verbose_logger
from litellm.exceptions import BadRequestError
from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base
from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
@ -21,6 +24,9 @@ from litellm.types.router import GenericLiteLLMParams
from litellm.types.vector_stores import (
VECTOR_STORE_OPENAI_PARAMS,
BaseVectorStoreAuthCredentials,
OpenAIComparisonFilter,
OpenAICompoundFilter,
OpenAIVectorStoreFilter,
VectorStoreIndexEndpoints,
VectorStoreResultContent,
VectorStoreSearchOptionalRequestParams,
@ -33,6 +39,39 @@ if TYPE_CHECKING:
else:
LiteLLMLoggingObj = Any
_AWS_COMPARISON_OPERATORS: Final = MappingProxyType(
{
"eq": "equals",
"ne": "notEquals",
"gt": "greaterThan",
"gte": "greaterThanOrEquals",
"lt": "lessThan",
"lte": "lessThanOrEquals",
"in": "in",
"nin": "notIn",
}
)
_AWS_LOGICAL_OPERATORS: Final = MappingProxyType({"and": "andAll", "or": "orAll"})
_OPENAI_FILTER: Final[TypeAdapter[OpenAIComparisonFilter | OpenAICompoundFilter]] = TypeAdapter(OpenAIVectorStoreFilter)
def _aws_retrieval_filter(openai_filter: OpenAIComparisonFilter | OpenAICompoundFilter) -> Mapping[str, object]:
if isinstance(openai_filter, OpenAIComparisonFilter):
value: Final = list(openai_filter.value) if isinstance(openai_filter.value, tuple) else openai_filter.value
return {_AWS_COMPARISON_OPERATORS[openai_filter.type]: {"key": openai_filter.key, "value": value}}
if len(openai_filter.filters) == 1:
return _aws_retrieval_filter(openai_filter.filters[0])
return {
_AWS_LOGICAL_OPERATORS[openai_filter.type]: [_aws_retrieval_filter(member) for member in openai_filter.filters]
}
def _validation_reason(error: ValidationError) -> str:
return "; ".join(
".".join(("filters", *(str(part) for part in detail["loc"]))) + f": {detail['msg']}"
for detail in error.errors()
)
class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
"""Vector store configuration for AWS Bedrock Knowledge Bases."""
@ -62,34 +101,7 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
return ["filters", "max_num_results", "ranking_options"]
def _map_operator_to_aws(self, operator: str) -> str:
"""
Map OpenAI-style operators to AWS Bedrock operator names.
OpenAI uses: eq, ne, gt, gte, lt, lte, in, nin
AWS uses: equals, notEquals, greaterThan, greaterThanOrEquals, lessThan, lessThanOrEquals, in, notIn, startsWith, listContains, stringContains
"""
operator_mapping: Final = {
"eq": "equals",
"ne": "notEquals",
"gt": "greaterThan",
"gte": "greaterThanOrEquals",
"lt": "lessThan",
"lte": "lessThanOrEquals",
"in": "in",
"nin": "notIn",
# AWS-specific operators (pass through)
"equals": "equals",
"notEquals": "notEquals",
"greaterThan": "greaterThan",
"greaterThanOrEquals": "greaterThanOrEquals",
"lessThan": "lessThan",
"lessThanOrEquals": "lessThanOrEquals",
"notIn": "notIn",
"startsWith": "startsWith",
"listContains": "listContains",
"stringContains": "stringContains",
}
return operator_mapping.get(operator, operator)
return _AWS_COMPARISON_OPERATORS.get(operator, operator)
def _map_operator_filter(self, filter_dict: dict) -> dict:
"""
@ -154,7 +166,7 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
def map_openai_params(
self,
non_default_params: dict,
non_default_params: Mapping[str, object],
optional_params: dict,
drop_params: bool,
) -> dict:
@ -162,27 +174,35 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
if param == "max_num_results":
optional_params["numberOfResults"] = value
elif param == "filters" and value is not None:
# map the openai filters to the aws kb filters format
# openai filters = {"key": <key>, "value": <value>, "operator": <operator>} OR {"and" | "or": [{"key": <key>, "value": <value>, "operator": <operator>}]}
# aws kb filters = {"operator": {"<key>": <value>}} OR {"andAll | orAll": [{"operator": {"<key>": <value>}}]}
# 1. check if filter is in openai format
# 2. if it is, map it to the aws kb filters format
# 3. if it is not, assume it is in aws kb filters format and add it to the optional_params
aws_filters: dict | None = None
if isinstance(value, dict):
if "operator" in value:
# Single operator - map directly (no wrapping needed)
aws_filters = self._map_operator_filter(value)
elif "and" in value or "or" in value:
aws_filters = self._map_and_or_filters(value)
else:
# Assume it's already in AWS KB format
aws_filters = value
optional_params["filters"] = aws_filters
optional_params["filters"] = self._aws_filters(value, drop_params=drop_params)
return optional_params
def _aws_filters(self, filters: object, drop_params: bool) -> Mapping[str, object] | None:
if not isinstance(filters, dict):
return self._invalid_filters(f"filters must be an object, got {type(filters).__name__}", drop_params)
if "type" in filters:
try:
return _aws_retrieval_filter(_OPENAI_FILTER.validate_python(filters))
except ValidationError as error:
return self._invalid_filters(_validation_reason(error), drop_params)
if "operator" in filters:
return self._map_operator_filter(filters)
if "and" in filters or "or" in filters:
return self._map_and_or_filters(filters)
return filters
@staticmethod
def _invalid_filters(reason: str, drop_params: bool) -> None:
if not drop_params:
raise BadRequestError(
message=f"Invalid filters for the Bedrock vector store: {reason}. "
"Fix the filter, or set drop_params to drop it.",
model=None,
llm_provider="bedrock",
)
verbose_logger.debug("Dropping invalid Bedrock vector store filters under drop_params: %s", reason)
def validate_environment(self, headers: dict, litellm_params: GenericLiteLLMParams | None) -> dict:
headers = headers or {}
headers.setdefault("Content-Type", "application/json")

24
litellm/types/vector_stores.py
View file

@ -2,8 +2,9 @@ from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime
from enum import Enum
from typing import Any, Literal
from typing import Annotated, Any, Literal, TypeAlias
from pydantic import ConfigDict, Field
from typing_extensions import ReadOnly, TypedDict
from litellm.types.llms.base import LiteLLMBaseModel
@ -323,6 +324,27 @@ VECTOR_STORE_OPENAI_PARAMS = Literal[
]
class OpenAIComparisonFilter(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True)
type: Literal["eq", "ne", "gt", "gte", "lt", "lte", "in", "nin"]
key: str
value: str | bool | int | float | tuple[str | int | float, ...]
class OpenAICompoundFilter(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True)
type: Literal["and", "or"]
filters: tuple["OpenAIVectorStoreFilter", ...] = Field(min_length=1)
OpenAIVectorStoreFilter: TypeAlias = Annotated[
OpenAIComparisonFilter | OpenAICompoundFilter,
Field(discriminator="type"),
]
@dataclass
class VectorStoreToolParams:
"""Parameters extracted from a file_search tool definition"""

3
litellm/vector_stores/utils.py
View file

@ -5,6 +5,7 @@ from typing import Final, cast, get_type_hints
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
from litellm.constants import ROUTING_REQUEST_TAGS_METADATA_KEY, VECTOR_STORE_CALLER_METADATA_KWARG
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params
from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig
from litellm.types.vector_stores import (
VectorStoreCreateOptionalRequestParams,
@ -111,7 +112,7 @@ class VectorStoreRequestUtils:
optional_params: Final = vector_store_provider_config.map_openai_params(
non_default_params=params,
optional_params=filtered_params,
drop_params=False,
drop_params=drops_unsupported_params(params.get("drop_params")),
)
return cast(VectorStoreSearchOptionalRequestParams, optional_params)

217
tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py
View file

@ -375,3 +375,220 @@ async def test_bedrock_search_with_credentials_managed_registry(
assert search_response["data"][0]["score"] == 0.9
assert search_response["data"][0]["content"][0]["text"] == "Mock search result"
async def _search_with_filters(respx_mock: respx.MockRouter, filters: object, **kwargs: object) -> respx.Route:
import litellm
retrieve_route: Final = respx_mock.post(BEDROCK_RETRIEVE_URL).mock(
return_value=httpx.Response(200, json=BEDROCK_RETRIEVE_RESPONSE)
)
await litellm.vector_stores.asearch(
vector_store_id="T37J8R4WTM",
query="vacation policy",
filters=filters,
custom_llm_provider="bedrock",
aws_access_key_id="test_access_key",
aws_secret_access_key="test_secret_key",
aws_region_name="us-east-1",
**kwargs,
)
return retrieve_route
def _sent_filter(retrieve_route: respx.Route) -> object:
body: Final = json.loads(retrieve_route.calls.last.request.read())
return body.get("retrievalConfiguration", {}).get("vectorSearchConfiguration", {}).get("filter")
@pytest.mark.parametrize(
("openai_filter", "aws_filter"),
[
pytest.param(
{"type": "eq", "key": "department", "value": "hr"},
{"equals": {"key": "department", "value": "hr"}},
id="eq",
),
pytest.param(
{"type": "ne", "key": "archived", "value": True},
{"notEquals": {"key": "archived", "value": True}},
id="ne-bool",
),
pytest.param(
{"type": "gt", "key": "year", "value": 2022},
{"greaterThan": {"key": "year", "value": 2022}},
id="gt-int",
),
pytest.param(
{"type": "gte", "key": "score", "value": 0.5},
{"greaterThanOrEquals": {"key": "score", "value": 0.5}},
id="gte-float",
),
pytest.param(
{"type": "lt", "key": "year", "value": 2024},
{"lessThan": {"key": "year", "value": 2024}},
id="lt",
),
pytest.param(
{"type": "lte", "key": "year", "value": 2024},
{"lessThanOrEquals": {"key": "year", "value": 2024}},
id="lte",
),
pytest.param(
{"type": "in", "key": "department", "value": ["hr", "sales"]},
{"in": {"key": "department", "value": ["hr", "sales"]}},
id="in",
),
pytest.param(
{"type": "nin", "key": "year", "value": [2022, 2023]},
{"notIn": {"key": "year", "value": [2022, 2023]}},
id="nin",
),
pytest.param(
{
"type": "and",
"filters": [
{"type": "eq", "key": "department", "value": "hr"},
{"type": "gte", "key": "year", "value": 2024},
],
},
{
"andAll": [
{"equals": {"key": "department", "value": "hr"}},
{"greaterThanOrEquals": {"key": "year", "value": 2024}},
]
},
id="and",
),
pytest.param(
{
"type": "or",
"filters": [
{"type": "eq", "key": "department", "value": "sales"},
{
"type": "and",
"filters": [
{"type": "eq", "key": "department", "value": "hr"},
{"type": "lt", "key": "year", "value": 2024},
],
},
],
},
{
"orAll": [
{"equals": {"key": "department", "value": "sales"}},
{
"andAll": [
{"equals": {"key": "department", "value": "hr"}},
{"lessThan": {"key": "year", "value": 2024}},
]
},
]
},
id="or-with-nested-and",
),
pytest.param(
{"type": "or", "filters": [{"type": "eq", "key": "department", "value": "hr"}]},
{"equals": {"key": "department", "value": "hr"}},
id="one-member-compound-unwrapped",
),
pytest.param(
{"andAll": [{"equals": {"key": "a", "value": 1}}, {"startsWith": {"key": "b", "value": "x"}}]},
{"andAll": [{"equals": {"key": "a", "value": 1}}, {"startsWith": {"key": "b", "value": "x"}}]},
id="aws-native-unchanged",
),
pytest.param(
{"key": "department", "value": "hr", "operator": "ne"},
{"notEquals": {"key": "department", "value": "hr"}},
id="operator-shape-still-mapped",
),
pytest.param(
{
"and": [
{"key": "department", "value": "hr", "operator": "eq"},
{"key": "year", "value": 2024, "operator": "gte"},
]
},
{
"andAll": [
{"equals": {"key": "department", "value": "hr"}},
{"greaterThanOrEquals": {"key": "year", "value": 2024}},
]
},
id="and-list-shape-still-mapped",
),
],
)
@pytest.mark.asyncio
async def test_search_sends_openai_filters_as_aws_retrieval_filters(
respx_mock: respx.MockRouter,
monkeypatch: pytest.MonkeyPatch,
openai_filter: dict[str, object],
aws_filter: dict[str, object],
):
import litellm
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
retrieve_route: Final = await _search_with_filters(respx_mock, openai_filter)
assert _sent_filter(retrieve_route) == aws_filter
INVALID_FILTERS: Final = [
pytest.param({"type": "regex", "key": "department", "value": "h.*"}, id="unknown-type"),
pytest.param({"type": "eq", "key": "department"}, id="missing-value"),
pytest.param({"type": "and", "filters": []}, id="empty-compound"),
pytest.param(
{"type": "and", "filters": [{"type": "eq", "key": "department", "value": "hr"}, {"key": "year"}]},
id="untyped-member",
),
pytest.param([{"type": "eq", "key": "department", "value": "hr"}], id="not-an-object"),
]
@pytest.mark.parametrize("filters", INVALID_FILTERS)
@pytest.mark.asyncio
async def test_search_rejects_invalid_filters_with_400_without_calling_aws(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, filters: object
):
import litellm
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setattr(litellm, "drop_params", False)
retrieve_route: Final = respx_mock.post(BEDROCK_RETRIEVE_URL).mock(
return_value=httpx.Response(200, json=BEDROCK_RETRIEVE_RESPONSE)
)
with pytest.raises(litellm.BadRequestError) as error:
await litellm.vector_stores.asearch(
vector_store_id="T37J8R4WTM",
query="vacation policy",
filters=filters,
custom_llm_provider="bedrock",
aws_access_key_id="test_access_key",
aws_secret_access_key="test_secret_key",
aws_region_name="us-east-1",
)
assert error.value.status_code == 400
assert "filters" in error.value.message
assert retrieve_route.call_count == 0
@pytest.mark.parametrize("filters", INVALID_FILTERS)
@pytest.mark.parametrize("drop_params_source", ["request", "global"])
@pytest.mark.asyncio
async def test_search_drops_invalid_filters_under_drop_params(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, filters: object, drop_params_source: str
):
import litellm
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setattr(litellm, "drop_params", drop_params_source == "global")
request_kwargs: Final = {"drop_params": True} if drop_params_source == "request" else {}
retrieve_route: Final = await _search_with_filters(respx_mock, filters, **request_kwargs)
assert retrieve_route.call_count == 1
assert _sent_filter(retrieve_route) is None