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

Fixed the exact token count outliving the conversation it was reported for

Browse Source
This commit is contained in:
Thorsten Sommer 2026-09-23 18:36:25 +02:00
parent 4918858650
commit c23ce9e12a
Signed by untrusted user who does not match committer: tsommer
GPG Key ID: 371BBA77A02C0108
6 changed files with 239 additions and 52 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

42
app/MindWork AI Studio/Chat/ChatThread.cs
View File

@ -392,6 +392,48 @@ public sealed record ChatThread
return true;
}
/// <summary>
/// Finds what a provider reported for this conversation, as long as the report still describes it.
/// </summary>
/// <remarks>
/// A report describes one request: the conversation up to the answer which carries it. It is
/// worth something only while the thread still is that conversation, so only the last block is
/// asked, and no earlier answer ever stands in for it. Whatever came after an older report -- a
/// message whose request was turned down, an answer which is still being written, an answer
/// without a report of its own -- is missing from that report's number, and the estimate is
/// closer to the truth than a figure which leaves it out. Answers written while tools were
/// offered carry no report yet, so for them the estimate always takes over.<br/><br/>
///
/// A report also stops counting when the thread holds a different number of blocks than the
/// request did, which means an earlier message was deleted, and when the next request goes to
/// another model, which counts the same conversation with another tokenizer. Editing the last
/// message or rolling the chat back makes an earlier answer the last block again, with exactly
/// the blocks it was reported for, so its report counts once more.<br/><br/>
///
/// What goes unnoticed is a change beside the messages: another system prompt, profile, or
/// selection of tools. That shows only with the next answer. Noticing it would take a
/// fingerprint of the system prompt as it was sent, after the data sources added to it, which
/// is a lot of machinery for a number which corrects itself one answer later.
/// </remarks>
/// <param name="model">The model the next request would go to.</param>
/// <returns>What the provider reported, or unknown when no report describes this conversation.</returns>
public TokenUsage ReportedUsageFor(Model model)
{
if (this.Blocks.Count is 0)
return TokenUsage.UNKNOWN;
if (this.Blocks[^1].Content is not ContentText { IsStreaming: false, ReportedUsage: { } reported })
return TokenUsage.UNKNOWN;
if (reported.BlockCount != this.Blocks.Count)
return TokenUsage.UNKNOWN;
if (!string.Equals(reported.ModelId, model.Id, StringComparison.Ordinal))
return TokenUsage.UNKNOWN;
return reported.ToTokenUsage();
}
private static void DeleteManagedAttachments(ContentBlock block)
{
if (block.Content is not ContentText textContent)

25
app/MindWork AI Studio/Chat/ContentText.cs
View File

@ -58,10 +58,12 @@ public sealed class ContentText : IContent
/// <remarks>
/// Kept on the answer rather than beside the chat, so that it is stored, loaded, and exported
/// with the message it belongs to -- and so that it goes away when the message does. An edited
/// or regenerated answer removes its block, which removes these numbers with it, and the chat
/// falls back to the estimate instead of carrying a figure for a conversation which no longer
/// exists. Null for every answer written before this was recorded, and at every provider which
/// reports nothing.
/// or regenerated answer removes its block, which removes these numbers with it. Null for every
/// answer written before this was recorded, and at every provider which reports nothing.
///
/// Having a report is not the same as the report still being true. Whether it still describes
/// the conversation is decided by ChatThread.ReportedUsageFor, which looks at the thread around
/// the answer, not just at the answer.
/// </remarks>
public ReportedTokenUsage? ReportedUsage { get; set; }
@ -121,7 +123,8 @@ public sealed class ContentText : IContent
/// </remarks>
/// <param name="usage">What the current chunk of the stream says, which is mostly nothing.</param>
/// <param name="modelId">The model the request went to.</param>
public void RecordReportedUsage(TokenUsage usage, string modelId)
/// <param name="blockCount">How many blocks the conversation had when the request went out, this answer included.</param>
public void RecordReportedUsage(TokenUsage usage, string modelId, int blockCount)
{
if (!usage.IsKnown)
return;
@ -131,6 +134,7 @@ public sealed class ContentText : IContent
PromptTokens = usage.PromptTokens,
CompletionTokens = usage.CompletionTokens,
ModelId = modelId,
BlockCount = blockCount,
};
}
@ -195,6 +199,9 @@ public sealed class ContentText : IContent
// Get the settings manager:
var settings = Program.SERVICE_PROVIDER.GetService<SettingsManager>()!;
// What the request carries, for telling later whether a report still describes this chat:
var blocksSent = chatThread.Blocks.Count;
// Start another thread by using a task to uncouple
// the UI thread from the AI processing:
try
@ -224,7 +231,7 @@ public sealed class ContentText : IContent
this.Sources.MergeSources(contentStreamChunk.Sources);
// Keep what the provider says the request cost:
this.RecordReportedUsage(contentStreamChunk.Usage, chatModel.Id);
this.RecordReportedUsage(contentStreamChunk.Usage, chatModel.Id, blocksSent);
// Notify the UI that the content has changed,
// depending on the energy saving mode:
@ -349,6 +356,12 @@ public sealed class ContentText : IContent
}
/// <inheritdoc />
/// <remarks>
/// The reported usage stays behind on purpose. A clone continues somewhere else -- as the
/// example conversation of a chat template, or as an assistant's conversation carried over into
/// a chat -- with another system prompt around it, and what the provider stated was for the
/// request this answer came out of.
/// </remarks>
public IContent DeepClone() => new ContentText
{
Text = this.Text,

14
app/MindWork AI Studio/Chat/ReportedTokenUsage.cs
View File

@ -37,6 +37,20 @@ public sealed record ReportedTokenUsage
/// </remarks>
public string ModelId { get; init; } = string.Empty;
/// <summary>
/// How many blocks the conversation had when the request went out, this answer included.
/// </summary>
/// <remarks>
/// What tells an outdated report apart without anybody having to remember anything about it.
/// Blocks are only ever added at the end, so while this answer is the last block, a thread with
/// the same count is the thread the provider saw. A lower count means an earlier message was
/// deleted, and that message is still inside the reported number.
///
/// Taken when the request is sent rather than when the report arrives: whatever is deleted
/// while the answer streams in was still part of what the provider counted.
/// </remarks>
public int BlockCount { get; init; }
/// <summary>
/// States the stored numbers as a usage again.
/// </summary>

42
app/MindWork AI Studio/Components/ChatComponent.razor.cs
View File

@ -1614,7 +1614,7 @@ public partial class ChatComponent : MSGComponentBase
var toolDefinitions = this.GetRunnableToolDefinitions();
provider = this.Provider;
parts = ConversationParts.Of(thread, this.BuildSystemPromptFor(thread, toolDefinitions), this.UserInput, this.ComposerState.FileAttachments, provider.SupportsImageInput(), toolDefinitions);
reported = LastReportedTokensOf(thread, provider.Model);
reported = thread.ReportedUsageFor(provider.Model);
});
var counted = await this.ConversationTokenCounter.CountAsync(provider, parts, reported, token);
@ -1631,46 +1631,6 @@ public partial class ChatComponent : MSGComponentBase
});
}
/// <summary>
/// Finds what a provider last said a request of this conversation cost.
/// </summary>
/// <remarks>
/// The last answer which carries such a number decides, and nothing else has to be remembered
/// for it: the number lives on the answer, so editing, regenerating, or deleting a message
/// takes it along and an earlier answer -- or none at all -- becomes the one which counts.
///
/// An answer still being written is passed over. Its number arrives with the last line of the
/// stream, and until then it states what the request before it cost, which is a conversation
/// shorter than the one on the screen.
/// </remarks>
/// <param name="thread">The conversation to look through.</param>
/// <param name="model">The model the next request would go to.</param>
/// <returns>What the provider reported, or unknown when none of them did.</returns>
private static TokenUsage LastReportedTokensOf(ChatThread thread, Model model)
{
for (var index = thread.Blocks.Count - 1; index >= 0; index--)
{
if (thread.Blocks[index].Content is not ContentText { IsStreaming: false, ReportedUsage: { } reported })
continue;
var usage = reported.ToTokenUsage();
if (!usage.IsKnown)
continue;
//
// A number charged for another model says nothing about this one: another model counts
// the same conversation with another tokenizer. Reading on would only find older
// answers of that same other model, so the search ends here and the estimate takes
// over until this model has answered once.
//
return string.Equals(reported.ModelId, model.Id, StringComparison.Ordinal)
? usage
: TokenUsage.UNKNOWN;
}
return TokenUsage.UNKNOWN;
}
/// <summary>
/// Works out the system prompt a thread would send.
/// </summary>

