No commit message

momentohq · Jan 23, 2024 · 3efc17f · 3efc17f
2 parents 0d62d55 + da3a044
commit 3efc17f
Show file tree

Hide file tree

Showing 10 changed files with 219 additions and 16 deletions.
diff --git a/README.md b/README.md
@@ -1,7 +1,7 @@
 <head>
   <meta name="Momento .NET Client Library Documentation" content=".NET 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/src/Momento.Sdk/IPreviewVectorIndexClient.cs b/src/Momento.Sdk/IPreviewVectorIndexClient.cs
@@ -105,6 +105,17 @@ public interface IPreviewVectorIndexClient : IDisposable
     ///</returns>
     public Task<DeleteIndexResponse> DeleteIndexAsync(string indexName);
 
+    /// <summary>
+    /// Gets the number of items in a vector index.
+    /// </summary>
+    /// <remarks>
+    /// In the event the index does not exist, the response will be an error.
+    /// A count of zero is reserved for an index that exists but has no items.
+    /// </remarks>
+    /// <param name="indexName">The name of the vector index to get the item count from.</param>
+    /// <returns>Task representing the result of the count items operation.</returns>
+    public Task<CountItemsResponse> CountItemsAsync(string indexName);
+
     /// <summary>
     /// 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.

diff --git a/src/Momento.Sdk/Internal/VectorIndexDataClient.cs b/src/Momento.Sdk/Internal/VectorIndexDataClient.cs
@@ -28,6 +28,30 @@ public VectorIndexDataClient(IVectorIndexConfiguration config, string authToken,
         _exceptionMapper = new CacheExceptionMapper(config.LoggerFactory);
     }
 
+    const string REQUEST_COUNT_ITEMS = "COUNT_ITEMS";
+    public async Task<CountItemsResponse> CountItemsAsync(string indexName)
+    {
+        try
+        {
+            _logger.LogTraceVectorIndexRequest(REQUEST_COUNT_ITEMS, indexName);
+            CheckValidIndexName(indexName);
+            var request = new _CountItemsRequest() { IndexName = indexName, All = new _CountItemsRequest.Types.All() };
+
+            var response =
+                await grpcManager.Client.CountItemsAsync(request, new CallOptions(deadline: CalculateDeadline()));
+            // To maintain CLS compliance we use a long here instead of a ulong.
+            // The max value of a long is still over 9 quintillion so we should be good for a while.
+            var itemCount = checked((long)response.ItemCount);
+            return _logger.LogTraceVectorIndexRequestSuccess(REQUEST_COUNT_ITEMS, indexName,
+                new CountItemsResponse.Success(itemCount));
+        }
+        catch (Exception e)
+        {
+            return _logger.LogTraceVectorIndexRequestError(REQUEST_COUNT_ITEMS, indexName,
+                new CountItemsResponse.Error(_exceptionMapper.Convert(e)));
+        }
+    }
+
     const string REQUEST_UPSERT_ITEM_BATCH = "UPSERT_ITEM_BATCH";
     public async Task<UpsertItemBatchResponse> UpsertItemBatchAsync(string indexName,
         IEnumerable<Item> items)

