background embed

app/MindWork AI Studio/Tools/Databases/Qdrant/QdrantClientImplementation.cs

@@ -1,10 +1,11 @@
 using Qdrant.Client;
 using Qdrant.Client.Grpc;
 using AIStudio.Tools.PluginSystem;
+using static Qdrant.Client.Grpc.Conditions;
 
 namespace AIStudio.Tools.Databases.Qdrant;
 
-public class QdrantClientImplementation : DatabaseClient
+public class QdrantClientImplementation : EmbeddingStore
 {
     private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(QdrantClientImplementation).Namespace, nameof(QdrantClientImplementation));
 
@@ -18,12 +19,12 @@ public class QdrantClientImplementation : DatabaseClient
     
     private string ApiToken { get; }
     
-    public QdrantClientImplementation(string name, string path, int httpPort, int grpcPort, string fingerprint, string apiToken): base(name, path)
+    public QdrantClientImplementation(string name, string path, int? httpPort, int? grpcPort, string? fingerprint, string? apiToken): base(name, path)
     {
-        this.HttpPort = httpPort;
-        this.GrpcPort = grpcPort;
-        this.Fingerprint = fingerprint;
-        this.ApiToken = apiToken;
+        this.HttpPort = httpPort ?? 0;
+        this.GrpcPort = grpcPort ?? 0;
+        this.Fingerprint = fingerprint ?? string.Empty;
+        this.ApiToken = apiToken ?? string.Empty;
         this.GrpcClient = this.CreateQdrantClient();
     }
     
@@ -62,5 +63,56 @@ public class QdrantClientImplementation : DatabaseClient
         yield return (TB("Number of collections"), await this.GetCollectionsAmount());
     }
 
+    public override async Task EnsureEmbeddingStoreExists(string collectionName, int vectorSize, CancellationToken token)
+    {
+        var exists = await this.GrpcClient.CollectionExistsAsync(collectionName, token);
+        if (exists)
+            return;
+
+        await this.GrpcClient.CreateCollectionAsync(
+            collectionName,
+            new VectorParams
+            {
+                Size = (ulong)vectorSize,
+                Distance = Distance.Cosine,
+            },
+            cancellationToken: token);
+    }
+
+    public override Task InsertEmbedding(string collectionName, IReadOnlyList<EmbeddingStoragePoint> points, CancellationToken token)
+    {
+        var qdrantPoints = points.Select(point => new PointStruct
+        {
+            Id = Guid.Parse(point.PointId),
+            Vectors = point.Vector.ToArray(),
+            Payload =
+            {
+                ["data_source_id"] = point.DataSourceId,
+                ["data_source_name"] = point.DataSourceName,
+                ["data_source_type"] = point.DataSourceType,
+                ["file_path"] = point.FilePath,
+                ["file_name"] = point.FileName,
+                ["relative_path"] = point.RelativePath,
+                ["chunk_index"] = (long)point.ChunkIndex,
+                ["text"] = point.Text,
+                ["fingerprint"] = point.Fingerprint,
+                ["last_write_utc"] = point.LastWriteUtc.ToString("O"),
+                ["embedded_at_utc"] = point.EmbeddedAtUtc.ToString("O"),
+            }
+        }).ToList();
+
+        return this.GrpcClient.UpsertAsync(collectionName, qdrantPoints, true, null, null, token);
+    }
+
+    public override Task DeleteEmbeddingByFile(string collectionName, string filePath, CancellationToken token)
+    {
+        return this.GrpcClient.DeleteAsync(collectionName, MatchKeyword("file_path", filePath), true, null, null, token);
+    }
+
+    public override Task DeleteEmbeddingStore(string collectionName, CancellationToken token)
+    {
+        return this.GrpcClient.DeleteCollectionAsync(collectionName, cancellationToken: token);
+    }
+
     public override void Dispose() => this.GrpcClient.Dispose();
-}
+}