made the chunking process more robust

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

View File

@ -1121,16 +1121,7 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable
return; return;
} }
var tokenizerResponse = await this.RustService.EnsureTokenizer(this.Provider.InstanceName, this.Provider.TokenizerPath); var response = await this.RustService.GetTokenCount(this.Provider.InstanceName, this.Provider.TokenizerPath, this.inputField.Value);
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);
if (response is null) if (response is null)
return; return;
if (!response.Value.Success) if (!response.Value.Success)

View File

@ -80,6 +80,7 @@ public partial class SettingsPanelEmbeddings : SettingsPanelProviderBase
{ x => x.IsEditing, true }, { x => x.IsEditing, true },
{ x => x.DataHost, embeddingProvider.Host }, { x => x.DataHost, embeddingProvider.Host },
{ x => x.DataTokenizerPath, embeddingProvider.TokenizerPath }, { x => x.DataTokenizerPath, embeddingProvider.TokenizerPath },
{ x => x.DataTokenLimit, embeddingProvider.EffectiveTokenLimit },
}; };
var dialogReference = await this.DialogService.ShowAsync<EmbeddingProviderDialog>(T("Edit Embedding Provider"), dialogParameters, DialogOptions.FULLSCREEN); var dialogReference = await this.DialogService.ShowAsync<EmbeddingProviderDialog>(T("Edit Embedding Provider"), dialogParameters, DialogOptions.FULLSCREEN);

View File

@ -128,6 +128,18 @@
<MudJustifiedText Typo="Typo.body1" Class="mb-3"> <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.") @T("For better embeddings and less storage usage, it's recommended to use a custom tokenizer to enable a more accurate token count.")
</MudJustifiedText> </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 <SelectFile
File="@this.dataFilePath" File="@this.dataFilePath"
FileChanged="@this.OnDataFilePathChanged" FileChanged="@this.OnDataFilePathChanged"

View File

@ -74,6 +74,9 @@ public partial class EmbeddingProviderDialog : MSGComponentBase, ISecretId
[Parameter] [Parameter]
public string DataTokenizerPath { get; set; } = string.Empty; public string DataTokenizerPath { get; set; } = string.Empty;
[Parameter]
public int DataTokenLimit { get; set; } = EmbeddingProvider.DEFAULT_TOKEN_LIMIT;
[Inject] [Inject]
private RustService RustService { get; init; } = null!; private RustService RustService { get; init; } = null!;
@ -148,6 +151,7 @@ public partial class EmbeddingProviderDialog : MSGComponentBase, ISecretId
IsEnterpriseConfiguration = false, IsEnterpriseConfiguration = false,
EnterpriseConfigurationPluginId = Guid.Empty, EnterpriseConfigurationPluginId = Guid.Empty,
TokenizerPath = this.dataFilePath, TokenizerPath = this.dataFilePath,
TokenLimit = this.DataTokenLimit,
}; };
} }
@ -277,6 +281,14 @@ public partial class EmbeddingProviderDialog : MSGComponentBase, ISecretId
return null; 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 void Cancel() => this.MudDialog.Cancel();
private async Task OnAPIKeyChanged(string apiKey) private async Task OnAPIKeyChanged(string apiKey)

View File

