From 9ab196c13c879069956f7f8578e31c51bb1a8a72 Mon Sep 17 00:00:00 2001 From: Paul Koudelka Date: Mon, 10 Aug 2026 19:42:39 +0200 Subject: [PATCH] ensured prompts fit embedding model --- .../DataSourceEmbeddingService.Files.cs | 13 ++--- .../DataSourceLocalRetrievalService.cs | 56 +++++++++++++++++++ .../Tools/Services/RustService.Tokenizer.cs | 2 + 3 files changed, 64 insertions(+), 7 deletions(-) diff --git a/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.Files.cs b/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.Files.cs index 52a50b65..ee4a2b07 100644 --- a/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.Files.cs +++ b/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.Files.cs @@ -15,7 +15,6 @@ public sealed partial class DataSourceEmbeddingService { private const string OFFICE_LOCK_FILE_PREFIX = "~$"; private const int DEFAULT_CHUNK_OVERLAP_TOKEN_LENGTH = 300; - private const int MAX_TOKEN_COUNT_REQUEST_TEXT_LENGTH = 200_000; private static readonly string[] RAG_DELIMITED_TABLE_FILE_EXTENSIONS = ["csv", "tsv"]; private static readonly string[] RAG_SPREADSHEET_FILE_EXTENSIONS = ["ods", "xlsm", "xlsb"]; @@ -101,7 +100,7 @@ public sealed partial class DataSourceEmbeddingService var tokenCount = estimatedTokenCount; var textWithOverlap = AddOverlapPrefix(text, requiredOverlapPrefix); - if (textWithOverlap.Length <= MAX_TOKEN_COUNT_REQUEST_TEXT_LENGTH && + if (textWithOverlap.Length <= RustService.MAX_TOKEN_COUNT_REQUEST_TEXT_LENGTH && (estimatedTokenCount is null || estimatedTokenCount <= options.MaxChunkTokenLength)) { tokenCount = await this.GetEmbeddingTokenCountAsync(embeddingProvider, textWithOverlap, token); @@ -240,7 +239,7 @@ public sealed partial class DataSourceEmbeddingService var candidateUnitCount = minimumCandidateUnitCount + (maximumCandidateUnitCount - minimumCandidateUnitCount) / 2; var candidateText = AddOverlapPrefix(string.Concat(units.Skip(startUnitIndex).Take(candidateUnitCount)).Trim(), overlapPrefix); - var candidateFits = candidateText.Length <= MAX_TOKEN_COUNT_REQUEST_TEXT_LENGTH && + var candidateFits = candidateText.Length <= RustService.MAX_TOKEN_COUNT_REQUEST_TEXT_LENGTH && await this.GetEmbeddingTokenCountAsync(embeddingProvider, candidateText, token) <= maxChunkTokenLength; if (candidateFits) { @@ -410,7 +409,7 @@ public sealed partial class DataSourceEmbeddingService yield break; var bestEndIndex = startIndex; - var maximumCandidateEndIndex = Math.Min(text.Length, startIndex + MAX_TOKEN_COUNT_REQUEST_TEXT_LENGTH); + var maximumCandidateEndIndex = Math.Min(text.Length, startIndex + RustService.MAX_TOKEN_COUNT_REQUEST_TEXT_LENGTH); if (estimatedTokenCount > options.MaxChunkTokenLength) { @@ -426,7 +425,7 @@ public sealed partial class DataSourceEmbeddingService { var candidateEndIndex = minimumCandidateEndIndex + (currentMaximumCandidateEndIndex - minimumCandidateEndIndex) / 2; var candidate = AddOverlapPrefix(text[startIndex..candidateEndIndex].Trim(), overlapPrefix); - var candidateFits = candidate.Length <= MAX_TOKEN_COUNT_REQUEST_TEXT_LENGTH && + var candidateFits = candidate.Length <= RustService.MAX_TOKEN_COUNT_REQUEST_TEXT_LENGTH && await this.GetEmbeddingTokenCountAsync(embeddingProvider, candidate, token) <= options.MaxChunkTokenLength; if (candidateFits) { @@ -438,12 +437,12 @@ public sealed partial class DataSourceEmbeddingService } if (bestEndIndex < maximumCandidateEndIndex || bestEndIndex >= text.Length || - maximumCandidateEndIndex - startIndex >= MAX_TOKEN_COUNT_REQUEST_TEXT_LENGTH) + maximumCandidateEndIndex - startIndex >= RustService.MAX_TOKEN_COUNT_REQUEST_TEXT_LENGTH) break; var previousCandidateLength = maximumCandidateEndIndex - startIndex; maximumCandidateEndIndex = (int)Math.Min( - Math.Min(text.Length, startIndex + (long)MAX_TOKEN_COUNT_REQUEST_TEXT_LENGTH), + Math.Min(text.Length, startIndex + (long)RustService.MAX_TOKEN_COUNT_REQUEST_TEXT_LENGTH), startIndex + Math.Max(previousCandidateLength + 1L, previousCandidateLength * 2L)); } diff --git a/app/MindWork AI Studio/Tools/Services/DataSourceLocalRetrievalService.cs b/app/MindWork AI Studio/Tools/Services/DataSourceLocalRetrievalService.cs index 569b48ec..e9dbfab7 100644 --- a/app/MindWork AI Studio/Tools/Services/DataSourceLocalRetrievalService.cs +++ b/app/MindWork AI Studio/Tools/Services/DataSourceLocalRetrievalService.cs @@ -6,11 +6,13 @@ using AIStudio.Tools.Databases; using AIStudio.Tools.Databases.EmbeddingState; using AIStudio.Tools.Databases.VectorStore; using AIStudio.Tools.RAG; +using AIStudio.Tools.Rust; namespace AIStudio.Tools.Services; public sealed class DataSourceLocalRetrievalService( SettingsManager settingsManager, + RustService rustService, DatabaseClientProvider databaseClientProvider, ILogger logger) { @@ -107,6 +109,9 @@ public sealed class DataSourceLocalRetrievalService( return []; } + if (!await this.QueryFitsEmbeddingProviderAsync(dataSource, embeddingProvider, query, token)) + return []; + var provider = embeddingProvider.CreateProvider(); var vectors = await provider.EmbedTextAsync(embeddingProvider.Model, settingsManager, token, [query]); token.ThrowIfCancellationRequested(); @@ -136,6 +141,57 @@ public sealed class DataSourceLocalRetrievalService( } } + private async Task QueryFitsEmbeddingProviderAsync( + IInternalDataSource dataSource, + EmbeddingProvider embeddingProvider, + string query, + CancellationToken token) + { + var providerTokenLimit = Math.Max(1, embeddingProvider.EffectiveTokenLimit); + if (query.Length > RustService.MAX_TOKEN_COUNT_REQUEST_TEXT_LENGTH) + { + logger.LogWarning( + "Skipping vector retrieval for data source '{DataSourceName}' ({DataSourceId}) because the latest prompt has {CharacterCount} characters and exceeds the safe tokenizer request length of {MaxCharacterCount}. ProviderTokenLimit={ProviderTokenLimit}.", + dataSource.Name, + dataSource.Id, + query.Length, + RustService.MAX_TOKEN_COUNT_REQUEST_TEXT_LENGTH, + providerTokenLimit); + return false; + } + + var tokenCountResponse = await rustService.GetTokenCount( + embeddingProvider.Name, + embeddingProvider.TokenizerPath, + query, + token); + if (tokenCountResponse is not { Success: true, Status: TokenizerStatus.AVAILABLE }) + { + logger.LogWarning( + "Skipping vector retrieval for data source '{DataSourceName}' ({DataSourceId}) because the token count for embedding provider '{EmbeddingProviderName}' could not be determined. Reason='{Reason}'.", + dataSource.Name, + dataSource.Id, + embeddingProvider.Name, + tokenCountResponse?.Message ?? "No response was returned by the tokenizer service."); + return false; + } + + var queryTokenCount = tokenCountResponse.Value.TokenCount; + if (queryTokenCount > providerTokenLimit) + { + logger.LogWarning( + "Skipping vector retrieval for data source '{DataSourceName}' ({DataSourceId}) because the latest prompt has {QueryTokenCount} tokens, exceeding embedding provider '{EmbeddingProviderName}' limit of {ProviderTokenLimit} tokens.", + dataSource.Name, + dataSource.Id, + queryTokenCount, + embeddingProvider.Name, + providerTokenLimit); + return false; + } + + return true; + } + private async Task> SearchBm25Async(IInternalDataSource dataSource, string query, int maxMatches, CancellationToken token) { try diff --git a/app/MindWork AI Studio/Tools/Services/RustService.Tokenizer.cs b/app/MindWork AI Studio/Tools/Services/RustService.Tokenizer.cs index e044ae1e..66da3766 100644 --- a/app/MindWork AI Studio/Tools/Services/RustService.Tokenizer.cs +++ b/app/MindWork AI Studio/Tools/Services/RustService.Tokenizer.cs @@ -5,6 +5,8 @@ namespace AIStudio.Tools.Services; public sealed partial class RustService { + internal const int MAX_TOKEN_COUNT_REQUEST_TEXT_LENGTH = 200_000; + private readonly SemaphoreSlim tokenizerLock = new(1, 1); private string currentTokenizerPath = string.Empty; private bool hasInitializedTokenizer;