Skip to content

Commit 4a57c6d

Browse files
rojiroji
committed
Change UpsertAsync return type to IReadOnlyList<TKey>
Closes #10692
1 parent 7172c56 commit 4a57c6d

File tree

44 files changed

+136
-193
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

44 files changed

+136
-193
lines changed

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

+2-7
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
// Copyright (c) Microsoft. All rights reserved.
22

33
using System.Reflection;
4-
using System.Runtime.CompilerServices;
54
using Microsoft.Extensions.VectorData;
65
using Microsoft.SemanticKernel.Embeddings;
76

@@ -108,15 +107,11 @@ public async Task<TKey> UpsertAsync(TRecord record, CancellationToken cancellati
108107
}
109108

110109
/// <inheritdoc />
111-
public async IAsyncEnumerable<TKey> UpsertAsync(IEnumerable<TRecord> records, [EnumeratorCancellation] CancellationToken cancellationToken = default)
110+
public async Task<IReadOnlyList<TKey>> UpsertAsync(IEnumerable<TRecord> records, CancellationToken cancellationToken = default)
112111
{
113112
var recordWithEmbeddingsTasks = records.Select(r => this.AddEmbeddingsAsync(r, cancellationToken));
114113
var recordWithEmbeddings = await Task.WhenAll(recordWithEmbeddingsTasks).ConfigureAwait(false);
115-
var upsertResults = this._decoratedVectorStoreRecordCollection.UpsertAsync(recordWithEmbeddings, cancellationToken);
116-
await foreach (var upsertResult in upsertResults.ConfigureAwait(false))
117-
{
118-
yield return upsertResult;
119-
}
114+
return await this._decoratedVectorStoreRecordCollection.UpsertAsync(recordWithEmbeddings, cancellationToken);
120115
}
121116

122117
/// <inheritdoc />

‎dotnet/samples/Concepts/Memory/VectorStore_EmbeddingGeneration.cs

+1-1
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ public async Task UseEmbeddingGenerationViaDecoratorAsync()
4545
await collection.CreateCollectionIfNotExistsAsync();
4646

4747
// Create and upsert glossary entries into the collection.
48-
await collection.UpsertAsync(CreateGlossaryEntries()).ToListAsync();
48+
await collection.UpsertAsync(CreateGlossaryEntries());
4949

5050
// Search the collection using a vectorizable text search.
5151
var search = collection as IVectorizableTextSearch<Glossary>;

‎dotnet/samples/Concepts/Optimization/FrugalGPTWithFilters.cs

+1-1
Original file line numberDiff line numberDiff line change
@@ -210,7 +210,7 @@ public async Task OnPromptRenderAsync(PromptRenderContext context, Func<PromptRe
210210
var collection = vectorStore.GetCollection<string, ExampleRecord>(CollectionName);
211211
await collection.CreateCollectionIfNotExistsAsync(context.CancellationToken);
212212

213-
await collection.UpsertAsync(exampleRecords, cancellationToken: context.CancellationToken).ToListAsync(context.CancellationToken);
213+
await collection.UpsertAsync(exampleRecords, cancellationToken: context.CancellationToken);
214214

215215
// Generate embedding for original request.
216216
var requestEmbedding = await textEmbeddingGenerationService.GenerateEmbeddingAsync(request, cancellationToken: context.CancellationToken);

‎dotnet/samples/Concepts/Optimization/PluginSelectionWithFilters.cs

+1-1
Original file line numberDiff line numberDiff line change
@@ -341,7 +341,7 @@ public async Task SaveAsync(string collectionName, KernelPluginCollection plugin
341341
var collection = vectorStore.GetCollection<string, FunctionRecord>(collectionName);
342342
await collection.CreateCollectionIfNotExistsAsync(cancellationToken);
343343

344-
await collection.UpsertAsync(functionRecords, cancellationToken: cancellationToken).ToListAsync(cancellationToken);
344+
await collection.UpsertAsync(functionRecords, cancellationToken: cancellationToken);
345345
}
346346

347347
private static List<(KernelFunction Function, string TextToVectorize)> GetFunctionsData(KernelPluginCollection plugins)

‎dotnet/samples/Demos/VectorStoreRAG/DataLoader.cs

+2-2
Original file line numberDiff line numberDiff line change
@@ -66,8 +66,8 @@ public async Task LoadPdf(string pdfPath, int batchSize, int betweenBatchDelayIn
6666

6767
// Upsert the records into the vector store.
6868
var records = await Task.WhenAll(recordTasks).ConfigureAwait(false);
69-
var upsertedKeys = vectorStoreRecordCollection.UpsertAsync(records, cancellationToken: cancellationToken);
70-
await foreach (var key in upsertedKeys.ConfigureAwait(false))
69+
var upsertedKeys = await vectorStoreRecordCollection.UpsertAsync(records, cancellationToken: cancellationToken).ConfigureAwait(false);
70+
foreach (var key in upsertedKeys)
7171
{
7272
Console.WriteLine($"Upserted record '{key}' into VectorDB");
7373
}

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -457,7 +457,7 @@ public async Task CanUpsertManyRecordsAsync(bool useDefinition)
457457
// Act.
458458
var actual = await sut.UpsertAsync(
459459
[model1, model2],
460-
cancellationToken: this._testCancellationToken).ToListAsync();
460+
cancellationToken: this._testCancellationToken);
461461

462462
// Assert.
463463
Assert.NotNull(actual);

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -418,7 +418,7 @@ public async Task UpsertBatchReturnsRecordKeysAsync()
418418
"collection");
419419

420420
// Act
421-
var results = await sut.UpsertAsync([hotel1, hotel2, hotel3]).ToListAsync();
421+
var results = await sut.UpsertAsync([hotel1, hotel2, hotel3]);
422422

423423
// Assert
424424
Assert.NotNull(results);

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -448,7 +448,7 @@ public async Task UpsertBatchReturnsRecordKeysAsync()
448448
"collection");
449449

450450
// Act
451-
var results = await sut.UpsertAsync([hotel1, hotel2, hotel3]).ToListAsync();
451+
var results = await sut.UpsertAsync([hotel1, hotel2, hotel3]);
452452

453453
// Assert
454454
Assert.NotNull(results);

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -257,7 +257,7 @@ public async Task CanUpsertManyRecordsAsync<TKey>(bool useDefinition, TKey testK
257257
// Act
258258
var actual = await sut.UpsertAsync(
259259
[record1, record2],
260-
cancellationToken: this._testCancellationToken).ToListAsync();
260+
cancellationToken: this._testCancellationToken);
261261

262262
// Assert
263263
Assert.NotNull(actual);

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

+3-7
Original file line numberDiff line numberDiff line change
@@ -279,12 +279,12 @@ public virtual async Task<TKey> UpsertAsync(TRecord record, CancellationToken ca
279279
}
280280

281281
/// <inheritdoc />
282-
public virtual async IAsyncEnumerable<TKey> UpsertAsync(IEnumerable<TRecord> records, [EnumeratorCancellation] CancellationToken cancellationToken = default)
282+
public virtual async Task<IReadOnlyList<TKey>> UpsertAsync(IEnumerable<TRecord> records, CancellationToken cancellationToken = default)
283283
{
284284
Verify.NotNull(records);
285285
if (!records.Any())
286286
{
287-
yield break;
287+
return [];
288288
}
289289

290290
// Create Options
@@ -293,11 +293,7 @@ public virtual async IAsyncEnumerable<TKey> UpsertAsync(IEnumerable<TRecord> rec
293293
// Upsert records
294294
var results = await this.MapToStorageModelAndUploadDocumentAsync(records, innerOptions, cancellationToken).ConfigureAwait(false);
295295

296-
// Get results
297-
foreach (var key in results.Value.Results.Select(x => x.Key))
298-
{
299-
yield return (TKey)(object)key;
300-
}
296+
return results.Value.Results.Select(x => (TKey)(object)x.Key).ToList();
301297
}
302298

303299
/// <inheritdoc />

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

+2-9
Original file line numberDiff line numberDiff line change
@@ -242,20 +242,13 @@ await this._mongoCollection
242242
}
243243

