fixed tokenization issue for long files

app/MindWork AI Studio/Assistants/I18N/allTexts.lua

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

app/MindWork AI Studio/Tools/ContentStreamSseEvent.cs

@@ -12,4 +12,7 @@ public sealed class ContentStreamSseEvent
 
     [JsonPropertyName("metadata")] 
     public ContentStreamSseMetadata? Metadata { get; init; }
+
+    [JsonPropertyName("token_count")]
+    public int? TokenCount { get; init; }
 }

app/MindWork AI Studio/Tools/ContentStreamSseHandler.cs

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

app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.Files.cs

@@ -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);
+        var tokenCount = estimatedTokenCount;
-        if (tokenCount <= options.MaxChunkTokenLength)
+        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);
+            tokenCount = await this.GetEmbeddingTokenCountAsync(embeddingProvider, textWithOverlap, token);
-            yield break;
+            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 minimumCandidateUnitCount = 1;
-        var high = units.Count - startIndex;
+        var availableUnitCount = units.Count - startUnitIndex;
-        var best = 0;
+        var maximumCandidateUnitCount = availableUnitCount;
+        var largestValidUnitCount = 0;
 
-        while (low <= high)
+        if (estimatedUnitTokenCounts is not null)
         {
-            token.ThrowIfCancellationRequested();
+            maximumCandidateUnitCount = 0;
-
+            var cumulativeEstimatedTokenCount = 0L;
-            var mid = low + (high - low) / 2;
+            for (var unitIndex = startUnitIndex; unitIndex < units.Count; unitIndex++)
-            var candidate = AddOverlapPrefix(string.Concat(units.Skip(startIndex).Take(mid)).Trim(), overlapPrefix);
-            var tokenCount = await this.GetEmbeddingTokenCountAsync(embeddingProvider, candidate, token);
-            if (tokenCount <= maxChunkTokenLength)
             {
-                best = mid;
+                cumulativeEstimatedTokenCount += estimatedUnitTokenCounts[unitIndex];
-                low = mid + 1;
+                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 estimatedChunkLength = Math.Max(1L, (long)text.Length * options.MaxChunkTokenLength / estimatedTokenCount.Value);
-                var candidate = AddOverlapPrefix(text[startIndex..mid].Trim(), overlapPrefix);
+                maximumCandidateEndIndex = (int)Math.Min(text.Length, startIndex + estimatedChunkLength);
-                var tokenCount = await this.GetEmbeddingTokenCountAsync(embeddingProvider, candidate, token);
+            }
-                if (tokenCount <= options.MaxChunkTokenLength)
+
+            while (true)
+            {
+                var minimumCandidateEndIndex = bestEndIndex + 1;
+                var currentMaximumCandidateEndIndex = maximumCandidateEndIndex;
+                while (minimumCandidateEndIndex <= currentMaximumCandidateEndIndex)
                 {
-                    bestEndIndex = mid;
+                    var candidateEndIndex = minimumCandidateEndIndex + (currentMaximumCandidateEndIndex - minimumCandidateEndIndex) / 2;
-                    low = mid + 1;
+                    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));

app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.cs

@@ -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);
+                await this.CleanupFailedFileAsync(embeddingState, vectorStore, dataSource, collectionName, file.FullName, optimizationTracker, token);
-                optimizationTracker.MarkChanged();
-                await embeddingState.DeleteFileAsync(dataSource.Id, file.FullName, 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,

app/MindWork AI Studio/Tools/Services/RustService.Retrieval.cs

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

runtime/src/file_data.rs

@@ -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) => {
+                                Ok(chunk) => {
-                                    chunk.set_stream_id(id_ref);
+                                    let chunks = chunk.into_bounded_text_segments();
-                                    yield Ok(Event::default().json_data(&chunk).unwrap_or_else(|e| Event::default().data(format!("Error: {e}"))));
+
+                                    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}'");

runtime/src/qdrant_edge_database.rs

@@ -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)?;
+        if path.exists() && !is_initialized {
-        let shard = if has_existing_store(&path) {
+            warn!(Source = "Qdrant Edge"; "Removing incompletely initialized vector store '{}' before recreating it.", store_name);
-            EdgeShard::load(&path, None)?
+            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)?;
+            fs::create_dir_all(&path).map_err(|error| {
-            EdgeShard::new(&path, edge_config(vector_size))?
+                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 path.exists() && !is_initialized {
-        if !has_existing_store(&path) {
+            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 {
+fn store_is_initialized(path: &Path, store_name: &str) -> QdrantEdgeResult<bool> {
-    path.join("edge_config.json").exists() || path.join("segments").exists()
+    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<()> {

runtime/src/tokenizer.rs

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