Add batching upsert support to Qdrant Memory

extensions/Qdrant/Qdrant/QdrantMemory.cs

@@ -22,7 +22,7 @@ namespace Microsoft.KernelMemory.MemoryDb.Qdrant;
 /// * allow using more Qdrant specific filtering logic
 /// </summary>
 [Experimental("KMEXP03")]
-public sealed class QdrantMemory : IMemoryDb
+public sealed class QdrantMemory : IMemoryDb, IMemoryDbBatchUpsert
 {
  private readonly ITextEmbeddingGenerator _embeddingGenerator;
  private readonly QdrantClient<DefaultQdrantPayload> _qdrantClient;
@@ -91,40 +91,60 @@ public async Task<string> UpsertAsync(
  string index,
  MemoryRecord record,
  CancellationToken cancellationToken = default)
+ {
+ var result = this.BatchUpsertAsync(index, [record], cancellationToken);
+ var id = await result.SingleAsync(cancellationToken).ConfigureAwait(false);
+ return id;
+ }
+
+ /// <inheritdoc />
+ public async IAsyncEnumerable<string> BatchUpsertAsync(string index, IEnumerable<MemoryRecord> records, [EnumeratorCancellation] CancellationToken cancellationToken = default)
  {
  index = NormalizeIndexName(index);
 
- QdrantPoint<DefaultQdrantPayload> qdrantPoint;
+ // Call ToList to avoid multiple enumerations (CA1851: Possible multiple enumerations of 'IEnumerable' collection. Consider using an implementation that avoids multiple enumerations).
+ var localRecords = records.ToList();
 
- if (string.IsNullOrEmpty(record.Id))
+ var qdrantPoints = new List<QdrantPoint<DefaultQdrantPayload>>();
+ foreach (var record in localRecords)
  {
- record.Id = Guid.NewGuid().ToString("N");
- qdrantPoint = QdrantPoint<DefaultQdrantPayload>.FromMemoryRecord(record);
- qdrantPoint.Id = Guid.NewGuid();
+ QdrantPoint<DefaultQdrantPayload> qdrantPoint;
 
- this._log.LogTrace("Generate new Qdrant point ID {0} and record ID {1}", qdrantPoint.Id, record.Id);
- }
- else
- {
- qdrantPoint = QdrantPoint<DefaultQdrantPayload>.FromMemoryRecord(record);
- QdrantPoint<DefaultQdrantPayload>? existingPoint = await this._qdrantClient
- .GetVectorByPayloadIdAsync(index, record.Id, cancellationToken: cancellationToken)
- .ConfigureAwait(false);
- if (existingPoint == null)
+ if (string.IsNullOrEmpty(record.Id))
  {
+ record.Id = Guid.NewGuid().ToString("N");
+ qdrantPoint = QdrantPoint<DefaultQdrantPayload>.FromMemoryRecord(record);
  qdrantPoint.Id = Guid.NewGuid();
- this._log.LogTrace("No record with ID {0} found, generated a new point ID {1}", record.Id, qdrantPoint.Id);
+
+ this._log.LogTrace("Generate new Qdrant point ID {0} and record ID {1}", qdrantPoint.Id, record.Id);
  }
  else
  {
- qdrantPoint.Id = existingPoint.Id;
- this._log.LogTrace("Point ID {0} found, updating...", qdrantPoint.Id);
+ qdrantPoint = QdrantPoint<DefaultQdrantPayload>.FromMemoryRecord(record);
+ QdrantPoint<DefaultQdrantPayload>? existingPoint = await this._qdrantClient
+ .GetVectorByPayloadIdAsync(index, record.Id, cancellationToken: cancellationToken)
+ .ConfigureAwait(false);
+ if (existingPoint == null)
+ {
+ qdrantPoint.Id = Guid.NewGuid();
+ this._log.LogTrace("No record with ID {0} found, generated a new point ID {1}", record.Id, qdrantPoint.Id);
+ }
+ else
+ {
+ qdrantPoint.Id = existingPoint.Id;
+ this._log.LogTrace("Point ID {0} found, updating...", qdrantPoint.Id);
+ }
  }
+
+ qdrantPoints.Add(qdrantPoint);
  }
 
- await this._qdrantClient.UpsertVectorsAsync(index, new[] { qdrantPoint }, cancellationToken).ConfigureAwait(false);
+ await this._qdrantClient.UpsertVectorsAsync(index, qdrantPoints, cancellationToken).ConfigureAwait(false);
 
- return record.Id;
+ foreach (var record in localRecords)
+ {
+ yield return record.Id;
+ }
  }
 
  /// <inheritdoc />