Fixed page numbers being guessed from the text instead of read from the metadata

This commit is contained in:
Thorsten Sommer committed 2026-09-15 13:05:20 +02:00
1 parent 6ce7d856a3
commit 01e09804f6
9 files changed
+285 -49

No files matched your search

@@ -1,3 +1,9 @@
namespace AIStudio.Tools.Services;
public sealed record ArbitraryFileDataSegment(string Content, int TokenCount);
/// <summary>
/// One piece of an extracted file, as the runtime delivered it.
/// </summary>
/// <param name="Content">The extracted text.</param>
/// <param name="TokenCount">The number of tokens of that text.</param>
/// <param name="PageNumber">The page that text came from, or null when it has none. Presentations and spreadsheets have none.</param>
public sealed record ArbitraryFileDataSegment(string Content, int TokenCount, int? PageNumber);
@@ -16,6 +16,22 @@ public sealed partial class DataSourceEmbeddingService
internal const int DEFAULT_CHUNK_OVERLAP_TOKEN_LENGTH = 300;
private const bool IMAGE_EMBEDDING_ENABLED = false;
/// <summary>
/// What this build writes next to a chunk besides its text. Raise it whenever that changes.
/// </summary>
/// <remarks>
/// A stored chunk keeps the metadata of the run which wrote it, and nothing recomputes it: the
/// fingerprint of a file says whether the file changed, not whether we got better at reading
/// it. Raising this number makes the embedding signature differ, which drops the index and
/// builds it again — the only way corrected page numbers reach a data source somebody indexed
/// earlier.
///
/// Version 2: the page of a chunk is taken from the runtime metadata instead of being read back
/// out of the chunk text, which is what left Word and OpenDocument files, and passages
/// continuing across a page break, without a page.
/// </remarks>
private const string CHUNK_METADATA_VERSION = "2";
private enum RagFileIndexingDecision
{
INDEXABLE,
@@ -23,10 +39,23 @@ public sealed partial class DataSourceEmbeddingService
UNSUPPORTED,
}
private sealed record ExtractedFileSegment(string Text, int? TokenCount);
private sealed record ExtractedFileSegment(string Text, int? TokenCount, int? PageNumber);
private sealed record ExtractedFileContent(string Text, IReadOnlyList<ExtractedFileSegment> SourceSegments);
/// <summary>
/// One chunk as the chunking produced it, together with the page it starts on.
/// </summary>
/// <remarks>
/// The page is carried rather than read back out of the chunk text. The runtime states it, and
/// the chunking knows which source segment a chunk begins in, so nothing has to be derived from
/// a marker in the text — which is what used to leave Word files and continued passages without
/// a page.
/// </remarks>
/// <param name="Text">The chunk itself, overlap prefix included.</param>
/// <param name="PageNumber">The page the chunk's own content starts on, or null when it has none.</param>
private sealed record EmbeddingChunk(string Text, int? PageNumber);
private sealed record EmbeddingChunkDraft(string ChunkId, string Text, int ChunkIndex, int? PageNumber);
private sealed record ChunkingOptions(int MaxChunkTokenLength, int OverlapTokenLength);
@@ -37,7 +66,7 @@ public sealed partial class DataSourceEmbeddingService
private sealed record DataSourceMetadataSnapshot(string SourceHash, IReadOnlyDictionary<string, string> FileHashes);
private async IAsyncEnumerable<string> StreamEmbeddingChunksAsync(string filePath, IDataSource dataSource, EmbeddingProvider embeddingProvider, [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token)
private async IAsyncEnumerable<EmbeddingChunk> StreamEmbeddingChunksAsync(string filePath, IDataSource dataSource, EmbeddingProvider embeddingProvider, [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token)
{
var options = this.GetChunkingOptions(dataSource, embeddingProvider);
var strategy = this.GetChunkingStrategy(filePath);
@@ -55,26 +84,31 @@ public sealed partial class DataSourceEmbeddingService
{
var normalized = NormalizeChunkSegment(segment.Content);
if (!string.IsNullOrWhiteSpace(normalized))
segments.Add(new(normalized, segment.TokenCount));
segments.Add(new(normalized, segment.TokenCount, segment.PageNumber));
}
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)
private async IAsyncEnumerable<EmbeddingChunk> SplitByChunkingStrategyAsync(ExtractedFileContent content, ChunkingStrategy strategy, ChunkingOptions options, EmbeddingProvider embeddingProvider, [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token)
{
var estimatedTokenCount = SumTokenCounts(content.SourceSegments);
await foreach (var chunk in this.SplitTextByRulesAsync(content.Text, content.SourceSegments, strategy, 0, options, embeddingProvider, token, estimatedTokenCount: estimatedTokenCount))
// The whole text starts where the first segment starts, so that is the page it is on until
// the splitting reaches a segment boundary:
var firstPageNumber = content.SourceSegments.Count > 0 ? content.SourceSegments[0].PageNumber : null;
await foreach (var chunk in this.SplitTextByRulesAsync(content.Text, content.SourceSegments, strategy, 0, options, embeddingProvider, firstPageNumber, token, estimatedTokenCount: estimatedTokenCount))
yield return chunk;
}
private async IAsyncEnumerable<string> SplitTextByRulesAsync(
private async IAsyncEnumerable<EmbeddingChunk> SplitTextByRulesAsync(
string text,
IReadOnlyList<ExtractedFileSegment> sourceSegments,
ChunkingStrategy strategy,
int ruleIndex,
ChunkingOptions options,
EmbeddingProvider embeddingProvider,
int? currentPageNumber,
[System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token,
string requiredOverlapPrefix = "",
int? estimatedTokenCount = null)
@@ -91,14 +125,14 @@ public sealed partial class DataSourceEmbeddingService
tokenCount = await this.GetEmbeddingTokenCountAsync(embeddingProvider, textWithOverlap, token);
if (tokenCount <= options.MaxChunkTokenLength)
{
yield return textWithOverlap;
yield return new(textWithOverlap, currentPageNumber);
yield break;
}
}
if (ruleIndex >= strategy.Rules.Count)
{
await foreach (var hardChunk in this.SplitTextByHardCutAsync(text, options, embeddingProvider, token, requiredOverlapPrefix, estimatedTokenCount))
await foreach (var hardChunk in this.SplitTextByHardCutAsync(text, options, embeddingProvider, currentPageNumber, token, requiredOverlapPrefix, estimatedTokenCount))
yield return hardChunk;
yield break;
@@ -107,7 +141,7 @@ 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, estimatedTokenCount))
await foreach (var hardChunk in this.SplitTextByHardCutAsync(text, options, embeddingProvider, currentPageNumber, token, requiredOverlapPrefix, estimatedTokenCount))
yield return hardChunk;
yield break;
@@ -116,7 +150,7 @@ public sealed partial class DataSourceEmbeddingService
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, estimatedTokenCount))
await foreach (var chunk in this.SplitTextByRulesAsync(text, sourceSegments, strategy, ruleIndex + 1, options, embeddingProvider, currentPageNumber, token, requiredOverlapPrefix, estimatedTokenCount))
yield return chunk;
yield break;
@@ -135,6 +169,15 @@ public sealed partial class DataSourceEmbeddingService
var overlapPrefix = requiredOverlapPrefix;
var unitTokenCounts = EstimateSplitUnitTokenCounts(units, sourceSegments, rule.UsesSourceSegmentCounts, estimatedTokenCount);
//
// The first rule of every strategy cuts along the segments the runtime delivered, so there
// a unit is a segment and carries that segment's page. Every later rule cuts inside a
// single segment, where all units share the page they were handed. This is what ties a
// chunk to a page without anybody reading the text.
//
var unitsAreSourceSegments = rule.UsesSourceSegmentCounts && sourceSegments.Count == units.Count;
int? PageOfUnit(int unitIndex) => unitsAreSourceSegments ? sourceSegments[unitIndex].PageNumber ?? currentPageNumber : currentPageNumber;
while (index < units.Count)
{
token.ThrowIfCancellationRequested();
@@ -145,8 +188,14 @@ public sealed partial class DataSourceEmbeddingService
var rawChunk = string.Concat(units.Skip(index).Take(unitCount)).Trim();
var chunk = AddOverlapPrefix(rawChunk, overlapPrefix);
overlapPrefix = string.Empty;
//
// The page of the first unit this chunk covers, not of the overlap prefix in front
// of it: the prefix repeats what the chunk before already said, while the page has
// to name where this chunk's own content begins.
//
if (!string.IsNullOrWhiteSpace(chunk))
yield return chunk;
yield return new(chunk, PageOfUnit(index));
var nextIndex = index + unitCount;
if (nextIndex >= units.Count)
@@ -178,9 +227,10 @@ public sealed partial class DataSourceEmbeddingService
string? lastSplitUnit = null;
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))
var unitPageNumber = PageOfUnit(index);
await foreach (var splitUnit in this.SplitTextByRulesAsync(units[index], [new(units[index], unitTokenCount, unitPageNumber)], strategy, ruleIndex + 1, options, embeddingProvider, unitPageNumber, token, overlapPrefix, unitTokenCount))
{
lastSplitUnit = splitUnit;
lastSplitUnit = splitUnit.Text;
yield return splitUnit;
}
@@ -372,10 +422,15 @@ public sealed partial class DataSourceEmbeddingService
return bestStartIndex <= chunkStartIndex ? chunkEndIndex : bestStartIndex;
}
private async IAsyncEnumerable<string> SplitTextByHardCutAsync(
/// <remarks>
/// The hard cut is only ever reached inside a single piece of text which no rule could split
/// any further, so every chunk it produces sits on the page that piece was handed.
/// </remarks>
private async IAsyncEnumerable<EmbeddingChunk> SplitTextByHardCutAsync(
string text,
ChunkingOptions options,
EmbeddingProvider embeddingProvider,
int? currentPageNumber,
[System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token,
string requiredOverlapPrefix = "",
int? estimatedTokenCount = null)
@@ -455,7 +510,7 @@ public sealed partial class DataSourceEmbeddingService
var chunk = AddOverlapPrefix(text[startIndex..bestEndIndex].Trim(), overlapPrefix);
if (!string.IsNullOrWhiteSpace(chunk))
yield return chunk;
yield return new(chunk, currentPageNumber);
if (bestEndIndex >= text.Length)
yield break;
@@ -934,6 +989,7 @@ public sealed partial class DataSourceEmbeddingService
private string BuildEmbeddingSignature(IDataSource dataSource, EmbeddingProvider embeddingProvider, ChunkingOptions chunkingOptions)
{
return string.Join('|',
CHUNK_METADATA_VERSION,
embeddingProvider.Id,
embeddingProvider.UsedLLMProvider,
embeddingProvider.Model.Id,
@@ -1086,14 +1142,6 @@ public sealed partial class DataSourceEmbeddingService
return string.IsNullOrWhiteSpace(extension) ? "unknown" : extension;
}
private static int? TryExtractPageNumber(string chunk)
{
var match = Regex.Match(chunk, @"^\s*#\s+Page\s+(\d+)\b", RegexOptions.CultureInvariant | RegexOptions.IgnoreCase);
return match.Success && int.TryParse(match.Groups[1].Value, out var pageNumber) && pageNumber > 0
? pageNumber
: null;
}
private string CreatePointId(string dataSourceId, string fingerprint, int chunkIndex) =>
CreateStableGuid($"{dataSourceId}:chunk:{fingerprint}:{chunkIndex}");
@@ -843,7 +843,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
await foreach (var chunk in this.StreamEmbeddingChunksAsync(file.FullName, dataSource, embeddingProvider, token))
{
batch.Add(new(this.CreatePointId(dataSource.Id, fingerprint, totalChunkCount), chunk, totalChunkCount, TryExtractPageNumber(chunk)));
batch.Add(new(this.CreatePointId(dataSource.Id, fingerprint, totalChunkCount), chunk.Text, totalChunkCount, chunk.PageNumber));
totalChunkCount++;
if (batch.Count >= embeddingBatchSize)
@@ -278,7 +278,7 @@ public sealed partial class RustService
{
if (segment.TokenCount is { } tokenCount)
{
yield return new(segment.Content, tokenCount);
yield return new(segment.Content, tokenCount, segment.PageNumber);
continue;
}
@@ -291,7 +291,7 @@ public sealed partial class RustService
var countedSegment = await this.GetTokenCount(embeddingProvider, segment.Content, token);
if (countedSegment is { Success: true } counted)
{
yield return new(segment.Content, counted.TokenCount);
yield return new(segment.Content, counted.TokenCount, segment.PageNumber);
continue;
}
@@ -303,7 +303,7 @@ public sealed partial class RustService
}
}
private async IAsyncEnumerable<(string Content, int? TokenCount)> StreamArbitraryFileDataCore(
private async IAsyncEnumerable<(string Content, int? TokenCount, int? PageNumber)> StreamArbitraryFileDataCore(
string path,
bool extractImages,
bool includeTokenCount,
@@ -420,12 +420,13 @@ public sealed partial class RustService
}
//
// The count comes from the processed event, not from the event which was just read:
// a reader may hold content back across several events, and the count of the content
// it releases is the count of that content, not of the event that released it.
// The count and the page come from the processed event, not from the event which
// was just read: a reader may hold content back across several events, and the
// count and page of the content it releases describe that content, not the event
// that released it.
//
if (!string.IsNullOrWhiteSpace(processedEvent.Content))
yield return (processedEvent.Content, processedEvent.TokenCount);
yield return (processedEvent.Content, processedEvent.TokenCount, processedEvent.PageNumber);
}
}
finally
@@ -434,7 +435,7 @@ public sealed partial class RustService
}
if (finalContentChunk is { } pendingContent && !string.IsNullOrWhiteSpace(pendingContent.Content))
yield return (pendingContent.Content, pendingContent.TokenCount);
yield return (pendingContent.Content, pendingContent.TokenCount, pendingContent.PageNumber);
if (promptInjectionRedactedCount is 0)
yield break;