Count the whole conversation against the model's window

app/MindWork AI Studio/Components/ChatComponent.razor.cs

@@ -1,3 +1,5 @@
+using System.Globalization;
+
 using AIStudio.Chat;
 using AIStudio.Dialogs;
 using AIStudio.Provider;
@@ -54,9 +56,10 @@ public partial class ChatComponent : MSGComponentBase
 
     [Inject]
     private IDialogService DialogService { get; init; } = null!;
+
+    [Inject]
+    private ConversationTokenCounter ConversationTokenCounter { get; init; } = null!;
     
-    [Inject] 
-    private RustService RustService { get; init; } = null!;
     [Inject]
     private IJSRuntime JsRuntime { get; init; } = null!;
 
@@ -93,11 +96,57 @@ public partial class ChatComponent : MSGComponentBase
     private Guid loadedParameterWorkspaceId = Guid.Empty;
     private Guid foregroundChatId = Guid.Empty;
     private int workspaceHeaderSyncVersion;
-    private string tokenCount = "0";
-    private bool HasCustomTokenizer => !string.IsNullOrWhiteSpace(this.Provider.TokenizerPath);
-    private string TokenCountMessage => this.HasCustomTokenizer
-        ? $"{this.T("Estimated amount of tokens:")} {this.tokenCount}"
-        : string.Empty;
+    private ConversationTokens conversationTokens = ConversationTokens.UNAVAILABLE;
+
+    /// <summary>
+    /// The culture the token numbers are written in.
+    /// </summary>
+    /// <remarks>
+    /// Taken from the language plugin the user chose, not from the machine. AI Studio's language is
+    /// a setting of its own, and a German who set German would otherwise read English separators
+    /// inside a German sentence -- where "1,234" means something a thousand times smaller.
+    /// </remarks>
+    private CultureInfo currentCulture = CultureInfo.InvariantCulture;
+
+    /// <summary>
+    /// What the helper text under the input field says about the token budget.
+    /// </summary>
+    /// <remarks>
+    /// Four sentences rather than one built from pieces, because a translator needs to see the
+    /// whole thing: which of the two numbers is the limit, and where the word for "about" belongs,
+    /// are decisions no language makes the same way.
+    ///
+    /// The images are named rather than counted. Every vendor charges a picture differently, and
+    /// none of those rules can be applied without decoding the file, so the honest answer is to say
+    /// how many of them the number does not include.
+    /// </remarks>
+    private string TokenCountMessage
+    {
+        get
+        {
+            if (!this.conversationTokens.IsKnown)
+                return string.Empty;
+
+            var used = TokenAmount.Format(this.conversationTokens.Tokens, this.currentCulture);
+            var budget = this.conversationTokens.Window.IsKnown
+                ? string.Format(this.conversationTokens.IsEstimate ? this.T("approx. {0} of {1} tokens") : this.T("{0} of {1} tokens"), used, TokenAmount.Format(this.conversationTokens.Window.DefaultTokens, this.currentCulture))
+                : string.Format(this.conversationTokens.IsEstimate ? this.T("approx. {0} tokens") : this.T("{0} tokens"), used);
+
+            if (this.conversationTokens.UncountedImages is 0)
+                return budget;
+
+            return $"{budget} {string.Format(this.T("plus {0} image(s), which cannot be counted"), this.conversationTokens.UncountedImages)}";
+        }
+    }
+
+    /// <summary>
+    /// Takes over the culture of the language the user chose for AI Studio.
+    /// </summary>
+    private async Task RefreshCulture()
+    {
+        var activeLanguagePlugin = await this.SettingsManager.GetActiveLanguagePlugin();
+        this.currentCulture = CommonTools.DeriveActiveCultureOrInvariant(activeLanguagePlugin.IETFTag);
+    }
 
     private MediaImportOwner CurrentMediaImportOwner => MediaImportOwner.ForChat(this.ChatThread?.ChatId ?? this.draftMediaOwnerId);
 
