@lobehub/chat 1.70.10 → 1.71.0

package/src/database/migrations/meta/_journal.json

@@ -112,6 +112,13 @@
       "when": 1739901891891,
       "tag": "0015_add_message_search_metadata",
       "breakpoints": true
+    },
+    {
+      "idx": 16,
+      "version": "7",
+      "when": 1741844738677,
+      "tag": "0016_add_message_index",
+      "breakpoints": true
     }
   ],
   "version": "6"

package/src/database/schemas/message.ts

@@ -73,6 +73,9 @@ export const messages = pgTable(
       table.clientId,
       table.userId,
     ),
+    topicIdIdx: index('messages_topic_id_idx').on(table.topicId),
+    parentIdIdx: index('messages_parent_id_idx').on(table.parentId),
+    quotaIdIdx: index('messages_quota_id_idx').on(table.quotaId),
   }),
 );
 

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

@@ -31,6 +31,9 @@ import { StreamingResponse } from '../utils/response';
 import { GoogleGenerativeAIStream, 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']);
+
 export interface GoogleModelCard {
   displayName: string;
   inputTokenLimit: number;
@@ -50,8 +53,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 +96,10 @@ 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,
             },

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

@@ -94,4 +94,103 @@ 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`,
+    ]);
+  });
 });

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,75 @@ 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 [
+          { 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',
+          },
+        ];
+      }
+      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/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/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) => {

package/src/services/message/type.ts

@@ -8,7 +8,7 @@ import {
   ChatTranslate,
   CreateMessageParams,
   MessageItem,
-  ModelRankItem,
+  ModelRankItem, UpdateMessageParams,
 } from '@/types/message';
 
 /* eslint-disable typescript-sort-keys/interface */
@@ -33,7 +33,7 @@ export interface IMessageService {
   rankModels(): Promise<ModelRankItem[]>;
   getHeatmaps(): Promise<HeatmapsProps['data']>;
   updateMessageError(id: string, error: ChatMessageError): Promise<any>;
-  updateMessage(id: string, message: Partial<MessageItem>): Promise<any>;
+  updateMessage(id: string, message: Partial<UpdateMessageParams>): Promise<any>;
   updateMessageTTS(id: string, tts: Partial<ChatTTS> | false): Promise<any>;
   updateMessageTranslate(id: string, translate: Partial<ChatTranslate> | false): Promise<any>;
   updateMessagePluginState(id: string, value: Record<string, any>): Promise<any>;

package/src/services/upload.ts

@@ -1,14 +1,95 @@
+import dayjs from 'dayjs';
+import { sha256 } from 'js-sha256';
+
 import { fileEnv } from '@/config/file';
+import { isServerMode } from '@/const/version';
+import { parseDataUri } from '@/libs/agent-runtime/utils/uriParser';
 import { edgeClient } from '@/libs/trpc/client';
 import { API_ENDPOINTS } from '@/services/_url';
 import { clientS3Storage } from '@/services/file/ClientS3';
-import { FileMetadata } from '@/types/files';
+import { FileMetadata, UploadBase64ToS3Result } from '@/types/files';
 import { FileUploadState, FileUploadStatus } from '@/types/files/upload';
 import { uuid } from '@/utils/uuid';
 
 export const UPLOAD_NETWORK_ERROR = 'NetWorkError';
 
+interface UploadFileToS3Options {
+  directory?: string;
+  filename?: string;
+  onProgress?: (status: FileUploadStatus, state: FileUploadState) => void;
+}
+
 class UploadService {
+  /**
+   * uniform upload method for both server and client
+   */
+  uploadFileToS3 = async (
+    file: File,
+    options: UploadFileToS3Options = {},
+  ): Promise<FileMetadata> => {
+    const { directory, onProgress } = options;
+
+    if (isServerMode) {
+      return this.uploadWithProgress(file, { directory, onProgress });
+    } else {
+      const fileArrayBuffer = await file.arrayBuffer();
+
+      // 1. check file hash
+      const hash = sha256(fileArrayBuffer);
+
+      return this.uploadToClientS3(hash, file);
+    }
+  };
+
+  uploadBase64ToS3 = async (
+    base64Data: string,
+    options: UploadFileToS3Options = {},
+  ): Promise<UploadBase64ToS3Result> => {
+    // 解析 base64 数据
+    const { base64, mimeType, type } = parseDataUri(base64Data);
+
+    if (!base64 || !mimeType || type !== 'base64') {
+      throw new Error('Invalid base64 data for image');
+    }
+
+    // 将 base64 转换为 Blob
+    const byteCharacters = atob(base64);
+    const byteArrays = [];
+
+    // 分块处理以避免内存问题
+    for (let offset = 0; offset < byteCharacters.length; offset += 1024) {
+      const slice = byteCharacters.slice(offset, offset + 1024);
+
+      const byteNumbers: number[] = Array.from({ length: slice.length });
+      for (let i = 0; i < slice.length; i++) {
+        byteNumbers[i] = slice.charCodeAt(i);
+      }
+
+      const byteArray = new Uint8Array(byteNumbers);
+      byteArrays.push(byteArray);
+    }
+
+    const blob = new Blob(byteArrays, { type: mimeType });
+
+    // 确定文件扩展名
+    const fileExtension = mimeType.split('/')[1] || 'png';
+    const fileName = `${options.filename || `image_${dayjs().format('YYYY-MM-DD-hh-mm-ss')}`}.${fileExtension}`;
+
+    // 创建文件对象
+    const file = new File([blob], fileName, { type: mimeType });
+
+    // 使用统一的上传方法
+    const metadata = await this.uploadFileToS3(file, options);
+    const hash = sha256(await file.arrayBuffer());
+
+    return {
+      fileType: mimeType,
+      hash,
+      metadata,
+      size: file.size,
+    };
+  };
+
   uploadWithProgress = async (
     file: File,
     {