background embed

app/MindWork AI Studio/Tools/Databases/EmbeddingStoragePoint.cs

@@ -0,0 +1,16 @@
+namespace AIStudio.Tools.Databases;
+
+public sealed record EmbeddingStoragePoint(
+    string PointId,
+    IReadOnlyList<float> Vector,
+    string DataSourceId,
+    string DataSourceName,
+    string DataSourceType,
+    string FilePath,
+    string FileName,
+    string RelativePath,
+    int ChunkIndex,
+    string Text,
+    string Fingerprint,
+    DateTime LastWriteUtc,
+    DateTime EmbeddedAtUtc);

app/MindWork AI Studio/Tools/Databases/EmbeddingStore.cs

@@ -1,6 +1,6 @@
 namespace AIStudio.Tools.Databases;
 
-public abstract class DatabaseClient(string name, string path)
+public abstract class EmbeddingStore(string name, string path)
 {
     public string Name => name;
 
@@ -8,7 +8,7 @@ public abstract class DatabaseClient(string name, string path)
     
     private string Path => path;
 
-    private ILogger<DatabaseClient>? logger;
+    private ILogger<EmbeddingStore>? logger;
     
     public abstract IAsyncEnumerable<(string Label, string Value)> GetDisplayInfo();
 
@@ -45,10 +45,19 @@ public abstract class DatabaseClient(string name, string path)
         return $"{size:0##} {suffixes[suffixIndex]}";
     }
     
-    public void SetLogger(ILogger<DatabaseClient> logService)
+    public void SetLogger(ILogger<EmbeddingStore> logService)
     {
         this.logger = logService;
     }
+    
+
+    public abstract Task EnsureEmbeddingStoreExists(string collectionName, int vectorSize, CancellationToken token);
+
+    public abstract Task InsertEmbedding(string collectionName, IReadOnlyList<EmbeddingStoragePoint> points, CancellationToken token);
+
+    public abstract Task DeleteEmbeddingByFile(string collectionName, string filePath, CancellationToken token);
+
+    public abstract Task DeleteEmbeddingStore(string collectionName, CancellationToken token);
 
     public abstract void Dispose();
 }

app/MindWork AI Studio/Tools/Databases/EmbeddingStoreFactory.cs

@@ -0,0 +1,33 @@
+using AIStudio.Tools.Databases.Qdrant;
+
+namespace AIStudio.Tools.Databases;
+
+public class EmbeddingStoreFactory
+{
+    public static EmbeddingStore Create(EmbeddingStoreConfiguration configuration) => configuration.Kind switch
+    {
+        EmbeddingStoreKind.NONE => new NoEmbeddingStore(configuration.Name, configuration.UnavailableReason ?? "unknown"),
+        _ when configuration.Location is null => new NoEmbeddingStore(configuration.Name, $"No location specified for {configuration.Name}"),
+        EmbeddingStoreKind.QDRANT_REMOTE when configuration.Location is RemoteLocation location=> new QdrantClientImplementation(configuration.Name, location.Path, location.HttpPort, location.GrpcPort, location.Fingerprint, location.ApiToken),
+        _ => throw new ArgumentException("Invalid configuration for " + configuration.Name, nameof(configuration)),
+    };
+}
+
+public enum EmbeddingStoreKind
+{
+    NONE,
+    QDRANT_EMBED,
+    QDRANT_REMOTE,
+}
+
+public abstract record EmbeddingStoreLocation;
+
+public sealed record EmbeddedLocation(string Path) : EmbeddingStoreLocation;
+
+public sealed record RemoteLocation(string Path, int? HttpPort, int? GrpcPort, string? Fingerprint, string? ApiToken) : EmbeddingStoreLocation;
+
+public sealed record EmbeddingStoreConfiguration(
+    EmbeddingStoreKind Kind,
+    string Name,
+    EmbeddingStoreLocation? Location,
+    string? UnavailableReason);

app/MindWork AI Studio/Tools/Databases/NoDatabaseClient.cs

@@ -1,24 +0,0 @@
-using AIStudio.Tools.PluginSystem;
-
-namespace AIStudio.Tools.Databases;
-
-public sealed class NoDatabaseClient(string name, string? unavailableReason) : DatabaseClient(name, string.Empty)
-{
-    private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(NoDatabaseClient).Namespace, nameof(NoDatabaseClient));
-    
-    public override bool IsAvailable => false;
-
-    public override async IAsyncEnumerable<(string Label, string Value)> GetDisplayInfo()
-    {
-        yield return (TB("Status"), TB("Unavailable"));
-
-        if (!string.IsNullOrWhiteSpace(unavailableReason))
-            yield return (TB("Reason"), unavailableReason);
-
-        await Task.CompletedTask;
-    }
-
-    public override void Dispose()
-    {
-    }
-}

app/MindWork AI Studio/Tools/Databases/NoEmbeddingStore.cs

@@ -0,0 +1,39 @@
+using AIStudio.Tools.PluginSystem;
+
+namespace AIStudio.Tools.Databases;
+
+public sealed class NoEmbeddingStore(string name, string? unavailableReason) : EmbeddingStore(name, string.Empty)
+{
+    private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(NoEmbeddingStore).Namespace, nameof(NoEmbeddingStore));
+    
+    public override bool IsAvailable => false;
+
+    public override async IAsyncEnumerable<(string Label, string Value)> GetDisplayInfo()
+    {
+        yield return (TB("Status"), TB("Unavailable"));
+
+        if (!string.IsNullOrWhiteSpace(unavailableReason))
+            yield return (TB("Reason"), unavailableReason);
+
+        await Task.CompletedTask;
+    }
+
+    public override Task EnsureEmbeddingStoreExists(string collectionName, int vectorSize, CancellationToken token) => throw this.BuildUnavailableException();
+
+    public override Task InsertEmbedding(string collectionName, IReadOnlyList<EmbeddingStoragePoint> points, CancellationToken token) => throw this.BuildUnavailableException();
+
+    public override Task DeleteEmbeddingByFile(string collectionName, string filePath, CancellationToken token) => Task.CompletedTask;
+
+    public override Task DeleteEmbeddingStore(string collectionName, CancellationToken token) => Task.CompletedTask;
+
+    public override void Dispose()
+    {
+    }
+
+    private InvalidOperationException BuildUnavailableException()
+    {
+        return new InvalidOperationException(string.IsNullOrWhiteSpace(unavailableReason)
+            ? "The vector database is not available."
+            : unavailableReason);
+    }
+}

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();
-}
+}