Pair buffered document content with its own token count

app/MindWork AI Studio/Tools/ContentStreamPendingContent.cs

@@ -0,0 +1,28 @@
+namespace AIStudio.Tools;
+
+/// <summary>
+/// Content which a reader held back, together with the token count of exactly that content.
+/// </summary>
+/// <remarks>
+/// 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.
+/// </remarks>
+/// <param name="Content">The assembled content.</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)
+{
+    /// <summary>
+    /// Adds up two token counts, where an unknown count makes the sum unknown as well.
+    /// </summary>
+    /// <remarks>
+    /// A partial sum would understate the whole and would let the chunking size a chunk by a part
+    /// of what it holds. Reporting the count as unknown is the honest answer, because the caller
+    /// can still count the content itself.
+    /// </remarks>
+    /// <param name="left">The first count, or null when it is unknown.</param>
+    /// <param name="right">The second count, or null when it is unknown.</param>
+    /// <returns>The sum, or null when either count is unknown.</returns>
+    public static int? AddTokenCounts(int? left, int? right) => left is null || right is null ? null : left + right;
+}

app/MindWork AI Studio/Tools/ContentStreamProcessedEvent.cs

@@ -11,14 +11,25 @@ namespace AIStudio.Tools;
 /// <param name="Content">The content to append, or null when this event carries none.</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>
-public readonly record struct ContentStreamProcessedEvent(string? Content, ContentStreamErrorDetails? Error, ContentStreamPromptInjectionDetails? PromptInjection = null)
+/// <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)
 {
     /// <summary>
     /// An event which neither produced content nor reported a failure.
     /// </summary>
     public static readonly ContentStreamProcessedEvent NOTHING = new(null, null);
 
-    public static ContentStreamProcessedEvent FromContent(string? content) => new(content, null);
+    /// <summary>
+    /// An event which produced content, with the token count of that very content.
+    /// </summary>
+    /// <remarks>
+    /// 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.
+    /// </remarks>
+    /// <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>
+    public static ContentStreamProcessedEvent FromContent(string? content, int? tokenCount = null) => new(content, null, TokenCount: tokenCount);
 
     public static ContentStreamProcessedEvent FromError(ContentStreamErrorDetails? error) => new(null, error);
 

app/MindWork AI Studio/Tools/ContentStreamSseHandler.cs

@@ -17,7 +17,7 @@ public static class ContentStreamSseHandler
                 switch (sseEvent.Metadata)
                 {
                     case ContentStreamTextMetadata:
-                        return ContentStreamProcessedEvent.FromContent(sseEvent.Content);
+                        return ContentStreamProcessedEvent.FromContent(sseEvent.Content, sseEvent.TokenCount);
 
                     case ContentStreamPdfMetadata pdfMetadata:
                         var pageNumber = pdfMetadata.Pdf?.PageNumber ?? 0;
@@ -25,7 +25,7 @@ public static class ContentStreamSseHandler
                                 # Page {pageNumber}
                                 {sseEvent.Content}
 
-                                """);
+                                """, sseEvent.TokenCount);
 
                     case ContentStreamSpreadsheetMetadata spreadsheetMetadata:
                         var sheetName = spreadsheetMetadata.Spreadsheet?.SheetName;
@@ -38,23 +38,27 @@ public static class ContentStreamSseHandler
                         }
 
                         spreadSheetResult.Append(sseEvent.Content);
-                        return ContentStreamProcessedEvent.FromContent(spreadSheetResult.ToString());
+                        return ContentStreamProcessedEvent.FromContent(spreadSheetResult.ToString(), sseEvent.TokenCount);
 
                     //
                     // Documents which the runtime reads page by page are buffered, so the images of
                     // 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.
+                    //
                     case ContentStreamDocumentMetadata documentMetadata:
                         if (documentMetadata.Document?.PageNumber is not > 0)
-                            return ContentStreamProcessedEvent.FromContent(sseEvent.Content);
+                            return ContentStreamProcessedEvent.FromContent(sseEvent.Content, sseEvent.TokenCount);
 
                         var documentManager = DOCUMENT_MANAGERS.GetOrAdd(sseEvent.StreamId!, _ => new());
-                        var documentContent = documentManager.AddPage(documentMetadata, sseEvent.Content, extractImages);
+                        var documentContent = documentManager.AddPage(documentMetadata, sseEvent.Content, sseEvent.TokenCount, extractImages);
-                        return documentContent is null ? ContentStreamProcessedEvent.NOTHING : ContentStreamProcessedEvent.FromContent(documentContent);
+                        return documentContent is null ? ContentStreamProcessedEvent.NOTHING : ContentStreamProcessedEvent.FromContent(documentContent.Value.Content, documentContent.Value.TokenCount);
 
                     case ContentStreamImageMetadata:
