diff --git a/app/MindWork AI Studio/Assistants/I18N/allTexts.lua b/app/MindWork AI Studio/Assistants/I18N/allTexts.lua index 43074e21..2dad213c 100644 --- a/app/MindWork AI Studio/Assistants/I18N/allTexts.lua +++ b/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." diff --git a/app/MindWork AI Studio/Tools/ContentStreamSseEvent.cs b/app/MindWork AI Studio/Tools/ContentStreamSseEvent.cs index 2c47551f..02d69c0d 100644 --- a/app/MindWork AI Studio/Tools/ContentStreamSseEvent.cs +++ b/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; } } \ No newline at end of file diff --git a/app/MindWork AI Studio/Tools/ContentStreamSseHandler.cs b/app/MindWork AI Studio/Tools/ContentStreamSseHandler.cs index 247d3ebf..d1ea0fa4 100644 --- a/app/MindWork AI Studio/Tools/ContentStreamSseHandler.cs +++ b/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; } -} \ No newline at end of file +} diff --git a/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.Files.cs b/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.Files.cs index 28a3539f..52a50b65 100644 --- a/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.Files.cs +++ b/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 SourceSegments); + private sealed record ExtractedFileSegment(string Text, int? TokenCount); + + private sealed record ExtractedFileContent(string Text, IReadOnlyList 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 Rules); - private sealed record ChunkingRule(string Name, Func, IReadOnlyList>? Split); + private sealed record ChunkingRule(string Name, Func, IReadOnlyList>? Split, bool UsesSourceSegmentCounts = false); private sealed record DataSourceMetadataSnapshot(string SourceHash, IReadOnlyDictionary 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 ReadExtractedFileContentAsync(string filePath, CancellationToken token) + private async Task ReadExtractedFileContentAsync(string filePath, EmbeddingProvider embeddingProvider, CancellationToken token) { - var segments = new List(); + var segments = new List(); - 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 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 SplitTextByRulesAsync( string text, - IReadOnlyList sourceSegments, + IReadOnlyList 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 FindLargestUnitCountWithinMaxChunkLengthAsync(IReadOnlyList units, int startIndex, EmbeddingProvider embeddingProvider, int maxChunkTokenLength, CancellationToken token, string overlapPrefix = "") + private async Task FindLargestUnitCountWithinMaxChunkLengthAsync(IReadOnlyList units, IReadOnlyList? 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 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? EstimateSplitUnitTokenCounts( + IReadOnlyList units, + IReadOnlyList 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(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 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 + { + 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)); diff --git a/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.cs b/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.cs index 2a8837b2..42ccb0bb 100644 --- a/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.cs +++ b/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); - 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(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, diff --git a/app/MindWork AI Studio/Tools/Services/RustService.Retrieval.cs b/app/MindWork AI Studio/Tools/Services/RustService.Retrieval.cs index a7275208..38b1e01f 100644 --- a/app/MindWork AI Studio/Tools/Services/RustService.Retrieval.cs +++ b/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 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 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 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) diff --git a/runtime/src/file_data.rs b/runtime/src/file_data.rs index ca8a1671..d783540d 100644 --- a/runtime/src/file_data.rs +++ b/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, } 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 { + 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 = std::result::Result>; type ChunkStream = Pin> + 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 @@ -145,11 +171,21 @@ pub async fn extract_data( match stream_result { Ok(mut stream) => { - while let Some(chunk) = stream.next().await { + 'stream_chunks: while let Some(chunk) = stream.next().await { match chunk { - Ok(mut chunk) => { - chunk.set_stream_id(id_ref); - yield Ok(Event::default().json_data(&chunk).unwrap_or_else(|e| Event::default().data(format!("Error: {e}")))); + Ok(chunk) => { + let chunks = chunk.into_bounded_text_segments(); + + for mut chunk in chunks { + chunk.set_stream_id(id_ref); + if query.include_token_count { + if let Err(e) = chunk.set_token_count() { + yield Ok(Event::default().json_data(format!("Error counting tokens: {e}")).unwrap_or_else(|_| Event::default().data(format!("Error counting tokens: {e}")))); + break 'stream_chunks; + } + } + yield Ok(Event::default().json_data(&chunk).unwrap_or_else(|e| Event::default().data(format!("Error: {e}")))); + } }, Err(e) => { @@ -175,6 +211,33 @@ pub async fn extract_data( Sse::new(stream) } +fn bounded_text_segment_ranges(content: &str) -> Vec<(usize, usize)> { + let mut ranges = Vec::new(); + let mut start = 0; + + while start < content.len() { + let remaining = &content[start..]; + let Some(maximum_end_offset) = remaining + .char_indices() + .nth(MAX_TEXT_SEGMENT_LENGTH_IN_CHARS) + .map(|(index, _)| index) + else { + ranges.push((start, content.len())); + break; + }; + + let end = start + maximum_end_offset; + ranges.push((start, end)); + start = end; + } + + if ranges.is_empty() { + ranges.push((0, 0)); + } + + ranges +} + async fn stream_data(file_path: &str, extract_images: bool) -> Result { if !Path::new(file_path).exists() { error!("File does not exist: '{file_path}'"); diff --git a/runtime/src/qdrant_edge_database.rs b/runtime/src/qdrant_edge_database.rs index 92acab01..9be183d2 100644 --- a/runtime/src/qdrant_edge_database.rs +++ b/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 = Result>; @@ -182,21 +185,52 @@ impl QdrantEdgeDatabase { fn store_path(&self, store_name: &str) -> QdrantEdgeResult { validate_store_name(store_name)?; - Ok(self.base_path.join("stores").join(store_name)) + Ok(self.base_path.join("stores").join(store_directory_name(store_name))) } // To ensure a shard exists and that you can insert a vector fn get_or_create_store(&mut self, store_name: &str, vector_size: usize) -> QdrantEdgeResult<&EdgeShard> { + let path = self.store_path(store_name)?; + let is_initialized = store_is_initialized(&path, store_name)?; if self.shards.contains_key(store_name) { - return Ok(self.shards.get(store_name).unwrap()); + if is_initialized { + return Ok(self.shards.get(store_name).unwrap()); + } + + warn!(Source = "Qdrant Edge"; "Removing stale cached vector store '{}' because its initialized data directory no longer exists.", store_name); + self.shards.remove(store_name); } - let path = self.store_path(store_name)?; - let shard = if has_existing_store(&path) { - EdgeShard::load(&path, None)? + if path.exists() && !is_initialized { + warn!(Source = "Qdrant Edge"; "Removing incompletely initialized vector store '{}' before recreating it.", store_name); + fs::remove_dir_all(&path).map_err(|error| { + format!("Failed to remove incomplete vector store '{store_name}' at '{}': {error}", path.display()) + })?; + } + + let shard = if is_initialized { + EdgeShard::load(&path, None).map_err(|error| { + format!("Failed to load vector store '{store_name}' from '{}': {error}", path.display()) + })? } else { - fs::create_dir_all(&path)?; - EdgeShard::new(&path, edge_config(vector_size))? + fs::create_dir_all(&path).map_err(|error| { + format!("Failed to create directory for vector store '{store_name}' at '{}': {error}", path.display()) + })?; + let shard = match EdgeShard::new(&path, edge_config(vector_size)) { + Ok(shard) => shard, + Err(error) => { + let cleanup_issue = remove_partial_store(&path); + return Err(format!("Failed to create vector store '{store_name}' at '{}': {error}{cleanup_issue}", path.display()).into()); + }, + }; + + if let Err(error) = write_store_initialization_marker(&path, store_name) { + drop(shard); + let cleanup_issue = remove_partial_store(&path); + return Err(format!("Failed to finalize vector store '{store_name}' at '{}': {error}{cleanup_issue}", path.display()).into()); + } + + shard }; self.shards.insert(store_name.to_string(), shard); @@ -205,16 +239,31 @@ impl QdrantEdgeDatabase { // To check whether a shard exists so you can delete a file from it fn get_existing_store(&mut self, store_name: &str) -> QdrantEdgeResult> { + let path = self.store_path(store_name)?; + let is_initialized = store_is_initialized(&path, store_name)?; if self.shards.contains_key(store_name) { - return Ok(self.shards.get(store_name)); + if is_initialized { + return Ok(self.shards.get(store_name)); + } + + warn!(Source = "Qdrant Edge"; "Removing stale cached vector store '{}' because its initialized data directory no longer exists.", store_name); + self.shards.remove(store_name); } - let path = self.store_path(store_name)?; - if !has_existing_store(&path) { + if path.exists() && !is_initialized { + warn!(Source = "Qdrant Edge"; "Removing incompletely initialized vector store '{}' before continuing.", store_name); + fs::remove_dir_all(&path).map_err(|error| { + format!("Failed to remove incomplete vector store '{store_name}' at '{}': {error}", path.display()) + })?; + } + + if !is_initialized { return Ok(None); } - let shard = EdgeShard::load(&path, None)?; + let shard = EdgeShard::load(&path, None).map_err(|error| { + format!("Failed to load vector store '{store_name}' from '{}': {error}", path.display()) + })?; self.shards.insert(store_name.to_string(), shard); Ok(self.shards.get(store_name)) } @@ -224,7 +273,7 @@ impl QdrantEdgeDatabase { let stores_count = if stores_path.exists() { fs::read_dir(stores_path)? .filter_map(Result::ok) - .filter(|entry| entry.path().is_dir()) + .filter(|entry| entry.path().join(STORE_INITIALIZATION_MARKER).is_file()) .count() } else { 0 @@ -339,6 +388,16 @@ impl QdrantEdgeDatabase { } } +fn store_directory_name(store_name: &str) -> String { + // Qdrant creates deeply nested files, so keep the physical path short on Windows. + let digest = Sha256::digest(store_name.as_bytes()); + let short_hash = digest[..12] + .iter() + .map(|byte| format!("{byte:02x}")) + .collect::(); + format!("store_{short_hash}") +} + fn qdrant_edge_base_path() -> QdrantEdgeResult { let data_directory = DATA_DIRECTORY .get() @@ -616,8 +675,38 @@ fn edge_optimizers_config() -> EdgeOptimizersConfig { } } -fn has_existing_store(path: &Path) -> bool { - path.join("edge_config.json").exists() || path.join("segments").exists() +fn store_is_initialized(path: &Path, store_name: &str) -> QdrantEdgeResult { + 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<()> { diff --git a/runtime/src/tokenizer.rs b/runtime/src/tokenizer.rs index bb7ca0c9..6d554760 100644 --- a/runtime/src/tokenizer.rs +++ b/runtime/src/tokenizer.rs @@ -212,6 +212,15 @@ fn handle_tokenizer_validate(path: &PathBuf) -> Result { } pub fn get_token_count(text: &str) -> Result { + get_token_count_internal(text, true) +} + +pub fn get_segment_token_count(text: &str) -> Result { + // 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 { if text.trim().is_empty() { return Ok(0); } @@ -234,7 +243,7 @@ pub fn get_token_count(text: &str) -> Result { 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}");