From 47af43fe1f8ba9f41bc523e7474bc3c009ddfb6c Mon Sep 17 00:00:00 2001 From: Andrew Beal Date: Fri, 17 Oct 2025 22:23:45 +0100 Subject: [PATCH] feat: add token counting display to status bar Implement real-time input/output token tracking in the status bar by integrating ITokenService with ChatService. Token counts are calculated from conversation history and updated on conversation load, reset, and message completion. --- Components/ChatWindow.svelte | 4 +++ Services/ChatService.ts | 44 ++++++++++++++++++++++++++++++--- Services/ServiceRegistration.ts | 3 ++- Services/StatusBarService.ts | 1 + 4 files changed, 48 insertions(+), 4 deletions(-) diff --git a/Components/ChatWindow.svelte b/Components/ChatWindow.svelte index 8d67ec8..0bca172 100644 --- a/Components/ChatWindow.svelte +++ b/Components/ChatWindow.svelte @@ -45,6 +45,7 @@ if (chatContainer) { plugin.registerDomEvent(chatContainer, 'click', handleLinkClick); } + chatService.setStatusBarTokens(0, 0); }); async function handleLinkClick(evt: MouseEvent) { @@ -119,6 +120,7 @@ onComplete: () => { isSubmitting = false; chatArea.onFinishedSubmitting(); + chatService.updateTokenDisplay(conversation); } }); } @@ -162,6 +164,7 @@ $: if ($conversationStore.shouldReset) { conversation = new Conversation(); + chatService.setStatusBarTokens(0, 0); conversationStore.clearResetFlag(); } @@ -170,6 +173,7 @@ conversation = loadedConversation; conversationService.setCurrentConversationPath(filePath); chatService.onNameChanged?.(loadedConversation.title); + chatService.updateTokenDisplay(loadedConversation); conversationStore.clearLoadFlag(); scrollToBottom(); } diff --git a/Services/ChatService.ts b/Services/ChatService.ts index 65cfe4b..5dd6e53 100644 --- a/Services/ChatService.ts +++ b/Services/ChatService.ts @@ -5,6 +5,9 @@ import type { IAIClass } from "AIClasses/IAIClass"; import type { ConversationFileSystemService } from "./ConversationFileSystemService"; import type { AIFunctionService } from "./AIFunctionService"; import type { ConversationNamingService } from "./ConversationNamingService"; +import type { ITokenService } from "AIClasses/ITokenService"; +import type { IPrompt } from "AIClasses/IPrompt"; +import type { StatusBarService } from "./StatusBarService"; import { Conversation } from "Conversations/Conversation"; import { ConversationContent } from "Conversations/ConversationContent"; import { Role } from "Enums/Role"; @@ -21,6 +24,9 @@ export class ChatService { private conversationService: ConversationFileSystemService; private aiFunctionService: AIFunctionService; private namingService: ConversationNamingService; + private tokenService: ITokenService | undefined; + private prompt: IPrompt; + private statusBarService: StatusBarService; private semaphore: Semaphore; private semaphoreHeld: boolean = false; @@ -30,16 +36,19 @@ export class ChatService { this.conversationService = Resolve(Services.ConversationFileSystemService); this.aiFunctionService = Resolve(Services.AIFunctionService); this.namingService = Resolve(Services.ConversationNamingService); + this.prompt = Resolve(Services.IPrompt); + this.statusBarService = Resolve(Services.StatusBarService); this.semaphore = new Semaphore(1, false); } public onNameChanged: ((name: string) => void) | undefined = undefined; - resolveAIProvider() { + public resolveAIProvider() { this.ai = Resolve(Services.IAIClass); + this.tokenService = Resolve(Services.ITokenService); } - async submit(conversation: Conversation, allowDestructiveActions: boolean, userRequest: string, callbacks: ChatServiceCallbacks): Promise { + public async submit(conversation: Conversation, allowDestructiveActions: boolean, userRequest: string, callbacks: ChatServiceCallbacks): Promise { if (!await this.semaphore.wait()) { return conversation; } @@ -96,7 +105,7 @@ export class ChatService { } } - stop(): void { + public stop(): void { if (this.abortController) { this.abortController.abort(); this.abortController = null; @@ -104,6 +113,35 @@ export class ChatService { this.semaphore.release(); } + public async updateTokenDisplay(conversation: Conversation): Promise { + if (this.tokenService === undefined) { + return; + } + + const systemInstruction = this.prompt.systemInstruction(); + const userInstruction = await this.prompt.userInstruction(); + + const inputMessages = conversation.contents + .filter(msg => msg.role === Role.User && !msg.isFunctionCallResponse) + .map(msg => msg.content) + .join("\n"); + + const outputMessages = conversation.contents + .filter(msg => msg.role === Role.Assistant && !msg.isFunctionCall) + .map(msg => msg.content) + .join("\n"); + + const inputText = systemInstruction + "\n" + userInstruction + "\n" + inputMessages; + const inputTokens = await this.tokenService.countTokens(inputText); + const outputTokens = await this.tokenService.countTokens(outputMessages); + + this.setStatusBarTokens(inputTokens, outputTokens); + } + + public setStatusBarTokens(inputTokens: number, outputTokens: number): void { + this.statusBarService.setStatusBarMessage(`Input Tokens: ${inputTokens} / Output Tokens: ${outputTokens}`); + } + private async streamRequestResponse( conversation: Conversation, allowDestructiveActions: boolean, callbacks: ChatServiceCallbacks ): Promise<{ functionCall: AIFunctionCall | null, shouldContinue: boolean }> { diff --git a/Services/ServiceRegistration.ts b/Services/ServiceRegistration.ts index 38d91ec..1ea329d 100644 --- a/Services/ServiceRegistration.ts +++ b/Services/ServiceRegistration.ts @@ -21,11 +21,12 @@ import { ConversationNamingService } from "./ConversationNamingService"; import { VaultService } from "./VaultService"; import type { ITokenService } from "AIClasses/ITokenService"; import { GeminiTokenService } from "AIClasses/Gemini/GeminiTokenService"; +import { StatusBarService } from "./StatusBarService"; export function RegisterDependencies(plugin: AIAgentPlugin) { RegisterSingleton(Services.AIAgentPlugin, plugin); RegisterSingleton(Services.FileManager, plugin.app.fileManager); - + RegisterSingleton(Services.StatusBarService, new StatusBarService()); RegisterSingleton(Services.VaultService, new VaultService()); RegisterSingleton(Services.WorkSpaceService, new WorkSpaceService()); RegisterSingleton(Services.FileSystemService, new FileSystemService()); diff --git a/Services/StatusBarService.ts b/Services/StatusBarService.ts index 5fccbfa..51282a9 100644 --- a/Services/StatusBarService.ts +++ b/Services/StatusBarService.ts @@ -1,6 +1,7 @@ import type AIAgentPlugin from "main"; import { Resolve } from "./DependencyService"; import { Services } from "./Services"; +import { Selector } from "Enums/Selector"; export class StatusBarService {