ensured prompts fit embedding model

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));
             }
 

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<DataSourceLocalRetrievalService> 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<bool> 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<IReadOnlyList<EmbeddingStateSearchResult>> SearchBm25Async(IInternalDataSource dataSource, string query, int maxMatches, CancellationToken token)
     {
         try

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;