diff --git a/app/MindWork AI Studio/Components/ChatComponent.razor.cs b/app/MindWork AI Studio/Components/ChatComponent.razor.cs index 5d5b9b70..7f4a3103 100644 --- a/app/MindWork AI Studio/Components/ChatComponent.razor.cs +++ b/app/MindWork AI Studio/Components/ChatComponent.razor.cs @@ -1121,16 +1121,7 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable return; } - var tokenizerResponse = await this.RustService.EnsureTokenizer(this.Provider.InstanceName, this.Provider.TokenizerPath); - if (tokenizerResponse is null) - return; - if (!tokenizerResponse.Value.Success) - { - this.Logger.LogWarning($"Failed to initialize the tokenizer for the provider: status='{tokenizerResponse.Value.Status}', reason='{tokenizerResponse.Value.Message}'"); - return; - } - - var response = await this.RustService.GetTokenCount(this.inputField.Value); + var response = await this.RustService.GetTokenCount(this.Provider.InstanceName, this.Provider.TokenizerPath, this.inputField.Value); if (response is null) return; if (!response.Value.Success) @@ -1224,4 +1215,4 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable } #endregion -} \ No newline at end of file +} diff --git a/app/MindWork AI Studio/Components/Settings/SettingsPanelEmbeddings.razor.cs b/app/MindWork AI Studio/Components/Settings/SettingsPanelEmbeddings.razor.cs index 90d0f24d..f4a67de4 100644 --- a/app/MindWork AI Studio/Components/Settings/SettingsPanelEmbeddings.razor.cs +++ b/app/MindWork AI Studio/Components/Settings/SettingsPanelEmbeddings.razor.cs @@ -80,6 +80,7 @@ public partial class SettingsPanelEmbeddings : SettingsPanelProviderBase { x => x.IsEditing, true }, { x => x.DataHost, embeddingProvider.Host }, { x => x.DataTokenizerPath, embeddingProvider.TokenizerPath }, + { x => x.DataTokenLimit, embeddingProvider.EffectiveTokenLimit }, }; var dialogReference = await this.DialogService.ShowAsync(T("Edit Embedding Provider"), dialogParameters, DialogOptions.FULLSCREEN); diff --git a/app/MindWork AI Studio/Dialogs/EmbeddingProviderDialog.razor b/app/MindWork AI Studio/Dialogs/EmbeddingProviderDialog.razor index 22d17e61..f9d3f500 100644 --- a/app/MindWork AI Studio/Dialogs/EmbeddingProviderDialog.razor +++ b/app/MindWork AI Studio/Dialogs/EmbeddingProviderDialog.razor @@ -128,6 +128,18 @@ @T("For better embeddings and less storage usage, it's recommended to use a custom tokenizer to enable a more accurate token count.") + - \ No newline at end of file + diff --git a/app/MindWork AI Studio/Dialogs/EmbeddingProviderDialog.razor.cs b/app/MindWork AI Studio/Dialogs/EmbeddingProviderDialog.razor.cs index a0ddc902..7907b45e 100644 --- a/app/MindWork AI Studio/Dialogs/EmbeddingProviderDialog.razor.cs +++ b/app/MindWork AI Studio/Dialogs/EmbeddingProviderDialog.razor.cs @@ -73,6 +73,9 @@ public partial class EmbeddingProviderDialog : MSGComponentBase, ISecretId [Parameter] public string DataTokenizerPath { get; set; } = string.Empty; + + [Parameter] + public int DataTokenLimit { get; set; } = EmbeddingProvider.DEFAULT_TOKEN_LIMIT; [Inject] private RustService RustService { get; init; } = null!; @@ -148,6 +151,7 @@ public partial class EmbeddingProviderDialog : MSGComponentBase, ISecretId IsEnterpriseConfiguration = false, EnterpriseConfigurationPluginId = Guid.Empty, TokenizerPath = this.dataFilePath, + TokenLimit = this.DataTokenLimit, }; } @@ -277,6 +281,14 @@ public partial class EmbeddingProviderDialog : MSGComponentBase, ISecretId return null; } + private string? ValidateTokenLimit(int tokenLimit) + { + if (tokenLimit < 1) + return T("Please enter a token limit greater than 0."); + + return null; + } + private void Cancel() => this.MudDialog.Cancel(); private async Task OnAPIKeyChanged(string apiKey) diff --git a/app/MindWork AI Studio/Plugins/configuration/plugin.lua b/app/MindWork AI Studio/Plugins/configuration/plugin.lua index 5dbe1f93..c7d9547d 100644 --- a/app/MindWork AI Studio/Plugins/configuration/plugin.lua +++ b/app/MindWork AI Studio/Plugins/configuration/plugin.lua @@ -136,6 +136,9 @@ CONFIG["EMBEDDING_PROVIDERS"] = {} -- -- Optional: tokenizer path for this provider relative to the plugin directory. -- -- ["TokenizerPath"] = "", -- +-- -- Optional: maximum number of tokens per embedding chunk. If omitted, AI Studio uses its default. +-- -- ["TokenLimit"] = 8191, +-- -- ["Model"] = { -- ["Id"] = "", -- ["DisplayName"] = "", diff --git a/app/MindWork AI Studio/Settings/EmbeddingProvider.cs b/app/MindWork AI Studio/Settings/EmbeddingProvider.cs index 3c02c07b..5b68898c 100644 --- a/app/MindWork AI Studio/Settings/EmbeddingProvider.cs +++ b/app/MindWork AI Studio/Settings/EmbeddingProvider.cs @@ -21,8 +21,11 @@ public sealed record EmbeddingProvider( Guid EnterpriseConfigurationPluginId = default, string Hostname = "http://localhost:1234", Host Host = Host.NONE, - string TokenizerPath = "") : ConfigurationBaseObject, ISecretId + string TokenizerPath = "", + int TokenLimit = 8_191) : ConfigurationBaseObject, ISecretId { + public const int DEFAULT_TOKEN_LIMIT = 8_191; + private static readonly ILogger LOGGER = Program.LOGGER_FACTORY.CreateLogger(); public static readonly EmbeddingProvider NONE = new(); @@ -51,6 +54,9 @@ public sealed record EmbeddingProvider( [JsonIgnore] public string SecretName => this.Name; + [JsonIgnore] + public int EffectiveTokenLimit => this.TokenLimit > 0 ? this.TokenLimit : DEFAULT_TOKEN_LIMIT; + #endregion public static bool TryParseEmbeddingProviderTable(int idx, LuaTable table, Guid configPluginId, out ConfigurationBaseObject provider) @@ -105,6 +111,13 @@ public sealed record EmbeddingProvider( tokenizerPath = string.Empty; } + var tokenLimit = DEFAULT_TOKEN_LIMIT; + if (table.TryGetValue("TokenLimit", out var tokenLimitValue) && (!tokenLimitValue.TryRead(out tokenLimit) || tokenLimit < 1)) + { + LOGGER.LogWarning($"The configured embedding provider {idx} does not contain a valid token limit. Falling back to {DEFAULT_TOKEN_LIMIT}. (Plugin ID: {configPluginId})"); + tokenLimit = DEFAULT_TOKEN_LIMIT; + } + provider = new EmbeddingProvider { Num = 0, // will be set later by the PluginConfigurationObject @@ -118,6 +131,7 @@ public sealed record EmbeddingProvider( Hostname = hostname, Host = host, TokenizerPath = tokenizerPath, + TokenLimit = tokenLimit, }; // Handle encrypted API key if present: @@ -192,6 +206,7 @@ public sealed record EmbeddingProvider( ["UsedLLMProvider"] = "{{this.UsedLLMProvider}}", ["TokenizerPath"] = "{{this.TokenizerPath}}", + ["TokenLimit"] = {{this.EffectiveTokenLimit}}, ["Host"] = "{{this.Host}}", ["Hostname"] = "{{LuaTools.EscapeLuaString(this.Hostname)}}", diff --git a/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.Files.cs b/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.Files.cs index 5c0d6de5..673eee24 100644 --- a/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.Files.cs +++ b/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.Files.cs @@ -1,5 +1,6 @@ using System.Security.Cryptography; using System.Text; +using System.Text.RegularExpressions; using AIStudio.Settings; using AIStudio.Settings.DataModel; @@ -24,11 +25,13 @@ public sealed partial class DataSourceEmbeddingService UNSUPPORTED, } - private async IAsyncEnumerable StreamEmbeddingChunksAsync(string filePath, [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token) + private async IAsyncEnumerable StreamEmbeddingChunksAsync(string filePath, EmbeddingProvider embeddingProvider, [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token) { if (this.IsImageFilePath(filePath)) { - yield return this.BuildImageIndexText(filePath); + await foreach (var imageChunk in this.SplitChunkByEmbeddingTokenLimitAsync(this.BuildImageIndexText(filePath), embeddingProvider, token)) + yield return imageChunk; + yield break; } @@ -46,7 +49,10 @@ public sealed partial class DataSourceEmbeddingService { var chunk = currentChunk.ToString().Trim(); if (!string.IsNullOrWhiteSpace(chunk)) - yield return chunk; + { + await foreach (var finalChunk in this.SplitChunkByEmbeddingTokenLimitAsync(chunk, embeddingProvider, token)) + yield return finalChunk; + } var overlap = chunk.Length > CHUNK_OVERLAP_LENGTH ? chunk[^CHUNK_OVERLAP_LENGTH..] @@ -68,7 +74,141 @@ public sealed partial class DataSourceEmbeddingService var finalChunk = currentChunk.ToString().Trim(); if (!string.IsNullOrWhiteSpace(finalChunk)) - yield return finalChunk; + { + await foreach (var chunk in this.SplitChunkByEmbeddingTokenLimitAsync(finalChunk, embeddingProvider, token)) + yield return chunk; + } + } + + private async IAsyncEnumerable SplitChunkByEmbeddingTokenLimitAsync(string chunk, EmbeddingProvider embeddingProvider, [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token) + { + var tokenLimit = embeddingProvider.EffectiveTokenLimit; + var tokenCount = await this.GetEmbeddingTokenCountAsync(embeddingProvider, chunk, token); + if (tokenCount <= tokenLimit) + { + yield return chunk; + yield break; + } + + logger.LogDebug( + "Splitting an embedding chunk for provider '{EmbeddingProviderName}' because it has {TokenCount} tokens and the configured limit is {TokenLimit}.", + embeddingProvider.Name, + tokenCount, + tokenLimit); + + await foreach (var splitChunk in this.SplitTextByTokenLimitAsync(chunk, embeddingProvider, tokenLimit, token)) + yield return splitChunk; + } + + private async IAsyncEnumerable SplitTextByTokenLimitAsync(string text, EmbeddingProvider embeddingProvider, int tokenLimit, [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token) + { + var units = SplitTextIntoTokenUnits(text); + var index = 0; + + while (index < units.Count) + { + token.ThrowIfCancellationRequested(); + + var unitCount = await this.FindLargestUnitCountWithinTokenLimitAsync(units, index, embeddingProvider, tokenLimit, token); + if (unitCount > 0) + { + var chunk = string.Concat(units.Skip(index).Take(unitCount)).Trim(); + if (!string.IsNullOrWhiteSpace(chunk)) + yield return chunk; + + index += unitCount; + continue; + } + + await foreach (var splitUnit in this.SplitOversizedTextUnitByTokenLimitAsync(units[index], embeddingProvider, tokenLimit, token)) + yield return splitUnit; + + index++; + } + } + + private async Task FindLargestUnitCountWithinTokenLimitAsync(IReadOnlyList units, int startIndex, EmbeddingProvider embeddingProvider, int tokenLimit, CancellationToken token) + { + var low = 1; + var high = units.Count - startIndex; + var best = 0; + + while (low <= high) + { + token.ThrowIfCancellationRequested(); + + var mid = low + (high - low) / 2; + var candidate = string.Concat(units.Skip(startIndex).Take(mid)).Trim(); + var tokenCount = await this.GetEmbeddingTokenCountAsync(embeddingProvider, candidate, token); + if (tokenCount <= tokenLimit) + { + best = mid; + low = mid + 1; + } + else + high = mid - 1; + } + + return best; + } + + private async IAsyncEnumerable SplitOversizedTextUnitByTokenLimitAsync(string text, EmbeddingProvider embeddingProvider, int tokenLimit, [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token) + { + var startIndex = 0; + while (startIndex < text.Length) + { + token.ThrowIfCancellationRequested(); + + var low = startIndex + 1; + var high = text.Length; + var bestEndIndex = startIndex; + + while (low <= high) + { + var mid = low + (high - low) / 2; + var candidate = text[startIndex..mid].Trim(); + var tokenCount = await this.GetEmbeddingTokenCountAsync(embeddingProvider, candidate, token); + if (tokenCount <= tokenLimit) + { + bestEndIndex = mid; + low = mid + 1; + } + else + high = mid - 1; + } + + if (bestEndIndex == startIndex) + { + var smallestCandidate = text[startIndex..Math.Min(startIndex + 1, text.Length)].Trim(); + var smallestCandidateTokenCount = await this.GetEmbeddingTokenCountAsync(embeddingProvider, smallestCandidate, token); + throw new InvalidOperationException($"The token limit for embedding provider '{embeddingProvider.Name}' is too low. The smallest possible split still has {smallestCandidateTokenCount} tokens, but the configured limit is {tokenLimit}."); + } + + var chunk = text[startIndex..bestEndIndex].Trim(); + if (!string.IsNullOrWhiteSpace(chunk)) + yield return chunk; + + startIndex = bestEndIndex; + } + } + + private async Task GetEmbeddingTokenCountAsync(EmbeddingProvider embeddingProvider, string text, CancellationToken token) + { + var response = await rustService.GetTokenCount(embeddingProvider.Name, embeddingProvider.TokenizerPath, text, token); + if (response is { Success: true, Status: TokenizerStatus.AVAILABLE }) + return response.Value.TokenCount; + + var message = response?.Message ?? "No response was returned by the tokenizer service."; + throw new InvalidOperationException($"Could not count tokens for embedding provider '{embeddingProvider.Name}'. {message}"); + } + + private static List SplitTextIntoTokenUnits(string text) + { + var matches = Regex.Matches(text, @"\S+\s*", RegexOptions.CultureInvariant); + if (matches.Count == 0) + return [text]; + + return matches.Cast().Select(match => match.Value).ToList(); } private FileEnumerationResult GetInputFiles(IDataSource dataSource) @@ -289,7 +429,8 @@ public sealed partial class DataSourceEmbeddingService embeddingProvider.Model.Id, embeddingProvider.Host, embeddingProvider.Hostname, - embeddingProvider.TokenizerPath); + embeddingProvider.TokenizerPath, + embeddingProvider.EffectiveTokenLimit); } private async Task BuildFingerprintAsync(FileInfo file, CancellationToken token) diff --git a/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.cs b/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.cs index f5434e44..b0299a85 100644 --- a/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.cs +++ b/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.cs @@ -378,7 +378,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM var batch = new List<(string Text, int ChunkIndex)>(EMBEDDING_BATCH_SIZE); var totalChunkCount = 0; - await foreach (var chunk in this.StreamEmbeddingChunksAsync(file.FullName, token)) + await foreach (var chunk in this.StreamEmbeddingChunksAsync(file.FullName, embeddingProvider, token)) { batch.Add((chunk, totalChunkCount)); totalChunkCount++; diff --git a/app/MindWork AI Studio/Tools/Services/RustService.Tokenizer.cs b/app/MindWork AI Studio/Tools/Services/RustService.Tokenizer.cs index c90b2c82..e044ae1e 100644 --- a/app/MindWork AI Studio/Tools/Services/RustService.Tokenizer.cs +++ b/app/MindWork AI Studio/Tools/Services/RustService.Tokenizer.cs @@ -84,11 +84,33 @@ public sealed partial class RustService return await result.Content.ReadFromJsonAsync(this.jsonRustSerializerOptions); } - public async Task GetTokenCount(string text) + public Task GetTokenCount(string text) + { + return this.GetTokenCountCoreAsync(text); + } + + public async Task GetTokenCount(string providerName, string path, string text, CancellationToken cancellationToken = default) + { + await this.tokenizerLock.WaitAsync(cancellationToken); + try + { + var tokenizerResponse = await this.EnsureTokenizerCoreAsync(providerName, path); + if (tokenizerResponse is not { Success: true, Status: TokenizerStatus.AVAILABLE }) + return tokenizerResponse; + + return await this.GetTokenCountCoreAsync(text, cancellationToken); + } + finally + { + this.tokenizerLock.Release(); + } + } + + private async Task GetTokenCountCoreAsync(string text, CancellationToken cancellationToken = default) { var result = await this.http.PostAsJsonAsync("/tokenizer/count", new { text = text, - }, this.jsonRustSerializerOptions); + }, this.jsonRustSerializerOptions, cancellationToken); if (!result.IsSuccessStatusCode) { @@ -130,26 +152,31 @@ public sealed partial class RustService await this.tokenizerLock.WaitAsync(); try { - if (this.hasInitializedTokenizer && this.currentTokenizerPath == path) - return new TokenizerResponse(true, 0, string.Empty, TokenizerStatus.AVAILABLE); - - var response = await this.SetTokenizer(providerName, path); - if (response is { Success: true, Status: TokenizerStatus.AVAILABLE }) - { - this.currentTokenizerPath = path; - this.hasInitializedTokenizer = true; - } - else - { - this.currentTokenizerPath = string.Empty; - this.hasInitializedTokenizer = false; - } - - return response; + return await this.EnsureTokenizerCoreAsync(providerName, path); } finally { this.tokenizerLock.Release(); } } + + private async Task EnsureTokenizerCoreAsync(string providerName, string path) + { + if (this.hasInitializedTokenizer && this.currentTokenizerPath == path) + return new TokenizerResponse(true, 0, string.Empty, TokenizerStatus.AVAILABLE); + + var response = await this.SetTokenizer(providerName, path); + if (response is { Success: true, Status: TokenizerStatus.AVAILABLE }) + { + this.currentTokenizerPath = path; + this.hasInitializedTokenizer = true; + } + else + { + this.currentTokenizerPath = string.Empty; + this.hasInitializedTokenizer = false; + } + + return response; + } }