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

app/MindWork AI Studio/Tools/Event.cs

@@ -37,6 +37,7 @@ public enum Event
     
     // RAG events:
     RAG_AUTO_DATA_SOURCES_SELECTED,
+    RAG_EMBEDDING_STATUS_CHANGED,
     
     // File attachment events:
     REGISTER_FILE_DROP_AREA,

app/MindWork AI Studio/Tools/Services/RustService.Databases.cs

@@ -1,9 +1,46 @@
-using AIStudio.Tools.Rust;
+using AIStudio.Tools.Databases;
+using AIStudio.Tools.Rust;
 
 namespace AIStudio.Tools.Services;
 
 public sealed partial class RustService
 {
+    public async Task<EmbeddingStoreConfiguration> GetEmbeddingStoreConfiguration(EmbeddingStoreKind kind)
+    {
+        switch (kind)
+        {
+            case EmbeddingStoreKind.QDRANT_REMOTE:
+            {
+                var qdrantInfo = await this.GetQdrantInfo();
+                var invalidFields = new List<string>();
+                if (!qdrantInfo.IsAvailable)
+                    invalidFields.Add(qdrantInfo.UnavailableReason ?? "unknown");
+                if (string.IsNullOrWhiteSpace(qdrantInfo.Path))
+                    invalidFields.Add("Path");
+                if (qdrantInfo.PortHttp == 0)
+                    invalidFields.Add("HttpPort");
+                if (qdrantInfo.PortGrpc == 0)
+                    invalidFields.Add("GrpcPort");
+                if (string.IsNullOrWhiteSpace(qdrantInfo.Fingerprint))
+                    invalidFields.Add("Fingerprint");
+                if (string.IsNullOrWhiteSpace(qdrantInfo.ApiToken))
+                    invalidFields.Add("ApiToken");
+                if (invalidFields.Count <= 0) return new EmbeddingStoreConfiguration(kind, "Qdrant", new RemoteLocation(qdrantInfo.Path, qdrantInfo.PortHttp, qdrantInfo.PortGrpc, qdrantInfo.Fingerprint, qdrantInfo.ApiToken), null);
+                var reason = string.Join(", ", invalidFields);
+                Console.WriteLine($"Warning: Qdrant is not available. Starting without vector database. Reason: '{reason}'.");
+
+                return new EmbeddingStoreConfiguration(
+                    EmbeddingStoreKind.NONE,
+                    "Qdrant",
+                    null,
+                    reason);
+
+            }
+            default:
+                return new EmbeddingStoreConfiguration(kind, kind.ToString(), null, $"No configuration available for {kind}");
+        }
+    }
+    
     public async Task<QdrantInfo> GetQdrantInfo()
     {
         try

app/MindWork AI Studio/Tools/Services/RustService.Retrieval.cs

@@ -1,5 +1,6 @@
 using System.Text;
 using System.Text.Json;
+using System.Runtime.CompilerServices;
 
 namespace AIStudio.Tools.Services;
 
@@ -48,6 +49,9 @@ public sealed partial class RustService
                 }
                 catch (JsonException)
                 {
+                    if (this.TryLogSseErrorMessage(jsonContent, path))
+                        continue;
+
                     this.logger?.LogError("Failed to deserialize SSE event: {JsonContent}", jsonContent);
                 }
             }
@@ -65,4 +69,77 @@ public sealed partial class RustService
         
         return resultBuilder.ToString();
     }
-}
+
+    public async IAsyncEnumerable<string> StreamArbitraryFileData(string path, bool extractImages = false, [EnumeratorCancellation] CancellationToken token = default)
+    {
+        var streamId = Guid.NewGuid().ToString();
+        var requestUri = $"/retrieval/fs/extract?path={Uri.EscapeDataString(path)}&stream_id={streamId}&extract_images={extractImages}";
+        using var request = new HttpRequestMessage(HttpMethod.Get, requestUri);
+        using var response = await this.http.SendAsync(request, HttpCompletionOption.ResponseHeadersRead, token);
+
+        if (!response.IsSuccessStatusCode)
+            yield break;
+
+        string? finalContentChunk = null;
+        try
+        {
+            await using var stream = await response.Content.ReadAsStreamAsync(token);
+            using var reader = new StreamReader(stream);
+
+            while (!reader.EndOfStream && !token.IsCancellationRequested)
+            {
+                var line = await reader.ReadLineAsync(token);
+                if (string.IsNullOrWhiteSpace(line))
+                    continue;
+
+                if (!line.StartsWith("data:", StringComparison.InvariantCulture))
+                    continue;
+
+                var jsonContent = line[5..];
+                ContentStreamSseEvent? sseEvent = null;
+                try
+                {
+                    sseEvent = JsonSerializer.Deserialize<ContentStreamSseEvent>(jsonContent);
+                }
+                catch (JsonException)
+                {
+                    if (this.TryLogSseErrorMessage(jsonContent, path))
+                        continue;
+
+                    this.logger?.LogError("Failed to deserialize SSE event: {JsonContent}", jsonContent);
+                }
+
+                if (sseEvent is null)
+                    continue;
+
+                var content = ContentStreamSseHandler.ProcessEvent(sseEvent, extractImages);
+                if (!string.IsNullOrWhiteSpace(content))
+                    yield return content;
+            }
+        }
+        finally
+        {
+            finalContentChunk = ContentStreamSseHandler.Clear(streamId);
+        }
+
+        if (!string.IsNullOrWhiteSpace(finalContentChunk))
+            yield return finalContentChunk;
+    }
+
+    private bool TryLogSseErrorMessage(string jsonContent, string path)
+    {
+        try
+        {
+            var errorMessage = JsonSerializer.Deserialize<string>(jsonContent);
+            if (string.IsNullOrWhiteSpace(errorMessage))
+                return false;
+
+            this.logger?.LogError("Rust retrieval stream error for '{Path}': {ErrorMessage}", path, errorMessage);
+            return true;
+        }
+        catch (JsonException)
+        {
+            return false;
+        }
+    }
+}