Change UpsertAsync return type to IReadOnlyList<TKey>

Closes #11401

Author: roji

dotnet/samples/Concepts/Memory/VectorStoreEmbeddingGeneration/TextEmbeddingVectorStoreRecordCollection.cs

@@ -108,15 +108,11 @@ public async Task<TKey> UpsertAsync(TRecord record, CancellationToken cancellati
     }
 
     /// <inheritdoc />
-    public async IAsyncEnumerable<TKey> UpsertAsync(IEnumerable<TRecord> records, [EnumeratorCancellation] CancellationToken cancellationToken = default)
+    public async Task<IReadOnlyList<TKey>> UpsertAsync(IEnumerable<TRecord> records, CancellationToken cancellationToken = default)
     {
         var recordWithEmbeddingsTasks = records.Select(r => this.AddEmbeddingsAsync(r, cancellationToken));
         var recordWithEmbeddings = await Task.WhenAll(recordWithEmbeddingsTasks).ConfigureAwait(false);
-        var upsertResults = this._decoratedVectorStoreRecordCollection.UpsertAsync(recordWithEmbeddings, cancellationToken);
-        await foreach (var upsertResult in upsertResults.ConfigureAwait(false))
-        {
-            yield return upsertResult;
-        }
+        return await this._decoratedVectorStoreRecordCollection.UpsertAsync(recordWithEmbeddings, cancellationToken);
     }
 
     /// <inheritdoc />

dotnet/samples/Concepts/Memory/VectorStore_EmbeddingGeneration.cs

@@ -45,7 +45,7 @@ public async Task UseEmbeddingGenerationViaDecoratorAsync()
         await collection.CreateCollectionIfNotExistsAsync();
 
         // Create and upsert glossary entries into the collection.
-        await collection.UpsertAsync(CreateGlossaryEntries()).ToListAsync();
+        await collection.UpsertAsync(CreateGlossaryEntries());
 
         // Search the collection using a vectorizable text search.
         var search = collection as IVectorizableTextSearch<Glossary>;

dotnet/samples/Concepts/Optimization/FrugalGPTWithFilters.cs

@@ -210,7 +210,7 @@ public async Task OnPromptRenderAsync(PromptRenderContext context, Func<PromptRe
                 var collection = vectorStore.GetCollection<string, ExampleRecord>(CollectionName);
                 await collection.CreateCollectionIfNotExistsAsync(context.CancellationToken);
 
-                await collection.UpsertAsync(exampleRecords, cancellationToken: context.CancellationToken).ToListAsync(context.CancellationToken);
+                await collection.UpsertAsync(exampleRecords, cancellationToken: context.CancellationToken);
 
                 // Generate embedding for original request.
                 var requestEmbedding = await textEmbeddingGenerationService.GenerateEmbeddingAsync(request, cancellationToken: context.CancellationToken);

dotnet/samples/Concepts/Optimization/PluginSelectionWithFilters.cs

@@ -341,7 +341,7 @@ public async Task SaveAsync(string collectionName, KernelPluginCollection plugin
             var collection = vectorStore.GetCollection<string, FunctionRecord>(collectionName);
             await collection.CreateCollectionIfNotExistsAsync(cancellationToken);
 
-            await collection.UpsertAsync(functionRecords, cancellationToken: cancellationToken).ToListAsync(cancellationToken);
+            await collection.UpsertAsync(functionRecords, cancellationToken: cancellationToken);
         }
 
         private static List<(KernelFunction Function, string TextToVectorize)> GetFunctionsData(KernelPluginCollection plugins)

dotnet/samples/Demos/VectorStoreRAG/DataLoader.cs

@@ -66,8 +66,8 @@ public async Task LoadPdf(string pdfPath, int batchSize, int betweenBatchDelayIn
 
             // Upsert the records into the vector store.
             var records = await Task.WhenAll(recordTasks).ConfigureAwait(false);
-            var upsertedKeys = vectorStoreRecordCollection.UpsertAsync(records, cancellationToken: cancellationToken);
-            await foreach (var key in upsertedKeys.ConfigureAwait(false))
+            var upsertedKeys = await vectorStoreRecordCollection.UpsertAsync(records, cancellationToken: cancellationToken).ConfigureAwait(false);
+            foreach (var key in upsertedKeys)
             {
                 Console.WriteLine($"Upserted record '{key}' into VectorDB");
             }

dotnet/src/Connectors/Connectors.AzureAISearch.UnitTests/AzureAISearchVectorStoreRecordCollectionTests.cs

@@ -457,7 +457,7 @@ public async Task CanUpsertManyRecordsAsync(bool useDefinition)
         // Act.
         var actual = await sut.UpsertAsync(
             [model1, model2],
-            cancellationToken: this._testCancellationToken).ToListAsync();
+            cancellationToken: this._testCancellationToken);
 
         // Assert.
         Assert.NotNull(actual);

dotnet/src/Connectors/Connectors.AzureCosmosDBMongoDB.UnitTests/AzureCosmosDBMongoDBVectorStoreRecordCollectionTests.cs

@@ -418,7 +418,7 @@ public async Task UpsertBatchReturnsRecordKeysAsync()
             "collection");
 
         // Act
