From 2444ca08113784ec6645645606f420a83fca2634 Mon Sep 17 00:00:00 2001 From: matt Date: Thu, 16 Oct 2025 13:24:49 -0700 Subject: [PATCH] feat: implement token counter with real-time breakdown MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add comprehensive token counting functionality to provide visibility into context usage: Backend (Rust): - Add tiktoken-rs dependency for OpenAI-compatible token counting - Implement get_token_count command with detailed breakdown - Count tokens for: system prompt, preset instructions, persona, world info, author's note, message history, and current input - Per-section token breakdown for optimization insights Frontend (JavaScript/HTML/CSS): - Add token counter widget in status bar - Real-time updates as user types (debounced 300ms) - Expandable breakdown tooltip showing per-section counts - Automatic update when chat history loads or changes - Clean, minimal UI with hover interactions Features: - Accurate token counting using cl100k_base tokenizer - Debounced updates for performance - Detailed breakdown by context section - Visual indicator with total token count - Click to expand/collapse detailed breakdown - Auto-hide when no character is active This completes the "Must-Have for Basic Roleplay" features from the roadmap: ✅ World Info/Lorebooks ✅ Author's Note ✅ Token Counter - Message Examples Usage (next) 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude --- src-tauri/Cargo.toml | 2 + src-tauri/src/lib.rs | 131 +++++++++++++++++++++++++++++++++++++++++++ src/index.html | 47 ++++++++++++++++ src/main.js | 61 ++++++++++++++++++++ src/styles.css | 91 ++++++++++++++++++++++++++++++ 5 files changed, 332 insertions(+) diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml index da1d4ea..ddcfff7 100644 --- a/src-tauri/Cargo.toml +++ b/src-tauri/Cargo.toml @@ -32,4 +32,6 @@ png = "0.17" base64 = "0.21" image = "0.24" regex = "1" +chrono = "0.4" +tiktoken-rs = "0.5" diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 8e7b15d..c1a703f 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -9,6 +9,7 @@ use base64::Engine; use regex::Regex; use std::sync::{Mutex, OnceLock}; use std::collections::HashMap; +use tiktoken_rs::cl100k_base; #[derive(Debug, Clone, Serialize, Deserialize)] struct ApiConfig { @@ -2542,6 +2543,135 @@ fn restore_builtin_preset(preset_id: String) -> Result { .ok_or_else(|| format!("Built-in preset '{}' not found", preset_id)) } +// Token Counting + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct TokenBreakdown { + total: usize, + system_prompt: usize, + preset_instructions: usize, + persona: usize, + world_info: usize, + authors_note: usize, + message_history: usize, + current_input: usize, + estimated_max_tokens: usize, +} + +// Helper function to count tokens in a string +fn count_tokens(text: &str) -> usize { + if text.is_empty() { + return 0; + } + + let bpe = cl100k_base().unwrap(); + bpe.encode_with_special_tokens(text).len() +} + +#[tauri::command] +fn get_token_count(character_id: Option, current_input: String) -> Result { + // Get character (either specified or active) + let character = if let Some(id) = character_id { + load_character(&id).ok_or_else(|| format!("Character '{}' not found", id))? + } else { + get_active_character() + }; + + let history = load_history(&character.id); + let roleplay_settings = load_roleplay_settings(&character.id); + + // Build the same context that would be sent to the API + let (_system_additions, authors_note, _note_depth) = build_roleplay_context(&character, &history.messages, &roleplay_settings); + + // Count system prompt (including template processing) + let processed_system_prompt = replace_template_variables(&character.system_prompt, &character, &roleplay_settings); + let system_prompt_tokens = count_tokens(&processed_system_prompt); + + // Parse system additions to break down by component + let mut preset_tokens = 0; + let mut persona_tokens = 0; + let mut world_info_tokens = 0; + + // Count preset instructions and system additions + if let Some(preset_id) = &roleplay_settings.active_preset_id { + if let Some(preset) = load_preset(preset_id) { + if !preset.system_additions.is_empty() { + let processed_additions = replace_template_variables(&preset.system_additions, &character, &roleplay_settings); + preset_tokens += count_tokens(&processed_additions); + } + + let mut enabled_instructions: Vec<_> = preset.instructions.iter() + .filter(|i| i.enabled) + .collect(); + enabled_instructions.sort_by_key(|i| i.order); + + for instruction in enabled_instructions { + let processed_content = replace_template_variables(&instruction.content, &character, &roleplay_settings); + preset_tokens += count_tokens(&processed_content); + } + } + } + + // Count persona + if roleplay_settings.persona_enabled { + if let Some(name) = &roleplay_settings.persona_name { + if let Some(desc) = &roleplay_settings.persona_description { + let processed_desc = replace_template_variables(desc, &character, &roleplay_settings); + let persona_text = format!("\n\n[{}'s Persona: {}]", name, processed_desc); + persona_tokens = count_tokens(&persona_text); + } + } + } + + // Count world info + let activated_entries = scan_for_world_info(&history.messages, &roleplay_settings.world_info, roleplay_settings.scan_depth, roleplay_settings.recursion_depth); + if !activated_entries.is_empty() { + let mut wi_text = String::from("\n\n[Relevant World Information:"); + for entry in activated_entries { + let processed_content = replace_template_variables(&entry.content, &character, &roleplay_settings); + wi_text.push_str(&format!("\n- {}", processed_content)); + } + wi_text.push_str("\n]"); + world_info_tokens = count_tokens(&wi_text); + } + + // Count author's note + let authors_note_tokens = if let Some(note) = authors_note { + let note_text = format!("[Author's Note: {}]", note); + count_tokens(¬e_text) + } else { + 0 + }; + + // Count message history + let mut history_tokens = 0; + for msg in &history.messages { + history_tokens += count_tokens(msg.get_content()); + } + + // Count current input + let input_tokens = count_tokens(¤t_input); + + // Calculate total + let total = system_prompt_tokens + preset_tokens + persona_tokens + world_info_tokens + + authors_note_tokens + history_tokens + input_tokens; + + // Estimate remaining tokens for response (assuming 16k context with 4k max response) + let estimated_max_tokens = if total < 12000 { 4096 } else { 16384 - total }; + + Ok(TokenBreakdown { + total, + system_prompt: system_prompt_tokens, + preset_instructions: preset_tokens, + persona: persona_tokens, + world_info: world_info_tokens, + authors_note: authors_note_tokens, + message_history: history_tokens, + current_input: input_tokens, + estimated_max_tokens, + }) +} + // World Info Commands #[tauri::command] @@ -2842,6 +2972,7 @@ pub fn run() { duplicate_preset, is_builtin_preset_modified, restore_builtin_preset, + get_token_count, add_world_info_entry, update_world_info_entry, delete_world_info_entry, diff --git a/src/index.html b/src/index.html index 438acfc..61c8d00 100644 --- a/src/index.html +++ b/src/index.html @@ -637,6 +637,53 @@
Ready + +
+ + diff --git a/src/main.js b/src/main.js index f8e167c..503442f 100644 --- a/src/main.js +++ b/src/main.js @@ -1366,9 +1366,67 @@ function setupKeyboardShortcuts() { messageInput.addEventListener('input', () => { autoResize(messageInput); + updateTokenCount(); }); } +// Token Counter +let tokenUpdateTimeout = null; + +async function updateTokenCount() { + // Debounce token count updates + if (tokenUpdateTimeout) { + clearTimeout(tokenUpdateTimeout); + } + + tokenUpdateTimeout = setTimeout(async () => { + try { + const currentInput = messageInput.value; + const tokenData = await invoke('get_token_count', { + characterId: null, // Use active character + currentInput + }); + + // Update total display + const tokenCounter = document.getElementById('token-counter'); + const tokenCountTotal = document.getElementById('token-count-total'); + tokenCountTotal.textContent = `${tokenData.total} tokens`; + tokenCounter.style.display = 'flex'; + + // Update breakdown + document.getElementById('token-system').textContent = tokenData.system_prompt; + document.getElementById('token-preset').textContent = tokenData.preset_instructions; + document.getElementById('token-persona').textContent = tokenData.persona; + document.getElementById('token-worldinfo').textContent = tokenData.world_info; + document.getElementById('token-authorsnote').textContent = tokenData.authors_note; + document.getElementById('token-history').textContent = tokenData.message_history; + document.getElementById('token-input').textContent = tokenData.current_input; + document.getElementById('token-total-detail').textContent = tokenData.total; + } catch (error) { + console.error('Failed to update token count:', error); + // Hide token counter on error + document.getElementById('token-counter').style.display = 'none'; + } + }, 300); // Update after 300ms of no typing +} + +// Toggle token breakdown display +document.getElementById('token-details-btn').addEventListener('click', (e) => { + e.stopPropagation(); + const breakdown = document.getElementById('token-breakdown'); + breakdown.style.display = breakdown.style.display === 'none' ? 'block' : 'none'; +}); + +// Close breakdown when clicking outside +document.addEventListener('click', (e) => { + const breakdown = document.getElementById('token-breakdown'); + const detailsBtn = document.getElementById('token-details-btn'); + + if (!breakdown.contains(e.target) && !detailsBtn.contains(e.target)) { + breakdown.style.display = 'none'; + } +}); + // Load characters and populate dropdown async function loadCharacters() { console.log('Loading characters...'); @@ -1535,6 +1593,9 @@ async function loadChatHistory() { messagesContainer.innerHTML = ''; addMessage('API configured. Ready to chat.', false, true); } + + // Update token count after loading history + updateTokenCount(); } // Clear chat history diff --git a/src/styles.css b/src/styles.css index 74ab542..39f9e51 100644 --- a/src/styles.css +++ b/src/styles.css @@ -761,6 +761,97 @@ body { color: #22c55e; } +/* Token Counter */ +.token-counter { + display: flex; + align-items: center; + gap: 8px; + font-size: 12px; + color: var(--text-secondary); +} + +.token-count { + padding: 4px 8px; + background: var(--bg-tertiary); + border: 1px solid var(--border); + border-radius: 6px; +} + +.token-details-btn { + width: 20px; + height: 20px; + border: none; + background: transparent; + color: var(--text-secondary); + border-radius: 4px; + cursor: pointer; + display: flex; + align-items: center; + justify-content: center; + transition: all 0.2s ease; + padding: 0; +} + +.token-details-btn:hover { + background: var(--bg-tertiary); + color: var(--accent); +} + +.token-breakdown { + position: absolute; + bottom: 52px; + right: 16px; + background: var(--bg-secondary); + border: 1px solid var(--border); + border-radius: 8px; + padding: 12px; + box-shadow: 0 4px 12px rgba(0, 0, 0, 0.3); + z-index: 100; + min-width: 250px; +} + +.token-breakdown-header { + font-size: 13px; + font-weight: 600; + color: var(--text-primary); + margin-bottom: 8px; + padding-bottom: 8px; + border-bottom: 1px solid var(--border); +} + +.token-breakdown-list { + display: flex; + flex-direction: column; + gap: 6px; +} + +.token-breakdown-item { + display: flex; + justify-content: space-between; + font-size: 12px; + color: var(--text-secondary); +} + +.token-breakdown-item span:first-child { + color: var(--text-secondary); +} + +.token-breakdown-item span:last-child { + color: var(--text-primary); + font-weight: 500; +} + +.token-breakdown-total { + display: flex; + justify-content: space-between; + font-size: 13px; + font-weight: 600; + color: var(--accent); + margin-top: 8px; + padding-top: 8px; + border-top: 1px solid var(--border); +} + @keyframes pulse { 0%, 100% { opacity: 1;