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,17 +1,21 @@
namespace AIStudio.Tools; namespace AIStudio.Tools;
/// <summary> /// <summary>
/// Content which a reader held back, together with the token count of exactly that content. /// Content which a reader held back, together with the token count and the page of exactly that
/// content.
/// </summary> /// </summary>
/// <remarks> /// <remarks>
/// Readers which assemble a page or a slide from several stream events cannot pass their content /// Readers which assemble a page or a slide from several stream events cannot pass their content
/// on right away. Its token count has to travel with it: the count describes the content, not the /// on right away. Its token count has to travel with it: the count describes the content, not the
/// event which happened to arrive at the moment the content was released. Keeping the two together /// event which happened to arrive at the moment the content was released. Keeping the two together
/// is what stops a page from being sized by the text of the page after it. /// is what stops a page from being sized by the text of the page after it. The page number travels
/// for the very same reason, and because a number the runtime already stated must not be derived
/// from the text again further down the line.
/// </remarks> /// </remarks>
/// <param name="Content">The assembled content.</param> /// <param name="Content">The assembled content.</param>
/// <param name="TokenCount">The number of tokens of that content, or null when it is unknown.</param> /// <param name="TokenCount">The number of tokens of that content, or null when it is unknown.</param>
public readonly record struct ContentStreamPendingContent(string Content, int? TokenCount) /// <param name="PageNumber">The page that content came from, or null when it has none.</param>
public readonly record struct ContentStreamPendingContent(string Content, int? TokenCount, int? PageNumber = null)
{ {
/// <summary> /// <summary>
/// Adds up two token counts, where an unknown count makes the sum unknown as well. /// Adds up two token counts, where an unknown count makes the sum unknown as well.
@@ -12,7 +12,8 @@ namespace AIStudio.Tools;
/// <param name="Error">The reported failure, or null when the event was processed successfully.</param> /// <param name="Error">The reported failure, or null when the event was processed successfully.</param>
/// <param name="PromptInjection">What the runtime filtered out of the content, or null when it filtered nothing.</param> /// <param name="PromptInjection">What the runtime filtered out of the content, or null when it filtered nothing.</param>
/// <param name="TokenCount">The number of tokens of the content, or null when it is unknown.</param> /// <param name="TokenCount">The number of tokens of the content, or null when it is unknown.</param>
public readonly record struct ContentStreamProcessedEvent(string? Content, ContentStreamErrorDetails? Error, ContentStreamPromptInjectionDetails? PromptInjection = null, int? TokenCount = null) /// <param name="PageNumber">The page the content came from, or null when it has none.</param>
public readonly record struct ContentStreamProcessedEvent(string? Content, ContentStreamErrorDetails? Error, ContentStreamPromptInjectionDetails? PromptInjection = null, int? TokenCount = null, int? PageNumber = null)
{ {
/// <summary> /// <summary>
/// An event which neither produced content nor reported a failure. /// An event which neither produced content nor reported a failure.
@@ -20,16 +21,18 @@ public readonly record struct ContentStreamProcessedEvent(string? Content, Conte
public static readonly ContentStreamProcessedEvent NOTHING = new(null, null); public static readonly ContentStreamProcessedEvent NOTHING = new(null, null);
/// <summary> /// <summary>
/// An event which produced content, with the token count of that very content. /// An event which produced content, with the token count and the page of that very content.
/// </summary> /// </summary>
/// <remarks> /// <remarks>
/// The count travels with the content because a reader may hold content back across several /// The count travels with the content because a reader may hold content back across several
/// events: pairing it with the count of the event which released it would size it by the /// events: pairing it with the count of the event which released it would size it by the
/// wrong text. /// wrong text. The page travels along for the same reason, and so that whoever indexes the
/// content is told where it came from instead of having to read it back out of the text.
/// </remarks> /// </remarks>
/// <param name="content">The content to append.</param> /// <param name="content">The content to append.</param>
/// <param name="tokenCount">The number of tokens of that content, or null when it is unknown.</param> /// <param name="tokenCount">The number of tokens of that content, or null when it is unknown.</param>
public static ContentStreamProcessedEvent FromContent(string? content, int? tokenCount = null) => new(content, null, TokenCount: tokenCount); /// <param name="pageNumber">The page that content came from, or null when it has none.</param>
public static ContentStreamProcessedEvent FromContent(string? content, int? tokenCount = null, int? pageNumber = null) => new(content, null, TokenCount: tokenCount, PageNumber: pageNumber);
public static ContentStreamProcessedEvent FromError(ContentStreamErrorDetails? error) => new(null, error); public static ContentStreamProcessedEvent FromError(ContentStreamErrorDetails? error) => new(null, error);
@@ -19,13 +19,19 @@ public static class ContentStreamSseHandler
case ContentStreamTextMetadata: case ContentStreamTextMetadata:
return ContentStreamProcessedEvent.FromContent(sseEvent.Content, sseEvent.TokenCount); return ContentStreamProcessedEvent.FromContent(sseEvent.Content, sseEvent.TokenCount);
//
// The heading tells the AI which page it is reading. The number is handed on
// separately as well, because whoever indexes this content needs it as a
// number: reading it back out of the heading would mean guessing at something
// the runtime already stated.
//
case ContentStreamPdfMetadata pdfMetadata: case ContentStreamPdfMetadata pdfMetadata:
var pageNumber = pdfMetadata.Pdf?.PageNumber ?? 0; var pageNumber = pdfMetadata.Pdf?.PageNumber ?? 0;
return ContentStreamProcessedEvent.FromContent($""" return ContentStreamProcessedEvent.FromContent($"""
# Page {pageNumber} # Page {pageNumber}
{sseEvent.Content} {sseEvent.Content}
""", sseEvent.TokenCount); """, sseEvent.TokenCount, pageNumber > 0 ? pageNumber : null);
case ContentStreamSpreadsheetMetadata spreadsheetMetadata: case ContentStreamSpreadsheetMetadata spreadsheetMetadata:
var sheetName = spreadsheetMetadata.Spreadsheet?.SheetName; var sheetName = spreadsheetMetadata.Spreadsheet?.SheetName;
@@ -45,9 +51,10 @@ public static class ContentStreamSseHandler
// a page can follow its Markdown. Documents converted as a whole, e.g. by Pandoc, // a page can follow its Markdown. Documents converted as a whole, e.g. by Pandoc,
// carry no page number and are passed on unchanged. // carry no page number and are passed on unchanged.
// //
// The buffering is why the count comes back from the reader rather than from // The buffering is why the count and the page come back from the reader rather
// this event: the page which is released here arrived one event ago, and this // than from this event: the page which is released here arrived one event ago,
// event's count belongs to the page which is now being buffered. // and this event's count and number belong to the page which is now being
// buffered.
// //
case ContentStreamDocumentMetadata documentMetadata: case ContentStreamDocumentMetadata documentMetadata:
if (documentMetadata.Document?.PageNumber is not > 0) if (documentMetadata.Document?.PageNumber is not > 0)
@@ -55,7 +62,7 @@ public static class ContentStreamSseHandler
var documentManager = DOCUMENT_MANAGERS.GetOrAdd(sseEvent.StreamId!, _ => new()); var documentManager = DOCUMENT_MANAGERS.GetOrAdd(sseEvent.StreamId!, _ => new());
var documentContent = documentManager.AddPage(documentMetadata, sseEvent.Content, sseEvent.TokenCount, extractImages); var documentContent = documentManager.AddPage(documentMetadata, sseEvent.Content, sseEvent.TokenCount, extractImages);
return documentContent is null ? ContentStreamProcessedEvent.NOTHING : ContentStreamProcessedEvent.FromContent(documentContent.Value.Content, documentContent.Value.TokenCount); return documentContent is null ? ContentStreamProcessedEvent.NOTHING : ContentStreamProcessedEvent.FromContent(documentContent.Value.Content, documentContent.Value.TokenCount, documentContent.Value.PageNumber);
case ContentStreamImageMetadata: case ContentStreamImageMetadata:
return ContentStreamProcessedEvent.FromContent(sseEvent.Content, sseEvent.TokenCount); return ContentStreamProcessedEvent.FromContent(sseEvent.Content, sseEvent.TokenCount);
@@ -184,7 +191,9 @@ public static class ContentStreamSseHandler
/// <remarks> /// <remarks>
/// The readers which assemble pages or slides always keep the last one of them: nothing tells /// The readers which assemble pages or slides always keep the last one of them: nothing tells
/// them that no further image is coming. It is released here, and it carries its own token /// them that no further image is coming. It is released here, and it carries its own token
/// count, because a chunk without one cannot be sized by the caller. /// count, because a chunk without one cannot be sized by the caller. Only the page reader
/// states a page; a stream is read by one of them, so there is no second number to weigh
/// against.
/// </remarks> /// </remarks>
/// <param name="streamId">The stream to release and forget.</param> /// <param name="streamId">The stream to release and forget.</param>
/// <returns>The content which was held back, or null when there was none.</returns> /// <returns>The content which was held back, or null when there was none.</returns>
@@ -195,6 +204,7 @@ public static class ContentStreamSseHandler
var finalContentChunk = new StringBuilder(); var finalContentChunk = new StringBuilder();
int? tokenCount = 0; int? tokenCount = 0;
int? pageNumber = null;
if(SLIDE_MANAGERS.TryGetValue(streamId, out var slideManager) if(SLIDE_MANAGERS.TryGetValue(streamId, out var slideManager)
&& slideManager.GetAllSlidesInOrder() is { } slides && slideManager.GetAllSlidesInOrder() is { } slides
&& !string.IsNullOrWhiteSpace(slides.Content)) && !string.IsNullOrWhiteSpace(slides.Content))
@@ -209,6 +219,7 @@ public static class ContentStreamSseHandler
{ {
finalContentChunk.Append(page.Content); finalContentChunk.Append(page.Content);
tokenCount = ContentStreamPendingContent.AddTokenCounts(tokenCount, page.TokenCount); tokenCount = ContentStreamPendingContent.AddTokenCounts(tokenCount, page.TokenCount);
pageNumber = page.PageNumber;
} }
SLIDE_MANAGERS.TryRemove(streamId, out _); SLIDE_MANAGERS.TryRemove(streamId, out _);
@@ -217,6 +228,6 @@ public static class ContentStreamSseHandler
foreach (var key in CHUNKED_IMAGES.Keys.Where(k => k.StartsWith(imageIdPrefix, StringComparison.InvariantCultureIgnoreCase))) foreach (var key in CHUNKED_IMAGES.Keys.Where(k => k.StartsWith(imageIdPrefix, StringComparison.InvariantCultureIgnoreCase)))
CHUNKED_IMAGES.TryRemove(key, out _); CHUNKED_IMAGES.TryRemove(key, out _);
return finalContentChunk.Length > 0 ? new ContentStreamPendingContent(finalContentChunk.ToString(), tokenCount) : null; return finalContentChunk.Length > 0 ? new ContentStreamPendingContent(finalContentChunk.ToString(), tokenCount, pageNumber) : null;
} }
} }
@@ -10,6 +10,7 @@ public sealed class DocumentManager
{ {
private StringBuilder? currentPageContent; private StringBuilder? currentPageContent;
private int? currentPageTokenCount; private int? currentPageTokenCount;
private int? currentPageNumber;
public ContentStreamPendingContent? AddPage(ContentStreamDocumentMetadata metadata, string? content, int? tokenCount, bool extractImages) public ContentStreamPendingContent? AddPage(ContentStreamDocumentMetadata metadata, string? content, int? tokenCount, bool extractImages)
{ {
@@ -36,9 +37,12 @@ public sealed class DocumentManager
// //
// The count waits here together with the page it belongs to. Handing it out along with // The count waits here together with the page it belongs to. Handing it out along with
// the page we just completed would size that page by the text of this one. // the page we just completed would size that page by the text of this one. The page
// number waits for the same reason: it belongs to the page being buffered, not to the
// one leaving here.
// //
this.currentPageTokenCount = tokenCount; this.currentPageTokenCount = tokenCount;
this.currentPageNumber = pageNumber;
return completedPage; return completedPage;
} }
@@ -72,8 +76,10 @@ public sealed class DocumentManager
var result = this.currentPageContent.ToString(); var result = this.currentPageContent.ToString();
var tokenCount = this.currentPageTokenCount; var tokenCount = this.currentPageTokenCount;
var pageNumber = this.currentPageNumber;
this.currentPageContent = null; this.currentPageContent = null;
this.currentPageTokenCount = null; this.currentPageTokenCount = null;
return string.IsNullOrWhiteSpace(result) ? null : new ContentStreamPendingContent(result, tokenCount); this.currentPageNumber = null;
return string.IsNullOrWhiteSpace(result) ? null : new ContentStreamPendingContent(result, tokenCount, pageNumber);
} }
} }
@@ -1,3 +1,9 @@
namespace AIStudio.Tools.Services; 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; internal const int DEFAULT_CHUNK_OVERLAP_TOKEN_LENGTH = 300;
private const bool IMAGE_EMBEDDING_ENABLED = false; 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 private enum RagFileIndexingDecision
{ {
INDEXABLE, INDEXABLE,
@@ -23,10 +39,23 @@ public sealed partial class DataSourceEmbeddingService
UNSUPPORTED, 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); 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 EmbeddingChunkDraft(string ChunkId, string Text, int ChunkIndex, int? PageNumber);
private sealed record ChunkingOptions(int MaxChunkTokenLength, int OverlapTokenLength); 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 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 options = this.GetChunkingOptions(dataSource, embeddingProvider);
var strategy = this.GetChunkingStrategy(filePath); var strategy = this.GetChunkingStrategy(filePath);
@@ -55,26 +84,31 @@ public sealed partial class DataSourceEmbeddingService
{ {
var normalized = NormalizeChunkSegment(segment.Content); var normalized = NormalizeChunkSegment(segment.Content);
if (!string.IsNullOrWhiteSpace(normalized)) 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); 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); 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; yield return chunk;
} }
private async IAsyncEnumerable<string> SplitTextByRulesAsync( private async IAsyncEnumerable<EmbeddingChunk> SplitTextByRulesAsync(
string text, string text,
IReadOnlyList<ExtractedFileSegment> sourceSegments, IReadOnlyList<ExtractedFileSegment> sourceSegments,
ChunkingStrategy strategy, ChunkingStrategy strategy,
int ruleIndex, int ruleIndex,
ChunkingOptions options, ChunkingOptions options,
EmbeddingProvider embeddingProvider, EmbeddingProvider embeddingProvider,
int? currentPageNumber,
[System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token, [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token,
string requiredOverlapPrefix = "", string requiredOverlapPrefix = "",
int? estimatedTokenCount = null) int? estimatedTokenCount = null)
@@ -91,14 +125,14 @@ public sealed partial class DataSourceEmbeddingService
tokenCount = await this.GetEmbeddingTokenCountAsync(embeddingProvider, textWithOverlap, token); tokenCount = await this.GetEmbeddingTokenCountAsync(embeddingProvider, textWithOverlap, token);
if (tokenCount <= options.MaxChunkTokenLength) if (tokenCount <= options.MaxChunkTokenLength)
{ {
yield return textWithOverlap; yield return new(textWithOverlap, currentPageNumber);
yield break; yield break;
} }
} }
if (ruleIndex >= strategy.Rules.Count) 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 return hardChunk;
yield break; yield break;
@@ -107,7 +141,7 @@ public sealed partial class DataSourceEmbeddingService
var rule = strategy.Rules[ruleIndex]; var rule = strategy.Rules[ruleIndex];
if (rule.Split is null) 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 return hardChunk;
yield break; yield break;
@@ -116,7 +150,7 @@ public sealed partial class DataSourceEmbeddingService
var units = NormalizeSplitUnits(rule.Split(text, sourceSegments.Select(segment => segment.Text).ToList()), text); var units = NormalizeSplitUnits(rule.Split(text, sourceSegments.Select(segment => segment.Text).ToList()), text);
if (units.Count <= 1) 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 return chunk;
yield break; yield break;
@@ -135,6 +169,15 @@ public sealed partial class DataSourceEmbeddingService
var overlapPrefix = requiredOverlapPrefix; var overlapPrefix = requiredOverlapPrefix;
var unitTokenCounts = EstimateSplitUnitTokenCounts(units, sourceSegments, rule.UsesSourceSegmentCounts, estimatedTokenCount); 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) while (index < units.Count)
{ {
token.ThrowIfCancellationRequested(); token.ThrowIfCancellationRequested();
@@ -145,8 +188,14 @@ public sealed partial class DataSourceEmbeddingService
var rawChunk = string.Concat(units.Skip(index).Take(unitCount)).Trim(); var rawChunk = string.Concat(units.Skip(index).Take(unitCount)).Trim();
var chunk = AddOverlapPrefix(rawChunk, overlapPrefix); var chunk = AddOverlapPrefix(rawChunk, overlapPrefix);
overlapPrefix = string.Empty; 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)) if (!string.IsNullOrWhiteSpace(chunk))
yield return chunk; yield return new(chunk, PageOfUnit(index));
var nextIndex = index + unitCount; var nextIndex = index + unitCount;
if (nextIndex >= units.Count) if (nextIndex >= units.Count)
@@ -178,9 +227,10 @@ public sealed partial class DataSourceEmbeddingService
string? lastSplitUnit = null; string? lastSplitUnit = null;
var unitTokenCount = unitTokenCounts?[index]; 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; yield return splitUnit;
} }
@@ -372,10 +422,15 @@ public sealed partial class DataSourceEmbeddingService
return bestStartIndex <= chunkStartIndex ? chunkEndIndex : bestStartIndex; 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, string text,
ChunkingOptions options, ChunkingOptions options,
EmbeddingProvider embeddingProvider, EmbeddingProvider embeddingProvider,
int? currentPageNumber,
[System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token, [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token,
string requiredOverlapPrefix = "", string requiredOverlapPrefix = "",
int? estimatedTokenCount = null) int? estimatedTokenCount = null)
@@ -455,7 +510,7 @@ public sealed partial class DataSourceEmbeddingService
var chunk = AddOverlapPrefix(text[startIndex..bestEndIndex].Trim(), overlapPrefix); var chunk = AddOverlapPrefix(text[startIndex..bestEndIndex].Trim(), overlapPrefix);
if (!string.IsNullOrWhiteSpace(chunk)) if (!string.IsNullOrWhiteSpace(chunk))
yield return chunk; yield return new(chunk, currentPageNumber);
if (bestEndIndex >= text.Length) if (bestEndIndex >= text.Length)
yield break; yield break;
@@ -934,6 +989,7 @@ public sealed partial class DataSourceEmbeddingService
private string BuildEmbeddingSignature(IDataSource dataSource, EmbeddingProvider embeddingProvider, ChunkingOptions chunkingOptions) private string BuildEmbeddingSignature(IDataSource dataSource, EmbeddingProvider embeddingProvider, ChunkingOptions chunkingOptions)
{ {
return string.Join('|', return string.Join('|',
CHUNK_METADATA_VERSION,
embeddingProvider.Id, embeddingProvider.Id,
embeddingProvider.UsedLLMProvider, embeddingProvider.UsedLLMProvider,
embeddingProvider.Model.Id, embeddingProvider.Model.Id,
@@ -1086,14 +1142,6 @@ public sealed partial class DataSourceEmbeddingService
return string.IsNullOrWhiteSpace(extension) ? "unknown" : extension; 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) => private string CreatePointId(string dataSourceId, string fingerprint, int chunkIndex) =>
CreateStableGuid($"{dataSourceId}:chunk:{fingerprint}:{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)) 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++; totalChunkCount++;
if (batch.Count >= embeddingBatchSize) if (batch.Count >= embeddingBatchSize)
@@ -278,7 +278,7 @@ public sealed partial class RustService
{ {
if (segment.TokenCount is { } tokenCount) if (segment.TokenCount is { } tokenCount)
{ {
yield return new(segment.Content, tokenCount); yield return new(segment.Content, tokenCount, segment.PageNumber);
continue; continue;
} }
@@ -291,7 +291,7 @@ public sealed partial class RustService
var countedSegment = await this.GetTokenCount(embeddingProvider, segment.Content, token); var countedSegment = await this.GetTokenCount(embeddingProvider, segment.Content, token);
if (countedSegment is { Success: true } counted) if (countedSegment is { Success: true } counted)
{ {
yield return new(segment.Content, counted.TokenCount); yield return new(segment.Content, counted.TokenCount, segment.PageNumber);
continue; 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, string path,
bool extractImages, bool extractImages,
bool includeTokenCount, 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: // The count and the page come from the processed event, not from the event which
// a reader may hold content back across several events, and the count of the content // was just read: a reader may hold content back across several events, and the
// it releases is the count of that content, not of the event that released it. // count and page of the content it releases describe that content, not the event
// that released it.
// //
if (!string.IsNullOrWhiteSpace(processedEvent.Content)) if (!string.IsNullOrWhiteSpace(processedEvent.Content))
yield return (processedEvent.Content, processedEvent.TokenCount); yield return (processedEvent.Content, processedEvent.TokenCount, processedEvent.PageNumber);
} }
} }
finally finally
@@ -434,7 +435,7 @@ public sealed partial class RustService
} }
if (finalContentChunk is { } pendingContent && !string.IsNullOrWhiteSpace(pendingContent.Content)) 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) if (promptInjectionRedactedCount is 0)
yield break; yield break;
@@ -0,0 +1,157 @@
using AIStudio.Tools;
namespace AIStudio.Tests.Tools;
/// <summary>
/// Checks that the page a passage came from is handed on as a number.
/// </summary>
/// <remarks>
/// The runtime states the page of every page it reads. That number used to be written into the
/// text as a heading and read back out of it further down, which left Word and OpenDocument files
/// without a page for good: they are marked with a comment, not with a heading, so the search for
/// a heading never found anything. The tests here pin the number to the metadata, which is the one
/// place it is actually stated.
/// </remarks>
[TestFixture]
public sealed class ContentStreamPageNumberTests
{
[Test]
public void APdfPageStatesItsNumber()
{
var processed = ContentStreamSseHandler.ProcessEvent(PdfEvent(7, "The mixing console is described here."));
Assert.Multiple(() =>
{
Assert.That(processed.PageNumber, Is.EqualTo(7), "The page comes from the metadata of the event.");
Assert.That(processed.Content, Does.Contain("# Page 7"), "The heading stays, because it is what tells the AI which page it reads.");
});
}
[Test]
public void APdfPageWithoutANumberStatesNone()
{
var processed = ContentStreamSseHandler.ProcessEvent(PdfEvent(null, "A page the runtime could not number."));
Assert.That(processed.PageNumber, Is.Null, "Without a number in the metadata there is no page to state.");
}
/// <remarks>
/// This is the case the old approach got wrong: a document which writes about page numbers
/// looks exactly like the marker that used to be searched for.
/// </remarks>
[Test]
public void ATextWhichReadsLikeAPageMarkerIsNotOne()
{
var processed = ContentStreamSseHandler.ProcessEvent(new()
{
Content = "# Page 42\nStill nothing but the text of the document.",
StreamId = NewStreamId(),
Metadata = new ContentStreamTextMetadata(),
});
Assert.Multiple(() =>
{
Assert.That(processed.PageNumber, Is.Null, "Nothing is read out of the text, so a line which looks like a marker stays text.");
Assert.That(processed.Content, Is.EqualTo("# Page 42\nStill nothing but the text of the document."), "The text itself is passed on untouched.");
});
}
/// <remarks>
/// A Word or OpenDocument page is held back until it is clear that no image follows it, so the
/// page leaving the reader is always the one before the event which released it. Its number has
/// to wait together with it; handing out the number of the arriving event would put every
/// passage one page too far ahead.
/// </remarks>
[Test]
public void ADocumentPageCarriesItsOwnNumberAndNotTheOneWhichReleasedIt()
{
var streamId = NewStreamId();
try
{
var first = ContentStreamSseHandler.ProcessEvent(DocumentEvent(streamId, 1, "What the first page says."));
var second = ContentStreamSseHandler.ProcessEvent(DocumentEvent(streamId, 2, "What the second page says."));
Assert.Multiple(() =>
{
Assert.That(first.Content, Is.Null, "The first page is still being buffered, so nothing is released yet.");
Assert.That(second.PageNumber, Is.EqualTo(1), "What is released here is the first page, so it carries page one.");
Assert.That(second.Content, Does.Contain("What the first page says."), "The content released belongs to the page whose number is stated.");
});
}
finally
{
ContentStreamSseHandler.Clear(streamId);
}
}
[Test]
public void TheLastDocumentPageIsReleasedWithItsNumber()
{
var streamId = NewStreamId();
ContentStreamSseHandler.ProcessEvent(DocumentEvent(streamId, 1, "What the first page says."));
ContentStreamSseHandler.ProcessEvent(DocumentEvent(streamId, 2, "What the second page says."));
var remainder = ContentStreamSseHandler.Clear(streamId);
Assert.That(remainder, Is.Not.Null, "The reader always keeps its last page, so there is something left to release.");
Assert.Multiple(() =>
{
Assert.That(remainder!.Value.PageNumber, Is.EqualTo(2), "The page kept back is the second one.");
Assert.That(remainder.Value.Content, Does.Contain("What the second page says."), "The content released belongs to the page whose number is stated.");
});
}
/// <remarks>
/// A slide is not a page, and no program can be told to open one. Stating none is what later
/// lets a click on such a source open the file and stop there.
/// </remarks>
[Test]
public void ASlideStatesNoPage()
{
var processed = ContentStreamSseHandler.ProcessEvent(new()
{
Content = "What the third slide says.",
StreamId = NewStreamId(),
Metadata = new ContentStreamPresentationMetadata { Presentation = new() { SlideNumber = 3 } },
}, extractImages: false);
Assert.Multiple(() =>
{
Assert.That(processed.PageNumber, Is.Null, "A slide number is not a page number.");
Assert.That(processed.Content, Does.Contain("# Slide 3"), "The heading stays, so the AI still knows which slide it reads.");
});
}
[Test]
public void ASpreadsheetRowStatesNoPage()
{
var processed = ContentStreamSseHandler.ProcessEvent(new()
{
Content = "| Console | Channels |",
StreamId = NewStreamId(),
Metadata = new ContentStreamSpreadsheetMetadata { Spreadsheet = new() { SheetName = "Inventory", RowNumber = 0 } },
});
Assert.That(processed.PageNumber, Is.Null, "A sheet has rows, not pages.");
}
private static ContentStreamSseEvent PdfEvent(int? pageNumber, string content) => new()
{
Content = content,
StreamId = NewStreamId(),
Metadata = new ContentStreamPdfMetadata { Pdf = new() { PageNumber = pageNumber } },
};
private static ContentStreamSseEvent DocumentEvent(string streamId, int pageNumber, string content) => new()
{
Content = content,
StreamId = streamId,
Metadata = new ContentStreamDocumentMetadata { Document = new() { PageNumber = pageNumber } },
};
//
// The readers are kept in static tables keyed by the stream. A test which reuses an ID would
// read the pages another test left behind.
//
private static string NewStreamId() => Guid.NewGuid().ToString();
}