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

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

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.

* fix(bedrock): reject booleans in OpenAI in/nin filter lists and map older filter shapes first

A list value now holds only strings and numbers, as in OpenAI's schema, so pydantic no longer turns true/false into 1/0 before the filter reaches AWS. The older operator and and/or shapes are checked before the type key, so an older filter that also carries a type key maps as it did before.

* fix(bedrock): pass AWS-native vector store filters through when they also carry a type key

* test(bedrock): check every AWS RetrievalFilter key passes through with a type key

* fix(bedrock): stop validating a vector store filter at its first bad item

A malformed OpenAI filter listed every pydantic error in its 400, so one 30 KB filter built a 1.4 MB message that stalled the worker while it was redacted. Validation now stops at the first bad list item or group member, the reason is cut at 2000 characters, and the 400 no longer chains the pydantic error into the logged traceback.

* fix(bedrock): report a bad vector store filter member without calling its group empty

An and/or group's non-empty check ran as min_length on the partly validated tuple after fail_fast stopped at a bad member, so every enclosing group also reported that it had no members. It now runs as an after-validator, which only sees a group whose members all validated, so one bad member gives one error at any depth.

A reason over 2000 characters now keeps its first and last 1000 characters, so a deeply nested path still ends on the bad field and its message.

* fix(bedrock): keep the vector store filter types importable on Python 3.10

Python 3.10's typing.Union hashes every member, and the FailFast() marker in the filter value's list member is an unhashable dataclass, so importing litellm raised TypeError: unhashable type: 'FailFast'. Field(fail_fast=True) sets the same fail-fast flag on the tuple schema and is hashable, so the value and the and/or group use it instead.

* fix(bedrock): check OpenAI vector store filter values per operator and send numeric in/nin as AWS accepts them

AWS in and notIn take only a non-empty list of strings, so an in or nin list holding a number now goes out as orAll of equals or andAll of notEquals, and a one-item list as a single equals or notEquals

Each OpenAI filter now refuses keys it does not define and checks its value by operator: eq and ne take a scalar, gt, gte, lt and lte take a number, in and nin take a non-empty list. An empty filters object is invalid too. Each of these answers 400 before any AWS call, or is dropped under drop_params, where before it reached AWS and failed there or, for a stray key, was silently ignored

_aws_retrieval_filter joins the recursion detector's ignore list, since it walks only a validated filter that pydantic-core stops at 255 levels

* fix(bedrock): join same-operator filter groups so OpenAI filters stay within the AWS nesting limit

AWS rejects a RetrievalFilter nested more than 2 levels deep, and the orAll or andAll a numeric in or nin turns into used to add one, so and(eq, or(in [2023, 2024], eq)) failed at AWS. A group whose member is the same kind of group, written by the caller or made from a numeric in or nin, now takes that member's items, so it goes out as andAll(equals, orAll(equals, equals, equals))

The mapping is unchanged otherwise, and a mix of and and or nested deeper than AWS allows still fails at AWS

_aws_group_members takes _aws_retrieval_filter's place in the recursion detector's ignore list, with the same bound

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
zoveki 2026-11-14 10:39:43 -08:00 • committed by GitHub
parent ec48e2491c
commit da67a9c8c3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 620 additions and 51 deletions

View file

