diff --git a/app/MindWork AI Studio/Tools/ContentStreamPendingContent.cs b/app/MindWork AI Studio/Tools/ContentStreamPendingContent.cs index 5795dd91..17e022cf 100644 --- a/app/MindWork AI Studio/Tools/ContentStreamPendingContent.cs +++ b/app/MindWork AI Studio/Tools/ContentStreamPendingContent.cs @@ -1,17 +1,21 @@ namespace AIStudio.Tools; /// -/// 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. /// /// /// 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 /// 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. /// /// The assembled content. /// The number of tokens of that content, or null when it is unknown. -public readonly record struct ContentStreamPendingContent(string Content, int? TokenCount) +/// The page that content came from, or null when it has none. +public readonly record struct ContentStreamPendingContent(string Content, int? TokenCount, int? PageNumber = null) { /// /// Adds up two token counts, where an unknown count makes the sum unknown as well. diff --git a/app/MindWork AI Studio/Tools/ContentStreamProcessedEvent.cs b/app/MindWork AI Studio/Tools/ContentStreamProcessedEvent.cs index 69da60d0..50469c0b 100644 --- a/app/MindWork AI Studio/Tools/ContentStreamProcessedEvent.cs +++ b/app/MindWork AI Studio/Tools/ContentStreamProcessedEvent.cs @@ -12,7 +12,8 @@ namespace AIStudio.Tools; /// The reported failure, or null when the event was processed successfully. /// What the runtime filtered out of the content, or null when it filtered nothing. /// The number of tokens of the content, or null when it is unknown. -public readonly record struct ContentStreamProcessedEvent(string? Content, ContentStreamErrorDetails? Error, ContentStreamPromptInjectionDetails? PromptInjection = null, int? TokenCount = null) +/// The page the content came from, or null when it has none. +public readonly record struct ContentStreamProcessedEvent(string? Content, ContentStreamErrorDetails? Error, ContentStreamPromptInjectionDetails? PromptInjection = null, int? TokenCount = null, int? PageNumber = null) { /// /// 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); /// - /// 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. /// /// /// 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 - /// 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. /// /// The content to append. /// The number of tokens of that content, or null when it is unknown. - public static ContentStreamProcessedEvent FromContent(string? content, int? tokenCount = null) => new(content, null, TokenCount: tokenCount); + /// The page that content came from, or null when it has none. + 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); diff --git a/app/MindWork AI Studio/Tools/ContentStreamSseHandler.cs b/app/MindWork AI Studio/Tools/ContentStreamSseHandler.cs index 6bfe6a5e..76fe614f 100644 --- a/app/MindWork AI Studio/Tools/ContentStreamSseHandler.cs +++ b/app/MindWork AI Studio/Tools/ContentStreamSseHandler.cs @@ -19,13 +19,19 @@ public static class ContentStreamSseHandler case ContentStreamTextMetadata: 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: var pageNumber = pdfMetadata.Pdf?.PageNumber ?? 0; return ContentStreamProcessedEvent.FromContent($""" # Page {pageNumber} {sseEvent.Content} - """, sseEvent.TokenCount); + """, sseEvent.TokenCount, pageNumber > 0 ? pageNumber : null); case ContentStreamSpreadsheetMetadata spreadsheetMetadata: 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, // carry no page number and are passed on unchanged. // - // The buffering is why the count comes back from the reader rather than from - // this event: the page which is released here arrived one event ago, and this - // event's count belongs to the page which is now being buffered. + // The buffering is why the count and the page come back from the reader rather + // than from this event: the page which is released here arrived one event ago, + // and this event's count and number belong to the page which is now being + // buffered. // case ContentStreamDocumentMetadata documentMetadata: if (documentMetadata.Document?.PageNumber is not > 0) @@ -55,7 +62,7 @@ public static class ContentStreamSseHandler var documentManager = DOCUMENT_MANAGERS.GetOrAdd(sseEvent.StreamId!, _ => new()); 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: return ContentStreamProcessedEvent.FromContent(sseEvent.Content, sseEvent.TokenCount); @@ -184,7 +191,9 @@ public static class ContentStreamSseHandler /// /// 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 - /// 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. /// /// The stream to release and forget. /// The content which was held back, or null when there was none. @@ -195,6 +204,7 @@ public static class ContentStreamSseHandler var finalContentChunk = new StringBuilder(); int? tokenCount = 0; + int? pageNumber = null; if(SLIDE_MANAGERS.TryGetValue(streamId, out var slideManager) && slideManager.GetAllSlidesInOrder() is { } slides && !string.IsNullOrWhiteSpace(slides.Content)) @@ -209,6 +219,7 @@ public static class ContentStreamSseHandler { finalContentChunk.Append(page.Content); tokenCount = ContentStreamPendingContent.AddTokenCounts(tokenCount, page.TokenCount); + pageNumber = page.PageNumber; } 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))) 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; } } \ No newline at end of file diff --git a/app/MindWork AI Studio/Tools/DocumentManager.cs b/app/MindWork AI Studio/Tools/DocumentManager.cs index 7af4c704..26672bbb 100644 --- a/app/MindWork AI Studio/Tools/DocumentManager.cs +++ b/app/MindWork AI Studio/Tools/DocumentManager.cs @@ -10,6 +10,7 @@ public sealed class DocumentManager { private StringBuilder? currentPageContent; private int? currentPageTokenCount; + private int? currentPageNumber; 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 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.currentPageNumber = pageNumber; return completedPage; } @@ -72,8 +76,10 @@ public sealed class DocumentManager var result = this.currentPageContent.ToString(); var tokenCount = this.currentPageTokenCount; + var pageNumber = this.currentPageNumber; this.currentPageContent = 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); } } diff --git a/app/MindWork AI Studio/Tools/Services/ArbitraryFileDataSegment.cs b/app/MindWork AI Studio/Tools/Services/ArbitraryFileDataSegment.cs index 99c1ce7f..9bf83135 100644 --- a/app/MindWork AI Studio/Tools/Services/ArbitraryFileDataSegment.cs +++ b/app/MindWork AI Studio/Tools/Services/ArbitraryFileDataSegment.cs @@ -1,3 +1,9 @@ namespace AIStudio.Tools.Services; -public sealed record ArbitraryFileDataSegment(string Content, int TokenCount); +/// +/// One piece of an extracted file, as the runtime delivered it. +/// +/// The extracted text. +/// The number of tokens of that text. +/// The page that text came from, or null when it has none. Presentations and spreadsheets have none. +public sealed record ArbitraryFileDataSegment(string Content, int TokenCount, int? PageNumber); diff --git a/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.Files.cs b/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.Files.cs index 15f107e5..0810f1b4 100644 --- a/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.Files.cs +++ b/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.Files.cs @@ -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; + /// + /// What this build writes next to a chunk besides its text. Raise it whenever that changes. + /// + /// + /// 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. + /// + 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 SourceSegments); + /// + /// One chunk as the chunking produced it, together with the page it starts on. + /// + /// + /// 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. + /// + /// The chunk itself, overlap prefix included. + /// The page the chunk's own content starts on, or null when it has none. + 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 FileHashes); - private async IAsyncEnumerable StreamEmbeddingChunksAsync(string filePath, IDataSource dataSource, EmbeddingProvider embeddingProvider, [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token) + private async IAsyncEnumerable 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 SplitByChunkingStrategyAsync(ExtractedFileContent content, ChunkingStrategy strategy, ChunkingOptions options, EmbeddingProvider embeddingProvider, [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token) + private async IAsyncEnumerable 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 SplitTextByRulesAsync( + private async IAsyncEnumerable SplitTextByRulesAsync( string text, IReadOnlyList 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 SplitTextByHardCutAsync( + /// + /// 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. + /// + private async IAsyncEnumerable 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}"); diff --git a/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.cs b/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.cs index 086b2f75..90fe8df1 100644 --- a/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.cs +++ b/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.cs @@ -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) diff --git a/app/MindWork AI Studio/Tools/Services/RustService.Retrieval.cs b/app/MindWork AI Studio/Tools/Services/RustService.Retrieval.cs index ccccb630..283ac634 100644 --- a/app/MindWork AI Studio/Tools/Services/RustService.Retrieval.cs +++ b/app/MindWork AI Studio/Tools/Services/RustService.Retrieval.cs @@ -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; diff --git a/app/Tests/Tools/ContentStreamPageNumberTests.cs b/app/Tests/Tools/ContentStreamPageNumberTests.cs new file mode 100644 index 00000000..f57a3a37 --- /dev/null +++ b/app/Tests/Tools/ContentStreamPageNumberTests.cs @@ -0,0 +1,157 @@ +using AIStudio.Tools; + +namespace AIStudio.Tests.Tools; + +/// +/// Checks that the page a passage came from is handed on as a number. +/// +/// +/// 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. +/// +[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."); + } + + /// + /// 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. + /// + [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."); + }); + } + + /// + /// 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. + /// + [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."); + }); + } + + /// + /// 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. + /// + [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(); +} \ No newline at end of file