@lobehub/chat 1.84.26 → 1.85.0

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

@@ -11,6 +11,7 @@ import {
   StreamToolCallChunkData,
   createCallbacksTransformer,
   createSSEProtocolTransformer,
+  createTokenSpeedCalculator,
   generateToolCallId,
 } from './protocol';
 
@@ -19,31 +20,62 @@ const transformGoogleGenerativeAIStream = (
   context: StreamContext,
 ): StreamProtocolChunk | StreamProtocolChunk[] => {
   // maybe need another structure to add support for multiple choices
+  const candidate = chunk.candidates?.[0];
+  const usage = chunk.usageMetadata;
+  const usageChunks: StreamProtocolChunk[] = [];
+  if (candidate?.finishReason && usage) {
+    const outputReasoningTokens = (usage as any).thoughtsTokenCount || undefined;
+    const totalOutputTokens = (usage.candidatesTokenCount ?? 0) + (outputReasoningTokens ?? 0);
+
+    usageChunks.push(
+      { data: candidate.finishReason, id: context?.id, type: 'stop' },
+      {
+        data: {
+          // TODO: Google SDK 0.24.0 don't have promptTokensDetails types
+          inputImageTokens: (usage as any).promptTokensDetails?.find(
+            (i: any) => i.modality === 'IMAGE',
+          )?.tokenCount,
+          inputTextTokens: (usage as any).promptTokensDetails?.find(
+            (i: any) => i.modality === 'TEXT',
+          )?.tokenCount,
+          outputReasoningTokens,
+          outputTextTokens: totalOutputTokens - (outputReasoningTokens ?? 0),
+          totalInputTokens: usage.promptTokenCount,
+          totalOutputTokens,
+          totalTokens: usage.totalTokenCount,
+        } as ModelTokensUsage,
+        id: context?.id,
+        type: 'usage',
+      },
+    );
+  }
+
   const functionCalls = chunk.functionCalls?.();
 
   if (functionCalls) {
-    return {
-      data: functionCalls.map(
-        (value, index): StreamToolCallChunkData => ({
-          function: {
-            arguments: JSON.stringify(value.args),
-            name: value.name,
-          },
-          id: generateToolCallId(index, value.name),
-          index: index,
-          type: 'function',
-        }),
-      ),
-      id: context.id,
-      type: 'tool_calls',
-    };
+    return [
+      {
+        data: functionCalls.map(
+          (value, index): StreamToolCallChunkData => ({
+            function: {
+              arguments: JSON.stringify(value.args),
+              name: value.name,
+            },
+            id: generateToolCallId(index, value.name),
+            index: index,
+            type: 'function',
+          }),
+        ),
+        id: context.id,
+        type: 'tool_calls',
+      },
+      ...usageChunks,
+    ];
   }
 
   const text = chunk.text?.();
 
-  if (chunk.candidates) {
-    const candidate = chunk.candidates[0];
-
+  if (candidate) {
     // return the grounding
     if (candidate.groundingMetadata) {
       const { webSearchQueries, groundingChunks } = candidate.groundingMetadata;
@@ -64,31 +96,15 @@ const transformGoogleGenerativeAIStream = (
           id: context.id,
           type: 'grounding',
         },
+        ...usageChunks,
       ];
     }
 
     if (candidate.finishReason) {
       if (chunk.usageMetadata) {
-        const usage = chunk.usageMetadata;
         return [
           !!text ? { data: text, id: context?.id, type: 'text' } : undefined,
-          { data: candidate.finishReason, id: context?.id, type: 'stop' },
-          {
-            data: {
-              // TODO: Google SDK 0.24.0 don't have promptTokensDetails types
-              inputImageTokens: (usage as any).promptTokensDetails?.find(
-                (i: any) => i.modality === 'IMAGE',
-              )?.tokenCount,
-              inputTextTokens: (usage as any).promptTokensDetails?.find(
-                (i: any) => i.modality === 'TEXT',
-              )?.tokenCount,
-              totalInputTokens: usage.promptTokenCount,
-              totalOutputTokens: usage.candidatesTokenCount,
-              totalTokens: usage.totalTokenCount,
-            } as ModelTokensUsage,
-            id: context?.id,
-            type: 'usage',
-          },
+          ...usageChunks,
         ].filter(Boolean) as StreamProtocolChunk[];
       }
       return { data: candidate.finishReason, id: context?.id, type: 'stop' };
@@ -117,13 +133,21 @@ const transformGoogleGenerativeAIStream = (
   };
 };
 
+export interface GoogleAIStreamOptions {
+  callbacks?: ChatStreamCallbacks;
+  inputStartAt?: number;
+}
+
 export const GoogleGenerativeAIStream = (
   rawStream: ReadableStream<EnhancedGenerateContentResponse>,
-  callbacks?: ChatStreamCallbacks,
+  { callbacks, inputStartAt }: GoogleAIStreamOptions = {},
 ) => {
   const streamStack: StreamContext = { id: 'chat_' + nanoid() };
 
   return rawStream
-    .pipeThrough(createSSEProtocolTransformer(transformGoogleGenerativeAIStream, streamStack))
+    .pipeThrough(
+      createTokenSpeedCalculator(transformGoogleGenerativeAIStream, { inputStartAt, streamStack }),
+    )
+    .pipeThrough(createSSEProtocolTransformer((c) => c, streamStack))
     .pipeThrough(createCallbacksTransformer(callbacks));
 };

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

@@ -298,17 +298,37 @@ export const TOKEN_SPEED_CHUNK_ID = 'output_speed';
  */
 export const createTokenSpeedCalculator = (
   transformer: (chunk: any, stack: StreamContext) => StreamProtocolChunk | StreamProtocolChunk[],
-  { streamStack, inputStartAt }: { inputStartAt?: number; streamStack?: StreamContext } = {},
+  { inputStartAt, streamStack }: { inputStartAt?: number; streamStack?: StreamContext } = {},
 ) => {
   let outputStartAt: number | undefined;
+  let outputThinking: boolean | undefined;
 
   const process = (chunk: StreamProtocolChunk) => {
     let result = [chunk];
-    // if the chunk is the first text chunk, set as output start
-    if (!outputStartAt && chunk.type === 'text') outputStartAt = Date.now();
+    // if the chunk is the first text or reasoning chunk, set as output start
+    if (!outputStartAt && (chunk.type === 'text' || chunk.type === 'reasoning')) {
+      outputStartAt = Date.now();
+    }
+
+    /**
+     * 部分 provider 在正式输出 reasoning 前，可能会先输出 content 为空字符串的 chunk，
+     * 其中 reasoning 可能为 null，会导致判断是否输出思考内容错误，所以过滤掉 null 或者空字符串。
+     * 也可能是某些特殊 token，所以不修改 outputStartAt 的逻辑。
+     */
+    if (
+      outputThinking === undefined &&
+      (chunk.type === 'text' || chunk.type === 'reasoning') &&
+      typeof chunk.data === 'string' &&
+      chunk.data.length > 0
+    ) {
+      outputThinking = chunk.type === 'reasoning';
+    }
     // if the chunk is the stop chunk, set as output finish
     if (inputStartAt && outputStartAt && chunk.type === 'usage') {
-      const outputTokens = chunk.data?.totalOutputTokens || chunk.data?.outputTextTokens;
+      const totalOutputTokens = chunk.data?.totalOutputTokens || chunk.data?.outputTextTokens;
+      const reasoningTokens = chunk.data?.outputReasoningTokens || 0;
+      const outputTokens =
+        (outputThinking ?? false) ? totalOutputTokens : totalOutputTokens - reasoningTokens;
       result.push({
         data: {
           tps: (outputTokens / (Date.now() - outputStartAt)) * 1000,

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

@@ -103,10 +103,12 @@ describe('VertexAIStream', () => {
     const onCompletionMock = vi.fn();
 
     const protocolStream = VertexAIStream(mockGoogleStream, {
-      onStart: onStartMock,
-      onText: onTextMock,
-      onToolsCalling: onToolCallMock,
-      onCompletion: onCompletionMock,
+      callbacks: {
+        onStart: onStartMock,
+        onText: onTextMock,
+        onToolsCalling: onToolCallMock,
+        onCompletion: onCompletionMock,
+      },
     });
 
     const decoder = new TextDecoder();
@@ -136,6 +138,7 @@ describe('VertexAIStream', () => {
 
   it('tool_calls', async () => {
     vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1');
+
     const rawChunks = [
       {
         candidates: [
@@ -204,10 +207,12 @@ describe('VertexAIStream', () => {
     const onCompletionMock = vi.fn();
 
     const protocolStream = VertexAIStream(mockGoogleStream, {
-      onStart: onStartMock,
-      onText: onTextMock,
-      onToolsCalling: onToolCallMock,
-      onCompletion: onCompletionMock,
+      callbacks: {
+        onStart: onStartMock,
+        onText: onTextMock,
+        onToolsCalling: onToolCallMock,
+        onCompletion: onCompletionMock,
+      },
     });
 
     const decoder = new TextDecoder();
@@ -223,10 +228,106 @@ describe('VertexAIStream', () => {
       'id: chat_1\n',
       'event: tool_calls\n',
       `data: [{"function":{"arguments":"{\\"city\\":\\"杭州\\"}","name":"realtime-weather____fetchCurrentWeather"},"id":"realtime-weather____fetchCurrentWeather_0","index":0,"type":"function"}]\n\n`,
+      'id: chat_1\n',
+      'event: stop\n',
+      'data: "STOP"\n\n',
+      'id: chat_1\n',
+      'event: usage\n',
+      'data: {"outputTextTokens":9,"totalInputTokens":95,"totalOutputTokens":9,"totalTokens":104}\n\n',
     ]);
 
     expect(onStartMock).toHaveBeenCalledTimes(1);
     expect(onToolCallMock).toHaveBeenCalledTimes(1);
     expect(onCompletionMock).toHaveBeenCalledTimes(1);
   });
+
+  it('should handle stop with content', async () => {
+    vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1');
+
+    const data = [
+      {
+        candidates: [
+          {
+            content: { parts: [{ text: '234' }], role: 'model' },
+            safetyRatings: [
+              { category: 'HARM_CATEGORY_HATE_SPEECH', probability: 'NEGLIGIBLE' },
+              { category: 'HARM_CATEGORY_DANGEROUS_CONTENT', probability: 'NEGLIGIBLE' },
+              { category: 'HARM_CATEGORY_HARASSMENT', probability: 'NEGLIGIBLE' },
+              { category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT', probability: 'NEGLIGIBLE' },
+            ],
+          },
+        ],
+        text: () => '234',
+        usageMetadata: {
+          promptTokenCount: 20,
+          totalTokenCount: 20,
+          promptTokensDetails: [{ modality: 'TEXT', tokenCount: 20 }],
+        },
+        modelVersion: 'gemini-2.0-flash-exp-image-generation',
+      },
+      {
+        text: () => '567890\n',
+        candidates: [
+          {
+            content: { parts: [{ text: '567890\n' }], role: 'model' },
+            finishReason: 'STOP',
+            safetyRatings: [
+              { category: 'HARM_CATEGORY_HATE_SPEECH', probability: 'NEGLIGIBLE' },
+              { category: 'HARM_CATEGORY_DANGEROUS_CONTENT', probability: 'NEGLIGIBLE' },
+              { category: 'HARM_CATEGORY_HARASSMENT', probability: 'NEGLIGIBLE' },
+              { category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT', probability: 'NEGLIGIBLE' },
+            ],
+          },
+        ],
+        usageMetadata: {
+          promptTokenCount: 19,
+          candidatesTokenCount: 11,
+          totalTokenCount: 30,
+          promptTokensDetails: [{ modality: 'TEXT', tokenCount: 19 }],
+          candidatesTokensDetails: [{ modality: 'TEXT', tokenCount: 11 }],
+        },
+        modelVersion: 'gemini-2.0-flash-exp-image-generation',
+      },
+    ];
+
+    const mockGoogleStream = new ReadableStream({
+      start(controller) {
+        data.forEach((item) => {
+          controller.enqueue(item);
+        });
+
+        controller.close();
+      },
+    });
+
+    const protocolStream = VertexAIStream(mockGoogleStream);
+
+    const decoder = new TextDecoder();
+    const chunks = [];
+
+    // @ts-ignore
+    for await (const chunk of protocolStream) {
+      chunks.push(decoder.decode(chunk, { stream: true }));
+    }
+
+    expect(chunks).toEqual(
+      [
+        'id: chat_1',
+        'event: text',
+        'data: "234"\n',
+
+        'id: chat_1',
+        'event: text',
+        `data: "567890\\n"\n`,
+        // stop
+        'id: chat_1',
+        'event: stop',
+        `data: "STOP"\n`,
+        // usage
+        'id: chat_1',
+        'event: usage',
+        `data: {"inputTextTokens":19,"outputTextTokens":11,"totalInputTokens":19,"totalOutputTokens":11,"totalTokens":30}\n`,
+      ].map((i) => i + '\n'),
+    );
+  });
 });

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

@@ -1,27 +1,58 @@
 import { EnhancedGenerateContentResponse, GenerateContentResponse } from '@google/generative-ai';
 
+import { ModelTokensUsage } from '@/types/message';
 import { nanoid } from '@/utils/uuid';
 
-import { ChatStreamCallbacks } from '../../types';
+import { type GoogleAIStreamOptions } from './google-ai';
 import {
   StreamContext,
   StreamProtocolChunk,
   createCallbacksTransformer,
   createSSEProtocolTransformer,
+  createTokenSpeedCalculator,
   generateToolCallId,
 } from './protocol';
 
 const transformVertexAIStream = (
   chunk: GenerateContentResponse,
-  stack: StreamContext,
-): StreamProtocolChunk => {
+  context: StreamContext,
+): StreamProtocolChunk | StreamProtocolChunk[] => {
   // maybe need another structure to add support for multiple choices
-  const candidates = chunk.candidates;
+  const candidate = chunk.candidates?.[0];
+  const usage = chunk.usageMetadata;
+  const usageChunks: StreamProtocolChunk[] = [];
+  if (candidate?.finishReason && usage) {
+    const outputReasoningTokens = (usage as any).thoughtsTokenCount || undefined;
+    const totalOutputTokens = (usage.candidatesTokenCount ?? 0) + (outputReasoningTokens ?? 0);
+
+    usageChunks.push(
+      { data: candidate.finishReason, id: context?.id, type: 'stop' },
+      {
+        data: {
+          // TODO: Google SDK 0.24.0 don't have promptTokensDetails types
+          inputImageTokens: (usage as any).promptTokensDetails?.find(
+            (i: any) => i.modality === 'IMAGE',
+          )?.tokenCount,
+          inputTextTokens: (usage as any).promptTokensDetails?.find(
+            (i: any) => i.modality === 'TEXT',
+          )?.tokenCount,
+          outputReasoningTokens,
+          outputTextTokens: totalOutputTokens - (outputReasoningTokens ?? 0),
+          totalInputTokens: usage.promptTokenCount,
+          totalOutputTokens,
+          totalTokens: usage.totalTokenCount,
+        } as ModelTokensUsage,
+        id: context?.id,
+        type: 'usage',
+      },
+    );
+  }
 
+  const candidates = chunk.candidates;
   if (!candidates)
     return {
       data: '',
-      id: stack?.id,
+      id: context?.id,
       type: 'text',
     };
 
@@ -32,44 +63,58 @@ const transformVertexAIStream = (
     if (part.functionCall) {
       const functionCall = part.functionCall;
 
-      return {
-        data: [
-          {
-            function: {
-              arguments: JSON.stringify(functionCall.args),
-              name: functionCall.name,
+      return [
+        {
+          data: [
+            {
+              function: {
+                arguments: JSON.stringify(functionCall.args),
+                name: functionCall.name,
+              },
+              id: generateToolCallId(0, functionCall.name),
+              index: 0,
+              type: 'function',
             },
-            id: generateToolCallId(0, functionCall.name),
-            index: 0,
-            type: 'function',
-          },
-        ],
-        id: stack?.id,
-        type: 'tool_calls',
-      };
+          ],
+          id: context?.id,
+          type: 'tool_calls',
+        },
+        ...usageChunks,
+      ];
+    }
+
+    if (item.finishReason) {
+      if (chunk.usageMetadata) {
+        return [
+          !!part.text ? { data: part.text, id: context?.id, type: 'text' } : undefined,
+          ...usageChunks,
+        ].filter(Boolean) as StreamProtocolChunk[];
+      }
+      return { data: item.finishReason, id: context?.id, type: 'stop' };
     }
 
     return {
       data: part.text,
-      id: stack?.id,
+      id: context?.id,
       type: 'text',
     };
   }
 
   return {
     data: '',
-    id: stack?.id,
+    id: context?.id,
     type: 'stop',
   };
 };
 
 export const VertexAIStream = (
   rawStream: ReadableStream<EnhancedGenerateContentResponse>,
-  callbacks?: ChatStreamCallbacks,
+  { callbacks, inputStartAt }: GoogleAIStreamOptions = {},
 ) => {
   const streamStack: StreamContext = { id: 'chat_' + nanoid() };
 
   return rawStream
-    .pipeThrough(createSSEProtocolTransformer(transformVertexAIStream, streamStack))
+    .pipeThrough(createTokenSpeedCalculator(transformVertexAIStream, { inputStartAt, streamStack }))
+    .pipeThrough(createSSEProtocolTransformer((c) => c, streamStack))
     .pipeThrough(createCallbacksTransformer(callbacks));
 };

package/src/libs/trpc/lambda/context.ts

@@ -57,6 +57,13 @@ export type LambdaContext = Awaited<ReturnType<typeof createContextInner>>;
  * @link https://trpc.io/docs/v11/context
  */
 export const createLambdaContext = async (request: NextRequest): Promise<LambdaContext> => {
+  // we have a special header to debug the api endpoint in development mode
+  // IT WON'T GO INTO PRODUCTION ANYMORE
+  const isDebugApi = request.headers.get('lobe-auth-dev-backend-api') === '1';
+  if (process.env.NODE_ENV === 'development' && isDebugApi) {
+    return { userId: process.env.MOCK_DEV_USER_ID };
+  }
+
   log('createLambdaContext called for request');
   // for API-response caching see https://trpc.io/docs/v11/caching
 

package/src/prompts/files/file.ts

@@ -1,9 +1,11 @@
 import { ChatFileItem } from '@/types/message';
 
-const filePrompt = (item: ChatFileItem, addUrl: boolean) =>
-  addUrl
-    ? `<file id="${item.id}" name="${item.name}" type="${item.fileType}" size="${item.size}" url="${item.url}"></file>`
-    : `<file id="${item.id}" name="${item.name}" type="${item.fileType}" size="${item.size}"></file>`;
+const filePrompt = (item: ChatFileItem, addUrl: boolean) => {
+  const content = item.content || '';
+  return addUrl
+    ? `<file id="${item.id}" name="${item.name}" type="${item.fileType}" size="${item.size}" url="${item.url}">${content}</file>`
+    : `<file id="${item.id}" name="${item.name}" type="${item.fileType}" size="${item.size}">${content}</file>`;
+};
 
 export const filePrompts = (fileList: ChatFileItem[], addUrl: boolean) => {
   if (fileList.length === 0) return '';