diff --git a/src/Momento.Sdk/Internal/VectorIndexDataGrpcManager.cs b/src/Momento.Sdk/Internal/VectorIndexDataGrpcManager.cs
@@ -19,6 +19,7 @@ namespace Momento.Sdk.Internal;
 
 public interface IVectorIndexDataClient
 {
+    public Task<_CountItemsResponse> CountItemsAsync(_CountItemsRequest request, CallOptions callOptions);
     public Task<_UpsertItemBatchResponse> UpsertItemBatchAsync(_UpsertItemBatchRequest request, CallOptions callOptions);
     public Task<_SearchResponse> SearchAsync(_SearchRequest request, CallOptions callOptions);
     public Task<_SearchAndFetchVectorsResponse> SearchAndFetchVectorsAsync(_SearchAndFetchVectorsRequest request, CallOptions callOptions);
@@ -47,6 +48,12 @@ public VectorIndexDataClientWithMiddleware(VectorIndex.VectorIndexClient generat
         _middlewares = middlewares;
     }
 
+    public async Task<_CountItemsResponse> CountItemsAsync(_CountItemsRequest request, CallOptions callOptions)
+    {
+        var wrapped = await _middlewares.WrapRequest(request, callOptions, (r, o) => _generatedClient.CountItemsAsync(r, o));
+        return await wrapped.ResponseAsync;
+    }
+
     public async Task<_UpsertItemBatchResponse> UpsertItemBatchAsync(_UpsertItemBatchRequest request, CallOptions callOptions)
     {
         var wrapped = await _middlewares.WrapRequest(request, callOptions, (r, o) => _generatedClient.UpsertItemBatchAsync(r, o));

diff --git a/src/Momento.Sdk/Momento.Sdk.csproj b/src/Momento.Sdk/Momento.Sdk.csproj
@@ -57,7 +57,7 @@
 	<ItemGroup>
 	  <PackageReference Include="Grpc.Net.Client" Version="2.49.0" />
 	  <PackageReference Include="Microsoft.Bcl.AsyncInterfaces" Version="7.0.0" />
-	  <PackageReference Include="Momento.Protos" Version="0.97.1" />
+	  <PackageReference Include="Momento.Protos" Version="0.102.1" />
 	  <PackageReference Include="JWT" Version="9.0.3" />
 	  <PackageReference Include="System.Threading.Channels" Version="6.0.0" />
 	  <PackageReference Include="Microsoft.Extensions.Logging" Version="6.0.0" />

diff --git a/src/Momento.Sdk/PreviewVectorIndexClient.cs b/src/Momento.Sdk/PreviewVectorIndexClient.cs
@@ -1,3 +1,4 @@
+using System;
 using System.Collections.Generic;
 using System.Threading.Tasks;
 using Momento.Sdk.Auth;
@@ -54,6 +55,12 @@ public async Task<DeleteIndexResponse> DeleteIndexAsync(string indexName)
         return await controlClient.DeleteIndexAsync(indexName);
     }
 
+    /// <inheritdoc />
+    public async Task<CountItemsResponse> CountItemsAsync(string indexName)
+    {
+        return await dataClient.CountItemsAsync(indexName);
+    }
+
     /// <inheritdoc />
     public async Task<UpsertItemBatchResponse> UpsertItemBatchAsync(string indexName,
         IEnumerable<Item> items)

diff --git a/src/Momento.Sdk/Responses/Vector/CountItemsResponse.cs b/src/Momento.Sdk/Responses/Vector/CountItemsResponse.cs
@@ -0,0 +1,80 @@
+using Momento.Sdk.Exceptions;
+
+namespace Momento.Sdk.Responses.Vector;
+
+/// <summary>
+/// Parent response type for a count items request. The
+/// response object is resolved to a type-safe object of one of
+/// the following subtypes:
+/// <list type="bullet">
+/// <item><description>CountItemsResponse.Success</description></item>
+/// <item><description>CountItemsResponse.Error</description></item>
+/// </list>
+/// Pattern matching can be used to operate on the appropriate subtype.
+/// For example:
+/// <code>
+/// if (response is CountItemsResponse.Success successResponse)
+/// {
+///     return successResponse.ItemCount;
+/// }
+/// else if (response is CountItemsResponse.Error errorResponse)
+/// {
+///     // handle error as appropriate
+/// }
+/// else
+/// {
+///     // handle unexpected response
+/// }
+/// </code>
+/// </summary>
+public abstract class CountItemsResponse
+{
+    /// <include file="../../docs.xml" path='docs/class[@name="Success"]/description/*' />
+    public class Success : CountItemsResponse
+    {
+        /// <summary>
+        /// The number of items in the vector index.
+        /// </summary>
+        public long ItemCount { get; }
+
+        /// <include file="../../docs.xml" path='docs/class[@name="Success"]/description/*' />
+        /// <param name="itemCount">The number of items in the vector index.</param>
+        public Success(long itemCount)
+        {
+            ItemCount = itemCount;
+        }
+
+        /// <inheritdoc />
+        public override string ToString()
+        {
+            return $"{base.ToString()}: {ItemCount}";
+        }
+
+    }
+
+    /// <include file="../../docs.xml" path='docs/class[@name="Error"]/description/*' />
+    public class Error : CountItemsResponse, IError
+    {
+        /// <include file="../../docs.xml" path='docs/class[@name="Error"]/constructor/*' />
+        public Error(SdkException error)
+        {
+            InnerException = error;
+        }
+
+        /// <inheritdoc />
+        public SdkException InnerException { get; }
+
+        /// <inheritdoc />
+        public MomentoErrorCode ErrorCode => InnerException.ErrorCode;
+
+        /// <inheritdoc />
+        public string Message => $"{InnerException.MessageWrapper}: {InnerException.Message}";
+
+        /// <inheritdoc />
+        public override string ToString()
+        {
+            return $"{base.ToString()}: {Message}";
+        }
+
+    }
+}
diff --git a/tests/Integration/Momento.Sdk.Tests/Utils.cs b/tests/Integration/Momento.Sdk.Tests/Utils.cs
@@ -11,7 +11,18 @@ public static class Utils
 
     public static string TestCacheName() => "dotnet-integration-" + NewGuidString();
 
-    public static string TestVectorIndexName() => "dotnet-integration-" + NewGuidString();
+    /// <summary>
+    /// Returns a test vector index name that is unique to this test run.
+    /// </summary>
+    /// <remarks>
+    /// This is useful for debugging leaking test vector indexes.
+    /// </remarks>
+    /// <param name="meaningfulIdentifier">A string that uniquely identifies the test, e.g. "test1".</param>
+    /// <returns>A test vector index name that is unique to this test run.</returns>
+    public static string TestVectorIndexName(string meaningfulIdentifier)
+    {
+        return $"dotnet-integration-{NewGuidString()}-{meaningfulIdentifier}";
+    }
 
     public static string NewGuidString() => Guid.NewGuid().ToString();
 

diff --git a/tests/Integration/Momento.Sdk.Tests/VectorIndexControlTest.cs b/tests/Integration/Momento.Sdk.Tests/VectorIndexControlTest.cs
@@ -21,9 +21,9 @@ public static IEnumerable<object[]> CreateAndListIndexTestData
         {
             return new List<object[]>
             {
-                new object[] { new IndexInfo(Utils.TestVectorIndexName(), 3, SimilarityMetric.CosineSimilarity) },
-                new object[] { new IndexInfo(Utils.TestVectorIndexName(), 3, SimilarityMetric.InnerProduct) },
-                new object[] { new IndexInfo(Utils.TestVectorIndexName(), 3, SimilarityMetric.EuclideanSimilarity) }
+                new object[] { new IndexInfo(Utils.TestVectorIndexName("control-create-and-list-1"), 3, SimilarityMetric.CosineSimilarity) },
+                new object[] { new IndexInfo(Utils.TestVectorIndexName("control-create-and-list-2"), 3, SimilarityMetric.InnerProduct) },
+                new object[] { new IndexInfo(Utils.TestVectorIndexName("control-create-and-list-3"), 3, SimilarityMetric.EuclideanSimilarity) }
             };
         }
     }
