diff --git a/app/MindWork AI Studio/Chat/ConversationParts.cs b/app/MindWork AI Studio/Chat/ConversationParts.cs index 6cd363af..3c6e8dc1 100644 --- a/app/MindWork AI Studio/Chat/ConversationParts.cs +++ b/app/MindWork AI Studio/Chat/ConversationParts.cs @@ -22,6 +22,17 @@ public sealed record ConversationParts /// public IReadOnlyList Texts { get; init; } = []; + /// + /// The texts which are still being written. + /// + /// + /// They cost exactly what the others cost; what sets them apart is that they will never be seen + /// again in this shape. The sentence somebody is typing changes with the next pause, and an + /// answer being streamed is a different text three seconds later -- so remembering what they + /// cost fills memory with answers nobody will ask for again. + /// + public IReadOnlyList GrowingTexts { get; init; } = []; + /// /// The documents whose content is put into the request. /// @@ -52,6 +63,7 @@ public sealed record ConversationParts public static ConversationParts Of(ChatThread? thread, string systemPrompt, string draft, IEnumerable? draftAttachments, bool imagesAreSent) { var texts = new List(); + var growing = new List(); var documents = new List(); var images = 0; @@ -70,13 +82,17 @@ public sealed record ConversationParts if (block.ContentType is not ContentType.TEXT || block.Content is not ContentText text || string.IsNullOrWhiteSpace(text.Text)) continue; - texts.Add(text.Text); + if (text.IsStreaming) + growing.Add(text.Text); + else + texts.Add(text.Text); + Sort(text.FileAttachments, documents, ref images); } } if (!string.IsNullOrWhiteSpace(draft)) - texts.Add(draft); + growing.Add(draft); if (draftAttachments is not null) Sort(draftAttachments, documents, ref images); @@ -84,6 +100,7 @@ public sealed record ConversationParts return new() { Texts = texts, + GrowingTexts = growing, Documents = documents, Images = imagesAreSent ? images : 0, }; diff --git a/app/MindWork AI Studio/Chat/ConversationTokenTracker.cs b/app/MindWork AI Studio/Chat/ConversationTokenTracker.cs new file mode 100644 index 00000000..3fbe4d0d --- /dev/null +++ b/app/MindWork AI Studio/Chat/ConversationTokenTracker.cs @@ -0,0 +1,172 @@ +namespace AIStudio.Chat; + +/// +/// Keeps a number up to date which nothing announces. +/// +/// +/// A conversation is a plain list of plain objects. Nothing raises an event when a block is added, +/// when a document is attached, or when an answer grows by another sentence -- so a number derived +/// from all of that cannot be wired to the places which change it. It was tried: fifteen call sites, +/// and four review rounds each found another one which was missing. +/// +/// So the number is recomputed instead of notified. Whoever thinks something may have changed nudges +/// this tracker, and the tracker decides when to do the work: many nudges in a row become one run, a +/// nudge arriving during a run becomes exactly one further run, and a minimum distance keeps a burst +/// of them from turning into a burst of counting. +/// +/// The heartbeat is not distrust of the nudges. Attachments are read from disk every time they are +/// sent, so a file somebody edits in another program changes what the next message costs without +/// anything happening in AI Studio which anyone could nudge from. +/// +/// Does the actual work. Gets a token which ends it when the tracker goes away. +/// +/// How long to stay quiet after a run before honouring the next nudge. Asked again each time, +/// because what is reasonable depends on what is going on: a person who just switched a profile is +/// waiting for the number, while an answer being written moves it with every word and wants a +/// slower pace than the words arrive at. +/// +/// How long to wait for a nudge before running anyway. +public sealed class ConversationTokenTracker(Func recount, Func quietTime, TimeSpan heartbeat) : IAsyncDisposable +{ + /// + /// How long a tracker which is going away waits for its own loop. + /// + /// + /// The loop ends on cancellation, so this is only ever reached when something it called does + /// not. Whoever is leaving the screen must not be the one who waits for that. + /// + private static readonly TimeSpan SHUTDOWN_PATIENCE = TimeSpan.FromSeconds(2); + + private readonly SemaphoreSlim wakeUp = new(0, 1); + private readonly CancellationTokenSource stopping = new(); + + private Task? loop; + + /// + /// Starts the loop. Calling this twice does nothing the second time. + /// + public void Start() => this.loop ??= Task.Run(this.RunAsync); + + /// + /// Says that something may have changed. + /// + /// + /// Cheap on purpose, because it is called from the render path. It says "maybe", never "yes": + /// asking for a run which turns out to change nothing costs a few lookups, while missing one is + /// the bug this whole class exists to make impossible. + /// + public void Nudge() + { + // + // One pending wake-up is all a loop can act on. A second one would only make it run again + // with the same answer. + // + if (this.wakeUp.CurrentCount > 0) + return; + + try + { + this.wakeUp.Release(); + } + catch (SemaphoreFullException) + { + // + // Two threads got past the check above at the same time. The one which won left the + // wake-up we wanted, so there is nothing left to do here. + // + } + catch (ObjectDisposedException) + { + // The tracker is going away, and a number nobody will look at needs no update. + } + } + + private async Task RunAsync() + { + var token = this.stopping.Token; + while (!token.IsCancellationRequested) + { + try + { + // + // Sleeps until somebody nudges -- or until the heartbeat is due, which is what the + // timeout returning false means. Both lead to the same run, so the result is not + // even looked at. + // + await this.wakeUp.WaitAsync(heartbeat, token); + if (token.IsCancellationRequested) + return; + + // + // Deliberately without draining further wake-ups first. A nudge which arrives while + // this run reads the conversation may well be about a change this run is already + // seeing -- and then the extra run costs a few lookups. Draining would risk the + // other case, where the change comes after the read and nobody asks again. + // + try + { + await recount(token); + } + catch (Exception) when (!token.IsCancellationRequested) + { + // + // One failed run must not end the loop: a tracker which died on a single bad + // answer would leave a stale number standing forever, which is the failure this + // class was built to rule out. Saying what went wrong is the job of the work + // itself, which is the only side that has a logger. + // + } + + // + // The quiet time is kept after the work, not before it: the first nudge of a burst + // is answered at once, and the rest of the burst collapses into the single run which + // follows this delay. + // + // It is also what paces a run which feeds itself. Showing a new number renders, and + // a render nudges -- so while something changes continuously, this delay is the + // whole cadence. + // + await Task.Delay(quietTime(), token); + } + catch (OperationCanceledException) + { + return; + } + catch (ObjectDisposedException) + { + // The tracker was disposed underneath this loop, which is another way of stopping. + return; + } + } + } + + #region Implementation of IAsyncDisposable + + public async ValueTask DisposeAsync() + { + await this.stopping.CancelAsync(); + + if (this.loop is not null) + { + try + { + // + // Awaited rather than abandoned, so that nothing is still counting into a component + // which is already gone. The counting itself takes the same token, so a run which + // sits in an IPC call ends with it -- and the patience is there for the case where + // it does not, because a chat being closed is not worth hanging on to. + // + await this.loop.WaitAsync(SHUTDOWN_PATIENCE); + } + catch (Exception) + { + // The loop ends on cancellation; whatever else it carries out is of no use here. + } + } + + this.stopping.Dispose(); + this.wakeUp.Dispose(); + } + + #endregion +} \ No newline at end of file diff --git a/app/MindWork AI Studio/Components/AttachDocuments.razor.cs b/app/MindWork AI Studio/Components/AttachDocuments.razor.cs index c3895fa2..6ea7ad26 100644 --- a/app/MindWork AI Studio/Components/AttachDocuments.razor.cs +++ b/app/MindWork AI Studio/Components/AttachDocuments.razor.cs @@ -258,8 +258,16 @@ public partial class AttachDocuments : MSGComponentBase this.DocumentPaths = await ReviewAttachmentsDialog.OpenDialogAsync(this.DialogService, this.DocumentPaths); foreach (var removedAttachment in previousAttachments.Except(this.DocumentPaths)) ManagedTranscriptAttachment.TryDeleteOwnedFile(removedAttachment); - + this.ReconcileOwnerPendingTranscripts(); + + // + // Said out loud, like every other path in this file. Removing a file in the dialog changed + // the attachments while whoever owns them heard nothing about it -- the chat then kept + // showing what the message no longer carries. + // + await this.DocumentPathsChanged.InvokeAsync(this.DocumentPaths); + await this.OnChange(this.DocumentPaths); } private async Task ClearAllFiles() diff --git a/app/MindWork AI Studio/Components/ChatComponent.razor b/app/MindWork AI Studio/Components/ChatComponent.razor index ad7bd9e0..addb3536 100644 --- a/app/MindWork AI Studio/Components/ChatComponent.razor +++ b/app/MindWork AI Studio/Components/ChatComponent.razor @@ -51,7 +51,6 @@ Disabled="@this.IsInputForbidden()" Immediate="@true" OnKeyUp="@this.InputKeyEvent" - WhenTextChangedAsync="@(_ =>this.CalculateTokenCount())" UserAttributes="@USER_INPUT_ATTRIBUTES" Class="@this.UserInputClass" DebounceTime="TimeSpan.FromSeconds(1)" diff --git a/app/MindWork AI Studio/Components/ChatComponent.razor.cs b/app/MindWork AI Studio/Components/ChatComponent.razor.cs index 3cd3bb62..b43799e2 100644 --- a/app/MindWork AI Studio/Components/ChatComponent.razor.cs +++ b/app/MindWork AI Studio/Components/ChatComponent.razor.cs @@ -98,6 +98,52 @@ public partial class ChatComponent : MSGComponentBase private int workspaceHeaderSyncVersion; private ConversationTokens conversationTokens = ConversationTokens.UNAVAILABLE; + /// + /// How much of the window must be used before the number starts saying so. + /// + private const double WINDOW_NEARLY_FULL = 0.8d; + + /// + /// How long the token count stays quiet after it ran, while nothing is being written. + /// + /// + /// A render is cheap to ask about and a count is not. This is what keeps a burst of renders -- + /// loading a chat touches several things in a row -- from turning into a burst of counting, + /// while staying short enough that switching a profile moves the number right away. + /// + private static readonly TimeSpan TOKEN_COUNT_QUIET_TIME = TimeSpan.FromMilliseconds(500); + + /// + /// How long the token count stays quiet while an answer is being written. + /// + /// + /// An answer grows with every word, so each count finds a new number, shows it, and thereby + /// renders -- which asks for the next count. That makes this the whole cadence while a model + /// writes, and three seconds is the pace the chat itself keeps: the job service hands its + /// progress to the screen no more often than that. + /// + private static readonly TimeSpan TOKEN_COUNT_STREAMING_QUIET_TIME = TimeSpan.FromSeconds(3); + + /// + /// How long the token count waits for a reason before counting anyway. + /// + /// + /// For what happens outside AI Studio: an attached document is read from disk every time it is + /// sent, so somebody editing it in another program changes what the next message costs without + /// anything here rendering. + /// + private static readonly TimeSpan TOKEN_COUNT_HEARTBEAT = TimeSpan.FromSeconds(10); + + /// + /// Recomputes the token count whenever something might have changed. + /// + private ConversationTokenTracker? tokenTracker; + + /// + /// How long to leave the token count alone after it ran. + /// + private TimeSpan TokenCountQuietTime() => this.IsCurrentChatStreaming ? TOKEN_COUNT_STREAMING_QUIET_TIME : TOKEN_COUNT_QUIET_TIME; + /// /// The culture the token numbers are written in. /// @@ -176,6 +222,14 @@ public partial class ChatComponent : MSGComponentBase this.MediaTranscriptionService.StateChanged += this.OnMediaImportStateChanged; await this.RefreshCulture(); + // + // The number under the input field follows from the conversation, and nothing in a + // conversation announces that it changed: blocks, attachments and the answer being written + // are plain objects somebody mutates. So it is recomputed rather than notified. + // + this.tokenTracker = new(this.RecountTokensAsync, this.TokenCountQuietTime, TOKEN_COUNT_HEARTBEAT); + this.tokenTracker.Start(); + // Apply the filters for the message bus: this.ApplyFilters([], [ Event.HAS_CHAT_UNSAVED_CHANGES, Event.RESET_CHAT_STATE, Event.CHAT_STREAMING_DONE, Event.AI_JOB_CHANGED, Event.AI_JOB_FINISHED, Event.CHAT_GENERATION_CHANGED, Event.WORKSPACE_RENAMED, Event.CONFIGURATION_CHANGED ]); @@ -424,6 +478,15 @@ public partial class ChatComponent : MSGComponentBase await this.inputField.FocusAsync(); this.previousInputForbidden = inputForbidden; + + // + // Everything which can move the token count also renders this component: the selections in + // the toolbar, the attachments and the composer all travel through an event callback whose + // receiver is this component, and the streamed answer arrives as a message which already + // asks for a render. So this one line stands in for the fifteen call sites which used to be + // spread over this file -- and which kept missing one. + // + this.tokenTracker?.Nudge(); await base.OnAfterRenderAsync(firstRender); } @@ -439,14 +502,6 @@ public partial class ChatComponent : MSGComponentBase await this.ApplyLoadedChatParameterAsync(); await this.SyncForegroundChatAsync(); - - // - // Both of these change the answer, and the chat is the reason the count is no longer about - // the draft alone: opening another conversation, or loading one, changes what the next - // message would carry along with it. - // - await this.CalculateTokenCount(); - await this.ConsumeMediaOutcomeAsync(); await base.OnParametersSetAsync(); } @@ -590,7 +645,34 @@ public partial class ChatComponent : MSGComponentBase private string UserInputStyle => this.SettingsManager.ConfigurationData.Confidence.ShowProviderConfidence ? this.Provider.UsedLLMProvider.GetConfidence(this.SettingsManager).SetColorStyle(this.SettingsManager) : string.Empty; - private string UserInputClass => this.SettingsManager.ConfigurationData.Confidence.ShowProviderConfidence ? "confidence-border" : string.Empty; + private string UserInputClass => $"{(this.SettingsManager.ConfigurationData.Confidence.ShowProviderConfidence ? "confidence-border" : string.Empty)} {this.TokenBudgetClass}".Trim(); + + /// + /// How much of the model's context window the conversation already takes. + /// + /// + /// Zero whenever nobody wrote the window down. There is then nothing to be full of, and a share + /// of an unknown total would be a number made up on the spot. + /// + private double TokenBudgetFill => this.conversationTokens is { IsKnown: true, Window.IsKnown: true } + ? (double) this.conversationTokens.Tokens / this.conversationTokens.Window.DefaultTokens + : 0d; + + /// + /// What the number under the input field is coloured with, if anything. + /// + /// + /// Two steps rather than a gradient: below four fifths there is nothing to do about it, above + /// it there is -- shorten the chat, start a new one, or pick a model which reads more -- and + /// past the window the request will be refused or trimmed by the provider. + /// + private string TokenBudgetClass => this.TokenBudgetFill switch + { + >= 1d => "token-budget-exceeded", + >= WINDOW_NEARLY_FULL => "token-budget-nearly-full", + + _ => string.Empty, + }; private void ApplyStandardDataSourceOptions() { @@ -657,8 +739,7 @@ public partial class ChatComponent : MSGComponentBase // // A thread which already exists has to carry the choice. Before the first message there is - // none, and the choice then travels in the thread a new chat is started with -- so the - // count below has to happen either way, which is what the early return here used to skip. + // none, and the choice then travels in the thread a new chat is started with. // if (this.ChatThread is not null) { @@ -669,9 +750,6 @@ public partial class ChatComponent : MSGComponentBase await this.ChatThreadChanged.InvokeAsync(this.ChatThread); } - - // A profile is a paragraph of the system prompt, so choosing another one changes the count: - await this.CalculateTokenCount(); } private async Task ChatTemplateWasChanged(ChatTemplate chatTemplate) @@ -685,13 +763,6 @@ public partial class ChatComponent : MSGComponentBase if (this.ChatThread is not null) await this.StartNewChat(true); - - // - // Counted in both cases. Without a thread nothing is started anew, but the template already - // decides the system prompt, the attachments and possibly an example conversation of the - // first message, and all of that costs before anything is sent. - // - await this.CalculateTokenCount(); } private void RefreshCurrentProfileAndChatTemplate() @@ -729,9 +800,6 @@ public partial class ChatComponent : MSGComponentBase this.ComposerState.ApplyTemplate(this.currentChatTemplate); await this.ApplyUpdatedStandardDataSourceOptionsAfterConfigurationChange(); - - // The provider, the profile and the template may all have moved, and each of them counts: - await this.CalculateTokenCount(); } private IReadOnlyList GetAgentSelectedDataSources() @@ -794,10 +862,7 @@ public partial class ChatComponent : MSGComponentBase // Was a modifier key pressed as well? var isModifier = keyEvent.AltKey || keyEvent.CtrlKey || keyEvent.MetaKey || keyEvent.ShiftKey; - - if (isEnter) - await this.CalculateTokenCount(); - + // Depending on the user's settings, might react to shortcuts: switch (this.SettingsManager.ConfigurationData.Chat.ShortcutSendBehavior) { @@ -825,20 +890,13 @@ public partial class ChatComponent : MSGComponentBase this.hasUnsavedChanges = true; } - private async Task ComposerAttachmentsChanged(HashSet attachments) + private void ComposerAttachmentsChanged(HashSet attachments) { if (!ReferenceEquals(this.ComposerState.FileAttachments, attachments)) this.ComposerState.ReplaceFileAttachments(attachments); this.ComposerState.MarkUserDraft(); this.hasUnsavedChanges = true; - - // - // A document is usually the largest thing a person attaches, so this is the moment the - // number matters most. It is also the expensive one: the file is read and measured here, - // once, and remembered afterwards. - // - await this.CalculateTokenCount(); } /// Creates and stores a stable draft immediately after media import confirmation. @@ -974,13 +1032,6 @@ public partial class ChatComponent : MSGComponentBase await this.inputField.BlurAsync(); - // - // The draft just became part of the conversation, so the number does not drop back: what - // was being typed a moment ago now travels with every further message. - // - await this.CalculateTokenCount(); - - // Enable the stream state for the chat component: this.hasUnsavedChanges = true; @@ -1025,7 +1076,7 @@ public partial class ChatComponent : MSGComponentBase private void ApplyToolSelectionOfLoadedChat() => this.selectedToolIds = ToolSelectionRules.NormalizeSelection(this.ChatThread?.SelectedToolIds ?? this.SettingsManager.GetDefaultToolIds(Tools.Components.CHAT)); - private async Task SelectedToolIdsChanged(HashSet updatedToolIds) + private void SelectedToolIdsChanged(HashSet updatedToolIds) { this.selectedToolIds = ToolSelectionRules.NormalizeSelection(updatedToolIds); @@ -1040,12 +1091,6 @@ public partial class ChatComponent : MSGComponentBase this.ChatThread.SelectedToolIds = [..this.selectedToolIds]; this.hasUnsavedChanges = true; } - - // - // Every tool the model may run describes itself in the system prompt, so picking tools costs - // tokens before a single one of them is called. - // - await this.CalculateTokenCount(); } private async Task SaveThread() @@ -1166,7 +1211,6 @@ public partial class ChatComponent : MSGComponentBase await this.SyncForegroundChatAsync(); this.MarkCurrentChatAsLoadedParameter(); await this.ChatThreadChanged.InvokeAsync(this.ChatThread); - await this.CalculateTokenCount(); } private async Task MoveChatToWorkspace() @@ -1270,9 +1314,8 @@ public partial class ChatComponent : MSGComponentBase await this.SyncForegroundChatAsync(); this.ApplyStandardDataSourceOptions(); await this.ChatThreadChanged.InvokeAsync(this.ChatThread); - await this.CalculateTokenCount(); } - + private async Task SelectProviderWhenLoadingChat() { var chatProvider = this.ChatThread?.SelectedProvider; @@ -1306,9 +1349,6 @@ public partial class ChatComponent : MSGComponentBase this.hasUnsavedChanges = true; await this.SaveThread(); this.StateHasChanged(); - - // One message less in the conversation is one message less in every request from here on: - await this.CalculateTokenCount(); } private async Task RegenerateBlock(IContent aiBlock) @@ -1326,46 +1366,40 @@ public partial class ChatComponent : MSGComponentBase await this.SendMessage(reuseLastUserPrompt: true); } - private async Task EditLastUserBlock(IContent block) + private Task EditLastUserBlock(IContent block) { if(this.ChatThread is null) - return; + return Task.CompletedTask; if (block is not ContentText textBlock) - return; + return Task.CompletedTask; var lastBlock = this.ChatThread.Blocks.Last(); var lastBlockContent = lastBlock.Content; if(lastBlockContent is null) - return; + return Task.CompletedTask; this.RestoreComposerFromTextBlock(textBlock); this.ChatThread.Remove(block); this.ChatThread.Remove(lastBlockContent); this.hasUnsavedChanges = true; this.StateHasChanged(); - - // - // The message moved out of the conversation and back into the composer, attachments and - // all. It costs the same either way, but nothing says so unless it is counted again. - // - await this.CalculateTokenCount(); + return Task.CompletedTask; } - private async Task EditLastBlock(IContent block) + private Task EditLastBlock(IContent block) { if(this.ChatThread is null) - return; + return Task.CompletedTask; if (block is not ContentText textBlock) - return; + return Task.CompletedTask; this.RestoreComposerFromTextBlock(textBlock); this.ChatThread.Remove(block); this.hasUnsavedChanges = true; this.StateHasChanged(); - - await this.CalculateTokenCount(); + return Task.CompletedTask; } private void RestoreComposerFromTextBlock(ContentText textBlock) @@ -1387,24 +1421,47 @@ public partial class ChatComponent : MSGComponentBase /// estimate, and it says so. /// /// Read the text from the bound property rather than from the input field: the field is a - /// component reference, which is only set once the component has rendered, while counting is - /// also triggered while parameters are set. + /// component reference, which is only set once the component has rendered. + /// + /// Called by the tracker, never directly. Whoever thinks something changed nudges it instead, + /// and it decides when the work is worth doing. /// - private async Task CalculateTokenCount() + /// Ends the count when the component goes away. + private async Task RecountTokensAsync(CancellationToken token) { + var provider = AIStudio.Settings.Provider.NONE; + var parts = ConversationParts.NOTHING; + // - // Before the first message there is no thread yet, so what is measured is the one a new - // chat would start with. A preselected profile or a chat template is already part of that, - // and it may even bring an example conversation along -- reporting nothing for all of it - // would tell a person their window is empty while their first message already is not. + // Collected on the render thread, counted off it. Counting may take an IPC call per text, + // and while it runs, the background job which writes the answer appends to the very list + // which is walked here. // - var thread = this.ChatThread ?? this.NewChatThread(string.Empty); - var counted = await this.ConversationTokenCounter.CountAsync(this.Provider, thread, this.BuildSystemPromptFor(thread), this.UserInput, this.ComposerState.FileAttachments); - if (counted == this.conversationTokens) + await this.InvokeAsync(() => + { + // + // Before the first message there is no thread yet, so what is measured is the one a new + // chat would start with. A preselected profile or a chat template is already part of + // that, and it may even bring an example conversation along -- reporting nothing for all + // of it would tell a person their window is empty while their first message is not. + // + var thread = this.ChatThread ?? this.NewChatThread(string.Empty); + provider = this.Provider; + parts = ConversationParts.Of(thread, this.BuildSystemPromptFor(thread), this.UserInput, this.ComposerState.FileAttachments, provider.SupportsImageInput()); + }); + + var counted = await this.ConversationTokenCounter.CountAsync(provider, parts, token); + if (token.IsCancellationRequested) return; - this.conversationTokens = counted; - this.StateHasChanged(); + await this.InvokeAsync(() => + { + if (counted == this.conversationTokens) + return; + + this.conversationTokens = counted; + this.StateHasChanged(); + }); } /// @@ -1477,9 +1534,6 @@ public partial class ChatComponent : MSGComponentBase this.hasUnsavedChanges = true; if(this.autoSaveEnabled) await this.SaveThread(); - - // The answer just became part of what every further message carries: - await this.CalculateTokenCount(); break; case Event.WORKSPACE_RENAMED: @@ -1537,6 +1591,10 @@ public partial class ChatComponent : MSGComponentBase protected override async ValueTask DisposeResourcesAsync() { this.MediaTranscriptionService.StateChanged -= this.OnMediaImportStateChanged; + + if (this.tokenTracker is not null) + await this.tokenTracker.DisposeAsync(); + if(this.SettingsManager.ConfigurationData.Workspace.StorageBehavior is WorkspaceStorageBehavior.STORE_CHATS_AUTOMATICALLY) { await this.SaveThread(); diff --git a/app/MindWork AI Studio/Tools/Services/ConversationTokenCounter.cs b/app/MindWork AI Studio/Tools/Services/ConversationTokenCounter.cs index 482482ad..fa491863 100644 Binary files a/app/MindWork AI Studio/Tools/Services/ConversationTokenCounter.cs and b/app/MindWork AI Studio/Tools/Services/ConversationTokenCounter.cs differ diff --git a/app/MindWork AI Studio/wwwroot/app.css b/app/MindWork AI Studio/wwwroot/app.css index 5c83cd77..223b0ba6 100644 --- a/app/MindWork AI Studio/wwwroot/app.css +++ b/app/MindWork AI Studio/wwwroot/app.css @@ -101,6 +101,19 @@ border-color: var(--confidence-color) !important; } +/* + * The token count under the chat input. It is a plain grey number until the conversation is near + * the model's context window, and then it says so by colour: there is nothing to do about four + * fifths of a window, and quite a lot to do about a full one. + */ +.token-budget-nearly-full .mud-input-helper-text { + color: var(--mud-palette-warning); +} + +.token-budget-exceeded .mud-input-helper-text { + color: var(--mud-palette-error); +} + :root { --custom-icon-color: #000000; } diff --git a/app/Tests/Chat/ConversationPartsTests.cs b/app/Tests/Chat/ConversationPartsTests.cs index 369eecf0..de1d3452 100644 --- a/app/Tests/Chat/ConversationPartsTests.cs +++ b/app/Tests/Chat/ConversationPartsTests.cs @@ -46,7 +46,47 @@ public sealed class ConversationPartsTests var parts = ConversationParts.Of(thread, "You are helpful.", "And of Italy?", null, imagesAreSent: true); - Assert.That(parts.Texts, Is.EqualTo(new[] { "You are helpful.", "What is the capital of France?", "Paris.", "And of Italy?" })); + Assert.Multiple(() => + { + Assert.That(parts.Texts, Is.EqualTo(new[] { "You are helpful.", "What is the capital of France?", "Paris." })); + Assert.That(parts.GrowingTexts, Is.EqualTo(new[] { "And of Italy?" })); + }); + } + + [Test] + public void WhatIsStillBeingWrittenIsKeptApartFromWhatStands() + { + // + // Both cost the same and both are counted. They are kept apart because of what happens + // afterwards: a message which stands says the same thing forever and its count is worth + // remembering, while the answer being streamed is a different text three seconds later. + // + var streaming = Block("The answer so far"); + ((ContentText)streaming.Content!).IsStreaming = true; + var thread = new ChatThread { Blocks = [Block("A question."), streaming] }; + + var parts = ConversationParts.Of(thread, string.Empty, "a draft", null, imagesAreSent: true); + + Assert.Multiple(() => + { + Assert.That(parts.Texts, Is.EqualTo(new[] { "A question." })); + Assert.That(parts.GrowingTexts, Is.EqualTo(new[] { "The answer so far", "a draft" })); + }); + } + + [Test] + public void AnAnswerWhichIsFinishedStandsLikeAnyOtherMessage() + { + var finished = Block("The whole answer."); + ((ContentText)finished.Content!).IsStreaming = false; + + var parts = ConversationParts.Of(new() { Blocks = [finished] }, string.Empty, string.Empty, null, imagesAreSent: true); + + Assert.Multiple(() => + { + Assert.That(parts.Texts, Is.EqualTo(new[] { "The whole answer." })); + Assert.That(parts.GrowingTexts, Is.Empty); + }); } [Test] @@ -85,7 +125,11 @@ public sealed class ConversationPartsTests { var parts = ConversationParts.Of(null, string.Empty, "Hello", null, imagesAreSent: true); - Assert.That(parts.Texts, Is.EqualTo(new[] { "Hello" })); + Assert.Multiple(() => + { + Assert.That(parts.Texts, Is.Empty); + Assert.That(parts.GrowingTexts, Is.EqualTo(new[] { "Hello" })); + }); } [TestCase("")] @@ -94,7 +138,11 @@ public sealed class ConversationPartsTests { var parts = ConversationParts.Of(null, string.Empty, draft, null, imagesAreSent: true); - Assert.That(parts.Texts, Is.Empty); + Assert.Multiple(() => + { + Assert.That(parts.Texts, Is.Empty); + Assert.That(parts.GrowingTexts, Is.Empty); + }); } [Test] diff --git a/app/Tests/Chat/ConversationTokenTrackerTests.cs b/app/Tests/Chat/ConversationTokenTrackerTests.cs new file mode 100644 index 00000000..3dcd3c1a --- /dev/null +++ b/app/Tests/Chat/ConversationTokenTrackerTests.cs @@ -0,0 +1,210 @@ +using AIStudio.Chat; + +namespace AIStudio.Tests.Chat; + +/// +/// Checks when the token count is recomputed and when it is not. +/// +/// +/// This is the part which kept going wrong. The number used to be wired to the places which change +/// the conversation -- fifteen of them in the end -- and four review rounds each found another place +/// which had been forgotten. So it is no longer wired to anything: whoever suspects a change nudges, +/// and what is checked here is that the tracker turns those nudges into the right amount of work. +/// +/// The waits are generous on purpose. What is asserted is the behaviour, not the clock, so every +/// interval here is far enough apart that a busy build machine cannot turn one into the other. +/// +[TestFixture] +public sealed class ConversationTokenTrackerTests +{ + /// + /// A heartbeat which never fires, for the tests which are about nudges alone. + /// + private static readonly TimeSpan NO_HEARTBEAT = Timeout.InfiniteTimeSpan; + + [Test] + public async Task ManyNudgesInARowCostOneCount() + { + // + // Loading a chat touches several things one after the other, and every one of them renders. + // Counting once per render would measure the same conversation half a dozen times. + // + var runs = 0; + var firstRun = new TaskCompletionSource(); + + await using var tracker = new ConversationTokenTracker(_ => + { + Interlocked.Increment(ref runs); + firstRun.TrySetResult(); + return Task.CompletedTask; + }, () => TimeSpan.FromSeconds(2), NO_HEARTBEAT); + + tracker.Start(); + for (var i = 0; i < 50; i++) + tracker.Nudge(); + + await firstRun.Task.WaitAsync(TimeSpan.FromSeconds(5)); + await Task.Delay(200); + + Assert.That(runs, Is.EqualTo(1)); + } + + [Test] + public async Task ANudgeArrivingDuringACountLeadsToExactlyOneMore() + { + // + // Something changed while we were reading, so the answer we just worked out may already be + // out of date -- but only one further count can be needed, however many nudges arrived. + // + var runs = 0; + var firstRunStarted = new TaskCompletionSource(); + var releaseFirstRun = new TaskCompletionSource(); + + await using var tracker = new ConversationTokenTracker(async _ => + { + if (Interlocked.Increment(ref runs) is not 1) + return; + + firstRunStarted.TrySetResult(); + await releaseFirstRun.Task; + }, () => TimeSpan.FromMilliseconds(100), NO_HEARTBEAT); + + tracker.Start(); + tracker.Nudge(); + await firstRunStarted.Task.WaitAsync(TimeSpan.FromSeconds(5)); + + // + // Nudged while the first run is held up, five times over, because a render storm is what + // this has to survive. + // + for (var i = 0; i < 5; i++) + tracker.Nudge(); + + releaseFirstRun.SetResult(); + await Task.Delay(1_000); + + Assert.That(runs, Is.EqualTo(2)); + } + + [Test] + public async Task WithoutAnyNudgeTheHeartbeatStillCounts() + { + // + // For what happens outside AI Studio: an attached file somebody edits in another program + // changes what the next message costs, and nothing here renders because of it. + // + var runs = 0; + + await using var tracker = new ConversationTokenTracker(_ => + { + Interlocked.Increment(ref runs); + return Task.CompletedTask; + }, () => TimeSpan.FromMilliseconds(50), TimeSpan.FromMilliseconds(200)); + + tracker.Start(); + await Task.Delay(1_000); + + Assert.That(runs, Is.GreaterThanOrEqualTo(2)); + } + + [Test] + public async Task TheQuietTimeIsAskedAnewAfterEveryRun() + { + // + // Because the right answer changes with what is going on. Showing a new number renders, and + // a render nudges, so while something moves continuously this delay is the entire cadence + // -- and a chat which is waiting for a model wants a slower one than a chat which is not. + // + var runs = 0; + var asked = 0; + + await using var tracker = new ConversationTokenTracker(_ => + { + Interlocked.Increment(ref runs); + return Task.CompletedTask; + }, () => + { + Interlocked.Increment(ref asked); + return TimeSpan.FromMilliseconds(50); + }, TimeSpan.FromMilliseconds(100)); + + tracker.Start(); + await Task.Delay(1_000); + + Assert.Multiple(() => + { + Assert.That(runs, Is.GreaterThanOrEqualTo(2)); + Assert.That(asked, Is.EqualTo(runs)); + }); + } + + [Test] + public async Task NothingIsCountedAfterTheTrackerIsGone() + { + var runs = 0; + var tracker = new ConversationTokenTracker(_ => + { + Interlocked.Increment(ref runs); + return Task.CompletedTask; + }, () => TimeSpan.FromMilliseconds(50), TimeSpan.FromMilliseconds(100)); + + tracker.Start(); + await Task.Delay(400); + await tracker.DisposeAsync(); + + var afterDisposal = runs; + await Task.Delay(400); + + Assert.Multiple(() => + { + Assert.That(afterDisposal, Is.GreaterThan(0), "The tracker never ran, so this proves nothing about stopping it."); + Assert.That(runs, Is.EqualTo(afterDisposal)); + }); + } + + [Test] + public async Task ACountWhichHangsDoesNotHoldUpDisposal() + { + // + // Counting ends in an IPC call to the runtime, and a component going away must not wait for + // one which is not coming back. The token handed to the work is the way out, and this is + // the test that it really is one. + // + var running = new TaskCompletionSource(); + var tracker = new ConversationTokenTracker(async token => + { + running.TrySetResult(); + await Task.Delay(Timeout.InfiniteTimeSpan, token); + }, () => TimeSpan.FromMilliseconds(50), NO_HEARTBEAT); + + tracker.Start(); + tracker.Nudge(); + await running.Task.WaitAsync(TimeSpan.FromSeconds(5)); + + var disposal = tracker.DisposeAsync().AsTask(); + var finishedInTime = await Task.WhenAny(disposal, Task.Delay(TimeSpan.FromSeconds(2))) == disposal; + + Assert.That(finishedInTime, Is.True); + } + + [Test] + public async Task AFailedCountDoesNotEndTheTracker() + { + // + // A tracker which died on one bad answer would leave a stale number standing forever, which + // is the one failure this whole mechanism exists to rule out. + // + var runs = 0; + + await using var tracker = new ConversationTokenTracker(_ => + { + Interlocked.Increment(ref runs); + throw new InvalidOperationException("The tokenizer did not answer."); + }, () => TimeSpan.FromMilliseconds(50), TimeSpan.FromMilliseconds(100)); + + tracker.Start(); + await Task.Delay(1_000); + + Assert.That(runs, Is.GreaterThanOrEqualTo(2)); + } +} \ No newline at end of file