244244
/// <inheritdoc />
245-
public virtual async IAsyncEnumerable<TKey> UpsertAsync(IEnumerable<TRecord> records, [EnumeratorCancellation] CancellationToken cancellationToken = default)
245+
public virtual async Task<IReadOnlyList<TKey>> UpsertAsync(IEnumerable<TRecord> records, CancellationToken cancellationToken = default)
246246
{
247247
Verify.NotNull(records);
248248

249249
var tasks = records.Select(record => this.UpsertAsync(record, cancellationToken));
250250
var results = await Task.WhenAll(tasks).ConfigureAwait(false);
251-
252-
foreach (var result in results)
253-
{
254-
if (result is not null)
255-
{
256-
yield return (TKey)(object)result;
257-
}
258-
}
251+
return results.Where(r => r is not null).ToList();
259252
}
260253

261254
/// <inheritdoc />

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

+2-10
Original file line numberDiff line numberDiff line change
@@ -276,22 +276,14 @@ await this.RunOperationAsync(OperationName, () =>
276276
}
277277

278278
/// <inheritdoc />
279-
public virtual async IAsyncEnumerable<TKey> UpsertAsync(IEnumerable<TRecord> records, [EnumeratorCancellation] CancellationToken cancellationToken = default)
279+
public virtual async Task<IReadOnlyList<TKey>> UpsertAsync(IEnumerable<TRecord> records, CancellationToken cancellationToken = default)
280280
{
281281
Verify.NotNull(records);
282282

283283
// TODO: Do proper bulk upsert rather than parallel single inserts, #11350
284284
var tasks = records.Select(record => this.UpsertAsync(record, cancellationToken));
285-
286285
var keys = await Task.WhenAll(tasks).ConfigureAwait(false);
287-
288-
foreach (var key in keys)
289-
{
290-
if (key is not null)
291-
{
292-
yield return key;
293-
}
294-
}
286+
return keys.Where(k => k is not null).ToList();
295287
}
296288

297289
/// <inheritdoc />

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

