mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-15 09:12:18 +00:00
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:
parent
ec48e2491c
commit
da67a9c8c3
5 changed files with 620 additions and 51 deletions
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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"""
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue