diff --git a/litellm/llms/bedrock/vector_stores/transformation.py b/litellm/llms/bedrock/vector_stores/transformation.py index b4cac70fee2..d4f27116144 100644 --- a/litellm/llms/bedrock/vector_stores/transformation.py +++ b/litellm/llms/bedrock/vector_stores/transformation.py @@ -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": , "value": , "operator": } OR {"and" | "or": [{"key": , "value": , "operator": }]} - # aws kb filters = {"operator": {"": }} OR {"andAll | orAll": [{"operator": {"": }}]} - # 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") diff --git a/litellm/types/vector_stores.py b/litellm/types/vector_stores.py index b154d161e13..6d67e0d957f 100644 --- a/litellm/types/vector_stores.py +++ b/litellm/types/vector_stores.py @@ -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""" diff --git a/litellm/vector_stores/utils.py b/litellm/vector_stores/utils.py index eb61491853f..61daabdedac 100644 --- a/litellm/vector_stores/utils.py +++ b/litellm/vector_stores/utils.py @@ -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) diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index fc153ee9447..9c6774e47ad 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -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. ] diff --git a/tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py b/tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py index 2309149b00c..28639153cfd 100644 --- a/tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py +++ b/tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py @@ -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: