🐛 fix: update convertUsage to handle XAI provider and adjust OpenAIStream to pass provider (#8557)

sxjeru · web-flow · commit d1e4a54b0101 · 2025-07-25T16:16:54.000+08:00
diff --git a/src/libs/model-runtime/utils/streams/openai/openai.ts b/src/libs/model-runtime/utils/streams/openai/openai.ts
@@ -23,6 +23,7 @@ import {
 const transformOpenAIStream = (
   chunk: OpenAI.ChatCompletionChunk,
   streamContext: StreamContext,
+  provider?: string,
 ): StreamProtocolChunk | StreamProtocolChunk[] => {
   // handle the first chunk error
   if (FIRST_CHUNK_ERROR_KEY in chunk) {
@@ -45,7 +46,7 @@ const transformOpenAIStream = (
     if (!item) {
       if (chunk.usage) {
         const usage = chunk.usage;
-        return { data: convertUsage(usage), id: chunk.id, type: 'usage' };
+        return { data: convertUsage(usage, provider), id: chunk.id, type: 'usage' };
       }
 
       return { data: chunk, id: chunk.id, type: 'data' };
@@ -155,7 +156,7 @@ const transformOpenAIStream = (
 
       if (chunk.usage) {
         const usage = chunk.usage;
-        return { data: convertUsage(usage), id: chunk.id, type: 'usage' };
+        return { data: convertUsage(usage, provider), id: chunk.id, type: 'usage' };
       }
 
       // xAI Live Search 功能返回引用源
@@ -274,7 +275,7 @@ const transformOpenAIStream = (
     // litellm 的返回结果中，存在 delta 为空，但是有 usage 的情况
     if (chunk.usage) {
       const usage = chunk.usage;
-      return { data: convertUsage(usage), id: chunk.id, type: 'usage' };
+      return { data: convertUsage(usage, provider), id: chunk.id, type: 'usage' };
     }
 
     // 其余情况下，返回 delta 和 index
@@ -321,6 +322,9 @@ export const OpenAIStream = (
 ) => {
   const streamStack: StreamContext = { id: '' };
 
+  const transformWithProvider = (chunk: OpenAI.ChatCompletionChunk, streamContext: StreamContext) =>
+    transformOpenAIStream(chunk, streamContext, provider);
+
   const readableStream =
     stream instanceof ReadableStream ? stream : convertIterableToStream(stream);
 
@@ -330,7 +334,7 @@ export const OpenAIStream = (
       // provider like huggingface or minimax will return error in the stream,
       // so in the first Transformer, we need to handle the error
       .pipeThrough(createFirstErrorHandleTransformer(bizErrorTypeTransformer, provider))
-      .pipeThrough(createTokenSpeedCalculator(transformOpenAIStream, { inputStartAt, streamStack }))
+      .pipeThrough(createTokenSpeedCalculator(transformWithProvider, { inputStartAt, streamStack }))
       .pipeThrough(createSSEProtocolTransformer((c) => c, streamStack))
       .pipeThrough(createCallbacksTransformer(callbacks))
   );
diff --git a/src/libs/model-runtime/utils/usageConverter.test.ts b/src/libs/model-runtime/utils/usageConverter.test.ts
@@ -1,7 +1,7 @@
 import OpenAI from 'openai';
 import { describe, expect, it } from 'vitest';
 
-import { convertUsage } from './usageConverter';
+import { convertUsage, convertResponseUsage } from './usageConverter';
 
 describe('convertUsage', () => {
   it('should convert basic OpenAI usage data correctly', () => {
@@ -246,4 +246,48 @@ describe('convertUsage', () => {
     expect(result).not.toHaveProperty('outputReasoningTokens');
     expect(result).not.toHaveProperty('outputAudioTokens');
   });
+
+  it('should handle XAI provider correctly where completion_tokens does not include reasoning_tokens', () => {
+    // Arrange
+    const xaiUsage: OpenAI.Completions.CompletionUsage = {
+      prompt_tokens: 6103,
+      completion_tokens: 66, // 这个不包含 reasoning_tokens
+      total_tokens: 6550,
+      prompt_tokens_details: {
+        audio_tokens: 0,
+        cached_tokens: 0,
+      },
+      completion_tokens_details: {
+        accepted_prediction_tokens: 0,
+        audio_tokens: 0,
+        reasoning_tokens: 381, // 这是额外的 reasoning tokens
+        rejected_prediction_tokens: 0,
+      },
+    };
+
+    // Act
+    const xaiResult = convertUsage(xaiUsage, 'xai');
+
+    // Assert
+    expect(xaiResult).toMatchObject({
+      totalInputTokens: 6103,
+      totalOutputTokens: 66,
+      outputTextTokens: 66, // 不减去 reasoning_tokens
+      outputReasoningTokens: 381,
+      totalTokens: 6550,
+    });
+
+    // 测试其他 provider（默认行为）
+    const defaultResult = convertUsage(xaiUsage);
+
+    // 默认行为: outputTextTokens 应该是 completion_tokens - reasoning_tokens - audio_tokens = 66 - 381 - 0 = -315
+    expect(defaultResult.outputTextTokens).toBe(-315);
+    expect(defaultResult).toMatchObject({
+      totalInputTokens: 6103,
+      totalOutputTokens: 66,
+      outputTextTokens: -315, // 负数确实会出现在结果中
+      outputReasoningTokens: 381,
+      totalTokens: 6550,
+    });
+  });
 });
diff --git a/src/libs/model-runtime/utils/usageConverter.ts b/src/libs/model-runtime/utils/usageConverter.ts
@@ -2,7 +2,7 @@ import OpenAI from 'openai';
 
 import { ModelTokensUsage } from '@/types/message';
 
-export const convertUsage = (usage: OpenAI.Completions.CompletionUsage): ModelTokensUsage => {
+export const convertUsage = (usage: OpenAI.Completions.CompletionUsage, provider?: string): ModelTokensUsage => {
   // 目前只有 pplx 才有 citation_tokens
   const inputTextTokens = usage.prompt_tokens || 0;
   const inputCitationTokens = (usage as any).citation_tokens || 0;
@@ -17,7 +17,11 @@ export const convertUsage = (usage: OpenAI.Completions.CompletionUsage): ModelTo
   const totalOutputTokens = usage.completion_tokens;
   const outputReasoning = usage.completion_tokens_details?.reasoning_tokens || 0;
   const outputAudioTokens = usage.completion_tokens_details?.audio_tokens || 0;
-  const outputTextTokens = totalOutputTokens - outputReasoning - outputAudioTokens;
+  
+  // XAI 的 completion_tokens 不包含 reasoning_tokens，需要特殊处理
+  const outputTextTokens = provider === 'xai' 
+    ? totalOutputTokens - outputAudioTokens
+    : totalOutputTokens - outputReasoning - outputAudioTokens;
 
   const totalTokens = inputCitationTokens + usage.total_tokens;
 
