mirror of
https://github.com/MindWorkAI/AI-Studio.git
synced 2026-08-11 16:32:10 +00:00
fixed tokenization issue for long files
This commit is contained in:
Paul Koudelka
parent
185a99c9ed
commit
3918b48409
@ -3112,9 +3112,6 @@ UI_TEXT_CONTENT["AISTUDIO::COMPONENTS::CHATCOMPONENT::T2991985411"] = "Delete th
|
||||
-- Move Chat to Workspace
|
||||
UI_TEXT_CONTENT["AISTUDIO::COMPONENTS::CHATCOMPONENT::T3045856778"] = "Move Chat to Workspace"
|
||||
|
||||
-- The selected provider is not allowed in this chat due to data security reasons.
|
||||
UI_TEXT_CONTENT["AISTUDIO::COMPONENTS::CHATCOMPONENT::T3403290862"] = "The selected provider is not allowed in this chat due to data security reasons."
|
||||
|
||||
-- Select a provider first
|
||||
UI_TEXT_CONTENT["AISTUDIO::COMPONENTS::CHATCOMPONENT::T3654197869"] = "Select a provider first"
|
||||
|
||||
@ -3124,6 +3121,9 @@ UI_TEXT_CONTENT["AISTUDIO::COMPONENTS::CHATCOMPONENT::T377990776"] = "Estimated
|
||||
-- Start new chat in workspace '{0}'
|
||||
UI_TEXT_CONTENT["AISTUDIO::COMPONENTS::CHATCOMPONENT::T3928697643"] = "Start new chat in workspace '{0}'"
|
||||
|
||||
-- The selected provider is not allowed in this chat due to data security or compliance reasons.
|
||||
UI_TEXT_CONTENT["AISTUDIO::COMPONENTS::CHATCOMPONENT::T393827158"] = "The selected provider is not allowed in this chat due to data security or compliance reasons."
|
||||
|
||||
-- Start temporary chat
|
||||
UI_TEXT_CONTENT["AISTUDIO::COMPONENTS::CHATCOMPONENT::T4113970938"] = "Start temporary chat"
|
||||
|
||||
@ -3247,6 +3247,9 @@ UI_TEXT_CONTENT["AISTUDIO::COMPONENTS::DATASOURCECLOUDEMBEDDINGWARNING::T9168792
|
||||
-- Yes, let the AI decide which data sources are needed.
|
||||
UI_TEXT_CONTENT["AISTUDIO::COMPONENTS::DATASOURCESELECTION::T1031370894"] = "Yes, let the AI decide which data sources are needed."
|
||||
|
||||
-- Your data sources cannot be used with the LLM provider you selected due to data privacy or compliance requirements, or they are currently unavailable.
|
||||
UI_TEXT_CONTENT["AISTUDIO::COMPONENTS::DATASOURCESELECTION::T1200866624"] = "Your data sources cannot be used with the LLM provider you selected due to data privacy or compliance requirements, or they are currently unavailable."
|
||||
|
||||
-- Yes, let the AI validate & filter the retrieved data.
|
||||
UI_TEXT_CONTENT["AISTUDIO::COMPONENTS::DATASOURCESELECTION::T1309929755"] = "Yes, let the AI validate & filter the retrieved data."
|
||||
|
||||
@ -3283,9 +3286,6 @@ UI_TEXT_CONTENT["AISTUDIO::COMPONENTS::DATASOURCESELECTION::T3100256862"] = "AI-
|
||||
-- No, I don't want to use data sources.
|
||||
UI_TEXT_CONTENT["AISTUDIO::COMPONENTS::DATASOURCESELECTION::T3135725655"] = "No, I don't want to use data sources."
|
||||
|
||||
-- Your data sources cannot be used with the LLM provider you selected due to data privacy, or they are currently unavailable.
|
||||
UI_TEXT_CONTENT["AISTUDIO::COMPONENTS::DATASOURCESELECTION::T3215374102"] = "Your data sources cannot be used with the LLM provider you selected due to data privacy, or they are currently unavailable."
|
||||
|
||||
-- No, I manually decide which data source to use.
|
||||
UI_TEXT_CONTENT["AISTUDIO::COMPONENTS::DATASOURCESELECTION::T3440789294"] = "No, I manually decide which data source to use."
|
||||
|
||||
@ -5077,9 +5077,6 @@ UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALDIRECTORYDIALOG::T3659673500"
|
||||
-- Compliance level
|
||||
UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALDIRECTORYDIALOG::T3995796156"] = "Compliance level"
|
||||
|
||||
-- Your security policy
|
||||
UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALDIRECTORYDIALOG::T4081226330"] = "Your security policy"
|
||||
|
||||
-- Select the base directory
|
||||
UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALDIRECTORYDIALOG::T562479068"] = "Select the base directory"
|
||||
|
||||
@ -5125,9 +5122,6 @@ UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALDIRECTORYINFODIALOG::T2717738
|
||||
-- The directory chosen for the data source does not exist anymore. Please edit the data source and correct the path.
|
||||
UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALDIRECTORYINFODIALOG::T2875614207"] = "The directory chosen for the data source does not exist anymore. Please edit the data source and correct the path."
|
||||
|
||||
-- your security policy
|
||||
UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALDIRECTORYINFODIALOG::T2879113658"] = "your security policy"
|
||||
|
||||
-- Maximum matches per query
|
||||
UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALDIRECTORYINFODIALOG::T2889706179"] = "Maximum matches per query"
|
||||
|
||||
@ -5155,9 +5149,6 @@ UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALDIRECTORYINFODIALOG::T3949388
|
||||
-- Compliance level
|
||||
UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALDIRECTORYINFODIALOG::T3995796156"] = "Compliance level"
|
||||
|
||||
-- Your security policy
|
||||
UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALDIRECTORYINFODIALOG::T4081226330"] = "Your security policy"
|
||||
|
||||
-- Number of files
|
||||
UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALDIRECTORYINFODIALOG::T417749210"] = "Number of files"
|
||||
|
||||
@ -5248,9 +5239,6 @@ UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALFILEDIALOG::T3740148848"] = "
|
||||
-- Compliance level
|
||||
UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALFILEDIALOG::T3995796156"] = "Compliance level"
|
||||
|
||||
-- Your security policy
|
||||
UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALFILEDIALOG::T4081226330"] = "Your security policy"
|
||||
|
||||
-- This data source already has indexed embeddings. Delete and recreate it to change the file path or embedding model.
|
||||
UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALFILEDIALOG::T458267883"] = "This data source already has indexed embeddings. Delete and recreate it to change the file path or embedding model."
|
||||
|
||||
@ -5293,9 +5281,6 @@ UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALFILEINFODIALOG::T2717738728"]
|
||||
-- the file size
|
||||
UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALFILEINFODIALOG::T2837935239"] = "the file size"
|
||||
|
||||
-- your security policy
|
||||
UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALFILEINFODIALOG::T2879113658"] = "your security policy"
|
||||
|
||||
-- File path
|
||||
UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALFILEINFODIALOG::T2879895266"] = "File path"
|
||||
|
||||
@ -5323,9 +5308,6 @@ UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALFILEINFODIALOG::T3688254408"]
|
||||
-- Compliance level
|
||||
UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALFILEINFODIALOG::T3995796156"] = "Compliance level"
|
||||
|
||||
-- Your security policy
|
||||
UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALFILEINFODIALOG::T4081226330"] = "Your security policy"
|
||||
|
||||
-- the compliance level
|
||||
UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALFILEINFODIALOG::T607609781"] = "the compliance level"
|
||||
|
||||
@ -10018,6 +10000,9 @@ UI_TEXT_CONTENT["AISTUDIO::TOOLS::VALIDATION::DATASOURCEVALIDATION::T2160507967"
|
||||
-- Please select your security policy.
|
||||
UI_TEXT_CONTENT["AISTUDIO::TOOLS::VALIDATION::DATASOURCEVALIDATION::T2250909198"] = "Please select your security policy."
|
||||
|
||||
-- The selected embedding provider is not allowed to process this data source due to its data security policy. Select a self-hosted or organization-trusted embedding provider.
|
||||
UI_TEXT_CONTENT["AISTUDIO::TOOLS::VALIDATION::DATASOURCEVALIDATION::T2616101497"] = "The selected embedding provider is not allowed to process this data source due to its data security policy. Select a self-hosted or organization-trusted embedding provider."
|
||||
|
||||
-- Please test the connection before saving.
|
||||
UI_TEXT_CONTENT["AISTUDIO::TOOLS::VALIDATION::DATASOURCEVALIDATION::T285470497"] = "Please test the connection before saving."
|
||||
|
||||
@ -10042,6 +10027,12 @@ UI_TEXT_CONTENT["AISTUDIO::TOOLS::VALIDATION::DATASOURCEVALIDATION::T3965971107"
|
||||
-- The name is already used by another data source. Please choose a different name.
|
||||
UI_TEXT_CONTENT["AISTUDIO::TOOLS::VALIDATION::DATASOURCEVALIDATION::T4001510395"] = "The name is already used by another data source. Please choose a different name."
|
||||
|
||||
-- Please select a compliance level.
|
||||
UI_TEXT_CONTENT["AISTUDIO::TOOLS::VALIDATION::DATASOURCEVALIDATION::T4066952091"] = "Please select a compliance level."
|
||||
|
||||
-- The selected embedding provider has confidence '{0}', but this data source requires provider confidence '{1}'. Select an embedding provider with equal or higher confidence or lower the compliance level.
|
||||
UI_TEXT_CONTENT["AISTUDIO::TOOLS::VALIDATION::DATASOURCEVALIDATION::T476537662"] = "The selected embedding provider has confidence '{0}', but this data source requires provider confidence '{1}'. Select an embedding provider with equal or higher confidence or lower the compliance level."
|
||||
|
||||
-- Please acknowledge that you are aware of the cloud embedding implications.
|
||||
UI_TEXT_CONTENT["AISTUDIO::TOOLS::VALIDATION::DATASOURCEVALIDATION::T490875633"] = "Please acknowledge that you are aware of the cloud embedding implications."
|
||||
|
||||
|
||||
@ -12,4 +12,7 @@ public sealed class ContentStreamSseEvent
|
||||
|
||||
[JsonPropertyName("metadata")]
|
||||
public ContentStreamSseMetadata? Metadata { get; init; }
|
||||
|
||||
[JsonPropertyName("token_count")]
|
||||
public int? TokenCount { get; init; }
|
||||
}
|
||||
@ -44,6 +44,14 @@ public static class ContentStreamSseHandler
|
||||
return sseEvent.Content;
|
||||
|
||||
case ContentStreamPresentationMetadata presentationMetadata:
|
||||
if (!extractImages)
|
||||
{
|
||||
var slideNumber = presentationMetadata.Presentation?.SlideNumber ?? 0;
|
||||
return slideNumber > 0
|
||||
? $"# Slide {slideNumber}\n{sseEvent.Content}"
|
||||
: sseEvent.Content;
|
||||
}
|
||||
|
||||
var slideManager = SLIDE_MANAGERS.GetOrAdd(
|
||||
sseEvent.StreamId!,
|
||||
_ => new()
|
||||
@ -131,4 +139,4 @@ public static class ContentStreamSseHandler
|
||||
|
||||
return finalContentChunk.Length > 0 ? finalContentChunk.ToString() : null;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@ -15,6 +15,7 @@ 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"];
|
||||
@ -28,7 +29,9 @@ public sealed partial class DataSourceEmbeddingService
|
||||
UNSUPPORTED,
|
||||
}
|
||||
|
||||
private sealed record ExtractedFileContent(string Text, IReadOnlyList<string> SourceSegments);
|
||||
private sealed record ExtractedFileSegment(string Text, int? TokenCount);
|
||||
|
||||
private sealed record ExtractedFileContent(string Text, IReadOnlyList<ExtractedFileSegment> SourceSegments);
|
||||
|
||||
private sealed record EmbeddingChunkDraft(string ChunkId, string Text, int ChunkIndex, int? PageNumber);
|
||||
|
||||
@ -36,7 +39,7 @@ public sealed partial class DataSourceEmbeddingService
|
||||
|
||||
private sealed record ChunkingStrategy(string Name, IReadOnlyList<ChunkingRule> Rules);
|
||||
|
||||
private sealed record ChunkingRule(string Name, Func<string, IReadOnlyList<string>, IReadOnlyList<string>>? Split);
|
||||
private sealed record ChunkingRule(string Name, Func<string, IReadOnlyList<string>, IReadOnlyList<string>>? Split, bool UsesSourceSegmentCounts = false);
|
||||
|
||||
private sealed record DataSourceMetadataSnapshot(string SourceHash, IReadOnlyDictionary<string, string> FileHashes);
|
||||
|
||||
@ -49,61 +52,69 @@ public sealed partial class DataSourceEmbeddingService
|
||||
if (this.IsImageFilePath(filePath))
|
||||
{
|
||||
var imageIndexText = this.BuildImageIndexText(filePath);
|
||||
content = new(imageIndexText, [imageIndexText]);
|
||||
content = new(imageIndexText, [new(imageIndexText, null)]);
|
||||
}
|
||||
else
|
||||
{
|
||||
content = await this.ReadExtractedFileContentAsync(filePath, token);
|
||||
content = await this.ReadExtractedFileContentAsync(filePath, embeddingProvider, token);
|
||||
}
|
||||
|
||||
await foreach (var chunk in this.SplitByChunkingStrategyAsync(content, strategy, options, embeddingProvider, token))
|
||||
yield return chunk;
|
||||
}
|
||||
|
||||
private async Task<ExtractedFileContent> ReadExtractedFileContentAsync(string filePath, CancellationToken token)
|
||||
private async Task<ExtractedFileContent> ReadExtractedFileContentAsync(string filePath, EmbeddingProvider embeddingProvider, CancellationToken token)
|
||||
{
|
||||
var segments = new List<string>();
|
||||
var segments = new List<ExtractedFileSegment>();
|
||||
|
||||
await foreach (var segment in rustService.StreamArbitraryFileData(filePath, token: token))
|
||||
await foreach (var segment in rustService.StreamArbitraryFileDataWithTokenCounts(filePath, embeddingProvider.Name, embeddingProvider.TokenizerPath, token))
|
||||
{
|
||||
var normalized = NormalizeChunkSegment(segment);
|
||||
var normalized = NormalizeChunkSegment(segment.Content);
|
||||
if (!string.IsNullOrWhiteSpace(normalized))
|
||||
segments.Add(normalized);
|
||||
segments.Add(new(normalized, segment.TokenCount));
|
||||
}
|
||||
|
||||
return new(string.Join("\n", segments).Trim(), segments);
|
||||
return new(string.Join("\n", segments.Select(segment => segment.Text)).Trim(), segments);
|
||||
}
|
||||
|
||||
private async IAsyncEnumerable<string> SplitByChunkingStrategyAsync(ExtractedFileContent content, ChunkingStrategy strategy, ChunkingOptions options, EmbeddingProvider embeddingProvider, [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token)
|
||||
{
|
||||
await foreach (var chunk in this.SplitTextByRulesAsync(content.Text, content.SourceSegments, strategy, 0, options, embeddingProvider, token))
|
||||
var estimatedTokenCount = SumTokenCounts(content.SourceSegments);
|
||||
await foreach (var chunk in this.SplitTextByRulesAsync(content.Text, content.SourceSegments, strategy, 0, options, embeddingProvider, token, estimatedTokenCount: estimatedTokenCount))
|
||||
yield return chunk;
|
||||
}
|
||||
|
||||
private async IAsyncEnumerable<string> SplitTextByRulesAsync(
|
||||
string text,
|
||||
IReadOnlyList<string> sourceSegments,
|
||||
IReadOnlyList<ExtractedFileSegment> sourceSegments,
|
||||
ChunkingStrategy strategy,
|
||||
int ruleIndex,
|
||||
ChunkingOptions options,
|
||||
EmbeddingProvider embeddingProvider,
|
||||
[System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token,
|
||||
string requiredOverlapPrefix = "")
|
||||
string requiredOverlapPrefix = "",
|
||||
int? estimatedTokenCount = null)
|
||||
{
|
||||
text = text.Trim();
|
||||
if (string.IsNullOrWhiteSpace(text))
|
||||
yield break;
|
||||
|
||||
var tokenCount = await this.GetEmbeddingTokenCountAsync(embeddingProvider, AddOverlapPrefix(text, requiredOverlapPrefix), token);
|
||||
if (tokenCount <= options.MaxChunkTokenLength)
|
||||
var tokenCount = estimatedTokenCount;
|
||||
var textWithOverlap = AddOverlapPrefix(text, requiredOverlapPrefix);
|
||||
if (textWithOverlap.Length <= MAX_TOKEN_COUNT_REQUEST_TEXT_LENGTH &&
|
||||
(estimatedTokenCount is null || estimatedTokenCount <= options.MaxChunkTokenLength))
|
||||
{
|
||||
yield return AddOverlapPrefix(text, requiredOverlapPrefix);
|
||||
yield break;
|
||||
tokenCount = await this.GetEmbeddingTokenCountAsync(embeddingProvider, textWithOverlap, token);
|
||||
if (tokenCount <= options.MaxChunkTokenLength)
|
||||
{
|
||||
yield return textWithOverlap;
|
||||
yield break;
|
||||
}
|
||||
}
|
||||
|
||||
if (ruleIndex >= strategy.Rules.Count)
|
||||
{
|
||||
await foreach (var hardChunk in this.SplitTextByHardCutAsync(text, options, embeddingProvider, token, requiredOverlapPrefix))
|
||||
await foreach (var hardChunk in this.SplitTextByHardCutAsync(text, options, embeddingProvider, token, requiredOverlapPrefix, estimatedTokenCount))
|
||||
yield return hardChunk;
|
||||
|
||||
yield break;
|
||||
@ -112,23 +123,23 @@ public sealed partial class DataSourceEmbeddingService
|
||||
var rule = strategy.Rules[ruleIndex];
|
||||
if (rule.Split is null)
|
||||
{
|
||||
await foreach (var hardChunk in this.SplitTextByHardCutAsync(text, options, embeddingProvider, token, requiredOverlapPrefix))
|
||||
await foreach (var hardChunk in this.SplitTextByHardCutAsync(text, options, embeddingProvider, token, requiredOverlapPrefix, estimatedTokenCount))
|
||||
yield return hardChunk;
|
||||
|
||||
yield break;
|
||||
}
|
||||
|
||||
var units = NormalizeSplitUnits(rule.Split(text, sourceSegments), text);
|
||||
var units = NormalizeSplitUnits(rule.Split(text, sourceSegments.Select(segment => segment.Text).ToList()), text);
|
||||
if (units.Count <= 1)
|
||||
{
|
||||
await foreach (var chunk in this.SplitTextByRulesAsync(text, sourceSegments, strategy, ruleIndex + 1, options, embeddingProvider, token, requiredOverlapPrefix))
|
||||
await foreach (var chunk in this.SplitTextByRulesAsync(text, sourceSegments, strategy, ruleIndex + 1, options, embeddingProvider, token, requiredOverlapPrefix, estimatedTokenCount))
|
||||
yield return chunk;
|
||||
|
||||
yield break;
|
||||
}
|
||||
|
||||
logger.LogDebug(
|
||||
"Splitting content for embedding provider '{EmbeddingProviderName}' with strategy '{ChunkingStrategy}' and rule '{ChunkingRule}'. TokenCount={TokenCount}, MaxChunkTokenLength={MaxChunkTokenLength}, OverlapTokenLength={OverlapTokenLength}.",
|
||||
"Splitting content for embedding provider '{EmbeddingProviderName}' with strategy '{ChunkingStrategy}' and rule '{ChunkingRule}'. EstimatedTokenCount={EstimatedTokenCount}, MaxChunkTokenLength={MaxChunkTokenLength}, OverlapTokenLength={OverlapTokenLength}.",
|
||||
embeddingProvider.Name,
|
||||
strategy.Name,
|
||||
rule.Name,
|
||||
@ -138,12 +149,13 @@ public sealed partial class DataSourceEmbeddingService
|
||||
|
||||
var index = 0;
|
||||
var overlapPrefix = requiredOverlapPrefix;
|
||||
var unitTokenCounts = EstimateSplitUnitTokenCounts(units, sourceSegments, rule.UsesSourceSegmentCounts, estimatedTokenCount);
|
||||
|
||||
while (index < units.Count)
|
||||
{
|
||||
token.ThrowIfCancellationRequested();
|
||||
|
||||
var unitCount = await this.FindLargestUnitCountWithinMaxChunkLengthAsync(units, index, embeddingProvider, options.MaxChunkTokenLength, token, overlapPrefix);
|
||||
var unitCount = await this.FindLargestUnitCountWithinMaxChunkLengthAsync(units, unitTokenCounts, index, embeddingProvider, options.MaxChunkTokenLength, token, overlapPrefix);
|
||||
if (unitCount > 0)
|
||||
{
|
||||
var rawChunk = string.Concat(units.Skip(index).Take(unitCount)).Trim();
|
||||
@ -181,7 +193,8 @@ public sealed partial class DataSourceEmbeddingService
|
||||
}
|
||||
|
||||
string? lastSplitUnit = null;
|
||||
await foreach (var splitUnit in this.SplitTextByRulesAsync(units[index], [units[index]], strategy, ruleIndex + 1, options, embeddingProvider, token, overlapPrefix))
|
||||
var unitTokenCount = unitTokenCounts?[index];
|
||||
await foreach (var splitUnit in this.SplitTextByRulesAsync(units[index], [new(units[index], unitTokenCount)], strategy, ruleIndex + 1, options, embeddingProvider, token, overlapPrefix, unitTokenCount))
|
||||
{
|
||||
lastSplitUnit = splitUnit;
|
||||
yield return splitUnit;
|
||||
@ -194,29 +207,105 @@ public sealed partial class DataSourceEmbeddingService
|
||||
}
|
||||
}
|
||||
|
||||
private async Task<int> FindLargestUnitCountWithinMaxChunkLengthAsync(IReadOnlyList<string> units, int startIndex, EmbeddingProvider embeddingProvider, int maxChunkTokenLength, CancellationToken token, string overlapPrefix = "")
|
||||
private async Task<int> FindLargestUnitCountWithinMaxChunkLengthAsync(IReadOnlyList<string> units, IReadOnlyList<int>? estimatedUnitTokenCounts, int startUnitIndex, EmbeddingProvider embeddingProvider, int maxChunkTokenLength, CancellationToken token, string overlapPrefix = "")
|
||||
{
|
||||
var low = 1;
|
||||
var high = units.Count - startIndex;
|
||||
var best = 0;
|
||||
var minimumCandidateUnitCount = 1;
|
||||
var availableUnitCount = units.Count - startUnitIndex;
|
||||
var maximumCandidateUnitCount = availableUnitCount;
|
||||
var largestValidUnitCount = 0;
|
||||
|
||||
while (low <= high)
|
||||
if (estimatedUnitTokenCounts is not null)
|
||||
{
|
||||
token.ThrowIfCancellationRequested();
|
||||
|
||||
var mid = low + (high - low) / 2;
|
||||
var candidate = AddOverlapPrefix(string.Concat(units.Skip(startIndex).Take(mid)).Trim(), overlapPrefix);
|
||||
var tokenCount = await this.GetEmbeddingTokenCountAsync(embeddingProvider, candidate, token);
|
||||
if (tokenCount <= maxChunkTokenLength)
|
||||
maximumCandidateUnitCount = 0;
|
||||
var cumulativeEstimatedTokenCount = 0L;
|
||||
for (var unitIndex = startUnitIndex; unitIndex < units.Count; unitIndex++)
|
||||
{
|
||||
best = mid;
|
||||
low = mid + 1;
|
||||
cumulativeEstimatedTokenCount += estimatedUnitTokenCounts[unitIndex];
|
||||
if (cumulativeEstimatedTokenCount > maxChunkTokenLength)
|
||||
break;
|
||||
|
||||
maximumCandidateUnitCount++;
|
||||
}
|
||||
else
|
||||
high = mid - 1;
|
||||
|
||||
if (maximumCandidateUnitCount == 0)
|
||||
maximumCandidateUnitCount = 1;
|
||||
}
|
||||
|
||||
return best;
|
||||
while (true)
|
||||
{
|
||||
var searchedMaximumCandidateUnitCount = maximumCandidateUnitCount;
|
||||
while (minimumCandidateUnitCount <= maximumCandidateUnitCount)
|
||||
{
|
||||
token.ThrowIfCancellationRequested();
|
||||
|
||||
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 &&
|
||||
await this.GetEmbeddingTokenCountAsync(embeddingProvider, candidateText, token) <= maxChunkTokenLength;
|
||||
if (candidateFits)
|
||||
{
|
||||
largestValidUnitCount = candidateUnitCount;
|
||||
minimumCandidateUnitCount = candidateUnitCount + 1;
|
||||
}
|
||||
else
|
||||
maximumCandidateUnitCount = candidateUnitCount - 1;
|
||||
}
|
||||
|
||||
if (largestValidUnitCount < searchedMaximumCandidateUnitCount || largestValidUnitCount >= availableUnitCount)
|
||||
break;
|
||||
|
||||
minimumCandidateUnitCount = searchedMaximumCandidateUnitCount + 1;
|
||||
maximumCandidateUnitCount = (int)Math.Min(
|
||||
availableUnitCount,
|
||||
Math.Max((long)minimumCandidateUnitCount, (long)searchedMaximumCandidateUnitCount * 2));
|
||||
}
|
||||
|
||||
return largestValidUnitCount;
|
||||
}
|
||||
|
||||
private static int? SumTokenCounts(IReadOnlyList<ExtractedFileSegment> segments)
|
||||
{
|
||||
var result = 0L;
|
||||
foreach (var segment in segments)
|
||||
{
|
||||
if (segment.TokenCount is null)
|
||||
return null;
|
||||
|
||||
result += segment.TokenCount.Value;
|
||||
}
|
||||
|
||||
return (int)Math.Min(result, int.MaxValue);
|
||||
}
|
||||
|
||||
private static IReadOnlyList<int>? EstimateSplitUnitTokenCounts(
|
||||
IReadOnlyList<string> units,
|
||||
IReadOnlyList<ExtractedFileSegment> sourceSegments,
|
||||
bool usesSourceSegmentCounts,
|
||||
int? sourceTokenCount)
|
||||
{
|
||||
if (usesSourceSegmentCounts && sourceSegments.Count == units.Count && sourceSegments.All(segment => segment.TokenCount is not null))
|
||||
return sourceSegments.Select(segment => segment.TokenCount.GetValueOrDefault()).ToList();
|
||||
|
||||
if (sourceTokenCount is null)
|
||||
return null;
|
||||
|
||||
var totalLength = Math.Max(1, units.Sum(unit => unit.Length));
|
||||
var result = new List<int>(units.Count);
|
||||
var allocatedTokenCount = 0;
|
||||
var consumedLength = 0L;
|
||||
|
||||
foreach (var unit in units)
|
||||
{
|
||||
consumedLength += unit.Length;
|
||||
var tokenCountAtBoundary = (int)Math.Min(sourceTokenCount.Value, (long)sourceTokenCount.Value * consumedLength / totalLength);
|
||||
result.Add(Math.Max(0, tokenCountAtBoundary - allocatedTokenCount));
|
||||
allocatedTokenCount = tokenCountAtBoundary;
|
||||
}
|
||||
|
||||
if (result.Count > 0 && allocatedTokenCount < sourceTokenCount.Value)
|
||||
result[^1] += sourceTokenCount.Value - allocatedTokenCount;
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
private async Task<string> CreateOverlapPrefixAsync(string chunk, ChunkingStrategy strategy, ChunkingRule rule, ChunkingOptions options, EmbeddingProvider embeddingProvider, CancellationToken token)
|
||||
@ -304,7 +393,8 @@ public sealed partial class DataSourceEmbeddingService
|
||||
ChunkingOptions options,
|
||||
EmbeddingProvider embeddingProvider,
|
||||
[System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token,
|
||||
string requiredOverlapPrefix = "")
|
||||
string requiredOverlapPrefix = "",
|
||||
int? estimatedTokenCount = null)
|
||||
{
|
||||
text = text.Trim();
|
||||
var startIndex = 0;
|
||||
@ -319,22 +409,42 @@ public sealed partial class DataSourceEmbeddingService
|
||||
if (startIndex >= text.Length)
|
||||
yield break;
|
||||
|
||||
var low = startIndex + 1;
|
||||
var high = text.Length;
|
||||
var bestEndIndex = startIndex;
|
||||
var maximumCandidateEndIndex = Math.Min(text.Length, startIndex + MAX_TOKEN_COUNT_REQUEST_TEXT_LENGTH);
|
||||
|
||||
while (low <= high)
|
||||
if (estimatedTokenCount > options.MaxChunkTokenLength)
|
||||
{
|
||||
var mid = low + (high - low) / 2;
|
||||
var candidate = AddOverlapPrefix(text[startIndex..mid].Trim(), overlapPrefix);
|
||||
var tokenCount = await this.GetEmbeddingTokenCountAsync(embeddingProvider, candidate, token);
|
||||
if (tokenCount <= options.MaxChunkTokenLength)
|
||||
var estimatedChunkLength = Math.Max(1L, (long)text.Length * options.MaxChunkTokenLength / estimatedTokenCount.Value);
|
||||
maximumCandidateEndIndex = (int)Math.Min(text.Length, startIndex + estimatedChunkLength);
|
||||
}
|
||||
|
||||
while (true)
|
||||
{
|
||||
var minimumCandidateEndIndex = bestEndIndex + 1;
|
||||
var currentMaximumCandidateEndIndex = maximumCandidateEndIndex;
|
||||
while (minimumCandidateEndIndex <= currentMaximumCandidateEndIndex)
|
||||
{
|
||||
bestEndIndex = mid;
|
||||
low = mid + 1;
|
||||
var candidateEndIndex = minimumCandidateEndIndex + (currentMaximumCandidateEndIndex - minimumCandidateEndIndex) / 2;
|
||||
var candidate = AddOverlapPrefix(text[startIndex..candidateEndIndex].Trim(), overlapPrefix);
|
||||
var candidateFits = candidate.Length <= MAX_TOKEN_COUNT_REQUEST_TEXT_LENGTH &&
|
||||
await this.GetEmbeddingTokenCountAsync(embeddingProvider, candidate, token) <= options.MaxChunkTokenLength;
|
||||
if (candidateFits)
|
||||
{
|
||||
bestEndIndex = candidateEndIndex;
|
||||
minimumCandidateEndIndex = candidateEndIndex + 1;
|
||||
}
|
||||
else
|
||||
currentMaximumCandidateEndIndex = candidateEndIndex - 1;
|
||||
}
|
||||
else
|
||||
high = mid - 1;
|
||||
|
||||
if (bestEndIndex < maximumCandidateEndIndex || bestEndIndex >= text.Length ||
|
||||
maximumCandidateEndIndex - startIndex >= 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),
|
||||
startIndex + Math.Max(previousCandidateLength + 1L, previousCandidateLength * 2L));
|
||||
}
|
||||
|
||||
if (bestEndIndex == startIndex)
|
||||
@ -466,7 +576,7 @@ public sealed partial class DataSourceEmbeddingService
|
||||
|
||||
if (this.IsPresentationFilePath(filePath))
|
||||
return new("presentation", [
|
||||
new("Slide", SplitBySourceSegments),
|
||||
new("Slide", SplitBySourceSegments, true),
|
||||
new("Line break", SplitByLineBreaks),
|
||||
new("Whitespace", SplitByWhitespace),
|
||||
new("Hard cut", null),
|
||||
@ -474,7 +584,7 @@ public sealed partial class DataSourceEmbeddingService
|
||||
|
||||
if (this.IsDelimitedTableFilePath(filePath) || this.IsSpreadsheetFilePath(filePath))
|
||||
return new("table", [
|
||||
new("Row or sheet", SplitBySourceSegments),
|
||||
new("Row or sheet", SplitBySourceSegments, true),
|
||||
new("Line break", SplitByLineBreaks),
|
||||
new("Whitespace", SplitByWhitespace),
|
||||
new("Hard cut", null),
|
||||
@ -484,8 +594,8 @@ public sealed partial class DataSourceEmbeddingService
|
||||
return GetSourceCodeChunkingStrategy(filePath);
|
||||
|
||||
return new("document", [
|
||||
new("Page or extracted section", SplitBySourceSegments, true),
|
||||
new("Heading", SplitByDocumentHeadings),
|
||||
new("Page or extracted section", SplitBySourceSegments),
|
||||
new("Paragraph", SplitByParagraphs),
|
||||
new("Line break", SplitByLineBreaks),
|
||||
new("Whitespace", SplitByWhitespace),
|
||||
@ -495,7 +605,11 @@ public sealed partial class DataSourceEmbeddingService
|
||||
|
||||
private static ChunkingStrategy GetSourceCodeChunkingStrategy(string filePath)
|
||||
{
|
||||
var rules = GetSourceCodeDelimiterRules(filePath).ToList();
|
||||
var rules = new List<ChunkingRule>
|
||||
{
|
||||
new("Extracted section", SplitBySourceSegments, true),
|
||||
};
|
||||
rules.AddRange(GetSourceCodeDelimiterRules(filePath));
|
||||
rules.Add(new("Line break", SplitByLineBreaks));
|
||||
rules.Add(new("Whitespace", SplitByWhitespace));
|
||||
rules.Add(new("Hard cut", null));
|
||||
|
||||
@ -570,6 +570,10 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
|
||||
var startedAtUtc = DateTime.UtcNow;
|
||||
var chunkCount = await this.IndexOneFileAsync(embeddingState, vectorStore, dataSource, file, fingerprint, embeddingProvider, provider, manifest, optimizationTracker, token);
|
||||
token.ThrowIfCancellationRequested();
|
||||
var fingerprintAfterEmbedding = BuildFileMetadataHash(file);
|
||||
if (!string.Equals(fingerprint, fingerprintAfterEmbedding, StringComparison.Ordinal))
|
||||
throw new IOException($"The file '{file.FullName}' changed while it was being embedded. Its partial embeddings will be discarded and the file will be retried on the next refresh.");
|
||||
|
||||
var embeddedAtUtc = DateTime.UtcNow;
|
||||
var record = new EmbeddedFileRecord(
|
||||
fingerprint,
|
||||
@ -606,9 +610,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
|
||||
lastError = exception.Message;
|
||||
failureDetails.Add(new DataSourceEmbeddingFailure(file.FullName, exception.Message));
|
||||
manifest.Files.Remove(file.FullName);
|
||||
await this.DeleteFilePointsAsync(vectorStore, collectionName, file.FullName, token);
|
||||
optimizationTracker.MarkChanged();
|
||||
await embeddingState.DeleteFileAsync(dataSource.Id, file.FullName, token);
|
||||
await this.CleanupFailedFileAsync(embeddingState, vectorStore, dataSource, collectionName, file.FullName, optimizationTracker, token);
|
||||
|
||||
logger.LogWarning(exception, "Failed to embed file '{FilePath}' for data source '{DataSourceName}'.", file.FullName, dataSource.Name);
|
||||
this.UpsertStatus(this.CreateStatus(dataSource, DataSourceEmbeddingState.RUNNING, totalFiles, skippedFiles + completedFiles, failedFiles, file.Name, exception.Message, failureDetails));
|
||||
@ -669,7 +671,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
|
||||
var parentFile = this.CreateEmbeddingStateFile(dataSource, file, fingerprint, 0, DateTime.UtcNow);
|
||||
await embeddingState.UpsertFileAsync(dataSource.Id, parentFile, token);
|
||||
|
||||
var embeddingBatchSize = embeddingProvider.EffectiveEmbeddingBatchSize;
|
||||
var embeddingBatchSize = Math.Max(1, embeddingProvider.EffectiveEmbeddingBatchSize);
|
||||
var batch = new List<EmbeddingChunkDraft>(embeddingBatchSize);
|
||||
var totalChunkCount = 0;
|
||||
|
||||
@ -743,15 +745,21 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
|
||||
if (vectorSize <= 0)
|
||||
throw new InvalidOperationException("The embedding provider returned an empty vector.");
|
||||
|
||||
if (vectors.Any(vector => vector.Count != vectorSize))
|
||||
throw new InvalidOperationException("The embedding provider returned vectors with inconsistent dimensions.");
|
||||
|
||||
if (vectors.Any(vector => vector.Any(value => !float.IsFinite(value))))
|
||||
throw new InvalidOperationException("The embedding provider returned a vector containing a non-finite value.");
|
||||
|
||||
if (manifest.VectorSize > 0 && manifest.VectorSize != vectorSize)
|
||||
throw new InvalidOperationException($"The embedding vector size changed from {manifest.VectorSize} to {vectorSize}. Please re-save the data source to trigger a clean re-index.");
|
||||
|
||||
if (manifest.VectorSize == 0)
|
||||
{
|
||||
token.ThrowIfCancellationRequested();
|
||||
manifest.VectorSize = vectorSize;
|
||||
await this.EnsureCollectionExistsAsync(vectorStore, collectionName, vectorSize, token);
|
||||
await embeddingState.UpdateVectorSizeAsync(dataSource.Id, vectorSize, token);
|
||||
manifest.VectorSize = vectorSize;
|
||||
logger.LogInformation(
|
||||
"Created embedding collection '{CollectionName}' with vector size {VectorSize} for data source '{DataSourceName}' ({DataSourceId}).",
|
||||
collectionName,
|
||||
@ -846,6 +854,53 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
|
||||
await vectorStore.DeleteEmbeddingByFile(collectionName, filePath, token);
|
||||
}
|
||||
|
||||
private async Task CleanupFailedFileAsync(
|
||||
EmbeddingStateClient embeddingState,
|
||||
VectorStoreClient vectorStore,
|
||||
IDataSource dataSource,
|
||||
string collectionName,
|
||||
string filePath,
|
||||
VectorStoreOptimizationTracker optimizationTracker,
|
||||
CancellationToken token)
|
||||
{
|
||||
try
|
||||
{
|
||||
await this.DeleteFilePointsAsync(vectorStore, collectionName, filePath, token);
|
||||
optimizationTracker.MarkChanged();
|
||||
}
|
||||
catch (OperationCanceledException) when (token.IsCancellationRequested)
|
||||
{
|
||||
throw;
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
logger.LogWarning(
|
||||
exception,
|
||||
"Could not remove vector points while cleaning up failed embedding for file '{FilePath}' in data source '{DataSourceName}' ({DataSourceId}).",
|
||||
filePath,
|
||||
dataSource.Name,
|
||||
dataSource.Id);
|
||||
}
|
||||
|
||||
try
|
||||
{
|
||||
await embeddingState.DeleteFileAsync(dataSource.Id, filePath, token);
|
||||
}
|
||||
catch (OperationCanceledException) when (token.IsCancellationRequested)
|
||||
{
|
||||
throw;
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
logger.LogWarning(
|
||||
exception,
|
||||
"Could not remove embedding state while cleaning up failed embedding for file '{FilePath}' in data source '{DataSourceName}' ({DataSourceId}).",
|
||||
filePath,
|
||||
dataSource.Name,
|
||||
dataSource.Id);
|
||||
}
|
||||
}
|
||||
|
||||
private async Task OptimizeCollectionIfNeededAsync(
|
||||
VectorStoreOptimizationTracker optimizationTracker,
|
||||
VectorStoreClient vectorStore,
|
||||
|
||||
@ -2,14 +2,18 @@ using System.Text;
|
||||
using System.Text.Json;
|
||||
using System.Runtime.CompilerServices;
|
||||
|
||||
using AIStudio.Tools.Rust;
|
||||
|
||||
namespace AIStudio.Tools.Services;
|
||||
|
||||
public sealed record ArbitraryFileDataSegment(string Content, int TokenCount);
|
||||
|
||||
public sealed partial class RustService
|
||||
{
|
||||
public async Task<string> ReadArbitraryFileData(string path, int maxChunks, bool extractImages = false)
|
||||
{
|
||||
var streamId = Guid.NewGuid().ToString();
|
||||
var requestUri = $"/retrieval/fs/extract?path={Uri.EscapeDataString(path)}&stream_id={streamId}&extract_images={extractImages}";
|
||||
var requestUri = $"/retrieval/fs/extract?path={Uri.EscapeDataString(path)}&stream_id={streamId}&extract_images={extractImages}&include_token_count=false";
|
||||
var request = new HttpRequestMessage(HttpMethod.Get, requestUri);
|
||||
var response = await this.http.SendAsync(request, HttpCompletionOption.ResponseHeadersRead);
|
||||
|
||||
@ -80,14 +84,67 @@ public sealed partial class RustService
|
||||
}
|
||||
|
||||
public async IAsyncEnumerable<string> StreamArbitraryFileData(string path, bool extractImages = false, [EnumeratorCancellation] CancellationToken token = default)
|
||||
{
|
||||
await foreach (var segment in this.StreamArbitraryFileDataCore(path, extractImages, false, token))
|
||||
yield return segment.Content;
|
||||
}
|
||||
|
||||
public async IAsyncEnumerable<ArbitraryFileDataSegment> StreamArbitraryFileDataWithTokenCounts(
|
||||
string path,
|
||||
string providerName,
|
||||
string tokenizerPath,
|
||||
[EnumeratorCancellation] CancellationToken token = default)
|
||||
{
|
||||
await this.tokenizerLock.WaitAsync(token);
|
||||
try
|
||||
{
|
||||
var tokenizerResponse = await this.EnsureTokenizerCoreAsync(providerName, tokenizerPath);
|
||||
if (tokenizerResponse is not { Success: true, Status: TokenizerStatus.AVAILABLE })
|
||||
{
|
||||
var message = tokenizerResponse?.Message ?? "No response was returned by the tokenizer service.";
|
||||
throw new InvalidOperationException($"Could not initialize tokenizer for provider '{providerName}'. {message}");
|
||||
}
|
||||
|
||||
await foreach (var segment in this.StreamArbitraryFileDataCore(path, false, true, token))
|
||||
{
|
||||
if (segment.TokenCount is null)
|
||||
throw new InvalidOperationException($"Rust did not return a token count for an extracted segment from '{path}'.");
|
||||
|
||||
yield return new(segment.Content, segment.TokenCount.Value);
|
||||
}
|
||||
}
|
||||
finally
|
||||
{
|
||||
this.tokenizerLock.Release();
|
||||
}
|
||||
}
|
||||
|
||||
private async IAsyncEnumerable<(string Content, int? TokenCount)> StreamArbitraryFileDataCore(
|
||||
string path,
|
||||
bool extractImages,
|
||||
bool includeTokenCount,
|
||||
[EnumeratorCancellation] CancellationToken token)
|
||||
{
|
||||
var streamId = Guid.NewGuid().ToString();
|
||||
var requestUri = $"/retrieval/fs/extract?path={Uri.EscapeDataString(path)}&stream_id={streamId}&extract_images={extractImages}";
|
||||
var requestUri = $"/retrieval/fs/extract?path={Uri.EscapeDataString(path)}&stream_id={streamId}&extract_images={extractImages}&include_token_count={includeTokenCount}";
|
||||
using var request = new HttpRequestMessage(HttpMethod.Get, requestUri);
|
||||
using var response = await this.http.SendAsync(request, HttpCompletionOption.ResponseHeadersRead, token);
|
||||
|
||||
if (!response.IsSuccessStatusCode)
|
||||
{
|
||||
var responseBody = await response.Content.ReadAsStringAsync(token);
|
||||
this.logger?.LogError(
|
||||
"Failed to stream arbitrary file data from Rust runtime. Status: {StatusCode}, reason: '{ReasonPhrase}', path: '{Path}', body: '{Body}'",
|
||||
response.StatusCode,
|
||||
response.ReasonPhrase,
|
||||
path,
|
||||
responseBody);
|
||||
|
||||
if (includeTokenCount)
|
||||
throw new InvalidOperationException($"Rust could not extract and count '{path}'. HTTP {(int)response.StatusCode} ({response.ReasonPhrase}): {responseBody}");
|
||||
|
||||
yield break;
|
||||
}
|
||||
|
||||
string? finalContentChunk = null;
|
||||
try
|
||||
@ -113,7 +170,12 @@ public sealed partial class RustService
|
||||
catch (JsonException)
|
||||
{
|
||||
if (this.TryLogSseErrorMessage(jsonContent, path))
|
||||
{
|
||||
if (includeTokenCount)
|
||||
throw new InvalidOperationException($"Rust could not extract and count a segment from '{path}'. See the runtime log for details.");
|
||||
|
||||
continue;
|
||||
}
|
||||
|
||||
this.logger?.LogError("Failed to deserialize SSE event: {JsonContent}", jsonContent);
|
||||
}
|
||||
@ -123,7 +185,7 @@ public sealed partial class RustService
|
||||
|
||||
var content = ContentStreamSseHandler.ProcessEvent(sseEvent, extractImages);
|
||||
if (!string.IsNullOrWhiteSpace(content))
|
||||
yield return content;
|
||||
yield return (content, sseEvent.TokenCount);
|
||||
}
|
||||
}
|
||||
finally
|
||||
@ -132,7 +194,7 @@ public sealed partial class RustService
|
||||
}
|
||||
|
||||
if (!string.IsNullOrWhiteSpace(finalContentChunk))
|
||||
yield return finalContentChunk;
|
||||
yield return (finalContentChunk, null);
|
||||
}
|
||||
|
||||
private bool TryLogSseErrorMessage(string jsonContent, string path)
|
||||
|
||||
@ -28,17 +28,40 @@ pub struct Chunk {
|
||||
pub content: String,
|
||||
pub stream_id: String,
|
||||
pub metadata: Metadata,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub token_count: Option<usize>,
|
||||
}
|
||||
|
||||
impl Chunk {
|
||||
pub fn new(content: String, metadata: Metadata) -> Self {
|
||||
Chunk { content, stream_id: String::new(), metadata }
|
||||
Chunk { content, stream_id: String::new(), metadata, token_count: None }
|
||||
}
|
||||
|
||||
pub fn set_stream_id(&mut self, stream_id: &str) { self.stream_id = stream_id.to_string(); }
|
||||
|
||||
pub fn set_token_count(&mut self) -> std::result::Result<(), String> {
|
||||
self.token_count = Some(crate::tokenizer::get_segment_token_count(&self.content)?);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn into_bounded_text_segments(self) -> Vec<Self> {
|
||||
if matches!(&self.metadata, Metadata::Image {}) {
|
||||
return vec![self];
|
||||
}
|
||||
|
||||
let ranges = bounded_text_segment_ranges(&self.content);
|
||||
if ranges.len() == 1 {
|
||||
return vec![self];
|
||||
}
|
||||
|
||||
ranges
|
||||
.into_iter()
|
||||
.map(|(start, end)| Chunk::new(self.content[start..end].to_string(), self.metadata.clone()))
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[derive(Clone, Debug, Serialize)]
|
||||
pub enum Metadata {
|
||||
Text {
|
||||
line_number: usize
|
||||
@ -62,7 +85,7 @@ pub enum Metadata {
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[derive(Clone, Debug, Serialize)]
|
||||
pub struct Base64Image {
|
||||
pub id: String,
|
||||
pub content: String,
|
||||
@ -80,6 +103,7 @@ const TO_MARKDOWN: &str = "markdown";
|
||||
const DOCX: &str = "docx";
|
||||
const ODT: &str = "odt";
|
||||
const IMAGE_SEGMENT_SIZE_IN_CHARS: usize = 8_192; // equivalent to ~ 5500 token
|
||||
const MAX_TEXT_SEGMENT_LENGTH_IN_CHARS: usize = 100_000;
|
||||
|
||||
type Result<T> = std::result::Result<T, Box<dyn std::error::Error + Send + Sync>>;
|
||||
type ChunkStream = Pin<Box<dyn Stream<Item = Result<Chunk>> + Send>>;
|
||||
@ -90,6 +114,8 @@ pub struct ExtractDataQuery {
|
||||
stream_id: String,
|
||||
#[serde(deserialize_with = "deserialize_bool_case_insensitive")]
|
||||
extract_images: bool,
|
||||
#[serde(default, deserialize_with = "deserialize_bool_case_insensitive")]
|
||||
include_token_count: bool,
|
||||
}
|
||||
|
||||
fn deserialize_bool_case_insensitive<'de, D>(deserializer: D) -> std::result::Result<bool, D::Error>
|
||||
@ -145,11 +171,21 @@ pub async fn extract_data(
|
||||
|
||||
match stream_result {
|
||||
Ok(mut stream) => {
|
||||
while let Some(chunk) = stream.next().await {
|
||||
'stream_chunks: while let Some(chunk) = stream.next().await {
|
||||
match chunk {
|
||||
Ok(mut chunk) => {
|
||||
chunk.set_stream_id(id_ref);
|
||||
yield Ok(Event::default().json_data(&chunk).unwrap_or_else(|e| Event::default().data(format!("Error: {e}"))));
|
||||
Ok(chunk) => {
|
||||
let chunks = chunk.into_bounded_text_segments();
|
||||
|
||||
for mut chunk in chunks {
|
||||
chunk.set_stream_id(id_ref);
|
||||
if query.include_token_count {
|
||||
if let Err(e) = chunk.set_token_count() {
|
||||
yield Ok(Event::default().json_data(format!("Error counting tokens: {e}")).unwrap_or_else(|_| Event::default().data(format!("Error counting tokens: {e}"))));
|
||||
break 'stream_chunks;
|
||||
}
|
||||
}
|
||||
yield Ok(Event::default().json_data(&chunk).unwrap_or_else(|e| Event::default().data(format!("Error: {e}"))));
|
||||
}
|
||||
},
|
||||
|
||||
Err(e) => {
|
||||
@ -175,6 +211,33 @@ pub async fn extract_data(
|
||||
Sse::new(stream)
|
||||
}
|
||||
|
||||
fn bounded_text_segment_ranges(content: &str) -> Vec<(usize, usize)> {
|
||||
let mut ranges = Vec::new();
|
||||
let mut start = 0;
|
||||
|
||||
while start < content.len() {
|
||||
let remaining = &content[start..];
|
||||
let Some(maximum_end_offset) = remaining
|
||||
.char_indices()
|
||||
.nth(MAX_TEXT_SEGMENT_LENGTH_IN_CHARS)
|
||||
.map(|(index, _)| index)
|
||||
else {
|
||||
ranges.push((start, content.len()));
|
||||
break;
|
||||
};
|
||||
|
||||
let end = start + maximum_end_offset;
|
||||
ranges.push((start, end));
|
||||
start = end;
|
||||
}
|
||||
|
||||
if ranges.is_empty() {
|
||||
ranges.push((0, 0));
|
||||
}
|
||||
|
||||
ranges
|
||||
}
|
||||
|
||||
async fn stream_data(file_path: &str, extract_images: bool) -> Result<ChunkStream> {
|
||||
if !Path::new(file_path).exists() {
|
||||
error!("File does not exist: '{file_path}'");
|
||||
|
||||
@ -15,6 +15,7 @@ use qdrant_edge::{
|
||||
UpdateOperation, ValueVariants, VectorInternal, Vectors, WithPayloadInterface, WithVector,
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sha2::{Digest, Sha256};
|
||||
use tauri::Manager;
|
||||
|
||||
use crate::api_token::APIToken;
|
||||
@ -27,6 +28,8 @@ const HNSW_EF_CONSTRUCT: usize = 100;
|
||||
const HNSW_FULL_SCAN_THRESHOLD_KB: usize = 10_000;
|
||||
const HNSW_MAX_INDEXING_THREADS: usize = 0;
|
||||
const VECTOR_INDEXING_THRESHOLD_KB: usize = 10_000;
|
||||
const STORE_INITIALIZATION_MARKER: &str = "store_name.txt";
|
||||
const STORE_INITIALIZATION_MARKER_TEMP: &str = "store_name.tmp";
|
||||
|
||||
type QdrantEdgeResult<T> = Result<T, Box<dyn std::error::Error + Send + Sync>>;
|
||||
|
||||
@ -182,21 +185,52 @@ impl QdrantEdgeDatabase {
|
||||
|
||||
fn store_path(&self, store_name: &str) -> QdrantEdgeResult<PathBuf> {
|
||||
validate_store_name(store_name)?;
|
||||
Ok(self.base_path.join("stores").join(store_name))
|
||||
Ok(self.base_path.join("stores").join(store_directory_name(store_name)))
|
||||
}
|
||||
|
||||
// To ensure a shard exists and that you can insert a vector
|
||||
fn get_or_create_store(&mut self, store_name: &str, vector_size: usize) -> QdrantEdgeResult<&EdgeShard> {
|
||||
let path = self.store_path(store_name)?;
|
||||
let is_initialized = store_is_initialized(&path, store_name)?;
|
||||
if self.shards.contains_key(store_name) {
|
||||
return Ok(self.shards.get(store_name).unwrap());
|
||||
if is_initialized {
|
||||
return Ok(self.shards.get(store_name).unwrap());
|
||||
}
|
||||
|
||||
warn!(Source = "Qdrant Edge"; "Removing stale cached vector store '{}' because its initialized data directory no longer exists.", store_name);
|
||||
self.shards.remove(store_name);
|
||||
}
|
||||
|
||||
let path = self.store_path(store_name)?;
|
||||
let shard = if has_existing_store(&path) {
|
||||
EdgeShard::load(&path, None)?
|
||||
if path.exists() && !is_initialized {
|
||||
warn!(Source = "Qdrant Edge"; "Removing incompletely initialized vector store '{}' before recreating it.", store_name);
|
||||
fs::remove_dir_all(&path).map_err(|error| {
|
||||
format!("Failed to remove incomplete vector store '{store_name}' at '{}': {error}", path.display())
|
||||
})?;
|
||||
}
|
||||
|
||||
let shard = if is_initialized {
|
||||
EdgeShard::load(&path, None).map_err(|error| {
|
||||
format!("Failed to load vector store '{store_name}' from '{}': {error}", path.display())
|
||||
})?
|
||||
} else {
|
||||
fs::create_dir_all(&path)?;
|
||||
EdgeShard::new(&path, edge_config(vector_size))?
|
||||
fs::create_dir_all(&path).map_err(|error| {
|
||||
format!("Failed to create directory for vector store '{store_name}' at '{}': {error}", path.display())
|
||||
})?;
|
||||
let shard = match EdgeShard::new(&path, edge_config(vector_size)) {
|
||||
Ok(shard) => shard,
|
||||
Err(error) => {
|
||||
let cleanup_issue = remove_partial_store(&path);
|
||||
return Err(format!("Failed to create vector store '{store_name}' at '{}': {error}{cleanup_issue}", path.display()).into());
|
||||
},
|
||||
};
|
||||
|
||||
if let Err(error) = write_store_initialization_marker(&path, store_name) {
|
||||
drop(shard);
|
||||
let cleanup_issue = remove_partial_store(&path);
|
||||
return Err(format!("Failed to finalize vector store '{store_name}' at '{}': {error}{cleanup_issue}", path.display()).into());
|
||||
}
|
||||
|
||||
shard
|
||||
};
|
||||
|
||||
self.shards.insert(store_name.to_string(), shard);
|
||||
@ -205,16 +239,31 @@ impl QdrantEdgeDatabase {
|
||||
|
||||
// To check whether a shard exists so you can delete a file from it
|
||||
fn get_existing_store(&mut self, store_name: &str) -> QdrantEdgeResult<Option<&EdgeShard>> {
|
||||
let path = self.store_path(store_name)?;
|
||||
let is_initialized = store_is_initialized(&path, store_name)?;
|
||||
if self.shards.contains_key(store_name) {
|
||||
return Ok(self.shards.get(store_name));
|
||||
if is_initialized {
|
||||
return Ok(self.shards.get(store_name));
|
||||
}
|
||||
|
||||
warn!(Source = "Qdrant Edge"; "Removing stale cached vector store '{}' because its initialized data directory no longer exists.", store_name);
|
||||
self.shards.remove(store_name);
|
||||
}
|
||||
|
||||
let path = self.store_path(store_name)?;
|
||||
if !has_existing_store(&path) {
|
||||
if path.exists() && !is_initialized {
|
||||
warn!(Source = "Qdrant Edge"; "Removing incompletely initialized vector store '{}' before continuing.", store_name);
|
||||
fs::remove_dir_all(&path).map_err(|error| {
|
||||
format!("Failed to remove incomplete vector store '{store_name}' at '{}': {error}", path.display())
|
||||
})?;
|
||||
}
|
||||
|
||||
if !is_initialized {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let shard = EdgeShard::load(&path, None)?;
|
||||
let shard = EdgeShard::load(&path, None).map_err(|error| {
|
||||
format!("Failed to load vector store '{store_name}' from '{}': {error}", path.display())
|
||||
})?;
|
||||
self.shards.insert(store_name.to_string(), shard);
|
||||
Ok(self.shards.get(store_name))
|
||||
}
|
||||
@ -224,7 +273,7 @@ impl QdrantEdgeDatabase {
|
||||
let stores_count = if stores_path.exists() {
|
||||
fs::read_dir(stores_path)?
|
||||
.filter_map(Result::ok)
|
||||
.filter(|entry| entry.path().is_dir())
|
||||
.filter(|entry| entry.path().join(STORE_INITIALIZATION_MARKER).is_file())
|
||||
.count()
|
||||
} else {
|
||||
0
|
||||
@ -339,6 +388,16 @@ impl QdrantEdgeDatabase {
|
||||
}
|
||||
}
|
||||
|
||||
fn store_directory_name(store_name: &str) -> String {
|
||||
// Qdrant creates deeply nested files, so keep the physical path short on Windows.
|
||||
let digest = Sha256::digest(store_name.as_bytes());
|
||||
let short_hash = digest[..12]
|
||||
.iter()
|
||||
.map(|byte| format!("{byte:02x}"))
|
||||
.collect::<String>();
|
||||
format!("store_{short_hash}")
|
||||
}
|
||||
|
||||
fn qdrant_edge_base_path() -> QdrantEdgeResult<PathBuf> {
|
||||
let data_directory = DATA_DIRECTORY
|
||||
.get()
|
||||
@ -616,8 +675,38 @@ fn edge_optimizers_config() -> EdgeOptimizersConfig {
|
||||
}
|
||||
}
|
||||
|
||||
fn has_existing_store(path: &Path) -> bool {
|
||||
path.join("edge_config.json").exists() || path.join("segments").exists()
|
||||
fn store_is_initialized(path: &Path, store_name: &str) -> QdrantEdgeResult<bool> {
|
||||
if !path.join("edge_config.json").is_file() || !path.join("segments").is_dir() {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
let marker_path = path.join(STORE_INITIALIZATION_MARKER);
|
||||
if !marker_path.exists() {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
let initialized_store_name = fs::read_to_string(&marker_path).map_err(|error| {
|
||||
format!("Failed to read vector store initialization marker '{}': {error}", marker_path.display())
|
||||
})?;
|
||||
if initialized_store_name != store_name {
|
||||
return Err(format!("Vector store path collision at '{}': expected store '{}', but the path belongs to '{}'.", path.display(), store_name, initialized_store_name).into());
|
||||
}
|
||||
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
fn write_store_initialization_marker(path: &Path, store_name: &str) -> std::io::Result<()> {
|
||||
let marker_path = path.join(STORE_INITIALIZATION_MARKER);
|
||||
let temporary_marker_path = path.join(STORE_INITIALIZATION_MARKER_TEMP);
|
||||
fs::write(&temporary_marker_path, store_name)?;
|
||||
fs::rename(temporary_marker_path, marker_path)
|
||||
}
|
||||
|
||||
fn remove_partial_store(path: &Path) -> String {
|
||||
match fs::remove_dir_all(path) {
|
||||
Ok(()) => String::new(),
|
||||
Err(error) => format!(" The incomplete store could not be removed: {error}"),
|
||||
}
|
||||
}
|
||||
|
||||
fn validate_vector_size(vector_size: usize) -> QdrantEdgeResult<()> {
|
||||
|
||||
@ -212,6 +212,15 @@ fn handle_tokenizer_validate(path: &PathBuf) -> Result<usize, String> {
|
||||
}
|
||||
|
||||
pub fn get_token_count(text: &str) -> Result<usize, String> {
|
||||
get_token_count_internal(text, true)
|
||||
}
|
||||
|
||||
pub fn get_segment_token_count(text: &str) -> Result<usize, String> {
|
||||
// Special tokens belong to the final encoding and would inflate sums across many segments.
|
||||
get_token_count_internal(text, false)
|
||||
}
|
||||
|
||||
fn get_token_count_internal(text: &str, add_special_tokens: bool) -> Result<usize, String> {
|
||||
if text.trim().is_empty() {
|
||||
return Ok(0);
|
||||
}
|
||||
@ -234,7 +243,7 @@ pub fn get_token_count(text: &str) -> Result<usize, String> {
|
||||
return Err(unavailable_with_status_update("Tokenizer not initialized."));
|
||||
}
|
||||
};
|
||||
let token_count = match tokenizer.encode(text, true) {
|
||||
let token_count = match tokenizer.encode(text, add_special_tokens) {
|
||||
Ok(enc) => enc.len(),
|
||||
Err(e) => {
|
||||
let reason = format!("Failed to tokenize text: {e}");
|
||||
|
||||
Loading…
Reference in New Issue
Block a user