MindWorkAI/AI-Studio
MindWorkAI/AI-Studio
1
0
Fork 0
mirror of https://github.com/MindWorkAI/AI-Studio.git synced 2026-09-27 01:53:36 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Pair buffered document content with its own token count

Browse Source
This commit is contained in:
Thorsten Sommer 2026-09-09 15:05:39 +02:00
parent f094cb6ac5
commit 7678508276
Signed by untrusted user who does not match committer: tsommer
GPG Key ID: 371BBA77A02C0108
7 changed files with 176 additions and 51 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

28
app/MindWork AI Studio/Tools/ContentStreamPendingContent.cs Normal file
View File

@ -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;
}

15
app/MindWork AI Studio/Tools/ContentStreamProcessedEvent.cs
View File

@ -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);

67
app/MindWork AI Studio/Tools/ContentStreamSseHandler.cs
View File

@ -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)
/// <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();
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;
}
}
}

24
app/MindWork AI Studio/Tools/DocumentManager.cs
View File

@ -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);
}
}

43
app/MindWork AI Studio/Tools/Services/RustService.Retrieval.cs
View File

@ -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<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))
yield return (finalContentChunk, null);
if (finalContentChunk is { } pendingContent && !string.IsNullOrWhiteSpace(pendingContent.Content))
yield return (pendingContent.Content, pendingContent.TokenCount);
if (promptInjectionRedactedCount is 0)
yield break;

12
app/MindWork AI Studio/Tools/Slide.cs
View File

@ -5,6 +5,16 @@ public sealed class Slide
public bool Delivered { get; set; }
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; }
}

38
app/MindWork AI Studio/Tools/SlideManager.cs
View File

@ -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,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<SlideTextContent>().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<SlideTextContent>())
{
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;
}
}