No commit message

momentohq · Jan 23, 2024 · 14f377b · 14f377b
2 parents 76acf10 + 5ba6988
commit 14f377b
Show file tree

Hide file tree

Showing 11 changed files with 240 additions and 8 deletions.
diff --git a/README.md b/README.md
@@ -1,7 +1,7 @@
 <head>
   <meta name="Momento Python Client Library Documentation" content="Python client software development kit for Momento Cache">
 </head>
-<img src="https://docs.momentohq.com/img/logo.svg" alt="logo" width="400"/>
+<img src="https://docs.momentohq.com/img/momento-logo-forest.svg" alt="logo" width="400"/>
 
 [![project status](https://momentohq.github.io/standards-and-practices/badges/project-status-official.svg)](https://github.com/momentohq/standards-and-practices/blob/main/docs/momento-on-github.md)
 [![project stability](https://momentohq.github.io/standards-and-practices/badges/project-stability-stable.svg)](https://github.com/momentohq/standards-and-practices/blob/main/docs/momento-on-github.md)

diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -30,7 +30,7 @@ exclude = ["src/momento/internal/codegen.py"]
 [tool.poetry.dependencies]
 python = "^3.7"
 
-momento-wire-types = "^0.98.1"
+momento-wire-types = "^0.102.1"
 grpcio = "^1.46.0"
 # note if you bump this presigned url test need be updated
 pyjwt = "^2.4.0"

diff --git a/src/momento/internal/aio/_vector_index_data_client.py b/src/momento/internal/aio/_vector_index_data_client.py
@@ -15,6 +15,8 @@
 from momento.internal.services import Service
 from momento.requests.vector_index import AllMetadata, FilterExpression, Item
 from momento.responses.vector_index import (
+    CountItems,
+    CountItemsResponse,
     DeleteItemBatch,
     DeleteItemBatchResponse,
     GetItemBatch,
@@ -50,6 +52,27 @@ def __init__(self, configuration: VectorIndexConfiguration, credential_provider:
     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,

diff --git a/src/momento/internal/synchronous/_vector_index_data_client.py b/src/momento/internal/synchronous/_vector_index_data_client.py
@@ -15,6 +15,8 @@
 from momento.internal.synchronous._vector_index_grpc_manager import _VectorIndexDataGrpcManager
 from momento.requests.vector_index import AllMetadata, FilterExpression, Item
 from momento.responses.vector_index import (
+    CountItems,
+    CountItemsResponse,
     DeleteItemBatch,
     DeleteItemBatchResponse,
     GetItemBatch,
@@ -50,6 +52,27 @@ def __init__(self, configuration: VectorIndexConfiguration, credential_provider:
     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,

diff --git a/src/momento/responses/vector_index/__init__.py b/src/momento/responses/vector_index/__init__.py
@@ -3,6 +3,7 @@
 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,
@@ -19,6 +20,8 @@
 from .response import VectorIndexResponse
 
 __all__ = [
+    "CountItems",
+    "CountItemsResponse",
     "CreateIndex",
     "CreateIndexResponse",
     "DeleteIndex",

diff --git a/src/momento/responses/vector_index/data/count_items.py b/src/momento/responses/vector_index/data/count_items.py
@@ -0,0 +1,37 @@
+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/vector_index_client.py b/src/momento/vector_index_client.py
@@ -36,6 +36,7 @@
 
 from momento.requests.vector_index import AllMetadata, FilterExpression, Item, SimilarityMetric
 from momento.responses.vector_index import (
+    CountItemsResponse,
     CreateIndexResponse,
     DeleteIndexResponse,
     DeleteItemBatchResponse,
@@ -116,6 +117,19 @@ def __exit__(
         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,

diff --git a/src/momento/vector_index_client_async.py b/src/momento/vector_index_client_async.py
@@ -36,6 +36,7 @@
 
 from momento.requests.vector_index import AllMetadata, FilterExpression, Item, SimilarityMetric
 from momento.responses.vector_index import (
+    CountItemsResponse,
     CreateIndexResponse,
     DeleteIndexResponse,
     DeleteItemBatchResponse,
@@ -116,6 +117,19 @@ async def __aexit__(
         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,

diff --git a/tests/momento/vector_index_client/test_data.py b/tests/momento/vector_index_client/test_data.py
@@ -8,6 +8,7 @@
 from momento.errors import MomentoErrorCode
 from momento.requests.vector_index import ALL_METADATA, Field, FilterExpression, Item, SimilarityMetric
 from momento.responses.vector_index import (
+    CountItems,
     CreateIndex,
     DeleteItemBatch,
     GetItemBatch,
@@ -19,7 +20,7 @@
 )
 
 from tests.conftest import TUniqueVectorIndexName
-from tests.utils import sleep, when_fetching_vectors_apply_vectors_to_hits
+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(
@@ -808,3 +809,61 @@ def test_get_items_by_id(
     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, ids=[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
@@ -8,6 +8,7 @@
 from momento.errors import MomentoErrorCode
 from momento.requests.vector_index import ALL_METADATA, Field, FilterExpression, Item, SimilarityMetric
 from momento.responses.vector_index import (
+    CountItems,
     CreateIndex,
     DeleteItemBatch,
     GetItemBatch,
@@ -19,7 +20,7 @@
 )
 
 from tests.conftest import TUniqueVectorIndexNameAsync
-from tests.utils import sleep_async, when_fetching_vectors_apply_vectors_to_hits
+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(
@@ -812,3 +813,61 @@ async def test_get_items_by_id(
     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, ids=[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