@@ -125,6 +174,7 @@ public partial class ChatComponent : MSGComponentBase
     protected override async Task OnInitializedAsync()
     {
         this.MediaTranscriptionService.StateChanged += this.OnMediaImportStateChanged;
+        await this.RefreshCulture();
 
         // 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 ]);
@@ -380,21 +430,22 @@ public partial class ChatComponent : MSGComponentBase
     protected override async Task OnParametersSetAsync()
     {
         var incomingChatId = this.ChatThread?.ChatId ?? Guid.Empty;
-        var providerChanged = this.Provider != this.lastSeenProvider;
         if (incomingChatId != this.lastSeenChatId || this.Provider != this.lastSeenProvider)
         {
             this.lastSeenChatId = incomingChatId;
             this.lastSeenProvider = this.Provider;
-            if (providerChanged)
-                this.tokenCount = "0";
-
             this.previousInputForbidden = true;
         }
 
         await this.ApplyLoadedChatParameterAsync();
         await this.SyncForegroundChatAsync();
-        if (providerChanged && this.HasCustomTokenizer)
-            await this.CalculateTokenCount();
+
+        //
+        // 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();
@@ -603,15 +654,24 @@ public partial class ChatComponent : MSGComponentBase
     private async Task ProfileWasChanged(Profile profile)
     {
         this.currentProfile = this.SettingsManager.GetProfileById(profile.Id);
-        if(this.ChatThread is null)
-            return;
 
-        this.ChatThread = this.ChatThread with
+        //
+        // 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.
+        //
+        if (this.ChatThread is not null)
         {
-            SelectedProfile = this.currentProfile.Id,
-        };
-        
-        await this.ChatThreadChanged.InvokeAsync(this.ChatThread);
+            this.ChatThread = this.ChatThread with
+            {
+                SelectedProfile = this.currentProfile.Id,
+            };
+
+            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)
@@ -623,10 +683,15 @@ public partial class ChatComponent : MSGComponentBase
         // Apply template's file attachments (replaces existing):
         this.ComposerState.ReplaceFileAttachments(this.currentChatTemplate.FileAttachments);
 
-        if(this.ChatThread is null)
-            return;
+        if (this.ChatThread is not null)
+            await this.StartNewChat(true);
 
-        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()
@@ -664,6 +729,9 @@ 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<DataSourceAgentSelected> GetAgentSelectedDataSources()
@@ -757,13 +825,20 @@ public partial class ChatComponent : MSGComponentBase
         this.hasUnsavedChanges = true;
     }
 
-    private void ComposerAttachmentsChanged(HashSet<FileAttachment> attachments)
+    private async Task ComposerAttachmentsChanged(HashSet<FileAttachment> 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();
     }
 
     /// <summary>Creates and stores a stable draft immediately after media import confirmation.</summary>
@@ -774,21 +849,13 @@ public partial class ChatComponent : MSGComponentBase
 
         this.RefreshCurrentProfileAndChatTemplate();
         var promptName = this.ExtractThreadName(this.ComposerState.UserInput);
-        this.ChatThread = new()
+        var threadName = string.IsNullOrWhiteSpace(this.ComposerState.UserInput)
+            ? $"Transkription: {Path.GetFileName(firstMediaPath)}"
+            : promptName;
+
+        this.ChatThread = this.NewChatThread(threadName) with
         {
-            IncludeDateTime = true,
-            SelectedProvider = this.Provider.Id,
-            SelectedProfile = this.currentProfile.Id,
-            SelectedChatTemplate = this.currentChatTemplate.Id,
-            SelectedToolIds = [..this.selectedToolIds],
-            SystemPrompt = SystemPrompts.DEFAULT,
-            WorkspaceId = this.currentWorkspaceId,
-            ChatId = Guid.NewGuid(),
             DataSourceOptions = this.earlyDataSourceOptions,
-            Name = string.IsNullOrWhiteSpace(this.ComposerState.UserInput)
-                ? $"Transkription: {Path.GetFileName(firstMediaPath)}"
-                : promptName,
-            Blocks = this.currentChatTemplate == ChatTemplate.NO_CHAT_TEMPLATE ? [] : this.currentChatTemplate.ExampleConversation.Select(block => block.DeepClone()).ToList(),
         };
 
         await WorkspaceBehaviour.StoreChatAsync(this.ChatThread);
@@ -819,21 +886,11 @@ public partial class ChatComponent : MSGComponentBase
         // Create a new chat thread if necessary:
         if (this.ChatThread is null)
         {
-            this.ChatThread = new()
+            this.ChatThread = this.NewChatThread(this.ExtractThreadName(this.ComposerState.UserInput)) with
             {
-                IncludeDateTime = true,
-                SelectedProvider = this.Provider.Id,
-                SelectedProfile = this.currentProfile.Id,
-                SelectedChatTemplate = this.currentChatTemplate.Id,
-                SelectedToolIds = [..this.selectedToolIds],
-                SystemPrompt = SystemPrompts.DEFAULT,
-                WorkspaceId = this.currentWorkspaceId,
-                ChatId = Guid.NewGuid(),
                 DataSourceOptions = this.earlyDataSourceOptions,
-                Name = this.ExtractThreadName(this.ComposerState.UserInput),
-                Blocks = this.currentChatTemplate == ChatTemplate.NO_CHAT_TEMPLATE ? [] : this.currentChatTemplate.ExampleConversation.Select(x => x.DeepClone()).ToList(),
             };
-            
+
             this.MarkCurrentChatAsLoadedParameter();
             await this.ChatThreadChanged.InvokeAsync(this.ChatThread);
         }
@@ -916,8 +973,14 @@ public partial class ChatComponent : MSGComponentBase
         this.ComposerState.Clear();
 
         await this.inputField.BlurAsync();
-        this.tokenCount = "0";
-        
+
+        //
+        // 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;
         
@@ -962,7 +1025,7 @@ public partial class ChatComponent : MSGComponentBase
     private void ApplyToolSelectionOfLoadedChat() =>
         this.selectedToolIds = ToolSelectionRules.NormalizeSelection(this.ChatThread?.SelectedToolIds ?? this.SettingsManager.GetDefaultToolIds(Tools.Components.CHAT));
 
-    private Task SelectedToolIdsChanged(HashSet<string> updatedToolIds)
+    private async Task SelectedToolIdsChanged(HashSet<string> updatedToolIds)
     {
         this.selectedToolIds = ToolSelectionRules.NormalizeSelection(updatedToolIds);
 
@@ -978,7 +1041,11 @@ public partial class ChatComponent : MSGComponentBase
             this.hasUnsavedChanges = true;
         }
 
-        return Task.CompletedTask;
+        //
+        // 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()
@@ -1087,19 +1154,7 @@ public partial class ChatComponent : MSGComponentBase
             // reset the chat thread only. The workspace id and the workspace name remain
             // the same:
             //
-            this.ChatThread = new()
-            {
-                IncludeDateTime = true,
-                SelectedProvider = this.Provider.Id,
-                SelectedProfile = this.currentProfile.Id,
-                SelectedChatTemplate = this.currentChatTemplate.Id,
-                SelectedToolIds = [..this.selectedToolIds],
-                SystemPrompt = SystemPrompts.DEFAULT,
-                WorkspaceId = this.currentWorkspaceId,
-                ChatId = Guid.NewGuid(),
-                Name = string.Empty,
-                Blocks = this.currentChatTemplate == ChatTemplate.NO_CHAT_TEMPLATE ? [] : this.currentChatTemplate.ExampleConversation.Select(x => x.DeepClone()).ToList(),
-            };
+            this.ChatThread = this.NewChatThread(string.Empty);
         }
 
         this.ComposerState.ApplyTemplate(this.currentChatTemplate);
