From ea87d79ca97435474b4bb33ba4d3ad1764fc2854 Mon Sep 17 00:00:00 2001 From: Paul Koudelka Date: Wed, 27 May 2026 20:02:43 +0200 Subject: [PATCH] current main is now successfully merged into branch --- .../Assistants/I18N/allTexts.lua | 52 ++- .../Components/ChatComponent.razor.cs | 6 +- .../Dialogs/EmbeddingProviderDialog.razor.cs | 2 +- .../Dialogs/ProviderDialog.razor.cs | 2 +- .../SettingsDialogDataSources.razor.cs | 2 + .../MindWork AI Studio.csproj | 6 +- .../Pages/Information.razor | 12 +- .../Pages/Information.razor.cs | 73 ++-- app/MindWork AI Studio/Program.cs | 14 +- .../Tools/Databases/DatabaseClient.cs | 6 +- .../DatabaseClientProvider.Qdrant.cs | 5 +- .../Tools/Databases/DatabaseClientProvider.cs | 94 +--- .../Tools/Databases/EmbeddingStoreFactory.cs | 33 -- .../Tools/Databases/NoDatabaseClient.cs | 4 +- ...gStoragePoint.cs => VectorStoragePoint.cs} | 2 +- .../VectorStore/IVectorStoreClient.cs | 20 + .../VectorStore/NoVectorStoreClient.cs | 39 ++ .../QdrantClientImplementation.cs | 14 +- .../MetaDataEmbeddingStoreAttribute.cs | 6 - .../Metadata/MetaDataVectorStoreAttribute.cs | 6 + .../PluginSystem/PluginConfigurationObject.cs | 2 +- .../Tools/Rust/TokenizerResponse.cs | 7 +- .../Tools/Rust/TokenizerStatus.cs | 8 + .../DataSourceEmbeddingService.Files.cs | 6 +- .../DataSourceEmbeddingService.State.cs | 11 +- .../DataSourceEmbeddingService.Watchers.cs | 22 +- .../Services/DataSourceEmbeddingService.cs | 131 +++--- .../Tools/Services/RustService.Tokenizer.cs | 83 ++-- runtime/Cargo.toml | 3 +- runtime/src/app_window.rs | 4 +- runtime/src/runtime_api.rs | 6 + runtime/src/tokenizer.rs | 404 +++++++++++------- 32 files changed, 607 insertions(+), 478 deletions(-) delete mode 100644 app/MindWork AI Studio/Tools/Databases/EmbeddingStoreFactory.cs rename app/MindWork AI Studio/Tools/Databases/{EmbeddingStoragePoint.cs => VectorStoragePoint.cs} (88%) create mode 100644 app/MindWork AI Studio/Tools/Databases/VectorStore/IVectorStoreClient.cs create mode 100644 app/MindWork AI Studio/Tools/Databases/VectorStore/NoVectorStoreClient.cs rename app/MindWork AI Studio/Tools/Databases/{Qdrant => VectorStore}/QdrantClientImplementation.cs (88%) delete mode 100644 app/MindWork AI Studio/Tools/Metadata/MetaDataEmbeddingStoreAttribute.cs create mode 100644 app/MindWork AI Studio/Tools/Metadata/MetaDataVectorStoreAttribute.cs create mode 100644 app/MindWork AI Studio/Tools/Rust/TokenizerStatus.cs diff --git a/app/MindWork AI Studio/Assistants/I18N/allTexts.lua b/app/MindWork AI Studio/Assistants/I18N/allTexts.lua index 26711122..763ac03a 100644 --- a/app/MindWork AI Studio/Assistants/I18N/allTexts.lua +++ b/app/MindWork AI Studio/Assistants/I18N/allTexts.lua @@ -4905,6 +4905,7 @@ UI_TEXT_CONTENT["AISTUDIO::DIALOGS::SETTINGS::SETTINGSDIALOGDATASOURCES::T146957 -- Refresh all UI_TEXT_CONTENT["AISTUDIO::DIALOGS::SETTINGS::SETTINGSDIALOGDATASOURCES::T1503082343"] = "Refresh all" + -- Kerberos/SSO ERI data sources cannot be exported yet. Please configure them manually in the configuration plugin. UI_TEXT_CONTENT["AISTUDIO::DIALOGS::SETTINGS::SETTINGSDIALOGDATASOURCES::T1577531115"] = "Kerberos/SSO ERI data sources cannot be exported yet. Please configure them manually in the configuration plugin." @@ -4968,11 +4969,12 @@ UI_TEXT_CONTENT["AISTUDIO::DIALOGS::SETTINGS::SETTINGSDIALOGDATASOURCES::T352566 -- No data sources configured yet. UI_TEXT_CONTENT["AISTUDIO::DIALOGS::SETTINGS::SETTINGSDIALOGDATASOURCES::T3549650120"] = "No data sources configured yet." --- Local data sources refresh when files change. -UI_TEXT_CONTENT["AISTUDIO::DIALOGS::SETTINGS::SETTINGSDIALOGDATASOURCES::T3687976654"] = "Local data sources refresh when files change." -- Export Access Token? UI_TEXT_CONTENT["AISTUDIO::DIALOGS::SETTINGS::SETTINGSDIALOGDATASOURCES::T3595669127"] = "Export Access Token?" +-- Local data sources refresh when files change. +UI_TEXT_CONTENT["AISTUDIO::DIALOGS::SETTINGS::SETTINGSDIALOGDATASOURCES::T3687976654"] = "Local data sources refresh when files change." + -- Export ERI Data Source UI_TEXT_CONTENT["AISTUDIO::DIALOGS::SETTINGS::SETTINGSDIALOGDATASOURCES::T3831281036"] = "Export ERI Data Source" @@ -4994,11 +4996,12 @@ UI_TEXT_CONTENT["AISTUDIO::DIALOGS::SETTINGS::SETTINGSDIALOGDATASOURCES::T750361 -- External Data (ERI-Server v1) UI_TEXT_CONTENT["AISTUDIO::DIALOGS::SETTINGS::SETTINGSDIALOGDATASOURCES::T774473996"] = "External Data (ERI-Server v1)" --- Local data sources refresh only when triggered manually. -UI_TEXT_CONTENT["AISTUDIO::DIALOGS::SETTINGS::SETTINGSDIALOGDATASOURCES::T854231603"] = "Local data sources refresh only when triggered manually." -- Cannot export this ERI data source because no authentication secret is configured. The issue was: {0} UI_TEXT_CONTENT["AISTUDIO::DIALOGS::SETTINGS::SETTINGSDIALOGDATASOURCES::T782820095"] = "Cannot export this ERI data source because no authentication secret is configured. The issue was: {0}" +-- Local data sources refresh only when triggered manually. +UI_TEXT_CONTENT["AISTUDIO::DIALOGS::SETTINGS::SETTINGSDIALOGDATASOURCES::T854231603"] = "Local data sources refresh only when triggered manually." + -- Local Directory UI_TEXT_CONTENT["AISTUDIO::DIALOGS::SETTINGS::SETTINGSDIALOGDATASOURCES::T926703547"] = "Local Directory" @@ -6130,6 +6133,9 @@ UI_TEXT_CONTENT["AISTUDIO::PAGES::INFORMATION::T1019424746"] = "Startup log file -- Browse AI Studio's source code on GitHub — we welcome your contributions. UI_TEXT_CONTENT["AISTUDIO::PAGES::INFORMATION::T1107156991"] = "Browse AI Studio's source code on GitHub — we welcome your contributions." +-- Vector store version +UI_TEXT_CONTENT["AISTUDIO::PAGES::INFORMATION::T1124039623"] = "Vector store version" + -- The Tokenizer library serves as the base framework for integrating the DeepSeek tokenizer. UI_TEXT_CONTENT["AISTUDIO::PAGES::INFORMATION::T1132433749"] = "The Tokenizer library serves as the base framework for integrating the DeepSeek tokenizer." @@ -6145,9 +6151,6 @@ UI_TEXT_CONTENT["AISTUDIO::PAGES::INFORMATION::T1290340974"] = "Unknown configur -- This library is used to read PDF files. This is necessary, e.g., for using PDFs as a data source for a chat. UI_TEXT_CONTENT["AISTUDIO::PAGES::INFORMATION::T1388816916"] = "This library is used to read PDF files. This is necessary, e.g., for using PDFs as a data source for a chat." --- Database version -UI_TEXT_CONTENT["AISTUDIO::PAGES::INFORMATION::T1420062548"] = "Database version" - -- This library is used to extend the MudBlazor library. It provides additional components that are not part of the MudBlazor library. UI_TEXT_CONTENT["AISTUDIO::PAGES::INFORMATION::T1421513382"] = "This library is used to extend the MudBlazor library. It provides additional components that are not part of the MudBlazor library." @@ -6295,6 +6298,9 @@ UI_TEXT_CONTENT["AISTUDIO::PAGES::INFORMATION::T2924964415"] = "AI Studio runs w -- Changelog UI_TEXT_CONTENT["AISTUDIO::PAGES::INFORMATION::T3017574265"] = "Changelog" +-- Vector store +UI_TEXT_CONTENT["AISTUDIO::PAGES::INFORMATION::T3046399223"] = "Vector store" + -- Enterprise configuration ID: UI_TEXT_CONTENT["AISTUDIO::PAGES::INFORMATION::T3092349641"] = "Enterprise configuration ID:" @@ -6370,9 +6376,6 @@ UI_TEXT_CONTENT["AISTUDIO::PAGES::INFORMATION::T3986423270"] = "Check Pandoc Ins -- Versions UI_TEXT_CONTENT["AISTUDIO::PAGES::INFORMATION::T4010195468"] = "Versions" --- Database -UI_TEXT_CONTENT["AISTUDIO::PAGES::INFORMATION::T4036243672"] = "Database" - -- This library is used by the Rust runtime to read the current user's username, e.g. when an organization-managed ERI server uses the OS username for authentication. UI_TEXT_CONTENT["AISTUDIO::PAGES::INFORMATION::T4060906280"] = "This library is used by the Rust runtime to read the current user's username, e.g. when an organization-managed ERI server uses the OS username for authentication." @@ -6396,6 +6399,7 @@ UI_TEXT_CONTENT["AISTUDIO::PAGES::INFORMATION::T585329785"] = "Used .NET SDK" -- We use the DeepSeek Tokenizer to estimate the number of tokens an input will generate. UI_TEXT_CONTENT["AISTUDIO::PAGES::INFORMATION::T591393704"] = "We use the DeepSeek Tokenizer to estimate the number of tokens an input will generate." + -- starting UI_TEXT_CONTENT["AISTUDIO::PAGES::INFORMATION::T594602073"] = "starting" @@ -7087,31 +7091,37 @@ UI_TEXT_CONTENT["AISTUDIO::TOOLS::CONFIDENCESCHEMESEXTENSIONS::T3893997203"] = " UI_TEXT_CONTENT["AISTUDIO::TOOLS::CONFIDENCESCHEMESEXTENSIONS::T4107860491"] = "Trust all LLM providers" -- Reason -UI_TEXT_CONTENT["AISTUDIO::TOOLS::DATABASES::NOEMBEDDINGSTORE::T1093747001"] = "Reason" - --- Starting -UI_TEXT_CONTENT["AISTUDIO::TOOLS::DATABASES::NODATABASECLIENT::T1233211769"] = "Starting" +UI_TEXT_CONTENT["AISTUDIO::TOOLS::DATABASES::NODATABASECLIENT::T1093747001"] = "Reason" -- Unavailable -UI_TEXT_CONTENT["AISTUDIO::TOOLS::DATABASES::NOEMBEDDINGSTORE::T3662391977"] = "Unavailable" +UI_TEXT_CONTENT["AISTUDIO::TOOLS::DATABASES::NODATABASECLIENT::T3662391977"] = "Unavailable" -- Status -UI_TEXT_CONTENT["AISTUDIO::TOOLS::DATABASES::NOEMBEDDINGSTORE::T6222351"] = "Status" +UI_TEXT_CONTENT["AISTUDIO::TOOLS::DATABASES::NODATABASECLIENT::T6222351"] = "Status" + +-- Reason +UI_TEXT_CONTENT["AISTUDIO::TOOLS::DATABASES::VECTORSTORE::NOVECTORSTORECLIENT::T1093747001"] = "Reason" + +-- Unavailable +UI_TEXT_CONTENT["AISTUDIO::TOOLS::DATABASES::VECTORSTORE::NOVECTORSTORECLIENT::T3662391977"] = "Unavailable" + +-- Status +UI_TEXT_CONTENT["AISTUDIO::TOOLS::DATABASES::VECTORSTORE::NOVECTORSTORECLIENT::T6222351"] = "Status" -- Storage size -UI_TEXT_CONTENT["AISTUDIO::TOOLS::DATABASES::QDRANT::QDRANTCLIENTIMPLEMENTATION::T1230141403"] = "Storage size" +UI_TEXT_CONTENT["AISTUDIO::TOOLS::DATABASES::VECTORSTORE::QDRANTCLIENTIMPLEMENTATION::T1230141403"] = "Storage size" -- HTTP port -UI_TEXT_CONTENT["AISTUDIO::TOOLS::DATABASES::QDRANT::QDRANTCLIENTIMPLEMENTATION::T1717573768"] = "HTTP port" +UI_TEXT_CONTENT["AISTUDIO::TOOLS::DATABASES::VECTORSTORE::QDRANTCLIENTIMPLEMENTATION::T1717573768"] = "HTTP port" -- Reported version -UI_TEXT_CONTENT["AISTUDIO::TOOLS::DATABASES::QDRANT::QDRANTCLIENTIMPLEMENTATION::T3556099842"] = "Reported version" +UI_TEXT_CONTENT["AISTUDIO::TOOLS::DATABASES::VECTORSTORE::QDRANTCLIENTIMPLEMENTATION::T3556099842"] = "Reported version" -- gRPC port -UI_TEXT_CONTENT["AISTUDIO::TOOLS::DATABASES::QDRANT::QDRANTCLIENTIMPLEMENTATION::T757840040"] = "gRPC port" +UI_TEXT_CONTENT["AISTUDIO::TOOLS::DATABASES::VECTORSTORE::QDRANTCLIENTIMPLEMENTATION::T757840040"] = "gRPC port" -- Number of collections -UI_TEXT_CONTENT["AISTUDIO::TOOLS::DATABASES::QDRANT::QDRANTCLIENTIMPLEMENTATION::T842647336"] = "Number of collections" +UI_TEXT_CONTENT["AISTUDIO::TOOLS::DATABASES::VECTORSTORE::QDRANTCLIENTIMPLEMENTATION::T842647336"] = "Number of collections" -- The related data is not allowed to be sent to any LLM provider. This means that this data source cannot be used at the moment. UI_TEXT_CONTENT["AISTUDIO::TOOLS::ERICLIENT::DATAMODEL::PROVIDERTYPEEXTENSIONS::T1555790630"] = "The related data is not allowed to be sent to any LLM provider. This means that this data source cannot be used at the moment." diff --git a/app/MindWork AI Studio/Components/ChatComponent.razor.cs b/app/MindWork AI Studio/Components/ChatComponent.razor.cs index 84d69add..9589ceb3 100644 --- a/app/MindWork AI Studio/Components/ChatComponent.razor.cs +++ b/app/MindWork AI Studio/Components/ChatComponent.razor.cs @@ -4,7 +4,7 @@ using AIStudio.Provider; using AIStudio.Settings; using AIStudio.Settings.DataModel; using AIStudio.Tools.AIJobs; - +using AIStudio.Tools.Services; using Microsoft.AspNetCore.Components; using Microsoft.AspNetCore.Components.Web; @@ -1051,7 +1051,7 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable return; if (!tokenizerResponse.Value.Success) { - this.Logger.LogWarning($"Failed to initialize the tokenizer for the provider: {tokenizerResponse.Value.Message}"); + this.Logger.LogWarning($"Failed to initialize the tokenizer for the provider: status='{tokenizerResponse.Value.Status}', reason='{tokenizerResponse.Value.Message}'"); return; } @@ -1060,7 +1060,7 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable return; if (!response.Value.Success) { - this.Logger.LogWarning($"Failed to calculate token count: {response.Value.Message}"); + this.Logger.LogWarning($"Failed to calculate token count: status='{response.Value.Status}', reason='{response.Value.Message}'"); return; } this.tokenCount = response.Value.TokenCount.ToString(); diff --git a/app/MindWork AI Studio/Dialogs/EmbeddingProviderDialog.razor.cs b/app/MindWork AI Studio/Dialogs/EmbeddingProviderDialog.razor.cs index 20eabb28..6a500323 100644 --- a/app/MindWork AI Studio/Dialogs/EmbeddingProviderDialog.razor.cs +++ b/app/MindWork AI Studio/Dialogs/EmbeddingProviderDialog.razor.cs @@ -249,7 +249,7 @@ public partial class EmbeddingProviderDialog : MSGComponentBase, ISecretId await this.form.Validate(); return; } - this.dataFilePath = response.Message; + this.dataFilePath = response.StoredPath; // Use the data model to store the provider. // We just return this data to the parent component: diff --git a/app/MindWork AI Studio/Dialogs/ProviderDialog.razor.cs b/app/MindWork AI Studio/Dialogs/ProviderDialog.razor.cs index d04d9a5c..993aabcf 100644 --- a/app/MindWork AI Studio/Dialogs/ProviderDialog.razor.cs +++ b/app/MindWork AI Studio/Dialogs/ProviderDialog.razor.cs @@ -276,7 +276,7 @@ public partial class ProviderDialog : MSGComponentBase, ISecretId await this.form.Validate(); return; } - this.dataFilePath = tokenizerResponse.Message; + this.dataFilePath = tokenizerResponse.StoredPath; // Use the data model to store the provider. // We just return this data to the parent component: diff --git a/app/MindWork AI Studio/Dialogs/Settings/SettingsDialogDataSources.razor.cs b/app/MindWork AI Studio/Dialogs/Settings/SettingsDialogDataSources.razor.cs index 83bca086..9353d299 100644 --- a/app/MindWork AI Studio/Dialogs/Settings/SettingsDialogDataSources.razor.cs +++ b/app/MindWork AI Studio/Dialogs/Settings/SettingsDialogDataSources.razor.cs @@ -2,6 +2,8 @@ using AIStudio.Settings; using AIStudio.Settings.DataModel; using AIStudio.Tools.ERIClient.DataModel; using AIStudio.Tools.PluginSystem; +using AIStudio.Tools.Services; +using Microsoft.AspNetCore.Components; namespace AIStudio.Dialogs.Settings; diff --git a/app/MindWork AI Studio/MindWork AI Studio.csproj b/app/MindWork AI Studio/MindWork AI Studio.csproj index 7cebafb9..5421ffc9 100644 --- a/app/MindWork AI Studio/MindWork AI Studio.csproj +++ b/app/MindWork AI Studio/MindWork AI Studio.csproj @@ -88,7 +88,7 @@ $([System.String]::Copy( $(Metadata) ).Split( ';' )[ 8 ]) $([System.String]::Copy( $(Metadata) ).Split( ';' )[ 9 ]) $([System.String]::Copy( $(Metadata) ).Split( ';' )[ 10 ]) - $([System.String]::Copy( $(Metadata) ).Split( ';' )[ 11 ]) + $([System.String]::Copy( $(Metadata) ).Split( ';' )[ 11 ]) true @@ -116,8 +116,8 @@ <_Parameter1>$(MetaPdfiumVersion) - - <_Parameter1>$(MetaQdrantVersion) + + <_Parameter1>$(MetaVectorStoreVersion) diff --git a/app/MindWork AI Studio/Pages/Information.razor b/app/MindWork AI Studio/Pages/Information.razor index 219e5896..ab2a7958 100644 --- a/app/MindWork AI Studio/Pages/Information.razor +++ b/app/MindWork AI Studio/Pages/Information.razor @@ -21,11 +21,11 @@ - @this.VersionDatabase + @this.VersionVectorStore - + - @foreach (var item in this.databaseDisplayInfo) + @foreach (var item in this.vectorStoreDisplayInfo) {
@@ -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())) }