fixed tokenization issue for long files

2026-10-07 03:29:40 +00:00 · 2026-08-10 18:21:31 +02:00
1 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)