@@ -1111,8 +1166,9 @@ public partial class ChatComponent : MSGComponentBase
         await this.SyncForegroundChatAsync();
         this.MarkCurrentChatAsLoadedParameter();
         await this.ChatThreadChanged.InvokeAsync(this.ChatThread);
+        await this.CalculateTokenCount();
     }
-    
+
     private async Task MoveChatToWorkspace()
     {
         if(this.ChatThread is null)
@@ -1214,6 +1270,7 @@ public partial class ChatComponent : MSGComponentBase
         await this.SyncForegroundChatAsync();
         this.ApplyStandardDataSourceOptions();
         await this.ChatThreadChanged.InvokeAsync(this.ChatThread);
+        await this.CalculateTokenCount();
     }
     
     private async Task SelectProviderWhenLoadingChat()
@@ -1249,6 +1306,9 @@ 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)
@@ -1266,42 +1326,46 @@ public partial class ChatComponent : MSGComponentBase
         await this.SendMessage(reuseLastUserPrompt: true);
     }
     
-    private Task EditLastUserBlock(IContent block)
+    private async Task EditLastUserBlock(IContent block)
     {
         if(this.ChatThread is null)
-            return Task.CompletedTask;
-        
+            return;
+
         if (block is not ContentText textBlock)
-            return Task.CompletedTask;
-        
+            return;
+
         var lastBlock = this.ChatThread.Blocks.Last();
         var lastBlockContent = lastBlock.Content;
         if(lastBlockContent is null)
-            return Task.CompletedTask;
-        
+            return;
+
         this.RestoreComposerFromTextBlock(textBlock);
         this.ChatThread.Remove(block);
         this.ChatThread.Remove(lastBlockContent);
         this.hasUnsavedChanges = true;
         this.StateHasChanged();
-        
-        return Task.CompletedTask;
+
+        //
+        // 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();
     }
-    
-    private Task EditLastBlock(IContent block)
+
+    private async Task EditLastBlock(IContent block)
     {
         if(this.ChatThread is null)
-            return Task.CompletedTask;
-        
+            return;
+
         if (block is not ContentText textBlock)
-            return Task.CompletedTask;
-        
+            return;
+
         this.RestoreComposerFromTextBlock(textBlock);
         this.ChatThread.Remove(block);
         this.hasUnsavedChanges = true;
         this.StateHasChanged();
-        
-        return Task.CompletedTask;
+
+        await this.CalculateTokenCount();
     }
 
     private void RestoreComposerFromTextBlock(ContentText textBlock)
@@ -1309,42 +1373,92 @@ public partial class ChatComponent : MSGComponentBase
         this.ComposerState.RestoreFromTextBlock(textBlock);
     }
     