+14-16
Original file line numberDiff line numberDiff line change
@@ -222,6 +222,17 @@ public Task DeleteAsync(IEnumerable<TKey> keys, CancellationToken cancellationTo
222222

223223
/// <inheritdoc />
224224
public Task<TKey> UpsertAsync(TRecord record, CancellationToken cancellationToken = default)
225+
=> Task.FromResult(this.Upsert(record));
226+
227+
/// <inheritdoc />
228+
public Task<IReadOnlyList<TKey>> UpsertAsync(IEnumerable<TRecord> records, CancellationToken cancellationToken = default)
229+
{
230+
Verify.NotNull(records);
231+
232+
return Task.FromResult<IReadOnlyList<TKey>>(records.Select(this.Upsert).ToList());
233+
}
234+
235+
private TKey Upsert(TRecord record)
225236
{
226237
Verify.NotNull(record);
227238

@@ -230,24 +241,11 @@ public Task<TKey> UpsertAsync(TRecord record, CancellationToken cancellationToke
230241
var key = (TKey)this._keyResolver(record)!;
231242
collectionDictionary.AddOrUpdate(key!, record, (key, currentValue) => record);
232243

233-
return Task.FromResult(key!);
234-
}
235-
236-
/// <inheritdoc />
237-
public async IAsyncEnumerable<TKey> UpsertAsync(IEnumerable<TRecord> records, [EnumeratorCancellation] CancellationToken cancellationToken = default)
238-
{
239-
Verify.NotNull(records);
240-
241-
foreach (var record in records)
242-
{
243-
yield return await this.UpsertAsync(record, cancellationToken).ConfigureAwait(false);
244-
}
244+
return key!;
245245
}
246246

247247
/// <inheritdoc />
248-
#pragma warning disable CS1998 // Async method lacks 'await' operators and will run synchronously - Need to satisfy the interface which returns IAsyncEnumerable
249-
public async Task<VectorSearchResults<TRecord>> VectorizedSearchAsync<TVector>(TVector vector, int top, VectorSearchOptions<TRecord>? options = null, CancellationToken cancellationToken = default)
250-
#pragma warning restore CS1998
248+
public Task<VectorSearchResults<TRecord>> VectorizedSearchAsync<TVector>(TVector vector, int top, VectorSearchOptions<TRecord>? options = null, CancellationToken cancellationToken = default)
251249
{
252250
Verify.NotNull(vector);
253251
Verify.NotLessThan(top, 1);
@@ -305,7 +303,7 @@ public async Task<VectorSearchResults<TRecord>> VectorizedSearchAsync<TVector>(T
305303

306304
// Build the response.
307305
var vectorSearchResultList = resultsPage.Select(x => new VectorSearchResult<TRecord>((TRecord)x.record, x.score)).ToAsyncEnumerable();
308-
return new VectorSearchResults<TRecord>(vectorSearchResultList) { TotalCount = count };
306+
return Task.FromResult(new VectorSearchResults<TRecord>(vectorSearchResultList) { TotalCount = count });
309307
}
310308

311309
/// <summary>

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

+2-9
Original file line numberDiff line numberDiff line change
@@ -248,20 +248,13 @@ await this._mongoCollection
248248
}
249249

250250
/// <inheritdoc />
251-
public virtual async IAsyncEnumerable<TKey> UpsertAsync(IEnumerable<TRecord> records, [EnumeratorCancellation] CancellationToken cancellationToken = default)
251+
public virtual async Task<IReadOnlyList<TKey>> UpsertAsync(IEnumerable<TRecord> records, CancellationToken cancellationToken = default)
252252
{
253253
Verify.NotNull(records);
254254

255255
var tasks = records.Select(record => this.UpsertAsync(record, cancellationToken));
256256
var results = await Task.WhenAll(tasks).ConfigureAwait(false);
257-
258-
foreach (var result in results)
259-
{
260-
if (result is not null)
261-
{
262-
yield return (TKey)(object)result;
263-
}
264-
}
257+
return results.Where(r => r is not null).ToList();
265258
}
266259

267260
/// <inheritdoc />

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

+3-6
Original file line numberDiff line numberDiff line change
@@ -298,7 +298,7 @@ await this.RunIndexOperationAsync(
298298
}
299299

300300
/// <inheritdoc />
301-
public virtual async IAsyncEnumerable<TKey> UpsertAsync(IEnumerable<TRecord> records, [EnumeratorCancellation] CancellationToken cancellationToken = default)
301+
public virtual async Task<IReadOnlyList<TKey>> UpsertAsync(IEnumerable<TRecord> records, CancellationToken cancellationToken = default)
302302
{
303303
Verify.NotNull(records);
304304

@@ -310,7 +310,7 @@ public virtual async IAsyncEnumerable<TKey> UpsertAsync(IEnumerable<TRecord> rec
310310

311311
if (vectors.Count == 0)
312312
{
313-
yield break;
313+
return [];
314314
}
315315

316316
Sdk.UpsertRequest request = new()
@@ -323,10 +323,7 @@ await this.RunIndexOperationAsync(
323323
"UpsertBatch",
324324
indexClient => indexClient.UpsertAsync(request, cancellationToken: cancellationToken)).ConfigureAwait(false);
325325

326-
foreach (var vector in vectors)
327-
{
328-
yield return (TKey)(object)vector.Id;
329-
}
326+
return vectors.Select(x => (TKey)(object)x.Id).ToList();
330327
}
331328

332329
/// <inheritdoc />

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

+3-4
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@
33
using System;
44
using System.Collections.Generic;
55
using System.Linq;
6-
using System.Runtime.CompilerServices;
76
using System.Threading;
87
using System.Threading.Tasks;
98
using Microsoft.Extensions.VectorData;
@@ -143,7 +142,7 @@ public virtual Task<TKey> UpsertAsync(TRecord record, CancellationToken cancella
143142
}
144143

145144
/// <inheritdoc/>
146-
public virtual async IAsyncEnumerable<TKey> UpsertAsync(IEnumerable<TRecord> records, [EnumeratorCancellation] CancellationToken cancellationToken = default)
145+
public virtual async Task<IReadOnlyList<TKey>> UpsertAsync(IEnumerable<TRecord> records, CancellationToken cancellationToken = default)
147146
{
148147
Verify.NotNull(records);
149148

@@ -157,7 +156,7 @@ public virtual async IAsyncEnumerable<TKey> UpsertAsync(IEnumerable<TRecord> rec
157156

158157
if (storageModels.Count == 0)
159158
{
160-
yield break;
159+
return [];
161160
}
162161

163162
var keys = storageModels.Select(model => model[this._model.KeyProperty.StorageName]!).ToList();
@@ -166,7 +165,7 @@ await this.RunOperationAsync(OperationName, () =>
166165
this._client.UpsertBatchAsync(this.CollectionName, storageModels, this._model.KeyProperty.StorageName, cancellationToken)
167166
).ConfigureAwait(false);
168167

169-
foreach (var key in keys) { yield return (TKey)key!; }
168+
return keys.Select(key => (TKey)key!).ToList();
170169
}
171170

172171
/// <inheritdoc/>

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

+8-22
Original file line numberDiff line numberDiff line change
@@ -374,7 +374,7 @@ await this.RunOperationAsync(
374374
}
375375

376376
/// <inheritdoc />
377-
public virtual async IAsyncEnumerable<TKey> UpsertAsync(IEnumerable<TRecord> records, [EnumeratorCancellation] CancellationToken cancellationToken = default)
377+
public virtual async Task<IReadOnlyList<TKey>> UpsertAsync(IEnumerable<TRecord> records, CancellationToken cancellationToken = default)
378378
{
379379
Verify.NotNull(records);
380380

@@ -390,28 +390,14 @@ await this.RunOperationAsync(
390390
UpsertName,
391391
() => this._qdrantClient.UpsertAsync(this._collectionName, pointStructs, true, cancellationToken: cancellationToken)).ConfigureAwait(false);
392392

393-
if (pointStructs.Count > 0)
394-
{
395-
switch (pointStructs[0].Id)
393+
return pointStructs.Count == 0
394+
? []
395+
: pointStructs[0].Id switch
396396
{
397-
case { HasNum: true }:
398-
foreach (var pointStruct in pointStructs)
399-
{
400-
yield return (TKey)(object)pointStruct.Id.Num;
401-
}
402-
break;
403-
404-
case { HasUuid: true }:
405-
foreach (var pointStruct in pointStructs)
406-
{
407-
yield return (TKey)(object)Guid.Parse(pointStruct.Id.Uuid);
408-
}
409-
break;
410-
411-
default:
412-
throw new UnreachableException("The Qdrant point ID is neither a number nor a UUID.");
413-
}
414-
}
397+
{ HasNum: true } => pointStructs.Select(pointStruct => (TKey)(object)pointStruct.Id.Num).ToList(),
398+
{ HasUuid: true } => pointStructs.Select(pointStruct => (TKey)(object)Guid.Parse(pointStruct.Id.Uuid)).ToList(),
399+
_ => throw new UnreachableException("The Qdrant point ID is neither a number nor a UUID.")
400+
};
415401
}
416402

417403
/// <inheritdoc />

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

+2-8
Original file line numberDiff line numberDiff line change
@@ -315,20 +315,14 @@ await this.RunOperationAsync(
315315
}
316316

317317
/// <inheritdoc />
318-
public virtual async IAsyncEnumerable<TKey> UpsertAsync(IEnumerable<TRecord> records, [EnumeratorCancellation] CancellationToken cancellationToken = default)
318+
public virtual async Task<IReadOnlyList<TKey>> UpsertAsync(IEnumerable<TRecord> records, CancellationToken cancellationToken = default)
319319
{
320320
Verify.NotNull(records);
321321

322322
// Upsert records in parallel.
323323
var tasks = records.Select(x => this.UpsertAsync(x, cancellationToken));
324324
var results = await Task.WhenAll(tasks).ConfigureAwait(false);
325-
foreach (var result in results)
326-
{
327-
if (result is not null)
328-
{
329-
yield return (TKey)(object)result;
330-
}
331-
}
325+
return results.Where(r => r is not null).ToList();
332326
}
333327

334328
/// <inheritdoc />

0 commit comments

Comments
 (0)