added token amount to chat (in progress)

app/MindWork AI Studio/Assistants/I18N/allTexts.lua

@@ -1618,6 +1618,9 @@ UI_TEXT_CONTENT["AISTUDIO::COMPONENTS::CHATCOMPONENT::T3403290862"] = "The selec
 -- Select a provider first
 UI_TEXT_CONTENT["AISTUDIO::COMPONENTS::CHATCOMPONENT::T3654197869"] = "Select a provider first"
 
+-- Estimated amount of tokens:
+UI_TEXT_CONTENT["AISTUDIO::COMPONENTS::CHATCOMPONENT::T377990776"] = "Estimated amount of tokens:"
+
 -- Start new chat in workspace '{0}'
 UI_TEXT_CONTENT["AISTUDIO::COMPONENTS::CHATCOMPONENT::T3928697643"] = "Start new chat in workspace '{0}'"
 

app/MindWork AI Studio/Components/ChatComponent.razor

@@ -123,7 +123,7 @@
                     <MudIconButton Icon="@Icons.Material.Filled.Error" Color="Color.Error"/>
                 </MudTooltip>
             }
-            <MudIconButton />
+            <MudText Typo="Typo.body2">@this.TokenCountMessage</MudText>
         </MudToolBar>
     </FooterContent>
 </InnerScrolling>

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

@@ -3,6 +3,7 @@ using AIStudio.Dialogs;
 using AIStudio.Provider;
 using AIStudio.Settings;
 using AIStudio.Settings.DataModel;
+using AIStudio.Tools.Services;
 
 using Microsoft.AspNetCore.Components;
 using Microsoft.AspNetCore.Components.Web;
@@ -37,6 +38,9 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable
     [Inject]
     private IDialogService DialogService { get; init; } = null!;
     
+    [Inject] 
+    private RustService RustService { get; init; } = null!;
+
     private const Placement TOOLBAR_TOOLTIP_PLACEMENT = Placement.Top;
     private static readonly Dictionary<string, object?> USER_INPUT_ATTRIBUTES = new();
 
@@ -58,6 +62,8 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable
     private Guid currentWorkspaceId = Guid.Empty;
     private CancellationTokenSource? cancellationTokenSource;
     private HashSet<FileAttachment> chatDocumentPaths = [];
+    private string tokenCount = "0";
+    private string TokenCountMessage => $"{this.T("Estimated amount of tokens:")} {this.tokenCount}";
 
     // Unfortunately, we need the input field reference to blur the focus away. Without
     // this, we cannot clear the input field.
@@ -405,6 +411,9 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable
         // 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)
         {
@@ -901,6 +910,18 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable
         return Task.CompletedTask;
     }
     
+    private async Task CalculateTokenCount()
+    {
+        if (this.inputField.Value is null)
+            return;
+        this.Logger.LogDebug($"Text to tokenize: '{this.inputField.Value}' ");
+        var response = await this.RustService.GetTokenCount(this.inputField.Value);
+        if (response is null)
+            return;
+        this.tokenCount = response.TokenCount.ToString();
+        this.Logger.LogDebug($"Token count: {this.tokenCount}");
+    }
+    
     #region Overrides of MSGComponentBase
     
     protected override async Task ProcessIncomingMessage<T>(ComponentBase? sendingComponent, Event triggeredEvent, T? data) where T : default

app/MindWork AI Studio/Tools/Rust/TokenCountInfo.cs

@@ -0,0 +1,6 @@
+namespace AIStudio.Tools.Rust;
+
+public sealed class TokenCountInfo
+{
+    public int TokenCount { get; set; }
+}

app/MindWork AI Studio/Tools/Services/RustService.Tokenizer.cs

@@ -0,0 +1,27 @@
+using AIStudio.Tools.Rust;
+
+namespace AIStudio.Tools.Services;
+
+public sealed partial class RustService
+{
+    public async Task<TokenCountInfo?> GetTokenCount(string text)
+    {
+        try
+        {
+            var cts = new CancellationTokenSource(TimeSpan.FromSeconds(5));
+            var payload = new { text };
+            var response = await this.http.PostAsJsonAsync("/system/tokenizer/count", payload, this.jsonRustSerializerOptions, cts.Token);
+            response.EnsureSuccessStatusCode();
+            return await response.Content.ReadFromJsonAsync<TokenCountInfo>(this.jsonRustSerializerOptions, cancellationToken: cts.Token);
+        }
+        catch (Exception e)
+        {
+            if(this.logger is not null)
+                this.logger.LogError(e, "Error while getting token count from Rust service.");
+            else
+                Console.WriteLine($"Error while getting token count from Rust service: '{e}'.");
+            
+            return null;
+        }
+    }
+}

runtime/Cargo.toml

@@ -42,6 +42,9 @@ pptx-to-md = "0.4.0"
 tempfile = "3.8"
 strum_macros = "0.27"
 sysinfo = "0.38.0"
+tiktoken-rs = "0.9.1"
+tokenizers = "0.22.2"
+hf-hub = "0.4.3"
 
 # Fixes security vulnerability downstream, where the upstream is not fixed yet:
 time = "0.3.47" # -> Rocket

runtime/src/lib.rs

@@ -18,3 +18,4 @@ pub mod certificate_factory;
 pub mod runtime_api_token;
 pub mod stale_process_cleanup;
 mod sidecar_types;
+pub mod tokenizer;

runtime/src/main.rs

@@ -11,7 +11,8 @@ use mindwork_ai_studio::environment::is_dev;
 use mindwork_ai_studio::log::init_logging;
 use mindwork_ai_studio::metadata::MetaData;
 use mindwork_ai_studio::runtime_api::start_runtime_api;
-
+use mindwork_ai_studio::stale_process_cleanup::kill_stale_process;
+use mindwork_ai_studio::tokenizer::{init_tokenizer, get_token_count};
 
 #[tokio::main]
 async fn main() {
@@ -43,8 +44,17 @@ async fn main() {
         info!("Running in production mode.");
     }
 
+    if let Err(e) = init_tokenizer() {
+        warn!(Source = "Tokenizer"; "Error during the initialisation of the tokenizer: {}", e);
+    }
+
+    let token_count = get_token_count("");
+    info!(".. Tokenizer contains {token_count} tokens.");
+
     generate_runtime_certificate();
     start_runtime_api();
     
     start_tauri();
+
+    Ok(())
 }

runtime/src/runtime_api.rs

@@ -91,6 +91,7 @@ pub fn start_runtime_api() {
                 crate::file_data::extract_data,
                 crate::log::get_log_paths,
                 crate::log::log_event,
+                crate::tokenizer::tokenizer_count,
                 crate::app_window::register_shortcut,
                 crate::app_window::validate_shortcut,
                 crate::app_window::suspend_shortcuts,