From 79bc60ef75ea098112403f91b11f2b221be58b14 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Sat, 23 Mar 2024 18:07:38 -0700 Subject: [PATCH 1/8] chore(deps): bump momento from 1.18.0 to 1.20.1 in /examples (#447) Bumps [momento](https://github.com/momentohq/client-sdk-python) from 1.18.0 to 1.20.1. - [Release notes](https://github.com/momentohq/client-sdk-python/releases) - [Commits](https://github.com/momentohq/client-sdk-python/compare/v1.18.0...v1.20.1) --- updated-dependencies: - dependency-name: momento dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- examples/poetry.lock | 18 +++++++++--------- examples/pyproject.toml | 2 +- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/examples/poetry.lock b/examples/poetry.lock index f842f904..276a2aeb 100644 --- a/examples/poetry.lock +++ b/examples/poetry.lock @@ -193,30 +193,30 @@ testing = ["flufl.flake8", "importlib-resources (>=1.3)", "packaging", "pep517", [[package]] name = "momento" -version = "1.18.0" +version = "1.20.1" description = "SDK for Momento" optional = false -python-versions = ">=3.7,<4.0" +python-versions = "<4.0,>=3.7" files = [ - {file = "momento-1.18.0-py3-none-any.whl", hash = "sha256:2e05a3b6dd46320a338303dfe067d517009500da5f90221cc4c9d916c16719fd"}, - {file = "momento-1.18.0.tar.gz", hash = "sha256:70b9a13a630090276ddd0431ba6e247544cc83f55a0b36f396e267b2c2669c82"}, + {file = "momento-1.20.1-py3-none-any.whl", hash = "sha256:453b2f1629f866d1ebec731894516368412ad8be8155f0b3ff890d497dda3520"}, + {file = "momento-1.20.1.tar.gz", hash = "sha256:756cd0bd15fabb7b1bebab5e704622f2aeecc5a6029c03b956b51644f30b7820"}, ] [package.dependencies] grpcio = ">=1.46.0,<2.0.0" importlib-metadata = {version = ">=4", markers = "python_version < \"3.8\""} -momento-wire-types = ">=0.105.3,<0.106.0" +momento-wire-types = ">=0.106.0,<0.107.0" pyjwt = ">=2.4.0,<3.0.0" [[package]] name = "momento-wire-types" -version = "0.105.3" +version = "0.106.2" description = "Momento Client Proto Generated Files" optional = false python-versions = ">=3.7,<4.0" files = [ - {file = "momento_wire_types-0.105.3-py3-none-any.whl", hash = "sha256:156c704eec560b9a102c16aa08283981f23fa6d8e61ddbb05b6206b979733672"}, - {file = "momento_wire_types-0.105.3.tar.gz", hash = "sha256:12a785c7950b8cf62c8516061bfa8b9124619e097f82c16d110e770385ad16bc"}, + {file = "momento_wire_types-0.106.2-py3-none-any.whl", hash = "sha256:1d7e8d199dff045cf8a777e1c1e469946c6ce3420b8da21d8969821466aa4ef9"}, + {file = "momento_wire_types-0.106.2.tar.gz", hash = "sha256:5846a09ac31b01b748371f5f44d0370b591a902f72a390d4c1e6bf347f0d9230"}, ] [package.dependencies] @@ -446,4 +446,4 @@ testing = ["big-O", "flake8 (<5)", "jaraco.functools", "jaraco.itertools", "more [metadata] lock-version = "2.0" python-versions = ">=3.7,<3.12" -content-hash = "97442992d861fa0db4d34f47a80e8103d834713d9206d75ae332e2994e86697b" +content-hash = "4f3802ee06829f8c13b058f8cbf94fc08d77cbe5ff4bdf3c822251c92798a33d" diff --git a/examples/pyproject.toml b/examples/pyproject.toml index 1cfb9688..cfdd1d2e 100644 --- a/examples/pyproject.toml +++ b/examples/pyproject.toml @@ -8,7 +8,7 @@ license = "Apache-2.0" [tool.poetry.dependencies] python = ">=3.7,<3.12" -momento = "1.18.0" +momento = "1.20.1" colorlog = "6.7.0" hdrhistogram = "^0.10.1" From 268068f0bdd315fafaccfb84fee4c71d8e67f023 Mon Sep 17 00:00:00 2001 From: Chris Price Date: Thu, 18 Apr 2024 17:19:30 -0700 Subject: [PATCH 2/8] feat: remove vector client (#450) * feat: remove vector client --- .github/workflows/on-pull-request-mvi.yml | 64 -- .github/workflows/on-pull-request.yml | 3 +- .../workflows/on-push-to-release-branch.yml | 3 +- Makefile | 5 - examples/README.md | 20 - examples/prepy310/vector_index.py | 149 --- examples/prepy310/vector_index_async.py | 156 --- examples/py310/doc-examples-python-apis.py | 239 ----- examples/py310/vector_index.py | 155 --- examples/py310/vector_index_async.py | 162 ---- pyproject.toml | 13 - src/momento/__init__.py | 7 +- src/momento/auth/credential_provider.py | 12 +- src/momento/auth/momento_endpoint_resolver.py | 5 - .../common_data/vector_index/__init__.py | 3 - src/momento/common_data/vector_index/item.py | 91 -- src/momento/config/__init__.py | 4 - .../config/vector_index_configuration.py | 83 -- .../config/vector_index_configurations.py | 30 - src/momento/internal/_utilities/__init__.py | 5 - .../_utilities/_channel_credentials.py | 4 +- .../_utilities/_vector_index_validation.py | 20 - .../aio/_vector_index_control_client.py | 106 -- .../internal/aio/_vector_index_data_client.py | 297 ------ .../aio/_vector_index_grpc_manager.py | 75 -- src/momento/internal/codegen.py | 8 +- .../_vector_index_control_client.py | 106 -- .../synchronous/_vector_index_data_client.py | 297 ------ .../synchronous/_vector_index_grpc_manager.py | 71 -- src/momento/requests/vector_index/__init__.py | 8 - .../requests/vector_index/filter_field.py | 71 -- src/momento/requests/vector_index/filters.py | 275 ------ src/momento/requests/vector_index/search.py | 8 - .../vector_index/similarity_metric.py | 15 - .../responses/vector_index/__init__.py | 48 - .../vector_index/control/__init__.py | 0 .../responses/vector_index/control/create.py | 57 -- .../responses/vector_index/control/delete.py | 49 - .../responses/vector_index/control/list.py | 101 -- .../responses/vector_index/data/__init__.py | 0 .../vector_index/data/count_items.py | 37 - .../vector_index/data/delete_item_batch.py | 30 - .../vector_index/data/get_item_batch.py | 56 -- .../data/get_item_metadata_batch.py | 48 - .../responses/vector_index/data/search.py | 74 -- .../data/search_and_fetch_vectors.py | 77 -- .../vector_index/data/upsert_item_batch.py | 30 - .../responses/vector_index/data/utils.py | 30 - .../responses/vector_index/response.py | 5 - src/momento/vector_index_client.py | 305 ------ src/momento/vector_index_client_async.py | 305 ------ tests/conftest.py | 98 +- .../config/test_vector_index_config.py | 97 -- .../momento/requests/vector_index/__init__.py | 0 .../requests/vector_index/test_filters.py | 140 --- .../requests/vector_index/test_item.py | 57 -- tests/momento/vector_index_client/__init__.py | 0 .../vector_index_client/test_control.py | 185 ---- .../vector_index_client/test_control_async.py | 191 ---- .../momento/vector_index_client/test_data.py | 908 ----------------- .../vector_index_client/test_data_async.py | 916 ------------------ tests/utils.py | 40 - 62 files changed, 12 insertions(+), 6442 deletions(-) delete mode 100644 .github/workflows/on-pull-request-mvi.yml delete mode 100644 examples/prepy310/vector_index.py delete mode 100644 examples/prepy310/vector_index_async.py delete mode 100644 examples/py310/vector_index.py delete mode 100644 examples/py310/vector_index_async.py delete mode 100644 src/momento/common_data/vector_index/__init__.py delete mode 100644 src/momento/common_data/vector_index/item.py delete mode 100644 src/momento/config/vector_index_configuration.py delete mode 100644 src/momento/config/vector_index_configurations.py delete mode 100644 src/momento/internal/_utilities/_vector_index_validation.py delete mode 100644 src/momento/internal/aio/_vector_index_control_client.py delete mode 100644 src/momento/internal/aio/_vector_index_data_client.py delete mode 100644 src/momento/internal/aio/_vector_index_grpc_manager.py delete mode 100644 src/momento/internal/synchronous/_vector_index_control_client.py delete mode 100644 src/momento/internal/synchronous/_vector_index_data_client.py delete mode 100644 src/momento/internal/synchronous/_vector_index_grpc_manager.py delete mode 100644 src/momento/requests/vector_index/__init__.py delete mode 100644 src/momento/requests/vector_index/filter_field.py delete mode 100644 src/momento/requests/vector_index/filters.py delete mode 100644 src/momento/requests/vector_index/search.py delete mode 100644 src/momento/requests/vector_index/similarity_metric.py delete mode 100644 src/momento/responses/vector_index/__init__.py delete mode 100644 src/momento/responses/vector_index/control/__init__.py delete mode 100644 src/momento/responses/vector_index/control/create.py delete mode 100644 src/momento/responses/vector_index/control/delete.py delete mode 100644 src/momento/responses/vector_index/control/list.py delete mode 100644 src/momento/responses/vector_index/data/__init__.py delete mode 100644 src/momento/responses/vector_index/data/count_items.py delete mode 100644 src/momento/responses/vector_index/data/delete_item_batch.py delete mode 100644 src/momento/responses/vector_index/data/get_item_batch.py delete mode 100644 src/momento/responses/vector_index/data/get_item_metadata_batch.py delete mode 100644 src/momento/responses/vector_index/data/search.py delete mode 100644 src/momento/responses/vector_index/data/search_and_fetch_vectors.py delete mode 100644 src/momento/responses/vector_index/data/upsert_item_batch.py delete mode 100644 src/momento/responses/vector_index/data/utils.py delete mode 100644 src/momento/responses/vector_index/response.py delete mode 100644 src/momento/vector_index_client.py delete mode 100644 src/momento/vector_index_client_async.py delete mode 100644 tests/momento/config/test_vector_index_config.py delete mode 100644 tests/momento/requests/vector_index/__init__.py delete mode 100644 tests/momento/requests/vector_index/test_filters.py delete mode 100644 tests/momento/requests/vector_index/test_item.py delete mode 100644 tests/momento/vector_index_client/__init__.py delete mode 100644 tests/momento/vector_index_client/test_control.py delete mode 100644 tests/momento/vector_index_client/test_control_async.py delete mode 100644 tests/momento/vector_index_client/test_data.py delete mode 100644 tests/momento/vector_index_client/test_data_async.py diff --git a/.github/workflows/on-pull-request-mvi.yml b/.github/workflows/on-pull-request-mvi.yml deleted file mode 100644 index 2c2a7217..00000000 --- a/.github/workflows/on-pull-request-mvi.yml +++ /dev/null @@ -1,64 +0,0 @@ -name: On Pull Request - -on: - pull_request: - branches: [main] - -jobs: - test: - runs-on: macos-latest - strategy: - max-parallel: 2 - matrix: - python-version: ["3.7", "3.8", "3.9", "3.10", "3.11"] - new-python-protobuf: ["true"] - include: - - python-version: "3.7" - new-python-protobuf: "false" - - env: - TEST_API_KEY: ${{ secrets.ALPHA_TEST_AUTH_TOKEN }} - TEST_CACHE_NAME: python-integration-test-${{ matrix.python-version }}-${{ matrix.new-python-protobuf }}-${{ github.sha }} - TEST_VECTOR_INDEX_NAME: python-integration-test-vector-${{ matrix.python-version }}-${{ matrix.new-python-protobuf }}-${{ github.sha }} - - steps: - - uses: actions/checkout@v3 - - - name: Commitlint and Other Shared Build Steps - uses: momentohq/standards-and-practices/github-actions/shared-build@gh-actions-v1 - env: - GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} - - - name: Setup Python ${{ matrix.python-version }} - uses: actions/setup-python@v4 - with: - python-version: ${{ matrix.python-version }} - - - name: Bootstrap poetry - run: | - curl -sL https://install.python-poetry.org | python - -y --version 1.3.1 - - - name: Configure poetry - run: /Users/runner/.local/bin/poetry config virtualenvs.in-project true - - - name: Install dependencies - run: /Users/runner/.local/bin/poetry install - - - name: Install Old Protobuf - # Exercises the wire types generated against the old protobuf library - if: matrix.new-python-protobuf == 'false' - run: /Users/runner/.local/bin/poetry add "protobuf<3.20" - - - name: Run mypy - # mypy has inconsistencies between 3.7 and the rest; default to lowest common denominator - if: matrix.python-version == '3.7' - run: /Users/runner/.local/bin/poetry run mypy src tests - - - name: Run ruff analyzer - run: /Users/runner/.local/bin/poetry run ruff check --no-fix src tests - - - name: Run ruff formatter - run: /Users/runner/.local/bin/poetry run ruff format --diff src tests - - - name: Run tests - run: /Users/runner/.local/bin/poetry run pytest tests/momento/vector_index_client -p no:sugar -q diff --git a/.github/workflows/on-pull-request.yml b/.github/workflows/on-pull-request.yml index f4ed58f9..e11e58ee 100644 --- a/.github/workflows/on-pull-request.yml +++ b/.github/workflows/on-pull-request.yml @@ -18,7 +18,6 @@ jobs: env: TEST_API_KEY: ${{ secrets.ALPHA_TEST_AUTH_TOKEN }} TEST_CACHE_NAME: python-integration-test-${{ matrix.python-version }}-${{ matrix.new-python-protobuf }}-${{ github.sha }} - TEST_VECTOR_INDEX_NAME: python-integration-test-vector-${{ matrix.python-version }}-${{ matrix.new-python-protobuf }}-${{ github.sha }} steps: - uses: actions/checkout@v3 @@ -60,7 +59,7 @@ jobs: run: poetry run ruff format --check --diff src tests - name: Run tests - run: poetry run pytest -p no:sugar -q --ignore=tests/momento/vector_index_client + run: poetry run pytest -p no:sugar -q test-examples: runs-on: ubuntu-20.04 diff --git a/.github/workflows/on-push-to-release-branch.yml b/.github/workflows/on-push-to-release-branch.yml index a3f3ff4e..ca0a4863 100644 --- a/.github/workflows/on-push-to-release-branch.yml +++ b/.github/workflows/on-push-to-release-branch.yml @@ -41,7 +41,6 @@ jobs: env: TEST_API_KEY: ${{ secrets.ALPHA_TEST_AUTH_TOKEN }} TEST_CACHE_NAME: python-integration-test-${{ matrix.python-version }}-${{ matrix.new-python-protobuf }}-${{ github.sha }} - TEST_VECTOR_INDEX_NAME: python-integration-test-vector-${{ matrix.python-version }}-${{ matrix.new-python-protobuf }}-${{ github.sha }} steps: - uses: actions/checkout@v3 @@ -78,7 +77,7 @@ jobs: run: poetry run ruff format --diff src tests - name: Run tests - run: poetry run pytest -p no:sugar -q --ignore=tests/momento/vector_index_client + run: poetry run pytest -p no:sugar -q publish: runs-on: ubuntu-20.04 diff --git a/Makefile b/Makefile index 55e3b2f6..aae55b70 100644 --- a/Makefile +++ b/Makefile @@ -27,10 +27,7 @@ lint: do-gen-sync: @poetry run python src/momento/internal/codegen.py src/momento/internal/aio/_scs_control_client.py src/momento/internal/synchronous/_scs_control_client.py @poetry run python src/momento/internal/codegen.py src/momento/internal/aio/_scs_data_client.py src/momento/internal/synchronous/_scs_data_client.py - @poetry run python src/momento/internal/codegen.py src/momento/internal/aio/_vector_index_control_client.py src/momento/internal/synchronous/_vector_index_control_client.py - @poetry run python src/momento/internal/codegen.py src/momento/internal/aio/_vector_index_data_client.py src/momento/internal/synchronous/_vector_index_data_client.py @poetry run python src/momento/internal/codegen.py src/momento/cache_client_async.py src/momento/cache_client.py - @poetry run python src/momento/internal/codegen.py src/momento/vector_index_client_async.py src/momento/vector_index_client.py @poetry run python src/momento/internal/codegen.py tests/momento/cache_client/shared_behaviors_async.py tests/momento/cache_client/shared_behaviors.py @poetry run python src/momento/internal/codegen.py tests/momento/cache_client/test_init_async.py tests/momento/cache_client/test_init.py @poetry run python src/momento/internal/codegen.py tests/momento/cache_client/test_control_async.py tests/momento/cache_client/test_control.py @@ -40,8 +37,6 @@ do-gen-sync: @poetry run python src/momento/internal/codegen.py tests/momento/cache_client/test_set_async.py tests/momento/cache_client/test_set.py @poetry run python src/momento/internal/codegen.py tests/momento/cache_client/test_sorted_set_async.py tests/momento/cache_client/test_sorted_set.py @poetry run python src/momento/internal/codegen.py tests/momento/cache_client/test_sorted_set_simple_async.py tests/momento/cache_client/test_sorted_set_simple.py - @poetry run python src/momento/internal/codegen.py tests/momento/vector_index_client/test_control_async.py tests/momento/vector_index_client/test_control.py - @poetry run python src/momento/internal/codegen.py tests/momento/vector_index_client/test_data_async.py tests/momento/vector_index_client/test_data.py .PHONY: gen-sync ## Generate synchronous code and tests from asynchronous code. diff --git a/examples/README.md b/examples/README.md index b8096140..4c2cc2c4 100644 --- a/examples/README.md +++ b/examples/README.md @@ -48,13 +48,6 @@ MOMENTO_API_KEY= poetry run python -m py310.example MOMENTO_API_KEY= poetry run python -m py310.example_async ``` -To run the python version 3.10+ vector index examples: - -```bash -MOMENTO_API_KEY= poetry run python -m py310.vector_index -MOMENTO_API_KEY= poetry run python -m py310.vector_index_async -``` - To run the examples with SDK debug logging enabled: ```bash @@ -69,13 +62,6 @@ MOMENTO_API_KEY= poetry run python -m prepy310.example MOMENTO_API_KEY= poetry run python -m prepy310.example_async ``` -To run the python version <3.10 vector index examples: - -```bash -MOMENTO_API_KEY= poetry run python -m prepy310.vector_index -MOMENTO_API_KEY= poetry run python -m prepy310.vector_index_async -``` - To run the examples with SDK debug logging enabled: ```bash @@ -119,12 +105,6 @@ DEBUG=true MOMENTO_API_KEY= python -m prepy310.example DEBUG=true MOMENTO_API_KEY= python -m prepy310.example_async ``` -To run the vector index example: - -```bash -MOMENTO_API_KEY= python -m vector_index.example -``` - ## Running the load generator example This repo includes a very basic load generator, to allow you to experiment diff --git a/examples/prepy310/vector_index.py b/examples/prepy310/vector_index.py deleted file mode 100644 index 7e3590f6..00000000 --- a/examples/prepy310/vector_index.py +++ /dev/null @@ -1,149 +0,0 @@ -import logging -from time import sleep - -from momento import ( - CredentialProvider, - PreviewVectorIndexClient, - VectorIndexConfigurations, -) -from momento.config import VectorIndexConfiguration -from momento.requests.vector_index import ALL_METADATA, Item, SimilarityMetric -from momento.responses.vector_index import ( - CreateIndex, - DeleteIndex, - DeleteItemBatch, - ListIndexes, - Search, - UpsertItemBatch, -) - -from example_utils.example_logging import initialize_logging - -_logger = logging.getLogger("vector-example") - -VECTOR_INDEX_CONFIGURATION: VectorIndexConfiguration = VectorIndexConfigurations.Default.latest() -VECTOR_AUTH_PROVIDER = CredentialProvider.from_environment_variable("MOMENTO_API_KEY") - - -def _print_start_banner() -> None: - _logger.info("******************************************************************") - _logger.info("* Momento Example Start *") - _logger.info("******************************************************************\n") - - -def _print_end_banner() -> None: - _logger.info("******************************************************************") - _logger.info("* Momento Example End *") - _logger.info("******************************************************************\n") - - -def create_index( - client: PreviewVectorIndexClient, - index_name: str, - num_dimensions: int, - similarity_metric: SimilarityMetric = SimilarityMetric.COSINE_SIMILARITY, -) -> None: - _logger.info(f"Creating index with name {index_name!r}") - create_index_response = client.create_index(index_name, num_dimensions, similarity_metric) - if isinstance(create_index_response, CreateIndex.Success): - _logger.info(f"Index with name {index_name!r} successfully created!") - elif isinstance(create_index_response, CreateIndex.IndexAlreadyExists): - _logger.info(f"Index with name {index_name!r} already exists") - elif isinstance(create_index_response, CreateIndex.Error): - _logger.error(f"Error while creating index {create_index_response.message}") - else: - _logger.error("Unreachable") - _logger.info("") - - -def list_indexes(client: PreviewVectorIndexClient) -> None: - _logger.info("Listing indexes:") - list_indexes_response = client.list_indexes() - if isinstance(list_indexes_response, ListIndexes.Success): - for index in list_indexes_response.indexes: - _logger.info(f"- {index!r}") - elif isinstance(list_indexes_response, ListIndexes.Error): - _logger.error(f"Error while listing indexes {list_indexes_response.message}") - else: - _logger.error("Unreachable") - _logger.info("") - - -def upsert_items(client: PreviewVectorIndexClient, index_name: str) -> None: - items = [ - Item(id="test_item_1", vector=[1.0, 2.0], metadata={"key1": "value1"}), - Item(id="test_item_2", vector=[3.0, 4.0], metadata={"key2": "value2"}), - Item(id="test_item_3", vector=[5.0, 6.0], metadata={"key1": "value3", "key3": "value3"}), - ] - _logger.info(f"Adding items {items}") - upsert_response = client.upsert_item_batch( - index_name, - items=items, - ) - if isinstance(upsert_response, UpsertItemBatch.Success): - _logger.info("Successfully added items") - elif isinstance(upsert_response, UpsertItemBatch.Error): - _logger.error(f"Error while adding items to index {index_name!r}: {upsert_response.message}") - else: - _logger.error("Unreachable") - _logger.info("") - - -def search(client: PreviewVectorIndexClient, index_name: str) -> None: - query_vector = [1.0, 2.0] - top_k = 3 - _logger.info(f"Searching index {index_name} with query_vector {query_vector} and top {top_k} elements") - search_response = client.search(index_name, query_vector=query_vector, top_k=top_k, metadata_fields=ALL_METADATA) - if isinstance(search_response, Search.Success): - _logger.info(f"Search succeeded with {len(search_response.hits)} matches:") - _logger.info(search_response.hits) - elif isinstance(search_response, Search.Error): - _logger.error(f"Error while searching on index {index_name}: {search_response.message}") - else: - _logger.error("Unreachable") - _logger.info("") - - -def delete_items(client: PreviewVectorIndexClient, index_name: str) -> None: - item_ids_to_delete = ["test_item_1", "test_item_3"] - _logger.info(f"Deleting items: {item_ids_to_delete}") - delete_response = client.delete_item_batch(index_name, filter=item_ids_to_delete) - if isinstance(delete_response, DeleteItemBatch.Success): - _logger.info("Successfully deleted items") - elif isinstance(delete_response, DeleteItemBatch.Error): - _logger.error(f"Error while deleting items {delete_response.message}") - else: - _logger.error("Unreachable") - _logger.info("") - - -def delete_index(client: PreviewVectorIndexClient, index_name: str) -> None: - _logger.info("Deleting index " + index_name) - delete_response = client.delete_index(index_name) - - if isinstance(delete_response, DeleteIndex.Success): - _logger.info(f"Index {index_name} deleted successfully!") - elif isinstance(delete_response, DeleteIndex.Error): - _logger.error(f"Failed to delete index {index_name} with error {delete_response.message}") - else: - _logger.error("Unreachable") - _logger.info("") - - -if __name__ == "__main__": - initialize_logging() - _print_start_banner() - with PreviewVectorIndexClient(VECTOR_INDEX_CONFIGURATION, VECTOR_AUTH_PROVIDER) as client: - index_name = "hello_momento_index" - - create_index(client, index_name, num_dimensions=2) - list_indexes(client) - upsert_items(client, index_name) - sleep(2) - search(client, index_name) - delete_items(client, index_name) - sleep(2) - _logger.info("Deleted two items; search will return 1 hit now") - search(client, index_name) - delete_index(client, index_name) - _print_end_banner() diff --git a/examples/prepy310/vector_index_async.py b/examples/prepy310/vector_index_async.py deleted file mode 100644 index b9adb0bf..00000000 --- a/examples/prepy310/vector_index_async.py +++ /dev/null @@ -1,156 +0,0 @@ -import asyncio -import logging -from time import sleep - -from momento import ( - CredentialProvider, - PreviewVectorIndexClientAsync, - VectorIndexConfigurations, -) -from momento.config import VectorIndexConfiguration -from momento.requests.vector_index import ALL_METADATA, Item, SimilarityMetric -from momento.responses.vector_index import ( - CreateIndex, - DeleteIndex, - DeleteItemBatch, - ListIndexes, - Search, - UpsertItemBatch, -) - -from example_utils.example_logging import initialize_logging - -_logger = logging.getLogger("vector-example") - -VECTOR_INDEX_CONFIGURATION: VectorIndexConfiguration = VectorIndexConfigurations.Default.latest() -VECTOR_AUTH_PROVIDER = CredentialProvider.from_environment_variable("MOMENTO_API_KEY") - - -def _print_start_banner() -> None: - _logger.info("******************************************************************") - _logger.info("* Momento Example Start *") - _logger.info("******************************************************************\n") - - -def _print_end_banner() -> None: - _logger.info("******************************************************************") - _logger.info("* Momento Example End *") - _logger.info("******************************************************************\n") - - -async def create_index( - client: PreviewVectorIndexClientAsync, - index_name: str, - num_dimensions: int, - similarity_metric: SimilarityMetric = SimilarityMetric.COSINE_SIMILARITY, -) -> None: - _logger.info(f"Creating index with name {index_name!r}") - create_index_response = await client.create_index(index_name, num_dimensions, similarity_metric) - if isinstance(create_index_response, CreateIndex.Success): - _logger.info(f"Index with name {index_name!r} successfully created!") - elif isinstance(create_index_response, CreateIndex.IndexAlreadyExists): - _logger.info(f"Index with name {index_name!r} already exists") - elif isinstance(create_index_response, CreateIndex.Error): - _logger.error(f"Error while creating index {create_index_response.message}") - else: - _logger.error("Unreachable") - _logger.info("") - - -async def list_indexes(client: PreviewVectorIndexClientAsync) -> None: - _logger.info("Listing indexes:") - list_indexes_response = await client.list_indexes() - if isinstance(list_indexes_response, ListIndexes.Success): - for index in list_indexes_response.indexes: - _logger.info(f"- {index!r}") - elif isinstance(list_indexes_response, ListIndexes.Error): - _logger.error(f"Error while listing indexes {list_indexes_response.message}") - else: - _logger.error("Unreachable") - _logger.info("") - - -async def upsert_items(client: PreviewVectorIndexClientAsync, index_name: str) -> None: - items = [ - Item(id="test_item_1", vector=[1.0, 2.0], metadata={"key1": "value1"}), - Item(id="test_item_2", vector=[3.0, 4.0], metadata={"key2": "value2"}), - Item(id="test_item_3", vector=[5.0, 6.0], metadata={"key1": "value3", "key3": "value3"}), - ] - _logger.info(f"Adding items {items}") - upsert_response = await client.upsert_item_batch( - index_name, - items=items, - ) - if isinstance(upsert_response, UpsertItemBatch.Success): - _logger.info("Successfully added items") - elif isinstance(upsert_response, UpsertItemBatch.Error): - _logger.error(f"Error while adding items to index {index_name!r}: {upsert_response.message}") - else: - _logger.error("Unreachable") - _logger.info("") - - -async def search(client: PreviewVectorIndexClientAsync, index_name: str) -> None: - query_vector = [1.0, 2.0] - top_k = 3 - _logger.info(f"Searching index {index_name} with query_vector {query_vector} and top {top_k} elements") - search_response = await client.search( - index_name, query_vector=query_vector, top_k=top_k, metadata_fields=ALL_METADATA - ) - if isinstance(search_response, Search.Success): - _logger.info(f"Search succeeded with {len(search_response.hits)} matches:") - _logger.info(search_response.hits) - elif isinstance(search_response, Search.Error): - _logger.error(f"Error while searching on index {index_name}: {search_response.message}") - else: - _logger.error("Unreachable") - _logger.info("") - - -async def delete_items(client: PreviewVectorIndexClientAsync, index_name: str) -> None: - item_ids_to_delete = ["test_item_1", "test_item_3"] - _logger.info(f"Deleting items: {item_ids_to_delete}") - delete_response = await client.delete_item_batch(index_name, filter=item_ids_to_delete) - if isinstance(delete_response, DeleteItemBatch.Success): - _logger.info("Successfully deleted items") - elif isinstance(delete_response, DeleteItemBatch.Error): - _logger.error(f"Error while deleting items {delete_response.message}") - else: - _logger.error("Unreachable") - _logger.info("") - - -async def delete_index(client: PreviewVectorIndexClientAsync, index_name: str) -> None: - _logger.info("Deleting index " + index_name) - delete_response = await client.delete_index(index_name) - - if isinstance(delete_response, DeleteIndex.Success): - _logger.info(f"Index {index_name} deleted successfully!") - elif isinstance(delete_response, DeleteIndex.Error): - _logger.error(f"Failed to delete index {index_name} with error {delete_response.message}") - else: - _logger.error("Unreachable") - _logger.info("") - - -async def main() -> None: - initialize_logging() - _print_start_banner() - async with PreviewVectorIndexClientAsync(VECTOR_INDEX_CONFIGURATION, VECTOR_AUTH_PROVIDER) as client: - index_name = "hello_momento_index" - - await create_index(client, index_name, num_dimensions=2) - await list_indexes(client) - await upsert_items(client, index_name) - sleep(2) - await search(client, index_name) - await delete_items(client, index_name) - sleep(2) - _logger.info("Deleted two items; search will return 1 hit now") - await search(client, index_name) - await delete_index(client, index_name) - _print_end_banner() - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/examples/py310/doc-examples-python-apis.py b/examples/py310/doc-examples-python-apis.py index 16fd6ffb..2094b9bf 100644 --- a/examples/py310/doc-examples-python-apis.py +++ b/examples/py310/doc-examples-python-apis.py @@ -5,13 +5,9 @@ CacheClientAsync, Configurations, CredentialProvider, - PreviewVectorIndexClientAsync, TopicClientAsync, TopicConfigurations, - VectorIndexConfigurations, ) -from momento.requests.vector_index import ALL_METADATA, Field, Item -from momento.requests.vector_index import filters as F from momento.responses import ( CacheDelete, CacheGet, @@ -23,18 +19,6 @@ TopicSubscribe, TopicSubscriptionItem, ) -from momento.responses.vector_index import ( - CountItems, - CreateIndex, - DeleteIndex, - DeleteItemBatch, - GetItemBatch, - GetItemMetadataBatch, - ListIndexes, - Search, - SearchAndFetchVectors, - UpsertItemBatch, -) def example_API_CredentialProviderFromEnvVar(): @@ -183,213 +167,6 @@ async def example_API_TopicPublish(topic_client: TopicClientAsync): # end example -async def example_API_InstantiateVectorClient(): - PreviewVectorIndexClientAsync( - VectorIndexConfigurations.Default.latest(), CredentialProvider.from_environment_variable("MOMENTO_API_KEY") - ) - - -# end example - - -async def example_API_CreateIndex(vector_client: PreviewVectorIndexClientAsync): - response = await vector_client.create_index("test-index", 2) - match response: - case CreateIndex.Success(): - print("Index 'test-index' created") - case CreateIndex.IndexAlreadyExists(): - print("Index 'test-index' already exists") - case CreateIndex.Error() as error: - print(f"Error creating index 'test-index': {error.message}") - - -# end example - - -async def example_API_ListIndexes(vector_client: PreviewVectorIndexClientAsync): - response = await vector_client.list_indexes() - match response: - case ListIndexes.Success() as success: - print(f"Indexes:\n{success.indexes}") - case CreateIndex.Error() as error: - print(f"Error listing indexes: {error.message}") - - -# end example - - -async def example_API_DeleteIndex(vector_client: PreviewVectorIndexClientAsync): - response = await vector_client.delete_index("test-index") - match response: - case DeleteIndex.Success(): - print("Index 'test-index' deleted") - case DeleteIndex.Error() as error: - print(f"Error deleting index 'test-index': {error.message}") - - -# end example - - -async def example_API_CountItems(vector_client: PreviewVectorIndexClientAsync): - response = await vector_client.count_items("test-index") - match response: - case CountItems.Success() as success: - print(f"Found {success.item_count} items") - case CountItems.Error() as error: - print(f"Error counting items in index 'test-index': {error.message}") - - -# end example - - -async def example_API_UpsertItemBatch(vector_client: PreviewVectorIndexClientAsync): - response = await vector_client.upsert_item_batch( - "test-index", - [ - Item(id="example_item_1", vector=[1.0, 2.0], metadata={"key1": "value1"}), - Item(id="example_item_2", vector=[3.0, 4.0], metadata={"key2": "value2"}), - ], - ) - match response: - case UpsertItemBatch.Success(): - print("Successfully added items to index 'test-index'") - case UpsertItemBatch.Error() as error: - print(f"Error adding items to index 'test-index': {error.message}") - - -# end example - - -async def example_API_DeleteItemBatch(vector_client: PreviewVectorIndexClientAsync): - response = await vector_client.delete_item_batch("test-index", ["example_item_1", "example_item_2"]) - match response: - case DeleteItemBatch.Success(): - print("Successfully deleted items from index 'test-index'") - case DeleteItemBatch.Error() as error: - print(f"Error deleting items from index 'test-index': {error.message}") - - -# end example - - -async def example_API_GetItemBatch(vector_client: PreviewVectorIndexClientAsync): - response = await vector_client.get_item_batch("test-index", ["example_item_1", "example_item_2"]) - match response: - case GetItemBatch.Success() as success: - print(f"Found {len(success.values)} items") - case GetItemBatch.Error() as error: - print(f"Error getting items from index 'test-index': {error.message}") - - -# end example - - -async def example_API_GetItemMetadataBatch(vector_client: PreviewVectorIndexClientAsync): - response = await vector_client.get_item_metadata_batch("test-index", ["example_item_1", "example_item_2"]) - match response: - case GetItemMetadataBatch.Success() as success: - print(f"Found metadata for {len(success.values)} items") - case GetItemMetadataBatch.Error() as error: - print(f"Error getting item metadata from index 'test-index': {error.message}") - - -# end example - - -async def example_API_Search(vector_client: PreviewVectorIndexClientAsync): - response = await vector_client.search("test-index", [1.0, 2.0], top_k=3, metadata_fields=ALL_METADATA) - match response: - case Search.Success() as success: - print(f"Found {len(success.hits)} matches") - case Search.Error() as error: - print(f"Error searching index 'test-index': {error.message}") - - -# end example - - -async def example_API_SearchAndFetchVectors(vector_client: PreviewVectorIndexClientAsync): - response = await vector_client.search_and_fetch_vectors( - "test-index", [1.0, 2.0], top_k=3, metadata_fields=ALL_METADATA - ) - match response: - case SearchAndFetchVectors.Success() as success: - print(f"Found {len(success.hits)} matches") - case SearchAndFetchVectors.Error() as error: - print(f"Error searching index 'test-index': {error.message}") - - -# end example - - -def example_API_FilterExpressionOverview() -> None: - # For convenience, the filter expressions classes can accessed with filters module: - # from momento.requests.vector_index import filters as F - # - # You can use the Field class to create a more idiomatic filter expression by using the - # overloaded comparison operators: - # from momento.requests.vector_index import Field - # - # Below we demonstrate both approaches to creating filter expressions. - - # Is the movie titled "The Matrix"? - F.Equals("movie_title", "The Matrix") - Field("movie_title") == "The Matrix" - - # Is the movie not titled "The Matrix"? - F.Not(F.Equals("movie_title", "The Matrix")) - Field("movie_title") != "The Matrix" - - # Was the movie released in 1999? - F.Equals("year", 1999) - Field("year") == 1999 - - # Did the movie gross 463.5 million dollars? - F.Equals("gross_revenue_millions", 463.5) - Field("gross_revenue_millions") == 463.5 - - # Was the movie in theaters? - F.Equals("in_theaters", True) - Field("in_theaters") - - # Was the movie released after 1990? - F.GreaterThan("year", 1990) - Field("year") > 1990 - - # Was the movie released in or after 2020? - F.GreaterThanOrEqual("year", 2020) - Field("year") >= 2020 - - # Was the movie released before 2000? - F.LessThan("year", 2000) - Field("year") < 2000 - - # Was the movie released in or before 2000? - F.LessThanOrEqual("year", 2000) - Field("year") <= 2000 - - # Was "Keanu Reeves" one of the actors? - F.ListContains("actors", "Keanu Reeves") - Field("actors").list_contains("Keanu Reeves") - - # Is the ID one of the following? - F.IdInSet(["tt0133093", "tt0234215", "tt0242653"]) - - # Was the movie directed by "Lana Wachowski" and released after 2000? - F.And(F.ListContains("directors", "Lana Wachowski"), F.GreaterThan("year", 2000)) - Field("directors").list_contains("Lana Wachowski") & (Field("year") > 2000) - - # Was the movie directed by "Lana Wachowski" or released after 2000? - F.Or(F.ListContains("directors", "Lana Wachowski"), F.GreaterThan("year", 2000)) - Field("directors").list_contains("Lana Wachowski") | (Field("year") > 2000) - - # Was "Keanu Reeves" not one of the actors? - F.Not(F.ListContains("actors", "Keanu Reeves")) - - -# end example - - async def main(): example_API_CredentialProviderFromEnvVar() @@ -419,22 +196,6 @@ async def main(): await example_API_TopicSubscribe(topic_client) await topic_client.close() - vector_client = PreviewVectorIndexClientAsync( - VectorIndexConfigurations.Default.latest(), CredentialProvider.from_environment_variable("MOMENTO_API_KEY") - ) - await example_API_InstantiateVectorClient() - await example_API_CreateIndex(vector_client) - await example_API_ListIndexes(vector_client) - await example_API_CountItems(vector_client) - await example_API_UpsertItemBatch(vector_client) - await example_API_GetItemBatch(vector_client) - await example_API_GetItemMetadataBatch(vector_client) - await example_API_Search(vector_client) - await example_API_SearchAndFetchVectors(vector_client) - example_API_FilterExpressionOverview() - await example_API_DeleteItemBatch(vector_client) - await example_API_DeleteIndex(vector_client) - if __name__ == "__main__": asyncio.run(main()) diff --git a/examples/py310/vector_index.py b/examples/py310/vector_index.py deleted file mode 100644 index 2fb4bcfd..00000000 --- a/examples/py310/vector_index.py +++ /dev/null @@ -1,155 +0,0 @@ -import logging -from time import sleep - -from momento import ( - CredentialProvider, - PreviewVectorIndexClient, - VectorIndexConfigurations, -) -from momento.config import VectorIndexConfiguration -from momento.requests.vector_index import ALL_METADATA, Item, SimilarityMetric -from momento.responses.vector_index import ( - CreateIndex, - DeleteIndex, - DeleteItemBatch, - ListIndexes, - Search, - UpsertItemBatch, -) - -from example_utils.example_logging import initialize_logging - -_logger = logging.getLogger("vector-example") - -VECTOR_INDEX_CONFIGURATION: VectorIndexConfiguration = VectorIndexConfigurations.Default.latest() -VECTOR_AUTH_PROVIDER = CredentialProvider.from_environment_variable("MOMENTO_API_KEY") - - -def _print_start_banner() -> None: - _logger.info("******************************************************************") - _logger.info("* Momento Example Start *") - _logger.info("******************************************************************\n") - - -def _print_end_banner() -> None: - _logger.info("******************************************************************") - _logger.info("* Momento Example End *") - _logger.info("******************************************************************\n") - - -def create_index( - client: PreviewVectorIndexClient, - index_name: str, - num_dimensions: int, - similarity_metric: SimilarityMetric = SimilarityMetric.COSINE_SIMILARITY, -) -> None: - _logger.info(f"Creating index with name {index_name!r}") - create_index_response = client.create_index(index_name, num_dimensions, similarity_metric) - match create_index_response: - case CreateIndex.Success(): - _logger.info(f"Index with name {index_name!r} successfully created!") - case CreateIndex.IndexAlreadyExists(): - _logger.info(f"Index with name {index_name!r} already exists") - case CreateIndex.Error() as create_index_error: - _logger.error(f"Error while creating index {create_index_error.message}") - case _: - _logger.error("Unreachable") - _logger.info("") - - -def list_indexes(client: PreviewVectorIndexClient) -> None: - _logger.info("Listing indexes:") - list_indexes_response = client.list_indexes() - match list_indexes_response: - case ListIndexes.Success() as success: - for index in success.indexes: - _logger.info(f"- {index!r}") - case ListIndexes.Error() as list_indexes_error: - _logger.error(f"Error while listing indexes {list_indexes_error.message}") - case _: - _logger.error("Unreachable") - _logger.info("") - - -def upsert_items(client: PreviewVectorIndexClient, index_name: str) -> None: - items = [ - Item(id="test_item_1", vector=[1.0, 2.0], metadata={"key1": "value1"}), - Item(id="test_item_2", vector=[3.0, 4.0], metadata={"key2": "value2"}), - Item(id="test_item_3", vector=[5.0, 6.0], metadata={"key1": "value3", "key3": "value3"}), - ] - _logger.info(f"Adding items {items}") - upsert_response = client.upsert_item_batch( - index_name, - items=items, - ) - match upsert_response: - case UpsertItemBatch.Success(): - _logger.info("Successfully added items") - case UpsertItemBatch.Error() as upsert_error: - _logger.error(f"Error while adding items to index {index_name!r}: {upsert_error.message}") - case _: - _logger.error("Unreachable") - _logger.info("") - - -def search(client: PreviewVectorIndexClient, index_name: str) -> None: - query_vector = [1.0, 2.0] - top_k = 3 - _logger.info(f"Searching index {index_name} with query_vector {query_vector} and top {top_k} elements") - search_response = client.search(index_name, query_vector=query_vector, top_k=top_k, metadata_fields=ALL_METADATA) - match search_response: - case Search.Success() as success: - _logger.info(f"Search succeeded with {len(success.hits)} matches:") - _logger.info(success.hits) - case Search.Error() as search_error: - _logger.error(f"Error while searching on index {index_name}: {search_error.message}") - case _: - _logger.error("Unreachable") - _logger.info("") - - -def delete_items(client: PreviewVectorIndexClient, index_name: str) -> None: - item_ids_to_delete = ["test_item_1", "test_item_3"] - _logger.info(f"Deleting items: {item_ids_to_delete}") - delete_response = client.delete_item_batch(index_name, filter=item_ids_to_delete) - match delete_response: - case DeleteItemBatch.Success(): - _logger.info("Successfully deleted items") - case DeleteItemBatch.Error() as delete_error: - _logger.error(f"Error while deleting items {delete_error.message}") - case _: - _logger.error("Unreachable") - _logger.info("") - - -def delete_index(client: PreviewVectorIndexClient, index_name: str) -> None: - _logger.info("Deleting index " + index_name) - delete_response = client.delete_index(index_name) - - match delete_response: - case DeleteIndex.Success(): - _logger.info(f"Index {index_name} deleted successfully!") - case DeleteIndex.Error() as delete_error: - _logger.error(f"Failed to delete index {index_name} with error {delete_error.message}") - case _: - _logger.error("Unreachable") - _logger.info("") - - -if __name__ == "__main__": - initialize_logging() - _print_start_banner() - with PreviewVectorIndexClient(VECTOR_INDEX_CONFIGURATION, VECTOR_AUTH_PROVIDER) as client: - index_name = "hello_momento_index" - - create_index(client, index_name, num_dimensions=2) - list_indexes(client) - upsert_items(client, index_name) - sleep(2) - search(client, index_name) - delete_items(client, index_name) - sleep(2) - _logger.info("Deleted two items; search will return 1 hit now") - search(client, index_name) - delete_index(client, index_name) - _print_end_banner() diff --git a/examples/py310/vector_index_async.py b/examples/py310/vector_index_async.py deleted file mode 100644 index 28478023..00000000 --- a/examples/py310/vector_index_async.py +++ /dev/null @@ -1,162 +0,0 @@ -import asyncio -import logging -from time import sleep - -from momento import ( - CredentialProvider, - PreviewVectorIndexClientAsync, - VectorIndexConfigurations, -) -from momento.config import VectorIndexConfiguration -from momento.requests.vector_index import ALL_METADATA, Item, SimilarityMetric -from momento.responses.vector_index import ( - CreateIndex, - DeleteIndex, - DeleteItemBatch, - ListIndexes, - Search, - UpsertItemBatch, -) - -from example_utils.example_logging import initialize_logging - -_logger = logging.getLogger("vector-example") - -VECTOR_INDEX_CONFIGURATION: VectorIndexConfiguration = VectorIndexConfigurations.Default.latest() -VECTOR_AUTH_PROVIDER = CredentialProvider.from_environment_variable("MOMENTO_API_KEY") - - -def _print_start_banner() -> None: - _logger.info("******************************************************************") - _logger.info("* Momento Example Start *") - _logger.info("******************************************************************\n") - - -def _print_end_banner() -> None: - _logger.info("******************************************************************") - _logger.info("* Momento Example End *") - _logger.info("******************************************************************\n") - - -async def create_index( - client: PreviewVectorIndexClientAsync, - index_name: str, - num_dimensions: int, - similarity_metric: SimilarityMetric = SimilarityMetric.COSINE_SIMILARITY, -) -> None: - _logger.info(f"Creating index with name {index_name!r}") - create_index_response = await client.create_index(index_name, num_dimensions, similarity_metric) - match create_index_response: - case CreateIndex.Success(): - _logger.info(f"Index with name {index_name!r} successfully created!") - case CreateIndex.IndexAlreadyExists(): - _logger.info(f"Index with name {index_name!r} already exists") - case CreateIndex.Error() as create_index_error: - _logger.error(f"Error while creating index {create_index_error.message}") - case _: - _logger.error("Unreachable") - _logger.info("") - - -async def list_indexes(client: PreviewVectorIndexClientAsync) -> None: - _logger.info("Listing indexes:") - list_indexes_response = await client.list_indexes() - match list_indexes_response: - case ListIndexes.Success() as success: - for index in success.indexes: - _logger.info(f"- {index!r}") - case ListIndexes.Error() as list_indexes_error: - _logger.error(f"Error while listing indexes {list_indexes_error.message}") - case _: - _logger.error("Unreachable") - _logger.info("") - - -async def upsert_items(client: PreviewVectorIndexClientAsync, index_name: str) -> None: - items = [ - Item(id="test_item_1", vector=[1.0, 2.0], metadata={"key1": "value1"}), - Item(id="test_item_2", vector=[3.0, 4.0], metadata={"key2": "value2"}), - Item(id="test_item_3", vector=[5.0, 6.0], metadata={"key1": "value3", "key3": "value3"}), - ] - _logger.info(f"Adding items {items}") - upsert_response = await client.upsert_item_batch( - index_name, - items=items, - ) - match upsert_response: - case UpsertItemBatch.Success(): - _logger.info("Successfully added items") - case UpsertItemBatch.Error() as upsert_error: - _logger.error(f"Error while adding items to index {index_name!r}: {upsert_error.message}") - case _: - _logger.error("Unreachable") - _logger.info("") - - -async def search(client: PreviewVectorIndexClientAsync, index_name: str) -> None: - query_vector = [1.0, 2.0] - top_k = 3 - _logger.info(f"Searching index {index_name} with query_vector {query_vector} and top {top_k} elements") - search_response = await client.search( - index_name, query_vector=query_vector, top_k=top_k, metadata_fields=ALL_METADATA - ) - match search_response: - case Search.Success() as success: - _logger.info(f"Search succeeded with {len(success.hits)} matches:") - _logger.info(success.hits) - case Search.Error() as search_error: - _logger.error(f"Error while searching on index {index_name}: {search_error.message}") - case _: - _logger.error("Unreachable") - _logger.info("") - - -async def delete_items(client: PreviewVectorIndexClientAsync, index_name: str) -> None: - item_ids_to_delete = ["test_item_1", "test_item_3"] - _logger.info(f"Deleting items: {item_ids_to_delete}") - delete_response = await client.delete_item_batch(index_name, filter=item_ids_to_delete) - match delete_response: - case DeleteItemBatch.Success(): - _logger.info("Successfully deleted items") - case DeleteItemBatch.Error() as delete_error: - _logger.error(f"Error while deleting items {delete_error.message}") - case _: - _logger.error("Unreachable") - _logger.info("") - - -async def delete_index(client: PreviewVectorIndexClientAsync, index_name: str) -> None: - _logger.info("Deleting index " + index_name) - delete_response = await client.delete_index(index_name) - - match delete_response: - case DeleteIndex.Success(): - _logger.info(f"Index {index_name} deleted successfully!") - case DeleteIndex.Error() as delete_error: - _logger.error(f"Failed to delete index {index_name} with error {delete_error.message}") - case _: - _logger.error("Unreachable") - _logger.info("") - - -async def main() -> None: - initialize_logging() - _print_start_banner() - async with PreviewVectorIndexClientAsync(VECTOR_INDEX_CONFIGURATION, VECTOR_AUTH_PROVIDER) as client: - index_name = "hello_momento_index" - - await create_index(client, index_name, num_dimensions=2) - await list_indexes(client) - await upsert_items(client, index_name) - sleep(2) - await search(client, index_name) - await delete_items(client, index_name) - sleep(2) - _logger.info("Deleted two items; search will return 1 hit now") - await search(client, index_name) - await delete_index(client, index_name) - _print_end_banner() - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/pyproject.toml b/pyproject.toml index 9eb7fc08..8e84aa2e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -97,9 +97,6 @@ module = [ "momento.internal.synchronous._scs_control_client", "momento.internal.synchronous._scs_data_client", "momento.internal.synchronous._scs_grpc_manager", - "momento.internal.synchronous._vector_index_control_client", - "momento.internal.synchronous._vector_index_data_client", - "momento.internal.synchronous._vector_index_grpc_manager", "momento.internal.synchronous._add_header_client_interceptor", "momento.internal.synchronous._retry_interceptor", "momento.internal.common._data_client_ops", @@ -109,16 +106,8 @@ module = [ "momento.internal.aio._scs_control_client", "momento.internal.aio._scs_data_client", "momento.internal.aio._scs_grpc_manager", - "momento.internal.aio._vector_index_control_client", - "momento.internal.aio._vector_index_data_client", - "momento.internal.aio._vector_index_grpc_manager", "momento.internal.aio._utilities", "momento.responses.control.signing_key.*", - "momento.responses.vector_index.data.get_item_batch", - "momento.responses.vector_index.data.get_item_metadata_batch", - "momento.responses.vector_index.data.search", - "momento.responses.vector_index.data.search_and_fetch_vectors", - "momento.responses.vector_index.data.utils", ] disallow_any_expr = false @@ -164,8 +153,6 @@ fix = true [tool.ruff.per-file-ignores] "__init__.py" = ["F401"] -"src/momento/common_data/vector_index/item.py" = ["E721"] -"src/momento/requests/vector_index/filters.py" = ["E721"] [tool.ruff.lint.pydocstyle] convention = "google" diff --git a/src/momento/__init__.py b/src/momento/__init__.py index 690e8d53..213ca4e5 100644 --- a/src/momento/__init__.py +++ b/src/momento/__init__.py @@ -12,11 +12,9 @@ from .auth import CredentialProvider from .cache_client import CacheClient from .cache_client_async import CacheClientAsync -from .config import Configurations, TopicConfigurations, VectorIndexConfigurations +from .config import Configurations, TopicConfigurations from .topic_client import TopicClient from .topic_client_async import TopicClientAsync -from .vector_index_client import PreviewVectorIndexClient -from .vector_index_client_async import PreviewVectorIndexClientAsync logging.getLogger("momentosdk").addHandler(logging.NullHandler()) logs.initialize_momento_logging() @@ -25,11 +23,8 @@ "CredentialProvider", "Configurations", "TopicConfigurations", - "VectorIndexConfigurations", "CacheClient", "CacheClientAsync", "TopicClient", "TopicClientAsync", - "PreviewVectorIndexClient", - "PreviewVectorIndexClientAsync", ] diff --git a/src/momento/auth/credential_provider.py b/src/momento/auth/credential_provider.py index 34d965d8..dbc26666 100644 --- a/src/momento/auth/credential_provider.py +++ b/src/momento/auth/credential_provider.py @@ -15,14 +15,12 @@ class CredentialProvider: auth_token: str control_endpoint: str cache_endpoint: str - vector_endpoint: str @staticmethod def from_environment_variable( env_var_name: str, control_endpoint: Optional[str] = None, cache_endpoint: Optional[str] = None, - vector_endpoint: Optional[str] = None, ) -> CredentialProvider: """Reads and parses a Momento auth token stored as an environment variable. @@ -32,8 +30,6 @@ def from_environment_variable( Defaults to None. cache_endpoint (Optional[str], optional): Optionally overrides the default cache endpoint. Defaults to None. - vector_endpoint (Optional[str], optional): Optionally overrides the default vector endpoint. - Defaults to None. Raises: RuntimeError: if the environment variable is missing @@ -44,14 +40,13 @@ def from_environment_variable( api_key = os.getenv(env_var_name) if not api_key: raise RuntimeError(f"Missing required environment variable {env_var_name}") - return CredentialProvider.from_string(api_key, control_endpoint, cache_endpoint, vector_endpoint) + return CredentialProvider.from_string(api_key, control_endpoint, cache_endpoint) @staticmethod def from_string( auth_token: str, control_endpoint: Optional[str] = None, cache_endpoint: Optional[str] = None, - vector_endpoint: Optional[str] = None, ) -> CredentialProvider: """Reads and parses a Momento auth token. @@ -61,8 +56,6 @@ def from_string( Defaults to None. cache_endpoint (Optional[str], optional): Optionally overrides the default cache endpoint. Defaults to None. - vector_endpoint (Optional[str], optional): Optionally overrides the default vector endpoint. - Defaults to None. Returns: CredentialProvider @@ -70,9 +63,8 @@ def from_string( token_and_endpoints = momento_endpoint_resolver.resolve(auth_token) control_endpoint = control_endpoint or token_and_endpoints.control_endpoint cache_endpoint = cache_endpoint or token_and_endpoints.cache_endpoint - vector_endpoint = vector_endpoint or token_and_endpoints.vector_endpoint auth_token = token_and_endpoints.auth_token - return CredentialProvider(auth_token, control_endpoint, cache_endpoint, vector_endpoint) + return CredentialProvider(auth_token, control_endpoint, cache_endpoint) def __repr__(self) -> str: attributes: Dict[str, str] = copy.copy(vars(self)) # type: ignore[misc] diff --git a/src/momento/auth/momento_endpoint_resolver.py b/src/momento/auth/momento_endpoint_resolver.py index e3385bc2..03bae987 100644 --- a/src/momento/auth/momento_endpoint_resolver.py +++ b/src/momento/auth/momento_endpoint_resolver.py @@ -11,17 +11,14 @@ _MOMENTO_CONTROL_ENDPOINT_PREFIX = "control." _MOMENTO_CACHE_ENDPOINT_PREFIX = "cache." -_MOMENTO_VECTOR_ENDPOINT_PREFIX = "vector." _CONTROL_ENDPOINT_CLAIM_ID = "cp" _CACHE_ENDPOINT_CLAIM_ID = "c" -_VECTOR_ENDPOINT_CLAIM_ID = "c" # we don't have a new claim here so defaulting to c @dataclass class _TokenAndEndpoints: control_endpoint: str cache_endpoint: str - vector_endpoint: str auth_token: str @@ -41,7 +38,6 @@ def resolve(auth_token: str) -> _TokenAndEndpoints: return _TokenAndEndpoints( control_endpoint=_MOMENTO_CONTROL_ENDPOINT_PREFIX + info["endpoint"], # type: ignore[misc] cache_endpoint=_MOMENTO_CACHE_ENDPOINT_PREFIX + info["endpoint"], # type: ignore[misc] - vector_endpoint=_MOMENTO_VECTOR_ENDPOINT_PREFIX + info["endpoint"], # type: ignore[misc] auth_token=info["api_key"], # type: ignore[misc] ) else: @@ -54,7 +50,6 @@ def _get_endpoint_from_token(auth_token: str) -> _TokenAndEndpoints: return _TokenAndEndpoints( control_endpoint=claims[_CONTROL_ENDPOINT_CLAIM_ID], # type: ignore[misc] cache_endpoint=claims[_CACHE_ENDPOINT_CLAIM_ID], # type: ignore[misc] - vector_endpoint=claims[_VECTOR_ENDPOINT_CLAIM_ID], # type: ignore[misc] auth_token=auth_token, ) except (DecodeError, KeyError) as e: diff --git a/src/momento/common_data/vector_index/__init__.py b/src/momento/common_data/vector_index/__init__.py deleted file mode 100644 index 7ecf161c..00000000 --- a/src/momento/common_data/vector_index/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .item import Item, ItemMetadata, Metadata - -__all__ = ["Item", "ItemMetadata", "Metadata"] diff --git a/src/momento/common_data/vector_index/item.py b/src/momento/common_data/vector_index/item.py deleted file mode 100644 index 0f6482cb..00000000 --- a/src/momento/common_data/vector_index/item.py +++ /dev/null @@ -1,91 +0,0 @@ -from __future__ import annotations - -from typing import Dict, List, Optional, Union - -from momento_wire_types import vectorindex_pb2 as pb - -from momento.errors.exceptions import InvalidArgumentException -from momento.internal.services import Service - -Metadata = Dict[str, Union[str, int, float, bool, List[str]]] -"""The metadata of an item.""" - - -class ItemMetadata: - """Represents the id-metadata portion of an entry in the vector index. - - This is used in requests to update only the metadata of an item. - """ - - def __init__(self, id: str, metadata: Optional[Metadata] = None) -> None: - """Represents the id-metadata portion of an entry in the vector index. - - Args: - id (str): The id of the item. - metadata (Optional[Metadata], optional): The metadata of the item. Defaults to None, ie empty metadata. - """ - self.id = id - self.metadata = metadata or {} - - def __eq__(self, other: object) -> bool: - if isinstance(other, ItemMetadata): - return self.id == other.id and self.metadata == other.metadata - return False - - def __hash__(self) -> int: - return hash((self.id, self.metadata)) - - def __repr__(self) -> str: - return f"{self.__class__.__name__}(id={self.id!r}, metadata={self.metadata!r})" - - -# TODO: support other datatypes for the vector (np.array, pd.Series, etc.) -class Item(ItemMetadata): - """Represents a full entry in the vector index. - - This is used for `upsert_item_batch` requests. - """ - - def __init__(self, id: str, vector: list[float], metadata: Optional[Metadata] = None) -> None: - """Represents an entry in the vector index. - - Args: - id (str): The id of the item. - vector (list[float]): The vector of the item. - metadata (Optional[Metadata], optional): The metadata of the item. Defaults to None, ie empty metadata. - """ - super().__init__(id, metadata) - self.vector = vector - - def __eq__(self, other: object) -> bool: - if isinstance(other, Item): - return super().__eq__(other) and self.vector == other.vector - return False - - def __hash__(self) -> int: - return hash((self.id, self.vector, self.metadata)) - - def __repr__(self) -> str: - return f"{self.__class__.__name__}(id={self.id!r}, vector={self.vector!r}, metadata={self.metadata!r})" - - def to_proto(self) -> pb._Item: - vector = pb._Vector(elements=self.vector) - metadata = [] - for k, v in self.metadata.items(): - if type(v) is str: - metadata.append(pb._Metadata(field=k, string_value=v)) - elif type(v) is int: - metadata.append(pb._Metadata(field=k, integer_value=v)) - elif type(v) is float: - metadata.append(pb._Metadata(field=k, double_value=v)) - elif type(v) is bool: - metadata.append(pb._Metadata(field=k, boolean_value=v)) - elif type(v) is list and all(type(x) is str for x in v): - list_of_strings = pb._Metadata._ListOfStrings(values=v) - metadata.append(pb._Metadata(field=k, list_of_strings_value=list_of_strings)) - else: - raise InvalidArgumentException( - f"Metadata values must be either str, int, float, bool, or list[str]. Field {k!r} has a value of type {type(v)!r} with value {v!r}.", # noqa: E501 - Service.INDEX, - ) - return pb._Item(id=self.id, vector=vector, metadata=metadata) diff --git a/src/momento/config/__init__.py b/src/momento/config/__init__.py index b60b4483..07b716d5 100644 --- a/src/momento/config/__init__.py +++ b/src/momento/config/__init__.py @@ -4,14 +4,10 @@ from .configurations import Configurations from .topic_configuration import TopicConfiguration from .topic_configurations import TopicConfigurations -from .vector_index_configuration import VectorIndexConfiguration -from .vector_index_configurations import VectorIndexConfigurations __all__ = [ "Configuration", "Configurations", "TopicConfiguration", "TopicConfigurations", - "VectorIndexConfiguration", - "VectorIndexConfigurations", ] diff --git a/src/momento/config/vector_index_configuration.py b/src/momento/config/vector_index_configuration.py deleted file mode 100644 index f311ff0b..00000000 --- a/src/momento/config/vector_index_configuration.py +++ /dev/null @@ -1,83 +0,0 @@ -from __future__ import annotations - -from abc import ABC, abstractmethod -from datetime import timedelta -from pathlib import Path - -from .transport.transport_strategy import TransportStrategy - - -class VectorIndexConfigurationBase(ABC): - @abstractmethod - def get_transport_strategy(self) -> TransportStrategy: - pass - - @abstractmethod - def with_transport_strategy(self, transport_strategy: TransportStrategy) -> VectorIndexConfigurationBase: - pass - - @abstractmethod - def with_client_timeout(self, client_timeout: timedelta) -> VectorIndexConfigurationBase: - pass - - @abstractmethod - def with_root_certificates_pem(self, root_certificate_path: Path) -> VectorIndexConfigurationBase: - pass - - -class VectorIndexConfiguration(VectorIndexConfigurationBase): - """Configuration options for Momento Vector Index Client.""" - - def __init__(self, transport_strategy: TransportStrategy): - """Instantiate a Configuration. - - Args: - transport_strategy (TransportStrategy): Configuration options for networking with - the Momento service. - """ - self._transport_strategy = transport_strategy - - def get_transport_strategy(self) -> TransportStrategy: - """Access the transport strategy. - - Returns: - TransportStrategy: the current configuration options for wire interactions with the Momento service. - """ - return self._transport_strategy - - def with_transport_strategy(self, transport_strategy: TransportStrategy) -> VectorIndexConfiguration: - """Copy constructor for overriding TransportStrategy. - - Args: - transport_strategy (TransportStrategy): the new TransportStrategy. - - Returns: - Configuration: the new Configuration with the specified TransportStrategy. - """ - return VectorIndexConfiguration(transport_strategy) - - def with_client_timeout(self, client_timeout: timedelta) -> VectorIndexConfiguration: - """Copies the Configuration and sets the new client-side timeout in the copy's TransportStrategy. - - Args: - client_timeout (timedelta): the new client-side timeout. - - Return: - Configuration: the new Configuration. - """ - return self.with_transport_strategy(self._transport_strategy.with_client_timeout(client_timeout)) - - def with_root_certificates_pem(self, root_certificates_pem_path: Path) -> VectorIndexConfiguration: - """Copies the Configuration and sets the new root certificates in the copy's TransportStrategy. - - Args: - root_certificates_pem_path (Path): the new root certificates. - - Returns: - VectorIndexConfigurationBase: the new Configuration. - """ - grpc_configuration = self._transport_strategy.get_grpc_configuration().with_root_certificates_pem( - root_certificates_pem_path - ) - transport_strategy = self._transport_strategy.with_grpc_configuration(grpc_configuration) - return self.with_transport_strategy(transport_strategy) diff --git a/src/momento/config/vector_index_configurations.py b/src/momento/config/vector_index_configurations.py deleted file mode 100644 index b0e9ee89..00000000 --- a/src/momento/config/vector_index_configurations.py +++ /dev/null @@ -1,30 +0,0 @@ -from __future__ import annotations - -from datetime import timedelta - -from .transport.transport_strategy import ( - StaticGrpcConfiguration, - StaticTransportStrategy, -) -from .vector_index_configuration import VectorIndexConfiguration - - -class VectorIndexConfigurations: - """Container class for pre-built configurations.""" - - class Default(VectorIndexConfiguration): - """Laptop config provides defaults suitable for a medium-to-high-latency dev environment. - - Permissive timeouts, retries, and relaxed latency and throughput targets. - """ - - @staticmethod - def latest() -> VectorIndexConfigurations.Default: - """Provides the latest recommended configuration for a laptop development environment. - - This configuration will be updated every time there is a new version of the laptop configuration. - """ - return VectorIndexConfigurations.Default( - # The deadline is high to account for time-intensive adds. - StaticTransportStrategy(StaticGrpcConfiguration(deadline=timedelta(seconds=120))) - ) diff --git a/src/momento/internal/_utilities/__init__.py b/src/momento/internal/_utilities/__init__.py index df180a98..3937cecb 100644 --- a/src/momento/internal/_utilities/__init__.py +++ b/src/momento/internal/_utilities/__init__.py @@ -15,8 +15,3 @@ ) from ._momento_version import momento_version from ._time import _timedelta_to_ms -from ._vector_index_validation import ( - _validate_index_name, - _validate_num_dimensions, - _validate_top_k, -) diff --git a/src/momento/internal/_utilities/_channel_credentials.py b/src/momento/internal/_utilities/_channel_credentials.py index 567c0f64..d4fbc1b5 100644 --- a/src/momento/internal/_utilities/_channel_credentials.py +++ b/src/momento/internal/_utilities/_channel_credentials.py @@ -6,11 +6,11 @@ import grpc -from momento.config import Configuration, VectorIndexConfiguration +from momento.config import Configuration def channel_credentials_from_root_certs_or_default( - config: Configuration | VectorIndexConfiguration, + config: Configuration, ) -> grpc.ChannelCredentials: """Create gRPC channel credentials from the root certificates or the default credentials. diff --git a/src/momento/internal/_utilities/_vector_index_validation.py b/src/momento/internal/_utilities/_vector_index_validation.py deleted file mode 100644 index 40f4a3de..00000000 --- a/src/momento/internal/_utilities/_vector_index_validation.py +++ /dev/null @@ -1,20 +0,0 @@ -from __future__ import annotations - -from momento.errors import InvalidArgumentException -from momento.internal.services import Service - -from ._data_validation import _validate_name - - -def _validate_index_name(cache_name: str) -> None: - _validate_name(cache_name, "Vector index name", Service.INDEX) - - -def _validate_num_dimensions(num_dimensions: int) -> None: - if num_dimensions < 1 or not isinstance(num_dimensions, int): - raise InvalidArgumentException("Number of dimensions must be a positive integer.", Service.INDEX) - - -def _validate_top_k(top_k: int) -> None: - if top_k < 1 or not isinstance(top_k, int): - raise InvalidArgumentException("Top k must be a positive integer.", Service.INDEX) diff --git a/src/momento/internal/aio/_vector_index_control_client.py b/src/momento/internal/aio/_vector_index_control_client.py deleted file mode 100644 index a9e2d516..00000000 --- a/src/momento/internal/aio/_vector_index_control_client.py +++ /dev/null @@ -1,106 +0,0 @@ -import grpc -from momento_wire_types import controlclient_pb2 as ctrl_pb -from momento_wire_types import controlclient_pb2_grpc as ctrl_grpc - -from momento import logs -from momento.auth import CredentialProvider -from momento.config import VectorIndexConfiguration -from momento.errors import InvalidArgumentException, convert_error -from momento.internal._utilities import _validate_index_name, _validate_num_dimensions -from momento.internal.aio._vector_index_grpc_manager import ( - _VectorIndexControlGrpcManager, -) -from momento.internal.services import Service -from momento.requests.vector_index import SimilarityMetric -from momento.responses.vector_index import ( - CreateIndex, - CreateIndexResponse, - DeleteIndex, - DeleteIndexResponse, - ListIndexes, - ListIndexesResponse, -) - -_DEADLINE_SECONDS = 60.0 # 1 minute - - -class _VectorIndexControlClient: - """Momento Internal.""" - - def __init__(self, configuration: VectorIndexConfiguration, credential_provider: CredentialProvider): - endpoint = credential_provider.control_endpoint - self._logger = logs.logger - self._logger.debug("Vector index control client instantiated with endpoint: %s", endpoint) - self._grpc_manager = _VectorIndexControlGrpcManager(configuration, credential_provider) - self._endpoint = endpoint - - @property - def endpoint(self) -> str: - return self._endpoint - - async def create_index( - self, index_name: str, num_dimensions: int, similarity_metric: SimilarityMetric - ) -> CreateIndexResponse: - try: - self._logger.info(f"Creating index with name: {index_name}") - _validate_index_name(index_name) - _validate_num_dimensions(num_dimensions) - - if similarity_metric == SimilarityMetric.EUCLIDEAN_SIMILARITY: - request = ctrl_pb._CreateIndexRequest( - index_name=index_name, - num_dimensions=num_dimensions, - similarity_metric=ctrl_pb._SimilarityMetric( - euclidean_similarity=ctrl_pb._SimilarityMetric._EuclideanSimilarity() - ), - ) - elif similarity_metric == SimilarityMetric.INNER_PRODUCT: - request = ctrl_pb._CreateIndexRequest( - index_name=index_name, - num_dimensions=num_dimensions, - similarity_metric=ctrl_pb._SimilarityMetric( - inner_product=ctrl_pb._SimilarityMetric._InnerProduct() - ), - ) - elif similarity_metric == SimilarityMetric.COSINE_SIMILARITY: - request = ctrl_pb._CreateIndexRequest( - index_name=index_name, - num_dimensions=num_dimensions, - similarity_metric=ctrl_pb._SimilarityMetric( - cosine_similarity=ctrl_pb._SimilarityMetric._CosineSimilarity() - ), - ) - else: - raise InvalidArgumentException(f"Invalid similarity metric `{similarity_metric}`", Service.INDEX) - await self._build_stub().CreateIndex(request, timeout=_DEADLINE_SECONDS) - except Exception as e: - self._logger.debug("Failed to create index: %s with exception: %s", index_name, e) - if isinstance(e, grpc.RpcError) and e.code() == grpc.StatusCode.ALREADY_EXISTS: - return CreateIndex.IndexAlreadyExists() - return CreateIndex.Error(convert_error(e, Service.INDEX)) - return CreateIndex.Success() - - async def delete_index(self, index_name: str) -> DeleteIndexResponse: - try: - self._logger.info(f"Deleting index with name: {index_name}") - _validate_index_name(index_name) - request = ctrl_pb._DeleteIndexRequest(index_name=index_name) - await self._build_stub().DeleteIndex(request, timeout=_DEADLINE_SECONDS) - except Exception as e: - self._logger.debug("Failed to delete index: %s with exception: %s", index_name, e) - return DeleteIndex.Error(convert_error(e, Service.INDEX)) - return DeleteIndex.Success() - - async def list_indexes(self) -> ListIndexesResponse: - try: - list_indexes_request = ctrl_pb._ListIndexesRequest() - response = await self._build_stub().ListIndexes(list_indexes_request, timeout=_DEADLINE_SECONDS) - return ListIndexes.Success.from_grpc_response(response) - except Exception as e: - return ListIndexes.Error(convert_error(e, Service.INDEX)) - - def _build_stub(self) -> ctrl_grpc.ScsControlStub: - return self._grpc_manager.async_stub() - - async def close(self) -> None: - await self._grpc_manager.close() diff --git a/src/momento/internal/aio/_vector_index_data_client.py b/src/momento/internal/aio/_vector_index_data_client.py deleted file mode 100644 index 1bacecb4..00000000 --- a/src/momento/internal/aio/_vector_index_data_client.py +++ /dev/null @@ -1,297 +0,0 @@ -from __future__ import annotations - -from datetime import timedelta -from typing import Optional - -from momento_wire_types import vectorindex_pb2 as vectorindex_pb -from momento_wire_types import vectorindex_pb2_grpc as vectorindex_grpc - -from momento import logs -from momento.auth import CredentialProvider -from momento.config import VectorIndexConfiguration -from momento.errors import convert_error -from momento.internal._utilities import _validate_index_name, _validate_top_k -from momento.internal.aio._vector_index_grpc_manager import _VectorIndexDataGrpcManager -from momento.internal.services import Service -from momento.requests.vector_index import AllMetadata, FilterExpression, Item -from momento.requests.vector_index import filters as F -from momento.responses.vector_index import ( - CountItems, - CountItemsResponse, - DeleteItemBatch, - DeleteItemBatchResponse, - GetItemBatch, - GetItemBatchResponse, - GetItemMetadataBatch, - GetItemMetadataBatchResponse, - Search, - SearchAndFetchVectors, - SearchAndFetchVectorsHit, - SearchAndFetchVectorsResponse, - SearchHit, - SearchResponse, - UpsertItemBatch, - UpsertItemBatchResponse, -) - - -class _VectorIndexDataClient: - """Internal vector index data client.""" - - def __init__(self, configuration: VectorIndexConfiguration, credential_provider: CredentialProvider): - endpoint = credential_provider.vector_endpoint - self._logger = logs.logger - self._logger.debug("Vector index data client instantiated with endpoint: %s", endpoint) - self._endpoint = endpoint - - default_deadline: timedelta = configuration.get_transport_strategy().get_grpc_configuration().get_deadline() - self._default_deadline_seconds = default_deadline.total_seconds() - - self._grpc_manager = _VectorIndexDataGrpcManager(configuration, credential_provider) - - @property - def endpoint(self) -> str: - return self._endpoint - - async def count_items( - self, - index_name: str, - ) -> CountItemsResponse: - try: - self._log_issuing_request("CountItems", {"index_name": index_name}) - _validate_index_name(index_name) - - request = vectorindex_pb._CountItemsRequest( - index_name=index_name, all=vectorindex_pb._CountItemsRequest.All() - ) - response: vectorindex_pb._CountItemsResponse = await self._build_stub().CountItems( - request, timeout=self._default_deadline_seconds - ) - - self._log_received_response("CountItems", {"index_name": index_name}) - return CountItems.Success(item_count=response.item_count) - except Exception as e: - self._log_request_error("count_items", e) - return CountItems.Error(convert_error(e, Service.INDEX)) - - async def upsert_item_batch( - self, - index_name: str, - items: list[Item], - ) -> UpsertItemBatchResponse: - try: - self._log_issuing_request("UpsertItemBatch", {"index_name": index_name}) - _validate_index_name(index_name) - request = vectorindex_pb._UpsertItemBatchRequest( - index_name=index_name, - items=[item.to_proto() for item in items], - ) - - await self._build_stub().UpsertItemBatch(request, timeout=self._default_deadline_seconds) - - self._log_received_response("UpsertItemBatch", {"index_name": index_name}) - return UpsertItemBatch.Success() - except Exception as e: - self._log_request_error("set", e) - return UpsertItemBatch.Error(convert_error(e, Service.INDEX)) - - async def delete_item_batch( - self, - index_name: str, - filter: FilterExpression | list[str], - ) -> DeleteItemBatchResponse: - try: - self._log_issuing_request("DeleteItemBatch", {"index_name": index_name}) - _validate_index_name(index_name) - - filter_expression: vectorindex_pb._FilterExpression - - if isinstance(filter, FilterExpression): - filter_expression = filter.to_filter_expression_proto() - else: - if len(filter) == 0: - return DeleteItemBatch.Success() - filter_expression = F.IdInSet(filter).to_filter_expression_proto() - - request = vectorindex_pb._DeleteItemBatchRequest(index_name=index_name, filter=filter_expression) - - await self._build_stub().DeleteItemBatch(request, timeout=self._default_deadline_seconds) - - self._log_received_response("DeleteItemBatch", {"index_name": index_name}) - return DeleteItemBatch.Success() - except Exception as e: - self._log_request_error("delete", e) - return DeleteItemBatch.Error(convert_error(e, Service.INDEX)) - - @staticmethod - def __build_metadata_request(metadata_fields: Optional[list[str]] | AllMetadata) -> vectorindex_pb._MetadataRequest: - if isinstance(metadata_fields, AllMetadata): - return vectorindex_pb._MetadataRequest(all=vectorindex_pb._MetadataRequest.All()) - else: - return vectorindex_pb._MetadataRequest( - some=vectorindex_pb._MetadataRequest.Some(fields=metadata_fields if metadata_fields else []) - ) - - @staticmethod - def __build_no_score_threshold(score_threshold: Optional[float]) -> Optional[vectorindex_pb._NoScoreThreshold]: - if score_threshold is None: - return vectorindex_pb._NoScoreThreshold() - return None - - @staticmethod - def __build_filter_expression( - filter_expression: Optional[FilterExpression], - ) -> Optional[vectorindex_pb._FilterExpression]: - if filter_expression is None: - return None - return filter_expression.to_filter_expression_proto() - - async def search( - self, - index_name: str, - query_vector: list[float], - top_k: int, - metadata_fields: Optional[list[str]] | AllMetadata = None, - score_threshold: Optional[float] = None, - filter: Optional[FilterExpression] = None, - ) -> SearchResponse: - try: - self._log_issuing_request("Search", {"index_name": index_name}) - _validate_index_name(index_name) - _validate_top_k(top_k) - - query_vector_pb = vectorindex_pb._Vector(elements=query_vector) - metadata_fields_pb = _VectorIndexDataClient.__build_metadata_request(metadata_fields) - no_score_threshold = _VectorIndexDataClient.__build_no_score_threshold(score_threshold) - filter_expression_pb = _VectorIndexDataClient.__build_filter_expression(filter) - - request = vectorindex_pb._SearchRequest( - index_name=index_name, - query_vector=query_vector_pb, - top_k=top_k, - metadata_fields=metadata_fields_pb, - score_threshold=score_threshold, - no_score_threshold=no_score_threshold, - filter=filter_expression_pb, - ) - - response: vectorindex_pb._SearchResponse = await self._build_stub().Search( - request, timeout=self._default_deadline_seconds - ) - - hits = [SearchHit.from_proto(hit) for hit in response.hits] - self._log_received_response("Search", {"index_name": index_name}) - return Search.Success(hits=hits) - except Exception as e: - self._log_request_error("search", e) - return Search.Error(convert_error(e, Service.INDEX)) - - async def search_and_fetch_vectors( - self, - index_name: str, - query_vector: list[float], - top_k: int, - metadata_fields: Optional[list[str]] | AllMetadata = None, - score_threshold: Optional[float] = None, - filter: Optional[FilterExpression] = None, - ) -> SearchAndFetchVectorsResponse: - try: - self._log_issuing_request("SearchAndFetchVectors", {"index_name": index_name}) - _validate_index_name(index_name) - _validate_top_k(top_k) - - query_vector_pb = vectorindex_pb._Vector(elements=query_vector) - metadata_fields_pb = _VectorIndexDataClient.__build_metadata_request(metadata_fields) - no_score_threshold = _VectorIndexDataClient.__build_no_score_threshold(score_threshold) - filter_expression_pb = _VectorIndexDataClient.__build_filter_expression(filter) - - request = vectorindex_pb._SearchAndFetchVectorsRequest( - index_name=index_name, - query_vector=query_vector_pb, - top_k=top_k, - metadata_fields=metadata_fields_pb, - score_threshold=score_threshold, - no_score_threshold=no_score_threshold, - filter=filter_expression_pb, - ) - - response: vectorindex_pb._SearchAndFetchVectorsResponse = await self._build_stub().SearchAndFetchVectors( - request, timeout=self._default_deadline_seconds - ) - - hits = [SearchAndFetchVectorsHit.from_proto(hit) for hit in response.hits] - self._log_received_response("SearchAndFetchVectors", {"index_name": index_name}) - return SearchAndFetchVectors.Success(hits=hits) - except Exception as e: - self._log_request_error("search_and_fetch_vectors", e) - return SearchAndFetchVectors.Error(convert_error(e, Service.INDEX)) - - async def get_item_batch( - self, - index_name: str, - filter: list[str], - ) -> GetItemBatchResponse: - try: - self._log_issuing_request("GetItemBatch", {"index_name": index_name}) - _validate_index_name(index_name) - - if len(filter) == 0: - return GetItemBatch.Success(values={}) - - request = vectorindex_pb._GetItemBatchRequest( - index_name=index_name, - filter=F.IdInSet(filter).to_filter_expression_proto(), - metadata_fields=vectorindex_pb._MetadataRequest(all=vectorindex_pb._MetadataRequest.All()), - ) - - batch_response: vectorindex_pb._GetItemBatchResponse = await self._build_stub().GetItemBatch( - request, timeout=self._default_deadline_seconds - ) - self._log_received_response("GetItemBatch", {"index_name": index_name}) - return GetItemBatch.Success.from_proto(batch_response) - except Exception as e: - self._log_request_error("get_item_batch", e) - return GetItemBatch.Error(convert_error(e, Service.INDEX)) - - async def get_item_metadata_batch( - self, - index_name: str, - filter: list[str], - ) -> GetItemMetadataBatchResponse: - try: - self._log_issuing_request("GetItemMetadataBatch", {"index_name": index_name}) - _validate_index_name(index_name) - - if len(filter) == 0: - return GetItemMetadataBatch.Success(values={}) - - request = vectorindex_pb._GetItemMetadataBatchRequest( - index_name=index_name, - filter=F.IdInSet(filter).to_filter_expression_proto(), - metadata_fields=vectorindex_pb._MetadataRequest(all=vectorindex_pb._MetadataRequest.All()), - ) - - batch_response: vectorindex_pb._GetItemMetadataBatchResponse = ( - await self._build_stub().GetItemMetadataBatch(request, timeout=self._default_deadline_seconds) - ) - self._log_received_response("GetItemMetadataBatch", {"index_name": index_name}) - return GetItemMetadataBatch.Success.from_proto(batch_response) - except Exception as e: - self._log_request_error("get_item_metadata_batch", e) - return GetItemMetadataBatch.Error(convert_error(e, Service.INDEX)) - - # TODO these were copied from the data client. Shouldn't use interpolation here for perf? - def _log_received_response(self, request_type: str, request_args: dict[str, str]) -> None: - self._logger.log(logs.TRACE, f"Received a {request_type} response for {request_args}") - - def _log_issuing_request(self, request_type: str, request_args: dict[str, str]) -> None: - self._logger.log(logs.TRACE, f"Issuing a {request_type} request with {request_args}") - - def _log_request_error(self, request_type: str, e: Exception) -> None: - self._logger.warning(f"{request_type} failed with exception: {e}") - - def _build_stub(self) -> vectorindex_grpc.VectorIndexStub: - return self._grpc_manager.async_stub() - - async def close(self) -> None: - await self._grpc_manager.close() diff --git a/src/momento/internal/aio/_vector_index_grpc_manager.py b/src/momento/internal/aio/_vector_index_grpc_manager.py deleted file mode 100644 index 0de125c3..00000000 --- a/src/momento/internal/aio/_vector_index_grpc_manager.py +++ /dev/null @@ -1,75 +0,0 @@ -from __future__ import annotations - -import grpc -from momento_wire_types import controlclient_pb2_grpc as control_client -from momento_wire_types import vectorindex_pb2_grpc as vector_index_client - -from momento.auth import CredentialProvider -from momento.config import VectorIndexConfiguration -from momento.internal._utilities import momento_version -from momento.internal._utilities._channel_credentials import ( - channel_credentials_from_root_certs_or_default, -) -from momento.internal._utilities._grpc_channel_options import ( - grpc_control_channel_options_from_grpc_config, - grpc_data_channel_options_from_grpc_config, -) - -from ._add_header_client_interceptor import AddHeaderClientInterceptor, Header - - -class _VectorIndexControlGrpcManager: - """Internal gRPC control manager.""" - - version = momento_version - - def __init__(self, configuration: VectorIndexConfiguration, credential_provider: CredentialProvider): - self._secure_channel = grpc.aio.secure_channel( - target=credential_provider.control_endpoint, - credentials=channel_credentials_from_root_certs_or_default(configuration), - interceptors=_interceptors(credential_provider.auth_token), - options=grpc_control_channel_options_from_grpc_config( - grpc_config=configuration.get_transport_strategy().get_grpc_configuration(), - ), - ) - - async def close(self) -> None: - await self._secure_channel.close() - - def async_stub(self) -> control_client.ScsControlStub: - return control_client.ScsControlStub(self._secure_channel) # type: ignore[no-untyped-call] - - -class _VectorIndexDataGrpcManager: - """Internal gRPC vector index data manager.""" - - version = momento_version - - def __init__(self, configuration: VectorIndexConfiguration, credential_provider: CredentialProvider): - self._secure_channel = grpc.aio.secure_channel( - target=credential_provider.vector_endpoint, - credentials=channel_credentials_from_root_certs_or_default(configuration), - interceptors=_interceptors(credential_provider.auth_token), - options=grpc_data_channel_options_from_grpc_config( - configuration.get_transport_strategy().get_grpc_configuration() - ), - ) - - async def close(self) -> None: - await self._secure_channel.close() - - def async_stub(self) -> vector_index_client.VectorIndexStub: - return vector_index_client.VectorIndexStub(self._secure_channel) # type: ignore[no-untyped-call] - - -def _interceptors(auth_token: str) -> list[grpc.aio.ClientInterceptor]: - headers = [ - Header("authorization", auth_token), - Header("agent", f"python:{_VectorIndexControlGrpcManager.version}"), - ] - return list( - filter( - None, - [AddHeaderClientInterceptor(headers)], - ) - ) diff --git a/src/momento/internal/codegen.py b/src/momento/internal/codegen.py index 42650c16..5cfcbb56 100644 --- a/src/momento/internal/codegen.py +++ b/src/momento/internal/codegen.py @@ -122,8 +122,8 @@ def transform(self, input_module: cst.Module) -> cst.Module: ("^async_(.*)", "\\1"), ("(.*?)_async_(.*)", "\\1_\\2"), ("(.*?)_async$", "\\1"), - ("^((?:Cache|(?:Preview)?VectorIndex)Client)Async$", "\\1"), - ("^(TUnique(?:Cache|VectorIndex)Name)Async$", "\\1"), + ("^((?:Cache)Client)Async$", "\\1"), + ("^(TUnique(?:Cache)Name)Async$", "\\1"), ("__aenter__", "__enter__"), ("__aexit__", "__exit__"), ("^aio$", "synchronous"), @@ -132,8 +132,8 @@ def transform(self, input_module: cst.Module) -> cst.Module: simple_string_replacements = SimpleStringReplacement( [ - ("((?:Cache|(?:Preview)?VectorIndex)Client)Async", "\\1"), - (r"(.*?)Async(\s+(?:Cache|Vector\s+Index)\s+Client.*?)", "\\1Synchronous\\2"), + ("((?:Cache)Client)Async", "\\1"), + (r"(.*?)Async(\s+(?:Cache)\s+Client.*?)", "\\1Synchronous\\2"), (r"(.*?)\bawait\s+(.*?)", "\\1\\2"), ] ) diff --git a/src/momento/internal/synchronous/_vector_index_control_client.py b/src/momento/internal/synchronous/_vector_index_control_client.py deleted file mode 100644 index f996fc4d..00000000 --- a/src/momento/internal/synchronous/_vector_index_control_client.py +++ /dev/null @@ -1,106 +0,0 @@ -import grpc -from momento_wire_types import controlclient_pb2 as ctrl_pb -from momento_wire_types import controlclient_pb2_grpc as ctrl_grpc - -from momento import logs -from momento.auth import CredentialProvider -from momento.config import VectorIndexConfiguration -from momento.errors import InvalidArgumentException, convert_error -from momento.internal._utilities import _validate_index_name, _validate_num_dimensions -from momento.internal.services import Service -from momento.internal.synchronous._vector_index_grpc_manager import ( - _VectorIndexControlGrpcManager, -) -from momento.requests.vector_index import SimilarityMetric -from momento.responses.vector_index import ( - CreateIndex, - CreateIndexResponse, - DeleteIndex, - DeleteIndexResponse, - ListIndexes, - ListIndexesResponse, -) - -_DEADLINE_SECONDS = 60.0 # 1 minute - - -class _VectorIndexControlClient: - """Momento Internal.""" - - def __init__(self, configuration: VectorIndexConfiguration, credential_provider: CredentialProvider): - endpoint = credential_provider.control_endpoint - self._logger = logs.logger - self._logger.debug("Vector index control client instantiated with endpoint: %s", endpoint) - self._grpc_manager = _VectorIndexControlGrpcManager(configuration, credential_provider) - self._endpoint = endpoint - - @property - def endpoint(self) -> str: - return self._endpoint - - def create_index( - self, index_name: str, num_dimensions: int, similarity_metric: SimilarityMetric - ) -> CreateIndexResponse: - try: - self._logger.info(f"Creating index with name: {index_name}") - _validate_index_name(index_name) - _validate_num_dimensions(num_dimensions) - - if similarity_metric == SimilarityMetric.EUCLIDEAN_SIMILARITY: - request = ctrl_pb._CreateIndexRequest( - index_name=index_name, - num_dimensions=num_dimensions, - similarity_metric=ctrl_pb._SimilarityMetric( - euclidean_similarity=ctrl_pb._SimilarityMetric._EuclideanSimilarity() - ), - ) - elif similarity_metric == SimilarityMetric.INNER_PRODUCT: - request = ctrl_pb._CreateIndexRequest( - index_name=index_name, - num_dimensions=num_dimensions, - similarity_metric=ctrl_pb._SimilarityMetric( - inner_product=ctrl_pb._SimilarityMetric._InnerProduct() - ), - ) - elif similarity_metric == SimilarityMetric.COSINE_SIMILARITY: - request = ctrl_pb._CreateIndexRequest( - index_name=index_name, - num_dimensions=num_dimensions, - similarity_metric=ctrl_pb._SimilarityMetric( - cosine_similarity=ctrl_pb._SimilarityMetric._CosineSimilarity() - ), - ) - else: - raise InvalidArgumentException(f"Invalid similarity metric `{similarity_metric}`", Service.INDEX) - self._build_stub().CreateIndex(request, timeout=_DEADLINE_SECONDS) - except Exception as e: - self._logger.debug("Failed to create index: %s with exception: %s", index_name, e) - if isinstance(e, grpc.RpcError) and e.code() == grpc.StatusCode.ALREADY_EXISTS: - return CreateIndex.IndexAlreadyExists() - return CreateIndex.Error(convert_error(e, Service.INDEX)) - return CreateIndex.Success() - - def delete_index(self, index_name: str) -> DeleteIndexResponse: - try: - self._logger.info(f"Deleting index with name: {index_name}") - _validate_index_name(index_name) - request = ctrl_pb._DeleteIndexRequest(index_name=index_name) - self._build_stub().DeleteIndex(request, timeout=_DEADLINE_SECONDS) - except Exception as e: - self._logger.debug("Failed to delete index: %s with exception: %s", index_name, e) - return DeleteIndex.Error(convert_error(e, Service.INDEX)) - return DeleteIndex.Success() - - def list_indexes(self) -> ListIndexesResponse: - try: - list_indexes_request = ctrl_pb._ListIndexesRequest() - response = self._build_stub().ListIndexes(list_indexes_request, timeout=_DEADLINE_SECONDS) - return ListIndexes.Success.from_grpc_response(response) - except Exception as e: - return ListIndexes.Error(convert_error(e, Service.INDEX)) - - def _build_stub(self) -> ctrl_grpc.ScsControlStub: - return self._grpc_manager.stub() - - def close(self) -> None: - self._grpc_manager.close() diff --git a/src/momento/internal/synchronous/_vector_index_data_client.py b/src/momento/internal/synchronous/_vector_index_data_client.py deleted file mode 100644 index 22088aa5..00000000 --- a/src/momento/internal/synchronous/_vector_index_data_client.py +++ /dev/null @@ -1,297 +0,0 @@ -from __future__ import annotations - -from datetime import timedelta -from typing import Optional - -from momento_wire_types import vectorindex_pb2 as vectorindex_pb -from momento_wire_types import vectorindex_pb2_grpc as vectorindex_grpc - -from momento import logs -from momento.auth import CredentialProvider -from momento.config import VectorIndexConfiguration -from momento.errors import convert_error -from momento.internal._utilities import _validate_index_name, _validate_top_k -from momento.internal.services import Service -from momento.internal.synchronous._vector_index_grpc_manager import _VectorIndexDataGrpcManager -from momento.requests.vector_index import AllMetadata, FilterExpression, Item -from momento.requests.vector_index import filters as F -from momento.responses.vector_index import ( - CountItems, - CountItemsResponse, - DeleteItemBatch, - DeleteItemBatchResponse, - GetItemBatch, - GetItemBatchResponse, - GetItemMetadataBatch, - GetItemMetadataBatchResponse, - Search, - SearchAndFetchVectors, - SearchAndFetchVectorsHit, - SearchAndFetchVectorsResponse, - SearchHit, - SearchResponse, - UpsertItemBatch, - UpsertItemBatchResponse, -) - - -class _VectorIndexDataClient: - """Internal vector index data client.""" - - def __init__(self, configuration: VectorIndexConfiguration, credential_provider: CredentialProvider): - endpoint = credential_provider.vector_endpoint - self._logger = logs.logger - self._logger.debug("Vector index data client instantiated with endpoint: %s", endpoint) - self._endpoint = endpoint - - default_deadline: timedelta = configuration.get_transport_strategy().get_grpc_configuration().get_deadline() - self._default_deadline_seconds = default_deadline.total_seconds() - - self._grpc_manager = _VectorIndexDataGrpcManager(configuration, credential_provider) - - @property - def endpoint(self) -> str: - return self._endpoint - - def count_items( - self, - index_name: str, - ) -> CountItemsResponse: - try: - self._log_issuing_request("CountItems", {"index_name": index_name}) - _validate_index_name(index_name) - - request = vectorindex_pb._CountItemsRequest( - index_name=index_name, all=vectorindex_pb._CountItemsRequest.All() - ) - response: vectorindex_pb._CountItemsResponse = self._build_stub().CountItems( - request, timeout=self._default_deadline_seconds - ) - - self._log_received_response("CountItems", {"index_name": index_name}) - return CountItems.Success(item_count=response.item_count) - except Exception as e: - self._log_request_error("count_items", e) - return CountItems.Error(convert_error(e, Service.INDEX)) - - def upsert_item_batch( - self, - index_name: str, - items: list[Item], - ) -> UpsertItemBatchResponse: - try: - self._log_issuing_request("UpsertItemBatch", {"index_name": index_name}) - _validate_index_name(index_name) - request = vectorindex_pb._UpsertItemBatchRequest( - index_name=index_name, - items=[item.to_proto() for item in items], - ) - - self._build_stub().UpsertItemBatch(request, timeout=self._default_deadline_seconds) - - self._log_received_response("UpsertItemBatch", {"index_name": index_name}) - return UpsertItemBatch.Success() - except Exception as e: - self._log_request_error("set", e) - return UpsertItemBatch.Error(convert_error(e, Service.INDEX)) - - def delete_item_batch( - self, - index_name: str, - filter: FilterExpression | list[str], - ) -> DeleteItemBatchResponse: - try: - self._log_issuing_request("DeleteItemBatch", {"index_name": index_name}) - _validate_index_name(index_name) - - filter_expression: vectorindex_pb._FilterExpression - - if isinstance(filter, FilterExpression): - filter_expression = filter.to_filter_expression_proto() - else: - if len(filter) == 0: - return DeleteItemBatch.Success() - filter_expression = F.IdInSet(filter).to_filter_expression_proto() - - request = vectorindex_pb._DeleteItemBatchRequest(index_name=index_name, filter=filter_expression) - - self._build_stub().DeleteItemBatch(request, timeout=self._default_deadline_seconds) - - self._log_received_response("DeleteItemBatch", {"index_name": index_name}) - return DeleteItemBatch.Success() - except Exception as e: - self._log_request_error("delete", e) - return DeleteItemBatch.Error(convert_error(e, Service.INDEX)) - - @staticmethod - def __build_metadata_request(metadata_fields: Optional[list[str]] | AllMetadata) -> vectorindex_pb._MetadataRequest: - if isinstance(metadata_fields, AllMetadata): - return vectorindex_pb._MetadataRequest(all=vectorindex_pb._MetadataRequest.All()) - else: - return vectorindex_pb._MetadataRequest( - some=vectorindex_pb._MetadataRequest.Some(fields=metadata_fields if metadata_fields else []) - ) - - @staticmethod - def __build_no_score_threshold(score_threshold: Optional[float]) -> Optional[vectorindex_pb._NoScoreThreshold]: - if score_threshold is None: - return vectorindex_pb._NoScoreThreshold() - return None - - @staticmethod - def __build_filter_expression( - filter_expression: Optional[FilterExpression], - ) -> Optional[vectorindex_pb._FilterExpression]: - if filter_expression is None: - return None - return filter_expression.to_filter_expression_proto() - - def search( - self, - index_name: str, - query_vector: list[float], - top_k: int, - metadata_fields: Optional[list[str]] | AllMetadata = None, - score_threshold: Optional[float] = None, - filter: Optional[FilterExpression] = None, - ) -> SearchResponse: - try: - self._log_issuing_request("Search", {"index_name": index_name}) - _validate_index_name(index_name) - _validate_top_k(top_k) - - query_vector_pb = vectorindex_pb._Vector(elements=query_vector) - metadata_fields_pb = _VectorIndexDataClient.__build_metadata_request(metadata_fields) - no_score_threshold = _VectorIndexDataClient.__build_no_score_threshold(score_threshold) - filter_expression_pb = _VectorIndexDataClient.__build_filter_expression(filter) - - request = vectorindex_pb._SearchRequest( - index_name=index_name, - query_vector=query_vector_pb, - top_k=top_k, - metadata_fields=metadata_fields_pb, - score_threshold=score_threshold, - no_score_threshold=no_score_threshold, - filter=filter_expression_pb, - ) - - response: vectorindex_pb._SearchResponse = self._build_stub().Search( - request, timeout=self._default_deadline_seconds - ) - - hits = [SearchHit.from_proto(hit) for hit in response.hits] - self._log_received_response("Search", {"index_name": index_name}) - return Search.Success(hits=hits) - except Exception as e: - self._log_request_error("search", e) - return Search.Error(convert_error(e, Service.INDEX)) - - def search_and_fetch_vectors( - self, - index_name: str, - query_vector: list[float], - top_k: int, - metadata_fields: Optional[list[str]] | AllMetadata = None, - score_threshold: Optional[float] = None, - filter: Optional[FilterExpression] = None, - ) -> SearchAndFetchVectorsResponse: - try: - self._log_issuing_request("SearchAndFetchVectors", {"index_name": index_name}) - _validate_index_name(index_name) - _validate_top_k(top_k) - - query_vector_pb = vectorindex_pb._Vector(elements=query_vector) - metadata_fields_pb = _VectorIndexDataClient.__build_metadata_request(metadata_fields) - no_score_threshold = _VectorIndexDataClient.__build_no_score_threshold(score_threshold) - filter_expression_pb = _VectorIndexDataClient.__build_filter_expression(filter) - - request = vectorindex_pb._SearchAndFetchVectorsRequest( - index_name=index_name, - query_vector=query_vector_pb, - top_k=top_k, - metadata_fields=metadata_fields_pb, - score_threshold=score_threshold, - no_score_threshold=no_score_threshold, - filter=filter_expression_pb, - ) - - response: vectorindex_pb._SearchAndFetchVectorsResponse = self._build_stub().SearchAndFetchVectors( - request, timeout=self._default_deadline_seconds - ) - - hits = [SearchAndFetchVectorsHit.from_proto(hit) for hit in response.hits] - self._log_received_response("SearchAndFetchVectors", {"index_name": index_name}) - return SearchAndFetchVectors.Success(hits=hits) - except Exception as e: - self._log_request_error("search_and_fetch_vectors", e) - return SearchAndFetchVectors.Error(convert_error(e, Service.INDEX)) - - def get_item_batch( - self, - index_name: str, - filter: list[str], - ) -> GetItemBatchResponse: - try: - self._log_issuing_request("GetItemBatch", {"index_name": index_name}) - _validate_index_name(index_name) - - if len(filter) == 0: - return GetItemBatch.Success(values={}) - - request = vectorindex_pb._GetItemBatchRequest( - index_name=index_name, - filter=F.IdInSet(filter).to_filter_expression_proto(), - metadata_fields=vectorindex_pb._MetadataRequest(all=vectorindex_pb._MetadataRequest.All()), - ) - - batch_response: vectorindex_pb._GetItemBatchResponse = self._build_stub().GetItemBatch( - request, timeout=self._default_deadline_seconds - ) - self._log_received_response("GetItemBatch", {"index_name": index_name}) - return GetItemBatch.Success.from_proto(batch_response) - except Exception as e: - self._log_request_error("get_item_batch", e) - return GetItemBatch.Error(convert_error(e, Service.INDEX)) - - def get_item_metadata_batch( - self, - index_name: str, - filter: list[str], - ) -> GetItemMetadataBatchResponse: - try: - self._log_issuing_request("GetItemMetadataBatch", {"index_name": index_name}) - _validate_index_name(index_name) - - if len(filter) == 0: - return GetItemMetadataBatch.Success(values={}) - - request = vectorindex_pb._GetItemMetadataBatchRequest( - index_name=index_name, - filter=F.IdInSet(filter).to_filter_expression_proto(), - metadata_fields=vectorindex_pb._MetadataRequest(all=vectorindex_pb._MetadataRequest.All()), - ) - - batch_response: vectorindex_pb._GetItemMetadataBatchResponse = self._build_stub().GetItemMetadataBatch( - request, timeout=self._default_deadline_seconds - ) - self._log_received_response("GetItemMetadataBatch", {"index_name": index_name}) - return GetItemMetadataBatch.Success.from_proto(batch_response) - except Exception as e: - self._log_request_error("get_item_metadata_batch", e) - return GetItemMetadataBatch.Error(convert_error(e, Service.INDEX)) - - # TODO these were copied from the data client. Shouldn't use interpolation here for perf? - def _log_received_response(self, request_type: str, request_args: dict[str, str]) -> None: - self._logger.log(logs.TRACE, f"Received a {request_type} response for {request_args}") - - def _log_issuing_request(self, request_type: str, request_args: dict[str, str]) -> None: - self._logger.log(logs.TRACE, f"Issuing a {request_type} request with {request_args}") - - def _log_request_error(self, request_type: str, e: Exception) -> None: - self._logger.warning(f"{request_type} failed with exception: {e}") - - def _build_stub(self) -> vectorindex_grpc.VectorIndexStub: - return self._grpc_manager.stub() - - def close(self) -> None: - self._grpc_manager.close() diff --git a/src/momento/internal/synchronous/_vector_index_grpc_manager.py b/src/momento/internal/synchronous/_vector_index_grpc_manager.py deleted file mode 100644 index 789d731a..00000000 --- a/src/momento/internal/synchronous/_vector_index_grpc_manager.py +++ /dev/null @@ -1,71 +0,0 @@ -from __future__ import annotations - -import grpc -from momento_wire_types import controlclient_pb2_grpc as control_client -from momento_wire_types import vectorindex_pb2_grpc as vector_index_client - -from momento.auth import CredentialProvider -from momento.config import VectorIndexConfiguration -from momento.internal._utilities import momento_version -from momento.internal._utilities._channel_credentials import ( - channel_credentials_from_root_certs_or_default, -) -from momento.internal._utilities._grpc_channel_options import ( - grpc_control_channel_options_from_grpc_config, - grpc_data_channel_options_from_grpc_config, -) -from momento.internal.synchronous._add_header_client_interceptor import ( - AddHeaderClientInterceptor, - Header, -) - - -class _VectorIndexControlGrpcManager: - """Internal gRPC control mananger.""" - - version = momento_version - - def __init__(self, configuration: VectorIndexConfiguration, credential_provider: CredentialProvider): - self._secure_channel = grpc.secure_channel( - target=credential_provider.control_endpoint, - credentials=channel_credentials_from_root_certs_or_default(configuration), - options=grpc_control_channel_options_from_grpc_config( - grpc_config=configuration.get_transport_strategy().get_grpc_configuration(), - ), - ) - intercept_channel = grpc.intercept_channel(self._secure_channel, *_interceptors(credential_provider.auth_token)) - self._stub = control_client.ScsControlStub(intercept_channel) # type: ignore[no-untyped-call] - - def close(self) -> None: - self._secure_channel.close() - - def stub(self) -> control_client.ScsControlStub: - return self._stub - - -class _VectorIndexDataGrpcManager: - """Internal gRPC vector index data manager.""" - - version = momento_version - - def __init__(self, configuration: VectorIndexConfiguration, credential_provider: CredentialProvider): - self._secure_channel = grpc.secure_channel( - target=credential_provider.vector_endpoint, - credentials=channel_credentials_from_root_certs_or_default(configuration), - options=grpc_data_channel_options_from_grpc_config( - configuration.get_transport_strategy().get_grpc_configuration() - ), - ) - intercept_channel = grpc.intercept_channel(self._secure_channel, *_interceptors(credential_provider.auth_token)) - self._stub = vector_index_client.VectorIndexStub(intercept_channel) # type: ignore[no-untyped-call] - - def close(self) -> None: - self._secure_channel.close() - - def stub(self) -> vector_index_client.VectorIndexStub: - return self._stub - - -def _interceptors(auth_token: str) -> list[grpc.UnaryUnaryClientInterceptor]: - headers = [Header("authorization", auth_token), Header("agent", f"python:{_VectorIndexControlGrpcManager.version}")] - return list(filter(None, [AddHeaderClientInterceptor(headers)])) diff --git a/src/momento/requests/vector_index/__init__.py b/src/momento/requests/vector_index/__init__.py deleted file mode 100644 index ef60c4cb..00000000 --- a/src/momento/requests/vector_index/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -from momento.common_data.vector_index.item import Item - -from .filter_field import Field -from .filters import FilterExpression -from .search import ALL_METADATA, AllMetadata -from .similarity_metric import SimilarityMetric - -__all__ = ["Item", "AllMetadata", "ALL_METADATA", "Field", "FilterExpression", "SimilarityMetric"] diff --git a/src/momento/requests/vector_index/filter_field.py b/src/momento/requests/vector_index/filter_field.py deleted file mode 100644 index bac5a042..00000000 --- a/src/momento/requests/vector_index/filter_field.py +++ /dev/null @@ -1,71 +0,0 @@ -"""Field class for filter expressions. - -This class is used to create filter expressions in a more readable way:: - - from momento.requests.vector_index.filters import Field - - Field("name") == "foo" - Field("age") >= 18 - Field("tags").list_contains("books") - (Field("year") > 2000) | (Field("year") < 1990) -""" - -from __future__ import annotations - -from dataclasses import dataclass - -from momento_wire_types import vectorindex_pb2 as vectorindex_pb - -from . import filters as F - - -@dataclass -class Field: - """Represents a field in a filter expression. - - Can be used to create filter expressions in a more readable way:: - - from momento.requests.vector_index.filters import Field - - Field("name") == "foo" - Field("age") >= 18 - Field("tags").list_contains("books") - (Field("year") > 2000) | (Field("year") < 1990) - """ - - name: str - """The name of the field.""" - - def to_filter_expression_proto(self) -> vectorindex_pb._FilterExpression: - """Converts the field to a protobuf filter expression. - - A bare field is equivalent to the field in a boolean context. - - Returns: - vectorindex_pb._FilterExpression: The protobuf filter expression. - """ - return F.Equals(self.name, True).to_filter_expression_proto() - - def __eq__(self, other: str | int | float | bool) -> F.Equals: # type: ignore - return F.Equals(self.name, other) - - def __invert__(self) -> F.Equals: - return F.Equals(self.name, False) - - def __ne__(self, other: str | int | float | bool) -> F.Not: # type: ignore - return F.Not(F.Equals(self.name, other)) - - def __gt__(self, other: int | float) -> F.GreaterThan: - return F.GreaterThan(self.name, other) - - def __ge__(self, other: int | float) -> F.GreaterThanOrEqual: - return F.GreaterThanOrEqual(self.name, other) - - def __lt__(self, other: int | float) -> F.LessThan: - return F.LessThan(self.name, other) - - def __le__(self, other: int | float) -> F.LessThanOrEqual: - return F.LessThanOrEqual(self.name, other) - - def list_contains(self, value: str) -> F.ListContains: - return F.ListContains(self.name, value) diff --git a/src/momento/requests/vector_index/filters.py b/src/momento/requests/vector_index/filters.py deleted file mode 100644 index bb8ced1b..00000000 --- a/src/momento/requests/vector_index/filters.py +++ /dev/null @@ -1,275 +0,0 @@ -"""Filter expressions for the vector index request. - -The filter expressions are used to filter the results of a vector index request. -This module contains the base class for all filter expressions, as well as -implementations for all the different types of filter expressions:: - And(Equals("foo", "bar"), GreaterThan("age", 18)) - Or(Equals("foo", "bar"), LessThan("age", 18)) - Not(Equals("foo", "bar")) - -The `Field` class is used to create filter expressions in a more idiomatic way. -""" - -from __future__ import annotations - -from abc import ABC, abstractmethod -from dataclasses import dataclass -from typing import Iterable - -from momento_wire_types import vectorindex_pb2 as vectorindex_pb - -from momento.errors.exceptions import InvalidArgumentException -from momento.internal.services import Service - - -@dataclass -class FilterExpression(ABC): - """Base class for all filter expressions.""" - - @abstractmethod - def to_filter_expression_proto(self) -> vectorindex_pb._FilterExpression: - """Converts the filter expression to a protobuf filter expression. - - Returns: - vectorindex_pb._FilterExpression: The protobuf filter expression. - """ - ... - - def __and__(self, other: FilterExpression) -> FilterExpression: - """Creates an AND expression between this expression and another. - - Args: - other (FilterExpression): The other expression to AND with. - - Returns: - FilterExpression: The AND expression. - """ - return And(self, other) - - def __or__(self, other: FilterExpression) -> FilterExpression: - """Creates an OR expression between this expression and another. - - Args: - other (FilterExpression): The other expression to OR with. - - Returns: - FilterExpression: The OR expression. - """ - return Or(self, other) - - def __invert__(self) -> FilterExpression: - """Creates a NOT expression of this expression. - - Returns: - FilterExpression: The NOT expression. - """ - return Not(self) - - -@dataclass -class And(FilterExpression): - """Represents an AND expression between two filter expressions.""" - - first_expression: FilterExpression - """The first expression to AND.""" - second_expression: FilterExpression - """The second expression to AND.""" - - def to_filter_expression_proto(self) -> vectorindex_pb._FilterExpression: - return vectorindex_pb._FilterExpression(and_expression=self.to_proto()) - - def to_proto(self) -> vectorindex_pb._AndExpression: - return vectorindex_pb._AndExpression( - first_expression=self.first_expression.to_filter_expression_proto(), - second_expression=self.second_expression.to_filter_expression_proto(), - ) - - -@dataclass -class Or(FilterExpression): - """Represents an OR expression between two filter expressions.""" - - first_expression: FilterExpression - """The first expression to OR.""" - second_expression: FilterExpression - """The second expression to OR.""" - - def to_filter_expression_proto(self) -> vectorindex_pb._FilterExpression: - return vectorindex_pb._FilterExpression(or_expression=self.to_proto()) - - def to_proto(self) -> vectorindex_pb._OrExpression: - return vectorindex_pb._OrExpression( - first_expression=self.first_expression.to_filter_expression_proto(), - second_expression=self.second_expression.to_filter_expression_proto(), - ) - - -@dataclass -class Not(FilterExpression): - """Represents a NOT expression of a filter expression.""" - - expression_to_negate: FilterExpression - """The expression to negate.""" - - def to_filter_expression_proto(self) -> vectorindex_pb._FilterExpression: - return vectorindex_pb._FilterExpression(not_expression=self.to_proto()) - - def to_proto(self) -> vectorindex_pb._NotExpression: - return vectorindex_pb._NotExpression( - expression_to_negate=self.expression_to_negate.to_filter_expression_proto() - ) - - -@dataclass -class Equals(FilterExpression): - """Represents an equals expression between a field and a value.""" - - field: str - """The field to compare.""" - value: str | int | float | bool - """The value to test equality with.""" - - def to_filter_expression_proto(self) -> vectorindex_pb._FilterExpression: - return vectorindex_pb._FilterExpression(equals_expression=self.to_proto()) - - def to_proto(self) -> vectorindex_pb._EqualsExpression: - if type(self.value) is str: - return vectorindex_pb._EqualsExpression(field=self.field, string_value=self.value) - elif type(self.value) is int: - return vectorindex_pb._EqualsExpression(field=self.field, integer_value=self.value) - elif type(self.value) is float: - return vectorindex_pb._EqualsExpression(field=self.field, float_value=self.value) - elif type(self.value) is bool: - return vectorindex_pb._EqualsExpression(field=self.field, boolean_value=self.value) - else: - raise InvalidArgumentException( - f"Invalid type for value: {type(self.value)} in equals expression. Must be one of str, int, float, bool.", - Service.INDEX, - ) - - -@dataclass -class GreaterThan(FilterExpression): - """Represents a greater than expression between a field and a value.""" - - field: str - """The field to compare.""" - value: int | float - """The value to test greater than with.""" - - def to_filter_expression_proto(self) -> vectorindex_pb._FilterExpression: - return vectorindex_pb._FilterExpression(greater_than_expression=self.to_proto()) - - def to_proto(self) -> vectorindex_pb._GreaterThanExpression: - if type(self.value) is int: - return vectorindex_pb._GreaterThanExpression(field=self.field, integer_value=self.value) - elif type(self.value) is float: - return vectorindex_pb._GreaterThanExpression(field=self.field, float_value=self.value) - else: - raise InvalidArgumentException( - f"Invalid type for value: {type(self.value)} in greater than expression. Must be one of int, float.", - Service.INDEX, - ) - - -@dataclass -class GreaterThanOrEqual(FilterExpression): - """Represents a greater than or equal expression between a field and a value.""" - - field: str - """The field to compare.""" - value: int | float - """The value to test greater than or equal with.""" - - def to_filter_expression_proto(self) -> vectorindex_pb._FilterExpression: - return vectorindex_pb._FilterExpression(greater_than_or_equal_expression=self.to_proto()) - - def to_proto(self) -> vectorindex_pb._GreaterThanOrEqualExpression: - if type(self.value) is int: - return vectorindex_pb._GreaterThanOrEqualExpression(field=self.field, integer_value=self.value) - elif type(self.value) is float: - return vectorindex_pb._GreaterThanOrEqualExpression(field=self.field, float_value=self.value) - else: - raise InvalidArgumentException( - f"Invalid type for value: {type(self.value)} in greater than or equal expression. Must be one of int, float.", - Service.INDEX, - ) - - -@dataclass -class LessThan(FilterExpression): - """Represents a less than expression between a field and a value.""" - - field: str - """The field to compare.""" - value: int | float - """The value to test less than with.""" - - def to_filter_expression_proto(self) -> vectorindex_pb._FilterExpression: - return vectorindex_pb._FilterExpression(less_than_expression=self.to_proto()) - - def to_proto(self) -> vectorindex_pb._LessThanExpression: - if type(self.value) is int: - return vectorindex_pb._LessThanExpression(field=self.field, integer_value=self.value) - elif type(self.value) is float: - return vectorindex_pb._LessThanExpression(field=self.field, float_value=self.value) - else: - raise InvalidArgumentException( - f"Invalid type for value: {type(self.value)} in less than expression. Must be one of int, float.", - Service.INDEX, - ) - - -@dataclass -class LessThanOrEqual(FilterExpression): - """Represents a less than or equal expression between a field and a value.""" - - field: str - """The field to compare.""" - value: int | float - """The value to test less than or equal with.""" - - def to_filter_expression_proto(self) -> vectorindex_pb._FilterExpression: - return vectorindex_pb._FilterExpression(less_than_or_equal_expression=self.to_proto()) - - def to_proto(self) -> vectorindex_pb._LessThanOrEqualExpression: - if type(self.value) is int: - return vectorindex_pb._LessThanOrEqualExpression(field=self.field, integer_value=self.value) - elif type(self.value) is float: - return vectorindex_pb._LessThanOrEqualExpression(field=self.field, float_value=self.value) - else: - raise InvalidArgumentException( - f"Invalid type for value: {type(self.value)} in less than or equal expression. Must be one of int, float.", - Service.INDEX, - ) - - -@dataclass -class ListContains(FilterExpression): - """Represents a list contains expression between a field and a value.""" - - field: str - """The list field to test.""" - value: str - """The value to test list contains with.""" - - def to_filter_expression_proto(self) -> vectorindex_pb._FilterExpression: - return vectorindex_pb._FilterExpression(list_contains_expression=self.to_proto()) - - 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/src/momento/requests/vector_index/search.py b/src/momento/requests/vector_index/search.py deleted file mode 100644 index 6a4f111f..00000000 --- a/src/momento/requests/vector_index/search.py +++ /dev/null @@ -1,8 +0,0 @@ -class AllMetadata: - """A class to represent a search for all metadata.""" - - pass - - -ALL_METADATA = AllMetadata() -"""A special value to pass to :meth:`VectorIndexClientAsync.search` to request all metadata fields.""" diff --git a/src/momento/requests/vector_index/similarity_metric.py b/src/momento/requests/vector_index/similarity_metric.py deleted file mode 100644 index f083922a..00000000 --- a/src/momento/requests/vector_index/similarity_metric.py +++ /dev/null @@ -1,15 +0,0 @@ -from enum import Enum - - -class SimilarityMetric(Enum): - """The similarity metric to use when comparing vectors in the index.""" - - EUCLIDEAN_SIMILARITY = "EUCLIDEAN_SIMILARITY" - """The Euclidean distance squared between two vectors, ie the sum of squared differences between each element. - Smaller is better. Ranges from 0 to infinity.""" - INNER_PRODUCT = "INNER_PRODUCT" - """The inner product between two vectors, ie the sum of the element-wise products. - Bigger is better. Ranges from 0 to infinity.""" - COSINE_SIMILARITY = "COSINE_SIMILARITY" - """The cosine similarity between two vectors, ie the cosine of the angle between them. - Bigger is better. Ranges from -1 to 1.""" diff --git a/src/momento/responses/vector_index/__init__.py b/src/momento/responses/vector_index/__init__.py deleted file mode 100644 index 36fae3fc..00000000 --- a/src/momento/responses/vector_index/__init__.py +++ /dev/null @@ -1,48 +0,0 @@ -from momento.common_data.vector_index.item import Item - -from .control.create import CreateIndex, CreateIndexResponse -from .control.delete import DeleteIndex, DeleteIndexResponse -from .control.list import IndexInfo, ListIndexes, ListIndexesResponse -from .data.count_items import CountItems, CountItemsResponse -from .data.delete_item_batch import DeleteItemBatch, DeleteItemBatchResponse -from .data.get_item_batch import ( - GetItemBatch, - GetItemBatchResponse, -) -from .data.get_item_metadata_batch import GetItemMetadataBatch, GetItemMetadataBatchResponse -from .data.search import Search, SearchHit, SearchResponse -from .data.search_and_fetch_vectors import ( - SearchAndFetchVectors, - SearchAndFetchVectorsHit, - SearchAndFetchVectorsResponse, -) -from .data.upsert_item_batch import UpsertItemBatch, UpsertItemBatchResponse -from .response import VectorIndexResponse - -__all__ = [ - "CountItems", - "CountItemsResponse", - "CreateIndex", - "CreateIndexResponse", - "DeleteIndex", - "DeleteIndexResponse", - "GetItemMetadataBatch", - "GetItemMetadataBatchResponse", - "GetItemBatch", - "GetItemBatchResponse", - "IndexInfo", - "Item", - "ListIndexes", - "ListIndexesResponse", - "SearchHit", - "Search", - "SearchAndFetchVectors", - "SearchAndFetchVectorsHit", - "SearchResponse", - "SearchAndFetchVectorsResponse", - "UpsertItemBatch", - "UpsertItemBatchResponse", - "DeleteItemBatch", - "DeleteItemBatchResponse", - "VectorIndexResponse", -] diff --git a/src/momento/responses/vector_index/control/__init__.py b/src/momento/responses/vector_index/control/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/src/momento/responses/vector_index/control/create.py b/src/momento/responses/vector_index/control/create.py deleted file mode 100644 index f95534ea..00000000 --- a/src/momento/responses/vector_index/control/create.py +++ /dev/null @@ -1,57 +0,0 @@ -from abc import ABC - -from ...mixins import ErrorResponseMixin -from ...response import ControlResponse - - -class CreateIndexResponse(ControlResponse): - """Parent response type for a create index request. - - The response object is resolved to a type-safe object of one of - the following subtypes: - - `CreateIndex.Success` - - `CreateIndex.IndexAlreadyExists` - - `CreateIndex.Error` - - Pattern matching can be used to operate on the appropriate subtype. - For example, in python 3.10+:: - - match response: - case CreateIndex.Success(): - ... - case CreateIndex.IndexAlreadyExists(): - ... - case CreateIndex.Error(): - ... - case _: - # Shouldn't happen - - or equivalently in earlier versions of python:: - - if isinstance(response, CreateIndex.Success): - ... - elif isinstance(response, CreateIndex.AlreadyExists): - ... - elif isinstance(response, CreateIndex.Error): - ... - else: - # Shouldn't happen - """ - - -class CreateIndex(ABC): - """Groups all `CreateIndexResponse` derived types under a common namespace.""" - - class Success(CreateIndexResponse): - """Indicates the request was successful.""" - - class IndexAlreadyExists(CreateIndexResponse): - """Indicates that a index with the requested name has already been created in the requesting account.""" - - class Error(CreateIndexResponse, ErrorResponseMixin): - """Contains information about an error returned from a request. - - This includes: - - `error_code`: `MomentoErrorCode` value for the error. - - `messsage`: a detailed error message. - """ diff --git a/src/momento/responses/vector_index/control/delete.py b/src/momento/responses/vector_index/control/delete.py deleted file mode 100644 index 77e5a446..00000000 --- a/src/momento/responses/vector_index/control/delete.py +++ /dev/null @@ -1,49 +0,0 @@ -from abc import ABC - -from ...mixins import ErrorResponseMixin -from ...response import ControlResponse - - -class DeleteIndexResponse(ControlResponse): - """Parent response type for a delete index request. - - The response object is resolved to a type-safe object of one of - the following subtypes: - - `DeleteIndex.Success` - - `DeleteIndex.Error` - - Pattern matching can be used to operate on the appropriate subtype. - For example, in python 3.10+:: - - match response: - case DeleteIndex.Success(): - ... - case DeleteIndex.Error(): - ... - case _: - # Shouldn't happen - - or equivalently in earlier versions of python:: - - if isinstance(response, DeleteIndex.Success): - ... - elif isinstance(response, DeleteIndex.Error): - ... - else: - # Shouldn't happen - """ - - -class DeleteIndex(ABC): - """Groups all `DeleteIndexResponse` derived types under a common namespace.""" - - class Success(DeleteIndexResponse): - """Indicates the request was successful.""" - - class Error(DeleteIndexResponse, ErrorResponseMixin): - """Contains information about an error returned from a request. - - This includes: - - `error_code`: `MomentoErrorCode` value for the error. - - `messsage`: a detailed error message. - """ diff --git a/src/momento/responses/vector_index/control/list.py b/src/momento/responses/vector_index/control/list.py deleted file mode 100644 index f193d896..00000000 --- a/src/momento/responses/vector_index/control/list.py +++ /dev/null @@ -1,101 +0,0 @@ -from __future__ import annotations - -from abc import ABC -from dataclasses import dataclass - -from momento_wire_types import controlclient_pb2 as ctrl_pb - -from momento.errors.exceptions import UnknownException -from momento.requests.vector_index import SimilarityMetric - -from ...mixins import ErrorResponseMixin -from ...response import ControlResponse - - -class ListIndexesResponse(ControlResponse): - """Parent response type for a list indexes request. - - The response object is resolved to a type-safe object of one of - the following subtypes: - - `ListIndexes.Success` - - `ListIndexes.Error` - - Pattern matching can be used to operate on the appropriate subtype. - For example, in python 3.10+:: - - match response: - case ListIndexes.Success(): - ... - case ListIndexes.Error(): - ... - case _: - # Shouldn't happen - - or equivalently in earlier versions of python:: - - if isinstance(response, ListIndexes.Success): - ... - elif isinstance(response, ListIndexes.Error): - ... - else: - # Shouldn't happen - """ - - -@dataclass -class IndexInfo: - """Contains a Momento index's info.""" - - name: str - num_dimensions: int - similarity_metric: SimilarityMetric - - @staticmethod - def from_grpc_response(grpc_index_info: ctrl_pb._ListIndexesResponse._Index) -> "IndexInfo": - metric_type: str = grpc_index_info.similarity_metric.WhichOneof("similarity_metric") - similarity_metric: SimilarityMetric - - if metric_type == "cosine_similarity": - similarity_metric = SimilarityMetric.COSINE_SIMILARITY - elif metric_type == "euclidean_similarity": - similarity_metric = SimilarityMetric.EUCLIDEAN_SIMILARITY - elif metric_type == "inner_product": - similarity_metric = SimilarityMetric.INNER_PRODUCT - else: - raise UnknownException(f"Unknown similarity metric: {metric_type}") - - return IndexInfo( - name=grpc_index_info.index_name, - num_dimensions=grpc_index_info.num_dimensions, - similarity_metric=similarity_metric, - ) - - -class ListIndexes(ABC): - """Groups all `ListIndexesResponse` derived types under a common namespace.""" - - @dataclass - class Success(ListIndexesResponse): - """Indicates the request was successful.""" - - indexes: list[IndexInfo] - """The list of indexes available to the user.""" - - @staticmethod - def from_grpc_response(grpc_list_index_response: ctrl_pb._ListIndexesResponse) -> ListIndexes.Success: - """Initializes ListIndexResponse to handle list index response. - - Args: - grpc_list_index_response: Protobuf based response returned by Scs. - """ - return ListIndexes.Success( - indexes=[IndexInfo.from_grpc_response(index) for index in grpc_list_index_response.indexes] # type: ignore[misc] # noqa: E501 - ) - - class Error(ListIndexesResponse, ErrorResponseMixin): - """Contains information about an error returned from a request. - - This includes: - - `error_code`: `MomentoErrorCode` value for the error. - - `messsage`: a detailed error message. - """ diff --git a/src/momento/responses/vector_index/data/__init__.py b/src/momento/responses/vector_index/data/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/src/momento/responses/vector_index/data/count_items.py b/src/momento/responses/vector_index/data/count_items.py deleted file mode 100644 index 47291f6b..00000000 --- a/src/momento/responses/vector_index/data/count_items.py +++ /dev/null @@ -1,37 +0,0 @@ -from __future__ import annotations - -from abc import ABC -from dataclasses import dataclass - -from ...mixins import ErrorResponseMixin -from ..response import VectorIndexResponse - - -class CountItemsResponse(VectorIndexResponse): - """Parent response type for a `count_items` request. - - Its subtypes are: - - `CountItems.Success` - - `CountItems.Error` - - See `PreviewVectorIndexClient` for how to work with responses. - """ - - -class CountItems(ABC): - """Groups all `CountItemsResponse` derived types under a common namespace.""" - - @dataclass - class Success(CountItemsResponse): - """Contains the result of a `count_items` request.""" - - item_count: int - """The number of items in the index.""" - - class Error(CountItemsResponse, ErrorResponseMixin): - """Contains information about an error returned from a request. - - This includes: - - `error_code`: `MomentoErrorCode` value for the error. - - `messsage`: a detailed error message. - """ diff --git a/src/momento/responses/vector_index/data/delete_item_batch.py b/src/momento/responses/vector_index/data/delete_item_batch.py deleted file mode 100644 index 2c94e3d5..00000000 --- a/src/momento/responses/vector_index/data/delete_item_batch.py +++ /dev/null @@ -1,30 +0,0 @@ -from abc import ABC - -from ...mixins import ErrorResponseMixin -from ..response import VectorIndexResponse - - -class DeleteItemBatchResponse(VectorIndexResponse): - """Parent response type for a vector index `delete_item_batch` request. - - Its subtypes are: - - `DeleteItemBatch.Success` - - `DeleteItemBatch.Error` - - See `VectorIndexClient` for how to work with responses. - """ - - -class DeleteItemBatch(ABC): - """Groups all `DeleteItemBatchResponse` derived types under a common namespace.""" - - class Success(DeleteItemBatchResponse): - """Indicates the request was successful.""" - - class Error(DeleteItemBatchResponse, ErrorResponseMixin): - """Contains information about an error returned from a request. - - This includes: - - `error_code`: `MomentoErrorCode` value for the error. - - `messsage`: a detailed error message. - """ diff --git a/src/momento/responses/vector_index/data/get_item_batch.py b/src/momento/responses/vector_index/data/get_item_batch.py deleted file mode 100644 index e20a2fd3..00000000 --- a/src/momento/responses/vector_index/data/get_item_batch.py +++ /dev/null @@ -1,56 +0,0 @@ -from __future__ import annotations - -from abc import ABC -from dataclasses import dataclass - -from momento_wire_types import vectorindex_pb2 as pb - -from momento.common_data.vector_index.item import Item - -from ...mixins import ErrorResponseMixin -from ..response import VectorIndexResponse -from .utils import pb_metadata_to_dict - - -class GetItemBatchResponse(VectorIndexResponse): - """Parent response type for a `get_item_batch` request. - - Its subtypes are: - - `GetItemBatch.Success` - - `GetItemBatch.Error` - - See `PreviewVectorIndexClient` for how to work with responses. - """ - - -class GetItemBatch(ABC): - """Groups all `GetItemBatchResponse` derived types under a common namespace.""" - - @dataclass - class Success(GetItemBatchResponse): - """Contains the result of a `get_item_batch` request.""" - - values: dict[str, Item] - """The items that were found.""" - - @staticmethod - def from_proto(response: pb._GetItemBatchResponse) -> "GetItemBatch.Success": - """Converts a sequence of proto _GetItemBatchResponse to a `GetItemBatch.Success`.""" - values = { - item.id: Item( - id=item.id, - vector=list(item.vector.elements), - metadata=pb_metadata_to_dict(item.metadata), - ) - for item in response.item_response - } - - return GetItemBatch.Success(values=values) - - class Error(GetItemBatchResponse, ErrorResponseMixin): - """Contains information about an error returned from a request. - - This includes: - - `error_code`: `MomentoErrorCode` value for the error. - - `messsage`: a detailed error message. - """ diff --git a/src/momento/responses/vector_index/data/get_item_metadata_batch.py b/src/momento/responses/vector_index/data/get_item_metadata_batch.py deleted file mode 100644 index 018bb650..00000000 --- a/src/momento/responses/vector_index/data/get_item_metadata_batch.py +++ /dev/null @@ -1,48 +0,0 @@ -from __future__ import annotations - -from abc import ABC -from dataclasses import dataclass - -from momento_wire_types import vectorindex_pb2 as pb - -from momento.common_data.vector_index.item import Metadata - -from ...mixins import ErrorResponseMixin -from ..response import VectorIndexResponse -from .utils import pb_metadata_to_dict - - -class GetItemMetadataBatchResponse(VectorIndexResponse): - """Parent response type for a `get_item_metadata_batch` request. - - Its subtypes are: - - `GetItemMetadataBatch.Success` - - `GetItemMetadataBatch.Error` - - See `PreviewVectorIndexClient` for how to work with responses. - """ - - -class GetItemMetadataBatch(ABC): - """Groups all `GetItemMetadataBatchResponse` derived types under a common namespace.""" - - @dataclass - class Success(GetItemMetadataBatchResponse): - """Contains the result of a `get_item_metadata_batch` request.""" - - values: dict[str, Metadata] - """The metadata of the items that were found.""" - - @staticmethod - def from_proto(response: pb._GetItemMetadataBatchResponse) -> "GetItemMetadataBatch.Success": - """Converts a proto _GetItemMetadataBatchResponse to a `GetItemMetadataBatch.Success`.""" - values = {item.id: pb_metadata_to_dict(item.metadata) for item in response.item_metadata_response} - return GetItemMetadataBatch.Success(values=values) - - class Error(GetItemMetadataBatchResponse, ErrorResponseMixin): - """Contains information about an error returned from a request. - - This includes: - - `error_code`: `MomentoErrorCode` value for the error. - - `messsage`: a detailed error message. - """ diff --git a/src/momento/responses/vector_index/data/search.py b/src/momento/responses/vector_index/data/search.py deleted file mode 100644 index 47b4d310..00000000 --- a/src/momento/responses/vector_index/data/search.py +++ /dev/null @@ -1,74 +0,0 @@ -from __future__ import annotations - -from abc import ABC -from dataclasses import dataclass -from typing import Optional - -from momento_wire_types import vectorindex_pb2 as vectorindex_pb - -from momento.common_data.vector_index.item import Metadata - -from ...mixins import ErrorResponseMixin -from ..response import VectorIndexResponse -from .utils import pb_metadata_to_dict - - -class SearchResponse(VectorIndexResponse): - """Parent response type for a vector index `search` request. - - Its subtypes are: - - `Search.Success` - - `Search.Error` - - See `VectorIndexClient` for how to work with responses. - """ - - -class SearchHit: - """A single search hit.""" - - def __init__(self, score: float, id: str, metadata: Optional[Metadata] = None): - """Initializes a new instance of `SearchHit`. - - Args: - score (float): The similarity of the hit to the query. - id (str): The id of the hit. - metadata (Optional[Metadata], optional): The metadata of the hit. Defaults to None, ie empty metadata. - """ - self.score = score - self.id = id - self.metadata = metadata or {} - - def __eq__(self, other: object) -> bool: - if isinstance(other, SearchHit): - return self.score == other.score and self.id == other.id and self.metadata == other.metadata - return False - - def __hash__(self) -> int: - return hash((self.score, self.id, self.metadata)) - - def __repr__(self) -> str: - return f"{self.__class__.__name__}(score={self.score!r}, id={self.id!r}, metadata={self.metadata!r})" - - @staticmethod - def from_proto(hit: vectorindex_pb._SearchHit) -> SearchHit: - metadata = pb_metadata_to_dict(hit.metadata) - return SearchHit(id=hit.id, score=hit.score, metadata=metadata) - - -class Search(ABC): - """Groups all `SearchResponse` derived types under a common namespace.""" - - @dataclass - class Success(SearchResponse): - """Indicates the request was successful.""" - - hits: list[SearchHit] - - class Error(SearchResponse, ErrorResponseMixin): - """Contains information about an error returned from a request. - - This includes: - - `error_code`: `MomentoErrorCode` value for the error. - - `messsage`: a detailed error message. - """ diff --git a/src/momento/responses/vector_index/data/search_and_fetch_vectors.py b/src/momento/responses/vector_index/data/search_and_fetch_vectors.py deleted file mode 100644 index 918c18a0..00000000 --- a/src/momento/responses/vector_index/data/search_and_fetch_vectors.py +++ /dev/null @@ -1,77 +0,0 @@ -from __future__ import annotations - -from abc import ABC -from dataclasses import dataclass -from typing import Optional - -from momento_wire_types import vectorindex_pb2 as vectorindex_pb - -from momento.common_data.vector_index.item import Metadata - -from ...mixins import ErrorResponseMixin -from ..response import VectorIndexResponse -from .search import SearchHit -from .utils import pb_metadata_to_dict - - -class SearchAndFetchVectorsResponse(VectorIndexResponse): - """Parent response type for a vector index `search_and_fetch_vectors` request. - - Its subtypes are: - - `SearchAndFetchVectors.Success` - - `SearchAndFetchVectors.Error` - - See `VectorIndexClient` for how to work with responses. - """ - - -class SearchAndFetchVectorsHit(SearchHit): - def __init__(self, score: float, id: str, vector: list[float], metadata: Optional[Metadata] = None): - """Initializes a new instance of `SearchAndFetchVectorsHit`. - - Args: - score (float): The similarity of the hit to the query. - id (str): The id of the hit. - vector (list[float]): The vector of the hit. - metadata (Optional[Metadata], optional): The metadata of the hit. Defaults to None, ie empty metadata. - """ - super().__init__(score, id, metadata) - self.vector = vector - - def __eq__(self, other: object) -> bool: - if isinstance(other, SearchAndFetchVectorsHit): - return super().__eq__(other) and self.vector == other.vector - return False - - def __hash__(self) -> int: - return hash((self.score, self.id, self.vector, self.metadata)) - - def __repr__(self) -> str: - return f"{self.__class__.__name__}(score={self.score!r}, id={self.id!r}, vector={self.vector!r}, metadata={self.metadata!r})" - - @staticmethod - def from_search_hit(hit: SearchHit, vector: list[float]) -> SearchAndFetchVectorsHit: - return SearchAndFetchVectorsHit(id=hit.id, score=hit.score, metadata=hit.metadata, vector=vector) - - @staticmethod - def from_proto(hit: vectorindex_pb._SearchAndFetchVectorsHit) -> SearchAndFetchVectorsHit: - metadata = pb_metadata_to_dict(hit.metadata) - return SearchAndFetchVectorsHit(id=hit.id, score=hit.score, metadata=metadata, vector=list(hit.vector.elements)) - - -class SearchAndFetchVectors(ABC): - """Groups all `SearchAndFetchVectorsResponse` derived types under a common namespace.""" - - @dataclass - class Success(SearchAndFetchVectorsResponse): - """Indicates the request was successful.""" - - hits: list[SearchAndFetchVectorsHit] - - class Error(SearchAndFetchVectorsResponse, ErrorResponseMixin): - """Contains information about an error returned from a request. - - This includes: - - `error_code`: `MomentoErrorCode` value for the error. - - `messsage`: a detailed error message. - """ diff --git a/src/momento/responses/vector_index/data/upsert_item_batch.py b/src/momento/responses/vector_index/data/upsert_item_batch.py deleted file mode 100644 index 354de4c4..00000000 --- a/src/momento/responses/vector_index/data/upsert_item_batch.py +++ /dev/null @@ -1,30 +0,0 @@ -from abc import ABC - -from ...mixins import ErrorResponseMixin -from ..response import VectorIndexResponse - - -class UpsertItemBatchResponse(VectorIndexResponse): - """Parent response type for a vector index `add_item_batch` request. - - Its subtypes are: - - `UpsertItemBatch.Success` - - `UpsertItemBatch.Error` - - See `VectorIndexClient` for how to work with responses. - """ - - -class UpsertItemBatch(ABC): - """Groups all `UpsertItemBatchResponse` derived types under a common namespace.""" - - class Success(UpsertItemBatchResponse): - """Indicates the request was successful.""" - - class Error(UpsertItemBatchResponse, ErrorResponseMixin): - """Contains information about an error returned from a request. - - This includes: - - `error_code`: `MomentoErrorCode` value for the error. - - `messsage`: a detailed error message. - """ diff --git a/src/momento/responses/vector_index/data/utils.py b/src/momento/responses/vector_index/data/utils.py deleted file mode 100644 index 29b6aa83..00000000 --- a/src/momento/responses/vector_index/data/utils.py +++ /dev/null @@ -1,30 +0,0 @@ -from __future__ import annotations - -from typing import Iterable - -from momento_wire_types import vectorindex_pb2 as pb - -from momento.common_data.vector_index.item import Metadata -from momento.errors import UnknownException - - -def pb_metadata_to_dict( - pb_metadata: Iterable[pb._Metadata], -) -> Metadata: - metadata: Metadata = {} - for item in pb_metadata: - type = item.WhichOneof("value") - field = item.field - if type == "string_value": - metadata[field] = item.string_value - elif type == "integer_value": - metadata[field] = item.integer_value - elif type == "double_value": - metadata[field] = item.double_value - elif type == "boolean_value": - metadata[field] = item.boolean_value - elif type == "list_of_strings_value": - metadata[field] = list(item.list_of_strings_value.values) - else: - raise UnknownException(f"Unknown metadata value: {type}") - return metadata diff --git a/src/momento/responses/vector_index/response.py b/src/momento/responses/vector_index/response.py deleted file mode 100644 index 8cb114a3..00000000 --- a/src/momento/responses/vector_index/response.py +++ /dev/null @@ -1,5 +0,0 @@ -from ..response import Response - - -class VectorIndexResponse(Response): - ... diff --git a/src/momento/vector_index_client.py b/src/momento/vector_index_client.py deleted file mode 100644 index 55b5a723..00000000 --- a/src/momento/vector_index_client.py +++ /dev/null @@ -1,305 +0,0 @@ -from __future__ import annotations - -from types import TracebackType -from typing import Optional, Type - -from momento import logs -from momento.auth import CredentialProvider -from momento.config import VectorIndexConfiguration - -try: - from momento.internal._utilities import _validate_request_timeout - from momento.internal.synchronous._vector_index_control_client import ( - _VectorIndexControlClient, - ) - from momento.internal.synchronous._vector_index_data_client import _VectorIndexDataClient -except ImportError as e: - if e.name == "cygrpc": - import sys - - print( - "There is an issue on M1 macs between GRPC native packaging and Python wheel tags. " - "See https://github.com/grpc/grpc/issues/28387", - file=sys.stderr, - ) - print("-".join("" for _ in range(99)), file=sys.stderr) - print(" TO WORK AROUND:", file=sys.stderr) - print(" * Install Rosetta 2", file=sys.stderr) - print( - " * Install Python from python.org (you might need to do this if you're using an arm-only build)", - file=sys.stderr, - ) - print(" * re-run with:", file=sys.stderr) - print("arch -x86_64 {} {}".format(sys.executable, *sys.argv), file=sys.stderr) - print("-".join("" for _ in range(99)), file=sys.stderr) - raise e - -from momento.requests.vector_index import AllMetadata, FilterExpression, Item, SimilarityMetric -from momento.responses.vector_index import ( - CountItemsResponse, - CreateIndexResponse, - DeleteIndexResponse, - DeleteItemBatchResponse, - GetItemBatchResponse, - GetItemMetadataBatchResponse, - ListIndexesResponse, - SearchAndFetchVectorsResponse, - SearchResponse, - UpsertItemBatchResponse, -) - - -class PreviewVectorIndexClient: - """Synchronous Vector Index Client. - - Vector and control methods return a response object unique to each request. - The response object is resolved to a type-safe object of one of several - sub-types. See the documentation for each response type for details. - - Pattern matching can be used to operate on the appropriate subtype. - For example, in python 3.10+ if you're deleting a key:: - - response = client.create_index(index_name, num_dimensions) - match response: - case CreateIndex.Success(): - ...the index was created... - case CreateIndex.IndexAlreadyExists(): - ... the index already exists... - case CreateIndex.Error(): - ...there was an error trying to delete the key... - - or equivalently in earlier versions of python:: - - response = client.create_index(index_name, num_dimensions) - if isinstance(response, CreateIndex.Success): - ... - case isinstance(response, CreateIndex.IndexAlreadyExists): - ... the index already exists... - elif isinstance(response, CreateIndex.Error): - ... - else: - raise Exception("This should never happen") - """ - - def __init__(self, configuration: VectorIndexConfiguration, credential_provider: CredentialProvider): - """Instantiate a client. - - Args: - configuration (VectorIndexConfiguration): An object holding configuration settings for communication - with the server. - credential_provider (CredentialProvider): An object holding the auth token and endpoint information. - - Raises: - IllegalArgumentException: If method arguments fail validations. - Example:: - - from momento import CredentialProvider, PreviewVectorIndexClient, VectorIndexConfigurations - - configuration = VectorIndexConfigurations.Laptop.latest() - credential_provider = CredentialProvider.from_environment_variable("MOMENTO_API_KEY") - client = PreviewVectorIndexClient(configuration, credential_provider) - """ - _validate_request_timeout(configuration.get_transport_strategy().get_grpc_configuration().get_deadline()) - self._logger = logs.logger - self._next_client_index = 0 - self._control_client = _VectorIndexControlClient(configuration, credential_provider) - self._data_client = _VectorIndexDataClient(configuration, credential_provider) - - def __enter__(self) -> PreviewVectorIndexClient: - return self - - def __exit__( - self, - exc_type: Optional[Type[BaseException]], - exc_value: Optional[BaseException], - traceback: Optional[TracebackType], - ) -> None: - self._control_client.close() - self._data_client.close() - - def count_items(self, index_name: str) -> CountItemsResponse: - """Gets the number of items in a vector index. - - Note that if the vector index does not exist, a `NOT_FOUND` error will be returned. - - Args: - index_name (str): Name of the index to count the items in. - - Returns: - CountItemsResponse: The result of a count items operation. - """ - return self._data_client.count_items(index_name) - - def create_index( - self, - index_name: str, - num_dimensions: int, - similarity_metric: SimilarityMetric = SimilarityMetric.COSINE_SIMILARITY, - ) -> CreateIndexResponse: - """Creates a vector index if it doesn't exist. - - Remark on the choice of similarity metric: - - Cosine similarity is appropriate for most embedding models as they tend to be optimized - for this metric. - - If the vectors are unit normalized, cosine similarity is equivalent to inner product. - If your vectors are already unit normalized, you can use inner product to improve - performance. - - Euclidean similarity, the sum of squared differences, is appropriate for datasets where - this metric is meaningful. For example, if the vectors represent images, and the - embedding model is trained to optimize the euclidean distance between images, then - euclidean similarity is appropriate. - - Args: - index_name (str): Name of the index to be created. - num_dimensions (int): Number of dimensions of the vectors to be indexed. - similarity_metric (SimilarityMetric): The similarity metric to use when comparing - vectors in the index. Defaults to SimilarityMetric.COSINE_SIMILARITY. - - Returns: - CreateIndexResponse: The result of a create index operation. - """ - return self._control_client.create_index(index_name, num_dimensions, similarity_metric) - - def delete_index(self, index_name: str) -> DeleteIndexResponse: - """Deletes a vector index and all of the items within it. - - Args: - index_name (str): Name of the index to be deleted. - - Returns: - DeleteIndexResponse: The result of a delete index operation. - """ - return self._control_client.delete_index(index_name) - - def list_indexes(self) -> ListIndexesResponse: - """Lists all vector indexes. - - Returns: - ListIndexesResponse: The result of a list indexes operation. - """ - return self._control_client.list_indexes() - - def upsert_item_batch(self, index_name: str, items: list[Item]) -> UpsertItemBatchResponse: - """Upserts a batch of items into a vector index. - - If an item with the same ID already exists in the index, it will be replaced. - Otherwise, it will be added to the index. - - Args: - index_name (str): Name of the index to add the items into. - items (list[Item]): The items to be added into the index. - - Returns: - AddItemBatchResponse: The result of an add item batch operation. - """ - return self._data_client.upsert_item_batch(index_name, items) - - def delete_item_batch(self, index_name: str, filter: FilterExpression | list[str]) -> DeleteItemBatchResponse: - """Deletes a batch of items from a vector index. - - Deletes any and all items with the given IDs from the index. - - Args: - index_name (str): Name of the index to delete the items from. - filter (FilterExpression | list[str]): A filter expression to match - items to be deleted, or list of item IDs to be deleted. - - Returns: - DeleteItemBatchResponse: The result of a delete item batch operation. - """ - return self._data_client.delete_item_batch(index_name, filter) - - def search( - self, - index_name: str, - query_vector: list[float], - top_k: int = 10, - metadata_fields: Optional[list[str]] | AllMetadata = None, - score_threshold: Optional[float] = None, - filter: Optional[FilterExpression] = None, - ) -> SearchResponse: - """Searches for the most similar vectors to the query vector in the index. - - Ranks the results using the similarity metric specified when the index was created. - - Args: - index_name (str): Name of the index to search in. - query_vector (list[float]): The vector to search for. - top_k (int): The number of results to return. Defaults to 10. - metadata_fields (Optional[list[str]] | AllMetadata): A list of metadata fields - to return with each result. If not provided, no metadata is returned. - If the special value `ALL_METADATA` is provided, all metadata is returned. - Defaults to None. - score_threshold (Optional[float]): A score threshold to filter results by. - For cosine similarity and inner product, scores lower than the threshold - are excluded. For euclidean similarity, scores higher than the threshold - are excluded. The threshold is exclusive. Defaults to None, ie no threshold. - filter (Optional[FilterExpression]): A filter expression to filter - results by. Defaults to None, ie no filter. - - Returns: - SearchResponse: The result of a search operation. - """ - return self._data_client.search(index_name, query_vector, top_k, metadata_fields, score_threshold, filter) - - def search_and_fetch_vectors( - self, - index_name: str, - query_vector: list[float], - top_k: int = 10, - metadata_fields: Optional[list[str]] | AllMetadata = None, - score_threshold: Optional[float] = None, - filter: Optional[FilterExpression] = None, - ) -> SearchAndFetchVectorsResponse: - """Searches for the most similar vectors to the query vector in the index. - - Ranks the results using the similarity metric specified when the index was created. - Also returns the vectors associated with each result. - - Args: - index_name (str): Name of the index to search in. - query_vector (list[float]): The vector to search for. - top_k (int): The number of results to return. Defaults to 10. - metadata_fields (Optional[list[str]] | AllMetadata): A list of metadata fields - to return with each result. If not provided, no metadata is returned. - If the special value `ALL_METADATA` is provided, all metadata is returned. - Defaults to None. - score_threshold (Optional[float]): A score threshold to filter results by. - For cosine similarity and inner product, scores lower than the threshold - are excluded. For euclidean similarity, scores higher than the threshold - are excluded. The threshold is exclusive. Defaults to None, ie no threshold. - filter (Optional[FilterExpression]): A filter expression to filter - results by. Defaults to None, ie no filter. - - Returns: - SearchResponse: The result of a search operation. - """ - return self._data_client.search_and_fetch_vectors( - index_name, query_vector, top_k, metadata_fields, score_threshold, filter - ) - - def get_item_batch(self, index_name: str, filter: list[str]) -> GetItemBatchResponse: - """Gets a batch of items from a vector index by ID. - - Args: - index_name (str): Name of the index to get the item from. - filter (list[str]): The IDs of the items to be retrieved from the index. - - Returns: - GetItemBatchResponse: The result of a get item batch operation. - """ - return self._data_client.get_item_batch(index_name, filter) - - def get_item_metadata_batch(self, index_name: str, filter: list[str]) -> GetItemMetadataBatchResponse: - """Gets metadata for a batch of items from a vector index by ID. - - Args: - index_name (str): Name of the index to get the items from. - filter (list[str]): The IDs of the item metadata to be retrieved from the index. - - Returns: - GetItemMetadataBatchResponse: The result of a get item metadata batch operation. - """ - return self._data_client.get_item_metadata_batch(index_name, filter) - - # TODO: repr diff --git a/src/momento/vector_index_client_async.py b/src/momento/vector_index_client_async.py deleted file mode 100644 index 94c236aa..00000000 --- a/src/momento/vector_index_client_async.py +++ /dev/null @@ -1,305 +0,0 @@ -from __future__ import annotations - -from types import TracebackType -from typing import Optional, Type - -from momento import logs -from momento.auth import CredentialProvider -from momento.config import VectorIndexConfiguration - -try: - from momento.internal._utilities import _validate_request_timeout - from momento.internal.aio._vector_index_control_client import ( - _VectorIndexControlClient, - ) - from momento.internal.aio._vector_index_data_client import _VectorIndexDataClient -except ImportError as e: - if e.name == "cygrpc": - import sys - - print( - "There is an issue on M1 macs between GRPC native packaging and Python wheel tags. " - "See https://github.com/grpc/grpc/issues/28387", - file=sys.stderr, - ) - print("-".join("" for _ in range(99)), file=sys.stderr) - print(" TO WORK AROUND:", file=sys.stderr) - print(" * Install Rosetta 2", file=sys.stderr) - print( - " * Install Python from python.org (you might need to do this if you're using an arm-only build)", - file=sys.stderr, - ) - print(" * re-run with:", file=sys.stderr) - print("arch -x86_64 {} {}".format(sys.executable, *sys.argv), file=sys.stderr) - print("-".join("" for _ in range(99)), file=sys.stderr) - raise e - -from momento.requests.vector_index import AllMetadata, FilterExpression, Item, SimilarityMetric -from momento.responses.vector_index import ( - CountItemsResponse, - CreateIndexResponse, - DeleteIndexResponse, - DeleteItemBatchResponse, - GetItemBatchResponse, - GetItemMetadataBatchResponse, - ListIndexesResponse, - SearchAndFetchVectorsResponse, - SearchResponse, - UpsertItemBatchResponse, -) - - -class PreviewVectorIndexClientAsync: - """Async Vector Index Client. - - Vector and control methods return a response object unique to each request. - The response object is resolved to a type-safe object of one of several - sub-types. See the documentation for each response type for details. - - Pattern matching can be used to operate on the appropriate subtype. - For example, in python 3.10+ if you're deleting a key:: - - response = await client.create_index(index_name, num_dimensions) - match response: - case CreateIndex.Success(): - ...the index was created... - case CreateIndex.IndexAlreadyExists(): - ... the index already exists... - case CreateIndex.Error(): - ...there was an error trying to delete the key... - - or equivalently in earlier versions of python:: - - response = await client.create_index(index_name, num_dimensions) - if isinstance(response, CreateIndex.Success): - ... - case isinstance(response, CreateIndex.IndexAlreadyExists): - ... the index already exists... - elif isinstance(response, CreateIndex.Error): - ... - else: - raise Exception("This should never happen") - """ - - def __init__(self, configuration: VectorIndexConfiguration, credential_provider: CredentialProvider): - """Instantiate a client. - - Args: - configuration (VectorIndexConfiguration): An object holding configuration settings for communication - with the server. - credential_provider (CredentialProvider): An object holding the auth token and endpoint information. - - Raises: - IllegalArgumentException: If method arguments fail validations. - Example:: - - from momento import CredentialProvider, PreviewVectorIndexClientAsync, VectorIndexConfigurations - - configuration = VectorIndexConfigurations.Laptop.latest() - credential_provider = CredentialProvider.from_environment_variable("MOMENTO_API_KEY") - client = PreviewVectorIndexClientAsync(configuration, credential_provider) - """ - _validate_request_timeout(configuration.get_transport_strategy().get_grpc_configuration().get_deadline()) - self._logger = logs.logger - self._next_client_index = 0 - self._control_client = _VectorIndexControlClient(configuration, credential_provider) - self._data_client = _VectorIndexDataClient(configuration, credential_provider) - - async def __aenter__(self) -> PreviewVectorIndexClientAsync: - return self - - async def __aexit__( - self, - exc_type: Optional[Type[BaseException]], - exc_value: Optional[BaseException], - traceback: Optional[TracebackType], - ) -> None: - await self._control_client.close() - await self._data_client.close() - - async def count_items(self, index_name: str) -> CountItemsResponse: - """Gets the number of items in a vector index. - - Note that if the vector index does not exist, a `NOT_FOUND` error will be returned. - - Args: - index_name (str): Name of the index to count the items in. - - Returns: - CountItemsResponse: The result of a count items operation. - """ - return await self._data_client.count_items(index_name) - - async def create_index( - self, - index_name: str, - num_dimensions: int, - similarity_metric: SimilarityMetric = SimilarityMetric.COSINE_SIMILARITY, - ) -> CreateIndexResponse: - """Creates a vector index if it doesn't exist. - - Remark on the choice of similarity metric: - - Cosine similarity is appropriate for most embedding models as they tend to be optimized - for this metric. - - If the vectors are unit normalized, cosine similarity is equivalent to inner product. - If your vectors are already unit normalized, you can use inner product to improve - performance. - - Euclidean similarity, the sum of squared differences, is appropriate for datasets where - this metric is meaningful. For example, if the vectors represent images, and the - embedding model is trained to optimize the euclidean distance between images, then - euclidean similarity is appropriate. - - Args: - index_name (str): Name of the index to be created. - num_dimensions (int): Number of dimensions of the vectors to be indexed. - similarity_metric (SimilarityMetric): The similarity metric to use when comparing - vectors in the index. Defaults to SimilarityMetric.COSINE_SIMILARITY. - - Returns: - CreateIndexResponse: The result of a create index operation. - """ - return await self._control_client.create_index(index_name, num_dimensions, similarity_metric) - - async def delete_index(self, index_name: str) -> DeleteIndexResponse: - """Deletes a vector index and all of the items within it. - - Args: - index_name (str): Name of the index to be deleted. - - Returns: - DeleteIndexResponse: The result of a delete index operation. - """ - return await self._control_client.delete_index(index_name) - - async def list_indexes(self) -> ListIndexesResponse: - """Lists all vector indexes. - - Returns: - ListIndexesResponse: The result of a list indexes operation. - """ - return await self._control_client.list_indexes() - - async def upsert_item_batch(self, index_name: str, items: list[Item]) -> UpsertItemBatchResponse: - """Upserts a batch of items into a vector index. - - If an item with the same ID already exists in the index, it will be replaced. - Otherwise, it will be added to the index. - - Args: - index_name (str): Name of the index to add the items into. - items (list[Item]): The items to be added into the index. - - Returns: - AddItemBatchResponse: The result of an add item batch operation. - """ - return await self._data_client.upsert_item_batch(index_name, items) - - async def delete_item_batch(self, index_name: str, filter: FilterExpression | list[str]) -> DeleteItemBatchResponse: - """Deletes a batch of items from a vector index. - - Deletes any and all items with the given IDs from the index. - - Args: - index_name (str): Name of the index to delete the items from. - filter (FilterExpression | list[str]): A filter expression to match - items to be deleted, or list of item IDs to be deleted. - - Returns: - DeleteItemBatchResponse: The result of a delete item batch operation. - """ - return await self._data_client.delete_item_batch(index_name, filter) - - async def search( - self, - index_name: str, - query_vector: list[float], - top_k: int = 10, - metadata_fields: Optional[list[str]] | AllMetadata = None, - score_threshold: Optional[float] = None, - filter: Optional[FilterExpression] = None, - ) -> SearchResponse: - """Searches for the most similar vectors to the query vector in the index. - - Ranks the results using the similarity metric specified when the index was created. - - Args: - index_name (str): Name of the index to search in. - query_vector (list[float]): The vector to search for. - top_k (int): The number of results to return. Defaults to 10. - metadata_fields (Optional[list[str]] | AllMetadata): A list of metadata fields - to return with each result. If not provided, no metadata is returned. - If the special value `ALL_METADATA` is provided, all metadata is returned. - Defaults to None. - score_threshold (Optional[float]): A score threshold to filter results by. - For cosine similarity and inner product, scores lower than the threshold - are excluded. For euclidean similarity, scores higher than the threshold - are excluded. The threshold is exclusive. Defaults to None, ie no threshold. - filter (Optional[FilterExpression]): A filter expression to filter - results by. Defaults to None, ie no filter. - - Returns: - SearchResponse: The result of a search operation. - """ - return await self._data_client.search(index_name, query_vector, top_k, metadata_fields, score_threshold, filter) - - async def search_and_fetch_vectors( - self, - index_name: str, - query_vector: list[float], - top_k: int = 10, - metadata_fields: Optional[list[str]] | AllMetadata = None, - score_threshold: Optional[float] = None, - filter: Optional[FilterExpression] = None, - ) -> SearchAndFetchVectorsResponse: - """Searches for the most similar vectors to the query vector in the index. - - Ranks the results using the similarity metric specified when the index was created. - Also returns the vectors associated with each result. - - Args: - index_name (str): Name of the index to search in. - query_vector (list[float]): The vector to search for. - top_k (int): The number of results to return. Defaults to 10. - metadata_fields (Optional[list[str]] | AllMetadata): A list of metadata fields - to return with each result. If not provided, no metadata is returned. - If the special value `ALL_METADATA` is provided, all metadata is returned. - Defaults to None. - score_threshold (Optional[float]): A score threshold to filter results by. - For cosine similarity and inner product, scores lower than the threshold - are excluded. For euclidean similarity, scores higher than the threshold - are excluded. The threshold is exclusive. Defaults to None, ie no threshold. - filter (Optional[FilterExpression]): A filter expression to filter - results by. Defaults to None, ie no filter. - - Returns: - SearchResponse: The result of a search operation. - """ - return await self._data_client.search_and_fetch_vectors( - index_name, query_vector, top_k, metadata_fields, score_threshold, filter - ) - - async def get_item_batch(self, index_name: str, filter: list[str]) -> GetItemBatchResponse: - """Gets a batch of items from a vector index by ID. - - Args: - index_name (str): Name of the index to get the item from. - filter (list[str]): The IDs of the items to be retrieved from the index. - - Returns: - GetItemBatchResponse: The result of a get item batch operation. - """ - return await self._data_client.get_item_batch(index_name, filter) - - async def get_item_metadata_batch(self, index_name: str, filter: list[str]) -> GetItemMetadataBatchResponse: - """Gets metadata for a batch of items from a vector index by ID. - - Args: - index_name (str): Name of the index to get the items from. - filter (list[str]): The IDs of the item metadata to be retrieved from the index. - - Returns: - GetItemMetadataBatchResponse: The result of a get item metadata batch operation. - """ - return await self._data_client.get_item_metadata_batch(index_name, filter) - - # TODO: repr diff --git a/tests/conftest.py b/tests/conftest.py index c5c3c7c6..becc4af1 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -13,14 +13,11 @@ CacheClientAsync, Configurations, CredentialProvider, - PreviewVectorIndexClient, - PreviewVectorIndexClientAsync, TopicClient, TopicClientAsync, TopicConfigurations, - VectorIndexConfigurations, ) -from momento.config import Configuration, TopicConfiguration, VectorIndexConfiguration +from momento.config import Configuration, TopicConfiguration from momento.typing import ( TCacheName, TDictionaryField, @@ -43,7 +40,6 @@ from tests.utils import ( unique_test_cache_name, - unique_test_vector_index_name, uuid_bytes, uuid_str, ) @@ -54,11 +50,9 @@ TEST_CONFIGURATION = Configurations.Laptop.latest() TEST_TOPIC_CONFIGURATION = TopicConfigurations.Default.latest() -TEST_VECTOR_CONFIGURATION: VectorIndexConfiguration = VectorIndexConfigurations.Default.latest() TEST_AUTH_PROVIDER = CredentialProvider.from_environment_variable("TEST_API_KEY") -TEST_VECTOR_AUTH_PROVIDER = CredentialProvider.from_environment_variable("TEST_API_KEY") TEST_CACHE_NAME: Optional[str] = os.getenv("TEST_CACHE_NAME") @@ -66,10 +60,6 @@ raise RuntimeError("Integration tests require TEST_CACHE_NAME env var; see README for more details.") TEST_TOPIC_NAME: Optional[str] = "my-topic" -TEST_VECTOR_INDEX_NAME: Optional[str] = os.getenv("TEST_VECTOR_INDEX_NAME") -if not TEST_VECTOR_INDEX_NAME: - raise RuntimeError("Integration tests require TEST_VECTOR_INDEX_NAME env var; see README for more details.") -TEST_VECTOR_DIMS = 2 DEFAULT_TTL_SECONDS: timedelta = timedelta(seconds=60) BAD_API_KEY: str = "eyJhbGciOiJIUzUxMiJ9.eyJzdWIiOiJpbnRlZ3JhdGlvbiIsImNwIjoiY29udHJvbC5jZWxsLWFscGhhLWRldi5wcmVwcm9kLmEubW9tZW50b2hxLmNvbSIsImMiOiJjYWNoZS5jZWxsLWFscGhhLWRldi5wcmVwcm9kLmEubW9tZW50b2hxLmNvbSJ9.gdghdjjfjyehhdkkkskskmmls76573jnajhjjjhjdhnndy" # noqa: E501 @@ -101,11 +91,6 @@ def topic_configuration() -> TopicConfiguration: return TEST_TOPIC_CONFIGURATION -@pytest.fixture(scope="session") -def vector_index_configuration() -> VectorIndexConfiguration: - return TEST_VECTOR_CONFIGURATION - - @pytest.fixture(scope="session") def cache_name() -> TCacheName: return cast(str, TEST_CACHE_NAME) @@ -116,16 +101,6 @@ def topic_name() -> TTopicName: return cast(str, TEST_TOPIC_NAME) -@pytest.fixture(scope="session") -def vector_index_name() -> str: - return cast(str, TEST_VECTOR_INDEX_NAME) - - -@pytest.fixture(scope="session") -def vector_index_dimensions() -> int: - return TEST_VECTOR_DIMS - - @pytest.fixture def list_name() -> TListName: return uuid_str() @@ -348,22 +323,6 @@ async def topic_client_async() -> AsyncIterator[TopicClientAsync]: yield _topic_client -@pytest.fixture(scope="session") -def vector_index_client() -> Iterator[PreviewVectorIndexClient]: - with PreviewVectorIndexClient(TEST_VECTOR_CONFIGURATION, TEST_VECTOR_AUTH_PROVIDER) as _client: - yield _client - - -@pytest.fixture(scope="session") -async def vector_index_client_async() -> AsyncIterator[PreviewVectorIndexClientAsync]: - async with PreviewVectorIndexClientAsync(TEST_VECTOR_CONFIGURATION, TEST_VECTOR_AUTH_PROVIDER) as _client: - await _client.create_index(cast(str, TEST_VECTOR_INDEX_NAME), TEST_VECTOR_DIMS) - try: - yield _client - finally: - await _client.delete_index(cast(str, TEST_VECTOR_INDEX_NAME)) - - TUniqueCacheName = Callable[[CacheClient], str] @@ -415,58 +374,3 @@ def _unique_cache_name_async(client: CacheClientAsync) -> str: finally: for cache_name in cache_names: await client_async.delete_cache(cache_name) - - -TUniqueVectorIndexName = Callable[[PreviewVectorIndexClient], str] - - -@pytest.fixture -def unique_vector_index_name( - vector_index_client: PreviewVectorIndexClient, -) -> Iterator[Callable[[PreviewVectorIndexClient], str]]: - """Synchronous version of unique_vector_index_name_async.""" - index_names = [] - - def _unique_vector_index_name(vector_index_client: PreviewVectorIndexClient) -> str: - index_name = unique_test_vector_index_name() - index_names.append(index_name) - return index_name - - try: - yield _unique_vector_index_name - finally: - for index_name in index_names: - vector_index_client.delete_index(index_name) - - -TUniqueVectorIndexNameAsync = Callable[[PreviewVectorIndexClientAsync], str] - - -@pytest_asyncio.fixture -async def unique_vector_index_name_async( - vector_index_client_async: PreviewVectorIndexClientAsync, -) -> AsyncIterator[Callable[[PreviewVectorIndexClientAsync], str]]: - """Returns unique vector index name for testing. - - Also ensures the index is deleted after the test, even if the test fails. - - It does not create the index for you. - - Args: - vector_index_client_async (VectorIndexClientAsync): The client to use to delete the index. - - Returns: - str: the unique index name - """ - index_names = [] - - def _unique_vector_index_name_async(vector_index_client_async: PreviewVectorIndexClientAsync) -> str: - index_name = unique_test_vector_index_name() - index_names.append(index_name) - return index_name - - try: - yield _unique_vector_index_name_async - finally: - for index_name in index_names: - await vector_index_client_async.delete_index(index_name) diff --git a/tests/momento/config/test_vector_index_config.py b/tests/momento/config/test_vector_index_config.py deleted file mode 100644 index 08af2e5a..00000000 --- a/tests/momento/config/test_vector_index_config.py +++ /dev/null @@ -1,97 +0,0 @@ -from pathlib import Path - -import pytest -from momento import ( - CredentialProvider, - PreviewVectorIndexClient, - PreviewVectorIndexClientAsync, - VectorIndexConfigurations, -) -from momento.config import VectorIndexConfiguration -from momento.config.transport.transport_strategy import StaticGrpcConfiguration -from momento.errors.error_details import MomentoErrorCode -from momento.responses.vector_index import ListIndexes, Search - -from tests.utils import unique_test_vector_index_name - - -def _with_root_cert(config: VectorIndexConfiguration, root_cert: bytes) -> VectorIndexConfiguration: - grpc_configuration = StaticGrpcConfiguration( - config.get_transport_strategy().get_grpc_configuration().get_deadline(), root_cert - ) - new_config = config.with_transport_strategy( - config.get_transport_strategy().with_grpc_configuration(grpc_configuration) - ) - return new_config - - -def test_missing_root_cert() -> None: - path = Path("bad") - with pytest.raises(FileNotFoundError) as e: - VectorIndexConfigurations.Default.latest().with_root_certificates_pem(path) - - assert f"Root certificate file not found at path: {path}" in str(e.value) - - -def test_bad_root_cert( - vector_index_configuration: VectorIndexConfiguration, credential_provider: CredentialProvider -) -> None: - root_cert = b"asdfasdf" - - config = _with_root_cert(vector_index_configuration, root_cert) - client = PreviewVectorIndexClient(config, credential_provider) - list_indexes_response = client.list_indexes() - assert isinstance(list_indexes_response, ListIndexes.Error) - assert list_indexes_response.error_code == MomentoErrorCode.SERVER_UNAVAILABLE - - search_response = client.search(unique_test_vector_index_name(), [1, 2]) - assert isinstance(search_response, Search.Error) - assert search_response.error_code == MomentoErrorCode.SERVER_UNAVAILABLE - - -async def test_bad_root_cert_async( - vector_index_configuration: VectorIndexConfiguration, credential_provider: CredentialProvider -) -> None: - root_cert = b"asdfasdf" - config = _with_root_cert(vector_index_configuration, root_cert) - - client = PreviewVectorIndexClientAsync(config, credential_provider) - list_indexes_response = await client.list_indexes() - assert isinstance(list_indexes_response, ListIndexes.Error) - assert list_indexes_response.error_code == MomentoErrorCode.SERVER_UNAVAILABLE - - search_response = await client.search(unique_test_vector_index_name(), [1, 2]) - assert isinstance(search_response, Search.Error) - assert search_response.error_code == MomentoErrorCode.SERVER_UNAVAILABLE - - -# def test_good_root_cert( -# vector_index_configuration: VectorIndexConfiguration, credential_provider: CredentialProvider -# ) -> None: -# # On my machine, this is the path to the root certificates pem file that is cached by the grpc library. -# root_cert_path = Path("./.venv/lib/python3.11/site-packages/grpc/_cython/_credentials/roots.pem") -# config = vector_index_configuration.with_root_certificates_pem(root_cert_path) - -# client = PreviewVectorIndexClient(config, credential_provider) -# list_indexes_response = client.list_indexes() -# assert isinstance(list_indexes_response, ListIndexes.Success) - -# search_response = client.search(unique_test_vector_index_name(), [1, 2]) -# assert isinstance(search_response, Search.Error) -# assert search_response.error_code == MomentoErrorCode.NOT_FOUND_ERROR - - -# async def test_good_root_cert_async( -# vector_index_configuration: VectorIndexConfiguration, credential_provider: CredentialProvider -# ) -> None: -# # On my machine, this is the path to the root certificates pem file that is cached by the grpc library. -# root_cert_path = Path("./.venv/lib/python3.11/site-packages/grpc/_cython/_credentials/roots.pem") -# config = vector_index_configuration.with_root_certificates_pem(root_cert_path) - -# client = PreviewVectorIndexClientAsync(config, credential_provider) -# list_indexes_response = await client.list_indexes() -# assert isinstance(list_indexes_response, ListIndexes.Success) - -# search_response = await client.search(unique_test_vector_index_name(), [1, 2]) -# assert isinstance(search_response, Search.Error) -# assert search_response.error_code == MomentoErrorCode.NOT_FOUND_ERROR diff --git a/tests/momento/requests/vector_index/__init__.py b/tests/momento/requests/vector_index/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/tests/momento/requests/vector_index/test_filters.py b/tests/momento/requests/vector_index/test_filters.py deleted file mode 100644 index 10063519..00000000 --- a/tests/momento/requests/vector_index/test_filters.py +++ /dev/null @@ -1,140 +0,0 @@ -import pytest -from momento.requests.vector_index import Field, FilterExpression -from momento.requests.vector_index import filters as F -from momento_wire_types import vectorindex_pb2 as pb - - -@pytest.mark.parametrize( - ["field_based_expression", "full_expression", "protobuf"], - [ - ( - Field("foo") == "bar", - F.Equals("foo", "bar"), - pb._FilterExpression(equals_expression=pb._EqualsExpression(field="foo", string_value="bar")), - ), - ( - Field("foo"), - F.Equals("foo", True), - pb._FilterExpression(equals_expression=pb._EqualsExpression(field="foo", boolean_value=True)), - ), - ( - ~Field("foo"), - F.Equals("foo", False), - pb._FilterExpression(equals_expression=pb._EqualsExpression(field="foo", boolean_value=False)), - ), - ( - Field("foo") == 5.5, - F.Equals("foo", 5.5), - pb._FilterExpression(equals_expression=pb._EqualsExpression(field="foo", float_value=5.5)), - ), - ( - Field("foo") != "bar", - F.Not(F.Equals("foo", "bar")), - pb._FilterExpression( - not_expression=pb._NotExpression( - expression_to_negate=pb._FilterExpression( - equals_expression=pb._EqualsExpression(field="foo", string_value="bar") - ) - ) - ), - ), - ( - ~(Field("foo") == "bar"), - F.Not(F.Equals("foo", "bar")), - pb._FilterExpression( - not_expression=pb._NotExpression( - expression_to_negate=pb._FilterExpression( - equals_expression=pb._EqualsExpression(field="foo", string_value="bar") - ) - ) - ), - ), - ( - Field("foo") > 5, - F.GreaterThan("foo", 5), - pb._FilterExpression(greater_than_expression=pb._GreaterThanExpression(field="foo", integer_value=5)), - ), - ( - Field("foo") > 5.5, - F.GreaterThan("foo", 5.5), - pb._FilterExpression(greater_than_expression=pb._GreaterThanExpression(field="foo", float_value=5.5)), - ), - ( - Field("foo") >= 5, - F.GreaterThanOrEqual("foo", 5), - pb._FilterExpression( - greater_than_or_equal_expression=pb._GreaterThanOrEqualExpression(field="foo", integer_value=5) - ), - ), - ( - Field("foo") >= 5.5, - F.GreaterThanOrEqual("foo", 5.5), - pb._FilterExpression( - greater_than_or_equal_expression=pb._GreaterThanOrEqualExpression(field="foo", float_value=5.5) - ), - ), - ( - Field("foo") < 5, - F.LessThan("foo", 5), - pb._FilterExpression(less_than_expression=pb._LessThanExpression(field="foo", integer_value=5)), - ), - ( - Field("foo") < 5.5, - F.LessThan("foo", 5.5), - pb._FilterExpression(less_than_expression=pb._LessThanExpression(field="foo", float_value=5.5)), - ), - ( - Field("foo") <= 5, - F.LessThanOrEqual("foo", 5), - pb._FilterExpression( - less_than_or_equal_expression=pb._LessThanOrEqualExpression(field="foo", integer_value=5) - ), - ), - ( - Field("foo") <= 5.5, - F.LessThanOrEqual("foo", 5.5), - pb._FilterExpression( - less_than_or_equal_expression=pb._LessThanOrEqualExpression(field="foo", float_value=5.5) - ), - ), - ( - Field("foo").list_contains("bar"), - F.ListContains("foo", "bar"), - pb._FilterExpression(list_contains_expression=pb._ListContainsExpression(field="foo", string_value="bar")), - ), - ( - (Field("foo") == "bar") & (Field("baz") == 5), - F.And(F.Equals("foo", "bar"), F.Equals("baz", 5)), - pb._FilterExpression( - and_expression=pb._AndExpression( - first_expression=pb._FilterExpression( - equals_expression=pb._EqualsExpression(field="foo", string_value="bar") - ), - second_expression=pb._FilterExpression( - equals_expression=pb._EqualsExpression(field="baz", integer_value=5) - ), - ) - ), - ), - ( - (Field("foo") == "bar") | (Field("baz") == 5), - F.Or(F.Equals("foo", "bar"), F.Equals("baz", 5)), - pb._FilterExpression( - or_expression=pb._OrExpression( - first_expression=pb._FilterExpression( - equals_expression=pb._EqualsExpression(field="foo", string_value="bar") - ), - second_expression=pb._FilterExpression( - equals_expression=pb._EqualsExpression(field="baz", integer_value=5) - ), - ) - ), - ), - ], -) -def test_field_shorthand( - field_based_expression: FilterExpression, full_expression: FilterExpression, protobuf: pb._FilterExpression -) -> None: - assert field_based_expression == full_expression - assert field_based_expression.to_filter_expression_proto() == protobuf - assert full_expression.to_filter_expression_proto() == protobuf diff --git a/tests/momento/requests/vector_index/test_item.py b/tests/momento/requests/vector_index/test_item.py deleted file mode 100644 index b77f2415..00000000 --- a/tests/momento/requests/vector_index/test_item.py +++ /dev/null @@ -1,57 +0,0 @@ -import pytest -from momento.errors.exceptions import InvalidArgumentException -from momento.internal.services import Service -from momento.requests.vector_index import Item - - -def test_serialize_item_with_diverse_metadata() -> None: - item = Item( - id="id", - vector=[1, 2, 3], - metadata={ - "string_key": "value", - "int_key": 1, - "double_key": 3.14, - "bool_key": True, - "list_of_strings_key": ["one", "two"], - }, - ).to_proto() - assert item.id == "id" - assert item.vector.elements == [1, 2, 3] - - assert item.metadata[0].field == "string_key" - assert item.metadata[0].string_value == "value" - - assert item.metadata[1].field == "int_key" - assert item.metadata[1].integer_value == 1 - - assert item.metadata[2].field == "double_key" - assert item.metadata[2].double_value == 3.14 - - assert item.metadata[3].field == "bool_key" - assert item.metadata[3].boolean_value is True - - assert item.metadata[4].field == "list_of_strings_key" - assert item.metadata[4].list_of_strings_value.values == ["one", "two"] - - -def test_serialize_item_with_bad_metadata() -> None: - with pytest.raises(InvalidArgumentException) as exc_info: - Item(id="id", vector=[1, 2, 3], metadata={"key": {"nested_key": "nested_value"}}).to_proto() # type: ignore - - assert exc_info.value.service == Service.INDEX - assert ( - exc_info.value.message - == "Metadata values must be either str, int, float, bool, or list[str]. Field 'key' has a value of type with value {'nested_key': 'nested_value'}." # noqa: E501 W503 - ) - - -def test_serialize_item_with_null_metadata() -> None: - with pytest.raises(InvalidArgumentException) as exc_info: - Item(id="id", vector=[1, 2, 3], metadata={"key": None}).to_proto() # type: ignore - - assert exc_info.value.service == Service.INDEX - assert ( - exc_info.value.message - == "Metadata values must be either str, int, float, bool, or list[str]. Field 'key' has a value of type with value None." # noqa: E501 W503 - ) diff --git a/tests/momento/vector_index_client/__init__.py b/tests/momento/vector_index_client/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/tests/momento/vector_index_client/test_control.py b/tests/momento/vector_index_client/test_control.py deleted file mode 100644 index 9cb3c505..00000000 --- a/tests/momento/vector_index_client/test_control.py +++ /dev/null @@ -1,185 +0,0 @@ -import pytest -from momento import CredentialProvider, PreviewVectorIndexClient -from momento.config import VectorIndexConfiguration -from momento.errors import MomentoErrorCode -from momento.requests.vector_index import SimilarityMetric -from momento.responses.vector_index import ( - CreateIndex, - DeleteIndex, - IndexInfo, - ListIndexes, -) - -from tests.conftest import TUniqueVectorIndexName -from tests.utils import unique_test_vector_index_name - - -def test_create_index_list_indexes_and_delete_index( - vector_index_client: PreviewVectorIndexClient, - unique_vector_index_name: TUniqueVectorIndexName, - vector_index_dimensions: int, -) -> None: - new_index_name = unique_vector_index_name(vector_index_client) - - create_index_response = vector_index_client.create_index(new_index_name, num_dimensions=vector_index_dimensions) - assert isinstance(create_index_response, CreateIndex.Success) - - list_indexes_response = vector_index_client.list_indexes() - assert isinstance(list_indexes_response, ListIndexes.Success) - assert any( - IndexInfo(new_index_name, vector_index_dimensions, SimilarityMetric.COSINE_SIMILARITY) == index - for index in list_indexes_response.indexes - ) - - delete_index_response = vector_index_client.delete_index(new_index_name) - assert isinstance(delete_index_response, DeleteIndex.Success) - - -def test_create_index_already_exists_when_creating_existing_index( - vector_index_client: PreviewVectorIndexClient, - unique_vector_index_name: TUniqueVectorIndexName, - vector_index_dimensions: int, -) -> None: - new_index_name = unique_vector_index_name(vector_index_client) - - response = vector_index_client.create_index(new_index_name, num_dimensions=vector_index_dimensions) - assert isinstance(response, CreateIndex.Success) - - response = vector_index_client.create_index(new_index_name, num_dimensions=vector_index_dimensions) - assert isinstance(response, CreateIndex.IndexAlreadyExists) - - del_response = vector_index_client.delete_index(new_index_name) - assert isinstance(del_response, DeleteIndex.Success) - - -def test_create_index_returns_error_for_bad_name( - vector_index_client: PreviewVectorIndexClient, -) -> None: - for bad_name, reason in [("", "not be empty"), (None, "be a string"), (1, "be a string")]: - response = vector_index_client.create_index(bad_name, num_dimensions=1) # type: ignore[arg-type] - assert isinstance(response, CreateIndex.Error) - assert response.error_code == MomentoErrorCode.INVALID_ARGUMENT_ERROR - assert response.inner_exception.message == f"Vector index name must {reason}" - assert response.message == f"Invalid argument passed to Momento client: {response.inner_exception.message}" - - -def test_create_index_returns_error_for_bad_num_dimensions( - vector_index_client: PreviewVectorIndexClient, - unique_vector_index_name: TUniqueVectorIndexName, -) -> None: - for bad_num_dimensions in [0, 1.1]: - response = vector_index_client.create_index( - unique_vector_index_name(vector_index_client), - num_dimensions=bad_num_dimensions, # type: ignore[arg-type] - ) - assert isinstance(response, CreateIndex.Error) - assert response.error_code == MomentoErrorCode.INVALID_ARGUMENT_ERROR - assert response.inner_exception.message == "Number of dimensions must be a positive integer." - assert response.message == f"Invalid argument passed to Momento client: {response.inner_exception.message}" - - -def test_create_index_returns_error_for_bad_similarity_metric( - vector_index_client: PreviewVectorIndexClient, -) -> None: - response = vector_index_client.create_index( - index_name="vector-index", - num_dimensions=2, - similarity_metric="ASDF", # type: ignore[arg-type] - ) - assert isinstance(response, CreateIndex.Error) - assert response.error_code == MomentoErrorCode.INVALID_ARGUMENT_ERROR - assert response.inner_exception.message == "Invalid similarity metric `ASDF`" - assert response.message == f"Invalid argument passed to Momento client: {response.inner_exception.message}" - - -# Delete index -def test_delete_index_succeeds(vector_index_client: PreviewVectorIndexClient, vector_index_dimensions: int) -> None: - index_name = unique_test_vector_index_name() - - response = vector_index_client.create_index(index_name, vector_index_dimensions) - assert isinstance(response, CreateIndex.Success) - - delete_response = vector_index_client.delete_index(index_name) - assert isinstance(delete_response, DeleteIndex.Success) - - delete_response = vector_index_client.delete_index(index_name) - assert isinstance(delete_response, DeleteIndex.Error) - assert delete_response.error_code == MomentoErrorCode.NOT_FOUND_ERROR - - -def test_delete_index_returns_not_found_error_when_deleting_unknown_index( - vector_index_client: PreviewVectorIndexClient, -) -> None: - index_name = unique_test_vector_index_name() - response = vector_index_client.delete_index(index_name) - assert isinstance(response, DeleteIndex.Error) - assert response.error_code == MomentoErrorCode.NOT_FOUND_ERROR - assert response.inner_exception.message == f'Index with name "{index_name}" does not exist' - print(response.message) - expected_resp_message = ( - f"A index with the specified name does not exist. To resolve this error, make sure you " - f"have created the index before attempting to use it: {response.inner_exception.message}" - ) - assert response.message == expected_resp_message - - -def test_delete_index_returns_error_for_bad_name( - vector_index_client: PreviewVectorIndexClient, -) -> None: - for bad_name, reason in [("", "not be empty"), (None, "be a string"), (1, "be a string")]: - response = vector_index_client.delete_index(bad_name) # type: ignore[arg-type] - assert isinstance(response, DeleteIndex.Error) - assert response.error_code == MomentoErrorCode.INVALID_ARGUMENT_ERROR - assert response.inner_exception.message == f"Vector index name must {reason}" - assert response.message == f"Invalid argument passed to Momento client: {response.inner_exception.message}" - - -def test_create_index_throws_authentication_exception_for_bad_token( - bad_token_credential_provider: CredentialProvider, - vector_index_configuration: VectorIndexConfiguration, - vector_index_dimensions: int, -) -> None: - index_name = unique_test_vector_index_name() - - with PreviewVectorIndexClient(vector_index_configuration, bad_token_credential_provider) as vector_index_client: - response = vector_index_client.create_index(index_name, num_dimensions=2) - assert isinstance(response, CreateIndex.Error) - assert response.error_code == MomentoErrorCode.AUTHENTICATION_ERROR - assert response.inner_exception.message == "Could not validate authorization token" - assert ( - response.message - == "Invalid authentication credentials to connect to index service: Could not validate authorization token" - ) - - -# List indexes -@pytest.mark.parametrize( - "num_dimensions, similarity_metric", - [ - (1, SimilarityMetric.COSINE_SIMILARITY), - (2, SimilarityMetric.EUCLIDEAN_SIMILARITY), - (3, SimilarityMetric.INNER_PRODUCT), - ], -) -def test_list_indexes_succeeds( - vector_index_client: PreviewVectorIndexClient, num_dimensions: int, similarity_metric: SimilarityMetric -) -> None: - index_name = unique_test_vector_index_name() - - initial_response = vector_index_client.list_indexes() - assert isinstance(initial_response, ListIndexes.Success) - - index_names = [index.name for index in initial_response.indexes] - assert index_name not in index_names - - try: - response = vector_index_client.create_index(index_name, num_dimensions, similarity_metric) - assert isinstance(response, CreateIndex.Success) - - list_cache_resp = vector_index_client.list_indexes() - assert isinstance(list_cache_resp, ListIndexes.Success) - - assert IndexInfo(index_name, num_dimensions, similarity_metric) in list_cache_resp.indexes - finally: - delete_response = vector_index_client.delete_index(index_name) - assert isinstance(delete_response, DeleteIndex.Success) diff --git a/tests/momento/vector_index_client/test_control_async.py b/tests/momento/vector_index_client/test_control_async.py deleted file mode 100644 index ee045f08..00000000 --- a/tests/momento/vector_index_client/test_control_async.py +++ /dev/null @@ -1,191 +0,0 @@ -import pytest -from momento import CredentialProvider, PreviewVectorIndexClientAsync -from momento.config import VectorIndexConfiguration -from momento.errors import MomentoErrorCode -from momento.requests.vector_index import SimilarityMetric -from momento.responses.vector_index import ( - CreateIndex, - DeleteIndex, - IndexInfo, - ListIndexes, -) - -from tests.conftest import TUniqueVectorIndexNameAsync -from tests.utils import unique_test_vector_index_name - - -async def test_create_index_list_indexes_and_delete_index( - vector_index_client_async: PreviewVectorIndexClientAsync, - unique_vector_index_name_async: TUniqueVectorIndexNameAsync, - vector_index_dimensions: int, -) -> None: - new_index_name = unique_vector_index_name_async(vector_index_client_async) - - create_index_response = await vector_index_client_async.create_index( - new_index_name, num_dimensions=vector_index_dimensions - ) - assert isinstance(create_index_response, CreateIndex.Success) - - list_indexes_response = await vector_index_client_async.list_indexes() - assert isinstance(list_indexes_response, ListIndexes.Success) - assert any( - IndexInfo(new_index_name, vector_index_dimensions, SimilarityMetric.COSINE_SIMILARITY) == index - for index in list_indexes_response.indexes - ) - - delete_index_response = await vector_index_client_async.delete_index(new_index_name) - assert isinstance(delete_index_response, DeleteIndex.Success) - - -async def test_create_index_already_exists_when_creating_existing_index( - vector_index_client_async: PreviewVectorIndexClientAsync, - unique_vector_index_name: TUniqueVectorIndexNameAsync, - vector_index_dimensions: int, -) -> None: - new_index_name = unique_vector_index_name(vector_index_client_async) - - response = await vector_index_client_async.create_index(new_index_name, num_dimensions=vector_index_dimensions) - assert isinstance(response, CreateIndex.Success) - - response = await vector_index_client_async.create_index(new_index_name, num_dimensions=vector_index_dimensions) - assert isinstance(response, CreateIndex.IndexAlreadyExists) - - del_response = await vector_index_client_async.delete_index(new_index_name) - assert isinstance(del_response, DeleteIndex.Success) - - -async def test_create_index_returns_error_for_bad_name( - vector_index_client_async: PreviewVectorIndexClientAsync, -) -> None: - for bad_name, reason in [("", "not be empty"), (None, "be a string"), (1, "be a string")]: - response = await vector_index_client_async.create_index(bad_name, num_dimensions=1) # type: ignore[arg-type] - assert isinstance(response, CreateIndex.Error) - assert response.error_code == MomentoErrorCode.INVALID_ARGUMENT_ERROR - assert response.inner_exception.message == f"Vector index name must {reason}" - assert response.message == f"Invalid argument passed to Momento client: {response.inner_exception.message}" - - -async def test_create_index_returns_error_for_bad_num_dimensions( - vector_index_client_async: PreviewVectorIndexClientAsync, - unique_vector_index_name_async: TUniqueVectorIndexNameAsync, -) -> None: - for bad_num_dimensions in [0, 1.1]: - response = await vector_index_client_async.create_index( - unique_vector_index_name_async(vector_index_client_async), - num_dimensions=bad_num_dimensions, # type: ignore[arg-type] - ) - assert isinstance(response, CreateIndex.Error) - assert response.error_code == MomentoErrorCode.INVALID_ARGUMENT_ERROR - assert response.inner_exception.message == "Number of dimensions must be a positive integer." - assert response.message == f"Invalid argument passed to Momento client: {response.inner_exception.message}" - - -async def test_create_index_returns_error_for_bad_similarity_metric( - vector_index_client_async: PreviewVectorIndexClientAsync, -) -> None: - response = await vector_index_client_async.create_index( - index_name="vector-index", - num_dimensions=2, - similarity_metric="ASDF", # type: ignore[arg-type] - ) - assert isinstance(response, CreateIndex.Error) - assert response.error_code == MomentoErrorCode.INVALID_ARGUMENT_ERROR - assert response.inner_exception.message == "Invalid similarity metric `ASDF`" - assert response.message == f"Invalid argument passed to Momento client: {response.inner_exception.message}" - - -# Delete index -async def test_delete_index_succeeds( - vector_index_client_async: PreviewVectorIndexClientAsync, vector_index_dimensions: int -) -> None: - index_name = unique_test_vector_index_name() - - response = await vector_index_client_async.create_index(index_name, vector_index_dimensions) - assert isinstance(response, CreateIndex.Success) - - delete_response = await vector_index_client_async.delete_index(index_name) - assert isinstance(delete_response, DeleteIndex.Success) - - delete_response = await vector_index_client_async.delete_index(index_name) - assert isinstance(delete_response, DeleteIndex.Error) - assert delete_response.error_code == MomentoErrorCode.NOT_FOUND_ERROR - - -async def test_delete_index_returns_not_found_error_when_deleting_unknown_index( - vector_index_client_async: PreviewVectorIndexClientAsync, -) -> None: - index_name = unique_test_vector_index_name() - response = await vector_index_client_async.delete_index(index_name) - assert isinstance(response, DeleteIndex.Error) - assert response.error_code == MomentoErrorCode.NOT_FOUND_ERROR - assert response.inner_exception.message == f'Index with name "{index_name}" does not exist' - print(response.message) - expected_resp_message = ( - f"A index with the specified name does not exist. To resolve this error, make sure you " - f"have created the index before attempting to use it: {response.inner_exception.message}" - ) - assert response.message == expected_resp_message - - -async def test_delete_index_returns_error_for_bad_name( - vector_index_client_async: PreviewVectorIndexClientAsync, -) -> None: - for bad_name, reason in [("", "not be empty"), (None, "be a string"), (1, "be a string")]: - response = await vector_index_client_async.delete_index(bad_name) # type: ignore[arg-type] - assert isinstance(response, DeleteIndex.Error) - assert response.error_code == MomentoErrorCode.INVALID_ARGUMENT_ERROR - assert response.inner_exception.message == f"Vector index name must {reason}" - assert response.message == f"Invalid argument passed to Momento client: {response.inner_exception.message}" - - -async def test_create_index_throws_authentication_exception_for_bad_token( - bad_token_credential_provider: CredentialProvider, - vector_index_configuration: VectorIndexConfiguration, - vector_index_dimensions: int, -) -> None: - index_name = unique_test_vector_index_name() - - async with PreviewVectorIndexClientAsync( - vector_index_configuration, bad_token_credential_provider - ) as vector_index_client: - response = await vector_index_client.create_index(index_name, num_dimensions=2) - assert isinstance(response, CreateIndex.Error) - assert response.error_code == MomentoErrorCode.AUTHENTICATION_ERROR - assert response.inner_exception.message == "Could not validate authorization token" - assert ( - response.message - == "Invalid authentication credentials to connect to index service: Could not validate authorization token" - ) - - -# List indexes -@pytest.mark.parametrize( - "num_dimensions, similarity_metric", - [ - (1, SimilarityMetric.COSINE_SIMILARITY), - (2, SimilarityMetric.EUCLIDEAN_SIMILARITY), - (3, SimilarityMetric.INNER_PRODUCT), - ], -) -async def test_list_indexes_succeeds( - vector_index_client_async: PreviewVectorIndexClientAsync, num_dimensions: int, similarity_metric: SimilarityMetric -) -> None: - index_name = unique_test_vector_index_name() - - initial_response = await vector_index_client_async.list_indexes() - assert isinstance(initial_response, ListIndexes.Success) - - index_names = [index.name for index in initial_response.indexes] - assert index_name not in index_names - - try: - response = await vector_index_client_async.create_index(index_name, num_dimensions, similarity_metric) - assert isinstance(response, CreateIndex.Success) - - list_cache_resp = await vector_index_client_async.list_indexes() - assert isinstance(list_cache_resp, ListIndexes.Success) - - assert IndexInfo(index_name, num_dimensions, similarity_metric) in list_cache_resp.indexes - finally: - delete_response = await vector_index_client_async.delete_index(index_name) - assert isinstance(delete_response, DeleteIndex.Success) diff --git a/tests/momento/vector_index_client/test_data.py b/tests/momento/vector_index_client/test_data.py deleted file mode 100644 index 1fe671b3..00000000 --- a/tests/momento/vector_index_client/test_data.py +++ /dev/null @@ -1,908 +0,0 @@ -from __future__ import annotations - -from typing import Optional, cast - -import pytest -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, filters -from momento.responses.vector_index import ( - CountItems, - CreateIndex, - DeleteItemBatch, - GetItemBatch, - GetItemMetadataBatch, - Search, - SearchAndFetchVectors, - SearchHit, - UpsertItemBatch, -) - -from tests.conftest import TUniqueVectorIndexName -from tests.utils import sleep, uuid_str, when_fetching_vectors_apply_vectors_to_hits - - -def test_create_index_with_inner_product_upsert_item_search_happy_path( - vector_index_client: PreviewVectorIndexClient, - unique_vector_index_name: TUniqueVectorIndexName, -) -> None: - index_name = unique_vector_index_name(vector_index_client) - create_response = vector_index_client.create_index( - index_name, num_dimensions=2, similarity_metric=SimilarityMetric.INNER_PRODUCT - ) - assert isinstance(create_response, CreateIndex.Success) - - upsert_response = vector_index_client.upsert_item_batch(index_name, items=[Item(id="test_item", vector=[1.0, 2.0])]) - assert isinstance(upsert_response, UpsertItemBatch.Success) - - sleep(2) - - search_response = vector_index_client.search(index_name, query_vector=[1.0, 2.0], top_k=1) - assert isinstance(search_response, Search.Success) - assert len(search_response.hits) == 1 - assert search_response.hits[0].id == "test_item" - assert search_response.hits[0].score == 5.0 - - -@pytest.mark.parametrize("similarity_metric", [SimilarityMetric.COSINE_SIMILARITY, None]) -def test_create_index_with_cosine_similarity_upsert_item_search_happy_path( - vector_index_client: PreviewVectorIndexClient, - unique_vector_index_name: TUniqueVectorIndexName, - similarity_metric: Optional[SimilarityMetric], -) -> None: - index_name = unique_vector_index_name(vector_index_client) - num_dimensions = 2 - if similarity_metric is not None: - create_response = vector_index_client.create_index( - index_name, num_dimensions=num_dimensions, similarity_metric=similarity_metric - ) - else: - create_response = vector_index_client.create_index(index_name, num_dimensions=num_dimensions) - assert isinstance(create_response, CreateIndex.Success) - - upsert_response = vector_index_client.upsert_item_batch( - index_name, - items=[ - Item(id="test_item_1", vector=[1.0, 1.0]), - Item(id="test_item_2", vector=[-1.0, 1.0]), - Item(id="test_item_3", vector=[-1.0, -1.0]), - ], - ) - assert isinstance(upsert_response, UpsertItemBatch.Success) - - sleep(2) - - search_response = vector_index_client.search(index_name, query_vector=[2.0, 2.0], top_k=3) - assert isinstance(search_response, Search.Success) - assert search_response.hits == [ - SearchHit(id="test_item_1", score=1.0), - SearchHit(id="test_item_2", score=0.0), - SearchHit(id="test_item_3", score=-1.0), - ] - - -def test_create_index_with_euclidean_similarity_upsert_item_search_happy_path( - vector_index_client: PreviewVectorIndexClient, - unique_vector_index_name: TUniqueVectorIndexName, -) -> None: - index_name = unique_vector_index_name(vector_index_client) - create_response = vector_index_client.create_index( - index_name, num_dimensions=2, similarity_metric=SimilarityMetric.EUCLIDEAN_SIMILARITY - ) - assert isinstance(create_response, CreateIndex.Success) - - upsert_response = vector_index_client.upsert_item_batch( - index_name, - items=[ - Item(id="test_item_1", vector=[1.0, 1.0]), - Item(id="test_item_2", vector=[-1.0, 1.0]), - Item(id="test_item_3", vector=[-1.0, -1.0]), - ], - ) - assert isinstance(upsert_response, UpsertItemBatch.Success) - - sleep(2) - - search_response = vector_index_client.search(index_name, query_vector=[1.0, 1.0], top_k=3) - assert isinstance(search_response, Search.Success) - assert search_response.hits == [ - SearchHit(id="test_item_1", score=0.0), - SearchHit(id="test_item_2", score=4.0), - SearchHit(id="test_item_3", score=8.0), - ] - - -def test_create_index_upsert_multiple_items_search_happy_path( - vector_index_client: PreviewVectorIndexClient, - unique_vector_index_name: TUniqueVectorIndexName, -) -> None: - index_name = unique_vector_index_name(vector_index_client) - create_response = vector_index_client.create_index( - index_name, num_dimensions=2, similarity_metric=SimilarityMetric.INNER_PRODUCT - ) - assert isinstance(create_response, CreateIndex.Success) - - upsert_response = vector_index_client.upsert_item_batch( - index_name, - items=[ - Item(id="test_item_1", vector=[1.0, 2.0]), - Item(id="test_item_2", vector=[3.0, 4.0]), - Item(id="test_item_3", vector=[5.0, 6.0]), - ], - ) - assert isinstance(upsert_response, UpsertItemBatch.Success) - - sleep(2) - - search_response = vector_index_client.search(index_name, query_vector=[1.0, 2.0], top_k=3) - assert isinstance(search_response, Search.Success) - assert len(search_response.hits) == 3 - - assert search_response.hits == [ - SearchHit(id="test_item_3", score=17.0), - SearchHit(id="test_item_2", score=11.0), - SearchHit(id="test_item_1", score=5.0), - ] - - -# A note on the parameterized search tests: -# The search tests are parameterized to test both the search and search_and_fetch_vectors methods. -# We pass both the name of the specific search method and the response type. -# The tests will run for both search methods, and the response type will be used to assert the type of the response. -# The vectors are attached to the hits if the response type is SearchAndFetchVectors.Success. -@pytest.mark.parametrize( - ["search_method_name", "response"], - [("search", Search.Success), ("search_and_fetch_vectors", SearchAndFetchVectors.Success)], -) -def test_create_index_upsert_multiple_items_search_with_top_k_happy_path( - vector_index_client: PreviewVectorIndexClient, - unique_vector_index_name: TUniqueVectorIndexName, - search_method_name: str, - response: type[Search.Success] | type[SearchAndFetchVectors.Success], -) -> None: - index_name = unique_vector_index_name(vector_index_client) - create_response = vector_index_client.create_index( - index_name, num_dimensions=2, similarity_metric=SimilarityMetric.INNER_PRODUCT - ) - assert isinstance(create_response, CreateIndex.Success) - - items = [ - Item(id="test_item_1", vector=[1.0, 2.0]), - Item(id="test_item_2", vector=[3.0, 4.0]), - Item(id="test_item_3", vector=[5.0, 6.0]), - ] - upsert_response = vector_index_client.upsert_item_batch( - index_name, - items=items, - ) - assert isinstance(upsert_response, UpsertItemBatch.Success) - - sleep(2) - - search = getattr(vector_index_client, search_method_name) - search_response = search(index_name, query_vector=[1.0, 2.0], top_k=2) - assert isinstance(search_response, response) - assert len(search_response.hits) == 2 - - hits = [SearchHit(id="test_item_3", score=17.0), SearchHit(id="test_item_2", score=11.0)] - hits = when_fetching_vectors_apply_vectors_to_hits(search_response, hits, items) - assert search_response.hits == hits - - -@pytest.mark.parametrize( - ["search_method_name", "response"], - [("search", Search.Success), ("search_and_fetch_vectors", SearchAndFetchVectors.Success)], -) -def test_upsert_and_search_with_metadata_happy_path( - vector_index_client: PreviewVectorIndexClient, - unique_vector_index_name: TUniqueVectorIndexName, - search_method_name: str, - response: type[Search.Success] | type[SearchAndFetchVectors.Success], -) -> None: - index_name = unique_vector_index_name(vector_index_client) - create_response = vector_index_client.create_index( - index_name, num_dimensions=2, similarity_metric=SimilarityMetric.INNER_PRODUCT - ) - assert isinstance(create_response, CreateIndex.Success) - - items = [ - Item(id="test_item_1", vector=[1.0, 2.0], metadata={"key1": "value1"}), - Item(id="test_item_2", vector=[3.0, 4.0], metadata={"key2": "value2"}), - Item(id="test_item_3", vector=[5.0, 6.0], metadata={"key1": "value3", "key3": "value3"}), - ] - upsert_response = vector_index_client.upsert_item_batch( - index_name, - items=items, - ) - assert isinstance(upsert_response, UpsertItemBatch.Success) - - sleep(2) - - search = getattr(vector_index_client, search_method_name) - search_response = search(index_name, query_vector=[1.0, 2.0], top_k=3) - assert isinstance(search_response, response) - assert len(search_response.hits) == 3 - - hits = [ - SearchHit(id="test_item_3", score=17.0), - SearchHit(id="test_item_2", score=11.0), - SearchHit(id="test_item_1", score=5.0), - ] - hits = when_fetching_vectors_apply_vectors_to_hits(search_response, hits, items) - assert search_response.hits == hits - - search_response = search(index_name, query_vector=[1.0, 2.0], top_k=3, metadata_fields=["key1"]) - assert isinstance(search_response, response) - assert len(search_response.hits) == 3 - - hits = [ - SearchHit(id="test_item_3", score=17.0, metadata={"key1": "value3"}), - SearchHit(id="test_item_2", score=11.0, metadata={}), - SearchHit(id="test_item_1", score=5.0, metadata={"key1": "value1"}), - ] - hits = when_fetching_vectors_apply_vectors_to_hits(search_response, hits, items) - assert search_response.hits == hits - - search_response = search( - index_name, query_vector=[1.0, 2.0], top_k=3, metadata_fields=["key1", "key2", "key3", "key4"] - ) - assert isinstance(search_response, response) - assert len(search_response.hits) == 3 - - hits = [ - SearchHit(id="test_item_3", score=17.0, metadata={"key1": "value3", "key3": "value3"}), - SearchHit(id="test_item_2", score=11.0, metadata={"key2": "value2"}), - SearchHit(id="test_item_1", score=5.0, metadata={"key1": "value1"}), - ] - hits = when_fetching_vectors_apply_vectors_to_hits(search_response, hits, items) - assert search_response.hits == hits - - -def test_upsert_with_bad_metadata(vector_index_client: PreviewVectorIndexClient) -> None: - response = vector_index_client.upsert_item_batch( - index_name="test_index", - items=[Item(id="test_item", vector=[1.0, 2.0], metadata={"key": {"subkey": "subvalue"}})], # type: ignore - ) - assert isinstance(response, UpsertItemBatch.Error) - assert response.error_code == MomentoErrorCode.INVALID_ARGUMENT_ERROR - assert ( - response.message - == "Invalid argument passed to Momento client: Metadata values must be either str, int, float, bool, or list[str]. Field 'key' has a value of type with value {'subkey': 'subvalue'}." # noqa: E501 W503 - ) - - -@pytest.mark.parametrize( - ["search_method_name", "response"], - [("search", Search.Success), ("search_and_fetch_vectors", SearchAndFetchVectors.Success)], -) -def test_upsert_and_search_with_all_metadata_happy_path( - vector_index_client: PreviewVectorIndexClient, - unique_vector_index_name: TUniqueVectorIndexName, - search_method_name: str, - response: type[Search.Success] | type[SearchAndFetchVectors.Success], -) -> None: - index_name = unique_vector_index_name(vector_index_client) - create_response = vector_index_client.create_index( - index_name, num_dimensions=2, similarity_metric=SimilarityMetric.INNER_PRODUCT - ) - assert isinstance(create_response, CreateIndex.Success) - - items = [ - Item(id="test_item_1", vector=[1.0, 2.0], metadata={"key1": "value1"}), - Item(id="test_item_2", vector=[3.0, 4.0], metadata={"key2": "value2"}), - Item(id="test_item_3", vector=[5.0, 6.0], metadata={"key1": "value3", "key3": "value3"}), - ] - upsert_response = vector_index_client.upsert_item_batch( - index_name, - items=items, - ) - assert isinstance(upsert_response, UpsertItemBatch.Success) - - sleep(2) - - search = getattr(vector_index_client, search_method_name) - search_response = search(index_name, query_vector=[1.0, 2.0], top_k=3) - assert isinstance(search_response, response) - assert len(search_response.hits) == 3 - - hits = [ - SearchHit(id="test_item_3", score=17.0), - SearchHit(id="test_item_2", score=11.0), - SearchHit(id="test_item_1", score=5.0), - ] - hits = when_fetching_vectors_apply_vectors_to_hits(search_response, hits, items) - assert search_response.hits == hits - - search_response = search(index_name, query_vector=[1.0, 2.0], top_k=3, metadata_fields=ALL_METADATA) - assert isinstance(search_response, response) - assert len(search_response.hits) == 3 - - hits = [ - SearchHit(id="test_item_3", score=17.0, metadata={"key1": "value3", "key3": "value3"}), - SearchHit(id="test_item_2", score=11.0, metadata={"key2": "value2"}), - SearchHit(id="test_item_1", score=5.0, metadata={"key1": "value1"}), - ] - hits = when_fetching_vectors_apply_vectors_to_hits(search_response, hits, items) - assert search_response.hits == hits - - -@pytest.mark.parametrize( - ["search_method_name", "response"], - [ - ("search", Search.Success), - ( - "search_and_fetch_vectors", - SearchAndFetchVectors.Success, - ), - ], -) -def test_upsert_and_search_with_diverse_metadata_happy_path( - vector_index_client: PreviewVectorIndexClient, - unique_vector_index_name: TUniqueVectorIndexName, - search_method_name: str, - response: type[Search.Success] | type[SearchAndFetchVectors.Success], -) -> None: - index_name = unique_vector_index_name(vector_index_client) - create_response = vector_index_client.create_index( - index_name, num_dimensions=2, similarity_metric=SimilarityMetric.INNER_PRODUCT - ) - assert isinstance(create_response, CreateIndex.Success) - - metadata: Metadata = { - "string": "value", - "bool": True, - "int": 1, - "float": 3.14, - "list": ["a", "b", "c"], - "empty_list": [], - } - items = [ - Item(id="test_item_1", vector=[1.0, 2.0], metadata=metadata), - ] - upsert_response = vector_index_client.upsert_item_batch(index_name, items=items) - assert isinstance(upsert_response, UpsertItemBatch.Success) - - sleep(2) - - search = getattr(vector_index_client, search_method_name) - search_response = search(index_name, query_vector=[1.0, 2.0], top_k=1, metadata_fields=ALL_METADATA) - assert isinstance(search_response, response) - assert len(search_response.hits) == 1 - - hits = [SearchHit(id="test_item_1", metadata=metadata, score=5.0)] - hits = when_fetching_vectors_apply_vectors_to_hits(search_response, hits, items) - assert search_response.hits == hits - - -def test_upsert_replaces_existing_items( - vector_index_client: PreviewVectorIndexClient, - unique_vector_index_name: TUniqueVectorIndexName, -) -> None: - index_name = unique_vector_index_name(vector_index_client) - create_response = vector_index_client.create_index( - index_name, num_dimensions=2, similarity_metric=SimilarityMetric.INNER_PRODUCT - ) - assert isinstance(create_response, CreateIndex.Success) - - upsert_response = vector_index_client.upsert_item_batch( - index_name, - items=[ - Item(id="test_item_1", vector=[1.0, 2.0], metadata={"key1": "value1"}), - Item(id="test_item_2", vector=[3.0, 4.0], metadata={"key2": "value2"}), - Item(id="test_item_3", vector=[5.0, 6.0], metadata={"key1": "value3", "key3": "value3"}), - ], - ) - assert isinstance(upsert_response, UpsertItemBatch.Success) - - sleep(2) - - upsert_response = vector_index_client.upsert_item_batch( - index_name, - items=[ - Item(id="test_item_1", vector=[2.0, 4.0], metadata={"key4": "value4"}), - ], - ) - assert isinstance(upsert_response, UpsertItemBatch.Success) - - sleep(2) - - search_response = vector_index_client.search( - index_name, query_vector=[1.0, 2.0], top_k=5, metadata_fields=["key1", "key2", "key3", "key4"] - ) - assert isinstance(search_response, Search.Success) - assert len(search_response.hits) == 3 - - assert search_response.hits == [ - SearchHit(id="test_item_3", score=17.0, metadata={"key1": "value3", "key3": "value3"}), - SearchHit(id="test_item_2", score=11.0, metadata={"key2": "value2"}), - SearchHit(id="test_item_1", score=10.0, metadata={"key4": "value4"}), - ] - - -def test_create_index_upsert_item_dimensions_different_than_num_dimensions_error( - vector_index_client: PreviewVectorIndexClient, - unique_vector_index_name: TUniqueVectorIndexName, -) -> None: - index_name = unique_vector_index_name(vector_index_client) - create_response = vector_index_client.create_index( - index_name, num_dimensions=2, similarity_metric=SimilarityMetric.INNER_PRODUCT - ) - assert isinstance(create_response, CreateIndex.Success) - - # upserting 3 dimensions - upsert_response = vector_index_client.upsert_item_batch( - index_name, items=[Item(id="test_item", vector=[1.0, 2.0, 3.0])] - ) - assert isinstance(upsert_response, UpsertItemBatch.Error) - - expected_inner_ex_message = "invalid parameter: vector, vector dimension has to match the index's dimension" - expected_message = f"Invalid argument passed to Momento client: {expected_inner_ex_message}" - assert upsert_response.message == expected_message - assert upsert_response.inner_exception.message == expected_inner_ex_message - - -@pytest.mark.parametrize( - ["search_method_name", "error_type"], - [("search", Search.Error), ("search_and_fetch_vectors", SearchAndFetchVectors.Error)], -) -def test_create_index_upsert_multiple_items_search_with_top_k_query_vector_dimensions_incorrect( - vector_index_client: PreviewVectorIndexClient, - unique_vector_index_name: TUniqueVectorIndexName, - search_method_name: str, - error_type: type[Search.Error] | type[SearchAndFetchVectors.Error], -) -> None: - index_name = unique_vector_index_name(vector_index_client) - create_response = vector_index_client.create_index( - index_name, num_dimensions=2, similarity_metric=SimilarityMetric.INNER_PRODUCT - ) - assert isinstance(create_response, CreateIndex.Success) - - upsert_response = vector_index_client.upsert_item_batch( - index_name, - items=[ - Item(id="test_item_1", vector=[1.0, 2.0]), - Item(id="test_item_2", vector=[3.0, 4.0]), - Item(id="test_item_3", vector=[5.0, 6.0]), - ], - ) - assert isinstance(upsert_response, UpsertItemBatch.Success) - - search = getattr(vector_index_client, search_method_name) - search_response = search(index_name, query_vector=[1.0, 2.0, 3.0], top_k=2) - assert isinstance(search_response, error_type) - - expected_inner_ex_message = "invalid parameter: query_vector, query vector dimension must match the index dimension" - expected_resp_message = f"Invalid argument passed to Momento client: {expected_inner_ex_message}" - - assert search_response.inner_exception.message == expected_inner_ex_message - assert search_response.message == expected_resp_message - - -def test_upsert_validates_index_name(vector_index_client: PreviewVectorIndexClient) -> None: - response = vector_index_client.upsert_item_batch(index_name="", items=[Item(id="test_item", vector=[1.0, 2.0])]) - assert isinstance(response, UpsertItemBatch.Error) - assert response.error_code == MomentoErrorCode.INVALID_ARGUMENT_ERROR - - -@pytest.mark.parametrize( - ["search_method_name", "error_type"], - [("search", Search.Error), ("search_and_fetch_vectors", SearchAndFetchVectors.Error)], -) -def test_search_validates_index_name( - vector_index_client: PreviewVectorIndexClient, - search_method_name: str, - error_type: type[Search.Error] | type[SearchAndFetchVectors.Error], -) -> None: - search = getattr(vector_index_client, search_method_name) - response = search(index_name="", query_vector=[1.0, 2.0]) - assert isinstance(response, error_type) - assert response.error_code == MomentoErrorCode.INVALID_ARGUMENT_ERROR - - -@pytest.mark.parametrize( - ["search_method_name", "error_type"], - [("search", Search.Error), ("search_and_fetch_vectors", SearchAndFetchVectors.Error)], -) -def test_search_validates_top_k( - vector_index_client: PreviewVectorIndexClient, - search_method_name: str, - error_type: type[Search.Error] | type[SearchAndFetchVectors.Error], -) -> None: - search = getattr(vector_index_client, search_method_name) - response = search(index_name="test_index", query_vector=[1.0, 2.0], top_k=0) - assert isinstance(response, error_type) - assert response.error_code == MomentoErrorCode.INVALID_ARGUMENT_ERROR - assert response.inner_exception.message == "Top k must be a positive integer." - - -@pytest.mark.parametrize( - ["similarity_metric", "distances", "thresholds", "search_method_name", "response"], - [ - # Distances are the distance to the same 3 data vectors from the same query vector. - # Thresholds are those that should: - # 1. exclude lowest two matches - # 2. keep all matches - # 3. exclude all matches - (SimilarityMetric.COSINE_SIMILARITY, [1.0, 0.0, -1.0], [0.5, -1.01, 1.0], "search", Search.Success), - ( - SimilarityMetric.COSINE_SIMILARITY, - [1.0, 0.0, -1.0], - [0.5, -1.01, 1.0], - "search_and_fetch_vectors", - SearchAndFetchVectors.Success, - ), - (SimilarityMetric.INNER_PRODUCT, [4.0, 0.0, -4.0], [0.0, -4.01, 4.0], "search", Search.Success), - ( - SimilarityMetric.INNER_PRODUCT, - [4.0, 0.0, -4.0], - [0.0, -4.01, 4.0], - "search_and_fetch_vectors", - SearchAndFetchVectors.Success, - ), - (SimilarityMetric.EUCLIDEAN_SIMILARITY, [2, 10, 18], [3, 20, -0.01], "search", Search.Success), - ( - SimilarityMetric.EUCLIDEAN_SIMILARITY, - [2, 10, 18], - [3, 20, -0.01], - "search_and_fetch_vectors", - SearchAndFetchVectors.Success, - ), - ], -) -def test_search_score_threshold_happy_path( - vector_index_client: PreviewVectorIndexClient, - unique_vector_index_name: TUniqueVectorIndexName, - similarity_metric: SimilarityMetric, - distances: list[float], - thresholds: list[float], - search_method_name: str, - response: type[Search.Success] | type[SearchAndFetchVectors.Success], -) -> None: - index_name = unique_vector_index_name(vector_index_client) - num_dimensions = 2 - create_response = vector_index_client.create_index( - index_name, num_dimensions=num_dimensions, similarity_metric=similarity_metric - ) - assert isinstance(create_response, CreateIndex.Success) - - items = [ - Item(id="test_item_1", vector=[1.0, 1.0]), - Item(id="test_item_2", vector=[-1.0, 1.0]), - Item(id="test_item_3", vector=[-1.0, -1.0]), - ] - upsert_response = vector_index_client.upsert_item_batch( - index_name, - items=items, - ) - assert isinstance(upsert_response, UpsertItemBatch.Success) - - sleep(2) - - query_vector = [2.0, 2.0] - search_hits = [SearchHit(id=f"test_item_{i+1}", score=distance) for i, distance in enumerate(distances)] - - search = getattr(vector_index_client, search_method_name) - search_response = search(index_name, query_vector=query_vector, top_k=3, score_threshold=thresholds[0]) - assert isinstance(search_response, response) - assert search_response.hits == when_fetching_vectors_apply_vectors_to_hits(search_response, [search_hits[0]], items) - - search_response2 = search(index_name, query_vector=query_vector, top_k=3, score_threshold=thresholds[1]) - assert isinstance(search_response2, response) - assert search_response2.hits == when_fetching_vectors_apply_vectors_to_hits(search_response, search_hits, items) - - search_response3 = search(index_name, query_vector=query_vector, top_k=3, score_threshold=thresholds[2]) - assert isinstance(search_response3, response) - assert search_response3.hits == [] - - -def test_search_with_filter_expression( - vector_index_client: PreviewVectorIndexClient, - unique_vector_index_name: TUniqueVectorIndexName, -) -> None: - index_name = unique_vector_index_name(vector_index_client) - num_dimensions = 2 - create_response = vector_index_client.create_index( - index_name, num_dimensions=num_dimensions, similarity_metric=SimilarityMetric.INNER_PRODUCT - ) - assert isinstance(create_response, CreateIndex.Success) - items = [ - Item( - id="test_item_1", - vector=[1.0, 1.0], - metadata={"str": "value1", "int": 0, "float": 0.0, "bool": True, "tags": ["a", "b", "c"]}, - ), - Item( - id="test_item_2", - vector=[-1.0, 1.0], - metadata={"str": "value2", "int": 5, "float": 5.0, "bool": False, "tags": ["a", "b"]}, - ), - Item( - id="test_item_3", - vector=[-1.0, -1.0], - metadata={"str": "value3", "int": 10, "float": 10.0, "bool": True, "tags": ["a", "d"]}, - ), - ] - upsert_response = vector_index_client.upsert_item_batch( - index_name, - items=items, - ) - assert isinstance(upsert_response, UpsertItemBatch.Success) - sleep(2) - - # Writing the test cases here instead of as a parameterized test because: - # 1. The search data is the same across tests, so no need to reindex each time. - # 2. It is 10x faster to run the tests this way. - for filter, expected_ids, test_case_name in [ - (Field("str") == "value1", ["test_item_1"], "string equality"), - (Field("str") != "value1", ["test_item_2", "test_item_3"], "string inequality"), - (Field("int") == 0, ["test_item_1"], "int equality"), - (Field("float") == 0.0, ["test_item_1"], "float equality"), - (Field("bool"), ["test_item_1", "test_item_3"], "bool equality"), - (~Field("bool"), ["test_item_2"], "bool inequality"), - (Field("int") > 5, ["test_item_3"], "int greater than"), - (Field("int") >= 5, ["test_item_2", "test_item_3"], "int greater than or equal to"), - (Field("float") > 5.0, ["test_item_3"], "float greater than"), - (Field("float") >= 5.0, ["test_item_2", "test_item_3"], "float greater than or equal to"), - (Field("int") < 5, ["test_item_1"], "int less than"), - (Field("int") <= 5, ["test_item_1", "test_item_2"], "int less than or equal to"), - (Field("float") < 5.0, ["test_item_1"], "float less than"), - (Field("float") <= 5.0, ["test_item_1", "test_item_2"], "float less than or equal to"), - (Field("tags").list_contains("a"), ["test_item_1", "test_item_2", "test_item_3"], "list contains a"), - (Field("tags").list_contains("b"), ["test_item_1", "test_item_2"], "list contains b"), - (Field("tags").list_contains("m"), [], "list contains m"), - ((Field("tags").list_contains("b")) & (Field("int") > 1), ["test_item_2"], "list contains b and int > 1"), - ( - (Field("tags").list_contains("b")) | (Field("int") > 1), - ["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 = cast(FilterExpression, filter) - search_response = vector_index_client.search(index_name, query_vector=[2.0, 2.0], filter=filter) - assert isinstance( - search_response, Search.Success - ), f"Expected search {test_case_name!r} to succeed but got {search_response!r}" - assert ( - [hit.id for hit in search_response.hits] == expected_ids - ), f"Expected search {test_case_name!r} to return {expected_ids!r} but got {search_response.hits!r}" - - search_and_fetch_vectors_response = vector_index_client.search_and_fetch_vectors( - index_name, query_vector=[2.0, 2.0], filter=filter - ) - assert isinstance( - search_and_fetch_vectors_response, SearchAndFetchVectors.Success - ), f"Expected search {test_case_name!r} to succeed but got {search_and_fetch_vectors_response!r}" - assert [hit.id for hit in search_and_fetch_vectors_response.hits] == expected_ids, ( - f"Expected search {test_case_name!r} to return {expected_ids!r} " - f"but got {search_and_fetch_vectors_response.hits!r}" - ) - - -def test_delete_validates_index_name(vector_index_client: PreviewVectorIndexClient) -> None: - response = vector_index_client.delete_item_batch(index_name="", filter=[]) - assert isinstance(response, DeleteItemBatch.Error) - assert response.error_code == MomentoErrorCode.INVALID_ARGUMENT_ERROR - - -def test_delete_items_by_id( - vector_index_client: PreviewVectorIndexClient, - unique_vector_index_name: TUniqueVectorIndexName, -) -> None: - index_name = unique_vector_index_name(vector_index_client) - create_response = vector_index_client.create_index( - index_name, num_dimensions=2, similarity_metric=SimilarityMetric.INNER_PRODUCT - ) - assert isinstance(create_response, CreateIndex.Success) - - upsert_response = vector_index_client.upsert_item_batch( - index_name, - items=[ - Item(id="test_item_1", vector=[1.0, 2.0], metadata={"key": "value1"}), - Item(id="test_item_2", vector=[3.0, 4.0], metadata={"key": "value2"}), - Item(id="test_item_3", vector=[5.0, 6.0], metadata={"key": "value3"}), - Item(id="test_item_4", vector=[7.0, 8.0], metadata={"key": "value4"}), - ], - ) - assert isinstance(upsert_response, UpsertItemBatch.Success) - - sleep(2) - - search_response = vector_index_client.search(index_name, query_vector=[1.0, 2.0], top_k=10) - assert isinstance(search_response, Search.Success) - assert len(search_response.hits) == 4 - - assert search_response.hits == [ - SearchHit(id="test_item_4", score=23.0), - SearchHit(id="test_item_3", score=17.0), - SearchHit(id="test_item_2", score=11.0), - SearchHit(id="test_item_1", score=5.0), - ] - - delete_response = vector_index_client.delete_item_batch(index_name, filter=["test_item_1", "test_item_3"]) - assert isinstance(delete_response, DeleteItemBatch.Success) - - sleep(2) - - search_response = vector_index_client.search(index_name, query_vector=[1.0, 2.0], top_k=10) - assert isinstance(search_response, Search.Success) - assert len(search_response.hits) == 2 - - assert search_response.hits == [ - SearchHit(id="test_item_4", score=23.0), - SearchHit(id="test_item_2", score=11.0), - ] - - -def test_delete_items_by_filter( - vector_index_client: PreviewVectorIndexClient, - unique_vector_index_name: TUniqueVectorIndexName, -) -> None: - index_name = unique_vector_index_name(vector_index_client) - create_response = vector_index_client.create_index(index_name, num_dimensions=2) - assert isinstance(create_response, CreateIndex.Success) - - items = [ - Item(id="test_item_1", vector=[1.0, 1.0], metadata={"key": "value1"}), - Item(id="test_item_2", vector=[3.0, 4.0], metadata={"key": "value2"}), - Item(id="test_item_3", vector=[5.0, 6.0], metadata={"key": "value1"}), - Item(id="test_item_4", vector=[7.0, 8.0], metadata={"key": "value2"}), - ] - upsert_response = vector_index_client.upsert_item_batch( - index_name, - items=items, - ) - assert isinstance(upsert_response, UpsertItemBatch.Success) - - sleep(2) - - search_response = vector_index_client.search(index_name, query_vector=[1.0, 1.0], top_k=10) - assert isinstance(search_response, Search.Success) - assert len(search_response.hits) == 4 - - delete_response = vector_index_client.delete_item_batch(index_name, filter=filters.Equals("key", "value1")) - assert isinstance(delete_response, DeleteItemBatch.Success) - - sleep(2) - - search_response = vector_index_client.search(index_name, query_vector=[1.0, 1.0], top_k=10) - assert isinstance(search_response, Search.Success) - assert len(search_response.hits) == 2 - - -@pytest.mark.parametrize( - [ - "get_item_method_name", - "ids", - "expected_get_item_response", - "expected_get_item_values", - ], - [ - ("get_item_batch", [], GetItemBatch.Success, {}), - ("get_item_metadata_batch", [], GetItemMetadataBatch.Success, {}), - ("get_item_batch", ["missing_id"], GetItemBatch.Success, {}), - ( - "get_item_metadata_batch", - ["test_item_1"], - GetItemMetadataBatch.Success, - {"test_item_1": {"key1": "value1"}}, - ), - ("get_item_metadata_batch", ["missing_id"], GetItemMetadataBatch.Success, {}), - ( - "get_item_batch", - ["test_item_1"], - GetItemBatch.Success, - { - "test_item_1": Item(id="test_item_1", vector=[1.0, 1.0], metadata={"key1": "value1"}), - }, - ), - ( - "get_item_batch", - ["test_item_1", "missing_id", "test_item_2"], - GetItemBatch.Success, - { - "test_item_1": Item(id="test_item_1", vector=[1.0, 1.0], metadata={"key1": "value1"}), - "test_item_2": Item(id="test_item_2", vector=[-1.0, 1.0], metadata={}), - }, - ), - ( - "get_item_metadata_batch", - ["test_item_1", "missing_id", "test_item_2"], - GetItemMetadataBatch.Success, - { - "test_item_1": {"key1": "value1"}, - "test_item_2": {}, - }, - ), - ], -) -def test_get_items_by_id( - vector_index_client: PreviewVectorIndexClient, - unique_vector_index_name: TUniqueVectorIndexName, - get_item_method_name: str, - ids: list[str], - expected_get_item_response: type[GetItemMetadataBatch.Success] | type[GetItemBatch.Success], - expected_get_item_values: dict[str, Metadata] | dict[str, Item], -) -> None: - index_name = unique_vector_index_name(vector_index_client) - create_response = vector_index_client.create_index(index_name, num_dimensions=2) - assert isinstance(create_response, CreateIndex.Success) - - items = [ - Item(id="test_item_1", vector=[1.0, 1.0], metadata={"key1": "value1"}), - Item(id="test_item_2", vector=[-1.0, 1.0]), - Item(id="test_item_3", vector=[-1.0, -1.0]), - ] - upsert_response = vector_index_client.upsert_item_batch( - index_name, - items=items, - ) - assert isinstance(upsert_response, UpsertItemBatch.Success) - - sleep(2) - - get_item = getattr(vector_index_client, get_item_method_name) - get_item_response = get_item(index_name, ids) - assert isinstance(get_item_response, expected_get_item_response) - assert get_item_response.values == expected_get_item_values - - -def test_count_items_on_missing_index( - vector_index_client: PreviewVectorIndexClient, -) -> None: - response = vector_index_client.count_items(index_name=uuid_str()) - assert isinstance(response, CountItems.Error) - assert response.error_code == MomentoErrorCode.NOT_FOUND_ERROR - - -def test_count_items_on_empty_index( - vector_index_client: PreviewVectorIndexClient, - unique_vector_index_name: TUniqueVectorIndexName, -) -> None: - index_name = unique_vector_index_name(vector_index_client) - - create_response = vector_index_client.create_index(index_name, num_dimensions=2) - assert isinstance(create_response, CreateIndex.Success) - - count_response = vector_index_client.count_items(index_name) - assert isinstance(count_response, CountItems.Success) - assert count_response.item_count == 0 - - -def test_count_items_with_items( - vector_index_client: PreviewVectorIndexClient, - unique_vector_index_name: TUniqueVectorIndexName, -) -> None: - num_items = 10 - index_name = unique_vector_index_name(vector_index_client) - - create_response = vector_index_client.create_index(index_name, num_dimensions=2) - assert isinstance(create_response, CreateIndex.Success) - - items = [Item(id=f"test_item_{i}", vector=[i, i]) for i in range(num_items)] # type: list[Item] - upsert_response = vector_index_client.upsert_item_batch( - index_name, - items=items, - ) - assert isinstance(upsert_response, UpsertItemBatch.Success) - - sleep(2) - - count_response = vector_index_client.count_items(index_name) - assert isinstance(count_response, CountItems.Success) - assert count_response.item_count == num_items - - num_items_to_delete = 5 - delete_response = vector_index_client.delete_item_batch( - index_name, filter=[item.id for item in items[:num_items_to_delete]] - ) - assert isinstance(delete_response, DeleteItemBatch.Success) - - sleep(2) - - count_response = vector_index_client.count_items(index_name) - assert isinstance(count_response, CountItems.Success) - assert count_response.item_count == num_items - num_items_to_delete diff --git a/tests/momento/vector_index_client/test_data_async.py b/tests/momento/vector_index_client/test_data_async.py deleted file mode 100644 index fc82f2b6..00000000 --- a/tests/momento/vector_index_client/test_data_async.py +++ /dev/null @@ -1,916 +0,0 @@ -from __future__ import annotations - -from typing import Optional, cast - -import pytest -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, filters -from momento.responses.vector_index import ( - CountItems, - CreateIndex, - DeleteItemBatch, - GetItemBatch, - GetItemMetadataBatch, - Search, - SearchAndFetchVectors, - SearchHit, - UpsertItemBatch, -) - -from tests.conftest import TUniqueVectorIndexNameAsync -from tests.utils import sleep_async, uuid_str, when_fetching_vectors_apply_vectors_to_hits - - -async def test_create_index_with_inner_product_upsert_item_search_happy_path( - vector_index_client_async: PreviewVectorIndexClientAsync, - unique_vector_index_name_async: TUniqueVectorIndexNameAsync, -) -> None: - index_name = unique_vector_index_name_async(vector_index_client_async) - create_response = await vector_index_client_async.create_index( - index_name, num_dimensions=2, similarity_metric=SimilarityMetric.INNER_PRODUCT - ) - assert isinstance(create_response, CreateIndex.Success) - - upsert_response = await vector_index_client_async.upsert_item_batch( - index_name, items=[Item(id="test_item", vector=[1.0, 2.0])] - ) - assert isinstance(upsert_response, UpsertItemBatch.Success) - - await sleep_async(2) - - search_response = await vector_index_client_async.search(index_name, query_vector=[1.0, 2.0], top_k=1) - assert isinstance(search_response, Search.Success) - assert len(search_response.hits) == 1 - assert search_response.hits[0].id == "test_item" - assert search_response.hits[0].score == 5.0 - - -@pytest.mark.parametrize("similarity_metric", [SimilarityMetric.COSINE_SIMILARITY, None]) -async def test_create_index_with_cosine_similarity_upsert_item_search_happy_path( - vector_index_client_async: PreviewVectorIndexClientAsync, - unique_vector_index_name_async: TUniqueVectorIndexNameAsync, - similarity_metric: Optional[SimilarityMetric], -) -> None: - index_name = unique_vector_index_name_async(vector_index_client_async) - num_dimensions = 2 - if similarity_metric is not None: - create_response = await vector_index_client_async.create_index( - index_name, num_dimensions=num_dimensions, similarity_metric=similarity_metric - ) - else: - create_response = await vector_index_client_async.create_index(index_name, num_dimensions=num_dimensions) - assert isinstance(create_response, CreateIndex.Success) - - upsert_response = await vector_index_client_async.upsert_item_batch( - index_name, - items=[ - Item(id="test_item_1", vector=[1.0, 1.0]), - Item(id="test_item_2", vector=[-1.0, 1.0]), - Item(id="test_item_3", vector=[-1.0, -1.0]), - ], - ) - assert isinstance(upsert_response, UpsertItemBatch.Success) - - await sleep_async(2) - - search_response = await vector_index_client_async.search(index_name, query_vector=[2.0, 2.0], top_k=3) - assert isinstance(search_response, Search.Success) - assert search_response.hits == [ - SearchHit(id="test_item_1", score=1.0), - SearchHit(id="test_item_2", score=0.0), - SearchHit(id="test_item_3", score=-1.0), - ] - - -async def test_create_index_with_euclidean_similarity_upsert_item_search_happy_path( - vector_index_client_async: PreviewVectorIndexClientAsync, - unique_vector_index_name_async: TUniqueVectorIndexNameAsync, -) -> None: - index_name = unique_vector_index_name_async(vector_index_client_async) - create_response = await vector_index_client_async.create_index( - index_name, num_dimensions=2, similarity_metric=SimilarityMetric.EUCLIDEAN_SIMILARITY - ) - assert isinstance(create_response, CreateIndex.Success) - - upsert_response = await vector_index_client_async.upsert_item_batch( - index_name, - items=[ - Item(id="test_item_1", vector=[1.0, 1.0]), - Item(id="test_item_2", vector=[-1.0, 1.0]), - Item(id="test_item_3", vector=[-1.0, -1.0]), - ], - ) - assert isinstance(upsert_response, UpsertItemBatch.Success) - - await sleep_async(2) - - search_response = await vector_index_client_async.search(index_name, query_vector=[1.0, 1.0], top_k=3) - assert isinstance(search_response, Search.Success) - assert search_response.hits == [ - SearchHit(id="test_item_1", score=0.0), - SearchHit(id="test_item_2", score=4.0), - SearchHit(id="test_item_3", score=8.0), - ] - - -async def test_create_index_upsert_multiple_items_search_happy_path( - vector_index_client_async: PreviewVectorIndexClientAsync, - unique_vector_index_name_async: TUniqueVectorIndexNameAsync, -) -> None: - index_name = unique_vector_index_name_async(vector_index_client_async) - create_response = await vector_index_client_async.create_index( - index_name, num_dimensions=2, similarity_metric=SimilarityMetric.INNER_PRODUCT - ) - assert isinstance(create_response, CreateIndex.Success) - - upsert_response = await vector_index_client_async.upsert_item_batch( - index_name, - items=[ - Item(id="test_item_1", vector=[1.0, 2.0]), - Item(id="test_item_2", vector=[3.0, 4.0]), - Item(id="test_item_3", vector=[5.0, 6.0]), - ], - ) - assert isinstance(upsert_response, UpsertItemBatch.Success) - - await sleep_async(2) - - search_response = await vector_index_client_async.search(index_name, query_vector=[1.0, 2.0], top_k=3) - assert isinstance(search_response, Search.Success) - assert len(search_response.hits) == 3 - - assert search_response.hits == [ - SearchHit(id="test_item_3", score=17.0), - SearchHit(id="test_item_2", score=11.0), - SearchHit(id="test_item_1", score=5.0), - ] - - -# A note on the parameterized search tests: -# The search tests are parameterized to test both the search and search_and_fetch_vectors methods. -# We pass both the name of the specific search method and the response type. -# The tests will run for both search methods, and the response type will be used to assert the type of the response. -# The vectors are attached to the hits if the response type is SearchAndFetchVectors.Success. -@pytest.mark.parametrize( - ["search_method_name", "response"], - [("search", Search.Success), ("search_and_fetch_vectors", SearchAndFetchVectors.Success)], -) -async def test_create_index_upsert_multiple_items_search_with_top_k_happy_path( - vector_index_client_async: PreviewVectorIndexClientAsync, - unique_vector_index_name_async: TUniqueVectorIndexNameAsync, - search_method_name: str, - response: type[Search.Success] | type[SearchAndFetchVectors.Success], -) -> None: - index_name = unique_vector_index_name_async(vector_index_client_async) - create_response = await vector_index_client_async.create_index( - index_name, num_dimensions=2, similarity_metric=SimilarityMetric.INNER_PRODUCT - ) - assert isinstance(create_response, CreateIndex.Success) - - items = [ - Item(id="test_item_1", vector=[1.0, 2.0]), - Item(id="test_item_2", vector=[3.0, 4.0]), - Item(id="test_item_3", vector=[5.0, 6.0]), - ] - upsert_response = await vector_index_client_async.upsert_item_batch( - index_name, - items=items, - ) - assert isinstance(upsert_response, UpsertItemBatch.Success) - - await sleep_async(2) - - search = getattr(vector_index_client_async, search_method_name) - search_response = await search(index_name, query_vector=[1.0, 2.0], top_k=2) - assert isinstance(search_response, response) - assert len(search_response.hits) == 2 - - hits = [SearchHit(id="test_item_3", score=17.0), SearchHit(id="test_item_2", score=11.0)] - hits = when_fetching_vectors_apply_vectors_to_hits(search_response, hits, items) - assert search_response.hits == hits - - -@pytest.mark.parametrize( - ["search_method_name", "response"], - [("search", Search.Success), ("search_and_fetch_vectors", SearchAndFetchVectors.Success)], -) -async def test_upsert_and_search_with_metadata_happy_path( - vector_index_client_async: PreviewVectorIndexClientAsync, - unique_vector_index_name_async: TUniqueVectorIndexNameAsync, - search_method_name: str, - response: type[Search.Success] | type[SearchAndFetchVectors.Success], -) -> None: - index_name = unique_vector_index_name_async(vector_index_client_async) - create_response = await vector_index_client_async.create_index( - index_name, num_dimensions=2, similarity_metric=SimilarityMetric.INNER_PRODUCT - ) - assert isinstance(create_response, CreateIndex.Success) - - items = [ - Item(id="test_item_1", vector=[1.0, 2.0], metadata={"key1": "value1"}), - Item(id="test_item_2", vector=[3.0, 4.0], metadata={"key2": "value2"}), - Item(id="test_item_3", vector=[5.0, 6.0], metadata={"key1": "value3", "key3": "value3"}), - ] - upsert_response = await vector_index_client_async.upsert_item_batch( - index_name, - items=items, - ) - assert isinstance(upsert_response, UpsertItemBatch.Success) - - await sleep_async(2) - - search = getattr(vector_index_client_async, search_method_name) - search_response = await search(index_name, query_vector=[1.0, 2.0], top_k=3) - assert isinstance(search_response, response) - assert len(search_response.hits) == 3 - - hits = [ - SearchHit(id="test_item_3", score=17.0), - SearchHit(id="test_item_2", score=11.0), - SearchHit(id="test_item_1", score=5.0), - ] - hits = when_fetching_vectors_apply_vectors_to_hits(search_response, hits, items) - assert search_response.hits == hits - - search_response = await search(index_name, query_vector=[1.0, 2.0], top_k=3, metadata_fields=["key1"]) - assert isinstance(search_response, response) - assert len(search_response.hits) == 3 - - hits = [ - SearchHit(id="test_item_3", score=17.0, metadata={"key1": "value3"}), - SearchHit(id="test_item_2", score=11.0, metadata={}), - SearchHit(id="test_item_1", score=5.0, metadata={"key1": "value1"}), - ] - hits = when_fetching_vectors_apply_vectors_to_hits(search_response, hits, items) - assert search_response.hits == hits - - search_response = await search( - index_name, query_vector=[1.0, 2.0], top_k=3, metadata_fields=["key1", "key2", "key3", "key4"] - ) - assert isinstance(search_response, response) - assert len(search_response.hits) == 3 - - hits = [ - SearchHit(id="test_item_3", score=17.0, metadata={"key1": "value3", "key3": "value3"}), - SearchHit(id="test_item_2", score=11.0, metadata={"key2": "value2"}), - SearchHit(id="test_item_1", score=5.0, metadata={"key1": "value1"}), - ] - hits = when_fetching_vectors_apply_vectors_to_hits(search_response, hits, items) - assert search_response.hits == hits - - -async def test_upsert_with_bad_metadata(vector_index_client_async: PreviewVectorIndexClientAsync) -> None: - response = await vector_index_client_async.upsert_item_batch( - index_name="test_index", - items=[Item(id="test_item", vector=[1.0, 2.0], metadata={"key": {"subkey": "subvalue"}})], # type: ignore - ) - assert isinstance(response, UpsertItemBatch.Error) - assert response.error_code == MomentoErrorCode.INVALID_ARGUMENT_ERROR - assert ( - response.message - == "Invalid argument passed to Momento client: Metadata values must be either str, int, float, bool, or list[str]. Field 'key' has a value of type with value {'subkey': 'subvalue'}." # noqa: E501 W503 - ) - - -@pytest.mark.parametrize( - ["search_method_name", "response"], - [("search", Search.Success), ("search_and_fetch_vectors", SearchAndFetchVectors.Success)], -) -async def test_upsert_and_search_with_all_metadata_happy_path( - vector_index_client_async: PreviewVectorIndexClientAsync, - unique_vector_index_name_async: TUniqueVectorIndexNameAsync, - search_method_name: str, - response: type[Search.Success] | type[SearchAndFetchVectors.Success], -) -> None: - index_name = unique_vector_index_name_async(vector_index_client_async) - create_response = await vector_index_client_async.create_index( - index_name, num_dimensions=2, similarity_metric=SimilarityMetric.INNER_PRODUCT - ) - assert isinstance(create_response, CreateIndex.Success) - - items = [ - Item(id="test_item_1", vector=[1.0, 2.0], metadata={"key1": "value1"}), - Item(id="test_item_2", vector=[3.0, 4.0], metadata={"key2": "value2"}), - Item(id="test_item_3", vector=[5.0, 6.0], metadata={"key1": "value3", "key3": "value3"}), - ] - upsert_response = await vector_index_client_async.upsert_item_batch( - index_name, - items=items, - ) - assert isinstance(upsert_response, UpsertItemBatch.Success) - - await sleep_async(2) - - search = getattr(vector_index_client_async, search_method_name) - search_response = await search(index_name, query_vector=[1.0, 2.0], top_k=3) - assert isinstance(search_response, response) - assert len(search_response.hits) == 3 - - hits = [ - SearchHit(id="test_item_3", score=17.0), - SearchHit(id="test_item_2", score=11.0), - SearchHit(id="test_item_1", score=5.0), - ] - hits = when_fetching_vectors_apply_vectors_to_hits(search_response, hits, items) - assert search_response.hits == hits - - search_response = await search(index_name, query_vector=[1.0, 2.0], top_k=3, metadata_fields=ALL_METADATA) - assert isinstance(search_response, response) - assert len(search_response.hits) == 3 - - hits = [ - SearchHit(id="test_item_3", score=17.0, metadata={"key1": "value3", "key3": "value3"}), - SearchHit(id="test_item_2", score=11.0, metadata={"key2": "value2"}), - SearchHit(id="test_item_1", score=5.0, metadata={"key1": "value1"}), - ] - hits = when_fetching_vectors_apply_vectors_to_hits(search_response, hits, items) - assert search_response.hits == hits - - -@pytest.mark.parametrize( - ["search_method_name", "response"], - [ - ("search", Search.Success), - ( - "search_and_fetch_vectors", - SearchAndFetchVectors.Success, - ), - ], -) -async def test_upsert_and_search_with_diverse_metadata_happy_path( - vector_index_client_async: PreviewVectorIndexClientAsync, - unique_vector_index_name_async: TUniqueVectorIndexNameAsync, - search_method_name: str, - response: type[Search.Success] | type[SearchAndFetchVectors.Success], -) -> None: - index_name = unique_vector_index_name_async(vector_index_client_async) - create_response = await vector_index_client_async.create_index( - index_name, num_dimensions=2, similarity_metric=SimilarityMetric.INNER_PRODUCT - ) - assert isinstance(create_response, CreateIndex.Success) - - metadata: Metadata = { - "string": "value", - "bool": True, - "int": 1, - "float": 3.14, - "list": ["a", "b", "c"], - "empty_list": [], - } - items = [ - Item(id="test_item_1", vector=[1.0, 2.0], metadata=metadata), - ] - upsert_response = await vector_index_client_async.upsert_item_batch(index_name, items=items) - assert isinstance(upsert_response, UpsertItemBatch.Success) - - await sleep_async(2) - - search = getattr(vector_index_client_async, search_method_name) - search_response = await search(index_name, query_vector=[1.0, 2.0], top_k=1, metadata_fields=ALL_METADATA) - assert isinstance(search_response, response) - assert len(search_response.hits) == 1 - - hits = [SearchHit(id="test_item_1", metadata=metadata, score=5.0)] - hits = when_fetching_vectors_apply_vectors_to_hits(search_response, hits, items) - assert search_response.hits == hits - - -async def test_upsert_replaces_existing_items( - vector_index_client_async: PreviewVectorIndexClientAsync, - unique_vector_index_name_async: TUniqueVectorIndexNameAsync, -) -> None: - index_name = unique_vector_index_name_async(vector_index_client_async) - create_response = await vector_index_client_async.create_index( - index_name, num_dimensions=2, similarity_metric=SimilarityMetric.INNER_PRODUCT - ) - assert isinstance(create_response, CreateIndex.Success) - - upsert_response = await vector_index_client_async.upsert_item_batch( - index_name, - items=[ - Item(id="test_item_1", vector=[1.0, 2.0], metadata={"key1": "value1"}), - Item(id="test_item_2", vector=[3.0, 4.0], metadata={"key2": "value2"}), - Item(id="test_item_3", vector=[5.0, 6.0], metadata={"key1": "value3", "key3": "value3"}), - ], - ) - assert isinstance(upsert_response, UpsertItemBatch.Success) - - await sleep_async(2) - - upsert_response = await vector_index_client_async.upsert_item_batch( - index_name, - items=[ - Item(id="test_item_1", vector=[2.0, 4.0], metadata={"key4": "value4"}), - ], - ) - assert isinstance(upsert_response, UpsertItemBatch.Success) - - await sleep_async(2) - - search_response = await vector_index_client_async.search( - index_name, query_vector=[1.0, 2.0], top_k=5, metadata_fields=["key1", "key2", "key3", "key4"] - ) - assert isinstance(search_response, Search.Success) - assert len(search_response.hits) == 3 - - assert search_response.hits == [ - SearchHit(id="test_item_3", score=17.0, metadata={"key1": "value3", "key3": "value3"}), - SearchHit(id="test_item_2", score=11.0, metadata={"key2": "value2"}), - SearchHit(id="test_item_1", score=10.0, metadata={"key4": "value4"}), - ] - - -async def test_create_index_upsert_item_dimensions_different_than_num_dimensions_error( - vector_index_client_async: PreviewVectorIndexClientAsync, - unique_vector_index_name_async: TUniqueVectorIndexNameAsync, -) -> None: - index_name = unique_vector_index_name_async(vector_index_client_async) - create_response = await vector_index_client_async.create_index( - index_name, num_dimensions=2, similarity_metric=SimilarityMetric.INNER_PRODUCT - ) - assert isinstance(create_response, CreateIndex.Success) - - # upserting 3 dimensions - upsert_response = await vector_index_client_async.upsert_item_batch( - index_name, items=[Item(id="test_item", vector=[1.0, 2.0, 3.0])] - ) - assert isinstance(upsert_response, UpsertItemBatch.Error) - - expected_inner_ex_message = "invalid parameter: vector, vector dimension has to match the index's dimension" - expected_message = f"Invalid argument passed to Momento client: {expected_inner_ex_message}" - assert upsert_response.message == expected_message - assert upsert_response.inner_exception.message == expected_inner_ex_message - - -@pytest.mark.parametrize( - ["search_method_name", "error_type"], - [("search", Search.Error), ("search_and_fetch_vectors", SearchAndFetchVectors.Error)], -) -async def test_create_index_upsert_multiple_items_search_with_top_k_query_vector_dimensions_incorrect( - vector_index_client_async: PreviewVectorIndexClientAsync, - unique_vector_index_name_async: TUniqueVectorIndexNameAsync, - search_method_name: str, - error_type: type[Search.Error] | type[SearchAndFetchVectors.Error], -) -> None: - index_name = unique_vector_index_name_async(vector_index_client_async) - create_response = await vector_index_client_async.create_index( - index_name, num_dimensions=2, similarity_metric=SimilarityMetric.INNER_PRODUCT - ) - assert isinstance(create_response, CreateIndex.Success) - - upsert_response = await vector_index_client_async.upsert_item_batch( - index_name, - items=[ - Item(id="test_item_1", vector=[1.0, 2.0]), - Item(id="test_item_2", vector=[3.0, 4.0]), - Item(id="test_item_3", vector=[5.0, 6.0]), - ], - ) - assert isinstance(upsert_response, UpsertItemBatch.Success) - - search = getattr(vector_index_client_async, search_method_name) - search_response = await search(index_name, query_vector=[1.0, 2.0, 3.0], top_k=2) - assert isinstance(search_response, error_type) - - expected_inner_ex_message = "invalid parameter: query_vector, query vector dimension must match the index dimension" - expected_resp_message = f"Invalid argument passed to Momento client: {expected_inner_ex_message}" - - assert search_response.inner_exception.message == expected_inner_ex_message - assert search_response.message == expected_resp_message - - -async def test_upsert_validates_index_name(vector_index_client_async: PreviewVectorIndexClientAsync) -> None: - response = await vector_index_client_async.upsert_item_batch( - index_name="", items=[Item(id="test_item", vector=[1.0, 2.0])] - ) - assert isinstance(response, UpsertItemBatch.Error) - assert response.error_code == MomentoErrorCode.INVALID_ARGUMENT_ERROR - - -@pytest.mark.parametrize( - ["search_method_name", "error_type"], - [("search", Search.Error), ("search_and_fetch_vectors", SearchAndFetchVectors.Error)], -) -async def test_search_validates_index_name( - vector_index_client_async: PreviewVectorIndexClientAsync, - search_method_name: str, - error_type: type[Search.Error] | type[SearchAndFetchVectors.Error], -) -> None: - search = getattr(vector_index_client_async, search_method_name) - response = await search(index_name="", query_vector=[1.0, 2.0]) - assert isinstance(response, error_type) - assert response.error_code == MomentoErrorCode.INVALID_ARGUMENT_ERROR - - -@pytest.mark.parametrize( - ["search_method_name", "error_type"], - [("search", Search.Error), ("search_and_fetch_vectors", SearchAndFetchVectors.Error)], -) -async def test_search_validates_top_k( - vector_index_client_async: PreviewVectorIndexClientAsync, - search_method_name: str, - error_type: type[Search.Error] | type[SearchAndFetchVectors.Error], -) -> None: - search = getattr(vector_index_client_async, search_method_name) - response = await search(index_name="test_index", query_vector=[1.0, 2.0], top_k=0) - assert isinstance(response, error_type) - assert response.error_code == MomentoErrorCode.INVALID_ARGUMENT_ERROR - assert response.inner_exception.message == "Top k must be a positive integer." - - -@pytest.mark.parametrize( - ["similarity_metric", "distances", "thresholds", "search_method_name", "response"], - [ - # Distances are the distance to the same 3 data vectors from the same query vector. - # Thresholds are those that should: - # 1. exclude lowest two matches - # 2. keep all matches - # 3. exclude all matches - (SimilarityMetric.COSINE_SIMILARITY, [1.0, 0.0, -1.0], [0.5, -1.01, 1.0], "search", Search.Success), - ( - SimilarityMetric.COSINE_SIMILARITY, - [1.0, 0.0, -1.0], - [0.5, -1.01, 1.0], - "search_and_fetch_vectors", - SearchAndFetchVectors.Success, - ), - (SimilarityMetric.INNER_PRODUCT, [4.0, 0.0, -4.0], [0.0, -4.01, 4.0], "search", Search.Success), - ( - SimilarityMetric.INNER_PRODUCT, - [4.0, 0.0, -4.0], - [0.0, -4.01, 4.0], - "search_and_fetch_vectors", - SearchAndFetchVectors.Success, - ), - (SimilarityMetric.EUCLIDEAN_SIMILARITY, [2, 10, 18], [3, 20, -0.01], "search", Search.Success), - ( - SimilarityMetric.EUCLIDEAN_SIMILARITY, - [2, 10, 18], - [3, 20, -0.01], - "search_and_fetch_vectors", - SearchAndFetchVectors.Success, - ), - ], -) -async def test_search_score_threshold_happy_path( - vector_index_client_async: PreviewVectorIndexClientAsync, - unique_vector_index_name_async: TUniqueVectorIndexNameAsync, - similarity_metric: SimilarityMetric, - distances: list[float], - thresholds: list[float], - search_method_name: str, - response: type[Search.Success] | type[SearchAndFetchVectors.Success], -) -> None: - index_name = unique_vector_index_name_async(vector_index_client_async) - num_dimensions = 2 - create_response = await vector_index_client_async.create_index( - index_name, num_dimensions=num_dimensions, similarity_metric=similarity_metric - ) - assert isinstance(create_response, CreateIndex.Success) - - items = [ - Item(id="test_item_1", vector=[1.0, 1.0]), - Item(id="test_item_2", vector=[-1.0, 1.0]), - Item(id="test_item_3", vector=[-1.0, -1.0]), - ] - upsert_response = await vector_index_client_async.upsert_item_batch( - index_name, - items=items, - ) - assert isinstance(upsert_response, UpsertItemBatch.Success) - - await sleep_async(2) - - query_vector = [2.0, 2.0] - search_hits = [SearchHit(id=f"test_item_{i+1}", score=distance) for i, distance in enumerate(distances)] - - search = getattr(vector_index_client_async, search_method_name) - search_response = await search(index_name, query_vector=query_vector, top_k=3, score_threshold=thresholds[0]) - assert isinstance(search_response, response) - assert search_response.hits == when_fetching_vectors_apply_vectors_to_hits(search_response, [search_hits[0]], items) - - search_response2 = await search(index_name, query_vector=query_vector, top_k=3, score_threshold=thresholds[1]) - assert isinstance(search_response2, response) - assert search_response2.hits == when_fetching_vectors_apply_vectors_to_hits(search_response, search_hits, items) - - search_response3 = await search(index_name, query_vector=query_vector, top_k=3, score_threshold=thresholds[2]) - assert isinstance(search_response3, response) - assert search_response3.hits == [] - - -async def test_search_with_filter_expression( - vector_index_client_async: PreviewVectorIndexClientAsync, - unique_vector_index_name_async: TUniqueVectorIndexNameAsync, -) -> None: - index_name = unique_vector_index_name_async(vector_index_client_async) - num_dimensions = 2 - create_response = await vector_index_client_async.create_index( - index_name, num_dimensions=num_dimensions, similarity_metric=SimilarityMetric.INNER_PRODUCT - ) - assert isinstance(create_response, CreateIndex.Success) - items = [ - Item( - id="test_item_1", - vector=[1.0, 1.0], - metadata={"str": "value1", "int": 0, "float": 0.0, "bool": True, "tags": ["a", "b", "c"]}, - ), - Item( - id="test_item_2", - vector=[-1.0, 1.0], - metadata={"str": "value2", "int": 5, "float": 5.0, "bool": False, "tags": ["a", "b"]}, - ), - Item( - id="test_item_3", - vector=[-1.0, -1.0], - metadata={"str": "value3", "int": 10, "float": 10.0, "bool": True, "tags": ["a", "d"]}, - ), - ] - upsert_response = await vector_index_client_async.upsert_item_batch( - index_name, - items=items, - ) - assert isinstance(upsert_response, UpsertItemBatch.Success) - await sleep_async(2) - - # Writing the test cases here instead of as a parameterized test because: - # 1. The search data is the same across tests, so no need to reindex each time. - # 2. It is 10x faster to run the tests this way. - for filter, expected_ids, test_case_name in [ - (Field("str") == "value1", ["test_item_1"], "string equality"), - (Field("str") != "value1", ["test_item_2", "test_item_3"], "string inequality"), - (Field("int") == 0, ["test_item_1"], "int equality"), - (Field("float") == 0.0, ["test_item_1"], "float equality"), - (Field("bool"), ["test_item_1", "test_item_3"], "bool equality"), - (~Field("bool"), ["test_item_2"], "bool inequality"), - (Field("int") > 5, ["test_item_3"], "int greater than"), - (Field("int") >= 5, ["test_item_2", "test_item_3"], "int greater than or equal to"), - (Field("float") > 5.0, ["test_item_3"], "float greater than"), - (Field("float") >= 5.0, ["test_item_2", "test_item_3"], "float greater than or equal to"), - (Field("int") < 5, ["test_item_1"], "int less than"), - (Field("int") <= 5, ["test_item_1", "test_item_2"], "int less than or equal to"), - (Field("float") < 5.0, ["test_item_1"], "float less than"), - (Field("float") <= 5.0, ["test_item_1", "test_item_2"], "float less than or equal to"), - (Field("tags").list_contains("a"), ["test_item_1", "test_item_2", "test_item_3"], "list contains a"), - (Field("tags").list_contains("b"), ["test_item_1", "test_item_2"], "list contains b"), - (Field("tags").list_contains("m"), [], "list contains m"), - ((Field("tags").list_contains("b")) & (Field("int") > 1), ["test_item_2"], "list contains b and int > 1"), - ( - (Field("tags").list_contains("b")) | (Field("int") > 1), - ["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 = cast(FilterExpression, filter) - search_response = await vector_index_client_async.search(index_name, query_vector=[2.0, 2.0], filter=filter) - assert isinstance( - search_response, Search.Success - ), f"Expected search {test_case_name!r} to succeed but got {search_response!r}" - assert ( - [hit.id for hit in search_response.hits] == expected_ids - ), f"Expected search {test_case_name!r} to return {expected_ids!r} but got {search_response.hits!r}" - - search_and_fetch_vectors_response = await vector_index_client_async.search_and_fetch_vectors( - index_name, query_vector=[2.0, 2.0], filter=filter - ) - assert isinstance( - search_and_fetch_vectors_response, SearchAndFetchVectors.Success - ), f"Expected search {test_case_name!r} to succeed but got {search_and_fetch_vectors_response!r}" - assert [hit.id for hit in search_and_fetch_vectors_response.hits] == expected_ids, ( - f"Expected search {test_case_name!r} to return {expected_ids!r} " - f"but got {search_and_fetch_vectors_response.hits!r}" - ) - - -async def test_delete_validates_index_name(vector_index_client_async: PreviewVectorIndexClientAsync) -> None: - response = await vector_index_client_async.delete_item_batch(index_name="", filter=[]) - assert isinstance(response, DeleteItemBatch.Error) - assert response.error_code == MomentoErrorCode.INVALID_ARGUMENT_ERROR - - -async def test_delete_items_by_id( - vector_index_client_async: PreviewVectorIndexClientAsync, - unique_vector_index_name_async: TUniqueVectorIndexNameAsync, -) -> None: - index_name = unique_vector_index_name_async(vector_index_client_async) - create_response = await vector_index_client_async.create_index( - index_name, num_dimensions=2, similarity_metric=SimilarityMetric.INNER_PRODUCT - ) - assert isinstance(create_response, CreateIndex.Success) - - upsert_response = await vector_index_client_async.upsert_item_batch( - index_name, - items=[ - Item(id="test_item_1", vector=[1.0, 2.0], metadata={"key": "value1"}), - Item(id="test_item_2", vector=[3.0, 4.0], metadata={"key": "value2"}), - Item(id="test_item_3", vector=[5.0, 6.0], metadata={"key": "value3"}), - Item(id="test_item_4", vector=[7.0, 8.0], metadata={"key": "value4"}), - ], - ) - assert isinstance(upsert_response, UpsertItemBatch.Success) - - await sleep_async(2) - - search_response = await vector_index_client_async.search(index_name, query_vector=[1.0, 2.0], top_k=10) - assert isinstance(search_response, Search.Success) - assert len(search_response.hits) == 4 - - assert search_response.hits == [ - SearchHit(id="test_item_4", score=23.0), - SearchHit(id="test_item_3", score=17.0), - SearchHit(id="test_item_2", score=11.0), - SearchHit(id="test_item_1", score=5.0), - ] - - delete_response = await vector_index_client_async.delete_item_batch( - index_name, filter=["test_item_1", "test_item_3"] - ) - assert isinstance(delete_response, DeleteItemBatch.Success) - - await sleep_async(2) - - search_response = await vector_index_client_async.search(index_name, query_vector=[1.0, 2.0], top_k=10) - assert isinstance(search_response, Search.Success) - assert len(search_response.hits) == 2 - - assert search_response.hits == [ - SearchHit(id="test_item_4", score=23.0), - SearchHit(id="test_item_2", score=11.0), - ] - - -async def test_delete_items_by_filter( - vector_index_client_async: PreviewVectorIndexClientAsync, - unique_vector_index_name_async: TUniqueVectorIndexNameAsync, -) -> None: - index_name = unique_vector_index_name_async(vector_index_client_async) - create_response = await vector_index_client_async.create_index(index_name, num_dimensions=2) - assert isinstance(create_response, CreateIndex.Success) - - items = [ - Item(id="test_item_1", vector=[1.0, 1.0], metadata={"key": "value1"}), - Item(id="test_item_2", vector=[3.0, 4.0], metadata={"key": "value2"}), - Item(id="test_item_3", vector=[5.0, 6.0], metadata={"key": "value1"}), - Item(id="test_item_4", vector=[7.0, 8.0], metadata={"key": "value2"}), - ] - upsert_response = await vector_index_client_async.upsert_item_batch( - index_name, - items=items, - ) - assert isinstance(upsert_response, UpsertItemBatch.Success) - - await sleep_async(2) - - search_response = await vector_index_client_async.search(index_name, query_vector=[1.0, 1.0], top_k=10) - assert isinstance(search_response, Search.Success) - assert len(search_response.hits) == 4 - - delete_response = await vector_index_client_async.delete_item_batch( - index_name, filter=filters.Equals("key", "value1") - ) - assert isinstance(delete_response, DeleteItemBatch.Success) - - await sleep_async(2) - - search_response = await vector_index_client_async.search(index_name, query_vector=[1.0, 1.0], top_k=10) - assert isinstance(search_response, Search.Success) - assert len(search_response.hits) == 2 - - -@pytest.mark.parametrize( - [ - "get_item_method_name", - "ids", - "expected_get_item_response", - "expected_get_item_values", - ], - [ - ("get_item_batch", [], GetItemBatch.Success, {}), - ("get_item_metadata_batch", [], GetItemMetadataBatch.Success, {}), - ("get_item_batch", ["missing_id"], GetItemBatch.Success, {}), - ( - "get_item_metadata_batch", - ["test_item_1"], - GetItemMetadataBatch.Success, - {"test_item_1": {"key1": "value1"}}, - ), - ("get_item_metadata_batch", ["missing_id"], GetItemMetadataBatch.Success, {}), - ( - "get_item_batch", - ["test_item_1"], - GetItemBatch.Success, - { - "test_item_1": Item(id="test_item_1", vector=[1.0, 1.0], metadata={"key1": "value1"}), - }, - ), - ( - "get_item_batch", - ["test_item_1", "missing_id", "test_item_2"], - GetItemBatch.Success, - { - "test_item_1": Item(id="test_item_1", vector=[1.0, 1.0], metadata={"key1": "value1"}), - "test_item_2": Item(id="test_item_2", vector=[-1.0, 1.0], metadata={}), - }, - ), - ( - "get_item_metadata_batch", - ["test_item_1", "missing_id", "test_item_2"], - GetItemMetadataBatch.Success, - { - "test_item_1": {"key1": "value1"}, - "test_item_2": {}, - }, - ), - ], -) -async def test_get_items_by_id( - vector_index_client_async: PreviewVectorIndexClientAsync, - unique_vector_index_name_async: TUniqueVectorIndexNameAsync, - get_item_method_name: str, - ids: list[str], - expected_get_item_response: type[GetItemMetadataBatch.Success] | type[GetItemBatch.Success], - expected_get_item_values: dict[str, Metadata] | dict[str, Item], -) -> None: - index_name = unique_vector_index_name_async(vector_index_client_async) - create_response = await vector_index_client_async.create_index(index_name, num_dimensions=2) - assert isinstance(create_response, CreateIndex.Success) - - items = [ - Item(id="test_item_1", vector=[1.0, 1.0], metadata={"key1": "value1"}), - Item(id="test_item_2", vector=[-1.0, 1.0]), - Item(id="test_item_3", vector=[-1.0, -1.0]), - ] - upsert_response = await vector_index_client_async.upsert_item_batch( - index_name, - items=items, - ) - assert isinstance(upsert_response, UpsertItemBatch.Success) - - await sleep_async(2) - - get_item = getattr(vector_index_client_async, get_item_method_name) - get_item_response = await get_item(index_name, ids) - assert isinstance(get_item_response, expected_get_item_response) - assert get_item_response.values == expected_get_item_values - - -async def test_count_items_on_missing_index( - vector_index_client_async: PreviewVectorIndexClientAsync, -) -> None: - response = await vector_index_client_async.count_items(index_name=uuid_str()) - assert isinstance(response, CountItems.Error) - assert response.error_code == MomentoErrorCode.NOT_FOUND_ERROR - - -async def test_count_items_on_empty_index( - vector_index_client_async: PreviewVectorIndexClientAsync, - unique_vector_index_name_async: TUniqueVectorIndexNameAsync, -) -> None: - index_name = unique_vector_index_name_async(vector_index_client_async) - - create_response = await vector_index_client_async.create_index(index_name, num_dimensions=2) - assert isinstance(create_response, CreateIndex.Success) - - count_response = await vector_index_client_async.count_items(index_name) - assert isinstance(count_response, CountItems.Success) - assert count_response.item_count == 0 - - -async def test_count_items_with_items( - vector_index_client_async: PreviewVectorIndexClientAsync, - unique_vector_index_name_async: TUniqueVectorIndexNameAsync, -) -> None: - num_items = 10 - index_name = unique_vector_index_name_async(vector_index_client_async) - - create_response = await vector_index_client_async.create_index(index_name, num_dimensions=2) - assert isinstance(create_response, CreateIndex.Success) - - items = [Item(id=f"test_item_{i}", vector=[i, i]) for i in range(num_items)] # type: list[Item] - upsert_response = await vector_index_client_async.upsert_item_batch( - index_name, - items=items, - ) - assert isinstance(upsert_response, UpsertItemBatch.Success) - - await sleep_async(2) - - count_response = await vector_index_client_async.count_items(index_name) - assert isinstance(count_response, CountItems.Success) - assert count_response.item_count == num_items - - num_items_to_delete = 5 - delete_response = await vector_index_client_async.delete_item_batch( - index_name, filter=[item.id for item in items[:num_items_to_delete]] - ) - assert isinstance(delete_response, DeleteItemBatch.Success) - - await sleep_async(2) - - count_response = await vector_index_client_async.count_items(index_name) - assert isinstance(count_response, CountItems.Success) - assert count_response.item_count == num_items - num_items_to_delete diff --git a/tests/utils.py b/tests/utils.py index 413580c5..85df2c09 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -4,22 +4,11 @@ import time import uuid -from momento.requests.vector_index import Item -from momento.responses.vector_index import Search, SearchAndFetchVectors -from momento.responses.vector_index.data.search import SearchHit -from momento.responses.vector_index.data.search_and_fetch_vectors import ( - SearchAndFetchVectorsHit, -) - def unique_test_cache_name() -> str: return f"python-test-{uuid_str()}" -def unique_test_vector_index_name() -> str: - return f"python-test-{uuid_str()}" - - def uuid_str() -> str: """Generate a UUID as a string. @@ -56,32 +45,3 @@ def sleep(seconds: int) -> None: async def sleep_async(seconds: int) -> None: await asyncio.sleep(seconds) - - -def when_fetching_vectors_apply_vectors_to_hits( - response: Search.Success | SearchAndFetchVectors.Success, - hits: list[SearchHit] | list[SearchAndFetchVectorsHit], - items: list[Item], -) -> list[SearchHit]: - """Normalizes the search hits according to the response. - - The tests will use this function to normalize the expected search hits - according to the response type. This allows us to specify one set of expected - search hits for both `search` and `search_and_fetch_vectors` requests. - - This method will augment the search hits with the original vectors when - the response is `SearchAndFetchVectors.Success`. Otherwise, it will return - the search hits as-is. - - Args: - response (Search.Success | SearchAndFetchVectors.Success): The search response. - hits (list[SearchHit] | list[SearchAndFetchVectorsHit]): The search hits. - items (list[Item]): The items that were indexed, used to fetch the original vectors. - - Returns: - list[SearchHit]: The normalized search hits. - """ - if isinstance(response, Search.Success): - return hits # type: ignore - item_index = {item.id: item for item in items} - return [SearchAndFetchVectorsHit.from_search_hit(hit, item_index[hit.id].vector) for hit in hits] From c6ade800ccd61cf701d3bceafb91863afed73f58 Mon Sep 17 00:00:00 2001 From: cprice404 Date: Fri, 19 Apr 2024 00:19:43 +0000 Subject: [PATCH 3/8] Update templated README.md file --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 16e679e2..d8f1bfa7 100644 --- a/README.md +++ b/README.md @@ -16,7 +16,7 @@ To get started with Momento you will need a Momento Auth Token. You can get one * Website: [https://www.gomomento.com/](https://www.gomomento.com/) * Momento Documentation: [https://docs.momentohq.com/](https://docs.momentohq.com/) * Getting Started: [https://docs.momentohq.com/getting-started](https://docs.momentohq.com/getting-started) -* Python SDK Documentation: [https://docs.momentohq.com/develop/sdks/python](https://docs.momentohq.com/develop/sdks/python) +* Python SDK Documentation: [https://docs.momentohq.com/sdks/python](https://docs.momentohq.com/sdks/python) * Discuss: [Momento Discord](https://discord.gg/3HkAKjUZGq) ## Packages From 9c3ef8c4dcfb5ddd2a7c406cec2a76129dde0ed1 Mon Sep 17 00:00:00 2001 From: Matt Straathof <11823378+bruuuuuuuce@users.noreply.github.com> Date: Tue, 7 May 2024 12:23:57 -0700 Subject: [PATCH 4/8] chore: add caching patterns section (#451) --- examples/py310/patterns.py | 51 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 51 insertions(+) create mode 100644 examples/py310/patterns.py diff --git a/examples/py310/patterns.py b/examples/py310/patterns.py new file mode 100644 index 00000000..93f85683 --- /dev/null +++ b/examples/py310/patterns.py @@ -0,0 +1,51 @@ +import asyncio +from datetime import timedelta + +from momento import ( + CacheClientAsync, + Configurations, + CredentialProvider, +) +from momento.responses import ( + CacheGet, + CacheSet, +) + +database: dict[str, str] = {} +database["test-key"] = "test-value" +async def example_patterns_WriteThroughCaching(cache_client: CacheClientAsync): + database.set("test-key", "test-value") + set_response = await cache_client.set("test-cache", "test-key", "test-value") + return + +# end example + +async def example_patterns_ReadAsideCaching(cache_client: CacheClientAsync): + get_response = await cache_client.get("test-cache", "test-key") + match get_response: + case CacheGet.Hit() as hit: + print(f"Retrieved value for key 'test-key': {hit.value_string}") + return + print(f"cache miss, fetching from database") + actual_value = database.get("test-key") + await cache_client.set("test-cache", "test-key", actual_value) + return + +# end example + +async def main(): + example_API_CredentialProviderFromEnvVar() + + await example_API_InstantiateCacheClient() + cache_client = await CacheClientAsync.create( + Configurations.Laptop.latest(), + CredentialProvider.from_environment_variable("MOMENTO_API_KEY"), + timedelta(seconds=60), + ) + + await example_patterns_ReadAsideCaching(cache_client) + await example_patterns_WriteThroughCaching(cache_client) + + +if __name__ == "__main__": + asyncio.run(main()) From c9c01ac66246610722b70614f7831985cf2c2445 Mon Sep 17 00:00:00 2001 From: Michael Landis Date: Thu, 16 May 2024 15:35:55 -0700 Subject: [PATCH 5/8] chore: add lambda zip example (#452) Adds a lambda example where we build the zip and upload to AWS via the console. Closes #819 --- examples/lambda/README.md | 24 +++++++++++-- examples/lambda/docker/Dockerfile | 4 +-- .../lambda/docker/lambda/aws_requirements.txt | 1 - .../lambda/docker/lambda/requirements.txt | 1 + examples/lambda/zip/.gitignore | 3 ++ examples/lambda/zip/Makefile | 27 ++++++++++++++ examples/lambda/zip/index.py | 36 +++++++++++++++++++ examples/lambda/zip/requirements.txt | 1 + 8 files changed, 92 insertions(+), 5 deletions(-) delete mode 100644 examples/lambda/docker/lambda/aws_requirements.txt create mode 100644 examples/lambda/docker/lambda/requirements.txt create mode 100644 examples/lambda/zip/.gitignore create mode 100644 examples/lambda/zip/Makefile create mode 100755 examples/lambda/zip/index.py create mode 100644 examples/lambda/zip/requirements.txt diff --git a/examples/lambda/README.md b/examples/lambda/README.md index 2da188ef..5ee5931e 100644 --- a/examples/lambda/README.md +++ b/examples/lambda/README.md @@ -19,7 +19,7 @@ The primary use is to provide a base for testing Momento in an AWS lambda enviro - Node version 14 or higher is required (for deploying the Cloudformation stack containing the Lambda) - To get started with Momento you will need a Momento Auth Token. You can get one from the [Momento Console](https://console.gomomento.com). Check out the [getting started](https://docs.momentohq.com/getting-started) guide for more information on obtaining an auth token. -## Deploying the Momento Python Lambda +## Deploying the Momento Python Lambda with Docker and AWS CDK The source code for the CDK application lives in the `infrastructure` directory. To build and deploy it you will first need to install the dependencies: @@ -42,4 +42,24 @@ npm run cdk deploy The lambda does not set up a way to access itself externally, so to run it, you will have to go to `MomentoDockerLambda` in AWS Lambda and run a test. -The lambda is set up to make set and get calls for the key 'key' in the cache 'cache'. You can play around with the code by changing the `docker/lambda/index.py` file. Remember to update `docker/lambda/aws_requirements.txt` file if you add additional Python dependencies. \ No newline at end of file +The lambda is set up to make set and get calls for the key 'key' in the cache 'cache'. You can play around with the code by changing the `docker/lambda/index.py` file. Remember to update `docker/lambda/aws_requirements.txt` file if you add additional Python dependencies. + +## Deploying the Momento Python Lambda as a Zip File on AWS Lambda + +Alternatively, we can deploy the Momento Python Lambda as a Zip File on AWS Lambda. We can do this using the `zip` directory in this example. + +Follow these steps to create the zip and deploy it to AWS Lambda using the AWS Management Console: + +1. Run `make dist` in the `zip` directory to package the lambda for upload as `dist.zip`. + :tipp: Check out the Makefile for important build steps to package for AWS Lambda. + +2. Create a new Lambda function by selecting "Author from scratch". +3. Set the function name to `momento-lambda-demo`. +4. Choose the runtime as `Python 3.8` (you can adjust this as desired). +5. Select the architecture as `x86_64`. +6. Click on "Create function" to create the Lambda function. +7. In the "Code" tab, choose "Upload from the zip" as the code source. +8. Under "Runtime settings", set the Handler to index.handler. +9. Switch to the "Configuration" tab. +10. Set the environment variable `MOMENTO_API_KEY` to your API key. +11. Finally, go to the "Test" tab to test your Lambda function. diff --git a/examples/lambda/docker/Dockerfile b/examples/lambda/docker/Dockerfile index 000cb80d..7c5b7769 100644 --- a/examples/lambda/docker/Dockerfile +++ b/examples/lambda/docker/Dockerfile @@ -4,10 +4,10 @@ WORKDIR /var/task # Copy the lambda and the requirements file COPY lambda/index.py . -COPY lambda/aws_requirements.txt . +COPY lambda/requirements.txt . # Install Python dependencies -RUN pip install -r aws_requirements.txt -t . +RUN pip install -r requirements.txt -t . # Set the CMD to your lambda (could also be done as a parameter override outside of the Dockerfile) CMD ["index.handler"] diff --git a/examples/lambda/docker/lambda/aws_requirements.txt b/examples/lambda/docker/lambda/aws_requirements.txt deleted file mode 100644 index bec69816..00000000 --- a/examples/lambda/docker/lambda/aws_requirements.txt +++ /dev/null @@ -1 +0,0 @@ -momento==1.8.0 \ No newline at end of file diff --git a/examples/lambda/docker/lambda/requirements.txt b/examples/lambda/docker/lambda/requirements.txt new file mode 100644 index 00000000..63eade95 --- /dev/null +++ b/examples/lambda/docker/lambda/requirements.txt @@ -0,0 +1 @@ +momento==1.8.0 diff --git a/examples/lambda/zip/.gitignore b/examples/lambda/zip/.gitignore new file mode 100644 index 00000000..06674e38 --- /dev/null +++ b/examples/lambda/zip/.gitignore @@ -0,0 +1,3 @@ +.venv +dist +dist.zip diff --git a/examples/lambda/zip/Makefile b/examples/lambda/zip/Makefile new file mode 100644 index 00000000..d2bc1127 --- /dev/null +++ b/examples/lambda/zip/Makefile @@ -0,0 +1,27 @@ +.PHONY: setup +setup: + python -m venv .venv + .venv/bin/pip install --upgrade pip + .venv/bin/pip install -r requirements.txt + +.PHONY: dist +dist: + rm -rf dist dist.zip +# Some important notes: +# 1. The `--platform` option is used to specify the target platform. In this case, we are targeting the manylinux2014_x86_64 platform for AWS Lambda. +# 2. The `--target` option is used to specify the target directory where the dependencies will be installed. +# 3. The `--implementation` option is used to specify the Python implementation. In this case, we are using the CPython implementation. +# 4. The `--python-version` option is used to specify the Python version. In this case, we are using Python 3.8. You can change this to match the Python version used by AWS Lambda. +# 5. The `--only-binary` option is used to specify that only binary distributions should be used. This is important because AWS Lambda does not support building from source. +# 6. The `--upgrade` option is used to ensure that the dependencies are up-to-date. +# 7. The `-r requirements.txt` option is used to specify the requirements file that contains the dependencies to install. + .venv/bin/pip install \ + --platform manylinux2014_x86_64 \ + --target=dist \ + --implementation cp \ + --python-version 3.8 \ + --only-binary=:all: \ + --upgrade \ + -r requirements.txt + cp index.py dist + cd dist && zip -9r ../dist.zip . diff --git a/examples/lambda/zip/index.py b/examples/lambda/zip/index.py new file mode 100755 index 00000000..986b4209 --- /dev/null +++ b/examples/lambda/zip/index.py @@ -0,0 +1,36 @@ +from datetime import timedelta + +from momento import CacheClient, Configurations, CredentialProvider +from momento.responses import CacheGet, CacheSet, CreateCache + + +def handler(event, lambda_context): + cache_name = "default-cache" + with CacheClient.create( + configuration=Configurations.Lambda.latest(), + credential_provider=CredentialProvider.from_environment_variable("MOMENTO_API_KEY"), + default_ttl=timedelta(seconds=60), + ) as cache_client: + create_cache_response = cache_client.create_cache(cache_name) + + if isinstance(create_cache_response, CreateCache.CacheAlreadyExists): + print(f"Cache with name: {cache_name} already exists.") + elif isinstance(create_cache_response, CreateCache.Error): + raise create_cache_response.inner_exception + + print("Setting Key: key to Value: value") + set_response = cache_client.set(cache_name, "key", "value") + + if isinstance(set_response, CacheSet.Error): + raise set_response.inner_exception + + print("Getting Key: key") + get_response = cache_client.get(cache_name, "key") + + if isinstance(get_response, CacheGet.Hit): + print(f"Look up resulted in a hit: {get_response}") + print(f"Looked up Value: {get_response.value_string!r}") + elif isinstance(get_response, CacheGet.Miss): + print("Look up resulted in a: miss. This is unexpected.") + elif isinstance(get_response, CacheGet.Error): + raise get_response.inner_exception diff --git a/examples/lambda/zip/requirements.txt b/examples/lambda/zip/requirements.txt new file mode 100644 index 00000000..f1aad823 --- /dev/null +++ b/examples/lambda/zip/requirements.txt @@ -0,0 +1 @@ +momento==1.20.1 From 8ceff0fbb7463e5a2e78c78192a6c7f07af76222 Mon Sep 17 00:00:00 2001 From: Michael Landis Date: Thu, 16 May 2024 16:56:08 -0700 Subject: [PATCH 6/8] chore: fix markdown in lambda zip example (#453) The previous PR had incorrect markdown for the tip section. This corrects that. --- examples/lambda/README.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/examples/lambda/README.md b/examples/lambda/README.md index 5ee5931e..e52404b7 100644 --- a/examples/lambda/README.md +++ b/examples/lambda/README.md @@ -44,14 +44,15 @@ The lambda does not set up a way to access itself externally, so to run it, you The lambda is set up to make set and get calls for the key 'key' in the cache 'cache'. You can play around with the code by changing the `docker/lambda/index.py` file. Remember to update `docker/lambda/aws_requirements.txt` file if you add additional Python dependencies. -## Deploying the Momento Python Lambda as a Zip File on AWS Lambda +## Deploying the Momento Python Lambda as a Zip File on AWS Lambda with the AWS Management Console Alternatively, we can deploy the Momento Python Lambda as a Zip File on AWS Lambda. We can do this using the `zip` directory in this example. Follow these steps to create the zip and deploy it to AWS Lambda using the AWS Management Console: 1. Run `make dist` in the `zip` directory to package the lambda for upload as `dist.zip`. - :tipp: Check out the Makefile for important build steps to package for AWS Lambda. + +> :bulb: **Tip**: Check out the Makefile for important build steps to package for AWS Lambda. 2. Create a new Lambda function by selecting "Author from scratch". 3. Set the function name to `momento-lambda-demo`. From 512d982256bffa9ae2680d210fb07c726b36ea6c Mon Sep 17 00:00:00 2001 From: Michael Landis Date: Mon, 1 Jul 2024 11:54:50 -0700 Subject: [PATCH 7/8] chore: use ubuntu-24.04 for ci where possible (#454) Updates the CI workflows to use ubuntu-24.04 for python 3.9+ and ubuntu-20.04 for python 3.7 and 3.8. Previously we could not run on ubuntu-22.04 due to a difficult-to-diagnose timeout problem. This timeout problem manifested on Azure VMs with ubuntu-22.04 but not on other platforms. This problem does not manifest on ubuntu-24.04, so we can use that where possible. Because python 3.7 and python 3.8 are not available on ubuntu-24.04, we still run those on ubuntu-20.04. At some point we can look in to installing via pyenv or manually on ubuntu-24.04. --- .github/workflows/on-pull-request.yml | 52 +++++++++++-------- .github/workflows/on-push-to-main-branch.yml | 2 +- .../workflows/on-push-to-release-branch.yml | 47 +++++++++++------ 3 files changed, 64 insertions(+), 37 deletions(-) diff --git a/.github/workflows/on-pull-request.yml b/.github/workflows/on-pull-request.yml index e11e58ee..8e0dbaa1 100644 --- a/.github/workflows/on-pull-request.yml +++ b/.github/workflows/on-pull-request.yml @@ -6,14 +6,27 @@ on: jobs: test: - runs-on: ubuntu-20.04 strategy: matrix: - python-version: ["3.7", "3.8", "3.9", "3.10", "3.11"] + # TODO: one of the examples dependencies does not support 3.11 + os: [ubuntu-24.04] + python-version: ["3.9", "3.10", "3.11"] new-python-protobuf: ["true"] include: + # 3.7 and 3.8 are no longer available on ubuntu-24.04 + # We run on 20.04 which was the last version where this worked. + # If support for 20.04 becomes an issue, we can install 3.7 and 3.8 + # with pyenv or manually. + - python-version: "3.7" + new-python-protobuf: "true" + os: ubuntu-20.04 - python-version: "3.7" new-python-protobuf: "false" + os: ubuntu-20.04 + - python-version: "3.8" + new-python-protobuf: "true" + os: ubuntu-20.04 + runs-on: ${{ matrix.os }} env: TEST_API_KEY: ${{ secrets.ALPHA_TEST_AUTH_TOKEN }} @@ -27,17 +40,17 @@ jobs: env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + - name: Install and configure Poetry + uses: snok/install-poetry@v1 + with: + version: 1.3.1 + virtualenvs-in-project: true + - name: Setup Python ${{ matrix.python-version }} uses: actions/setup-python@v4 with: python-version: ${{ matrix.python-version }} - - - name: Bootstrap poetry - run: | - curl -sL https://install.python-poetry.org | python - -y --version 1.3.1 - - - name: Configure poetry - run: poetry config virtualenvs.in-project true + cache: poetry - name: Install dependencies run: poetry install @@ -62,15 +75,11 @@ jobs: run: poetry run pytest -p no:sugar -q test-examples: - runs-on: ubuntu-20.04 + runs-on: ubuntu-24.04 strategy: matrix: # TODO: one of the examples dependencies does not support 3.11 include: - - python-version: "3.7" - package: prepy310 - - python-version: "3.8" - package: prepy310 - python-version: "3.9" package: prepy310 - python-version: "3.10" @@ -88,10 +97,18 @@ jobs: MOMENTO_API_KEY: ${{ secrets.ALPHA_TEST_AUTH_TOKEN }} steps: - uses: actions/checkout@v3 + + - name: Install and configure Poetry + uses: snok/install-poetry@v1 + with: + version: 1.3.1 + virtualenvs-in-project: true + - name: Python SDK sample uses: actions/setup-python@v4 with: python-version: ${{ matrix.python-version }} + cache: poetry - name: Verify README generation uses: momentohq/standards-and-practices/github-actions/oss-readme-template@gh-actions-v2 @@ -105,13 +122,6 @@ jobs: template_file: ./README.template.md output_file: ./README.md - - name: Bootstrap poetry - run: | - curl -sL https://install.python-poetry.org | python - -y --version 1.3.1 - - - name: Configure poetry - run: poetry config virtualenvs.in-project true - - name: Install dependencies working-directory: ./examples run: poetry install diff --git a/.github/workflows/on-push-to-main-branch.yml b/.github/workflows/on-push-to-main-branch.yml index 3080db3c..17620e66 100644 --- a/.github/workflows/on-push-to-main-branch.yml +++ b/.github/workflows/on-push-to-main-branch.yml @@ -6,7 +6,7 @@ on: jobs: build: - runs-on: ubuntu-20.04 + runs-on: ubuntu-24.04 steps: - name: Setup repo diff --git a/.github/workflows/on-push-to-release-branch.yml b/.github/workflows/on-push-to-release-branch.yml index ca0a4863..a809cd24 100644 --- a/.github/workflows/on-push-to-release-branch.yml +++ b/.github/workflows/on-push-to-release-branch.yml @@ -6,7 +6,7 @@ on: jobs: release: - runs-on: ubuntu-20.04 + runs-on: ubuntu-24.04 outputs: version: ${{ steps.release.outputs.release }} steps: @@ -30,14 +30,30 @@ jobs: run: echo "::set-output name=release::${{ steps.semrel.outputs.version }}" test: - runs-on: ubuntu-20.04 strategy: matrix: - python-version: ["3.7", "3.8", "3.9", "3.10", "3.11"] + os: [ubuntu-20.04] + python-version: ["3.9", "3.10", "3.11"] new-python-protobuf: ["true"] include: + # 3.7 and 3.8 are no longer available on ubuntu-24.04 + # We run on 20.04 which was the last version where this worked. + # If support for 20.04 becomes an issue, we can install 3.7 and 3.8 + # with pyenv or manually. - python-version: "3.7" + new-python-protobuf: "true" + os: ubuntu-20.04 + - python-version: "3.7" + new-python-protobuf: "false" + os: ubuntu-20.04 + - python-version: "3.8" + new-python-protobuf: "true" + os: ubuntu-20.04 + - python-version: "3.8" new-python-protobuf: "false" + os: ubuntu-20.04 + runs-on: ${{ matrix.os }} + env: TEST_API_KEY: ${{ secrets.ALPHA_TEST_AUTH_TOKEN }} TEST_CACHE_NAME: python-integration-test-${{ matrix.python-version }}-${{ matrix.new-python-protobuf }}-${{ github.sha }} @@ -45,18 +61,17 @@ jobs: steps: - uses: actions/checkout@v3 + - name: Install and configure Poetry + uses: snok/install-poetry@v1 + with: + version: 1.3.1 + virtualenvs-in-project: true + - name: Setup Python ${{ matrix.python-version }} uses: actions/setup-python@v4 with: python-version: ${{ matrix.python-version }} - - name: Bootstrap poetry - run: | - curl -sL https://install.python-poetry.org | python - -y --version 1.3.1 - - - name: Configure poetry - run: poetry config virtualenvs.in-project true - - name: Install dependencies run: poetry install @@ -80,21 +95,23 @@ jobs: run: poetry run pytest -p no:sugar -q publish: - runs-on: ubuntu-20.04 + runs-on: ubuntu-24.04 needs: [release, test] steps: - uses: actions/checkout@v3 + - name: Install and configure Poetry + uses: snok/install-poetry@v1 + with: + version: 1.3.1 + virtualenvs-in-project: true + - name: Setup Python 3.10 uses: actions/setup-python@v4 with: python-version: "3.10" - - name: Bootstrap poetry - run: | - curl -sL https://install.python-poetry.org | python - -y --version 1.3.1 - - name: Bump version run: poetry version ${{ needs.release.outputs.version }} From 7a59011ffbe8aa2dbebc398a936807ba13be0ecd Mon Sep 17 00:00:00 2001 From: Michael Landis Date: Mon, 1 Jul 2024 16:34:22 -0700 Subject: [PATCH 8/8] chore: add agent details and new runtime version header (#456) Agent details and runtime-version header Adds the service client type to the agent header details, and adds a new runtime-version header for the cache and topics clients. The runtime version information looks like: major.minor.patch (releasetype serial) for example: python 3.11.2 (final 0) The releasetype covers alpha vs beta vs experimental vs final versions. The serial covers particular alpha/beta release versioning. CI fix We also fix the poetry installation in CI for python 3.7. Previously we installed poetry first, then installed python setting the cache as poetry. On older versions of ubuntu this does not work since it uses the wrong python version to install poetry. So we change the installation order to how it was before: first install python, then install poetry. --- .github/workflows/on-pull-request.yml | 11 ++++---- .../workflows/on-push-to-release-branch.yml | 10 +++---- src/momento/internal/_utilities/__init__.py | 2 ++ .../internal/_utilities/_client_type.py | 13 ++++++++++ .../_utilities/_python_runtime_version.py | 10 +++++++ .../aio/_add_header_client_interceptor.py | 2 +- src/momento/internal/aio/_scs_grpc_manager.py | 26 ++++++++++++------- .../internal/synchronous/_scs_grpc_manager.py | 25 +++++++++++------- 8 files changed, 69 insertions(+), 30 deletions(-) create mode 100644 src/momento/internal/_utilities/_client_type.py create mode 100644 src/momento/internal/_utilities/_python_runtime_version.py diff --git a/.github/workflows/on-pull-request.yml b/.github/workflows/on-pull-request.yml index 8e0dbaa1..29337d50 100644 --- a/.github/workflows/on-pull-request.yml +++ b/.github/workflows/on-pull-request.yml @@ -40,18 +40,17 @@ jobs: env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + - name: Setup Python ${{ matrix.python-version }} + uses: actions/setup-python@v4 + with: + python-version: ${{ matrix.python-version }} + - name: Install and configure Poetry uses: snok/install-poetry@v1 with: version: 1.3.1 virtualenvs-in-project: true - - name: Setup Python ${{ matrix.python-version }} - uses: actions/setup-python@v4 - with: - python-version: ${{ matrix.python-version }} - cache: poetry - - name: Install dependencies run: poetry install diff --git a/.github/workflows/on-push-to-release-branch.yml b/.github/workflows/on-push-to-release-branch.yml index a809cd24..5b0793f4 100644 --- a/.github/workflows/on-push-to-release-branch.yml +++ b/.github/workflows/on-push-to-release-branch.yml @@ -61,17 +61,17 @@ jobs: steps: - uses: actions/checkout@v3 + - name: Setup Python ${{ matrix.python-version }} + uses: actions/setup-python@v4 + with: + python-version: ${{ matrix.python-version }} + - name: Install and configure Poetry uses: snok/install-poetry@v1 with: version: 1.3.1 virtualenvs-in-project: true - - name: Setup Python ${{ matrix.python-version }} - uses: actions/setup-python@v4 - with: - python-version: ${{ matrix.python-version }} - - name: Install dependencies run: poetry install diff --git a/src/momento/internal/_utilities/__init__.py b/src/momento/internal/_utilities/__init__.py index 3937cecb..5f64b76c 100644 --- a/src/momento/internal/_utilities/__init__.py +++ b/src/momento/internal/_utilities/__init__.py @@ -1,3 +1,4 @@ +from ._client_type import ClientType from ._data_validation import ( _as_bytes, _gen_dictionary_fields_as_bytes, @@ -14,4 +15,5 @@ _validate_ttl, ) from ._momento_version import momento_version +from ._python_runtime_version import PYTHON_RUNTIME_VERSION from ._time import _timedelta_to_ms diff --git a/src/momento/internal/_utilities/_client_type.py b/src/momento/internal/_utilities/_client_type.py new file mode 100644 index 00000000..297d9266 --- /dev/null +++ b/src/momento/internal/_utilities/_client_type.py @@ -0,0 +1,13 @@ +"""Enumerates the types of clients that can be used. + +Used to populate the agent header in gRPC requests. +""" + +from enum import Enum + + +class ClientType(Enum): + """Describes the type of client that is being used.""" + + CACHE = "cache" + TOPIC = "topic" diff --git a/src/momento/internal/_utilities/_python_runtime_version.py b/src/momento/internal/_utilities/_python_runtime_version.py new file mode 100644 index 00000000..bdab59ff --- /dev/null +++ b/src/momento/internal/_utilities/_python_runtime_version.py @@ -0,0 +1,10 @@ +"""Python runtime version information. + +Used to populate the `runtime-version` header in gRPC requests. +""" +import sys + +PYTHON_RUNTIME_VERSION = ( + f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro} " + f"({sys.version_info.releaselevel} {sys.version_info.serial})" +) diff --git a/src/momento/internal/aio/_add_header_client_interceptor.py b/src/momento/internal/aio/_add_header_client_interceptor.py index 6aa11957..1bb2edfe 100644 --- a/src/momento/internal/aio/_add_header_client_interceptor.py +++ b/src/momento/internal/aio/_add_header_client_interceptor.py @@ -10,7 +10,7 @@ class Header: - once_only_headers = ["agent"] + once_only_headers = ["agent", "runtime-version"] def __init__(self, name: str, value: str): self.name = name diff --git a/src/momento/internal/aio/_scs_grpc_manager.py b/src/momento/internal/aio/_scs_grpc_manager.py index f67be248..6cca62d7 100644 --- a/src/momento/internal/aio/_scs_grpc_manager.py +++ b/src/momento/internal/aio/_scs_grpc_manager.py @@ -13,7 +13,7 @@ from momento.config import Configuration, TopicConfiguration from momento.config.transport.transport_strategy import StaticGrpcConfiguration from momento.errors.exceptions import ConnectionException -from momento.internal._utilities import momento_version +from momento.internal._utilities import PYTHON_RUNTIME_VERSION, ClientType, momento_version from momento.internal._utilities._channel_credentials import ( channel_credentials_from_root_certs_or_default, ) @@ -42,7 +42,9 @@ def __init__(self, configuration: Configuration, credential_provider: Credential self._secure_channel = grpc.aio.secure_channel( target=credential_provider.control_endpoint, credentials=channel_credentials_from_root_certs_or_default(configuration), - interceptors=_interceptors(credential_provider.auth_token, configuration.get_retry_strategy()), + interceptors=_interceptors( + credential_provider.auth_token, ClientType.CACHE, configuration.get_retry_strategy() + ), options=grpc_control_channel_options_from_grpc_config( grpc_config=configuration.get_transport_strategy().get_grpc_configuration(), ), @@ -65,7 +67,9 @@ def __init__(self, configuration: Configuration, credential_provider: Credential self._secure_channel = grpc.aio.secure_channel( target=credential_provider.cache_endpoint, credentials=channel_credentials_from_root_certs_or_default(configuration), - interceptors=_interceptors(credential_provider.auth_token, configuration.get_retry_strategy()), + interceptors=_interceptors( + credential_provider.auth_token, ClientType.CACHE, configuration.get_retry_strategy() + ), # Here is where you would pass override configuration to the underlying C gRPC layer. # However, I have tried several different tuning options here and did not see any # performance improvements, so sticking with the defaults for now. @@ -142,7 +146,7 @@ def __init__(self, configuration: TopicConfiguration, credential_provider: Crede self._secure_channel = grpc.aio.secure_channel( target=credential_provider.cache_endpoint, credentials=grpc.ssl_channel_credentials(), - interceptors=_interceptors(credential_provider.auth_token, None), + interceptors=_interceptors(credential_provider.auth_token, ClientType.TOPIC, None), options=grpc_data_channel_options_from_grpc_config(grpc_config), ) @@ -165,7 +169,7 @@ def __init__(self, configuration: TopicConfiguration, credential_provider: Crede self._secure_channel = grpc.aio.secure_channel( target=credential_provider.cache_endpoint, credentials=grpc.ssl_channel_credentials(), - interceptors=_stream_interceptors(credential_provider.auth_token), + interceptors=_stream_interceptors(credential_provider.auth_token, ClientType.TOPIC), options=grpc_data_channel_options_from_grpc_config(grpc_config), ) @@ -176,10 +180,13 @@ def async_stub(self) -> pubsub_client.PubsubStub: return pubsub_client.PubsubStub(self._secure_channel) # type: ignore[no-untyped-call] -def _interceptors(auth_token: str, retry_strategy: Optional[RetryStrategy] = None) -> list[grpc.aio.ClientInterceptor]: +def _interceptors( + auth_token: str, client_type: ClientType, retry_strategy: Optional[RetryStrategy] = None +) -> list[grpc.aio.ClientInterceptor]: headers = [ Header("authorization", auth_token), - Header("agent", f"python:{_ControlGrpcManager.version}"), + Header("agent", f"python:{client_type.value}:{_ControlGrpcManager.version}"), + Header("runtime-version", f"python {PYTHON_RUNTIME_VERSION}"), ] return list( filter( @@ -192,9 +199,10 @@ def _interceptors(auth_token: str, retry_strategy: Optional[RetryStrategy] = Non ) -def _stream_interceptors(auth_token: str) -> list[grpc.aio.UnaryStreamClientInterceptor]: +def _stream_interceptors(auth_token: str, client_type: ClientType) -> list[grpc.aio.UnaryStreamClientInterceptor]: headers = [ Header("authorization", auth_token), - Header("agent", f"python:{_PubsubGrpcStreamManager.version}"), + Header("agent", f"python:{client_type.value}:{_PubsubGrpcStreamManager.version}"), + Header("runtime-version", f"python {PYTHON_RUNTIME_VERSION}"), ] return [AddHeaderStreamingClientInterceptor(headers)] diff --git a/src/momento/internal/synchronous/_scs_grpc_manager.py b/src/momento/internal/synchronous/_scs_grpc_manager.py index e67b78f3..c52bbd78 100644 --- a/src/momento/internal/synchronous/_scs_grpc_manager.py +++ b/src/momento/internal/synchronous/_scs_grpc_manager.py @@ -14,7 +14,7 @@ from momento.config import Configuration, TopicConfiguration from momento.config.transport.transport_strategy import StaticGrpcConfiguration from momento.errors.exceptions import ConnectionException -from momento.internal._utilities import momento_version +from momento.internal._utilities import PYTHON_RUNTIME_VERSION, ClientType, momento_version from momento.internal._utilities._channel_credentials import ( channel_credentials_from_root_certs_or_default, ) @@ -46,7 +46,8 @@ def __init__(self, configuration: Configuration, credential_provider: Credential ), ) intercept_channel = grpc.intercept_channel( - self._secure_channel, *_interceptors(credential_provider.auth_token, configuration.get_retry_strategy()) + self._secure_channel, + *_interceptors(credential_provider.auth_token, ClientType.CACHE, configuration.get_retry_strategy()), ) self._stub = control_client.ScsControlStub(intercept_channel) # type: ignore[no-untyped-call] @@ -73,7 +74,8 @@ def __init__(self, configuration: Configuration, credential_provider: Credential ) intercept_channel = grpc.intercept_channel( - self._secure_channel, *_interceptors(credential_provider.auth_token, configuration.get_retry_strategy()) + self._secure_channel, + *_interceptors(credential_provider.auth_token, ClientType.CACHE, configuration.get_retry_strategy()), ) self._stub = cache_client.ScsStub(intercept_channel) # type: ignore[no-untyped-call] @@ -164,7 +166,7 @@ def __init__(self, configuration: TopicConfiguration, credential_provider: Crede options=grpc_data_channel_options_from_grpc_config(grpc_config), ) intercept_channel = grpc.intercept_channel( - self._secure_channel, *_interceptors(credential_provider.auth_token, None) + self._secure_channel, *_interceptors(credential_provider.auth_token, ClientType.TOPIC, None) ) self._stub = pubsub_client.PubsubStub(intercept_channel) # type: ignore[no-untyped-call] @@ -190,7 +192,7 @@ def __init__(self, configuration: TopicConfiguration, credential_provider: Crede options=grpc_data_channel_options_from_grpc_config(grpc_config), ) intercept_channel = grpc.intercept_channel( - self._secure_channel, *_stream_interceptors(credential_provider.auth_token) + self._secure_channel, *_stream_interceptors(credential_provider.auth_token, ClientType.TOPIC) ) self._stub = pubsub_client.PubsubStub(intercept_channel) # type: ignore[no-untyped-call] @@ -202,9 +204,13 @@ def stub(self) -> pubsub_client.PubsubStub: def _interceptors( - auth_token: str, retry_strategy: Optional[RetryStrategy] = None + auth_token: str, client_type: ClientType, retry_strategy: Optional[RetryStrategy] = None ) -> list[grpc.UnaryUnaryClientInterceptor]: - headers = [Header("authorization", auth_token), Header("agent", f"python:{_ControlGrpcManager.version}")] + headers = [ + Header("authorization", auth_token), + Header("agent", f"python:{client_type.value}:{_ControlGrpcManager.version}"), + Header("runtime-version", f"python {PYTHON_RUNTIME_VERSION}"), + ] return list( filter( None, [AddHeaderClientInterceptor(headers), RetryInterceptor(retry_strategy) if retry_strategy else None] @@ -212,9 +218,10 @@ def _interceptors( ) -def _stream_interceptors(auth_token: str) -> list[grpc.UnaryStreamClientInterceptor]: +def _stream_interceptors(auth_token: str, client_type: ClientType) -> list[grpc.UnaryStreamClientInterceptor]: headers = [ Header("authorization", auth_token), - Header("agent", f"python:{_PubsubGrpcStreamManager.version}"), + Header("agent", f"python:{client_type.value}:{_PubsubGrpcStreamManager.version}"), + Header("runtime-version", f"python {PYTHON_RUNTIME_VERSION}"), ] return [AddHeaderStreamingClientInterceptor(headers)]