From 1295167d9486631f8dc97b0891b33eb9a46559e5 Mon Sep 17 00:00:00 2001 From: Michael Landis Date: Tue, 6 Feb 2024 10:09:49 -0800 Subject: [PATCH] feat: mvi id in set filter expression (#430) Bumps the python proto dependency and implements the MVI filter expression IdInSet. This filter expression is satisfied when an item has an id that is in the set provided by the user. We also add integration tests for this filter. --- poetry.lock | 20 +++++++++++++++---- pyproject.toml | 2 +- src/momento/requests/vector_index/filters.py | 15 ++++++++++++++ .../momento/vector_index_client/test_data.py | 5 ++++- .../vector_index_client/test_data_async.py | 5 ++++- 5 files changed, 40 insertions(+), 7 deletions(-) diff --git a/poetry.lock b/poetry.lock index 955398be..9b088bbf 100644 --- a/poetry.lock +++ b/poetry.lock @@ -158,6 +158,7 @@ files = [ {file = "libcst-0.4.10-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f3e9d9fdd9a9b9b8991936ff1c07527ce7ef396c8233280ba9a7137e72c2e48e"}, {file = "libcst-0.4.10-cp39-cp39-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:7e1b4cbaf7b1cdad5fa3eababe42d5b46c0d52afe13c5ba4eac2495fc57630ea"}, {file = "libcst-0.4.10-cp39-cp39-win_amd64.whl", hash = "sha256:bcbd07cec3d7a7be6f0299b0c246e085e3d6cc8af367e2c96059183b97c2e2fe"}, + {file = "libcst-0.4.10.tar.gz", hash = "sha256:b98a829d96e8b209fb761b00cd1bacc27c70eae77d00e57976e5ae2c718c3f81"}, ] [package.dependencies] @@ -170,13 +171,13 @@ dev = ["Sphinx (>=5.1.1)", "black (==23.1.0)", "build (>=0.10.0)", "coverage (>= [[package]] name = "momento-wire-types" -version = "0.102.1" +version = "0.105.1" description = "Momento Client Proto Generated Files" optional = false python-versions = ">=3.7,<4.0" files = [ - {file = "momento_wire_types-0.102.1-py3-none-any.whl", hash = "sha256:f9474cf96d7c36770a165a6ad24b5448b0fe519ea40132fd6b6d1257b46ec0c0"}, - {file = "momento_wire_types-0.102.1.tar.gz", hash = "sha256:47ee4a2db7e5b9e5e01c969c8e6926c17f3fcb7bfef7b642b774a09dca80d625"}, + {file = "momento_wire_types-0.105.1-py3-none-any.whl", hash = "sha256:b27260c9745a93d919ca0de2791c1dfe46b2f021f6532cf8a3dd15f54c23b305"}, + {file = "momento_wire_types-0.105.1.tar.gz", hash = "sha256:1db1b86b7c9b26a542041ae634974294d51e1adc7ed1912f4286e0fb1d441ecb"}, ] [package.dependencies] @@ -398,6 +399,7 @@ files = [ {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:69b023b2b4daa7548bcfbd4aa3da05b3a74b772db9e23b982788168117739938"}, {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:81e0b275a9ecc9c0c0c07b4b90ba548307583c125f54d5b6946cfee6360c733d"}, {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ba336e390cd8e4d1739f42dfe9bb83a3cc2e80f567d8805e11b46f4a943f5515"}, + {file = "PyYAML-6.0.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:326c013efe8048858a6d312ddd31d56e468118ad4cdeda36c719bf5bb6192290"}, {file = "PyYAML-6.0.1-cp310-cp310-win32.whl", hash = "sha256:bd4af7373a854424dabd882decdc5579653d7868b8fb26dc7d0e99f823aa5924"}, {file = "PyYAML-6.0.1-cp310-cp310-win_amd64.whl", hash = "sha256:fd1592b3fdf65fff2ad0004b5e363300ef59ced41c2e6b3a99d4089fa8c5435d"}, {file = "PyYAML-6.0.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:6965a7bc3cf88e5a1c3bd2e0b5c22f8d677dc88a455344035f03399034eb3007"}, @@ -405,8 +407,16 @@ files = [ {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:42f8152b8dbc4fe7d96729ec2b99c7097d656dc1213a3229ca5383f973a5ed6d"}, {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:062582fca9fabdd2c8b54a3ef1c978d786e0f6b3a1510e0ac93ef59e0ddae2bc"}, {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d2b04aac4d386b172d5b9692e2d2da8de7bfb6c387fa4f801fbf6fb2e6ba4673"}, + {file = "PyYAML-6.0.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:e7d73685e87afe9f3b36c799222440d6cf362062f78be1013661b00c5c6f678b"}, {file = "PyYAML-6.0.1-cp311-cp311-win32.whl", hash = "sha256:1635fd110e8d85d55237ab316b5b011de701ea0f29d07611174a1b42f1444741"}, {file = "PyYAML-6.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:bf07ee2fef7014951eeb99f56f39c9bb4af143d8aa3c21b1677805985307da34"}, + {file = "PyYAML-6.0.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:855fb52b0dc35af121542a76b9a84f8d1cd886ea97c84703eaa6d88e37a2ad28"}, + {file = "PyYAML-6.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40df9b996c2b73138957fe23a16a4f0ba614f4c0efce1e9406a184b6d07fa3a9"}, + {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a08c6f0fe150303c1c6b71ebcd7213c2858041a7e01975da3a99aed1e7a378ef"}, + {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6c22bec3fbe2524cde73d7ada88f6566758a8f7227bfbf93a408a9d86bcc12a0"}, + {file = "PyYAML-6.0.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8d4e9c88387b0f5c7d5f281e55304de64cf7f9c0021a3525bd3b1c542da3b0e4"}, + {file = "PyYAML-6.0.1-cp312-cp312-win32.whl", hash = "sha256:d483d2cdf104e7c9fa60c544d92981f12ad66a457afae824d146093b8c294c54"}, + {file = "PyYAML-6.0.1-cp312-cp312-win_amd64.whl", hash = "sha256:0d3304d8c0adc42be59c5f8a4d9e3d7379e6955ad754aa9d6ab7a398b59dd1df"}, {file = "PyYAML-6.0.1-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:50550eb667afee136e9a77d6dc71ae76a44df8b3e51e41b77f6de2932bfe0f47"}, {file = "PyYAML-6.0.1-cp36-cp36m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1fe35611261b29bd1de0070f0b2f47cb6ff71fa6595c077e42bd0c419fa27b98"}, {file = "PyYAML-6.0.1-cp36-cp36m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:704219a11b772aea0d8ecd7058d0082713c3562b4e271b849ad7dc4a5c90c13c"}, @@ -423,6 +433,7 @@ files = [ {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a0cd17c15d3bb3fa06978b4e8958dcdc6e0174ccea823003a106c7d4d7899ac5"}, {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:28c119d996beec18c05208a8bd78cbe4007878c6dd15091efb73a30e90539696"}, {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7e07cbde391ba96ab58e532ff4803f79c4129397514e1413a7dc761ccd755735"}, + {file = "PyYAML-6.0.1-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:49a183be227561de579b4a36efbb21b3eab9651dd81b1858589f796549873dd6"}, {file = "PyYAML-6.0.1-cp38-cp38-win32.whl", hash = "sha256:184c5108a2aca3c5b3d3bf9395d50893a7ab82a38004c8f61c258d4428e80206"}, {file = "PyYAML-6.0.1-cp38-cp38-win_amd64.whl", hash = "sha256:1e2722cc9fbb45d9b87631ac70924c11d3a401b2d7f410cc0e3bbf249f2dca62"}, {file = "PyYAML-6.0.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:9eb6caa9a297fc2c2fb8862bc5370d0303ddba53ba97e71f08023b6cd73d16a8"}, @@ -430,6 +441,7 @@ files = [ {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:5773183b6446b2c99bb77e77595dd486303b4faab2b086e7b17bc6bef28865f6"}, {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:b786eecbdf8499b9ca1d697215862083bd6d2a99965554781d0d8d1ad31e13a0"}, {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bc1bf2925a1ecd43da378f4db9e4f799775d6367bdb94671027b73b393a7c42c"}, + {file = "PyYAML-6.0.1-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:04ac92ad1925b2cff1db0cfebffb6ffc43457495c9b3c39d3fcae417d7125dc5"}, {file = "PyYAML-6.0.1-cp39-cp39-win32.whl", hash = "sha256:faca3bdcf85b2fc05d06ff3fbc1f83e1391b3e724afa3feba7d13eeab355484c"}, {file = "PyYAML-6.0.1-cp39-cp39-win_amd64.whl", hash = "sha256:510c9deebc5c0225e8c96813043e62b680ba2f9c50a08d3724c7f28a747d1486"}, {file = "PyYAML-6.0.1.tar.gz", hash = "sha256:bfdf460b1736c775f2ba9f6a92bca30bc2095067b8a9d77876d1fad6cc3b4a43"}, @@ -605,4 +617,4 @@ testing = ["big-O", "flake8 (<5)", "jaraco.functools", "jaraco.itertools", "more [metadata] lock-version = "2.0" python-versions = "^3.7" -content-hash = "66719b0ab5216e77ebbb7e448c0eee8a4d8d5b71ce44f69cfc1d71c1faf24528" +content-hash = "c77cae9d0aa0f8a51a272111bb0f3629c55572d151b4928a988609ebbb472079" diff --git a/pyproject.toml b/pyproject.toml index ff8affaa..fc276dd0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,7 +30,7 @@ exclude = ["src/momento/internal/codegen.py"] [tool.poetry.dependencies] python = "^3.7" -momento-wire-types = "^0.102.1" +momento-wire-types = "^0.105.1" grpcio = "^1.46.0" # note if you bump this presigned url test need be updated pyjwt = "^2.4.0" diff --git a/src/momento/requests/vector_index/filters.py b/src/momento/requests/vector_index/filters.py index 592c131b..bb8ced1b 100644 --- a/src/momento/requests/vector_index/filters.py +++ b/src/momento/requests/vector_index/filters.py @@ -14,6 +14,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass +from typing import Iterable from momento_wire_types import vectorindex_pb2 as vectorindex_pb @@ -258,3 +259,17 @@ def to_filter_expression_proto(self) -> vectorindex_pb._FilterExpression: def to_proto(self) -> vectorindex_pb._ListContainsExpression: # todo should make oneof defensively return vectorindex_pb._ListContainsExpression(field=self.field, string_value=self.value) + + +@dataclass +class IdInSet(FilterExpression): + """Represents an expression to test if an item id is in a set of ids.""" + + ids: Iterable[str] + """The set of ids to test id in set with.""" + + def to_filter_expression_proto(self) -> vectorindex_pb._FilterExpression: + return vectorindex_pb._FilterExpression(id_in_set_expression=self.to_proto()) + + def to_proto(self) -> vectorindex_pb._IdInSetExpression: + return vectorindex_pb._IdInSetExpression(ids=list(self.ids)) diff --git a/tests/momento/vector_index_client/test_data.py b/tests/momento/vector_index_client/test_data.py index 342f5503..304613e9 100644 --- a/tests/momento/vector_index_client/test_data.py +++ b/tests/momento/vector_index_client/test_data.py @@ -6,7 +6,7 @@ from momento import PreviewVectorIndexClient from momento.common_data.vector_index.item import Metadata from momento.errors import MomentoErrorCode -from momento.requests.vector_index import ALL_METADATA, Field, FilterExpression, Item, SimilarityMetric +from momento.requests.vector_index import ALL_METADATA, Field, FilterExpression, Item, SimilarityMetric, filters from momento.responses.vector_index import ( CountItems, CreateIndex, @@ -657,6 +657,9 @@ def test_search_with_filter_expression( ["test_item_1", "test_item_2", "test_item_3"], "list contains b or int > 1", ), + (filters.IdInSet({}), [], "id in empty set"), + (filters.IdInSet({"not there"}), [], "id in set not there"), + (filters.IdInSet({"test_item_1", "test_item_3"}), ["test_item_1", "test_item_3"], "id in set"), ]: filter_expression = cast(FilterExpression, filter_expression) search_response = vector_index_client.search( diff --git a/tests/momento/vector_index_client/test_data_async.py b/tests/momento/vector_index_client/test_data_async.py index 8d5a24fe..b9214450 100644 --- a/tests/momento/vector_index_client/test_data_async.py +++ b/tests/momento/vector_index_client/test_data_async.py @@ -6,7 +6,7 @@ from momento import PreviewVectorIndexClientAsync from momento.common_data.vector_index.item import Metadata from momento.errors import MomentoErrorCode -from momento.requests.vector_index import ALL_METADATA, Field, FilterExpression, Item, SimilarityMetric +from momento.requests.vector_index import ALL_METADATA, Field, FilterExpression, Item, SimilarityMetric, filters from momento.responses.vector_index import ( CountItems, CreateIndex, @@ -661,6 +661,9 @@ async def test_search_with_filter_expression( ["test_item_1", "test_item_2", "test_item_3"], "list contains b or int > 1", ), + (filters.IdInSet({}), [], "id in empty set"), + (filters.IdInSet({"not there"}), [], "id in set not there"), + (filters.IdInSet({"test_item_1", "test_item_3"}), ["test_item_1", "test_item_3"], "id in set"), ]: filter_expression = cast(FilterExpression, filter_expression) search_response = await vector_index_client_async.search(