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

ensured prompts fit embedding model

Browse Source
This commit is contained in:
Paul Koudelka 2026-08-10 19:42:39 +02:00
parent 3918b48409
commit 9ab196c13c
3 changed files with 64 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

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

56
app/MindWork AI Studio/Tools/Services/DataSourceLocalRetrievalService.cs
View File

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

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

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