@lobehub/chat 1.70.11 → 1.71.1

package/src/database/server/models/message.ts

@@ -21,6 +21,7 @@ import {
   CreateMessageParams,
   MessageItem,
   ModelRankItem,
+  UpdateMessageParams,
 } from '@/types/message';
 import { merge } from '@/utils/merge';
 import { today } from '@/utils/time';
@@ -497,15 +498,25 @@ export class MessageModel {
   };
   // **************** Update *************** //
 
-  update = async (id: string, message: Partial<MessageItem>) => {
-    return this.db
-      .update(messages)
-      .set({
-        ...message,
-        // TODO: need a better way to handle this
-        role: message.role as any,
-      })
-      .where(and(eq(messages.id, id), eq(messages.userId, this.userId)));
+  update = async (id: string, { imageList, ...message }: Partial<UpdateMessageParams>) => {
+    return this.db.transaction(async (trx) => {
+      // 1. insert message files
+      if (imageList && imageList.length > 0) {
+        await trx
+          .insert(messagesFiles)
+          .values(imageList.map((file) => ({ fileId: file.id, messageId: id })));
+      }
+
+      return trx
+        .update(messages)
+        .set({
+          ...message,
+          // TODO: need a better way to handle this
+          // TODO: but I forget why 🤡
+          role: message.role as any,
+        })
+        .where(and(eq(messages.id, id), eq(messages.userId, this.userId)));
+    });
   };
 
   updatePluginState = async (id: string, state: Record<string, any>) => {

package/src/database/server/models/user.test.ts

@@ -0,0 +1,58 @@
+// @vitest-environment node
+import { TRPCError } from '@trpc/server';
+import { describe, expect, it, vi } from 'vitest';
+
+import { UserModel, UserNotFoundError } from '@/database/server/models/user';
+
+describe('UserNotFoundError', () => {
+  it('should extend TRPCError with correct code and message', () => {
+    const error = new UserNotFoundError();
+
+    expect(error).toBeInstanceOf(TRPCError);
+    expect(error.code).toBe('UNAUTHORIZED');
+    expect(error.message).toBe('user not found');
+  });
+});
+
+describe('UserModel', () => {
+  const mockDb = {
+    query: {
+      users: {
+        findFirst: vi.fn(),
+      },
+    },
+  };
+
+  const mockUserId = 'test-user-id';
+  const userModel = new UserModel(mockDb as any, mockUserId);
+
+  describe('getUserRegistrationDuration', () => {
+    it('should return default values when user not found', async () => {
+      mockDb.query.users.findFirst.mockResolvedValue(null);
+
+      const result = await userModel.getUserRegistrationDuration();
+
+      expect(result).toEqual({
+        createdAt: expect.any(String),
+        duration: 1,
+        updatedAt: expect.any(String),
+      });
+    });
+
+    it('should calculate duration correctly for existing user', async () => {
+      const createdAt = new Date('2024-01-01');
+      mockDb.query.users.findFirst.mockResolvedValue({
+        createdAt,
+      });
+
+      const result = await userModel.getUserRegistrationDuration();
+
+      expect(result).toEqual({
+        createdAt: '2024-01-01',
+        duration: expect.any(Number),
+        updatedAt: expect.any(String),
+      });
+      expect(result.duration).toBeGreaterThan(0);
+    });
+  });
+});

package/src/features/AlertBanner/CloudBanner.tsx

@@ -61,7 +61,7 @@ const CloudBanner = memo<{ mobile?: boolean }>(({ mobile }) => {
       <b>{t('alert.cloud.title', { name: LOBE_CHAT_CLOUD })}:</b>
       <span>
         {t(mobile ? 'alert.cloud.descOnMobile' : 'alert.cloud.desc', {
-          credit: new Intl.NumberFormat('en-US').format(450_000),
+          credit: new Intl.NumberFormat('en-US').format(500_000),
           name: LOBE_CHAT_CLOUD,
         })}
       </span>

package/src/features/Conversation/Messages/Assistant/index.tsx

@@ -2,6 +2,7 @@ import { ReactNode, memo } from 'react';
 import { Flexbox } from 'react-layout-kit';
 
 import { LOADING_FLAT } from '@/const/message';
+import ImageFileListViewer from '@/features/Conversation/Messages/User/ImageFileListViewer';
 import { useChatStore } from '@/store/chat';
 import { aiChatSelectors, chatSelectors } from '@/store/chat/selectors';
 import { ChatMessage } from '@/types/message';
@@ -17,7 +18,7 @@ export const AssistantMessage = memo<
   ChatMessage & {
     editableContent: ReactNode;
   }
->(({ id, tools, content, chunksList, search, ...props }) => {
+>(({ id, tools, content, chunksList, search, imageList, ...props }) => {
   const editing = useChatStore(chatSelectors.isMessageEditing(id));
   const generating = useChatStore(chatSelectors.isMessageGenerating(id));
 
@@ -28,6 +29,7 @@ export const AssistantMessage = memo<
   const isIntentUnderstanding = useChatStore(aiChatSelectors.isIntentUnderstanding(id));
 
   const showSearch = !!search && !!search.citations?.length;
+  const showImageItems = !!imageList && imageList.length > 0;
 
   // remove \n to avoid empty content
   // refs: https://github.com/lobehub/lobe-chat/pull/6153
@@ -64,6 +66,7 @@ export const AssistantMessage = memo<
           />
         )
       )}
+      {showImageItems && <ImageFileListViewer items={imageList} />}
       {tools && (
         <Flexbox gap={8}>
           {tools.map((toolCall, index) => (

package/src/features/Conversation/Messages/User/index.tsx

@@ -12,16 +12,16 @@ export const UserMessage = memo<
   ChatMessage & {
     editableContent: ReactNode;
   }
->(({ id, editableContent, content, ...res }) => {
+>(({ id, editableContent, content, imageList, fileList }) => {
   if (content === LOADING_FLAT) return <BubblesLoading />;
 
   return (
     <Flexbox gap={8} id={id}>
       {editableContent}
-      {res.imageList && res.imageList?.length > 0 && <ImageFileListViewer items={res.imageList} />}
-      {res.fileList && res.fileList?.length > 0 && (
+      {imageList && imageList?.length > 0 && <ImageFileListViewer items={imageList} />}
+      {fileList && fileList?.length > 0 && (
         <div style={{ marginTop: 8 }}>
-          <FileListViewer items={res.fileList} />
+          <FileListViewer items={fileList} />
         </div>
       )}
     </Flexbox>

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

@@ -10,7 +10,6 @@ import {
   SchemaType,
 } from '@google/generative-ai';
 
-import { VertexAIStream } from '@/libs/agent-runtime/utils/streams/vertex-ai';
 import type { ChatModelCard } from '@/types/llm';
 import { imageUrlToBase64 } from '@/utils/imageToBase64';
 import { safeParseJSON } from '@/utils/safeParseJSON';
@@ -28,9 +27,25 @@ import { ModelProvider } from '../types/type';
 import { AgentRuntimeError } from '../utils/createError';
 import { debugStream } from '../utils/debugStream';
 import { StreamingResponse } from '../utils/response';
-import { GoogleGenerativeAIStream, convertIterableToStream } from '../utils/streams';
+import {
+  GoogleGenerativeAIStream,
+  VertexAIStream,
+  convertIterableToStream,
+} from '../utils/streams';
 import { parseDataUri } from '../utils/uriParser';
 
+const modelsOffSafetySettings = new Set(['gemini-2.0-flash-exp']);
+
+const modelsWithModalities = new Set([
+  'gemini-2.0-flash-exp',
+  'gemini-2.0-flash-exp-image-generation',
+]);
+
+const modelsDisableInstuction = new Set([
+  'gemini-2.0-flash-exp',
+  'gemini-2.0-flash-exp-image-generation',
+]);
+
 export interface GoogleModelCard {
   displayName: string;
   inputTokenLimit: number;
@@ -50,8 +65,7 @@ enum HarmBlockThreshold {
 }
 
 function getThreshold(model: string): HarmBlockThreshold {
-  const useOFF = ['gemini-2.0-flash-exp'];
-  if (useOFF.includes(model)) {
+  if (modelsOffSafetySettings.has(model)) {
     return 'OFF' as HarmBlockThreshold; // https://discuss.ai.google.dev/t/59352
   }
   return HarmBlockThreshold.BLOCK_NONE;
@@ -94,6 +108,8 @@ export class LobeGoogleAI implements LobeRuntimeAI {
           {
             generationConfig: {
               maxOutputTokens: payload.max_tokens,
+              // @ts-expect-error - Google SDK 0.24.0 doesn't have this property for now with
+              response_modalities: modelsWithModalities.has(model) ? ['Text', 'Image'] : undefined,
               temperature: payload.temperature,
               topP: payload.top_p,
             },
@@ -123,7 +139,9 @@ export class LobeGoogleAI implements LobeRuntimeAI {
         )
         .generateContentStream({
           contents,
-          systemInstruction: payload.system as string,
+          systemInstruction: modelsDisableInstuction.has(model)
+            ? undefined
+            : (payload.system as string),
           tools: this.buildGoogleTools(payload.tools, payload),
         });
 

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

@@ -94,4 +94,193 @@ describe('GoogleGenerativeAIStream', () => {
 
     expect(chunks).toEqual([]);
   });
+
+  it('should handle image', async () => {
+    vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1');
+
+    const data = {
+      candidates: [
+        {
+          content: {
+            parts: [{ inlineData: { mimeType: 'image/png', data: 'iVBORw0KGgoAA' } }],
+            role: 'model',
+          },
+          index: 0,
+        },
+      ],
+      usageMetadata: {
+        promptTokenCount: 6,
+        totalTokenCount: 6,
+        promptTokensDetails: [{ modality: 'TEXT', tokenCount: 6 }],
+      },
+      modelVersion: 'gemini-2.0-flash-exp',
+    };
+    const mockGenerateContentResponse = (text: string, functionCalls?: any[]) =>
+      ({
+        text: () => text,
+        functionCall: () => functionCalls?.[0],
+        functionCalls: () => functionCalls,
+      }) as EnhancedGenerateContentResponse;
+
+    const mockGoogleStream = new ReadableStream({
+      start(controller) {
+        controller.enqueue(data);
+
+        controller.close();
+      },
+    });
+
+    const protocolStream = GoogleGenerativeAIStream(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([
+      // image
+      'id: chat_1\n',
+      'event: base64_image\n',
+      `data: "data:image/png;base64,iVBORw0KGgoAA"\n\n`,
+    ]);
+  });
+
+  it('should handle token count', async () => {
+    vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1');
+
+    const data = {
+      candidates: [{ content: { role: 'model' }, finishReason: 'STOP', index: 0 }],
+      usageMetadata: {
+        promptTokenCount: 266,
+        totalTokenCount: 266,
+        promptTokensDetails: [
+          { modality: 'TEXT', tokenCount: 8 },
+          { modality: 'IMAGE', tokenCount: 258 },
+        ],
+      },
+      modelVersion: 'gemini-2.0-flash-exp',
+    };
+
+    const mockGoogleStream = new ReadableStream({
+      start(controller) {
+        controller.enqueue(data);
+
+        controller.close();
+      },
+    });
+
+    const protocolStream = GoogleGenerativeAIStream(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([
+      // stop
+      'id: chat_1\n',
+      'event: stop\n',
+      `data: "STOP"\n\n`,
+      // usage
+      'id: chat_1\n',
+      'event: usage\n',
+      `data: {"inputImageTokens":258,"inputTextTokens":8,"totalInputTokens":266,"totalTokens":266}\n\n`,
+    ]);
+  });
+
+  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 = GoogleGenerativeAIStream(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,"totalInputTokens":19,"totalOutputTokens":11,"totalTokens":30}\n`,
+      ].map((i) => i + '\n'),
+    );
+  });
 });

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

@@ -1,5 +1,6 @@
 import { EnhancedGenerateContentResponse } from '@google/generative-ai';
 
+import { ModelTokensUsage } from '@/types/message';
 import { GroundingSearch } from '@/types/search';
 import { nanoid } from '@/utils/uuid';
 
@@ -18,7 +19,7 @@ const transformGoogleGenerativeAIStream = (
   context: StreamContext,
 ): StreamProtocolChunk | StreamProtocolChunk[] => {
   // maybe need another structure to add support for multiple choices
-  const functionCalls = chunk.functionCalls();
+  const functionCalls = chunk.functionCalls?.();
 
   if (functionCalls) {
     return {
@@ -37,30 +38,76 @@ const transformGoogleGenerativeAIStream = (
       type: 'tool_calls',
     };
   }
-  const text = chunk.text();
 
-  if (chunk.candidates && chunk.candidates[0].groundingMetadata) {
-    const { webSearchQueries, groundingSupports, groundingChunks } =
-      chunk.candidates[0].groundingMetadata;
-    console.log({ groundingChunks, groundingSupports, webSearchQueries });
+  const text = chunk.text?.();
 
-    return [
-      { data: text, id: context.id, type: 'text' },
-      {
-        data: {
-          citations: groundingChunks?.map((chunk) => ({
-            // google 返回的 uri 是经过 google 自己处理过的 url，因此无法展现真实的 favicon
-            // 需要使用 title 作为替换
-            favicon: chunk.web?.title,
-            title: chunk.web?.title,
-            url: chunk.web?.uri,
-          })),
-          searchQueries: webSearchQueries,
-        } as GroundingSearch,
-        id: context.id,
-        type: 'grounding',
-      },
-    ];
+  if (chunk.candidates) {
+    const candidate = chunk.candidates[0];
+
+    // return the grounding
+    if (candidate.groundingMetadata) {
+      const { webSearchQueries, groundingChunks } = candidate.groundingMetadata;
+
+      return [
+        { data: text, id: context.id, type: 'text' },
+        {
+          data: {
+            citations: groundingChunks?.map((chunk) => ({
+              // google 返回的 uri 是经过 google 自己处理过的 url，因此无法展现真实的 favicon
+              // 需要使用 title 作为替换
+              favicon: chunk.web?.title,
+              title: chunk.web?.title,
+              url: chunk.web?.uri,
+            })),
+            searchQueries: webSearchQueries,
+          } as GroundingSearch,
+          id: context.id,
+          type: 'grounding',
+        },
+      ];
+    }
+
+    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',
+          },
+        ].filter(Boolean) as StreamProtocolChunk[];
+      }
+      return { data: candidate.finishReason, id: context?.id, type: 'stop' };
+    }
+
+    if (!!text?.trim()) return { data: text, id: context?.id, type: 'text' };
+
+    // streaming the image
+    if (Array.isArray(candidate.content.parts) && candidate.content.parts.length > 0) {
+      const part = candidate.content.parts[0];
+
+      if (part && part.inlineData && part.inlineData.data && part.inlineData.mimeType) {
+        return {
+          data: `data:${part.inlineData.mimeType};base64,${part.inlineData.data}`,
+          id: context.id,
+          type: 'base64_image',
+        };
+      }
+    }
   }
 
   return {

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

@@ -6,3 +6,4 @@ export * from './openai';
 export * from './protocol';
 export * from './qwen';
 export * from './spark';
+export * from './vertex-ai';

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

@@ -32,6 +32,8 @@ export interface StreamProtocolChunk {
   id?: string;
   type: // pure text
   | 'text'
+    // base64 format image
+    | 'base64_image'
     // Tools use
     | 'tool_calls'
     // Model Thinking

package/src/locales/default/components.ts

@@ -78,6 +78,7 @@ export default {
       custom: '自定义模型，默认设定同时支持函数调用与视觉识别，请根据实际情况验证上述能力的可用性',
       file: '该模型支持上传文件读取与识别',
       functionCall: '该模型支持函数调用（Function Call）',
+      imageOutput: '该模型支持生成图片',
       reasoning: '该模型支持深度思考',
       search: '该模型支持联网搜索',
       tokens: '该模型单个会话最多支持 {{tokens}} Tokens',

package/src/services/chat.ts

@@ -438,6 +438,8 @@ class ChatService {
         provider: params.provider!,
       });
 
+      // remove plugins
+      delete params.plugins;
       await this.getChatCompletion(
         { ...params, messages: oaiMessages, tools },
         {
@@ -474,7 +476,7 @@ class ChatService {
     // handle content type for vision model
     // for the models with visual ability, add image url to content
     // refs: https://platform.openai.com/docs/guides/vision/quick-start
-    const getContent = (m: ChatMessage) => {
+    const getUserContent = (m: ChatMessage) => {
       // only if message doesn't have images and files, then return the plain content
       if ((!m.imageList || m.imageList.length === 0) && (!m.fileList || m.fileList.length === 0))
         return m.content;
@@ -490,27 +492,43 @@ class ChatService {
       ] as UserMessageContentPart[];
     };
 
+    const getAssistantContent = (m: ChatMessage) => {
+      // signature is a signal of anthropic thinking mode
+      const shouldIncludeThinking = m.reasoning && !!m.reasoning?.signature;
+
+      if (shouldIncludeThinking) {
+        return [
+          {
+            signature: m.reasoning!.signature,
+            thinking: m.reasoning!.content,
+            type: 'thinking',
+          },
+          { text: m.content, type: 'text' },
+        ] as UserMessageContentPart[];
+      }
+      // only if message doesn't have images and files, then return the plain content
+
+      if (m.imageList && m.imageList.length > 0) {
+        return [
+          !!m.content ? { text: m.content, type: 'text' } : undefined,
+          ...m.imageList.map(
+            (i) => ({ image_url: { detail: 'auto', url: i.url }, type: 'image_url' }) as const,
+          ),
+        ].filter(Boolean) as UserMessageContentPart[];
+      }
+
+      return m.content;
+    };
+
     let postMessages = messages.map((m): OpenAIChatMessage => {
       const supportTools = isCanUseFC(model, provider);
       switch (m.role) {
         case 'user': {
-          return { content: getContent(m), role: m.role };
+          return { content: getUserContent(m), role: m.role };
         }
 
         case 'assistant': {
-          // signature is a signal of anthropic thinking mode
-          const shouldIncludeThinking = m.reasoning && !!m.reasoning?.signature;
-
-          const content = shouldIncludeThinking
-            ? [
-                {
-                  signature: m.reasoning!.signature,
-                  thinking: m.reasoning!.content,
-                  type: 'thinking',
-                } as any,
-                { text: m.content, type: 'text' },
-              ]
-            : m.content;
+          const content = getAssistantContent(m);
 
           if (!supportTools) {
             return { content, role: m.role };

package/src/services/file/client.ts

@@ -11,6 +11,8 @@ export class ClientService extends BaseClientService implements IFileService {
   }
 
   createFile: IFileService['createFile'] = async (file) => {
+    const { isExist } = await this.fileModel.checkHash(file.hash!);
+
     // save to local storage
     // we may want to save to a remote server later
     const res = await this.fileModel.create(
@@ -23,7 +25,7 @@ export class ClientService extends BaseClientService implements IFileService {
         size: file.size,
         url: file.url!,
       },
-      true,
+      !isExist,
     );
 
     // get file to base64 url

package/src/services/message/server.ts

@@ -64,8 +64,8 @@ export class ServerService implements IMessageService {
     return lambdaClient.message.updateMessagePlugin.mutate({ id, value: { arguments: args } });
   };
 
-  updateMessage: IMessageService['updateMessage'] = async (id, message) => {
-    return lambdaClient.message.update.mutate({ id, value: message });
+  updateMessage: IMessageService['updateMessage'] = async (id, value) => {
+    return lambdaClient.message.update.mutate({ id, value });
   };
 
   updateMessageTranslate: IMessageService['updateMessageTranslate'] = async (id, translate) => {