@@ -35,11 +35,11 @@
}
-
- @(this.showDatabaseDetails ? T("Hide Details") : T("Show Details"))
+ OnClick="@this.ToggleVectorStoreDetails">
+ @(this.showVectorStoreDetails ? T("Hide Details") : T("Show Details"))
diff --git a/app/MindWork AI Studio/Pages/Information.razor.cs b/app/MindWork AI Studio/Pages/Information.razor.cs
index 0cc199e0..7c018da2 100644
--- a/app/MindWork AI Studio/Pages/Information.razor.cs
+++ b/app/MindWork AI Studio/Pages/Information.razor.cs
@@ -29,13 +29,13 @@ public partial class Information : MSGComponentBase
private ISnackbar Snackbar { get; init; } = null!;
[Inject]
- private DatabaseClient DatabaseClient { get; init; } = null!;
+ private DatabaseClientProvider DatabaseClientProvider { get; init; } = null!;
private static readonly Assembly ASSEMBLY = Assembly.GetExecutingAssembly();
private static readonly MetaDataAttribute META_DATA = ASSEMBLY.GetCustomAttribute
()!;
private static readonly MetaDataArchitectureAttribute META_DATA_ARCH = ASSEMBLY.GetCustomAttribute()!;
private static readonly MetaDataLibrariesAttribute META_DATA_LIBRARIES = ASSEMBLY.GetCustomAttribute()!;
- private static readonly MetaDataEmbeddingStoreAttribute META_DATA_EMBEDDING_STORE = ASSEMBLY.GetCustomAttribute()!;
+ private static readonly MetaDataVectorStoreAttribute META_DATA_VECTOR_STORE = ASSEMBLY.GetCustomAttribute()!;
private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(Information).Namespace, nameof(Information));
@@ -62,18 +62,18 @@ public partial class Information : MSGComponentBase
private string VersionPdfium => $"{T("Used PDFium version")}: v{META_DATA_LIBRARIES.PdfiumVersion}";
- private string VersionEmbeddingStore
+ private string VersionVectorStore
{
get
{
- if (this.embeddingStore is null)
- return $"{T("Embedding store")}: {T("checking availability")}";
+ if (this.vectorStore is null)
+ return $"{T("Vector store")}: {T("checking availability")}";
- return this.embeddingStore.Status switch
+ return this.vectorStore.Status switch
{
- EmbeddingStoreStatus.AVAILABLE => $"{T("Embedding store version")}: {this.embeddingStore.Name} v{META_DATA_EMBEDDING_STORE.DatabaseVersion}",
- EmbeddingStoreStatus.STARTING => $"{T("Embedding store")}: {this.embeddingStore.Name} - {T("starting")}",
- _ => $"{T("Embedding store")}: {this.embeddingStore.Name} - {T("not available")}"
+ DatabaseClientStatus.AVAILABLE => $"{T("Vector store version")}: {this.vectorStore.Name} v{META_DATA_VECTOR_STORE.VectorStoreVersion}",
+ DatabaseClientStatus.STARTING => $"{T("Vector store")}: {this.vectorStore.Name} - {T("starting")}",
+ _ => $"{T("Vector store")}: {this.vectorStore.Name} - {T("not available")}"
};
}
}
@@ -85,7 +85,7 @@ public partial class Information : MSGComponentBase
private bool showEnterpriseConfigDetails;
- private bool showEmbeddingStoreDetails;
+ private bool showVectorStoreDetails;
private List configPlugins = PluginFactory.AvailablePlugins
.Where(x => x.Type is PluginType.CONFIGURATION)
@@ -95,14 +95,13 @@ public partial class Information : MSGComponentBase
private List enterpriseEnvironments = EnterpriseEnvironmentService.CURRENT_ENVIRONMENTS.ToList();
private List mandatoryInfoPanels = [];
-
- private sealed record EmbeddingStoreDisplayInfo(string Label, string Value);
private sealed record MandatoryInfoPanelData(string HeaderText, string PluginName, DataMandatoryInfo Info, DataMandatoryInfoAcceptance? Acceptance);
-
- private readonly List embeddingStoreDisplayInfo = new();
- private DatabaseClient? embeddingStore;
- private CancellationTokenSource? databaseRefreshCancellationTokenSource;
+
+ private sealed record VectorStoreDisplayInfo(string Label, string Value);
+ private readonly List vectorStoreDisplayInfo = new();
+ private DatabaseClient? vectorStore;
+ private CancellationTokenSource? vectorStoreRefreshCancellationTokenSource;
private bool HasAnyActiveEnvironment => this.enterpriseEnvironments.Any(e => e.IsActive);
@@ -148,9 +147,9 @@ public partial class Information : MSGComponentBase
this.osUserName = await this.RustService.ReadUserName();
this.logPaths = await this.RustService.GetLogPaths();
- await this.RefreshDatabaseInfo(CancellationToken.None);
- if (this.databaseClient?.Status is DatabaseClientStatus.STARTING)
- this.StartShortDatabaseRefreshLoop();
+ await this.RefreshVectorStoreInfo(CancellationToken.None);
+ if (this.vectorStore?.Status is DatabaseClientStatus.STARTING)
+ this.StartShortVectorStoreRefreshLoop();
// Determine the Pandoc version may take some time, so we start it here
// without waiting for the result:
@@ -249,22 +248,22 @@ public partial class Information : MSGComponentBase
this.showEnterpriseConfigDetails = !this.showEnterpriseConfigDetails;
}
- private void ToggleDatabaseDetails()
+ private void ToggleVectorStoreDetails()
{
- this.showEmbeddingStoreDetails = !this.showEmbeddingStoreDetails;
+ this.showVectorStoreDetails = !this.showVectorStoreDetails;
}
- private async Task RefreshDatabaseInfo(CancellationToken cancellationToken)
+ private async Task RefreshVectorStoreInfo(CancellationToken cancellationToken)
{
var refreshedClient = await this.DatabaseClientProvider.RefreshClientAsync(DatabaseRole.VECTOR_STORE, cancellationToken);
- this.databaseClient = refreshedClient;
- this.embeddingStoreDisplayInfo.Clear();
+ this.vectorStore = refreshedClient;
+ this.vectorStoreDisplayInfo.Clear();
try
{
await foreach (var (label, value) in refreshedClient.GetDisplayInfo().WithCancellation(cancellationToken))
{
- this.embeddingStoreDisplayInfo.Add(new EmbeddingStoreDisplayInfo(label, value));
+ this.vectorStoreDisplayInfo.Add(new VectorStoreDisplayInfo(label, value));
}
}
catch (OperationCanceledException)
@@ -273,20 +272,20 @@ public partial class Information : MSGComponentBase
}
catch (Exception e)
{
- this.databaseClient = new NoDatabaseClient(refreshedClient.Name, e.Message, DatabaseClientStatus.STARTING);
- await foreach (var (label, value) in this.databaseClient.GetDisplayInfo().WithCancellation(cancellationToken))
+ this.vectorStore = new NoDatabaseClient(refreshedClient.Name, e.Message, DatabaseClientStatus.STARTING);
+ await foreach (var (label, value) in this.vectorStore.GetDisplayInfo().WithCancellation(cancellationToken))
{
- this.embeddingStoreDisplayInfo.Add(new EmbeddingStoreDisplayInfo(label, value));
+ this.vectorStoreDisplayInfo.Add(new VectorStoreDisplayInfo(label, value));
}
}
}
- private void StartShortDatabaseRefreshLoop()
+ private void StartShortVectorStoreRefreshLoop()
{
- this.databaseRefreshCancellationTokenSource?.Cancel();
- this.databaseRefreshCancellationTokenSource?.Dispose();
- this.databaseRefreshCancellationTokenSource = new CancellationTokenSource();
- var cancellationToken = this.databaseRefreshCancellationTokenSource.Token;
+ this.vectorStoreRefreshCancellationTokenSource?.Cancel();
+ this.vectorStoreRefreshCancellationTokenSource?.Dispose();
+ this.vectorStoreRefreshCancellationTokenSource = new CancellationTokenSource();
+ var cancellationToken = this.vectorStoreRefreshCancellationTokenSource.Token;
_ = Task.Run(async () =>
{
@@ -298,11 +297,11 @@ public partial class Information : MSGComponentBase
await Task.Delay(TimeSpan.FromSeconds(1), cancellationToken);
await this.InvokeAsync(async () =>
{
- await this.RefreshDatabaseInfo(cancellationToken);
+ await this.RefreshVectorStoreInfo(cancellationToken);
this.StateHasChanged();
});
- if (this.databaseClient?.Status is not DatabaseClientStatus.STARTING)
+ if (this.vectorStore?.Status is not DatabaseClientStatus.STARTING)
return;
}
catch (OperationCanceledException)
@@ -331,8 +330,8 @@ public partial class Information : MSGComponentBase
protected override void DisposeResources()
{
- this.databaseRefreshCancellationTokenSource?.Cancel();
- this.databaseRefreshCancellationTokenSource?.Dispose();
+ this.vectorStoreRefreshCancellationTokenSource?.Cancel();
+ this.vectorStoreRefreshCancellationTokenSource?.Dispose();
base.DisposeResources();
}
diff --git a/app/MindWork AI Studio/Program.cs b/app/MindWork AI Studio/Program.cs
index b6ac9c21..21073095 100644
--- a/app/MindWork AI Studio/Program.cs
+++ b/app/MindWork AI Studio/Program.cs
@@ -28,7 +28,7 @@ internal sealed class Program
public static string API_TOKEN = null!;
public static IServiceProvider SERVICE_PROVIDER = null!;
public static ILoggerFactory LOGGER_FACTORY = null!;
- public static EmbeddingStore EMBEDDING_STORE = null!;
+ public static DatabaseClientProvider DATABASE_CLIENT_PROVIDER = null!;
public static async Task Main()
{
@@ -87,10 +87,6 @@ internal sealed class Program
return;
}
- var embeddingStoreConfig = await rust.GetEmbeddingStoreConfiguration(EmbeddingStoreKind.QDRANT_REMOTE);
-
- var embeddingStore = EmbeddingStoreFactory.Create(embeddingStoreConfig);
-
var builder = WebApplication.CreateBuilder();
builder.WebHost.ConfigureKestrel(kestrelServerOptions =>
{
@@ -148,7 +144,7 @@ internal sealed class Program
builder.Services.AddHostedService();
builder.Services.AddHostedService();
builder.Services.AddHostedService(sp => sp.GetRequiredService());
- builder.Services.AddSingleton(embeddingStore);
+ builder.Services.AddSingleton();
builder.Services.AddHostedService();
builder.Services.AddHostedService();
@@ -208,9 +204,7 @@ internal sealed class Program
RUST_SERVICE = rust;
ENCRYPTION = encryption;
- var databaseLogger = app.Services.GetRequiredService>();
- embeddingStore.SetLogger(databaseLogger);
- EMBEDDING_STORE = embeddingStore;
+ DATABASE_CLIENT_PROVIDER = app.Services.GetRequiredService();
programLogger.LogInformation("Initialize internal file system.");
app.Use(Redirect.HandlerContentAsync);
@@ -248,7 +242,7 @@ internal sealed class Program
await serverTask;
RUST_SERVICE.Dispose();
- EMBEDDING_STORE.Dispose();
+ DATABASE_CLIENT_PROVIDER.Dispose();
PluginFactory.Dispose();
programLogger.LogInformation("The AI Studio server was stopped.");
}
diff --git a/app/MindWork AI Studio/Tools/Databases/DatabaseClient.cs b/app/MindWork AI Studio/Tools/Databases/DatabaseClient.cs
index dd016fec..7186ed21 100644
--- a/app/MindWork AI Studio/Tools/Databases/DatabaseClient.cs
+++ b/app/MindWork AI Studio/Tools/Databases/DatabaseClient.cs
@@ -1,6 +1,6 @@
namespace AIStudio.Tools.Databases;
-public abstract class EmbeddingStore(string name, string path)
+public abstract class DatabaseClient(string name, string path)
{
public string Name => name;
@@ -12,7 +12,7 @@ public abstract class EmbeddingStore(string name, string path)
private string Path => path;
- private ILogger? logger;
+ private ILogger? logger;
public abstract IAsyncEnumerable<(string Label, string Value)> GetDisplayInfo();
@@ -50,7 +50,7 @@ public abstract class EmbeddingStore(string name, string path)
return $"{convertedSize:0.##} {suffixes[suffixIndex]}";
}
- public void SetLogger(ILogger logService)
+ public void SetLogger(ILogger logService)
{
this.logger = logService;
}
diff --git a/app/MindWork AI Studio/Tools/Databases/DatabaseClientProvider.Qdrant.cs b/app/MindWork AI Studio/Tools/Databases/DatabaseClientProvider.Qdrant.cs
index c01cdc13..f3676bb3 100644
--- a/app/MindWork AI Studio/Tools/Databases/DatabaseClientProvider.Qdrant.cs
+++ b/app/MindWork AI Studio/Tools/Databases/DatabaseClientProvider.Qdrant.cs
@@ -1,4 +1,5 @@
-using AIStudio.Tools.Rust;
+using AIStudio.Tools.Databases.VectorStore;
+using AIStudio.Tools.Rust;
namespace AIStudio.Tools.Databases;
@@ -6,7 +7,7 @@ public sealed partial class DatabaseClientProvider
{
private async Task CreateQdrantClientAsync(CancellationToken cancellationToken)
{
- var qdrantInfo = await this.rustService.GetQdrantInfo(cancellationToken);
+ var qdrantInfo = await rustService.GetQdrantInfo(cancellationToken);
if (qdrantInfo.Status is QdrantStatus.STARTING)
{
return this.CreateNoDatabaseClient(
diff --git a/app/MindWork AI Studio/Tools/Databases/DatabaseClientProvider.cs b/app/MindWork AI Studio/Tools/Databases/DatabaseClientProvider.cs
index a6f98e60..55087784 100644
--- a/app/MindWork AI Studio/Tools/Databases/DatabaseClientProvider.cs
+++ b/app/MindWork AI Studio/Tools/Databases/DatabaseClientProvider.cs
@@ -1,15 +1,14 @@
-using AIStudio.Tools.Databases.Qdrant;
-using AIStudio.Tools.Rust;
using AIStudio.Tools.Services;
+using AIStudio.Tools.Databases.VectorStore;
namespace AIStudio.Tools.Databases;
-public sealed class EmbeddingStoreProvider(RustService rustService, ILoggerFactory loggerFactory) : IDisposable
+public sealed partial class DatabaseClientProvider(RustService rustService, ILoggerFactory loggerFactory) : IDisposable
{
- private readonly Dictionary clients = new();
+ private readonly Dictionary clients = new();
private readonly Dictionary locks = new();
private readonly Lock locksLock = new();
- private readonly ILogger logger = loggerFactory.CreateLogger();
+ private readonly ILogger logger = loggerFactory.CreateLogger();
private readonly ILogger databaseClientLogger = loggerFactory.CreateLogger();
public async Task GetClientAsync(DatabaseRole databaseRole, CancellationToken cancellationToken = default)
@@ -45,6 +44,18 @@ public sealed class EmbeddingStoreProvider(RustService rustService, ILoggerFacto
}
}
+ public async Task GetVectorStoreAsync(CancellationToken cancellationToken = default)
+ {
+ var client = await this.GetClientAsync(DatabaseRole.VECTOR_STORE, cancellationToken);
+ if (client is IVectorStoreClient vectorStore)
+ return vectorStore;
+
+ return new NoVectorStoreClient(
+ client.Name,
+ "The configured database client does not support vector store operations.",
+ client.Status);
+ }
+
private DatabaseClient CacheIfAvailable(DatabaseRole databaseRole, DatabaseClient client)
{
if (!client.IsAvailable)
@@ -84,79 +95,6 @@ public sealed class EmbeddingStoreProvider(RustService rustService, ILoggerFacto
_ => new NoDatabaseClient(databaseRole.ToString(), "The requested database role is not supported.")
};
- private async Task CreateQdrantClientAsync(CancellationToken cancellationToken)
- {
- var qdrantInfo = await rustService.GetQdrantInfo(cancellationToken);
- if (qdrantInfo.Status is QdrantStatus.STARTING)
- {
- return this.CreateNoDatabaseClient(
- "Qdrant",
- "Qdrant is starting. Details will appear shortly.",
- DatabaseClientStatus.STARTING);
- }
-
- if (!qdrantInfo.IsAvailable || qdrantInfo.Status is QdrantStatus.UNAVAILABLE)
- {
- var reason = qdrantInfo.UnavailableReason ?? "unknown";
- this.logger.LogWarning("Qdrant is not available. Starting without vector database. Reason: '{Reason}'.", reason);
- return this.CreateNoDatabaseClient("Qdrant", qdrantInfo.UnavailableReason, DatabaseClientStatus.UNAVAILABLE);
- }
-
- if (!HasValidQdrantConnectionInfo(qdrantInfo, out var invalidReason))
- return this.CreateNoDatabaseClient("Qdrant", invalidReason, DatabaseClientStatus.UNAVAILABLE);
-
- var client = new QdrantClientImplementation("Qdrant", qdrantInfo.Path, qdrantInfo.PortHttp, qdrantInfo.PortGrpc, qdrantInfo.Fingerprint, qdrantInfo.ApiToken);
- client.SetLogger(this.databaseClientLogger);
-
- try
- {
- await client.CheckAvailabilityAsync();
- return client;
- }
- catch (Exception e)
- {
- client.Dispose();
- this.logger.LogWarning(e, "Qdrant reported as available by Rust, but the health check failed.");
- return this.CreateNoDatabaseClient("Qdrant", e.Message, DatabaseClientStatus.STARTING);
- }
- }
-
- private static bool HasValidQdrantConnectionInfo(QdrantInfo qdrantInfo, out string invalidReason)
- {
- if (qdrantInfo.Path == string.Empty)
- {
- invalidReason = "Failed to get the Qdrant path from Rust.";
- return false;
- }
-
- if (qdrantInfo.PortHttp == 0)
- {
- invalidReason = "Failed to get the Qdrant HTTP port from Rust.";
- return false;
- }
-
- if (qdrantInfo.PortGrpc == 0)
- {
- invalidReason = "Failed to get the Qdrant gRPC port from Rust.";
- return false;
- }
-
- if (qdrantInfo.Fingerprint == string.Empty)
- {
- invalidReason = "Failed to get the Qdrant fingerprint from Rust.";
- return false;
- }
-
- if (qdrantInfo.ApiToken == string.Empty)
- {
- invalidReason = "Failed to get the Qdrant API token from Rust.";
- return false;
- }
-
- invalidReason = string.Empty;
- return true;
- }
-
private NoDatabaseClient CreateNoDatabaseClient(string name, string? unavailableReason, DatabaseClientStatus status)
{
var client = new NoDatabaseClient(name, unavailableReason, status);
diff --git a/app/MindWork AI Studio/Tools/Databases/EmbeddingStoreFactory.cs b/app/MindWork AI Studio/Tools/Databases/EmbeddingStoreFactory.cs
deleted file mode 100644
index 9fa51305..00000000
--- a/app/MindWork AI Studio/Tools/Databases/EmbeddingStoreFactory.cs
+++ /dev/null
@@ -1,33 +0,0 @@
-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);
\ No newline at end of file
diff --git a/app/MindWork AI Studio/Tools/Databases/NoDatabaseClient.cs b/app/MindWork AI Studio/Tools/Databases/NoDatabaseClient.cs
index 230ec1fd..804a8948 100644
--- a/app/MindWork AI Studio/Tools/Databases/NoDatabaseClient.cs
+++ b/app/MindWork AI Studio/Tools/Databases/NoDatabaseClient.cs
@@ -2,9 +2,9 @@ using AIStudio.Tools.PluginSystem;
namespace AIStudio.Tools.Databases;
-public sealed class NoEmbeddingStore(string name, string? unavailableReason, DatabaseClientStatus status = DatabaseClientStatus.UNAVAILABLE) : EmbeddingStore(name, string.Empty)
+public sealed class NoDatabaseClient(string name, string? unavailableReason, DatabaseClientStatus status = DatabaseClientStatus.UNAVAILABLE) : DatabaseClient(name, string.Empty)
{
- private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(NoEmbeddingStore).Namespace, nameof(NoEmbeddingStore));
+ private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(NoDatabaseClient).Namespace, nameof(NoDatabaseClient));
public override DatabaseClientStatus Status => status;
diff --git a/app/MindWork AI Studio/Tools/Databases/EmbeddingStoragePoint.cs b/app/MindWork AI Studio/Tools/Databases/VectorStoragePoint.cs
similarity index 88%
rename from app/MindWork AI Studio/Tools/Databases/EmbeddingStoragePoint.cs
rename to app/MindWork AI Studio/Tools/Databases/VectorStoragePoint.cs
index 0b14fc2b..a614d664 100644
--- a/app/MindWork AI Studio/Tools/Databases/EmbeddingStoragePoint.cs
+++ b/app/MindWork AI Studio/Tools/Databases/VectorStoragePoint.cs
@@ -1,6 +1,6 @@
namespace AIStudio.Tools.Databases;
-public sealed record EmbeddingStoragePoint(
+public sealed record VectorStoragePoint(
string PointId,
IReadOnlyList Vector,
string DataSourceId,
diff --git a/app/MindWork AI Studio/Tools/Databases/VectorStore/IVectorStoreClient.cs b/app/MindWork AI Studio/Tools/Databases/VectorStore/IVectorStoreClient.cs
new file mode 100644
index 00000000..f1e96623
--- /dev/null
+++ b/app/MindWork AI Studio/Tools/Databases/VectorStore/IVectorStoreClient.cs
@@ -0,0 +1,20 @@
+namespace AIStudio.Tools.Databases.VectorStore;
+
+public interface IVectorStoreClient
+{
+ string Name { get; }
+
+ DatabaseClientStatus Status { get; }
+
+ bool IsAvailable { get; }
+
+ IAsyncEnumerable<(string Label, string Value)> GetDisplayInfo();
+
+ Task EnsureVectorStoreExists(string storeName, int vectorSize, CancellationToken token);
+
+ Task InsertEmbedding(string storeName, IReadOnlyList points, CancellationToken token);
+
+ Task DeleteEmbeddingByFile(string storeName, string filePath, CancellationToken token);
+
+ Task DeleteVectorStore(string storeName, CancellationToken token);
+}
diff --git a/app/MindWork AI Studio/Tools/Databases/VectorStore/NoVectorStoreClient.cs b/app/MindWork AI Studio/Tools/Databases/VectorStore/NoVectorStoreClient.cs
new file mode 100644
index 00000000..6f9eaf87
--- /dev/null
+++ b/app/MindWork AI Studio/Tools/Databases/VectorStore/NoVectorStoreClient.cs
@@ -0,0 +1,39 @@
+using AIStudio.Tools.PluginSystem;
+
+namespace AIStudio.Tools.Databases.VectorStore;
+
+public sealed class NoVectorStoreClient(string name, string? unavailableReason, DatabaseClientStatus status = DatabaseClientStatus.UNAVAILABLE) : IVectorStoreClient
+{
+ private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(NoVectorStoreClient).Namespace, nameof(NoVectorStoreClient));
+
+ public string Name => name;
+
+ public DatabaseClientStatus Status => status;
+
+ public bool IsAvailable => false;
+
+ public 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 Task EnsureVectorStoreExists(string storeName, int vectorSize, CancellationToken token) =>
+ Task.FromException(this.CreateUnavailableException());
+
+ public Task InsertEmbedding(string storeName, IReadOnlyList points, CancellationToken token) =>
+ Task.FromException(this.CreateUnavailableException());
+
+ public Task DeleteEmbeddingByFile(string storeName, string filePath, CancellationToken token) =>
+ Task.FromException(this.CreateUnavailableException());
+
+ public Task DeleteVectorStore(string storeName, CancellationToken token) =>
+ Task.FromException(this.CreateUnavailableException());
+
+ private InvalidOperationException CreateUnavailableException() =>
+ new(unavailableReason ?? "The vector store is not available.");
+}
diff --git a/app/MindWork AI Studio/Tools/Databases/Qdrant/QdrantClientImplementation.cs b/app/MindWork AI Studio/Tools/Databases/VectorStore/QdrantClientImplementation.cs
similarity index 88%
rename from app/MindWork AI Studio/Tools/Databases/Qdrant/QdrantClientImplementation.cs
rename to app/MindWork AI Studio/Tools/Databases/VectorStore/QdrantClientImplementation.cs
index e3265940..9a2768c8 100644
--- a/app/MindWork AI Studio/Tools/Databases/Qdrant/QdrantClientImplementation.cs
+++ b/app/MindWork AI Studio/Tools/Databases/VectorStore/QdrantClientImplementation.cs
@@ -4,9 +4,9 @@ using Grpc.Core;
using AIStudio.Tools.PluginSystem;
using static Qdrant.Client.Grpc.Conditions;
-namespace AIStudio.Tools.Databases.Qdrant;
+namespace AIStudio.Tools.Databases.VectorStore;
-public class QdrantClientImplementation : EmbeddingStore
+public class QdrantClientImplementation : DatabaseClient, IVectorStoreClient
{
private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(QdrantClientImplementation).Namespace, nameof(QdrantClientImplementation));
@@ -71,7 +71,7 @@ public class QdrantClientImplementation : EmbeddingStore
yield return (TB("Number of collections"), await this.GetCollectionsAmount());
}
- public override async Task EnsureEmbeddingStoreExists(string collectionName, int vectorSize, CancellationToken token)
+ public async Task EnsureVectorStoreExists(string collectionName, int vectorSize, CancellationToken token)
{
var exists = await this.GrpcClient.CollectionExistsAsync(collectionName, token);
if (exists)
@@ -87,7 +87,7 @@ public class QdrantClientImplementation : EmbeddingStore
cancellationToken: token);
}
- public override Task InsertEmbedding(string collectionName, IReadOnlyList points, CancellationToken token)
+ public Task InsertEmbedding(string collectionName, IReadOnlyList points, CancellationToken token)
{
var qdrantPoints = points.Select(point => new PointStruct
{
@@ -112,7 +112,7 @@ public class QdrantClientImplementation : EmbeddingStore
return this.GrpcClient.UpsertAsync(collectionName, qdrantPoints, true, null, null, token);
}
- public override async Task DeleteEmbeddingByFile(string collectionName, string filePath, CancellationToken token)
+ public async Task DeleteEmbeddingByFile(string collectionName, string filePath, CancellationToken token)
{
try
{
@@ -120,11 +120,10 @@ public class QdrantClientImplementation : EmbeddingStore
}
catch (RpcException exception) when (exception.StatusCode is StatusCode.NotFound)
{
- return;
}
}
- public override async Task DeleteEmbeddingStore(string collectionName, CancellationToken token)
+ public async Task DeleteVectorStore(string collectionName, CancellationToken token)
{
var exists = await this.GrpcClient.CollectionExistsAsync(collectionName, token);
if (!exists)
@@ -136,7 +135,6 @@ public class QdrantClientImplementation : EmbeddingStore
}
catch (RpcException exception) when (exception.StatusCode is StatusCode.NotFound)
{
- return;
}
}
diff --git a/app/MindWork AI Studio/Tools/Metadata/MetaDataEmbeddingStoreAttribute.cs b/app/MindWork AI Studio/Tools/Metadata/MetaDataEmbeddingStoreAttribute.cs
deleted file mode 100644
index 5ef6064b..00000000
--- a/app/MindWork AI Studio/Tools/Metadata/MetaDataEmbeddingStoreAttribute.cs
+++ /dev/null
@@ -1,6 +0,0 @@
-namespace AIStudio.Tools.Metadata;
-
-public class MetaDataDatabasesAttribute(string databaseVersion) : Attribute
-{
- public string DatabaseVersion => databaseVersion;
-}
\ No newline at end of file
diff --git a/app/MindWork AI Studio/Tools/Metadata/MetaDataVectorStoreAttribute.cs b/app/MindWork AI Studio/Tools/Metadata/MetaDataVectorStoreAttribute.cs
new file mode 100644
index 00000000..e3ba1b75
--- /dev/null
+++ b/app/MindWork AI Studio/Tools/Metadata/MetaDataVectorStoreAttribute.cs
@@ -0,0 +1,6 @@
+namespace AIStudio.Tools.Metadata;
+
+public class MetaDataVectorStoreAttribute(string vectorStoreVersion) : Attribute
+{
+ public string VectorStoreVersion => vectorStoreVersion;
+}
\ No newline at end of file
diff --git a/app/MindWork AI Studio/Tools/PluginSystem/PluginConfigurationObject.cs b/app/MindWork AI Studio/Tools/PluginSystem/PluginConfigurationObject.cs
index fa9b9610..90ce305f 100644
--- a/app/MindWork AI Studio/Tools/PluginSystem/PluginConfigurationObject.cs
+++ b/app/MindWork AI Studio/Tools/PluginSystem/PluginConfigurationObject.cs
@@ -438,7 +438,7 @@ public sealed record PluginConfigurationObject
return string.Empty;
}
- return storeResult.Message;
+ return storeResult.StoredPath;
}
private static string? ResolvePluginTokenizerPath(string configuredTokenizerPath, string pluginPath)
diff --git a/app/MindWork AI Studio/Tools/Rust/TokenizerResponse.cs b/app/MindWork AI Studio/Tools/Rust/TokenizerResponse.cs
index 54f0b61c..f3ef0893 100644
--- a/app/MindWork AI Studio/Tools/Rust/TokenizerResponse.cs
+++ b/app/MindWork AI Studio/Tools/Rust/TokenizerResponse.cs
@@ -1,3 +1,8 @@
namespace AIStudio.Tools.Rust;
-public readonly record struct TokenizerResponse(bool Success, int TokenCount, string Message);
+public readonly record struct TokenizerResponse(
+ bool Success,
+ int TokenCount,
+ string Message,
+ TokenizerStatus Status = TokenizerStatus.UNAVAILABLE,
+ string StoredPath = "");
diff --git a/app/MindWork AI Studio/Tools/Rust/TokenizerStatus.cs b/app/MindWork AI Studio/Tools/Rust/TokenizerStatus.cs
new file mode 100644
index 00000000..fd1b9dfc
--- /dev/null
+++ b/app/MindWork AI Studio/Tools/Rust/TokenizerStatus.cs
@@ -0,0 +1,8 @@
+namespace AIStudio.Tools.Rust;
+
+public enum TokenizerStatus
+{
+ UNAVAILABLE,
+ RUNNING,
+ AVAILABLE,
+}
diff --git a/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.Files.cs b/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.Files.cs
index 80ab1ef4..4b8f2c17 100644
--- a/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.Files.cs
+++ b/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.Files.cs
@@ -23,7 +23,7 @@ public sealed partial class DataSourceEmbeddingService
var currentChunk = new StringBuilder();
- await foreach (var segment in this.rustService.StreamArbitraryFileData(filePath, token: token))
+ await foreach (var segment in rustService.StreamArbitraryFileData(filePath, token: token))
{
var normalized = NormalizeChunkSegment(segment);
if (string.IsNullOrWhiteSpace(normalized))
@@ -118,7 +118,7 @@ public sealed partial class DataSourceEmbeddingService
}
catch (Exception exception)
{
- this.logger.LogWarning(exception, "Cannot access directory '{DirectoryPath}' while indexing.", currentPath);
+ logger.LogWarning(exception, "Cannot access directory '{DirectoryPath}' while indexing.", currentPath);
result.FailedFiles++;
result.LastError = $"The directory '{currentPath}' could not be accessed.";
continue;
@@ -135,7 +135,7 @@ public sealed partial class DataSourceEmbeddingService
}
catch (Exception exception)
{
- this.logger.LogWarning(exception, "Cannot inspect file '{FilePath}' while indexing.", filePath);
+ logger.LogWarning(exception, "Cannot inspect file '{FilePath}' while indexing.", filePath);
result.FailedFiles++;
result.LastError = $"The file '{filePath}' could not be inspected.";
continue;
diff --git a/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.State.cs b/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.State.cs
index b08bb3d4..3bc00949 100644
--- a/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.State.cs
+++ b/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.State.cs
@@ -1,6 +1,7 @@
using System.Text.Json;
using AIStudio.Settings;
+using AIStudio.Tools.Databases.VectorStore;
namespace AIStudio.Tools.Services;
@@ -70,13 +71,13 @@ public sealed partial class DataSourceEmbeddingService
await File.WriteAllTextAsync(statePath, json, token);
}
- private async Task ResetPersistedStateAsync(string dataSourceId)
+ private async Task ResetPersistedStateAsync(string dataSourceId, IVectorStoreClient? vectorStore, CancellationToken token)
{
- await this.EnsureStateLoadedAsync(CancellationToken.None);
+ await this.EnsureStateLoadedAsync(token);
this.manifests.Remove(dataSourceId);
- await this.DeleteCollectionAsync(this.GetCollectionName(dataSourceId));
- await this.SaveStateAsync(CancellationToken.None);
- this.logger.LogInformation("Reset persisted embedding state for data source '{DataSourceId}'.", dataSourceId);
+ await this.DeleteCollectionAsync(this.GetCollectionName(dataSourceId), vectorStore, token);
+ await this.SaveStateAsync(token);
+ logger.LogInformation("Reset persisted embedding state for data source '{DataSourceId}'.", dataSourceId);
}
private string GetStatePath()
diff --git a/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.Watchers.cs b/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.Watchers.cs
index 75cfa154..ad6302b3 100644
--- a/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.Watchers.cs
+++ b/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.Watchers.cs
@@ -14,13 +14,13 @@ public sealed partial class DataSourceEmbeddingService
private void RefreshWatchers()
{
- if (!this.settingsManager.ConfigurationData.DataSourceIndexing.AutomaticRefresh)
+ if (!settingsManager.ConfigurationData.DataSourceIndexing.AutomaticRefresh)
{
this.RemoveAllWatchers();
return;
}
- var supportedSources = this.settingsManager.ConfigurationData.DataSources
+ var supportedSources = settingsManager.ConfigurationData.DataSources
.Where(this.IsSupportedInternalDataSource)
.ToDictionary(source => source.Id, StringComparer.OrdinalIgnoreCase);
@@ -33,7 +33,7 @@ public sealed partial class DataSourceEmbeddingService
private void EnsureWatcher(IDataSource dataSource)
{
- if (!this.settingsManager.ConfigurationData.DataSourceIndexing.AutomaticRefresh)
+ if (!settingsManager.ConfigurationData.DataSourceIndexing.AutomaticRefresh)
return;
var configuration = GetWatchConfiguration(dataSource);
@@ -73,7 +73,7 @@ public sealed partial class DataSourceEmbeddingService
watcher.Renamed += (_, _) => this.OnWatchedDataSourceChanged(dataSourceId);
watcher.Error += (_, args) =>
{
- this.logger.LogWarning(args.GetException(), "The file watcher for data source '{DataSourceId}' failed. Recreating it.", dataSourceId);
+ logger.LogWarning(args.GetException(), "The file watcher for data source '{DataSourceId}' failed. Recreating it.", dataSourceId);
this.RemoveWatcher(dataSourceId);
this.EnsureWatcher(dataSourceId);
this.OnWatchedDataSourceChanged(dataSourceId);
@@ -83,7 +83,7 @@ public sealed partial class DataSourceEmbeddingService
}
catch (Exception exception)
{
- this.logger.LogWarning(exception, "Failed to create file watcher for data source '{DataSourceId}' at '{RootPath}'.", dataSourceId, configuration.RootPath);
+ logger.LogWarning(exception, "Failed to create file watcher for data source '{DataSourceId}' at '{RootPath}'.", dataSourceId, configuration.RootPath);
return null;
}
}
@@ -114,10 +114,10 @@ public sealed partial class DataSourceEmbeddingService
private void OnWatchedDataSourceChanged(string dataSourceId)
{
- if (!this.settingsManager.ConfigurationData.DataSourceIndexing.AutomaticRefresh)
+ if (!settingsManager.ConfigurationData.DataSourceIndexing.AutomaticRefresh)
return;
- this.logger.LogDebug("Detected file system change for data source '{DataSourceId}'. Scheduling a debounced embedding run.", dataSourceId);
+ logger.LogDebug("Detected file system change for data source '{DataSourceId}'. Scheduling a debounced embedding run.", dataSourceId);
var debounceToken = new CancellationTokenSource();
lock (this.watcherDebounceLock)
@@ -136,12 +136,12 @@ public sealed partial class DataSourceEmbeddingService
if (!this.TryCompletePendingWatcherRefresh(dataSourceId, debounceToken))
return;
- var dataSource = this.settingsManager.ConfigurationData.DataSources
+ var dataSource = settingsManager.ConfigurationData.DataSources
.FirstOrDefault(source => source.Id.Equals(dataSourceId, StringComparison.OrdinalIgnoreCase));
if (dataSource is not null)
{
- this.logger.LogInformation("Queueing data source '{DataSourceName}' ({DataSourceId}) after file system changes settled.", dataSource.Name, dataSource.Id);
+ logger.LogInformation("Queueing data source '{DataSourceName}' ({DataSourceId}) after file system changes settled.", dataSource.Name, dataSource.Id);
await this.QueueDataSourceAsync(dataSource);
}
}
@@ -150,7 +150,7 @@ public sealed partial class DataSourceEmbeddingService
}
catch (Exception exception)
{
- this.logger.LogWarning(exception, "Failed to queue watched data source '{DataSourceId}' after a file system change.", dataSourceId);
+ logger.LogWarning(exception, "Failed to queue watched data source '{DataSourceId}' after a file system change.", dataSourceId);
}
finally
{
@@ -161,7 +161,7 @@ public sealed partial class DataSourceEmbeddingService
private void EnsureWatcher(string dataSourceId)
{
- var dataSource = this.settingsManager.ConfigurationData.DataSources
+ var dataSource = settingsManager.ConfigurationData.DataSources
.FirstOrDefault(source => source.Id.Equals(dataSourceId, StringComparison.OrdinalIgnoreCase));
if (dataSource is not null)
diff --git a/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.cs b/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.cs
index c82de777..d4a1ee92 100644
--- a/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.cs
+++ b/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.cs
@@ -6,22 +6,19 @@ using AIStudio.Provider;
using AIStudio.Settings;
using AIStudio.Settings.DataModel;
using AIStudio.Tools.Databases;
+using AIStudio.Tools.Databases.VectorStore;
using AIStudio.Tools.PluginSystem;
-using AIStudio.Tools.Rust;
namespace AIStudio.Tools.Services;
-public sealed partial class DataSourceEmbeddingService : BackgroundService
+public sealed partial class DataSourceEmbeddingService(SettingsManager settingsManager, RustService rustService, DatabaseClientProvider databaseClientProvider, ILogger logger)
+ : BackgroundService
{
private const int MAX_CHUNK_LENGTH = 3_200;
private const int MIN_CHUNK_LENGTH = 800;
private const int CHUNK_OVERLAP_LENGTH = 320;
private const int EMBEDDING_BATCH_SIZE = 16;
- private readonly SettingsManager settingsManager;
- private readonly RustService rustService;
- private readonly EmbeddingStore embeddingStore;
- private readonly ILogger logger;
private readonly Channel queue = Channel.CreateUnbounded();
private readonly ConcurrentDictionary queuedIds = new(StringComparer.OrdinalIgnoreCase);
private readonly ConcurrentDictionary statuses = new(StringComparer.OrdinalIgnoreCase);
@@ -32,14 +29,6 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
private Dictionary manifests = new(StringComparer.OrdinalIgnoreCase);
private bool stateLoaded;
- public DataSourceEmbeddingService(SettingsManager settingsManager, RustService rustService, EmbeddingStore embeddingStore, ILogger logger)
- {
- this.settingsManager = settingsManager;
- this.rustService = rustService;
- this.embeddingStore = embeddingStore;
- this.logger = logger;
- }
-
public IReadOnlyList GetStatuses()
{
return this.statuses.Values
@@ -78,7 +67,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
{
this.RefreshWatchers();
- var tasks = this.settingsManager.ConfigurationData.DataSources
+ var tasks = settingsManager.ConfigurationData.DataSources
.Where(this.IsSupportedInternalDataSource)
.Select(this.QueueDataSourceAsync);
@@ -87,7 +76,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
public Task QueueAllInternalDataSourcesIfAutomaticRefreshAsync()
{
- if (!this.settingsManager.ConfigurationData.DataSourceIndexing.AutomaticRefresh)
+ if (!settingsManager.ConfigurationData.DataSourceIndexing.AutomaticRefresh)
{
this.RefreshWatchers();
return Task.CompletedTask;
@@ -106,16 +95,16 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
if (!this.IsSupportedInternalDataSource(dataSource))
return;
- this.logger.LogInformation("Queueing data source '{DataSourceName}' ({DataSourceId}) for background embeddings.", dataSource.Name, dataSource.Id);
+ logger.LogInformation("Queueing data source '{DataSourceName}' ({DataSourceId}) for background embeddings.", dataSource.Name, dataSource.Id);
this.RefreshWatchers();
- this.logger.LogDebug("Adding watcher for data source '{DataSourceName}' ({DataSourceId}).", dataSource.Name, dataSource.Id);
+ logger.LogDebug("Adding watcher for data source '{DataSourceName}' ({DataSourceId}).", dataSource.Name, dataSource.Id);
if (!this.statuses.TryGetValue(dataSource.Id, out var currentStatus) || currentStatus.State is not DataSourceEmbeddingState.RUNNING)
this.UpsertStatus(this.CreateStatus(dataSource, DataSourceEmbeddingState.QUEUED, currentStatus?.TotalFiles ?? 0, currentStatus?.IndexedFiles ?? 0, currentStatus?.FailedFiles ?? 0));
- this.logger.LogDebug("Upserting status for data source '{DataSourceName}' ({DataSourceId}).", dataSource.Name, dataSource.Id);
+ logger.LogDebug("Upserting status for data source '{DataSourceName}' ({DataSourceId}).", dataSource.Name, dataSource.Id);
if (this.queuedIds.TryAdd(dataSource.Id, 0))
await this.queue.Writer.WriteAsync(dataSource.Id);
- this.logger.LogDebug("Queued data source '{DataSourceName}' ({DataSourceId}).", dataSource.Name, dataSource.Id);
+ logger.LogDebug("Queued data source '{DataSourceName}' ({DataSourceId}).", dataSource.Name, dataSource.Id);
}
public async Task RemoveDataSourceAsync(IDataSource dataSource)
@@ -125,7 +114,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
this.RemoveWatcher(dataSource.Id);
this.statuses.TryRemove(dataSource.Id, out _);
- await this.ResetPersistedStateAsync(dataSource.Id);
+ await this.ResetPersistedStateAsync(dataSource.Id, null, CancellationToken.None);
this.PublishStatusChanged();
}
@@ -138,7 +127,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
var dataSourceId = await this.queue.Reader.ReadAsync(stoppingToken);
this.queuedIds.TryRemove(dataSourceId, out _);
- var dataSource = this.settingsManager.ConfigurationData.DataSources
+ var dataSource = settingsManager.ConfigurationData.DataSources
.FirstOrDefault(source => source.Id.Equals(dataSourceId, StringComparison.OrdinalIgnoreCase));
if (dataSource is null || !this.IsSupportedInternalDataSource(dataSource))
@@ -154,7 +143,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
}
catch (Exception exception)
{
- this.logger.LogError(exception, "Background embedding failed for data source '{DataSourceName}' ({DataSourceId}).", dataSource.Name, dataSource.Id);
+ logger.LogError(exception, "Background embedding failed for data source '{DataSourceName}' ({DataSourceId}).", dataSource.Name, dataSource.Id);
this.UpsertStatus(this.GetFallbackStatus(dataSource, exception.Message));
}
}
@@ -170,15 +159,17 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
private async Task ProcessDataSourceAsync(IDataSource dataSource, CancellationToken token)
{
await this.EnsureStateLoadedAsync(token);
- this.logger.LogInformation("Starting background embeddings for data source '{DataSourceName}' ({DataSourceId}).", dataSource.Name, dataSource.Id);
+ logger.LogInformation("Starting background embeddings for data source '{DataSourceName}' ({DataSourceId}).", dataSource.Name, dataSource.Id);
- if (!this.embeddingStore.IsAvailable)
+ var vectorStore = await databaseClientProvider.GetVectorStoreAsync(token);
+
+ if (!vectorStore.IsAvailable)
{
- this.logger.LogWarning(
+ logger.LogWarning(
"Skipping background embeddings for data source '{DataSourceName}' ({DataSourceId}) because the database client '{DatabaseName}' is unavailable.",
dataSource.Name,
dataSource.Id,
- this.embeddingStore.Name);
+ vectorStore.Name);
this.UpsertStatus(this.GetFallbackStatus(dataSource, "The vector database is not available."));
return;
}
@@ -189,7 +180,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
return;
}
- this.logger.LogInformation(
+ logger.LogInformation(
"Using embedding provider '{EmbeddingProviderId}' with model '{EmbeddingModelId}' for data source '{DataSourceName}' ({DataSourceId}).",
embeddingProvider.Id,
embeddingProvider.Model.Id,
@@ -197,12 +188,12 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
dataSource.Id);
var collectionName = this.GetCollectionName(dataSource.Id);
- var manifest = await this.EnsureCompatibleManifestAsync(dataSource, embeddingProvider, collectionName, token);
+ var manifest = await this.EnsureCompatibleManifestAsync(dataSource, embeddingProvider, collectionName, vectorStore, token);
var inputFiles = this.GetInputFiles(dataSource);
var indexedFiles = inputFiles.Files;
var totalFiles = indexedFiles.Count + inputFiles.FailedFiles;
- this.logger.LogInformation(
+ logger.LogInformation(
"Prepared data source '{DataSourceName}' ({DataSourceId}) for embedding. AccessibleFiles={AccessibleFiles}, FailedFiles={FailedFiles}, Collection='{CollectionName}'.",
dataSource.Name,
dataSource.Id,
@@ -210,7 +201,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
inputFiles.FailedFiles,
collectionName);
- await this.RemoveMissingFileEmbeddingsAsync(dataSource, collectionName, manifest, indexedFiles, token);
+ await this.RemoveMissingFileEmbeddingsAsync(vectorStore, dataSource, collectionName, manifest, indexedFiles, token);
await this.SaveStateAsync(token);
this.UpsertStatus(this.CreateStatus(
@@ -235,7 +226,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
if (manifest.Files.TryGetValue(file.FullName, out var existingRecord) &&
string.Equals(existingRecord.Fingerprint, fingerprint, StringComparison.Ordinal))
{
- this.logger.LogDebug(
+ logger.LogDebug(
"Skipping unchanged file '{FilePath}' for data source '{DataSourceName}' ({DataSourceId}).",
file.FullName,
dataSource.Name,
@@ -249,7 +240,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
try
{
- this.logger.LogInformation(
+ logger.LogInformation(
"Embedding file '{FilePath}' for data source '{DataSourceName}' ({DataSourceId}). Progress={CompletedFiles}/{TotalFiles}.",
file.FullName,
dataSource.Name,
@@ -257,7 +248,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
skippedFiles + completedFiles + 1,
totalFiles);
var startedAtUtc = DateTime.UtcNow;
- var chunkCount = await this.IndexOneFileAsync(dataSource, file, fingerprint, embeddingProvider, provider, manifest, token);
+ var chunkCount = await this.IndexOneFileAsync(vectorStore, dataSource, file, fingerprint, embeddingProvider, provider, manifest, token);
manifest.Files[file.FullName] = new EmbeddedFileRecord(
fingerprint,
file.Length,
@@ -266,7 +257,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
chunkCount);
await this.SaveStateAsync(token);
completedFiles++;
- this.logger.LogInformation(
+ logger.LogInformation(
"Embedded file '{FilePath}' for data source '{DataSourceName}' ({DataSourceId}) successfully. Chunks={ChunkCount}, DurationMs={DurationMs}.",
file.FullName,
dataSource.Name,
@@ -279,16 +270,16 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
failedFiles++;
lastError = exception.Message;
manifest.Files.Remove(file.FullName);
- await this.DeleteFilePointsAsync(collectionName, file.FullName, token);
+ await this.DeleteFilePointsAsync(vectorStore, collectionName, file.FullName, token);
await this.SaveStateAsync(token);
- this.logger.LogWarning(exception, "Failed to embed file '{FilePath}' for data source '{DataSourceName}'.", file.FullName, dataSource.Name);
+ logger.LogWarning(exception, "Failed to embed file '{FilePath}' for data source '{DataSourceName}'.", file.FullName, dataSource.Name);
this.UpsertStatus(this.CreateStatus(dataSource, DataSourceEmbeddingState.RUNNING, totalFiles, skippedFiles + completedFiles, failedFiles, file.Name, exception.Message));
}
}
this.UpsertStatus(this.CreateCompletedStatus(dataSource, totalFiles, skippedFiles + completedFiles, failedFiles, lastError));
- this.logger.LogInformation(
+ logger.LogInformation(
"Finished background embeddings for data source '{DataSourceName}' ({DataSourceId}). Indexed={IndexedFiles}, Failed={FailedFiles}, Total={TotalFiles}.",
dataSource.Name,
dataSource.Id,
@@ -298,6 +289,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
}
private async Task IndexOneFileAsync(
+ IVectorStoreClient vectorStore,
IDataSource dataSource,
FileInfo file,
string fingerprint,
@@ -307,11 +299,11 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
CancellationToken token)
{
var collectionName = this.GetCollectionName(dataSource.Id);
- this.logger.LogDebug(
+ logger.LogDebug(
"Resetting stored embeddings for file '{FilePath}' in collection '{CollectionName}' before re-indexing.",
file.FullName,
collectionName);
- await this.DeleteFilePointsAsync(collectionName, file.FullName, token);
+ await this.DeleteFilePointsAsync(vectorStore, collectionName, file.FullName, token);
var batch = new List<(string Text, int ChunkIndex)>(EMBEDDING_BATCH_SIZE);
var totalChunkCount = 0;
@@ -322,16 +314,16 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
totalChunkCount++;
if (batch.Count >= EMBEDDING_BATCH_SIZE)
- await this.FlushBatchAsync(dataSource, file, fingerprint, embeddingProvider, provider, manifest, collectionName, batch, token);
+ await this.FlushBatchAsync(vectorStore, dataSource, file, fingerprint, embeddingProvider, provider, manifest, collectionName, batch, token);
}
if (batch.Count > 0)
- await this.FlushBatchAsync(dataSource, file, fingerprint, embeddingProvider, provider, manifest, collectionName, batch, token);
+ await this.FlushBatchAsync(vectorStore, dataSource, file, fingerprint, embeddingProvider, provider, manifest, collectionName, batch, token);
if (totalChunkCount == 0)
throw new InvalidOperationException($"The file '{file.Name}' did not yield any text chunks.");
- this.logger.LogDebug(
+ logger.LogDebug(
"Generated {ChunkCount} chunks for file '{FilePath}' in data source '{DataSourceName}' ({DataSourceId}).",
totalChunkCount,
file.FullName,
@@ -342,6 +334,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
}
private async Task FlushBatchAsync(
+ IVectorStoreClient vectorStore,
IDataSource dataSource,
FileInfo file,
string fingerprint,
@@ -352,7 +345,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
List<(string Text, int ChunkIndex)> batch,
CancellationToken token)
{
- this.logger.LogDebug(
+ logger.LogDebug(
"Requesting embeddings for batch of {ChunkCount} chunks from file '{FilePath}' in data source '{DataSourceName}' ({DataSourceId}).",
batch.Count,
file.FullName,
@@ -360,7 +353,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
dataSource.Id);
var texts = batch.Select(item => item.Text).ToList();
- var vectors = await provider.EmbedTextAsync(embeddingProvider.Model, this.settingsManager, token, texts);
+ var vectors = await provider.EmbedTextAsync(embeddingProvider.Model, settingsManager, token, texts);
if (vectors.Count != batch.Count)
throw new InvalidOperationException($"The embedding provider returned {vectors.Count} vectors for {batch.Count} text chunks.");
@@ -374,9 +367,9 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
if (manifest.VectorSize == 0)
{
manifest.VectorSize = vectorSize;
- await this.EnsureCollectionExistsAsync(collectionName, vectorSize, token);
+ await this.EnsureCollectionExistsAsync(vectorStore, collectionName, vectorSize, token);
await this.SaveStateAsync(token);
- this.logger.LogInformation(
+ logger.LogInformation(
"Created embedding collection '{CollectionName}' with vector size {VectorSize} for data source '{DataSourceName}' ({DataSourceId}).",
collectionName,
vectorSize,
@@ -385,6 +378,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
}
await this.UpsertPointsAsync(
+ vectorStore,
collectionName,
dataSource,
file,
@@ -394,7 +388,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
this.TryGetRelativePath(dataSource, file),
token);
- this.logger.LogDebug(
+ logger.LogDebug(
"Stored {ChunkCount} embedded chunks for file '{FilePath}' in collection '{CollectionName}'.",
batch.Count,
file.FullName,
@@ -403,12 +397,13 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
batch.Clear();
}
- private async Task EnsureCollectionExistsAsync(string collectionName, int vectorSize, CancellationToken token)
+ private async Task EnsureCollectionExistsAsync(IVectorStoreClient vectorStore, string collectionName, int vectorSize, CancellationToken token)
{
- await this.embeddingStore.EnsureEmbeddingStoreExists(collectionName, vectorSize, token);
+ await vectorStore.EnsureVectorStoreExists(collectionName, vectorSize, token);
}
private async Task UpsertPointsAsync(
+ IVectorStoreClient vectorStore,
string collectionName,
IDataSource dataSource,
FileInfo file,
@@ -419,7 +414,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
CancellationToken token)
{
var embeddedAtUtc = DateTime.UtcNow;
- var points = batch.Select((item, index) => new EmbeddingStoragePoint(
+ var points = batch.Select((item, index) => new VectorStoragePoint(
this.CreatePointId(dataSource.Id, fingerprint, item.ChunkIndex),
vectors[index],
dataSource.Id,
@@ -434,24 +429,31 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
file.LastWriteTimeUtc,
embeddedAtUtc)).ToList();
- await this.embeddingStore.InsertEmbedding(collectionName, points, token);
+ await vectorStore.InsertEmbedding(collectionName, points, token);
}
- private async Task DeleteFilePointsAsync(string collectionName, string filePath, CancellationToken token)
+ private async Task DeleteFilePointsAsync(IVectorStoreClient vectorStore, string collectionName, string filePath, CancellationToken token)
{
- await this.embeddingStore.DeleteEmbeddingByFile(collectionName, filePath, token);
+ await vectorStore.DeleteEmbeddingByFile(collectionName, filePath, token);
}
- private async Task DeleteCollectionAsync(string collectionName)
+ private async Task DeleteCollectionAsync(string collectionName, IVectorStoreClient? vectorStore, CancellationToken token)
{
- await this.embeddingStore.DeleteEmbeddingStore(collectionName, CancellationToken.None);
+ vectorStore ??= await databaseClientProvider.GetVectorStoreAsync(token);
+ if (!vectorStore.IsAvailable)
+ {
+ logger.LogWarning("Could not delete embedding collection '{CollectionName}' because the vector store '{VectorStoreName}' is unavailable.", collectionName, vectorStore.Name);
+ return;
+ }
+
+ await vectorStore.DeleteVectorStore(collectionName, token);
}
private async Task WaitForInitialSettingsAndBootstrapAsync(CancellationToken token)
{
while (!token.IsCancellationRequested)
{
- if (this.settingsManager.HasCompletedInitialSettingsLoad
+ if (settingsManager.HasCompletedInitialSettingsLoad
&& !string.IsNullOrWhiteSpace(SettingsManager.ConfigDirectory)
&& !string.IsNullOrWhiteSpace(SettingsManager.DataDirectory))
{
@@ -463,7 +465,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
token.ThrowIfCancellationRequested();
- this.logger.LogInformation("Embedding background service is ready. Checking whether automatic data source refresh is enabled.");
+ logger.LogInformation("Embedding background service is ready. Checking whether automatic data source refresh is enabled.");
await this.QueueAllInternalDataSourcesIfAutomaticRefreshAsync();
}
@@ -474,26 +476,26 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
private bool TryResolveEmbeddingProvider(IDataSource dataSource, [NotNullWhen(true)] out EmbeddingProvider? embeddingProvider)
{
- embeddingProvider = this.settingsManager.ConfigurationData.EmbeddingProviders.FirstOrDefault(provider =>
+ embeddingProvider = settingsManager.ConfigurationData.EmbeddingProviders.FirstOrDefault(provider =>
dataSource is IInternalDataSource internalDataSource &&
provider.Id.Equals(internalDataSource.EmbeddingId, StringComparison.OrdinalIgnoreCase));
return embeddingProvider != default && embeddingProvider.UsedLLMProvider is not LLMProviders.NONE;
}
- private async Task EnsureCompatibleManifestAsync(IDataSource dataSource, EmbeddingProvider embeddingProvider, string collectionName, CancellationToken token)
+ private async Task EnsureCompatibleManifestAsync(IDataSource dataSource, EmbeddingProvider embeddingProvider, string collectionName, IVectorStoreClient vectorStore, CancellationToken token)
{
var embeddingSignature = this.BuildEmbeddingSignature(embeddingProvider);
var manifest = await this.GetManifestAsync(dataSource.Id, token);
if (!string.Equals(manifest.EmbeddingSignature, embeddingSignature, StringComparison.Ordinal))
{
- this.logger.LogInformation(
+ logger.LogInformation(
"Embedding configuration changed for data source '{DataSourceName}' ({DataSourceId}). Resetting persisted state and collection '{CollectionName}'.",
dataSource.Name,
dataSource.Id,
collectionName);
- await this.ResetPersistedStateAsync(dataSource.Id);
+ await this.ResetPersistedStateAsync(dataSource.Id, vectorStore, token);
manifest = await this.GetManifestAsync(dataSource.Id, token);
}
@@ -509,6 +511,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
}
private async Task RemoveMissingFileEmbeddingsAsync(
+ IVectorStoreClient vectorStore,
IDataSource dataSource,
string collectionName,
DataSourceEmbeddingManifest manifest,
@@ -521,9 +524,9 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
foreach (var removedFilePath in manifest.Files.Keys.Except(existingPaths, StringComparer.OrdinalIgnoreCase).ToList())
{
- await this.DeleteFilePointsAsync(collectionName, removedFilePath, token);
+ await this.DeleteFilePointsAsync(vectorStore, collectionName, removedFilePath, token);
manifest.Files.Remove(removedFilePath);
- this.logger.LogInformation(
+ logger.LogInformation(
"Removed stale embeddings for deleted file '{FilePath}' from data source '{DataSourceName}' ({DataSourceId}).",
removedFilePath,
dataSource.Name,
diff --git a/app/MindWork AI Studio/Tools/Services/RustService.Tokenizer.cs b/app/MindWork AI Studio/Tools/Services/RustService.Tokenizer.cs
index 76125dc7..c90b2c82 100644
--- a/app/MindWork AI Studio/Tools/Services/RustService.Tokenizer.cs
+++ b/app/MindWork AI Studio/Tools/Services/RustService.Tokenizer.cs
@@ -9,6 +9,31 @@ public sealed partial class RustService
private string currentTokenizerPath = string.Empty;
private bool hasInitializedTokenizer;
+ private static TokenizerResponse CreateUnavailableTokenizerResponse(string message) => new(
+ false,
+ 0,
+ message,
+ TokenizerStatus.UNAVAILABLE,
+ string.Empty);
+
+ public async Task GetTokenizerInfo(CancellationToken cancellationToken = default)
+ {
+ try
+ {
+ return await this.http.GetFromJsonAsync("/system/tokenizer/info", this.jsonRustSerializerOptions, cancellationToken);
+ }
+ catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
+ {
+ this.logger?.LogWarning("Fetching tokenizer info from Rust service was cancelled by caller.");
+ return CreateUnavailableTokenizerResponse("Operation cancelled by caller.");
+ }
+ catch (Exception e)
+ {
+ this.logger?.LogError(e, "Error while fetching tokenizer info from Rust service.");
+ return CreateUnavailableTokenizerResponse(e.Message);
+ }
+ }
+
public async Task ValidateTokenizer(string filePath)
{
var result = await this.http.PostAsJsonAsync("/tokenizer/validate", new {
@@ -18,15 +43,12 @@ public sealed partial class RustService
if (!result.IsSuccessStatusCode)
{
this.logger!.LogError($"Failed to validate the tokenizer '{result.StatusCode}'");
- return new TokenizerResponse
- {
- Success = false,
- Message = "An error occured while sending the path to the Rust framework for validation: "+result.StatusCode,
- TokenCount = 0
- };
+ return CreateUnavailableTokenizerResponse("An error occured while sending the path to the Rust framework for validation: "+result.StatusCode);
}
- return await result.Content.ReadFromJsonAsync(this.jsonRustSerializerOptions);
+ var response = await result.Content.ReadFromJsonAsync(this.jsonRustSerializerOptions);
+
+ return response;
}
public async Task StoreTokenizer(string modelId, string filePath)
@@ -40,11 +62,7 @@ public sealed partial class RustService
if (!result.IsSuccessStatusCode)
{
this.logger!.LogError($"Failed to store the tokenizer '{result.StatusCode}'");
- return new TokenizerResponse{
- Success = false,
- Message = "An error occured while sending the path to the Rust framework for storing: "+result.StatusCode,
- TokenCount = 0
- };
+ return CreateUnavailableTokenizerResponse("An error occured while sending the path to the Rust framework for storing: "+result.StatusCode);
}
return await result.Content.ReadFromJsonAsync(this.jsonRustSerializerOptions);
@@ -60,11 +78,7 @@ public sealed partial class RustService
if (!result.IsSuccessStatusCode)
{
this.logger!.LogError($"Failed to delete the tokenizer '{result.StatusCode}'");
- return new TokenizerResponse{
- Success = false,
- Message = "An error occured while sending the tokenizer delete request to the Rust framework: "+result.StatusCode,
- TokenCount = 0
- };
+ return CreateUnavailableTokenizerResponse("An error occured while sending the tokenizer delete request to the Rust framework: "+result.StatusCode);
}
return await result.Content.ReadFromJsonAsync(this.jsonRustSerializerOptions);
@@ -79,14 +93,15 @@ public sealed partial class RustService
if (!result.IsSuccessStatusCode)
{
this.logger!.LogError($"Failed to get the token count '{result.StatusCode}'");
- return new TokenizerResponse{
- Success = false,
- Message = "Error while getting token count from Rust service: "+result.StatusCode,
- TokenCount = 0
- };
+ this.hasInitializedTokenizer = false;
+ return CreateUnavailableTokenizerResponse("Error while getting token count from Rust service: "+result.StatusCode);
}
- return await result.Content.ReadFromJsonAsync(this.jsonRustSerializerOptions);
+ var response = await result.Content.ReadFromJsonAsync(this.jsonRustSerializerOptions);
+ if (response is not { Status: TokenizerStatus.AVAILABLE })
+ this.hasInitializedTokenizer = false;
+
+ return response;
}
public async Task SetTokenizer(string providerName, string path)
@@ -99,14 +114,15 @@ public sealed partial class RustService
if (!result.IsSuccessStatusCode)
{
this.logger!.LogError($"Failed to set the tokenizer '{result.StatusCode}'");
- return new TokenizerResponse{
- Success = false,
- Message = "An error occured while sending the path to the Rust framework for setting a tokenizer: "+result.StatusCode,
- TokenCount = 0
- };
+ this.hasInitializedTokenizer = false;
+ return CreateUnavailableTokenizerResponse("An error occured while sending the path to the Rust framework for setting a tokenizer: "+result.StatusCode);
}
- return await result.Content.ReadFromJsonAsync(this.jsonRustSerializerOptions);
+ var response = await result.Content.ReadFromJsonAsync(this.jsonRustSerializerOptions);
+ if (response is not { Success: true, Status: TokenizerStatus.AVAILABLE })
+ this.hasInitializedTokenizer = false;
+
+ return response;
}
public async Task EnsureTokenizer(string providerName, string path)
@@ -115,14 +131,19 @@ public sealed partial class RustService
try
{
if (this.hasInitializedTokenizer && this.currentTokenizerPath == path)
- return new TokenizerResponse(true, 0, "Success");
+ return new TokenizerResponse(true, 0, string.Empty, TokenizerStatus.AVAILABLE);
var response = await this.SetTokenizer(providerName, path);
- if (response is { Success: true })
+ if (response is { Success: true, Status: TokenizerStatus.AVAILABLE })
{
this.currentTokenizerPath = path;
this.hasInitializedTokenizer = true;
}
+ else
+ {
+ this.currentTokenizerPath = string.Empty;
+ this.hasInitializedTokenizer = false;
+ }
return response;
}
diff --git a/runtime/Cargo.toml b/runtime/Cargo.toml
index 5bb0bf50..8552913a 100644
--- a/runtime/Cargo.toml
+++ b/runtime/Cargo.toml
@@ -1,6 +1,6 @@
[package]
name = "mindwork-ai-studio"
-version = "26.4.1"
+version = "26.5.5"
edition = "2024"
description = "MindWork AI Studio"
authors = ["Thorsten Sommer"]
@@ -48,6 +48,7 @@ tempfile = "3.27.0"
strum_macros = "0.28.0"
sysinfo = "0.39.1"
bytes = "1.11.1"
+tokenizers = "0.23.1"
[target.'cfg(target_os = "windows")'.dependencies]
windows-registry = "0.6.1"
diff --git a/runtime/src/app_window.rs b/runtime/src/app_window.rs
index 22d57791..f03f2102 100644
--- a/runtime/src/app_window.rs
+++ b/runtime/src/app_window.rs
@@ -28,7 +28,7 @@ use crate::pdfium::PDFIUM_LIB_PATH;
use crate::qdrant::{start_qdrant_server, stop_qdrant_server};
#[cfg(debug_assertions)]
use crate::dotnet::create_startup_env_file;
-use crate::tokenizer::set_path_resolver;
+use crate::tokenizer::set_default_tokenizer_path;
/// The Tauri main window.
pub static MAIN_WINDOW: Lazy>> = Lazy::new(|| Mutex::new(None));
@@ -151,7 +151,7 @@ pub fn start_tauri() {
start_qdrant_server(app.handle().clone());
- set_path_resolver(app.path_resolver());
+ set_default_tokenizer_path(app.handle().clone());
info!(Source = "Bootloader Tauri"; "Reconfigure the file logger to use the app data directory {data_path:?}");
switch_to_file_logging(data_path).map_err(|e| error!("Failed to switch logging to file: {e}")).unwrap();
diff --git a/runtime/src/runtime_api.rs b/runtime/src/runtime_api.rs
index 89f6cec0..35b8da3b 100644
--- a/runtime/src/runtime_api.rs
+++ b/runtime/src/runtime_api.rs
@@ -33,6 +33,7 @@ pub fn start_runtime_api() {
.route("/system/dotnet/port", get(crate::dotnet::dotnet_port))
.route("/system/dotnet/ready", get(crate::dotnet::dotnet_ready))
.route("/system/qdrant/info", get(crate::qdrant::qdrant_port))
+ .route("/system/tokenizer/info", get(crate::tokenizer::tokenizer_info))
.route("/clipboard/set", post(crate::clipboard::set_clipboard))
.route("/events", get(crate::app_window::get_event_stream))
.route("/updates/check", get(crate::app_window::check_for_update))
@@ -56,6 +57,11 @@ pub fn start_runtime_api() {
.route("/retrieval/fs/extract", get(crate::file_data::extract_data))
.route("/log/paths", get(crate::log::get_log_paths))
.route("/log/event", post(crate::log::log_event))
+ .route("/tokenizer/count", post(crate::tokenizer::token_count))
+ .route("/tokenizer/validate", post(crate::tokenizer::validate_tokenizer))
+ .route("/tokenizer/store", post(crate::tokenizer::store_tokenizer))
+ .route("/tokenizer/delete", post(crate::tokenizer::delete_tokenizer))
+ .route("/tokenizer/set", post(crate::tokenizer::set_tokenizer))
.route("/shortcuts/register", post(crate::app_window::register_shortcut))
.route("/shortcuts/validate", post(crate::app_window::validate_shortcut))
.route("/shortcuts/suspend", post(crate::app_window::suspend_shortcuts))
diff --git a/runtime/src/tokenizer.rs b/runtime/src/tokenizer.rs
index b98ef474..bb7ca0c9 100644
--- a/runtime/src/tokenizer.rs
+++ b/runtime/src/tokenizer.rs
@@ -1,24 +1,43 @@
use std::fs;
use std::path::PathBuf;
-use std::sync::{OnceLock, RwLock};
-use log::warn;
-use rocket::post;
-use rocket::serde::json::Json;
-use rocket::serde::Serialize;
-use serde::Deserialize;
-use tauri::PathResolver;
-use tokenizers::Error;
-use tokenizers::tokenizer::{Tokenizer, Error as TokenizerError};
+use std::sync::{Mutex, OnceLock, RwLock};
+
+use axum::Json;
+use log::{error, warn};
+use once_cell::sync::Lazy;
+use serde::{Deserialize, Serialize};
+use tauri::path::BaseDirectory;
+use tauri::Manager;
+use tokenizers::tokenizer::Tokenizer;
+
use crate::api_token::APIToken;
use crate::environment::DATA_DIRECTORY;
-static TOKENIZER: OnceLock>> = OnceLock::new();
+const DEFAULT_TOKENIZER_RESOURCE_PATH: &str = "resources/tokenizers/tokenizer.json";
+const NO_TOKENIZER_LOADED_MESSAGE: &str = "Tokenizer must be set before counting tokens.";
-static TOKENIZER_PATH_RESOLVER: OnceLock = OnceLock::new();
+static TOKENIZER: OnceLock>> = OnceLock::new();
+static DEFAULT_TOKENIZER_PATH: OnceLock = OnceLock::new();
+static TOKENIZER_STATUS: Lazy> = Lazy::new(|| Mutex::new(TokenizerStatusInfo::default()));
+static TOKENIZER_OPERATION_LOCK: Lazy> = Lazy::new(|| Mutex::new(()));
+
+#[derive(Clone, Copy, Default, Serialize, PartialEq, Eq)]
+pub enum TokenizerStatus {
+ #[default]
+ Unavailable,
+ Running,
+ Available,
+}
+
+#[derive(Default)]
+struct TokenizerStatusInfo {
+ status: TokenizerStatus,
+ unavailable_reason: Option,
+}
#[derive(Deserialize)]
pub struct SetTokenText {
- pub text: String,
+ text: String,
}
#[derive(Clone, Deserialize)]
@@ -42,94 +61,205 @@ pub struct TokenizerResponse {
success: bool,
token_count: usize,
message: String,
+ status: TokenizerStatus,
+ stored_path: String,
}
-impl From> for TokenizerResponse {
- fn from(result: Result) -> Self {
- match result {
- Ok(count) => TokenizerResponse {
- success: true,
- token_count: count,
- message: "Success".to_string(),
- },
- Err(e) => TokenizerResponse {
- success: false,
- token_count: 0,
- message: e.to_string(),
- },
+impl TokenizerResponse {
+ fn available(token_count: usize) -> Self {
+ TokenizerResponse {
+ success: true,
+ token_count,
+ message: String::new(),
+ status: TokenizerStatus::Available,
+ stored_path: String::new(),
+ }
+ }
+
+ fn stored(stored_path: String) -> Self {
+ TokenizerResponse {
+ success: true,
+ token_count: 0,
+ message: String::new(),
+ status: TokenizerStatus::Available,
+ stored_path,
+ }
+ }
+
+ fn unavailable(reason: String) -> Self {
+ TokenizerResponse {
+ success: false,
+ token_count: 0,
+ message: reason,
+ status: TokenizerStatus::Unavailable,
+ stored_path: String::new(),
}
}
}
-pub fn set_path_resolver(path_resolver: PathResolver) {
- match TOKENIZER_PATH_RESOLVER.set(path_resolver) {
+pub fn set_default_tokenizer_path(app_handle: tauri::AppHandle) {
+ let tokenizer_path = match app_handle
+ .path()
+ .resolve(DEFAULT_TOKENIZER_RESOURCE_PATH, BaseDirectory::Resource)
+ {
+ Ok(path) => path,
+ Err(e) => {
+ let reason = format!("The default tokenizer file '{DEFAULT_TOKENIZER_RESOURCE_PATH}' could not be resolved: {e}");
+ error!(Source = "Tokenizer"; "{reason}");
+ set_tokenizer_unavailable(reason);
+ return;
+ }
+ };
+
+ if !tokenizer_path.is_file() {
+ let reason = format!("The default tokenizer file was not found: {}", tokenizer_path.display());
+ error!(Source = "Tokenizer"; "{reason}");
+ set_tokenizer_unavailable(reason);
+ return;
+ }
+
+ match DEFAULT_TOKENIZER_PATH.set(tokenizer_path) {
Ok(_) => (),
- Err(e) => warn!(Source = "Tokenizer"; "Could not set the path resolver: {:?}", e),
+ Err(e) => warn!(Source = "Tokenizer"; "Could not set the default tokenizer path: {:?}", e),
}
}
-fn tokenizer_state() -> &'static RwLock> {
- TOKENIZER.get_or_init(|| RwLock::new(None))
+pub async fn tokenizer_info(_token: APIToken) -> Json {
+ let status = TOKENIZER_STATUS.lock().unwrap();
+ match status.status {
+ TokenizerStatus::Available => Json(TokenizerResponse::available(0)),
+ TokenizerStatus::Running => Json(TokenizerResponse {
+ success: false,
+ token_count: 0,
+ message: String::new(),
+ status: TokenizerStatus::Running,
+ stored_path: String::new(),
+ }),
+ TokenizerStatus::Unavailable => Json(TokenizerResponse::unavailable(status.unavailable_reason.clone().unwrap_or_default())),
+ }
}
-pub fn handle_tokenizer_set(path: &str) -> Result<(), Error> {
- let tokenizer_path = if path.trim().is_empty() {
- let relative_source_path = String::from("resources/tokenizers/tokenizer.json");
- let path_resolver = TOKENIZER_PATH_RESOLVER
- .get()
- .ok_or_else(|| Error::from("Tokenizer path resolver is not initialized"))?;
- path_resolver
- .resolve_resource(relative_source_path)
- .ok_or_else(|| Error::from("Failed to resolve default tokenizer resource path"))?
- } else {
- PathBuf::from(path)
- };
+pub async fn token_count(_token: APIToken, req: Json) -> Json {
+ match get_token_count(&req.text) {
+ Ok(count) => Json(TokenizerResponse::available(count)),
+ Err(e) => Json(TokenizerResponse::unavailable(e)),
+ }
+}
- let tokenizer = Tokenizer::from_file(tokenizer_path)?;
- let mut tokenizer_guard = tokenizer_state()
- .write()
- .map_err(|_| Error::from("Tokenizer state lock is poisoned"))?;
- *tokenizer_guard = Some(tokenizer);
+pub async fn validate_tokenizer(_token: APIToken, payload: Json) -> Json {
+ match handle_tokenizer_validate(&PathBuf::from(payload.file_path.clone())) {
+ Ok(count) => Json(TokenizerResponse::available(count)),
+ Err(e) => Json(TokenizerResponse::unavailable(e)),
+ }
+}
+pub async fn store_tokenizer(_token: APIToken, payload: Json) -> Json {
+ match handle_tokenizer_store(&payload) {
+ Ok(dest_path) => Json(TokenizerResponse::stored(dest_path)),
+ Err(e) => Json(TokenizerResponse::unavailable(e.to_string())),
+ }
+}
+
+pub async fn delete_tokenizer(_token: APIToken, payload: Json) -> Json {
+ match handle_tokenizer_delete(&payload) {
+ Ok(_) => Json(TokenizerResponse::stored(String::new())),
+ Err(e) => Json(TokenizerResponse::unavailable(e.to_string())),
+ }
+}
+
+pub async fn set_tokenizer(_token: APIToken, payload: Json) -> Json {
+ match handle_tokenizer_set(&payload.file_path) {
+ Ok(_) => Json(TokenizerResponse::available(0)),
+ Err(e) => Json(TokenizerResponse::unavailable(e)),
+ }
+}
+
+pub fn handle_tokenizer_set(path: &str) -> Result<(), String> {
+ let _operation_guard = begin_tokenizer_operation()?;
+ set_tokenizer_running();
+
+ let tokenizer_path = resolve_tokenizer_path(path).map_err(|e| {
+ error!(Source = "Tokenizer"; "{e} Starting the app without a tokenizer.");
+ unavailable_with_status_update(&e)
+ })?;
+
+ let tokenizer = load_tokenizer_from_file(&tokenizer_path).map_err(|e| {
+ error!(Source = "Tokenizer"; "{e}");
+ unavailable_with_status_update(&e)
+ })?;
+
+ match tokenizer_state().write() {
+ Ok(mut tokenizer_guard) => *tokenizer_guard = Some(tokenizer),
+ Err(_) => return Err(unavailable_with_status_update("Tokenizer state lock is poisoned.")),
+ }
+
+ set_tokenizer_available();
Ok(())
}
-fn handle_tokenizer_validate(path: &PathBuf) -> Result {
- if !path.is_file() {
- return Err(TokenizerError::from(format!(
- "Tokenizer file was not found: {}",
- path.display()
- )));
+fn handle_tokenizer_validate(path: &PathBuf) -> Result {
+ let _operation_guard = begin_tokenizer_operation()?;
+ set_tokenizer_running();
+
+ let result = validate_tokenizer_file(path);
+ match tokenizer_state().read() {
+ Ok(tokenizer_guard) if tokenizer_guard.is_some() => set_tokenizer_available(),
+ Ok(_) => set_tokenizer_unavailable(NO_TOKENIZER_LOADED_MESSAGE.to_string()),
+ Err(_) => set_tokenizer_unavailable("Tokenizer state lock is poisoned.".to_string()),
}
- let tokenizer = Tokenizer::from_file(path).map_err(|e| {
- TokenizerError::from(format!(
- "Failed to load tokenizer from '{}': {}",
- path.display(),
- e
- ))
- })?;
+ result
+}
+pub fn get_token_count(text: &str) -> Result {
+ if text.trim().is_empty() {
+ return Ok(0);
+ }
+
+ let _operation_guard = begin_tokenizer_operation()?;
+ {
+ let status = TOKENIZER_STATUS.lock().unwrap();
+ if status.status != TokenizerStatus::Available {
+ return Err(status.unavailable_reason.clone().unwrap_or_else(|| NO_TOKENIZER_LOADED_MESSAGE.to_string()));
+ }
+ }
+
+ let tokenizer_guard = tokenizer_state()
+ .read()
+ .map_err(|_| unavailable_with_status_update("Tokenizer state lock is poisoned."))?;
+ let tokenizer = match tokenizer_guard.as_ref() {
+ Some(tokenizer) => tokenizer,
+ None => {
+ drop(tokenizer_guard);
+ return Err(unavailable_with_status_update("Tokenizer not initialized."));
+ }
+ };
+ let token_count = match tokenizer.encode(text, true) {
+ Ok(enc) => enc.len(),
+ Err(e) => {
+ let reason = format!("Failed to tokenize text: {e}");
+ drop(tokenizer_guard);
+ return Err(unavailable_with_status_update(&reason));
+ }
+ };
+
+ Ok(token_count)
+}
+
+fn validate_tokenizer_file(path: &PathBuf) -> Result {
+ let tokenizer = load_tokenizer_from_file(path)?;
let test_string = "Hello, world! This is a test string for tokenizer validation.";
-
- let encoding = tokenizer.encode(test_string, true).map_err(|e| {
- TokenizerError::from(format!(
- "Tokenizer failed to encode validation string: {}",
- e
- ))
- })?;
+ let encoding = tokenizer
+ .encode(test_string, true)
+ .map_err(|e| format!("Tokenizer failed to encode validation string: {e}"))?;
let token_count = encoding.len();
if token_count == 0 {
- return Err(TokenizerError::from(
- "Tokenizer produced 0 tokens for test string. The tokenizer is likely invalid or misconfigured.",
- ));
+ return Err("Tokenizer produced 0 tokens for test string. The tokenizer is likely invalid or misconfigured.".to_string());
}
if encoding.get_tokens().iter().any(|t| t.is_empty()) {
- return Err(TokenizerError::from(
- "Tokenizer produced empty tokens. The tokenizer is invalid.",
- ));
+ return Err("Tokenizer produced empty tokens. The tokenizer is invalid.".to_string());
}
Ok(token_count)
@@ -147,26 +277,25 @@ fn handle_tokenizer_store(payload: &TokenizerStorage) -> Result fs::remove_dir_all(model_path.clone())?,
- false => (),
+ if model_path.try_exists()? {
+ fs::remove_dir_all(&model_path)?;
}
if payload.file_path.trim().is_empty() {
- return Ok(String::from(""));
+ return Ok(String::new());
}
- fs::create_dir_all(model_path)?;
+ fs::create_dir_all(&model_path)?;
fs::copy(&source_path, &destination_path)?;
- Ok(destination_path.to_str().unwrap().to_string())
+ Ok(destination_path.to_string_lossy().to_string())
}
fn handle_tokenizer_delete(payload: &TokenizerDelete) -> Result<(), std::io::Error> {
@@ -189,74 +318,61 @@ fn handle_tokenizer_delete(payload: &TokenizerDelete) -> Result<(), std::io::Err
Ok(())
}
-pub fn get_token_count(text: &str) -> Result {
- if text.trim().is_empty() {
- return Err(TokenizerError::from("Input text is empty"));
+fn tokenizer_state() -> &'static RwLock> {
+ TOKENIZER.get_or_init(|| RwLock::new(None))
+}
+
+fn begin_tokenizer_operation() -> Result, String> {
+ TOKENIZER_OPERATION_LOCK
+ .lock()
+ .map_err(|_| unavailable_with_status_update("Tokenizer operation lock is poisoned."))
+}
+
+fn set_tokenizer_available() {
+ let mut status = TOKENIZER_STATUS.lock().unwrap();
+ status.status = TokenizerStatus::Available;
+ status.unavailable_reason = None;
+}
+
+fn set_tokenizer_running() {
+ let mut status = TOKENIZER_STATUS.lock().unwrap();
+ status.status = TokenizerStatus::Running;
+ status.unavailable_reason = None;
+}
+
+fn set_tokenizer_unavailable(reason: String) {
+ let mut status = TOKENIZER_STATUS.lock().unwrap();
+ status.status = TokenizerStatus::Unavailable;
+ status.unavailable_reason = Some(reason);
+}
+
+fn unavailable_with_status_update(reason: &str) -> String {
+ let reason = reason.to_string();
+ match tokenizer_state().write() {
+ Ok(mut tokenizer_guard) => *tokenizer_guard = None,
+ Err(_) => set_tokenizer_unavailable("Tokenizer state lock is poisoned.".to_string()),
}
- let tokenizer = tokenizer_state()
- .read()
- .map_err(|_| TokenizerError::from("Tokenizer state lock is poisoned"))?
- .clone()
- .ok_or_else(|| TokenizerError::from("Tokenizer not initialized"))?;
- let enc = tokenizer.encode(text, true)?;
- Ok(enc.len())
+ set_tokenizer_unavailable(reason.clone());
+ reason
}
-#[post("/tokenizer/count", data = "")]
-pub fn token_count(_token: APIToken, req: Json) -> Json {
- Json(get_token_count(&req.text).into())
-}
-
-#[post("/tokenizer/validate", data = "")]
-pub fn validate_tokenizer(_token: APIToken, payload: Json) -> Json {
- Json(handle_tokenizer_validate(&PathBuf::from(payload.file_path.clone())).into())
-}
-
-#[post("/tokenizer/store", data = "")]
-pub fn store_tokenizer(_token: APIToken, payload: Json) -> Json {
- match handle_tokenizer_store(&payload) {
- Ok(dest_path) => Json(TokenizerResponse {
- success: true,
- token_count: 0,
- message: dest_path,
- }),
- Err(e) => Json(TokenizerResponse {
- success: false,
- token_count: 0,
- message: e.to_string(),
- }),
+fn resolve_tokenizer_path(path: &str) -> Result {
+ if !path.trim().is_empty() {
+ return Ok(PathBuf::from(path));
}
+
+ DEFAULT_TOKENIZER_PATH
+ .get()
+ .cloned()
+ .ok_or_else(|| "Default tokenizer path is not initialized.".to_string())
}
-#[post("/tokenizer/delete", data = "")]
-pub fn delete_tokenizer(_token: APIToken, payload: Json) -> Json {
- match handle_tokenizer_delete(&payload) {
- Ok(_) => Json(TokenizerResponse {
- success: true,
- token_count: 0,
- message: "Success".to_string(),
- }),
- Err(e) => Json(TokenizerResponse {
- success: false,
- token_count: 0,
- message: e.to_string(),
- }),
+fn load_tokenizer_from_file(path: &PathBuf) -> Result {
+ if !path.is_file() {
+ return Err(format!("Tokenizer file was not found: {}", path.display()));
}
-}
-#[post("/tokenizer/set", data = "")]
-pub fn set_tokenizer(_token: APIToken, payload: Json) -> Json {
- match handle_tokenizer_set(&payload.file_path) {
- Ok(_) => Json(TokenizerResponse {
- success: true,
- token_count: 0,
- message: "Success".to_string(),
- }),
- Err(e) => Json(TokenizerResponse {
- success: false,
- token_count: 0,
- message: e.to_string(),
- }),
- }
+ Tokenizer::from_file(path)
+ .map_err(|e| format!("Failed to load tokenizer from '{}': {e}", path.display()))
}