fix: token usage statistics

Author: cyberelf

src/components/chat/messageHandler.ts

@@ -14,6 +14,11 @@ interface OpenAIResponse {
       content?: string;
     };
   }>;
+  usage?: {
+    total_tokens: number;
+    prompt_tokens: number;
+    completion_tokens: number;
+  };
   error?: {
     message: string;
   };
@@ -29,6 +34,27 @@ interface OpenAIRequest {
   max_tokens?: number;
 }
 
+interface TokenUsage {
+  totalTokens: number;
+  requestCount: number;
+}
+
+async function updateTokenUsage(newTokens: number): Promise<void> {
+  const result = await new Promise<{ tokenUsage?: TokenUsage }>((resolve) => {
+    chrome.storage.local.get(['tokenUsage'], resolve);
+  });
+  
+  const currentUsage = result.tokenUsage || { totalTokens: 0, requestCount: 0 };
+  const updatedUsage = {
+    totalTokens: currentUsage.totalTokens + newTokens,
+    requestCount: currentUsage.requestCount + 1
+  };
+  
+  await new Promise<void>((resolve) => {
+    chrome.storage.local.set({ tokenUsage: updatedUsage }, resolve);
+  });
+}
+
 export async function handleQuestion(question: string, context: string, model?: ModelType): Promise<void> {
   try {
     // Check if the question is a shortcut
@@ -50,6 +76,32 @@ export async function handleQuestion(question: string, context: string, model?:
     const selectedModel = model || settings.model;
     await addMessage('assistant', '', selectedModel);
 
+    // First make a non-streaming request to get token usage
+    const nonStreamingResponse = await fetch(`${settings.apiUrl}/chat/completions`, {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json',
+        'Authorization': `Bearer ${settings.apiKey}`
+      },
+      body: JSON.stringify({
+        model: selectedModel,
+        messages: [
+          { role: 'system', content: systemMessage },
+          { role: 'user', content: finalQuestion }
+        ],
+        stream: false
+      } as OpenAIRequest)
+    });
+
+    if (!nonStreamingResponse.ok) {
+      const data = await nonStreamingResponse.json();
+      throw new Error(data.error?.message || 'Failed to get response from OpenAI');
+    }
+
+    const nonStreamingData = await nonStreamingResponse.json();
+    const totalTokens = nonStreamingData.usage?.total_tokens || 0;
+
+    // Now make the streaming request for the actual response
     const response = await fetch(`${settings.apiUrl}/chat/completions`, {
       method: 'POST',
       headers: {
@@ -100,15 +152,17 @@ export async function handleQuestion(question: string, context: string, model?:
           }
         }
       }
+
+      // Update token usage after successful completion
+      await updateTokenUsage(totalTokens);
     } catch (error) {
       console.error('Error reading stream:', error);
       throw error;
     } finally {
       reader.releaseLock();
     }
-
   } catch (error) {
-    console.error('Error handling question:', error);
-    await addMessage('assistant', `Error: ${error instanceof Error ? error.message : 'Unknown error occurred'}`, model || (await getSettings()).model);
+    console.error('Error in handleQuestion:', error);
+    throw error;
   }
 } 

src/popup.html

@@ -53,6 +53,7 @@ <h3>Custom Prompts</h3>
       <p>Total Tokens: <span id="total-tokens">0</span></p>
       <p>Total Requests: <span id="request-count">0</span></p>
       <p>Average Tokens/Request: <span id="avg-tokens">0</span></p>
+      <button id="clear-usage" class="ai-button-base ai-danger-button">Clear Usage Stats</button>
     </div>
   </div>
   <script src="popup.js"></script>

src/popup.ts

@@ -1,7 +1,7 @@
 /// <reference types="chrome"/>
 
 import { MODELS, MODEL_DISPLAY_NAMES, DEFAULT_MODEL, ModelType } from './config';
-import { Settings as GlobalSettings, getSettings, updateSettings } from './settings';
+import { Settings as GlobalSettings, getSettings, updateSettings, clearTokenUsage } from './settings';
 
 interface TokenUsage {
   totalTokens: number;
@@ -12,6 +12,15 @@ interface CustomPrompts {
   [key: string]: string;
 }
 
+// Load token usage from storage
+async function loadTokenUsage(): Promise<TokenUsage> {
+  return new Promise((resolve) => {
+    chrome.storage.local.get(['tokenUsage'], (result) => {
+      resolve(result.tokenUsage || { totalTokens: 0, requestCount: 0 });
+    });
+  });
+}
+
 function initializeSettings(): void {
   // Initialize model selector
   const modelSelector = document.getElementById('model-selector') as HTMLSelectElement;
@@ -55,6 +64,9 @@ function initializeSettings(): void {
       showIconToggle.checked = settings.showIcon;
     }
   });
+
+  // Load token usage
+  loadTokenUsage().then(updateUsageStats);
 }
 
 function loadCustomPrompts(prompts: CustomPrompts): void {
@@ -169,6 +181,15 @@ function saveSettings(): void {
 document.addEventListener('DOMContentLoaded', () => {
   initializeSettings();
 
+  // Add event listener for clear usage button
+  const clearUsageButton = document.getElementById('clear-usage');
+  if (clearUsageButton) {
+    clearUsageButton.addEventListener('click', async () => {
+      await clearTokenUsage();
+      updateUsageStats({ totalTokens: 0, requestCount: 0 });
+    });
+  }
+
   // Load custom prompts
   chrome.storage.sync.get(['customPrompts'], (result: { customPrompts?: Record<string, string> }) => {
     const customPrompts = result.customPrompts || {};

src/settings.css

@@ -188,7 +188,9 @@ label {
 /* Save button */
 .save-button {
   width: 100%;
+  height: 32px;
   margin-top: 12px;
+  border-radius: 6px;
 }
 
 /* Token usage section */

src/settings.ts

@@ -74,6 +74,12 @@ export async function clearApiKey(): Promise<void> {
   return updateSettings(settings);
 }
 
+export async function clearTokenUsage(): Promise<void> {
+  return new Promise((resolve) => {
+    chrome.storage.local.set({ tokenUsage: { totalTokens: 0, requestCount: 0 } }, resolve);
+  });
+}
+
 function updateTokenDisplay(usage: TokenUsage): void {
   const totalTokensEl = document.getElementById('total-tokens');
   const requestCountEl = document.getElementById('request-count');