-        var results = await sut.UpsertAsync([hotel1, hotel2, hotel3]).ToListAsync();
+        var results = await sut.UpsertAsync([hotel1, hotel2, hotel3]);
 
         // Assert
         Assert.NotNull(results);

dotnet/src/Connectors/Connectors.AzureCosmosDBNoSQL.UnitTests/AzureCosmosDBNoSQLVectorStoreRecordCollectionTests.cs

@@ -448,7 +448,7 @@ public async Task UpsertBatchReturnsRecordKeysAsync()
             "collection");
 
         // Act
-        var results = await sut.UpsertAsync([hotel1, hotel2, hotel3]).ToListAsync();
+        var results = await sut.UpsertAsync([hotel1, hotel2, hotel3]);
 
         // Assert
         Assert.NotNull(results);

dotnet/src/Connectors/Connectors.InMemory.UnitTests/InMemoryVectorStoreRecordCollectionTests.cs

@@ -257,7 +257,7 @@ public async Task CanUpsertManyRecordsAsync<TKey>(bool useDefinition, TKey testK
         // Act
         var actual = await sut.UpsertAsync(
             [record1, record2],
-            cancellationToken: this._testCancellationToken).ToListAsync();
+            cancellationToken: this._testCancellationToken);
 
         // Assert
         Assert.NotNull(actual);

dotnet/src/Connectors/Connectors.Memory.AzureAISearch/AzureAISearchVectorStoreRecordCollection.cs

@@ -279,12 +279,12 @@ public virtual async Task<TKey> UpsertAsync(TRecord record, CancellationToken ca
     }
 
     /// <inheritdoc />
-    public virtual async IAsyncEnumerable<TKey> UpsertAsync(IEnumerable<TRecord> records, [EnumeratorCancellation] CancellationToken cancellationToken = default)
+    public virtual async Task<IReadOnlyList<TKey>> UpsertAsync(IEnumerable<TRecord> records, CancellationToken cancellationToken = default)
     {
         Verify.NotNull(records);
         if (!records.Any())
         {
-            yield break;
+            return [];
         }
 
         // Create Options
@@ -293,11 +293,7 @@ public virtual async IAsyncEnumerable<TKey> UpsertAsync(IEnumerable<TRecord> rec
         // Upsert records
         var results = await this.MapToStorageModelAndUploadDocumentAsync(records, innerOptions, cancellationToken).ConfigureAwait(false);
 
-        // Get results
-        foreach (var key in results.Value.Results.Select(x => x.Key))
-        {
-            yield return (TKey)(object)key;
-        }
+        return results.Value.Results.Select(x => (TKey)(object)x.Key).ToList();
     }
 
     /// <inheritdoc />

dotnet/src/Connectors/Connectors.Memory.AzureCosmosDBMongoDB/AzureCosmosDBMongoDBVectorStoreRecordCollection.cs

@@ -242,20 +242,13 @@ await this._mongoCollection
     }
 
     /// <inheritdoc />
