Small fixes for Chroma memory connector.

microsoft · Jul 24, 2023 · 3727341 · 3727341
1 parent 6a977f4
commit 3727341
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 16 deletions.
diff --git a/dotnet/src/Connectors/Connectors.Memory.Chroma/ChromaMemoryStore.cs b/dotnet/src/Connectors/Connectors.Memory.Chroma/ChromaMemoryStore.cs
@@ -67,15 +67,7 @@ public async Task DeleteCollectionAsync(string collectionName, CancellationToken
     {
         Verify.NotNullOrWhiteSpace(collectionName);
 
-        try
-        {
-            await this._chromaClient.DeleteCollectionAsync(collectionName, cancellationToken).ConfigureAwait(false);
-        }
-        catch (SKException e) when (e.Message.Contains(DeleteNonExistentCollectionErrorMessage))
-        {
-            this._logger.LogError("Cannot delete non-existent collection {0}", collectionName);
-            throw new SKException($"Cannot delete non-existent collection {collectionName}", e);
-        }
+        await this._chromaClient.DeleteCollectionAsync(collectionName, cancellationToken).ConfigureAwait(false);
     }
 
     /// <inheritdoc />
@@ -227,7 +219,6 @@ public async IAsyncEnumerable<string> UpsertBatchAsync(string collectionName, IE
     private const string IncludeEmbeddings = "embeddings";
     private const string IncludeDistances = "distances";
     private const string CollectionDoesNotExistErrorFormat = "Collection {0} does not exist";
-    private const string DeleteNonExistentCollectionErrorMessage = "list index out of range";
 
     private readonly ILogger _logger;
     private readonly IChromaClient _chromaClient;
@@ -251,7 +242,7 @@ await this.GetCollectionAsync(collectionName, cancellationToken).ConfigureAwait(
         {
             return await this._chromaClient.GetCollectionAsync(collectionName, cancellationToken).ConfigureAwait(false);
         }
-        catch (SKException e) when (e.Message.Contains(string.Format(CultureInfo.InvariantCulture, CollectionDoesNotExistErrorFormat, collectionName)))
+        catch (HttpOperationException e) when (e.ResponseContent?.Contains(string.Format(CultureInfo.InvariantCulture, CollectionDoesNotExistErrorFormat, collectionName)) ?? false)
         {
             this._logger.LogError("Collection {0} does not exist", collectionName);
 

diff --git a/dotnet/src/Connectors/Connectors.UnitTests/Memory/Chroma/ChromaMemoryStoreTests.cs b/dotnet/src/Connectors/Connectors.UnitTests/Memory/Chroma/ChromaMemoryStoreTests.cs
@@ -147,7 +147,7 @@ public async Task ItReturnsFalseWhenCollectionDoesNotExistAsync()
 
         this._chromaClientMock
             .Setup(client => client.GetCollectionAsync(collectionName, CancellationToken.None))
-            .Throws(new SKException(collectionDoesNotExistErrorMessage));
+            .Throws(new HttpOperationException() { ResponseContent = collectionDoesNotExistErrorMessage });
 
         var store = new ChromaMemoryStore(this._chromaClientMock.Object);
 

diff --git a/dotnet/src/IntegrationTests/Connectors/Memory/Chroma/ChromaMemoryStoreTests.cs b/dotnet/src/IntegrationTests/Connectors/Memory/Chroma/ChromaMemoryStoreTests.cs
@@ -117,13 +117,12 @@ public async Task ItThrowsExceptionOnNonExistentCollectionDeletionAsync()
         var collectionName = this.GetRandomCollectionName();
 
         // Act
-        var exception = await Record.ExceptionAsync(() => this._chromaMemoryStore.DeleteCollectionAsync(collectionName));
+        var exception = await Assert.ThrowsAsync<HttpOperationException>(() => this._chromaMemoryStore.DeleteCollectionAsync(collectionName));
 
         // Assert
-        Assert.IsType<SKException>(exception);
         Assert.Contains(
-            $"Cannot delete non-existent collection {collectionName}",
-            exception.Message,
+            $"Collection {collectionName} does not exist.",
+            exception.ResponseContent,
             StringComparison.InvariantCulture);
     }