improved stability of embedding chunks

This commit is contained in:
PaulKoudelka 2026-07-28 17:36:06 +02:00
parent 1bf9328fb9
commit 759c7f9d89
3 changed files with 37 additions and 3 deletions

View File

@ -25,6 +25,7 @@ public sealed record EmbeddingProvider(
int TokenLimit = 8_191) : ConfigurationBaseObject, ISecretId
{
public const int DEFAULT_TOKEN_LIMIT = 8_191;
private const int ASSUMED_TOKEN_LIMIT_SAFETY_PERCENT = 80;
private static readonly ILogger<EmbeddingProvider> LOGGER = Program.LOGGER_FACTORY.CreateLogger<EmbeddingProvider>();
@ -57,6 +58,16 @@ public sealed record EmbeddingProvider(
[JsonIgnore]
public int EffectiveTokenLimit => this.TokenLimit > 0 ? this.TokenLimit : DEFAULT_TOKEN_LIMIT;
[JsonIgnore]
public bool UsesAssumedTokenSizing => string.IsNullOrWhiteSpace(this.TokenizerPath)
|| this.TokenLimit <= 0
|| this.TokenLimit == DEFAULT_TOKEN_LIMIT;
[JsonIgnore]
public int EffectiveChunkTokenLimit => this.UsesAssumedTokenSizing
? Math.Max(1, (int)Math.Ceiling(this.EffectiveTokenLimit * ASSUMED_TOKEN_LIMIT_SAFETY_PERCENT / 100d) - 1)
: this.EffectiveTokenLimit;
#endregion
public static bool TryParseEmbeddingProviderTable(int idx, LuaTable table, Guid configPluginId, out ConfigurationBaseObject provider)

View File

@ -82,7 +82,7 @@ public sealed partial class DataSourceEmbeddingService
private async IAsyncEnumerable<string> SplitChunkByEmbeddingTokenLimitAsync(string chunk, EmbeddingProvider embeddingProvider, [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token)
{
var tokenLimit = embeddingProvider.EffectiveTokenLimit;
var tokenLimit = embeddingProvider.EffectiveChunkTokenLimit;
var tokenCount = await this.GetEmbeddingTokenCountAsync(embeddingProvider, chunk, token);
if (tokenCount <= tokenLimit)
{
@ -90,6 +90,15 @@ public sealed partial class DataSourceEmbeddingService
yield break;
}
if (embeddingProvider.UsesAssumedTokenSizing)
{
logger.LogDebug(
"Using conservative embedding chunk limit {ChunkTokenLimit} for provider '{EmbeddingProviderName}' because tokenizer or token limit sizing is assumed. ConfiguredTokenLimit={ConfiguredTokenLimit}.",
tokenLimit,
embeddingProvider.Name,
embeddingProvider.EffectiveTokenLimit);
}
logger.LogDebug(
"Splitting an embedding chunk for provider '{EmbeddingProviderName}' because it has {TokenCount} tokens and the configured limit is {TokenLimit}.",
embeddingProvider.Name,
@ -430,7 +439,8 @@ public sealed partial class DataSourceEmbeddingService
embeddingProvider.Host,
embeddingProvider.Hostname,
embeddingProvider.TokenizerPath,
embeddingProvider.EffectiveTokenLimit);
embeddingProvider.EffectiveTokenLimit,
embeddingProvider.EffectiveChunkTokenLimit);
}
private async Task<string> BuildFingerprintAsync(FileInfo file, CancellationToken token)

View File

@ -424,7 +424,20 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
dataSource.Id);
var texts = batch.Select(item => item.Text).ToList();
var vectors = await provider.EmbedTextAsync(embeddingProvider.Model, settingsManager, token, texts);
IReadOnlyList<IReadOnlyList<float>> vectors;
try
{
vectors = await provider.EmbedTextAsync(embeddingProvider.Model, settingsManager, token, texts);
}
catch (OperationCanceledException) when (token.IsCancellationRequested)
{
throw;
}
catch (Exception exception)
{
throw new InvalidOperationException($"The embedding provider failed to embed {batch.Count} chunk(s) for file '{file.Name}'. Provider message: {exception.Message}", exception);
}
if (vectors.Count != batch.Count)
throw new InvalidOperationException($"The embedding provider returned {vectors.Count} vectors for {batch.Count} text chunks.");