9
app/MindWork AI Studio/Tools/AIJobs/AIJobService.cs
View File

@ -301,10 +301,13 @@ public sealed class AIJobService(SettingsManager settingsManager, MessageBus mes
if (!TrySetWaitingForRemote(state, token))
return;
// What the request carries, for telling later whether a report still describes this chat:
var blocksSent = chatThread.Blocks.Count;
await this.NotifyChangedAsync(state);
await foreach (var contentStreamChunk in provider.StreamChatCompletion(request.ProviderSettings.Model, chatThread, settingsManager, token))
{
if (!TryApplyStreamChunk(state, contentStreamChunk, token))
if (!TryApplyStreamChunk(state, contentStreamChunk, blocksSent, token))
break;
var now = DateTimeOffset.Now;
@ -443,7 +446,7 @@ public sealed class AIJobService(SettingsManager settingsManager, MessageBus mes
}
}
private static bool TryApplyStreamChunk(AIJobState state, ContentStreamChunk contentStreamChunk, CancellationToken token)
private static bool TryApplyStreamChunk(AIJobState state, ContentStreamChunk contentStreamChunk, int blocksSent, CancellationToken token)
{
lock (state.SyncRoot)
{
@ -455,7 +458,7 @@ public sealed class AIJobService(SettingsManager settingsManager, MessageBus mes
aiText.IsStreaming = true;
aiText.Text += contentStreamChunk;
aiText.Sources.MergeSources(contentStreamChunk.Sources);
aiText.RecordReportedUsage(contentStreamChunk.Usage, state.ChatGenerationRequest.ProviderSettings.Model.Id);
aiText.RecordReportedUsage(contentStreamChunk.Usage, state.ChatGenerationRequest.ProviderSettings.Model.Id, blocksSent);
if (state.Snapshot.Status is not AIJobStatus.RUNNING)
{

155
app/Tests/Chat/ChatThreadReportedUsageTests.cs Normal file
View File

@ -0,0 +1,155 @@
using AIStudio.Chat;
using AIStudio.Provider;
namespace AIStudio.Tests.Chat;
/// <summary>
/// Checks when what a provider reported for a conversation still describes it.
/// </summary>
/// <remarks>
/// A reported number is shown as exact, so a wrong one does more harm than no number at all: it
/// tells a person that their chat fits while it does not. Every case below is one where the thread
/// moved on after the report, and the estimate has to take over -- or one where it came back to
/// exactly the state the report was taken in, and the report counts again.
/// </remarks>
[TestFixture]
public sealed class ChatThreadReportedUsageTests
{
private static readonly DateTimeOffset START = new(2026, 9, 23, 10, 0, 0, TimeSpan.Zero);
private static readonly Model MODEL = new("model-a", null);
private static readonly Model OTHER_MODEL = new("model-b", null);
[Test]
public void TheReportOfTheLastAnswerCounts()
{
var thread = Thread(Question(1), Answer(2, promptTokens: 1200, blockCount: 2));
var usage = thread.ReportedUsageFor(MODEL);
Assert.Multiple(() =>
{
Assert.That(usage.IsKnown, Is.True);
Assert.That(usage.PromptTokens, Is.EqualTo(1200));
});
}
[Test]
public void AMessageWhoseRequestFailedIsNotLeftOut()
{
//
// A request which was turned down leaves the question standing and removes the empty
// answer. The report before it knows nothing about that question -- which may well be the
// very message that made the chat too large.
//
var thread = Thread(Question(1), Answer(2, promptTokens: 1200, blockCount: 2), Question(3));
Assert.That(thread.ReportedUsageFor(MODEL).IsKnown, Is.False);
}
[Test]
public void AnAnswerBeingWrittenDoesNotBorrowTheReportBeforeIt()
{
var streaming = Block(ChatRole.AI, new ContentText { Text = "Half an ans", IsStreaming = true }, 4);
var thread = Thread(Question(1), Answer(2, promptTokens: 1200, blockCount: 2), Question(3), streaming);
Assert.That(thread.ReportedUsageFor(MODEL).IsKnown, Is.False);
}
[Test]
public void AnAnswerWithoutAReportDoesNotBorrowTheReportBeforeIt()
{
//
// What an answer written while tools were offered looks like today: finished, but without
// a report of its own.
//
var withoutReport = Block(ChatRole.AI, new ContentText { Text = "Second answer" }, 4);
var thread = Thread(Question(1), Answer(2, promptTokens: 1200, blockCount: 2), Question(3), withoutReport);
Assert.That(thread.ReportedUsageFor(MODEL).IsKnown, Is.False);
}
[Test]
public void DeletingAnEarlierMessageOutdatesTheReport()
{
//
// Deleting a large message to make room is exactly when somebody watches the number, and
// the deleted message is still inside what the provider reported.
//
var firstQuestion = Question(1);
var thread = Thread(firstQuestion, Answer(2, promptTokens: 1200, blockCount: 2), Question(3), Answer(4, promptTokens: 2600, blockCount: 4));
thread.Remove(firstQuestion.Content!);
Assert.That(thread.ReportedUsageFor(MODEL).IsKnown, Is.False);
}
[Test]
public void AnotherModelOutdatesTheReport()
{
var thread = Thread(Question(1), Answer(2, promptTokens: 1200, blockCount: 2));
Assert.That(thread.ReportedUsageFor(OTHER_MODEL).IsKnown, Is.False);
}
[Test]
public void EditingTheLastMessageBringsTheReportBeforeItBack()
{
var lastQuestion = Question(3);
var lastAnswer = Answer(4, promptTokens: 2600, blockCount: 4);
var thread = Thread(Question(1), Answer(2, promptTokens: 1200, blockCount: 2), lastQuestion, lastAnswer);
//
// What editing the last message does to the thread: the question goes back into the
// composer, and its answer goes with it.
//
thread.Remove(lastQuestion.Content!);
thread.Remove(lastAnswer.Content!);
var usage = thread.ReportedUsageFor(MODEL);
Assert.Multiple(() =>
{
Assert.That(usage.IsKnown, Is.True, "The first answer is the last block again, with exactly the blocks it was reported for.");
Assert.That(usage.PromptTokens, Is.EqualTo(1200));
});
}
[Test]
public void RollingBackToAnAnswerBringsItsReportBack()
{
var firstAnswer = Answer(2, promptTokens: 1200, blockCount: 2);
var thread = Thread(Question(1), firstAnswer, Question(3), Answer(4, promptTokens: 2600, blockCount: 4));
thread.RollBackTo(firstAnswer.Content!);
var usage = thread.ReportedUsageFor(MODEL);
Assert.Multiple(() =>
{
Assert.That(usage.IsKnown, Is.True);
Assert.That(usage.PromptTokens, Is.EqualTo(1200));
});
}
private static ChatThread Thread(params ContentBlock[] blocks) => new()
{
Blocks = [..blocks],
};
private static ContentBlock Question(int minute) => Block(ChatRole.USER, new ContentText { Text = $"Question {minute}" }, minute);
private static ContentBlock Answer(int minute, int promptTokens, int blockCount)
{
var answer = new ContentText { Text = $"Answer {minute}" };
answer.RecordReportedUsage(TokenUsage.Of(promptTokens, 100), MODEL.Id, blockCount);
return Block(ChatRole.AI, answer, minute);
}
private static ContentBlock Block(ChatRole role, ContentText content, int minute) => new()
{
Time = START.AddMinutes(minute),
ContentType = ContentType.TEXT,
Content = content,
Role = role,
};
}