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

This commit is contained in:
mateo-berri 2026-11-06 14:37:49 -08:00
parent 0bc2f59d62
commit 742e7de640
2 changed files with 35 additions and 5 deletions

View file

@ -52,6 +52,15 @@ _AWS_COMPARISON_OPERATORS: Final = MappingProxyType(
}
)
_AWS_LOGICAL_OPERATORS: Final = MappingProxyType({"and": "andAll", "or": "orAll"})
_AWS_FILTER_KEYS: Final = frozenset(
(
*_AWS_COMPARISON_OPERATORS.values(),
*_AWS_LOGICAL_OPERATORS.values(),
"startsWith",
"listContains",
"stringContains",
)
)
_OPENAI_FILTER: Final[TypeAdapter[OpenAIComparisonFilter | OpenAICompoundFilter]] = TypeAdapter(OpenAIVectorStoreFilter)
@ -185,7 +194,7 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
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:
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:

View file

@ -497,6 +497,11 @@ def _sent_filter(retrieve_route: respx.Route) -> object:
{"andAll": [{"equals": {"key": "a", "value": 1}}, {"startsWith": {"key": "b", "value": "x"}}]},
id="aws-native-unchanged",
),
pytest.param(
{"equals": {"key": "department", "value": "hr"}, "type": "eq"},
{"equals": {"key": "department", "value": "hr"}, "type": "eq"},
id="aws-native-with-type-key-unchanged",
),
pytest.param(
{"key": "department", "value": "hr", "operator": "ne"},
{"notEquals": {"key": "department", "value": "hr"}},
@ -522,6 +527,22 @@ def _sent_filter(retrieve_route: respx.Route) -> object:
},
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
@ -541,13 +562,13 @@ async def test_search_sends_openai_filters_as_aws_retrieval_filters(
INVALID_FILTERS: Final = (
("unknown-type", {"type": "regex", "key": "department", "value": "h.*"}, "Input tag 'regex'"),
("missing-value", {"type": "eq", "key": "department"}, "filters.eq.value: Field required"),
("empty-compound", {"type": "and", "filters": []}, "filters.and.filters: Tuple should have at least 1 item"),
("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: Unable to extract tag",
"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"),