diff --git a/python/packages/core/agent_framework/_in_memory.py b/python/packages/core/agent_framework/_in_memory.py index cc8f8046cb7..665c2918455 100644 --- a/python/packages/core/agent_framework/_in_memory.py +++ b/python/packages/core/agent_framework/_in_memory.py @@ -5,11 +5,12 @@ from __future__ import annotations import math -from collections.abc import Mapping, Sequence +import re +from collections.abc import Callable, Mapping, Sequence from copy import deepcopy from dataclasses import dataclass from datetime import date, datetime, time, timedelta -from decimal import Decimal +from decimal import ROUND_HALF_EVEN, Decimal from typing import Any, ClassVar, Generic, cast from uuid import UUID, uuid4 @@ -18,8 +19,10 @@ from ._feature_stage import ExperimentalFeature, experimental from ._telemetry import FeatureIndex, mark_feature_used from ._vector_filters import ( + _TEXT_VALUE_OPERATORS, # pyright: ignore[reportPrivateUsage] Filter, FilterExpression, + FilterGroup, validate_filter, ) from ._vectors import ( @@ -152,6 +155,120 @@ def _calculate_score(left: Vector, right: Vector, distance_function: DistanceFun return score +_DURATION_PATTERN = re.compile( + r"^(?P-)?P(?:(?P[0-9]+(?:\.[0-9]+)?)D)?" + r"(?:T(?:(?P[0-9]+(?:\.[0-9]+)?)H)?(?:(?P[0-9]+(?:\.[0-9]+)?)M)?" + r"(?:(?P[0-9]+(?:\.[0-9]+)?)S)?)?$" +) + + +def _restore_datetime(value: str) -> datetime: + """Parse a timestamp the shared serializer wrote.""" + if value.endswith("Z"): + value = f"{value[:-1]}+00:00" + return datetime.fromisoformat(value) + + +def _restore_timedelta(value: str) -> timedelta: + """Parse the ISO 8601 duration the shared serializer wrote. + + The components are summed as exact decimals of microseconds: the string + the serializer writes carries every microsecond, while a ``float`` of the + total seconds loses them on the largest durations (``timedelta.max`` even + rounds past the range the type supports). + """ + match = _DURATION_PATTERN.match(value) + if match is None: + raise ValueError(f"Unsupported duration '{value}'.") + parts = match.groupdict() + microseconds = ( + Decimal(parts["days"] or "0") * 86_400_000_000 + + Decimal(parts["hours"] or "0") * 3_600_000_000 + + Decimal(parts["minutes"] or "0") * 60_000_000 + + Decimal(parts["seconds"] or "0") * 1_000_000 + ) + delta = timedelta(microseconds=int(microseconds.quantize(Decimal(1), rounding=ROUND_HALF_EVEN))) + return -delta if parts["sign"] else delta + + +def _restore_time(value: str) -> time: + """Parse a time of day the shared serializer wrote.""" + if value.endswith("Z"): + # The serializer writes a UTC time with a trailing "Z", which + # ``time.fromisoformat`` accepts only from Python 3.11. + value = f"{value[:-1]}+00:00" + return time.fromisoformat(value) + + +_RESTORE_IN_MEMORY_VALUE_TYPES: dict[str, Callable[[str], Any]] = { + "datetime": _restore_datetime, + "date": date.fromisoformat, + "time": _restore_time, + "timedelta": _restore_timedelta, + "UUID": UUID, + "Decimal": Decimal, +} +_SERIALIZED_SCALAR_TYPE_NAMES = frozenset(_RESTORE_IN_MEMORY_VALUE_TYPES) +_COMPARISON_OPERATORS = frozenset({"eq", "ne", "gt", "gte", "lt", "lte"}) +_SEQUENCE_OPERAND_OPERATORS = frozenset({"between", "in", "not_in"}) + + +def _restore_in_memory_value(value: Any, type_tag: str | None) -> Any: + """Return a stored value in the type its field declares. + + In-memory records are stored through the shared serializer, so a native + scalar such as a datetime is held as its serialized value. A filter + operand arrives native, so the stored value is restored before the + comparison and the comparison keeps the field's semantics. + """ + if not isinstance(value, str) or type_tag is None: + return value + restore = _RESTORE_IN_MEMORY_VALUE_TYPES.get(type_tag) + if restore is None: + return value + try: + return restore(value) + except (TypeError, ValueError): + return value + + +def _evaluate_in_memory_filter( + expression: FilterExpression, + record: Mapping[str, Any], + definition: VectorStoreCollectionDefinition, +) -> bool: + """Evaluate a filter against the record as it is stored, one leaf at a time. + + A comparison operator restores its own leaf's field to the type the field + declares, so the native operand compares in the field's semantics. The + text operators stay on the stored string, which lets one group mix a text + operator and a native comparison over the same field. + """ + if isinstance(expression, FilterGroup): + values = (_evaluate_in_memory_filter(item, record, definition) for item in expression.filters) + match expression.operator: + case "and": + return all(values) + case "or": + return any(values) + case "not": + return not next(values) + case _: + raise ValueError(f"Unknown filter group operator '{expression.operator}'.") + field = definition.try_get_field(expression.field_name) + if ( + field is None + or field.type_ not in _SERIALIZED_SCALAR_TYPE_NAMES + or expression.operator in _TEXT_VALUE_OPERATORS + ): + return _evaluate_filter(expression, record, definition) + restored = dict(record) + storage_name = field.storage_name or field.name + if storage_name in restored: + restored[storage_name] = _restore_in_memory_value(restored[storage_name], field.type_) + return _evaluate_filter(expression, restored, definition) + + def _validate_in_memory_filter_value(value: Any) -> None: if value is None or isinstance(value, _SCALAR_FILTER_TYPES): if isinstance(value, float) and not math.isfinite(value): @@ -307,7 +424,8 @@ async def _inner_get( records = list(self._state.records.values()) if filter is not None: _validate_in_memory_filter(filter, self.definition) - records = [record for record in records if _evaluate_filter(filter, record, self.definition)] + _reject_incompatible_filter_operands(filter, self.definition) + records = [record for record in records if _evaluate_in_memory_filter(filter, record, self.definition)] if order_by: for field_name, ascending in reversed(tuple(order_by.items())): if not isinstance(ascending, bool): @@ -362,6 +480,7 @@ async def _inner_search( _validate_vector_dimensions(query_vector, vector_field) if filter is not None: _validate_in_memory_filter(filter, self.definition) + _reject_incompatible_filter_operands(filter, self.definition) distance_function = vector_field.distance_function or "DEFAULT" if distance_function == "DEFAULT": @@ -375,7 +494,7 @@ async def _inner_search( results: list[dict[str, Any]] = [] key_storage_name = self.definition.key_field_storage_name for record in self._state.records.values(): - if filter is not None and not _evaluate_filter(filter, record, self.definition): + if filter is not None and not _evaluate_in_memory_filter(filter, record, self.definition): continue stored_vector = record.get(storage_name) if stored_vector is None: @@ -434,6 +553,57 @@ def _validate_in_memory_filter( _validate_in_memory_filter_value(item.value) +def _filter_operand_values(item: Filter) -> tuple[Any, ...]: + """Return the operands a filter compares against the field value.""" + if item.operator in _COMPARISON_OPERATORS: + return (item.value,) + if item.operator in _SEQUENCE_OPERAND_OPERATORS: + value = item.value + if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + return tuple(cast(Sequence[Any], value)) + return (value,) + return () + + +def _operand_matches_field_type(field_type: str, operand: Any) -> bool: + """Return True when the operand can be compared in the type the field declares.""" + if operand is None: + return True + operand_type = type(operand).__name__ + if field_type in _SERIALIZED_SCALAR_TYPE_NAMES: + return operand_type == field_type + return operand_type not in _SERIALIZED_SCALAR_TYPE_NAMES + + +def _reject_incompatible_filter_operands( + expression: FilterExpression, + definition: VectorStoreCollectionDefinition, +) -> None: + """Reject an operand that cannot be compared in the type the field declares. + + A record is stored through the shared serializer, so the scalar types it + writes as text are restored before the comparison. An operand of another + type cannot be compared in the field's own semantics, so it is rejected + instead of quietly matching no record. The text operators are exempt: + they compare the stored string itself, and their operand is already + required to be a string by the shared filter validation. + """ + for item in _walk_filters(expression): + field = definition.try_get_field(item.field_name) + if field is None or field.type_ is None: + continue + if item.operator in _TEXT_VALUE_OPERATORS: + continue + for operand in _filter_operand_values(item): + if _operand_matches_field_type(field.type_, operand): + continue + raise TypeError( + f"InMemoryCollection cannot compare field '{item.field_name}' of type '{field.type_}' " + f"with a filter operand of type '{type(operand).__name__}'. The in-memory connector " + "compares native scalars in the type the field declares." + ) + + @experimental(feature_id=ExperimentalFeature.VECTOR_STORES) class InMemoryStore(BaseVectorStore): """Create in-memory collection clients that share process-local state.""" diff --git a/python/packages/core/tests/core/test_in_memory.py b/python/packages/core/tests/core/test_in_memory.py index 1253a189519..8a30897a040 100644 --- a/python/packages/core/tests/core/test_in_memory.py +++ b/python/packages/core/tests/core/test_in_memory.py @@ -8,7 +8,10 @@ import warnings from collections.abc import Iterator, Sequence from dataclasses import dataclass +from datetime import date, datetime, time, timedelta, timezone +from decimal import Decimal from typing import Annotated, Any, cast +from uuid import UUID import pytest @@ -866,3 +869,397 @@ def test_in_memory_filter_implementation_does_not_compile_or_evaluate_source() - def test_in_memory_and_filter_apis_are_experimental() -> None: for api in (Filter, FilterGroup, Param, InMemoryCollection, InMemoryStore): assert getattr(api, "__feature_stage__", None) == "experimental" + + +@vectorstoremodel(collection_name="scalar_filter_records") +@dataclass +class ScalarFilterRecord: + id: Annotated[str, VectorStoreField("key")] + created: Annotated[datetime, VectorStoreField("data")] + day: Annotated[date, VectorStoreField("data")] + token: Annotated[UUID, VectorStoreField("data")] + amount: Annotated[Decimal, VectorStoreField("data")] + span: Annotated[timedelta, VectorStoreField("data")] + label: Annotated[str, VectorStoreField("data")] + count: Annotated[int, VectorStoreField("data")] + flag: Annotated[bool, VectorStoreField("data")] + vector: Annotated[list[float] | None, VectorStoreField("vector", dimensions=2)] = None + + +_SCALAR_FILTER_RECORD = ScalarFilterRecord( + "one", + datetime(2026, 10, 16, 1, 0, tzinfo=timezone.utc), + date(2026, 10, 16), + UUID("00000000-0000-0000-0000-000000000001"), + Decimal("1.5"), + timedelta(hours=2), + "featured", + 5, + True, + [1.0, 0.0], +) + + +async def _scalar_filter_collection() -> InMemoryCollection[str, ScalarFilterRecord]: + collection = InMemoryCollection(ScalarFilterRecord) + await collection.ensure_collection_exists() + await collection.upsert([_SCALAR_FILTER_RECORD], generate_vectors=False) + return collection + + +@pytest.mark.parametrize( + ("field_name", "value"), + [ + ("created", datetime(2026, 10, 16, 1, 0, tzinfo=timezone.utc)), + ("day", date(2026, 10, 16)), + ("token", UUID("00000000-0000-0000-0000-000000000001")), + ("amount", Decimal("1.5")), + ("span", timedelta(hours=2)), + ], +) +async def test_in_memory_filter_matches_a_native_scalar_against_the_stored_value( + field_name: str, + value: Any, +) -> None: + """A native scalar filter must match the record it was stored from. + + The stored value is the serialized form of the field, so the operand has to + be compared in the same representation. + """ + collection = await _scalar_filter_collection() + + matched = await collection.get(filter=Filter(field_name, "eq", value)) + + assert [record.id for record in matched] == ["one"] + + +@pytest.mark.parametrize( + ("field_name", "value"), + [ + ("label", "featured"), + ("count", 5), + ("flag", True), + ], +) +async def test_in_memory_filter_still_matches_the_scalars_it_stored_unchanged( + field_name: str, + value: Any, +) -> None: + """The controls the issue lists as working must keep working.""" + collection = await _scalar_filter_collection() + + matched = await collection.get(filter=Filter(field_name, "eq", value)) + + assert [record.id for record in matched] == ["one"] + + +async def test_in_memory_vector_search_matches_a_native_scalar_filter() -> None: + """The search path compares the operands the same way as the get path.""" + collection = await _scalar_filter_collection() + value = datetime(2026, 10, 16, 1, 0, tzinfo=timezone.utc) + + results = await collection.search(vector=[1.0, 0.0], filter=Filter("created", "eq", value), top=10) + + assert [result["record"].id async for result in results] == ["one"] + + +async def test_in_memory_ordered_filter_compares_a_native_scalar() -> None: + """An ordered comparison on a native scalar must compare, not raise.""" + collection = await _scalar_filter_collection() + value = datetime(2026, 10, 16, 1, 0, tzinfo=timezone.utc) + + matched = await collection.get(filter=Filter("created", "gte", value)) + + assert [record.id for record in matched] == ["one"] + + +async def _scalar_amount_collection() -> InMemoryCollection[str, ScalarFilterRecord]: + """Two records whose amounts and timestamps sort the other way round as text.""" + collection = InMemoryCollection(ScalarFilterRecord) + await collection.ensure_collection_exists() + await collection.upsert( + [ + ScalarFilterRecord( + id="two", + created=datetime(2026, 10, 16, 23, 0, tzinfo=timezone(timedelta(hours=2))), + day=date(2026, 10, 16), + token=UUID("00000000-0000-0000-0000-000000000002"), + amount=Decimal("2"), + span=timedelta(hours=2), + label="small", + count=2, + flag=True, + vector=[1.0, 0.0], + ), + ScalarFilterRecord( + id="ten", + created=datetime(2026, 10, 17, 1, 0, tzinfo=timezone.utc), + day=date(2026, 10, 17), + token=UUID("00000000-0000-0000-0000-000000000010"), + amount=Decimal("10"), + span=timedelta(hours=10), + label="large", + count=10, + flag=True, + vector=[0.0, 1.0], + ), + ], + generate_vectors=False, + ) + return collection + + +async def test_in_memory_ordered_filter_compares_decimals_numerically() -> None: + """An ordered comparison on a decimal must compare numbers, not digits. + + The stored form of the two values sorts the other way round as text, so a + comparison on the serialized form would pick the wrong record. + """ + collection = await _scalar_amount_collection() + + matched = await collection.get(filter=Filter("amount", "gte", Decimal("10"))) + + assert [record.id for record in matched] == ["ten"] + + +async def test_in_memory_filter_does_not_match_a_different_native_scalar() -> None: + """A different value must still not match.""" + collection = await _scalar_filter_collection() + + matched = await collection.get(filter=Filter("created", "eq", datetime(2025, 1, 1, tzinfo=timezone.utc))) + + assert matched == [] + + +async def test_in_memory_filter_does_not_match_a_different_decimal() -> None: + """A different decimal must not match the stored one.""" + collection = await _scalar_amount_collection() + + matched = await collection.get(filter=Filter("amount", "eq", Decimal("2"))) + + assert [record.id for record in matched] == ["two"] + + +async def test_in_memory_filter_matches_the_same_instant_written_in_another_offset() -> None: + """The same instant written with another UTC offset is still the same value.""" + collection = await _scalar_filter_collection() + value = datetime(2026, 10, 16, 3, 0, tzinfo=timezone(timedelta(hours=2))) + + matched = await collection.get(filter=Filter("created", "eq", value)) + + assert [record.id for record in matched] == ["one"] + + +async def test_in_memory_ordered_filter_compares_timestamps_written_in_another_offset() -> None: + """The order of two timestamps must not depend on their UTC offsets. + + The stored form of the earlier timestamp sorts after the later one as text, + so a comparison on the serialized form would pick the wrong record. + """ + collection = await _scalar_amount_collection() + value = datetime(2026, 10, 16, 22, 0, tzinfo=timezone.utc) + + matched = await collection.get(filter=Filter("created", "gte", value)) + + assert [record.id for record in matched] == ["ten"] + + +@pytest.mark.parametrize( + ("field_name", "value"), + [ + ("created", "2026-10-16T01:00:00+00:00"), + ("day", datetime(2026, 10, 16)), + ("amount", 1.5), + ("span", 7200), + ("token", "00000000-0000-0000-0000-000000000001"), + ], +) +async def test_in_memory_filter_rejects_an_operand_of_another_type(field_name: str, value: Any) -> None: + """An operand that cannot be compared in the field's own type is rejected. + + The operand validator accepts every one of these operands, so without a + rejection the filter would quietly match no record at all. + """ + collection = await _scalar_filter_collection() + + with pytest.raises(TypeError, match="cannot compare field"): + await collection.get(filter=Filter(field_name, "eq", value)) + + +@pytest.mark.parametrize( + ("field_name", "value"), + [ + ("label", datetime(2026, 10, 16, 1, 0, tzinfo=timezone.utc)), + ("count", Decimal("5")), + ], +) +async def test_in_memory_filter_rejects_a_native_scalar_operand_for_another_field_type( + field_name: str, + value: Any, +) -> None: + """A native scalar operand only fits the field that declares it.""" + collection = await _scalar_filter_collection() + + with pytest.raises(TypeError, match="cannot compare field"): + await collection.get(filter=Filter(field_name, "eq", value)) + + +async def test_in_memory_search_rejects_an_operand_of_another_type() -> None: + """The search path rejects the operands the get path rejects.""" + collection = await _scalar_filter_collection() + + with pytest.raises(TypeError, match="cannot compare field"): + await collection.search(vector=[1.0, 0.0], filter=Filter("created", "eq", "2026-10-16"), top=10) + + +async def test_in_memory_filter_rejects_an_incompatible_operand_inside_a_collection() -> None: + """Every operand of a collection filter is checked.""" + collection = await _scalar_filter_collection() + value = datetime(2026, 10, 16, 1, 0, tzinfo=timezone.utc) + + with pytest.raises(TypeError, match="cannot compare field"): + await collection.get(filter=Filter("created", "in", (value, "2026-10-16"))) + with pytest.raises(TypeError, match="cannot compare field"): + await collection.get(filter=Filter("created", "between", (value, "2026-10-17"))) + + +async def test_in_memory_filter_rejects_an_incompatible_operand_inside_a_group() -> None: + """Every operand of a filter group is checked.""" + collection = await _scalar_filter_collection() + + with pytest.raises(TypeError, match="cannot compare field"): + await collection.get( + filter=FilterGroup( + "and", + (Filter("label", "eq", "featured"), Filter("created", "eq", "2026-10-16")), + ) + ) + + +async def test_in_memory_filter_still_matches_a_collection_of_native_scalars() -> None: + """A collection of operands that fit the field still matches.""" + collection = await _scalar_filter_collection() + value = datetime(2026, 10, 16, 1, 0, tzinfo=timezone.utc) + + matched = await collection.get(filter=Filter("created", "in", (value,))) + + assert [record.id for record in matched] == ["one"] + + +@vectorstoremodel(collection_name="time_filter_records") +@dataclass +class TimeFilterRecord: + id: Annotated[str, VectorStoreField("key")] + starts: Annotated[time, VectorStoreField("data")] + vector: Annotated[list[float] | None, VectorStoreField("vector", dimensions=2)] = None + + +async def _time_filter_collection() -> InMemoryCollection[str, TimeFilterRecord]: + collection = InMemoryCollection(TimeFilterRecord) + await collection.ensure_collection_exists() + await collection.upsert([TimeFilterRecord("one", time(12, 30, tzinfo=timezone.utc))], generate_vectors=False) + return collection + + +async def test_in_memory_filter_matches_a_utc_time_stored_with_a_z_suffix() -> None: + """A UTC time is stored with a trailing "Z" and must still match. + + ``time.fromisoformat`` accepts the "Z" suffix only from Python 3.11, so + the restore normalizes it the way the datetime restore does. + """ + collection = await _time_filter_collection() + + matched = await collection.get(filter=Filter("starts", "eq", time(12, 30, tzinfo=timezone.utc))) + + assert [record.id for record in matched] == ["one"] + + +async def test_in_memory_text_operator_compares_the_stored_scalar_string() -> None: + """A text operator reads the stored string, not the restored scalar. + + The serialized timestamp is what the text operators have always compared + these fields by. + """ + collection = await _scalar_filter_collection() + + matched = await collection.get(filter=Filter("created", "starts_with", "2026-10-16T01")) + + assert [record.id for record in matched] == ["one"] + + +async def test_in_memory_group_mixes_a_text_operator_and_a_native_comparison() -> None: + """One field can carry a text operator and a native comparison in a group. + + The two leaves need different forms of the stored value: the text operator + reads the serialized string while the comparison restores the scalar. + """ + collection = await _scalar_amount_collection() + + matched = await collection.get( + filter=FilterGroup( + "and", + ( + Filter("created", "starts_with", "2026-10-17"), + Filter("created", "gte", datetime(2026, 10, 17, 1, 0, tzinfo=timezone.utc)), + ), + ) + ) + + assert [record.id for record in matched] == ["ten"] + + +async def _long_span_collection(span: timedelta) -> InMemoryCollection[str, ScalarFilterRecord]: + collection = InMemoryCollection(ScalarFilterRecord) + await collection.ensure_collection_exists() + await collection.upsert( + [ + ScalarFilterRecord( + id="long", + created=datetime(2026, 10, 16, 1, 0, tzinfo=timezone.utc), + day=date(2026, 10, 16), + token=UUID("00000000-0000-0000-0000-000000000001"), + amount=Decimal("1.5"), + span=span, + label="long", + count=1, + flag=True, + ) + ], + generate_vectors=False, + ) + return collection + + +async def test_in_memory_restores_a_large_duration_exactly() -> None: + """A microsecond inside a large duration survives the restore. + + Summing the duration through ``float`` loses this microsecond, so the + equality would quietly match nothing. + """ + collection = await _long_span_collection(timedelta(days=100_000, microseconds=1)) + + matched = await collection.get(filter=Filter("span", "eq", timedelta(days=100_000, microseconds=1))) + + assert [record.id for record in matched] == ["long"] + + +async def test_in_memory_distinguishes_durations_a_microsecond_apart() -> None: + """Two durations a microsecond apart must not collapse into one.""" + collection = await _long_span_collection(timedelta(days=100_000, microseconds=1)) + + matched = await collection.get(filter=Filter("span", "eq", timedelta(days=100_000, microseconds=2))) + + assert matched == [] + + +async def test_in_memory_restores_the_largest_duration_exactly() -> None: + """``timedelta.max`` is restored, not rounded out of the type's range. + + A ``float`` of the total seconds rounds the stored duration past the + largest span ``timedelta`` supports. + """ + collection = await _long_span_collection(timedelta.max) + + matched = await collection.get(filter=Filter("span", "eq", timedelta.max)) + + assert [record.id for record in matched] == ["long"]