Added an option to cache embedding calculation results for the Embedding Generator

Author: YuriyIvon

src/DatabaseBenchmark/Generators/EmbeddingGenerator.cs

@@ -3,6 +3,7 @@
 using DatabaseBenchmark.Generators.Options;
 using DatabaseBenchmark.Plugins.Interfaces;
 using DatabaseBenchmark.Plugins.TextEmbedding;
+using System.Collections.Concurrent;
 
 namespace DatabaseBenchmark.Generators
 {
@@ -11,6 +12,8 @@ public class EmbeddingGenerator : IGenerator
         private readonly IGenerator _sourceGenerator;
         private readonly ITextEmbeddingModel _embeddingModel;
         private readonly int? _dimensions;
+        private readonly bool _cache;
+        private readonly ConcurrentDictionary<string, float[]> _embeddingsCache;
 
         public object Current { get; private set; }
 
@@ -30,6 +33,12 @@ public EmbeddingGenerator(EmbeddingGeneratorOptions options, IGenerator sourceGe
             //TODO: refactor resolution and interfaces when non-text embedding models are added
             _embeddingModel = pluginRepository.GetPlugin<ITextEmbeddingModel>(options.ModelName, PluginType.TextEmbeddingModel);
             _dimensions = options.Dimensions;
+            _cache = options.Cache;
+
+            if (_cache)
+            {
+                _embeddingsCache = new ConcurrentDictionary<string, float[]>();
+            }
         }
 
         public bool Next()
@@ -45,8 +54,7 @@ public bool Next()
             {
                 //TODO: refactor when non-text embedding models are added
                 var text = (string)sourceValue;
-                var embedding = _embeddingModel.GenerateEmbedding(text, _dimensions);
-                Current = embedding;
+                Current = GetOrComputeEmbedding(text);
             }
             else
             {
@@ -55,5 +63,10 @@ public bool Next()
 
             return true;
         }
+
+        private float[] GetOrComputeEmbedding(string text) =>
+            _cache
+                ? _embeddingsCache.GetOrAdd(text, t => _embeddingModel.GenerateEmbedding(t, _dimensions))
+                : _embeddingModel.GenerateEmbedding(text, _dimensions);
     }
 }

src/DatabaseBenchmark/Generators/Options/EmbeddingGeneratorOptions.cs

@@ -14,6 +14,8 @@ public class EmbeddingGeneratorOptions : GeneratorOptionsBase
 
         public int? Dimensions { get; set; }
 
+        public bool Cache { get; set; } = false;
+
         public IGeneratorOptions SourceGeneratorOptions { get; set; }
 
         public enum GeneratorKind