@ -136,6 +136,9 @@ CONFIG["EMBEDDING_PROVIDERS"] = {}
-- -- Optional: tokenizer path for this provider relative to the plugin directory. -- -- Optional: tokenizer path for this provider relative to the plugin directory.
-- -- ["TokenizerPath"] = "", -- -- ["TokenizerPath"] = "",
-- --
-- -- Optional: maximum number of tokens per embedding chunk. If omitted, AI Studio uses its default.
-- -- ["TokenLimit"] = 8191,
--
-- ["Model"] = { -- ["Model"] = {
-- ["Id"] = "<the model ID, e.g., nomic-embed-text>", -- ["Id"] = "<the model ID, e.g., nomic-embed-text>",
-- ["DisplayName"] = "<user-friendly name of the model>", -- ["DisplayName"] = "<user-friendly name of the model>",

View File

@ -21,8 +21,11 @@ public sealed record EmbeddingProvider(
Guid EnterpriseConfigurationPluginId = default, Guid EnterpriseConfigurationPluginId = default,
string Hostname = "http://localhost:1234", string Hostname = "http://localhost:1234",
Host Host = Host.NONE, 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>(); private static readonly ILogger<EmbeddingProvider> LOGGER = Program.LOGGER_FACTORY.CreateLogger<EmbeddingProvider>();
public static readonly EmbeddingProvider NONE = new(); public static readonly EmbeddingProvider NONE = new();
@ -51,6 +54,9 @@ public sealed record EmbeddingProvider(
[JsonIgnore] [JsonIgnore]
public string SecretName => this.Name; public string SecretName => this.Name;
[JsonIgnore]
public int EffectiveTokenLimit => this.TokenLimit > 0 ? this.TokenLimit : DEFAULT_TOKEN_LIMIT;
#endregion #endregion
public static bool TryParseEmbeddingProviderTable(int idx, LuaTable table, Guid configPluginId, out ConfigurationBaseObject provider) public static bool TryParseEmbeddingProviderTable(int idx, LuaTable table, Guid configPluginId, out ConfigurationBaseObject provider)
@ -105,6 +111,13 @@ public sealed record EmbeddingProvider(
tokenizerPath = string.Empty; 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 provider = new EmbeddingProvider
{ {
Num = 0, // will be set later by the PluginConfigurationObject Num = 0, // will be set later by the PluginConfigurationObject
@ -118,6 +131,7 @@ public sealed record EmbeddingProvider(
Hostname = hostname, Hostname = hostname,
Host = host, Host = host,
TokenizerPath = tokenizerPath, TokenizerPath = tokenizerPath,
TokenLimit = tokenLimit,
}; };
// Handle encrypted API key if present: // Handle encrypted API key if present:
@ -192,6 +206,7 @@ public sealed record EmbeddingProvider(
["UsedLLMProvider"] = "{{this.UsedLLMProvider}}", ["UsedLLMProvider"] = "{{this.UsedLLMProvider}}",
["TokenizerPath"] = "{{this.TokenizerPath}}", ["TokenizerPath"] = "{{this.TokenizerPath}}",
["TokenLimit"] = {{this.EffectiveTokenLimit}},
["Host"] = "{{this.Host}}", ["Host"] = "{{this.Host}}",
["Hostname"] = "{{LuaTools.EscapeLuaString(this.Hostname)}}", ["Hostname"] = "{{LuaTools.EscapeLuaString(this.Hostname)}}",

View File

@ -1,5 +1,6 @@
using System.Security.Cryptography; using System.Security.Cryptography;
using System.Text; using System.Text;
using System.Text.RegularExpressions;
using AIStudio.Settings; using AIStudio.Settings;
using AIStudio.Settings.DataModel; using AIStudio.Settings.DataModel;
@ -24,11 +25,13 @@ public sealed partial class DataSourceEmbeddingService
UNSUPPORTED, 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)) 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; yield break;
} }
@ -46,7 +49,10 @@ public sealed partial class DataSourceEmbeddingService
{ {
var chunk = currentChunk.ToString().Trim(); var chunk = currentChunk.ToString().Trim();
if (!string.IsNullOrWhiteSpace(chunk)) 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 var overlap = chunk.Length > CHUNK_OVERLAP_LENGTH
? chunk[^CHUNK_OVERLAP_LENGTH..] ? chunk[^CHUNK_OVERLAP_LENGTH..]
@ -68,7 +74,141 @@ public sealed partial class DataSourceEmbeddingService
var finalChunk = currentChunk.ToString().Trim(); var finalChunk = currentChunk.ToString().Trim();
if (!string.IsNullOrWhiteSpace(finalChunk)) 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) private FileEnumerationResult GetInputFiles(IDataSource dataSource)
@ -289,7 +429,8 @@ public sealed partial class DataSourceEmbeddingService
embeddingProvider.Model.Id, embeddingProvider.Model.Id,
embeddingProvider.Host, embeddingProvider.Host,
embeddingProvider.Hostname, embeddingProvider.Hostname,
embeddingProvider.TokenizerPath); embeddingProvider.TokenizerPath,
embeddingProvider.EffectiveTokenLimit);
} }
private async Task<string> BuildFingerprintAsync(FileInfo file, CancellationToken token) private async Task<string> BuildFingerprintAsync(FileInfo file, CancellationToken token)

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 batch = new List<(string Text, int ChunkIndex)>(EMBEDDING_BATCH_SIZE);
var totalChunkCount = 0; 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)); batch.Add((chunk, totalChunkCount));
totalChunkCount++; totalChunkCount++;

View File

@ -84,11 +84,33 @@ public sealed partial class RustService
return await result.Content.ReadFromJsonAsync<TokenizerResponse>(this.jsonRustSerializerOptions); 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 { var result = await this.http.PostAsJsonAsync("/tokenizer/count", new {
text = text, text = text,
}, this.jsonRustSerializerOptions); }, this.jsonRustSerializerOptions, cancellationToken);
if (!result.IsSuccessStatusCode) if (!result.IsSuccessStatusCode)
{ {
@ -130,26 +152,31 @@ public sealed partial class RustService
await this.tokenizerLock.WaitAsync(); await this.tokenizerLock.WaitAsync();
try try
{ {
if (this.hasInitializedTokenizer && this.currentTokenizerPath == path) return await this.EnsureTokenizerCoreAsync(providerName, 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;
} }
finally finally
{ {
this.tokenizerLock.Release(); 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;
}
} }