-                        return ContentStreamProcessedEvent.FromContent(sseEvent.Content);
+                        return ContentStreamProcessedEvent.FromContent(sseEvent.Content, sseEvent.TokenCount);
 
                     case ContentStreamPresentationMetadata presentationMetadata:
                         if (!extractImages)
@@ -62,7 +66,7 @@ public static class ContentStreamSseHandler
                             var slideNumber = presentationMetadata.Presentation?.SlideNumber ?? 0;
                             return ContentStreamProcessedEvent.FromContent(slideNumber > 0
                                 ? $"# Slide {slideNumber}\n{sseEvent.Content}"
-                                : sseEvent.Content);
+                                : sseEvent.Content, sseEvent.TokenCount);
                         }
 
                         var slideManager = SLIDE_MANAGERS.GetOrAdd(
@@ -70,7 +74,7 @@ public static class ContentStreamSseHandler
                             _ => new()
                         );
 
-                        slideManager.AddSlide(presentationMetadata, sseEvent.Content, extractImages);
+                        slideManager.AddSlide(presentationMetadata, sseEvent.Content, sseEvent.TokenCount, extractImages);
                         return ContentStreamProcessedEvent.NOTHING;
 
                     //
@@ -90,11 +94,11 @@ public static class ContentStreamSseHandler
                         return ContentStreamProcessedEvent.FromPromptInjection(promptInjectionMetadata.PromptInjection);
 
                     default:
-                        return ContentStreamProcessedEvent.FromContent(sseEvent.Content);
+                        return ContentStreamProcessedEvent.FromContent(sseEvent.Content, sseEvent.TokenCount);
                 }
 
             case { Content: not null, Metadata: null }:
-                return ContentStreamProcessedEvent.FromContent(sseEvent.Content);
+                return ContentStreamProcessedEvent.FromContent(sseEvent.Content, sseEvent.TokenCount);
 
             default:
                 return ContentStreamProcessedEvent.NOTHING;
@@ -174,24 +178,37 @@ public static class ContentStreamSseHandler
         return $"![Image](data:{imageMediaType};base64,{base64Image})";
     }
 
-    public static string? Clear(string streamId)
+    /// <summary>
+    /// Releases what the readers of a stream still hold back and forgets the stream.
+    /// </summary>
+    /// <remarks>
+    /// 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.
+    /// </remarks>
+    /// <param name="streamId">The stream to release and forget.</param>
+    /// <returns>The content which was held back, or null when there was none.</returns>
+    public static ContentStreamPendingContent? Clear(string streamId)
     {
         if (string.IsNullOrWhiteSpace(streamId))
             return null;
 
         var finalContentChunk = new StringBuilder();
-        if(SLIDE_MANAGERS.TryGetValue(streamId, out var slideManager))
+        int? tokenCount = 0;
+        if(SLIDE_MANAGERS.TryGetValue(streamId, out var slideManager)
+           && slideManager.GetAllSlidesInOrder() is { } slides
+           && !string.IsNullOrWhiteSpace(slides.Content))
         {
-            var result = slideManager.GetAllSlidesInOrder();
+            finalContentChunk.Append(slides.Content);
-            if (!string.IsNullOrWhiteSpace(result))
+            tokenCount = ContentStreamPendingContent.AddTokenCounts(tokenCount, slides.TokenCount);
-                finalContentChunk.Append(result);
         }
 
-        if (DOCUMENT_MANAGERS.TryGetValue(streamId, out var documentManager))
+        if (DOCUMENT_MANAGERS.TryGetValue(streamId, out var documentManager)
+            && documentManager.Flush() is { } page
+            && !string.IsNullOrWhiteSpace(page.Content))
         {
-            var result = documentManager.Flush();
+            finalContentChunk.Append(page.Content);
-            if (!string.IsNullOrWhiteSpace(result))
+            tokenCount = ContentStreamPendingContent.AddTokenCounts(tokenCount, page.TokenCount);
-                finalContentChunk.Append(result);
         }
 
         SLIDE_MANAGERS.TryRemove(streamId, out _);
@@ -200,6 +217,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 ? finalContentChunk.ToString() : null;
+        return finalContentChunk.Length > 0 ? new ContentStreamPendingContent(finalContentChunk.ToString(), tokenCount) : null;
     }
 }

app/MindWork AI Studio/Tools/DocumentManager.cs

@@ -9,12 +9,13 @@ namespace AIStudio.Tools;
 public sealed class DocumentManager
 {
     private StringBuilder? currentPageContent;
+    private int? currentPageTokenCount;
 
-    public string? AddPage(ContentStreamDocumentMetadata metadata, string? content, bool extractImages)
+    public ContentStreamPendingContent? AddPage(ContentStreamDocumentMetadata metadata, string? content, int? tokenCount, bool extractImages)
     {
         var pageNumber = metadata.Document?.PageNumber ?? 0;
         if (pageNumber == 0)
-            return content;
+            return content is null ? null : new ContentStreamPendingContent(content, tokenCount);
 
         var image = metadata.Document?.Image;
         if (image is null)
@@ -32,6 +33,12 @@ public sealed class DocumentManager
             this.currentPageContent.AppendLine($"<!-- Estimated page {pageNumber} -->");
             this.currentPageContent.AppendLine();
             this.currentPageContent.Append(content);
+
+            //
+            // 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.
+            //
+            this.currentPageTokenCount = tokenCount;
             return completedPage;
         }
 
@@ -45,19 +52,28 @@ public sealed class DocumentManager
             {
                 this.currentPageContent.AppendLine();
                 this.currentPageContent.AppendLine(markdownImage);
+
+                //
+                // The runtime counted the text of this page, not the image we just embedded into it.
+                // A data URI is orders of magnitude larger than that text, so the count no longer
+                // describes the page: we drop it, and whoever needs one counts the page itself.
+                //
+                this.currentPageTokenCount = null;
             }
         }
 
         return null;
     }
 
-    public string? Flush()
+    public ContentStreamPendingContent? Flush()
     {
         if (this.currentPageContent is null)
             return null;
 
         var result = this.currentPageContent.ToString();
+        var tokenCount = this.currentPageTokenCount;
         this.currentPageContent = null;
-        return string.IsNullOrWhiteSpace(result) ? null : result;
+        this.currentPageTokenCount = null;
+        return string.IsNullOrWhiteSpace(result) ? null : new ContentStreamPendingContent(result, tokenCount);
     }
 }

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

@@ -208,9 +208,9 @@ public sealed partial class RustService
         }
         finally
         {
-            var finalContentChunk = ContentStreamSseHandler.Clear(streamId);
+            // Reading the whole file at once needs no token counts, so only the content is used here:
-            if (!string.IsNullOrWhiteSpace(finalContentChunk))
+            if (ContentStreamSseHandler.Clear(streamId) is { } finalContentChunk && !string.IsNullOrWhiteSpace(finalContentChunk.Content))
-                resultBuilder.AppendLine(finalContentChunk);
+                resultBuilder.AppendLine(finalContentChunk.Content);
         }
 
         if (failureCode is not FileExtractionErrorCode.NONE)
@@ -267,14 +267,30 @@ public sealed partial class RustService
     {
         await foreach (var segment in this.StreamArbitraryFileDataCore(path, false, true, embeddingProvider.TokenizerPath, token))
         {
+            if (segment.TokenCount is { } tokenCount)
+            {
+                yield return new(segment.Content, tokenCount);
+                continue;
+            }
+
+            //
+            // A segment the runtime did not count, e.g. a page which carries an embedded image on
+            // top of its text. The runtime leaves such a count out on purpose instead of failing
+            // the extraction, because we can count the segment ourselves. Without this, a document
+            // would be dropped over a number we are able to produce.
+            //
+            var countedSegment = await this.GetTokenCount(embeddingProvider, segment.Content, token);
+            if (countedSegment is { Success: true } counted)
+            {
+                yield return new(segment.Content, counted.TokenCount);
+                continue;
+            }
+
             //
             // Carries a code so callers can classify it: the file itself is fine, the answer of
             // the runtime was not, which makes this worth another attempt.
             //
-            if (segment.TokenCount is null)
+            throw new FileExtractionException(FileExtractionErrorCode.INVALID_RESPONSE, $"Rust did not return a token count for an extracted segment from '{path}' using provider '{embeddingProvider.Name}', and counting it afterwards failed as well: {countedSegment?.Message}");
-                throw new FileExtractionException(FileExtractionErrorCode.INVALID_RESPONSE, $"Rust did not return a token count for an extracted segment from '{path}' using provider '{embeddingProvider.Name}'.");
-
-            yield return new(segment.Content, segment.TokenCount.Value);
         }
     }
 
@@ -309,7 +325,7 @@ public sealed partial class RustService
         var promptInjectionFindings = new List<PromptInjectionFinding>();
         var promptInjectionRedactedCount = 0;
 
-        string? finalContentChunk;
+        ContentStreamPendingContent? finalContentChunk;
         try
         {
             await using var stream = await response.Content.ReadAsStreamAsync(token);
@@ -394,8 +410,13 @@ public sealed partial class RustService
                     continue;
                 }
 
+                //
+                // 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.
+                //
                 if (!string.IsNullOrWhiteSpace(processedEvent.Content))
-                    yield return (processedEvent.Content, sseEvent.TokenCount);
+                    yield return (processedEvent.Content, processedEvent.TokenCount);
             }
         }
         finally
@@ -403,8 +424,8 @@ public sealed partial class RustService
             finalContentChunk = ContentStreamSseHandler.Clear(streamId);
         }
 
-        if (!string.IsNullOrWhiteSpace(finalContentChunk))
+        if (finalContentChunk is { } pendingContent && !string.IsNullOrWhiteSpace(pendingContent.Content))
-            yield return (finalContentChunk, null);
+            yield return (pendingContent.Content, pendingContent.TokenCount);
 
         if (promptInjectionRedactedCount is 0)
             yield break;

app/MindWork AI Studio/Tools/Slide.cs

@@ -7,4 +7,14 @@ public sealed class Slide
     public int Position { get; init; }
 
     public List<ISlideContent> Content { get; } = new();
+
+    /// <summary>
+    /// The number of tokens of everything this slide holds, or null when it is unknown.
+    /// </summary>
+    /// <remarks>
+    /// A slide grows across several stream events, so its count grows with it. It becomes unknown
+    /// as soon as an image is embedded: the runtime counted the text of the slide, and a data URI
+    /// is orders of magnitude larger than that.
+    /// </remarks>
+    public int? TokenCount { get; set; }
 }

app/MindWork AI Studio/Tools/SlideManager.cs

@@ -6,7 +6,7 @@ public sealed class SlideManager
 {
     private readonly Dictionary<int, Slide> slides = new();
 
-    public void AddSlide(ContentStreamPresentationMetadata metadata, string? content, bool extractImages = false)
+    public void AddSlide(ContentStreamPresentationMetadata metadata, string? content, int? tokenCount, bool extractImages = false)
     {
         var slideNumber = metadata.Presentation?.SlideNumber ?? 0;
         if(slideNumber is 0)
@@ -42,7 +42,11 @@ public sealed class SlideManager
             var createdSlide = new Slide
             {
                 Delivered = false,
-                Position = slideNumber
+                Position = slideNumber,
+
+                // The count of the text we just added. It travels with the slide, because the slide
+                // is delivered long after this event:
+                TokenCount = tokenCount
             };
 
             createdSlide.Content.Add(slideText);
@@ -54,7 +58,12 @@ public sealed class SlideManager
             {
                 var markdownImage = ContentStreamSseHandler.BuildImageMarkdown(image!.Id!, image.MediaType);
                 if (markdownImage is not null)
+                {
                     createdSlide.Content.Add(new SlideImageContent(markdownImage));
+
+                    // The runtime counted the text of the slide, not the data URI we just added:
+                    createdSlide.TokenCount = null;
+                }
             }
 
             this.slides[slideNumber] = createdSlide;
@@ -70,6 +79,7 @@ public sealed class SlideManager
             {
                 var textContent = slide.Content.OfType<SlideTextContent>().First();
                 textContent.Text.AppendLine(content);
+                slide.TokenCount = ContentStreamPendingContent.AddTokenCounts(slide.TokenCount, tokenCount);
             }
 
             // Add any image content?
@@ -77,17 +87,29 @@ public sealed class SlideManager
             {
                 var markdownImage = ContentStreamSseHandler.BuildImageMarkdown(image!.Id!, image.MediaType);
                 if (markdownImage is not null)
+                {
                     slide.Content.Add(new SlideImageContent(markdownImage));
+
+                    // The runtime counted the text of the slide, not the data URI we just added:
+                    slide.TokenCount = null;
+                }
             }
         }
     }
 
-    public string? GetAllSlidesInOrder()
+    public ContentStreamPendingContent? GetAllSlidesInOrder()
     {
         var content = new StringBuilder();
+
+        // Starts at zero and stays a number only as long as every slide contributes a count of its
+        // own. One slide without one makes the total unknown, which is what the caller has to know:
+        int? tokenCount = 0;
+
         foreach (var slide in this.slides.Values.Where(s => !s.Delivered).OrderBy(s => s.Position))
         {
             slide.Delivered = true;
+            tokenCount = ContentStreamPendingContent.AddTokenCounts(tokenCount, slide.TokenCount);
+
             foreach (var text in slide.Content.OfType<SlideTextContent>())
             {
                 content.AppendLine(text.Text.ToString());
@@ -101,6 +123,6 @@ public sealed class SlideManager
             }
         }
 
-        return content.Length > 0 ? content.ToString() : null;
+        return content.Length > 0 ? new ContentStreamPendingContent(content.ToString(), tokenCount) : null;
     }
 }