MindWorkAI/AI-Studio
MindWorkAI/AI-Studio
1
0
Fork 0
mirror of https://github.com/MindWorkAI/AI-Studio.git synced 2026-08-11 19:52:10 +00:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

made the chunking process more robust

Browse Source
This commit is contained in:
PaulKoudelka 2026-07-28 16:37:07 +02:00
parent 95ca74809b
commit 1bf9328fb9
9 changed files with 239 additions and 37 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
app/MindWork AI Studio/Components/ChatComponent.razor.cs
View File

@ -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)

1
app/MindWork AI Studio/Components/Settings/SettingsPanelEmbeddings.razor.cs
View File

@ -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<EmbeddingProviderDialog>(T("Edit Embedding Provider"), dialogParameters, DialogOptions.FULLSCREEN);

12
app/MindWork AI Studio/Dialogs/EmbeddingProviderDialog.razor
View File

@ -128,6 +128,18 @@
<MudJustifiedText Typo="Typo.body1" Class="mb-3">
@T("For better embeddings and less storage usage, it's recommended to use a custom tokenizer to enable a more accurate token count.")
</MudJustifiedText>
<MudNumericField
T="int"
@bind-Value="@this.DataTokenLimit"
Label="@T("Token limit")"
Class="mb-3"
Min="1"
Immediate="@true"
Adornment="Adornment.Start"
AdornmentIcon="@Icons.Material.Filled.Numbers"
AdornmentColor="Color.Info"
Validation="@this.ValidateTokenLimit"
HelperText="@T("Maximum number of tokens sent to the embedding model per chunk.")"/>
<SelectFile
File="@this.dataFilePath"
FileChanged="@this.OnDataFilePathChanged"

12
app/MindWork AI Studio/Dialogs/EmbeddingProviderDialog.razor.cs
View File

@ -74,6 +74,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)

3
app/MindWork AI Studio/Plugins/configuration/plugin.lua
View File

@ -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"] = "<the model ID, e.g., nomic-embed-text>",
-- ["DisplayName"] = "<user-friendly name of the model>",

17
app/MindWork AI Studio/Settings/EmbeddingProvider.cs
View File

@ -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<EmbeddingProvider> LOGGER = Program.LOGGER_FACTORY.CreateLogger<EmbeddingProvider>();
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<int>(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)}}",

151
app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.Files.cs
View File

@ -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<string> StreamEmbeddingChunksAsync(string filePath, [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token)
private async IAsyncEnumerable<string> 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<string> 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<string> 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<int> FindLargestUnitCountWithinTokenLimitAsync(IReadOnlyList<string> 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<string> 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<int> 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<string> SplitTextIntoTokenUnits(string text)
{
var matches = Regex.Matches(text, @"\S+\s*", RegexOptions.CultureInvariant);
if (matches.Count == 0)
return [text];
return matches.Cast<Match>().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<string> BuildFingerprintAsync(FileInfo file, CancellationToken token)

2
app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.cs
View File

@ -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++;

63
app/MindWork AI Studio/Tools/Services/RustService.Tokenizer.cs
View File

@ -84,11 +84,33 @@ public sealed partial class RustService
return await result.Content.ReadFromJsonAsync<TokenizerResponse>(this.jsonRustSerializerOptions);
}
public async Task<TokenizerResponse?> GetTokenCount(string text)
public Task<TokenizerResponse?> GetTokenCount(string text)
{
return this.GetTokenCountCoreAsync(text);
}
public async Task<TokenizerResponse?> 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<TokenizerResponse?> 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<TokenizerResponse?> 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;
}
}