@ -1,13 +1,13 @@
import re
from collections.abc import Mapping
from copy import deepcopy
from itertools import product
from itertools import chain, product
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Protocol, cast, runtime_checkable
from urllib.parse import urlparse
import httpx
from pydantic import TypeAdapter
from pydantic import TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_logger
@ -30,6 +30,12 @@ from litellm.types.vector_stores import (
RETRIEVAL_QUERY_PARAMS,
VECTOR_STORE_OPENAI_PARAMS,
BaseVectorStoreAuthCredentials,
OpenAIComparisonFilter,
OpenAICompoundFilter,
OpenAIEqualityFilter,
OpenAIMembershipFilter,
OpenAIRangeFilter,
OpenAIVectorStoreFilter,
VectorStoreIndexEndpoints,
VectorStoreResultContent,
VectorStoreSearchImageMediaType,
@ -44,6 +50,88 @@ 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"})
_NUMERIC_MEMBERSHIP: Final = MappingProxyType({"in": ("or", "eq"), "nin": ("and", "ne")})
_AWS_FILTER_KEYS: Final = frozenset(
(
*_AWS_COMPARISON_OPERATORS.values(),
*_AWS_LOGICAL_OPERATORS.values(),
"startsWith",
"listContains",
"stringContains",
)
)
_OPENAI_FILTER: Final[TypeAdapter[OpenAIComparisonFilter | OpenAICompoundFilter]] = TypeAdapter(OpenAIVectorStoreFilter)
_MAX_REASON_CHARS: Final = 2000
_REASON_END_CHARS: Final = _MAX_REASON_CHARS // 2
def _aws_retrieval_filter(openai_filter: OpenAIComparisonFilter | OpenAICompoundFilter) -> Mapping[str, object]:
match openai_filter:
case OpenAICompoundFilter():
return _aws_logical_filter(openai_filter.type, openai_filter.filters)
case OpenAIMembershipFilter() if _expands_to_logical_filter(openai_filter):
return _aws_logical_filter(_NUMERIC_MEMBERSHIP[openai_filter.type][0], (openai_filter,))
case OpenAIMembershipFilter():
return _aws_comparison(openai_filter.type, openai_filter.key, list(openai_filter.value))
case OpenAIEqualityFilter() | OpenAIRangeFilter():
return _aws_comparison(openai_filter.type, openai_filter.key, openai_filter.value)
def _aws_logical_filter(
operator: str, members: tuple[OpenAIComparisonFilter | OpenAICompoundFilter, ...]
) -> Mapping[str, object]:
aws_members: Final = tuple(chain.from_iterable(_aws_group_members(operator, member) for member in members))
if len(aws_members) == 1:
return aws_members[0]
return {_AWS_LOGICAL_OPERATORS[operator]: list(aws_members)}
def _aws_group_members(
operator: str, member: OpenAIComparisonFilter | OpenAICompoundFilter
) -> tuple[Mapping[str, object], ...]:
match member:
case OpenAICompoundFilter() if member.type == operator:
return tuple(chain.from_iterable(_aws_group_members(operator, nested) for nested in member.filters))
case OpenAIMembershipFilter() if (
_expands_to_logical_filter(member) and _NUMERIC_MEMBERSHIP[member.type][0] == operator
):
return tuple(
_aws_comparison(_NUMERIC_MEMBERSHIP[member.type][1], member.key, item) for item in member.value
)
case _:
return (_aws_retrieval_filter(member),)
def _expands_to_logical_filter(membership: OpenAIMembershipFilter) -> bool:
return not all(isinstance(item, str) for item in membership.value)
def _aws_comparison(operator: str, key: str, value: str | float | list[str | float]) -> Mapping[str, object]:
return {_AWS_COMPARISON_OPERATORS[operator]: {"key": key, "value": value}}
def _validation_reason(error: ValidationError) -> str:
reason: Final = "; ".join(
".".join(("filters", *(str(part) for part in detail["loc"]))) + f": {detail['msg']}"
for detail in error.errors()
)
return (
reason if len(reason) <= _MAX_REASON_CHARS else f"{reason[:_REASON_END_CHARS]}...{reason[-_REASON_END_CHARS:]}"
)
@runtime_checkable
class _SupportsSetDefault(Protocol):
@ -164,34 +252,7 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
return ["filters", "max_num_results", "ranking_options", "query_image"]
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:
"""
@ -256,7 +317,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:
@ -264,27 +325,37 @@ 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 not filters:
return self._invalid_filters("filters must not be empty", 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)
if "type" in filters and filters.keys().isdisjoint(_AWS_FILTER_KEYS):
try:
return _aws_retrieval_filter(_OPENAI_FILTER.validate_python(filters))
except ValidationError as error:
return self._invalid_filters(_validation_reason(error), drop_params)
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",
) from None
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")

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, Final, Literal, TypeAlias
from typing import Annotated, Any, Final, Literal, TypeAlias
from pydantic import AfterValidator, ConfigDict, Field, StrictFloat, StrictInt, StrictStr
from typing_extensions import ReadOnly, TypedDict
from litellm.types.llms.base import LiteLLMBaseModel
@ -345,6 +346,54 @@ VECTOR_STORE_OPENAI_PARAMS = Literal[
]
def _require_items(items: tuple[object, ...]) -> tuple[object, ...]:
if not items:
raise ValueError("must not be empty")
return items
class OpenAIEqualityFilter(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
type: Literal["eq", "ne"]
key: str
value: str | bool | int | float
class OpenAIRangeFilter(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
type: Literal["gt", "gte", "lt", "lte"]
key: str
value: StrictInt | StrictFloat
class OpenAIMembershipFilter(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
type: Literal["in", "nin"]
key: str
value: Annotated[
tuple[StrictStr | StrictInt | StrictFloat, ...], Field(fail_fast=True), AfterValidator(_require_items)
]
OpenAIComparisonFilter: TypeAlias = OpenAIEqualityFilter | OpenAIRangeFilter | OpenAIMembershipFilter
class OpenAICompoundFilter(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
type: Literal["and", "or"]
filters: Annotated[tuple["OpenAIVectorStoreFilter", ...], Field(fail_fast=True), AfterValidator(_require_items)]
OpenAIVectorStoreFilter: TypeAlias = Annotated[
OpenAIComparisonFilter | OpenAICompoundFilter,
Field(discriminator="type"),
]
@dataclass
class VectorStoreToolParams:
"""Parameters extracted from a file_search tool definition"""

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)

View file

@ -99,6 +99,7 @@ IGNORE_FUNCTIONS = [
"_json", # router_rust index views: max depth set (DEFAULT_MAX_RECURSE_DEPTH); fails closed by raising ValueError at the cap.
"to_wire", # router_rust model info: bounded by the nesting depth of a stored deployment or cost-map entry (finite JSON trees: pydantic refuses a cycle when it dumps a deployment).
"from_wire", # router_rust model info: walks only what router-core merged from to_wire's output, so it never goes deeper than to_wire did.
"_aws_group_members", # bedrock vector store filters: walks only a pydantic-validated filter, which pydantic-core stops at 255 levels of nesting.
"_valid_frontmatter_node", # max depth 32, manifest size capped at 1 MiB, and YAML aliases rejected before traversal.
"_gemini_wire_value", # max depth set (DEFAULT_MAX_RECURSE_DEPTH); returns the subtree as plain JSON at the cap.
]

View file

@ -1,8 +1,10 @@
import json
import traceback
from collections.abc import Mapping
from typing import Final
from unittest.mock import MagicMock
import botocore.session
import httpx
import openai
import pytest
@ -383,6 +385,451 @@ async def test_bedrock_search_with_credentials_managed_registry(
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": "department", "value": ["hr", "sales"]},
{"notIn": {"key": "department", "value": ["hr", "sales"]}},
id="nin",
),
pytest.param(
{"type": "in", "key": "year", "value": [2023, 2024]},
{"orAll": [{"equals": {"key": "year", "value": 2023}}, {"equals": {"key": "year", "value": 2024}}]},
id="in-numbers-as-or-of-equals",
),
pytest.param(
{"type": "nin", "key": "year", "value": [2022, 2023]},
{"andAll": [{"notEquals": {"key": "year", "value": 2022}}, {"notEquals": {"key": "year", "value": 2023}}]},
id="nin-numbers-as-and-of-not-equals",
),
pytest.param(
{"type": "in", "key": "code", "value": ["hr", 7]},
{"orAll": [{"equals": {"key": "code", "value": "hr"}}, {"equals": {"key": "code", "value": 7}}]},
id="in-strings-and-numbers-as-or-of-equals",
),
pytest.param(
{"type": "in", "key": "year", "value": [2023]},
{"equals": {"key": "year", "value": 2023}},
id="in-one-number-unwrapped",
),
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(
{
"type": "and",
"filters": [
{"type": "eq", "key": "department", "value": "hr"},
{
"type": "and",
"filters": [
{"type": "gte", "key": "year", "value": 2023},
{"type": "lt", "key": "year", "value": 2025},
],
},
],
},
{
"andAll": [
{"equals": {"key": "department", "value": "hr"}},
{"greaterThanOrEquals": {"key": "year", "value": 2023}},
{"lessThan": {"key": "year", "value": 2025}},
]
},
id="and-inside-and-joins-the-group",
),
pytest.param(
{
"type": "and",
"filters": [
{"type": "eq", "key": "department", "value": "hr"},
{"type": "nin", "key": "year", "value": [2022, 2023]},
],
},
{
"andAll": [
{"equals": {"key": "department", "value": "hr"}},
{"notEquals": {"key": "year", "value": 2022}},
{"notEquals": {"key": "year", "value": 2023}},
]
},
id="nin-numbers-inside-and-joins-the-group",
),
pytest.param(
{
"type": "and",
"filters": [
{"type": "eq", "key": "department", "value": "hr"},
{
"type": "or",
"filters": [
{"type": "in", "key": "year", "value": [2023, 2024]},
{"type": "eq", "key": "department", "value": "sales"},
],
},
],
},
{
"andAll": [
{"equals": {"key": "department", "value": "hr"}},
{
"orAll": [
{"equals": {"key": "year", "value": 2023}},
{"equals": {"key": "year", "value": 2024}},
{"equals": {"key": "department", "value": "sales"}},
]
},
]
},
id="in-numbers-inside-or-joins-the-group",
),
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(
{"key": "department", "value": "hr", "operator": "eq", "type": "string"},
{"equals": {"key": "department", "value": "hr"}},
id="operator-shape-with-type-key-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.param(
{
"type": "and",
"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-with-type-key-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
AWS_RETRIEVAL_FILTER_KEYS: Final = tuple(
botocore.session.get_session().get_service_model("bedrock-agent-runtime").shape_for("RetrievalFilter").members
)
@pytest.mark.parametrize("aws_key", AWS_RETRIEVAL_FILTER_KEYS)
@pytest.mark.asyncio
async def test_search_sends_aws_native_filters_with_a_type_key_unchanged(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, aws_key: str
):
import litellm
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
aws_filter: Final = {aws_key: {"key": "department", "value": "hr"}, "type": "eq"}
retrieve_route: Final = await _search_with_filters(respx_mock, aws_filter)
assert _sent_filter(retrieve_route) == aws_filter
INVALID_FILTERS: Final = (
("unknown-type", {"type": "regex", "key": "department", "value": "h.*"}, "'regex'"),
("missing-value", {"type": "eq", "key": "department"}, "filters.eq.value"),
("empty-compound", {"type": "and", "filters": []}, "filters.and.filters:"),
(
"untyped-member",
{"type": "and", "filters": [{"type": "eq", "key": "department", "value": "hr"}, {"key": "year"}]},
"filters.and.filters.1",
),
("bools-in-list", {"type": "in", "key": "remote", "value": [True, False]}, "filters.in.value"),
("not-an-object", [{"type": "eq", "key": "department", "value": "hr"}], "filters must be an object, got list"),
("empty-object", {}, "filters must not be empty"),
(
"extra-key-on-comparison",
{"type": "eq", "key": "department", "value": "hr", "filters": [{"type": "eq", "key": "year", "value": 2024}]},
"filters.eq.filters: Extra inputs are not permitted",
),
(
"extra-key-on-compound",
{"type": "or", "filters": [{"type": "eq", "key": "department", "value": "hr"}], "key": "year"},
"filters.or.key: Extra inputs are not permitted",
),
("in-scalar", {"type": "in", "key": "department", "value": "hr"}, "filters.in.value"),
("in-empty-list", {"type": "in", "key": "department", "value": []}, "filters.in.value: Value error, must not be"),
("eq-list", {"type": "eq", "key": "department", "value": ["hr"]}, "filters.eq.value"),
("gt-string", {"type": "gt", "key": "year", "value": "2023"}, "filters.gt.value"),
("gte-bool", {"type": "gte", "key": "year", "value": True}, "filters.gte.value"),
)
@pytest.mark.parametrize(
("filters", "reason"), [pytest.param(filters, reason, id=case) for case, filters, reason in 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, reason: str
):
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 reason in error.value.message
assert retrieve_route.call_count == 0
@pytest.mark.parametrize("filters", [pytest.param(filters, id=case) for case, filters, _ in 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
async def _rejection(respx_mock: respx.MockRouter, filters: object) -> BaseException:
import litellm
with pytest.raises(litellm.BadRequestError) as error:
await _search_with_filters(respx_mock, filters)
return error.value
def _nested_compound(depth: int) -> Mapping[str, object]:
if depth == 0:
return {"type": "eq", "key": "year"}
return {"type": "and", "filters": [_nested_compound(depth - 1)]}
@pytest.mark.parametrize(
("one_bad_item", "many_bad_items"),
[
pytest.param(
{"type": "in", "key": "remote", "value": [True]},
{"type": "in", "key": "remote", "value": [True] * 1000},
id="bools-in-list",
),
pytest.param(
{"type": "and", "filters": [{"key": "year"}]},
{"type": "and", "filters": [{"key": "year"}] * 1000},
id="untyped-members",
),
],
)
@pytest.mark.asyncio
async def test_search_rejection_names_the_first_bad_item_however_many_follow(
respx_mock: respx.MockRouter,
monkeypatch: pytest.MonkeyPatch,
one_bad_item: dict[str, object],
many_bad_items: dict[str, object],
):
import litellm
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setattr(litellm, "drop_params", False)
assert str(await _rejection(respx_mock, many_bad_items)) == str(await _rejection(respx_mock, one_bad_item))
@pytest.mark.asyncio
async def test_search_rejection_and_its_traceback_do_not_grow_with_nesting_depth(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
):
import litellm
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setattr(litellm, "drop_params", False)
shallow: Final = await _rejection(respx_mock, _nested_compound(160))
deep: Final = await _rejection(respx_mock, _nested_compound(240))
assert len(str(deep)) == len(str(shallow))
assert "eq.value" in str(deep)
assert len("".join(traceback.format_exception(deep))) == len("".join(traceback.format_exception(shallow)))
@pytest.mark.asyncio
async def test_search_rejection_names_a_bad_member_the_same_wherever_it_sits(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
):
import litellm
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setattr(litellm, "drop_params", False)
good: Final = {"type": "eq", "key": "department", "value": "hr"}
bad: Final = {"key": "year"}
bad_first: Final = str(await _rejection(respx_mock, {"type": "and", "filters": [bad, good]}))
bad_last: Final = str(await _rejection(respx_mock, {"type": "and", "filters": [good, bad]}))
assert bad_first.replace("filters.and.filters.0", "filters.and.filters.1") == bad_last
assert "..." not in bad_last
async def _search_knowledge_base(
asynchronous: bool, respond: MagicMock, extra_headers: Mapping[str, object]
) -> VectorStoreSearchResponse: