diff --git a/app/MindWork AI Studio/Tools/ContentStreamPendingContent.cs b/app/MindWork AI Studio/Tools/ContentStreamPendingContent.cs new file mode 100644 index 00000000..5795dd91 --- /dev/null +++ b/app/MindWork AI Studio/Tools/ContentStreamPendingContent.cs @@ -0,0 +1,28 @@ +namespace AIStudio.Tools; + +/// +/// Content which a reader held back, together with the token count 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. +/// +/// 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) +{ + /// + /// Adds up two token counts, where an unknown count makes the sum unknown as well. + /// + /// + /// 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. + /// + /// The first count, or null when it is unknown. + /// The second count, or null when it is unknown. + /// The sum, or null when either count is unknown. + public static int? AddTokenCounts(int? left, int? right) => left is null || right is null ? null : left + right; +} \ No newline at end of file diff --git a/app/MindWork AI Studio/Tools/ContentStreamProcessedEvent.cs b/app/MindWork AI Studio/Tools/ContentStreamProcessedEvent.cs index aed4f29c..69da60d0 100644 --- a/app/MindWork AI Studio/Tools/ContentStreamProcessedEvent.cs +++ b/app/MindWork AI Studio/Tools/ContentStreamProcessedEvent.cs @@ -11,14 +11,25 @@ namespace AIStudio.Tools; /// The content to append, or null when this event carries none. /// 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. -public readonly record struct ContentStreamProcessedEvent(string? Content, ContentStreamErrorDetails? Error, ContentStreamPromptInjectionDetails? PromptInjection = null) +/// 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) { /// /// An event which neither produced content nor reported a failure. /// public static readonly ContentStreamProcessedEvent NOTHING = new(null, null); - public static ContentStreamProcessedEvent FromContent(string? content) => new(content, null); + /// + /// An event which produced content, with the token count 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. + /// + /// 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); 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 ad099f35..6bfe6a5e 100644 --- a/app/MindWork AI Studio/Tools/ContentStreamSseHandler.cs +++ b/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); - return documentContent is null ? ContentStreamProcessedEvent.NOTHING : ContentStreamProcessedEvent.FromContent(documentContent); + var documentContent = documentManager.AddPage(documentMetadata, sseEvent.Content, sseEvent.TokenCount, extractImages); + 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,32 +178,45 @@ public static class ContentStreamSseHandler return $"![Image](data:{imageMediaType};base64,{base64Image})"; } - public static string? Clear(string streamId) + /// + /// Releases what the readers of a stream still hold back and forgets the stream. + /// + /// + /// 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. + /// + /// The stream to release and forget. + /// The content which was held back, or null when there was none. + 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(); - if (!string.IsNullOrWhiteSpace(result)) - finalContentChunk.Append(result); + finalContentChunk.Append(slides.Content); + tokenCount = ContentStreamPendingContent.AddTokenCounts(tokenCount, slides.TokenCount); } - 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(); - if (!string.IsNullOrWhiteSpace(result)) - finalContentChunk.Append(result); + finalContentChunk.Append(page.Content); + tokenCount = ContentStreamPendingContent.AddTokenCounts(tokenCount, page.TokenCount); } - + SLIDE_MANAGERS.TryRemove(streamId, out _); DOCUMENT_MANAGERS.TryRemove(streamId, out _); var imageIdPrefix = $"{streamId}-"; 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; } -} +} \ 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 018b4509..7af4c704 100644 --- a/app/MindWork AI Studio/Tools/DocumentManager.cs +++ b/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($""); 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); } } diff --git a/app/MindWork AI Studio/Tools/Services/RustService.Retrieval.cs b/app/MindWork AI Studio/Tools/Services/RustService.Retrieval.cs index 89cda150..2747cef7 100644 --- a/app/MindWork AI Studio/Tools/Services/RustService.Retrieval.cs +++ b/app/MindWork AI Studio/Tools/Services/RustService.Retrieval.cs @@ -208,9 +208,9 @@ public sealed partial class RustService } finally { - var finalContentChunk = ContentStreamSseHandler.Clear(streamId); - if (!string.IsNullOrWhiteSpace(finalContentChunk)) - resultBuilder.AppendLine(finalContentChunk); + // Reading the whole file at once needs no token counts, so only the content is used here: + if (ContentStreamSseHandler.Clear(streamId) is { } finalContentChunk && !string.IsNullOrWhiteSpace(finalContentChunk.Content)) + 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}'."); - - yield return new(segment.Content, segment.TokenCount.Value); + 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}"); } } @@ -309,7 +325,7 @@ public sealed partial class RustService var promptInjectionFindings = new List(); 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)) - yield return (finalContentChunk, null); + if (finalContentChunk is { } pendingContent && !string.IsNullOrWhiteSpace(pendingContent.Content)) + yield return (pendingContent.Content, pendingContent.TokenCount); if (promptInjectionRedactedCount is 0) yield break; diff --git a/app/MindWork AI Studio/Tools/Slide.cs b/app/MindWork AI Studio/Tools/Slide.cs index d071cf7e..e45792cc 100644 --- a/app/MindWork AI Studio/Tools/Slide.cs +++ b/app/MindWork AI Studio/Tools/Slide.cs @@ -5,6 +5,16 @@ public sealed class Slide public bool Delivered { get; set; } public int Position { get; init; } - + public List Content { get; } = new(); + + /// + /// The number of tokens of everything this slide holds, or null when it is unknown. + /// + /// + /// 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. + /// + public int? TokenCount { get; set; } } \ No newline at end of file diff --git a/app/MindWork AI Studio/Tools/SlideManager.cs b/app/MindWork AI Studio/Tools/SlideManager.cs index f6ed1ea6..b3c64841 100644 --- a/app/MindWork AI Studio/Tools/SlideManager.cs +++ b/app/MindWork AI Studio/Tools/SlideManager.cs @@ -6,7 +6,7 @@ public sealed class SlideManager { private readonly Dictionary 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,11 +42,15 @@ 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); - + // // Add image content to the slide? // @@ -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,24 +79,37 @@ public sealed class SlideManager { var textContent = slide.Content.OfType().First(); textContent.Text.AppendLine(content); + slide.TokenCount = ContentStreamPendingContent.AddTokenCounts(slide.TokenCount, tokenCount); } - + // Add any image content? if (addImage) { 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()) { content.AppendLine(text.Text.ToString()); @@ -100,7 +122,7 @@ public sealed class SlideManager content.AppendLine(); } } - - return content.Length > 0 ? content.ToString() : null; + + return content.Length > 0 ? new ContentStreamPendingContent(content.ToString(), tokenCount) : null; } } \ No newline at end of file