fix usage

Author: zerob13

src/main/presenter/llmProviderPresenter/providers/anthropicProvider.ts

@@ -14,6 +14,7 @@ import { presenter } from '@/presenter'
 import { HttpsProxyAgent } from 'https-proxy-agent'
 import { proxyConfig } from '../../proxyConfig'
 import { getModelConfig } from '../modelConfigs'
+import { Usage } from '@anthropic-ai/sdk/resources'
 
 export class AnthropicProvider extends BaseLLMProvider {
   private anthropic!: Anthropic
@@ -710,19 +711,12 @@ ${context}
       let currentToolId = ''
       let currentToolName = ''
       let currentToolInputs: Record<string, unknown> = {}
-
+      let usageMetadata: Usage | undefined
       // 处理流中的各种事件
       for await (const chunk of stream) {
         // 处理使用统计
         if (chunk.type === 'message_start' && chunk.message.usage) {
-          yield {
-            type: 'usage',
-            usage: {
-              prompt_tokens: chunk.message.usage.input_tokens,
-              completion_tokens: chunk.message.usage.output_tokens,
-              total_tokens: chunk.message.usage.input_tokens + chunk.message.usage.output_tokens
-            }
-          }
+          usageMetadata = chunk.message.usage
         }
 
         // 处理工具调用开始
@@ -904,7 +898,16 @@ ${context}
           continue
         }
       }
-
+      if (usageMetadata) {
+        yield {
+          type: 'usage',
+          usage: {
+            prompt_tokens: usageMetadata.input_tokens,
+            completion_tokens: usageMetadata.output_tokens,
+            total_tokens: usageMetadata.input_tokens + usageMetadata.output_tokens
+          }
+        }
+      }
       // 发送停止事件
       yield {
         type: 'stop',

src/main/presenter/llmProviderPresenter/providers/geminiProvider.ts

@@ -13,7 +13,8 @@ import {
   GenerativeModel,
   Part,
   Content,
-  GenerationConfig
+  GenerationConfig,
+  UsageMetadata
 } from '@google/generative-ai'
 import { ConfigPresenter } from '../../configPresenter'
 import { presenter } from '@/presenter'
@@ -808,19 +809,12 @@ export class GeminiProvider extends BaseLLMProvider {
     let buffer = ''
     let isInThinkTag = false
     let toolUseDetected = false
-
+    let usageMetadata: UsageMetadata | undefined
     // 流处理循环
     for await (const chunk of result.stream) {
       // 处理用量统计
       if (chunk.usageMetadata) {
-        yield {
-          type: 'usage',
-          usage: {
-            prompt_tokens: chunk.usageMetadata.promptTokenCount,
-            completion_tokens: chunk.usageMetadata.candidatesTokenCount,
-            total_tokens: chunk.usageMetadata.totalTokenCount
-          }
-        }
+        usageMetadata = chunk.usageMetadata
       }
 
       // 检查是否包含函数调用
@@ -948,7 +942,16 @@ export class GeminiProvider extends BaseLLMProvider {
       // 内容已经发送，清空buffer避免重复
       buffer = ''
     }
-
+    if (usageMetadata) {
+      yield {
+        type: 'usage',
+        usage: {
+          prompt_tokens: usageMetadata.promptTokenCount,
+          completion_tokens: usageMetadata.candidatesTokenCount,
+          total_tokens: usageMetadata.totalTokenCount
+        }
+      }
+    }
     // 处理剩余缓冲区内容
     if (buffer) {
       if (isInThinkTag) {

src/main/presenter/llmProviderPresenter/providers/ollamaProvider.ts

@@ -495,7 +495,6 @@ export class OllamaProvider extends BaseLLMProvider {
             completion_tokens: chunk.eval_count || 0,
             total_tokens: (chunk.prompt_eval_count || 0) + (chunk.eval_count || 0)
           }
-          yield { type: 'usage', usage }
         }
 
         // 处理原生工具调用
@@ -643,6 +642,9 @@ export class OllamaProvider extends BaseLLMProvider {
 
             continue
           }
+          if (usage) {
+            yield { type: 'usage', usage }
+          }
 
           // --- 思考标签处理 ---
           if (thinkState === 'inside') {