diff --git a/litellm/llms/bedrock/vector_stores/transformation.py b/litellm/llms/bedrock/vector_stores/transformation.py index 96d1da19bcc..116d8ea6997 100644 --- a/litellm/llms/bedrock/vector_stores/transformation.py +++ b/litellm/llms/bedrock/vector_stores/transformation.py @@ -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": , "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 "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") diff --git a/litellm/types/vector_stores.py b/litellm/types/vector_stores.py index 2259339e479..7298611d0ef 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, 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""" 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/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 87c57023c17..3ce910cdcc6 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 @@ -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