+    /// <summary>
+    /// Works out what the next request would take out of the model's context window.
+    /// </summary>
+    /// <remarks>
+    /// The whole conversation, not only what is being typed. A number counting the draft alone
+    /// answers a question nobody asks: what decides whether the next message fits is everything
+    /// which travels with it, and in a chat of any age the draft is the smallest part of that.
+    ///
+    /// This used to run only for providers with a tokenizer of their own, which is almost nobody,
+    /// so almost nobody ever saw a number. The runtime falls back to the tokenizer shipped with AI
+    /// Studio when a provider names none, so the count is available everywhere -- it is then an
+    /// 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.
+    /// </remarks>
     private async Task CalculateTokenCount()
     {
-        if (!this.HasCustomTokenizer)
-        {
-            if (this.tokenCount != "0")
-            {
-                this.tokenCount = "0";
-                this.StateHasChanged();
-            }
-
-            return;
-        }
-
         //
-        // 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. Counting is also
-        // triggered while parameters are set, which happens before that.
+        // 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.
         //
-        var currentInput = this.UserInput;
-        if (string.IsNullOrEmpty(currentInput))
-        {
-            this.tokenCount = "0";
+        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)
             return;
-        }
 
-        var response = await this.RustService.GetTokenCount(this.Provider, currentInput);
-        if (response is null)
-            return;
-        if (!response.Value.Success)
-        {
-            this.Logger.LogWarning("Failed to calculate token count: reason='{Reason}'", response.Value.Message);
-            return;
-        }
-        this.tokenCount = response.Value.TokenCount.ToString();
+        this.conversationTokens = counted;
         this.StateHasChanged();
     }
+
+    /// <summary>
+    /// Works out the system prompt a thread would send.
+    /// </summary>
+    /// <remarks>
+    /// Not the prompt a person typed: a chat template may replace it, the retrieved data of a data
+    /// source is appended to it, the selected profile adds a paragraph, and the policy of the
+    /// selected tools adds another. Switching a profile while writing therefore moves the number,
+    /// which is the whole reason this is asked rather than read off the thread.
+    ///
+    /// The tools are filtered for the provider the same way they are before sending, so that a tool
+    /// the provider is not trusted enough to receive does not count either.
+    /// </remarks>
+    /// <param name="thread">The thread to build the prompt for.</param>
+    /// <returns>The system prompt as it would be sent.</returns>
+    private string BuildSystemPromptFor(ChatThread thread)
+    {
+        var toolDefinitions = this.ToolRegistry.FilterToolIdsForProvider(this.Provider, this.selectedToolIds)
+            .Select(this.ToolRegistry.GetDefinition)
+            .Where(definition => definition is not null)
+            .Select(definition => definition!)
+            .ToList();
+
+        return thread.BuildSystemPrompt(this.SettingsManager, toolDefinitions).Text;
+    }
+
+    /// <summary>
+    /// The thread a new chat starts with, as the selections made so far decide it.
+    /// </summary>
+    /// <remarks>
+    /// In one place because three code paths used to write it out, and because the token count has
+    /// to measure the same thing they build. A count against a thread assembled differently from
+    /// the one which is then sent would be wrong in exactly the moment a person looks at it: before
+    /// they send their first message.
+    ///
+    /// The data source options are left out on purpose: two of the three callers set them and the
+    /// third replaces them right afterwards, so this stays the part they agree on.
+    /// </remarks>
+    /// <param name="name">The name of the thread.</param>
+    /// <returns>The new thread.</returns>
+    private ChatThread NewChatThread(string name) => new()
+    {
+        IncludeDateTime = true,
+        SelectedProvider = this.Provider.Id,
+        SelectedProfile = this.currentProfile.Id,
+        SelectedChatTemplate = this.currentChatTemplate.Id,
+        SelectedToolIds = [..this.selectedToolIds],
+        SystemPrompt = SystemPrompts.DEFAULT,
+        WorkspaceId = this.currentWorkspaceId,
+        ChatId = Guid.NewGuid(),
+        Name = name,
+        Blocks = this.currentChatTemplate == ChatTemplate.NO_CHAT_TEMPLATE ? [] : this.currentChatTemplate.ExampleConversation.Select(block => block.DeepClone()).ToList(),
+    };
     
     #region Overrides of MSGComponentBase
     
@@ -1363,6 +1477,9 @@ 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:
@@ -1372,6 +1489,7 @@ public partial class ChatComponent : MSGComponentBase
 
             case Event.CONFIGURATION_CHANGED:
             case Event.PLUGINS_RELOADED:
+                await this.RefreshCulture();
                 await this.RefreshChatSelectionsAfterConfigurationChange();
                 this.StateHasChanged();
                 break;