@lobehub/chat 1.79.4 → 1.79.6

package/src/libs/agent-runtime/utils/streams/openai.test.ts

@@ -44,14 +44,12 @@ describe('OpenAIStream', () => {
 
     const onStartMock = vi.fn();
     const onTextMock = vi.fn();
-    const onTokenMock = vi.fn();
     const onCompletionMock = vi.fn();
 
     const protocolStream = OpenAIStream(mockOpenAIStream, {
       callbacks: {
         onStart: onStartMock,
         onText: onTextMock,
-        onToken: onTokenMock,
         onCompletion: onCompletionMock,
       },
     });
@@ -77,9 +75,8 @@ describe('OpenAIStream', () => {
     ]);
 
     expect(onStartMock).toHaveBeenCalledTimes(1);
-    expect(onTextMock).toHaveBeenNthCalledWith(1, '"Hello"');
-    expect(onTextMock).toHaveBeenNthCalledWith(2, '" world!"');
-    expect(onTokenMock).toHaveBeenCalledTimes(2);
+    expect(onTextMock).toHaveBeenNthCalledWith(1, 'Hello');
+    expect(onTextMock).toHaveBeenNthCalledWith(2, ' world!');
     expect(onCompletionMock).toHaveBeenCalledTimes(1);
   });
 