-    public virtual async IAsyncEnumerable<TKey> UpsertAsync(IEnumerable<TRecord> records, [EnumeratorCancellation] CancellationToken cancellationToken = default)
+    public virtual async Task<IReadOnlyList<TKey>> UpsertAsync(IEnumerable<TRecord> records, CancellationToken cancellationToken = default)
     {
         Verify.NotNull(records);
 
         var tasks = records.Select(record => this.UpsertAsync(record, cancellationToken));
         var results = await Task.WhenAll(tasks).ConfigureAwait(false);
-
-        foreach (var result in results)
-        {
-            if (result is not null)
-            {
-                yield return (TKey)(object)result;
-            }
-        }
+        return results.Where(r => r is not null).ToList();
     }
 
     /// <inheritdoc />

dotnet/src/Connectors/Connectors.Memory.AzureCosmosDBNoSQL/AzureCosmosDBNoSQLVectorStoreRecordCollection.cs

@@ -276,22 +276,14 @@ await this.RunOperationAsync(OperationName, () =>
     }
 
     /// <inheritdoc />
-    public virtual async IAsyncEnumerable<TKey> UpsertAsync(IEnumerable<TRecord> records, [EnumeratorCancellation] CancellationToken cancellationToken = default)
+    public virtual async Task<IReadOnlyList<TKey>> UpsertAsync(IEnumerable<TRecord> records, CancellationToken cancellationToken = default)
     {
         Verify.NotNull(records);
 
         // TODO: Do proper bulk upsert rather than parallel single inserts, #11350
         var tasks = records.Select(record => this.UpsertAsync(record, cancellationToken));
-
         var keys = await Task.WhenAll(tasks).ConfigureAwait(false);
-
-        foreach (var key in keys)
-        {
-            if (key is not null)
-            {
-                yield return key;
-            }
-        }
+        return keys.Where(k => k is not null).ToList();
     }
 
     /// <inheritdoc />

dotnet/src/Connectors/Connectors.Memory.InMemory/InMemoryVectorStoreRecordCollection.cs

