Automatic embedding store invalidation

Author: ButterscotchV

BingusApi/EmbeddingServices/RocksDbStore.cs

@@ -5,8 +5,10 @@
 
 namespace BingusApi.EmbeddingServices
 {
-    public class RocksDbStore : IEmbeddingStore
+    public class RocksDbStore : IEmbeddingStore, IDisposable
     {
+        public const string ModelUidKey = "<model_uid>";
+
         private readonly RocksDb rocksDb;
 
         public RocksDbStore(RocksDb rocksDb)
@@ -16,14 +18,21 @@ public RocksDbStore(RocksDb rocksDb)
 
         private static byte[] SerializeString(string value) => Encoding.UTF8.GetBytes(value);
 
+        private static string DeserializeString(byte[] value) => Encoding.UTF8.GetString(value);
+
         public void Add(string key, Vector<float> embedding)
         {
-            rocksDb.Put(SerializeString(key), RocksDbSerializer.SerializeVector(embedding));
+            Add(key, RocksDbSerializer.SerializeVector(embedding));
         }
 
         public void Add(string key, float[] embedding)
         {
-            rocksDb.Put(SerializeString(key), RocksDbSerializer.SerializeArray(embedding));
+            Add(key, RocksDbSerializer.SerializeArray(embedding));
+        }
+
+        public void Add(string key, byte[] data)
+        {
+            rocksDb.Put(SerializeString(key), data);
         }
 
         public byte[]? GetBytes(string key)
@@ -47,5 +56,61 @@ public bool Has(string key)
         {
             return rocksDb.HasKey(SerializeString(key));
         }
+
+        public void Dispose()
+        {
+            rocksDb?.Dispose();
+            GC.SuppressFinalize(this);
+        }
+
+        public static RocksDbStore Create(
+            string dbPath,
+            string modelUid,
+            ILogger<RocksDbStore>? logger
+        )
+        {
+            var preexisting = Directory.Exists(dbPath) && Directory.GetFiles(dbPath).Length > 0;
+
+            var options = new DbOptions().SetCreateIfMissing(true);
+            var db = RocksDb.Open(options, dbPath);
+            var store = new RocksDbStore(db);
+
+            var lastModelUidB = store.GetBytes(ModelUidKey);
+            if (lastModelUidB != null)
+            {
+                var lastModelUid = DeserializeString(lastModelUidB);
+                if (lastModelUid != modelUid)
+                {
+                    logger?.LogWarning(
+                        "Model UID has changed from \"{LastModelUid}\" to \"{ModelUid}\". Deleting the store.",
+                        lastModelUid,
+                        modelUid
+                    );
+
+                    // Delete the store
+                    store.Dispose();
+                    foreach (var file in Directory.GetFiles(dbPath))
+                    {
+                        File.Delete(file);
+                    }
+
+                    // Recreate the store
+                    return Create(dbPath, modelUid, logger);
+                }
+            }
+            else
+            {
+                if (preexisting)
+                {
+                    logger?.LogWarning(
+                        "No model UID found in existing store. Writing current model UID \"{ModelUid}\".",
+                        modelUid
+                    );
+                }
+                store.Add(ModelUidKey, SerializeString(modelUid));
+            }
+
+            return store;
+        }
     }
 }

BingusApi/Program.cs

@@ -4,7 +4,6 @@
 using BingusLib.FaqHandling;
 using BingusLib.HNSW;
 using HNSW.Net;
-using RocksDbSharp;
 
 var builder = WebApplication.CreateBuilder(args);
 
@@ -45,9 +44,6 @@ string GetConfig(string fileName)
 builder.Services.AddSingleton<IRateLimitConfiguration, RateLimitConfiguration>();
 
 // Add dependencies
-builder.Services.AddSingleton<IEmbeddingStore>(sp => new RocksDbStore(
-    RocksDb.Open(new DbOptions().SetCreateIfMissing(true), "embedding_cache")
-));
 builder.Services.AddSingleton(sp =>
     new JsonConfigHandler<BingusConfig>(GetConfig(bingusConfig)).InitializeConfig(
         BingusConfig.Default,
@@ -61,6 +57,13 @@ string GetConfig(string fileName)
     )
 );
 
+// Create persistent embedding store dependent on model hash
+builder.Services.AddSingleton<IEmbeddingStore>(sp =>
+{
+    var modelUid = sp.GetRequiredService<BingusConfig>().GetModelUid();
+    return RocksDbStore.Create("embedding_cache", modelUid, sp.GetService<ILogger<RocksDbStore>>());
+});
+
 // Initialize Bingus library dependencies
 builder.Services.AddHttpClient();
 builder.Services.AddSingleton(sp => sp.GetRequiredService<BingusConfig>().GetSentenceEncoder(sp));

BingusLib/Config/BingusConfig.cs

@@ -57,5 +57,21 @@ public SentenceEncoder GetSentenceEncoder(IServiceProvider serviceProvider)
                     throw new JsonConfigException("No valid sentence encoder type was selected.");
             }
         }
+
+        public string GetModelUri()
+        {
+            return EncoderType.ToLower() switch
+            {
+                "use" => UseModelPath,
+                "api" => ApiUri,
+                "llama" => LlamaModelPath,
+                _ => throw new JsonConfigException("No valid sentence encoder type was selected."),
+            };
+        }
+
+        public string GetModelUid()
+        {
+            return $"{EncoderType.ToLower()}@{GetModelUri()}";
+        }
     }
 }