@@ -195,7 +192,7 @@ describe('OpenAIStream', () => {
 
     const protocolStream = OpenAIStream(mockOpenAIStream, {
       callbacks: {
-        onToolCall: onToolCallMock,
+        onToolsCalling: onToolCallMock,
       },
     });
 
@@ -578,7 +575,7 @@ describe('OpenAIStream', () => {
 
       const protocolStream = OpenAIStream(mockOpenAIStream, {
         callbacks: {
-          onToolCall: onToolCallMock,
+          onToolsCalling: onToolCallMock,
         },
       });
 
@@ -711,7 +708,7 @@ describe('OpenAIStream', () => {
 
       const protocolStream = OpenAIStream(mockOpenAIStream, {
         callbacks: {
-          onToolCall: onToolCallMock,
+          onToolsCalling: onToolCallMock,
         },
       });
 

package/src/libs/agent-runtime/utils/streams/protocol.ts

@@ -1,7 +1,8 @@
-import { ChatStreamCallbacks } from '@/libs/agent-runtime';
 import { ModelTokensUsage } from '@/types/message';
 
 import { AgentRuntimeErrorType } from '../../error';
+import { parseToolCalls } from '../../helpers';
+import { ChatStreamCallbacks } from '../../types';
 
 /**
  * context in the stream to save temporarily data
@@ -140,18 +141,31 @@ export const createSSEProtocolTransformer = (
 
 export function createCallbacksTransformer(cb: ChatStreamCallbacks | undefined) {
   const textEncoder = new TextEncoder();
-  let aggregatedResponse = '';
-  let currentType = '';
+  let aggregatedText = '';
+  let aggregatedThinking: string | undefined = undefined;
+  let usage: ModelTokensUsage | undefined;
+  let grounding: any;
+  let toolsCalling: any;
+
+  let currentType = '' as unknown as StreamProtocolChunk['type'];
   const callbacks = cb || {};
 
   return new TransformStream({
     async flush(): Promise<void> {
+      const data = {
+        grounding,
+        text: aggregatedText,
+        thinking: aggregatedThinking,
+        toolsCalling,
+        usage,
+      };
+
       if (callbacks.onCompletion) {
-        await callbacks.onCompletion(aggregatedResponse);
+        await callbacks.onCompletion(data);
       }
 
       if (callbacks.onFinal) {
-        await callbacks.onFinal(aggregatedResponse);
+        await callbacks.onFinal(data);
       }
     },
 
@@ -164,22 +178,53 @@ export function createCallbacksTransformer(cb: ChatStreamCallbacks | undefined)
 
       // track the type of the chunk
       if (chunk.startsWith('event:')) {
-        currentType = chunk.split('event:')[1].trim();
+        currentType = chunk.split('event:')[1].trim() as unknown as StreamProtocolChunk['type'];
       }
       // if the message is a data chunk, handle the callback
       else if (chunk.startsWith('data:')) {
         const content = chunk.split('data:')[1].trim();
 
+        let data: any = undefined;
+        try {
+          data = JSON.parse(content);
+        } catch {}
+
+        if (!data) return;
+
         switch (currentType) {
           case 'text': {
-            await callbacks.onText?.(content);
-            await callbacks.onToken?.(JSON.parse(content));
+            aggregatedText += data;
+            await callbacks.onText?.(data);
+            break;
+          }
+
+          case 'reasoning': {
+            if (!aggregatedThinking) {
+              aggregatedThinking = '';
+            }
+
+            aggregatedThinking += data;
+            await callbacks.onThinking?.(data);
+            break;
+          }
+
+          case 'usage': {
+            usage = data;
+            await callbacks.onUsage?.(data);
+            break;
+          }
+
+          case 'grounding': {
+            grounding = data;
+            await callbacks.onGrounding?.(data);
             break;
           }
 
           case 'tool_calls': {
-            // TODO: make on ToolCall callback
-            await callbacks.onToolCall?.();
+            if (!toolsCalling) toolsCalling = [];
+            toolsCalling = parseToolCalls(toolsCalling, data);
+
+            await callbacks.onToolsCalling?.({ chunk: data, toolsCalling });
           }
         }
       }

package/src/libs/agent-runtime/utils/streams/qwen.test.ts

@@ -43,13 +43,11 @@ describe('QwenAIStream', () => {
 
     const onStartMock = vi.fn();
     const onTextMock = vi.fn();
-    const onTokenMock = vi.fn();
     const onCompletionMock = vi.fn();
 
     const protocolStream = QwenAIStream(mockOpenAIStream, {
       onStart: onStartMock,
       onText: onTextMock,
-      onToken: onTokenMock,
       onCompletion: onCompletionMock,
     });
 
@@ -74,9 +72,8 @@ describe('QwenAIStream', () => {
     ]);
 
     expect(onStartMock).toHaveBeenCalledTimes(1);
-    expect(onTextMock).toHaveBeenNthCalledWith(1, '"Hello"');
-    expect(onTextMock).toHaveBeenNthCalledWith(2, '" world!"');
-    expect(onTokenMock).toHaveBeenCalledTimes(2);
+    expect(onTextMock).toHaveBeenNthCalledWith(1, 'Hello');
+    expect(onTextMock).toHaveBeenNthCalledWith(2, ' world!');
     expect(onCompletionMock).toHaveBeenCalledTimes(1);
   });
 
@@ -114,7 +111,7 @@ describe('QwenAIStream', () => {
     const onToolCallMock = vi.fn();
 
     const protocolStream = QwenAIStream(mockOpenAIStream, {
-      onToolCall: onToolCallMock,
+      onToolsCalling: onToolCallMock,
     });
 
     const decoder = new TextDecoder();

package/src/libs/agent-runtime/utils/streams/spark.test.ts

@@ -1,41 +1,42 @@
+import type OpenAI from 'openai';
 import { beforeAll, describe, expect, it, vi } from 'vitest';
+
 import { SparkAIStream, transformSparkResponseToStream } from './spark';
-import type OpenAI from 'openai';
 
 describe('SparkAIStream', () => {
   beforeAll(() => {});
 
   it('should transform non-streaming response to stream', async () => {
     const mockResponse = {
-      id: "cha000ceba6@dx193d200b580b8f3532",
-      object: "chat.completion",
+      id: 'cha000ceba6@dx193d200b580b8f3532',
+      object: 'chat.completion',
       created: 1734395014,
-      model: "max-32k",
+      model: 'max-32k',
       choices: [
         {
           message: {
-            role: "assistant",
-            content: "",
+            role: 'assistant',
+            content: '',
             refusal: null,
             tool_calls: {
-              type: "function",
+              type: 'function',
               function: {
                 arguments: '{"city":"Shanghai"}',
-                name: "realtime-weather____fetchCurrentWeather"
+                name: 'realtime-weather____fetchCurrentWeather',
               },
-              id: "call_1"
-            }
+              id: 'call_1',
+            },
           },
           index: 0,
           logprobs: null,
-          finish_reason: "tool_calls"
-        }
+          finish_reason: 'tool_calls',
+        },
       ],
       usage: {
         prompt_tokens: 8,
         completion_tokens: 0,
-        total_tokens: 8
-      }
+        total_tokens: 8,
+      },
     } as unknown as OpenAI.ChatCompletion;
 
     const stream = transformSparkResponseToStream(mockResponse);
@@ -48,15 +49,17 @@ describe('SparkAIStream', () => {
     }
 
     expect(chunks).toHaveLength(2);
-    expect(chunks[0].choices[0].delta.tool_calls).toEqual([{
-      function: {
-        arguments: '{"city":"Shanghai"}',
-        name: "realtime-weather____fetchCurrentWeather"
+    expect(chunks[0].choices[0].delta.tool_calls).toEqual([
+      {
+        function: {
+          arguments: '{"city":"Shanghai"}',
+          name: 'realtime-weather____fetchCurrentWeather',
+        },
+        id: 'call_1',
+        index: 0,
+        type: 'function',
       },
-      id: "call_1",
-      index: 0,
-      type: "function"
-    }]);
+    ]);
     expect(chunks[1].choices[0].finish_reason).toBeDefined();
   });
 
@@ -64,36 +67,36 @@ describe('SparkAIStream', () => {
     const mockStream = new ReadableStream({
       start(controller) {
         controller.enqueue({
-          id: "cha000b0bf9@dx193d1ffa61cb894532",
-          object: "chat.completion.chunk",
+          id: 'cha000b0bf9@dx193d1ffa61cb894532',
+          object: 'chat.completion.chunk',
           created: 1734395014,
-          model: "max-32k",
+          model: 'max-32k',
           choices: [
             {
               delta: {
-                role: "assistant",
-                content: "",
+                role: 'assistant',
+                content: '',
                 tool_calls: {
-                  type: "function",
+                  type: 'function',
                   function: {
                     arguments: '{"city":"Shanghai"}',
-                    name: "realtime-weather____fetchCurrentWeather"
+                    name: 'realtime-weather____fetchCurrentWeather',
                   },
-                  id: "call_1"
-                }
+                  id: 'call_1',
+                },
               },
-              index: 0
-            }
-          ]
+              index: 0,
+            },
+          ],
         } as unknown as OpenAI.ChatCompletionChunk);
         controller.close();
-      }
+      },
     });
 
     const onToolCallMock = vi.fn();
 
     const protocolStream = SparkAIStream(mockStream, {
-      onToolCall: onToolCallMock
+      onToolsCalling: onToolCallMock,
     });
 
     const decoder = new TextDecoder();
@@ -107,7 +110,7 @@ describe('SparkAIStream', () => {
     expect(chunks).toEqual([
       'id: cha000b0bf9@dx193d1ffa61cb894532\n',
       'event: tool_calls\n',
-      `data: [{\"function\":{\"arguments\":\"{\\\"city\\\":\\\"Shanghai\\\"}\",\"name\":\"realtime-weather____fetchCurrentWeather\"},\"id\":\"call_1\",\"index\":0,\"type\":\"function\"}]\n\n`
+      `data: [{\"function\":{\"arguments\":\"{\\\"city\\\":\\\"Shanghai\\\"}\",\"name\":\"realtime-weather____fetchCurrentWeather\"},\"id\":\"call_1\",\"index\":0,\"type\":\"function\"}]\n\n`,
     ]);
 
     expect(onToolCallMock).toHaveBeenCalledTimes(1);
@@ -117,43 +120,43 @@ describe('SparkAIStream', () => {
     const mockStream = new ReadableStream({
       start(controller) {
         controller.enqueue({
-          id: "test-id",
-          object: "chat.completion.chunk",
+          id: 'test-id',
+          object: 'chat.completion.chunk',
           created: 1734395014,
-          model: "max-32k",
+          model: 'max-32k',
           choices: [
             {
               delta: {
-                content: "Hello",
-                role: "assistant"
+                content: 'Hello',
+                role: 'assistant',
               },
-              index: 0
-            }
-          ]
+              index: 0,
+            },
+          ],
         } as OpenAI.ChatCompletionChunk);
         controller.enqueue({
-          id: "test-id",
-          object: "chat.completion.chunk",
+          id: 'test-id',
+          object: 'chat.completion.chunk',
           created: 1734395014,
-          model: "max-32k",
+          model: 'max-32k',
           choices: [
             {
               delta: {
-                content: " World",
-                role: "assistant"
+                content: ' World',
+                role: 'assistant',
               },
-              index: 0
-            }
-          ]
+              index: 0,
+            },
+          ],
         } as OpenAI.ChatCompletionChunk);
         controller.close();
-      }
+      },
     });
 
     const onTextMock = vi.fn();
-    
+
     const protocolStream = SparkAIStream(mockStream, {
-      onText: onTextMock
+      onText: onTextMock,
     });
 
     const decoder = new TextDecoder();
@@ -170,18 +173,18 @@ describe('SparkAIStream', () => {
       'data: "Hello"\n\n',
       'id: test-id\n',
       'event: text\n',
-      'data: " World"\n\n'
+      'data: " World"\n\n',
     ]);
 
-    expect(onTextMock).toHaveBeenNthCalledWith(1, '"Hello"');
-    expect(onTextMock).toHaveBeenNthCalledWith(2, '" World"');
+    expect(onTextMock).toHaveBeenNthCalledWith(1, 'Hello');
+    expect(onTextMock).toHaveBeenNthCalledWith(2, ' World');
   });
 
   it('should handle empty stream', async () => {
     const mockStream = new ReadableStream({
       start(controller) {
         controller.close();
-      }
+      },
     });
 
     const protocolStream = SparkAIStream(mockStream);

package/src/libs/agent-runtime/utils/streams/vertex-ai.test.ts

@@ -99,15 +99,13 @@ describe('VertexAIStream', () => {
 
     const onStartMock = vi.fn();
     const onTextMock = vi.fn();
-    const onTokenMock = vi.fn();
     const onToolCallMock = vi.fn();
     const onCompletionMock = vi.fn();
 
     const protocolStream = VertexAIStream(mockGoogleStream, {
       onStart: onStartMock,
       onText: onTextMock,
-      onToken: onTokenMock,
-      onToolCall: onToolCallMock,
+      onToolsCalling: onToolCallMock,
       onCompletion: onCompletionMock,
     });
 
@@ -132,7 +130,7 @@ describe('VertexAIStream', () => {
     ]);
 
     expect(onStartMock).toHaveBeenCalledTimes(1);
-    expect(onTokenMock).toHaveBeenCalledTimes(2);
+    expect(onTextMock).toHaveBeenCalledTimes(2);
     expect(onCompletionMock).toHaveBeenCalledTimes(1);
   });
 
@@ -202,15 +200,13 @@ describe('VertexAIStream', () => {
 
     const onStartMock = vi.fn();
     const onTextMock = vi.fn();
-    const onTokenMock = vi.fn();
     const onToolCallMock = vi.fn();
     const onCompletionMock = vi.fn();
 
     const protocolStream = VertexAIStream(mockGoogleStream, {
       onStart: onStartMock,
       onText: onTextMock,
-      onToken: onTokenMock,
-      onToolCall: onToolCallMock,
+      onToolsCalling: onToolCallMock,
       onCompletion: onCompletionMock,
     });
 

package/src/libs/agent-runtime/xai/index.ts

@@ -12,6 +12,16 @@ export const LobeXAI = LobeOpenAICompatibleFactory({
   chatCompletion: {
     // xAI API does not support stream_options: { include_usage: true }
     excludeUsage: true,
+    handlePayload: (payload) => {
+      const { frequency_penalty, model, presence_penalty, ...rest } = payload;
+
+      return {
+        ...rest,
+        frequency_penalty: model.includes('grok-3-mini') ? undefined : frequency_penalty,
+        model,
+        presence_penalty: model.includes('grok-3-mini') ? undefined : presence_penalty,
+      } as any;
+    },
   },
   debug: {
     chatCompletion: () => process.env.DEBUG_XAI_CHAT_COMPLETION === '1',

package/src/libs/agent-runtime/zhipu/index.test.ts

@@ -63,7 +63,7 @@ describe('LobeZhipuAI', () => {
       // 准备 callback 和 headers
       const mockCallback: ChatStreamCallbacks = {
         onStart: vi.fn(),
-        onToken: vi.fn(),
+        onText: vi.fn(),
       };
       const mockHeaders = { 'Custom-Header': 'TestValue' };
 
@@ -80,7 +80,7 @@ describe('LobeZhipuAI', () => {
       // 验证 callback 被调用
       await result.text(); // 确保流被消费
       expect(mockCallback.onStart).toHaveBeenCalled();
-      expect(mockCallback.onToken).toHaveBeenCalledWith('hello');
+      expect(mockCallback.onText).toHaveBeenCalledWith('hello');
 
       // 验证 headers 被正确传递
       expect(result.headers.get('Custom-Header')).toEqual('TestValue');

package/src/server/modules/AgentRuntime/index.ts

@@ -1,22 +1,10 @@
 import { getLLMConfig } from '@/config/llm';
 import { JWTPayload } from '@/const/auth';
-import { INBOX_SESSION_ID } from '@/const/session';
-import {
-  LOBE_CHAT_OBSERVATION_ID,
-  LOBE_CHAT_TRACE_ID,
-  TracePayload,
-  TraceTagMap,
-} from '@/const/trace';
-import { AgentRuntime, ChatStreamPayload, ModelProvider } from '@/libs/agent-runtime';
-import { TraceClient } from '@/libs/traces';
+import { AgentRuntime, ModelProvider } from '@/libs/agent-runtime';
 
 import apiKeyManager from './apiKeyManager';
 
-export interface AgentChatOptions {
-  enableTrace?: boolean;
-  provider: string;
-  trace?: TracePayload;
-}
+export * from './trace';
 
 /**
  * Retrieves the options object from environment and apikeymanager
@@ -26,7 +14,7 @@ export interface AgentChatOptions {
  * @param payload - The JWT payload.
  * @returns The options object.
  */
-const getLlmOptionsFromPayload = (provider: string, payload: JWTPayload) => {
+const getParamsFromPayload = (provider: string, payload: JWTPayload) => {
   const llmConfig = getLLMConfig() as Record<string, any>;
 
   switch (provider) {
@@ -131,66 +119,7 @@ export const initAgentRuntimeWithUserPayload = (
   params: any = {},
 ) => {
   return AgentRuntime.initializeWithProvider(provider, {
-    ...getLlmOptionsFromPayload(provider, payload),
+    ...getParamsFromPayload(provider, payload),
     ...params,
   });
 };
-
-export const createTraceOptions = (
-  payload: ChatStreamPayload,
-  { trace: tracePayload, provider }: AgentChatOptions,
-) => {
-  const { messages, model, tools, ...parameters } = payload;
-  // create a trace to monitor the completion
-  const traceClient = new TraceClient();
-  const trace = traceClient.createTrace({
-    id: tracePayload?.traceId,
-    input: messages,
-    metadata: { provider },
-    name: tracePayload?.traceName,
-    sessionId: `${tracePayload?.sessionId || INBOX_SESSION_ID}@${tracePayload?.topicId || 'start'}`,
-    tags: tracePayload?.tags,
-    userId: tracePayload?.userId,
-  });
-
-  const generation = trace?.generation({
-    input: messages,
-    metadata: { provider },
-    model,
-    modelParameters: parameters as any,
-    name: `Chat Completion (${provider})`,
-    startTime: new Date(),
-  });
-
-  return {
-    callback: {
-      experimental_onToolCall: async () => {
-        trace?.update({
-          tags: [...(tracePayload?.tags || []), TraceTagMap.ToolsCall],
-        });
-      },
-
-      onCompletion: async (completion: string) => {
-        generation?.update({
-          endTime: new Date(),
-          metadata: { provider, tools },
-          output: completion,
-        });
-
-        trace?.update({ output: completion });
-      },
-
-      onFinal: async () => {
-        await traceClient.shutdownAsync();
-      },
-
-      onStart: () => {
-        generation?.update({ completionStartTime: new Date() });
-      },
-    },
-    headers: {
-      [LOBE_CHAT_OBSERVATION_ID]: generation?.id,
-      [LOBE_CHAT_TRACE_ID]: trace?.id,
-    },
-  };
-};