@@ -222,6 +222,17 @@ public Task DeleteAsync(IEnumerable<TKey> keys, CancellationToken cancellationTo
 
     /// <inheritdoc />
     public Task<TKey> UpsertAsync(TRecord record, CancellationToken cancellationToken = default)
+        => Task.FromResult(this.Upsert(record));
+
+    /// <inheritdoc />
+    public Task<IReadOnlyList<TKey>> UpsertAsync(IEnumerable<TRecord> records, CancellationToken cancellationToken = default)
+    {
+        Verify.NotNull(records);
+
+        return Task.FromResult<IReadOnlyList<TKey>>(records.Select(this.Upsert).ToList());
+    }
+
+    private TKey Upsert(TRecord record)
     {
         Verify.NotNull(record);
 
@@ -230,24 +241,11 @@ public Task<TKey> UpsertAsync(TRecord record, CancellationToken cancellationToke
         var key = (TKey)this._keyResolver(record)!;
         collectionDictionary.AddOrUpdate(key!, record, (key, currentValue) => record);
 
-        return Task.FromResult(key!);
-    }
-
-    /// <inheritdoc />
-    public async IAsyncEnumerable<TKey> UpsertAsync(IEnumerable<TRecord> records, [EnumeratorCancellation] CancellationToken cancellationToken = default)
-    {
-        Verify.NotNull(records);
-
-        foreach (var record in records)
-        {
-            yield return await this.UpsertAsync(record, cancellationToken).ConfigureAwait(false);
-        }
+        return key!;
     }
 
     /// <inheritdoc />
-#pragma warning disable CS1998 // Async method lacks 'await' operators and will run synchronously - Need to satisfy the interface which returns IAsyncEnumerable
-    public async Task<VectorSearchResults<TRecord>> VectorizedSearchAsync<TVector>(TVector vector, int top, VectorSearchOptions<TRecord>? options = null, CancellationToken cancellationToken = default)
-#pragma warning restore CS1998
+    public Task<VectorSearchResults<TRecord>> VectorizedSearchAsync<TVector>(TVector vector, int top, VectorSearchOptions<TRecord>? options = null, CancellationToken cancellationToken = default)
     {
         Verify.NotNull(vector);
         Verify.NotLessThan(top, 1);
@@ -305,7 +303,7 @@ public async Task<VectorSearchResults<TRecord>> VectorizedSearchAsync<TVector>(T
 
         // Build the response.
         var vectorSearchResultList = resultsPage.Select(x => new VectorSearchResult<TRecord>((TRecord)x.record, x.score)).ToAsyncEnumerable();
-        return new VectorSearchResults<TRecord>(vectorSearchResultList) { TotalCount = count };
+        return Task.FromResult(new VectorSearchResults<TRecord>(vectorSearchResultList) { TotalCount = count });
     }
 
     /// <summary>

dotnet/src/Connectors/Connectors.Memory.MongoDB/MongoDBVectorStoreRecordCollection.cs

@@ -248,20 +248,13 @@ await this._mongoCollection
     }
 
     /// <inheritdoc />
-    public virtual async IAsyncEnumerable<TKey> UpsertAsync(IEnumerable<TRecord> records, [EnumeratorCancellation] CancellationToken cancellationToken = default)
+    public virtual async Task<IReadOnlyList<TKey>> UpsertAsync(IEnumerable<TRecord> records, CancellationToken cancellationToken = default)
     {
         Verify.NotNull(records);
 
         var tasks = records.Select(record => this.UpsertAsync(record, cancellationToken));
         var results = await Task.WhenAll(tasks).ConfigureAwait(false);
-
-        foreach (var result in results)
-        {
-            if (result is not null)
-            {
-                yield return (TKey)(object)result;
-            }
-        }
+        return results.Where(r => r is not null).ToList();
     }
 
     /// <inheritdoc />

dotnet/src/Connectors/Connectors.Memory.Pinecone/PineconeVectorStoreRecordCollection.cs

@@ -298,7 +298,7 @@ await this.RunIndexOperationAsync(
     }
 
     /// <inheritdoc />
-    public virtual async IAsyncEnumerable<TKey> UpsertAsync(IEnumerable<TRecord> records, [EnumeratorCancellation] CancellationToken cancellationToken = default)
+    public virtual async Task<IReadOnlyList<TKey>> UpsertAsync(IEnumerable<TRecord> records, CancellationToken cancellationToken = default)
     {
         Verify.NotNull(records);
 
@@ -310,7 +310,7 @@ public virtual async IAsyncEnumerable<TKey> UpsertAsync(IEnumerable<TRecord> rec
 
         if (vectors.Count == 0)
         {
-            yield break;
+            return [];
         }
 
         Sdk.UpsertRequest request = new()
@@ -323,10 +323,7 @@ await this.RunIndexOperationAsync(
             "UpsertBatch",
             indexClient => indexClient.UpsertAsync(request, cancellationToken: cancellationToken)).ConfigureAwait(false);
 
-        foreach (var vector in vectors)
-        {
-            yield return (TKey)(object)vector.Id;
-        }
+        return vectors.Select(x => (TKey)(object)x.Id).ToList();
     }
 
     /// <inheritdoc />

dotnet/src/Connectors/Connectors.Memory.Postgres/PostgresVectorStoreRecordCollection.cs

@@ -143,7 +143,7 @@ public virtual Task<TKey> UpsertAsync(TRecord record, CancellationToken cancella
     }
 
     /// <inheritdoc/>
-    public virtual async IAsyncEnumerable<TKey> UpsertAsync(IEnumerable<TRecord> records, [EnumeratorCancellation] CancellationToken cancellationToken = default)
+    public virtual async Task<IReadOnlyList<TKey>> UpsertAsync(IEnumerable<TRecord> records, CancellationToken cancellationToken = default)
     {
         Verify.NotNull(records);
 
@@ -157,7 +157,7 @@ public virtual async IAsyncEnumerable<TKey> UpsertAsync(IEnumerable<TRecord> rec
 
         if (storageModels.Count == 0)
         {
-            yield break;
+            return [];
         }
 
         var keys = storageModels.Select(model => model[this._model.KeyProperty.StorageName]!).ToList();
@@ -166,7 +166,7 @@ await this.RunOperationAsync(OperationName, () =>
             this._client.UpsertBatchAsync(this.CollectionName, storageModels, this._model.KeyProperty.StorageName, cancellationToken)
         ).ConfigureAwait(false);
 
-        foreach (var key in keys) { yield return (TKey)key!; }
+        return keys.Select(key => (TKey)key!).ToList();
     }
 
     /// <inheritdoc/>

dotnet/src/Connectors/Connectors.Memory.Qdrant/QdrantVectorStoreRecordCollection.cs

@@ -374,7 +374,7 @@ await this.RunOperationAsync(
     }
 
     /// <inheritdoc />
-    public virtual async IAsyncEnumerable<TKey> UpsertAsync(IEnumerable<TRecord> records, [EnumeratorCancellation] CancellationToken cancellationToken = default)
+    public virtual async Task<IReadOnlyList<TKey>> UpsertAsync(IEnumerable<TRecord> records, CancellationToken cancellationToken = default)
     {
         Verify.NotNull(records);
 
@@ -390,28 +390,14 @@ await this.RunOperationAsync(
             UpsertName,
             () => this._qdrantClient.UpsertAsync(this._collectionName, pointStructs, true, cancellationToken: cancellationToken)).ConfigureAwait(false);
 
-        if (pointStructs.Count > 0)
-        {
-            switch (pointStructs[0].Id)
+        return pointStructs.Count == 0
+            ? []
+            : pointStructs[0].Id switch
             {
-                case { HasNum: true }:
-                    foreach (var pointStruct in pointStructs)
-                    {
-                        yield return (TKey)(object)pointStruct.Id.Num;
-                    }
-                    break;
-
-                case { HasUuid: true }:
-                    foreach (var pointStruct in pointStructs)
-                    {
-                        yield return (TKey)(object)Guid.Parse(pointStruct.Id.Uuid);
-                    }
-                    break;
-
-                default:
-                    throw new UnreachableException("The Qdrant point ID is neither a number nor a UUID.");
-            }
-        }
+                { HasNum: true } => pointStructs.Select(pointStruct => (TKey)(object)pointStruct.Id.Num).ToList(),
+                { HasUuid: true } => pointStructs.Select(pointStruct => (TKey)(object)Guid.Parse(pointStruct.Id.Uuid)).ToList(),
+                _ => throw new UnreachableException("The Qdrant point ID is neither a number nor a UUID.")
+            };
     }
 
     /// <inheritdoc />

dotnet/src/Connectors/Connectors.Memory.Redis/RedisHashSetVectorStoreRecordCollection.cs

@@ -315,20 +315,14 @@ await this.RunOperationAsync(
     }
 
     /// <inheritdoc />
-    public virtual async IAsyncEnumerable<TKey> UpsertAsync(IEnumerable<TRecord> records, [EnumeratorCancellation] CancellationToken cancellationToken = default)
+    public virtual async Task<IReadOnlyList<TKey>> UpsertAsync(IEnumerable<TRecord> records, CancellationToken cancellationToken = default)
     {
         Verify.NotNull(records);
 
         // Upsert records in parallel.
         var tasks = records.Select(x => this.UpsertAsync(x, cancellationToken));
         var results = await Task.WhenAll(tasks).ConfigureAwait(false);
-        foreach (var result in results)
-        {
-            if (result is not null)
-            {
-                yield return (TKey)(object)result;
-            }
-        }
+        return results.Where(r => r is not null).ToList();
     }
 
     /// <inheritdoc />

dotnet/src/Connectors/Connectors.Memory.Redis/RedisJsonVectorStoreRecordCollection.cs

@@ -369,7 +369,7 @@ await this.RunOperationAsync(
     }
 
     /// <inheritdoc />
-    public virtual async IAsyncEnumerable<TKey> UpsertAsync(IEnumerable<TRecord> records, [EnumeratorCancellation] CancellationToken cancellationToken = default)
+    public virtual async Task<IReadOnlyList<TKey>> UpsertAsync(IEnumerable<TRecord> records, CancellationToken cancellationToken = default)
     {
         Verify.NotNull(records);
 
@@ -400,11 +400,7 @@ await this.RunOperationAsync(
                 .JSON()
                 .MSetAsync(keyPathValues)).ConfigureAwait(false);
 
-        // Return keys of upserted records.
-        foreach (var record in redisRecords)
-        {
-            yield return (TKey)(object)record.originalKey;
-        }
+        return redisRecords.Select(x => (TKey)(object)x.originalKey).ToList();
     }
 
     /// <inheritdoc />