@@ -44,7 +44,7 @@ public async Task CreateListDelete_HappyPath(IndexInfo indexInfo)
     [Fact]
     public async Task CreateIndexAsync_AlreadyExistsError()
     {
-        var indexName = Utils.TestVectorIndexName();
+        var indexName = Utils.TestVectorIndexName("control-create-index-already-exists");
         const int numDimensions = 3;
         using (Utils.WithVectorIndex(vectorIndexClient, indexName, numDimensions))
         {
@@ -79,7 +79,7 @@ public async Task CreateIndexAsync_InvalidNumDimensions()
     [Fact]
     public async Task DeleteIndexAsync_DoesntExistError()
     {
-        var indexName = Utils.TestVectorIndexName();
+        var indexName = Utils.TestVectorIndexName("control-delete-index-doesnt-exist");
         var deleteResponse = await vectorIndexClient.DeleteIndexAsync(indexName);
         Assert.True(deleteResponse is DeleteIndexResponse.Error, $"Unexpected response: {deleteResponse}");
         var deleteErr = (DeleteIndexResponse.Error)deleteResponse;

diff --git a/tests/Integration/Momento.Sdk.Tests/VectorIndexDataTest.cs b/tests/Integration/Momento.Sdk.Tests/VectorIndexDataTest.cs
@@ -109,7 +109,7 @@ public static IEnumerable<object[]> UpsertAndSearchTestData
     public async Task UpsertAndSearch_InnerProduct<T>(SearchDelegate<T> searchDelegate,
         AssertOnSearchResponse<T> assertOnSearchResponse)
     {
-        var indexName = Utils.TestVectorIndexName();
+        var indexName = Utils.TestVectorIndexName("data-upsert-and-search-inner-product");
         using (Utils.WithVectorIndex(vectorIndexClient, indexName, 2, SimilarityMetric.InnerProduct))
         {
             var items = new List<Item>
@@ -137,7 +137,7 @@ public async Task UpsertAndSearch_InnerProduct<T>(SearchDelegate<T> searchDelega
     public async Task UpsertAndSearch_CosineSimilarity<T>(SearchDelegate<T> searchDelegate,
         AssertOnSearchResponse<T> assertOnSearchResponse)
     {
-        var indexName = Utils.TestVectorIndexName();
+        var indexName = Utils.TestVectorIndexName("data-upsert-and-search-cosine-similarity");
         using (Utils.WithVectorIndex(vectorIndexClient, indexName, 2))
         {
             var items = new List<Item>
@@ -168,7 +168,7 @@ public async Task UpsertAndSearch_CosineSimilarity<T>(SearchDelegate<T> searchDe
     public async Task UpsertAndSearch_EuclideanSimilarity<T>(SearchDelegate<T> searchDelegate,
         AssertOnSearchResponse<T> assertOnSearchResponse)
     {
-        var indexName = Utils.TestVectorIndexName();
+        var indexName = Utils.TestVectorIndexName("data-upsert-and-search-euclidean-similarity");
         using (Utils.WithVectorIndex(vectorIndexClient, indexName, 2, SimilarityMetric.EuclideanSimilarity))
         {
             var items = new List<Item>
@@ -199,7 +199,7 @@ public async Task UpsertAndSearch_EuclideanSimilarity<T>(SearchDelegate<T> searc
     public async Task UpsertAndSearch_TopKLimit<T>(SearchDelegate<T> searchDelegate,
         AssertOnSearchResponse<T> assertOnSearchResponse)
     {
-        var indexName = Utils.TestVectorIndexName();
+        var indexName = Utils.TestVectorIndexName("data-upsert-and-search-top-k-limit");
         using (Utils.WithVectorIndex(vectorIndexClient, indexName, 2, SimilarityMetric.InnerProduct))
         {
             var items = new List<Item>
@@ -233,7 +233,7 @@ public async Task UpsertAndSearch_TopKLimit<T>(SearchDelegate<T> searchDelegate,
     public async Task UpsertAndSearch_WithMetadata<T>(SearchDelegate<T> searchDelegate,
         AssertOnSearchResponse<T> assertOnSearchResponse)
     {
-        var indexName = Utils.TestVectorIndexName();
+        var indexName = Utils.TestVectorIndexName("data-upsert-and-search-with-metadata");
         using (Utils.WithVectorIndex(vectorIndexClient, indexName, 2, SimilarityMetric.InnerProduct))
         {
             var items = new List<Item>
@@ -305,7 +305,7 @@ public async Task UpsertAndSearch_WithMetadata<T>(SearchDelegate<T> searchDelega
     public async Task UpsertAndSearch_WithDiverseMetadata<T>(SearchDelegate<T> searchDelegate,
         AssertOnSearchResponse<T> assertOnSearchResponse)
     {
-        var indexName = Utils.TestVectorIndexName();
+        var indexName = Utils.TestVectorIndexName("data-upsert-and-search-with-diverse-metadata");
         using (Utils.WithVectorIndex(vectorIndexClient, indexName, 2, SimilarityMetric.InnerProduct))
         {
             var metadata = new Dictionary<string, MetadataValue>
@@ -373,7 +373,7 @@ public async Task UpsertAndSearch_WithDiverseMetadata<T>(SearchDelegate<T> searc
     public async Task Search_PruneBasedOnThreshold<T>(SimilarityMetric similarityMetric, List<float> scores,
         List<float> thresholds, SearchDelegate<T> searchDelegate, AssertOnSearchResponse<T> assertOnSearchResponse)
     {
-        var indexName = Utils.TestVectorIndexName();
+        var indexName = Utils.TestVectorIndexName("data-search-prune-based-on-threshold");
         using (Utils.WithVectorIndex(vectorIndexClient, indexName, 2, similarityMetric))
         {
             var items = new List<Item>
@@ -500,7 +500,7 @@ public static IEnumerable<object[]> GetItemAndGetItemMetadataTestData
     [MemberData(nameof(GetItemAndGetItemMetadataTestData))]
     public async Task GetItemAndGetItemMetadata_HappyPath<T>(GetItemDelegate<T> getItemDelegate, AssertOnGetItemResponse<T> assertOnGetItemResponse, IEnumerable<string> ids, Object expected)
     {
-        var indexName = Utils.TestVectorIndexName();
+        var indexName = Utils.TestVectorIndexName("data-get-item-and-get-item-metadata-happy-path");
         using (Utils.WithVectorIndex(vectorIndexClient, indexName, 2, SimilarityMetric.InnerProduct))
         {
             var items = new List<Item>
@@ -520,4 +520,67 @@ public async Task GetItemAndGetItemMetadata_HappyPath<T>(GetItemDelegate<T> getI
             assertOnGetItemResponse.Invoke(getResponse, expected);
         }
     }
+
+    [Fact]
+    public async Task CountItemsAsync_OnMissingIndex_ReturnsError()
+    {
+        var indexName = Utils.NewGuidString();
+        var response = await vectorIndexClient.CountItemsAsync(indexName);
+        Assert.True(response is CountItemsResponse.Error, $"Unexpected response: {response}");
+        var error = (CountItemsResponse.Error)response;
+        Assert.Equal(MomentoErrorCode.NOT_FOUND_ERROR, error.InnerException.ErrorCode);
+    }
+
+    [Fact]
+    public async Task CountItemsAsync_OnEmptyIndex_ReturnsZero()
+    {
+        var indexName = Utils.TestVectorIndexName("data-count-items-on-empty-index");
+        using (Utils.WithVectorIndex(vectorIndexClient, indexName, 2, SimilarityMetric.InnerProduct))
+        {
+            var response = await vectorIndexClient.CountItemsAsync(indexName);
+            Assert.True(response is CountItemsResponse.Success, $"Unexpected response: {response}");
+            var successResponse = (CountItemsResponse.Success)response;
+            Assert.Equal(0, successResponse.ItemCount);
+        }
+    }
+
+    [Fact]
+    public async Task CountItemsAsync_HasItems_CountsCorrectly()
+    {
+        var indexName = Utils.TestVectorIndexName("data-count-items-has-items-counts-correctly");
+        using (Utils.WithVectorIndex(vectorIndexClient, indexName, 2, SimilarityMetric.InnerProduct))
+        {
+            var items = new List<Item>
+            {
+                new("test_item_1", new List<float> { 1.0f, 2.0f }),
+                new("test_item_2", new List<float> { 3.0f, 4.0f }),
+                new("test_item_3", new List<float> { 5.0f, 6.0f }),
+                new("test_item_4", new List<float> { 7.0f, 8.0f }),
+                new("test_item_5", new List<float> { 9.0f, 10.0f }),
+            };
+
+            var upsertResponse = await vectorIndexClient.UpsertItemBatchAsync(indexName, items);
+            Assert.True(upsertResponse is UpsertItemBatchResponse.Success,
+                $"Unexpected response: {upsertResponse}");
+
+            await Task.Delay(2_000);
+
+            var response = await vectorIndexClient.CountItemsAsync(indexName);
+            Assert.True(response is CountItemsResponse.Success, $"Unexpected response: {response}");
+            var successResponse = (CountItemsResponse.Success)response;
+            Assert.Equal(5, successResponse.ItemCount);
+
+            // Delete two items
+            var deleteResponse = await vectorIndexClient.DeleteItemBatchAsync(indexName,
+                new List<string> { "test_item_1", "test_item_2" });
+            Assert.True(deleteResponse is DeleteItemBatchResponse.Success, $"Unexpected response: {deleteResponse}");
+
+            await Task.Delay(2_000);
+
+            response = await vectorIndexClient.CountItemsAsync(indexName);
+            Assert.True(response is CountItemsResponse.Success, $"Unexpected response: {response}");
+            successResponse = (CountItemsResponse.Success)response;
+            Assert.Equal(3, successResponse.ItemCount);
+        }
+    }
 }