@lobehub/chat 1.121.1 → 1.122.0

package/src/server/routers/lambda/message.ts

@@ -7,6 +7,7 @@ import { authedProcedure, publicProcedure, router } from '@/libs/trpc/lambda';
 import { serverDatabase } from '@/libs/trpc/lambda/middleware';
 import { FileService } from '@/server/services/file';
 import { ChatMessage } from '@/types/message';
+import { UpdateMessageRAGParamsSchema } from '@/types/message/rag';
 import { BatchTaskResult } from '@/types/service';
 
 type ChatMessageList = ChatMessage[];
@@ -174,6 +175,12 @@ export const messageRouter = router({
       return ctx.messageModel.updateMessagePlugin(input.id, input.value);
     }),
 
+  updateMessageRAG: messageProcedure
+    .input(UpdateMessageRAGParamsSchema)
+    .mutation(async ({ input, ctx }) => {
+      await ctx.messageModel.updateMessageRAG(input.id, input.value);
+    }),
+
   updatePluginError: messageProcedure
     .input(
       z.object({

package/src/server/services/aiChat/index.test.ts

@@ -0,0 +1,57 @@
+import { describe, expect, it, vi } from 'vitest';
+
+import { MessageModel } from '@/database/models/message';
+import { TopicModel } from '@/database/models/topic';
+import { LobeChatDatabase } from '@/database/type';
+import { FileService } from '@/server/services/file';
+
+import { AiChatService } from '.';
+
+vi.mock('@/database/models/message');
+vi.mock('@/database/models/topic');
+vi.mock('@/server/services/file');
+
+describe('AiChatService', () => {
+  it('getMessagesAndTopics should fetch messages and topics concurrently', async () => {
+    const serverDB = {} as unknown as LobeChatDatabase;
+
+    const mockQueryMessages = vi.fn().mockResolvedValue([{ id: 'm1' }]);
+    const mockQueryTopics = vi.fn().mockResolvedValue([{ id: 't1' }]);
+
+    vi.mocked(MessageModel).mockImplementation(() => ({ query: mockQueryMessages }) as any);
+    vi.mocked(TopicModel).mockImplementation(() => ({ query: mockQueryTopics }) as any);
+    vi.mocked(FileService).mockImplementation(
+      () => ({ getFullFileUrl: vi.fn().mockResolvedValue('url') }) as any,
+    );
+
+    const service = new AiChatService(serverDB, 'u1');
+
+    const res = await service.getMessagesAndTopics({ includeTopic: true, sessionId: 's1' });
+
+    expect(mockQueryMessages).toHaveBeenCalledWith(
+      { includeTopic: true, sessionId: 's1' },
+      expect.objectContaining({ postProcessUrl: expect.any(Function) }),
+    );
+    expect(mockQueryTopics).toHaveBeenCalledWith({ sessionId: 's1' });
+    expect(res.messages).toEqual([{ id: 'm1' }]);
+    expect(res.topics).toEqual([{ id: 't1' }]);
+  });
+
+  it('getMessagesAndTopics should not query topics when includeTopic is false', async () => {
+    const serverDB = {} as unknown as LobeChatDatabase;
+
+    const mockQueryMessages = vi.fn().mockResolvedValue([]);
+    vi.mocked(MessageModel).mockImplementation(() => ({ query: mockQueryMessages }) as any);
+    vi.mocked(TopicModel).mockImplementation(() => ({ query: vi.fn() }) as any);
+    vi.mocked(FileService).mockImplementation(
+      () => ({ getFullFileUrl: vi.fn().mockResolvedValue('url') }) as any,
+    );
+
+    const service = new AiChatService(serverDB, 'u1');
+
+    const res = await service.getMessagesAndTopics({ includeTopic: false, topicId: 't1' });
+
+    expect(mockQueryMessages).toHaveBeenCalled();
+    expect(res.topics).toBeUndefined();
+  });
+});

package/src/server/services/aiChat/index.ts

@@ -0,0 +1,36 @@
+import { MessageModel } from '@/database/models/message';
+import { TopicModel } from '@/database/models/topic';
+import { LobeChatDatabase } from '@/database/type';
+import { FileService } from '@/server/services/file';
+
+export class AiChatService {
+  private userId: string;
+  private messageModel: MessageModel;
+  private fileService: FileService;
+  private topicModel: TopicModel;
+
+  constructor(serverDB: LobeChatDatabase, userId: string) {
+    this.userId = userId;
+
+    this.messageModel = new MessageModel(serverDB, userId);
+    this.topicModel = new TopicModel(serverDB, userId);
+    this.fileService = new FileService(serverDB, userId);
+  }
+
+  async getMessagesAndTopics(params: {
+    current?: number;
+    includeTopic?: boolean;
+    pageSize?: number;
+    sessionId?: string;
+    topicId?: string;
+  }) {
+    const [messages, topics] = await Promise.all([
+      this.messageModel.query(params, {
+        postProcessUrl: (path) => this.fileService.getFullFileUrl(path),
+      }),
+      params.includeTopic ? this.topicModel.query({ sessionId: params.sessionId }) : undefined,
+    ]);
+
+    return { messages, topics };
+  }
+}

package/src/services/aiChat.ts

@@ -0,0 +1,12 @@
+import { SendMessageServerParams } from '@lobechat/types';
+import { cleanObject } from '@lobechat/utils';
+
+import { lambdaClient } from '@/libs/trpc/client';
+
+class AiChatService {
+  sendMessageInServer = async (params: SendMessageServerParams) => {
+    return lambdaClient.aiChat.sendMessageInServer.mutate(cleanObject(params));
+  };
+}
+
+export const aiChatService = new AiChatService();

package/src/services/message/_deprecated.ts

@@ -143,4 +143,8 @@ export class ClientService implements IMessageService {
   async updateMessagePluginError() {
     throw new Error('Method not implemented.');
   }
+
+  async updateMessageRAG(): Promise<void> {
+    throw new Error('Method not implemented.');
+  }
 }

package/src/services/message/client.ts

@@ -96,6 +96,11 @@ export class ClientService extends BaseClientService implements IMessageService
     return this.messageModel.updateMessagePlugin(id, { error: value });
   };
 
+  updateMessageRAG: IMessageService['updateMessageRAG'] = async (id, value) => {
+    console.log(id, value);
+    throw new Error('not implemented');
+  };
+
   updateMessagePluginArguments: IMessageService['updateMessagePluginArguments'] = async (
     id,
     value,

package/src/services/message/server.ts

@@ -84,6 +84,10 @@ export class ServerService implements IMessageService {
     return lambdaClient.message.updatePluginError.mutate({ id, value: error as any });
   };
 
+  updateMessageRAG: IMessageService['updateMessageRAG'] = async (id, data) => {
+    return lambdaClient.message.updateMessageRAG.mutate({ id, value: data });
+  };
+
   removeMessage: IMessageService['removeMessage'] = async (id) => {
     return lambdaClient.message.removeMessage.mutate({ id });
   };

package/src/services/message/type.ts

@@ -11,6 +11,7 @@ import {
   ModelRankItem,
   UpdateMessageParams,
 } from '@/types/message';
+import { UpdateMessageRAGParams } from '@/types/message/rag';
 
 /* eslint-disable typescript-sort-keys/interface */
 
@@ -39,6 +40,7 @@ export interface IMessageService {
   updateMessageTranslate(id: string, translate: Partial<ChatTranslate> | false): Promise<any>;
   updateMessagePluginState(id: string, value: Record<string, any>): Promise<any>;
   updateMessagePluginError(id: string, value: ChatMessagePluginError | null): Promise<any>;
+  updateMessageRAG(id: string, value: UpdateMessageRAGParams): Promise<void>;
   updateMessagePluginArguments(id: string, value: string | Record<string, any>): Promise<any>;
   removeMessage(id: string): Promise<any>;
   removeMessages(ids: string[]): Promise<any>;

package/src/store/chat/slices/aiChat/actions/generateAIChat.ts

@@ -152,7 +152,13 @@ export const generateAIChat: StateCreator<
   },
 
   sendMessage: async ({ message, files, onlyAddUserMessage, isWelcomeQuestion }) => {
-    const { internal_coreProcessMessage, activeTopicId, activeId, activeThreadId } = get();
+    const {
+      internal_coreProcessMessage,
+      activeTopicId,
+      activeId,
+      activeThreadId,
+      sendMessageInServer,
+    } = get();
     if (!activeId) return;
 
     const fileIdList = files?.map((f) => f.id);
@@ -162,6 +168,10 @@ export const generateAIChat: StateCreator<
     // if message is empty or no files, then stop
     if (!message && !hasFile) return;
 
+    // router to server mode send message
+    if (isServerMode)
+      return sendMessageInServer({ message, files, onlyAddUserMessage, isWelcomeQuestion });
+
     set({ isCreatingMessage: true }, false, n('creatingMessage/start'));
 
     const newMessage: CreateMessageParams = {

package/src/store/chat/slices/aiChat/actions/generateAIChatV2.ts

@@ -0,0 +1,410 @@
+/* eslint-disable sort-keys-fix/sort-keys-fix, typescript-sort-keys/interface */
+// Disable the auto sort key eslint rule to make the code more logic and readable
+import { INBOX_SESSION_ID, isDesktop } from '@lobechat/const';
+import { knowledgeBaseQAPrompts } from '@lobechat/prompts';
+import {
+  ChatMessage,
+  ChatTopic,
+  MessageSemanticSearchChunk,
+  SendMessageParams,
+  TraceNameMap,
+} from '@lobechat/types';
+import { t } from 'i18next';
+import { StateCreator } from 'zustand/vanilla';
+
+import { aiChatService } from '@/services/aiChat';
+import { chatService } from '@/services/chat';
+import { messageService } from '@/services/message';
+import { getAgentStoreState } from '@/store/agent';
+import { agentChatConfigSelectors, agentSelectors } from '@/store/agent/slices/chat';
+import { aiModelSelectors, aiProviderSelectors, getAiInfraStoreState } from '@/store/aiInfra';
+import type { ChatStore } from '@/store/chat/store';
+import { getSessionStoreState } from '@/store/session';
+import { WebBrowsingManifest } from '@/tools/web-browsing';
+import { setNamespace } from '@/utils/storeDebug';
+
+import { chatSelectors, topicSelectors } from '../../../selectors';
+import { messageMapKey } from '../../../utils/messageMapKey';
+
+const n = setNamespace('ai');
+
+export interface AIGenerateV2Action {
+  /**
+   * Sends a new message to the AI chat system
+   */
+  sendMessageInServer: (params: SendMessageParams) => Promise<void>;
+  internal_refreshAiChat: (params: {
+    topics?: ChatTopic[];
+    messages: ChatMessage[];
+    sessionId: string;
+    topicId?: string;
+  }) => void;
+  /**
+   * Executes the core processing logic for AI messages
+   * including preprocessing and postprocessing steps
+   */
+  internal_execAgentRuntime: (params: {
+    messages: ChatMessage[];
+    userMessageId: string;
+    assistantMessageId: string;
+    isWelcomeQuestion?: boolean;
+    inSearchWorkflow?: boolean;
+    /**
+     * the RAG query content, should be embedding and used in the semantic search
+     */
+    ragQuery?: string;
+    threadId?: string;
+    inPortalThread?: boolean;
+    traceId?: string;
+  }) => Promise<void>;
+}
+
+export const generateAIChatV2: StateCreator<
+  ChatStore,
+  [['zustand/devtools', never]],
+  [],
+  AIGenerateV2Action
+> = (set, get) => ({
+  sendMessageInServer: async ({ message, files, onlyAddUserMessage, isWelcomeQuestion }) => {
+    const { activeTopicId, activeId, activeThreadId, internal_execAgentRuntime } = get();
+    if (!activeId) return;
+
+    const fileIdList = files?.map((f) => f.id);
+
+    const hasFile = !!fileIdList && fileIdList.length > 0;
+
+    // if message is empty or no files, then stop
+    if (!message && !hasFile) return;
+
+    if (onlyAddUserMessage) {
+      await get().addUserMessage({ message, fileList: fileIdList });
+
+      return;
+    }
+
+    const messages = chatSelectors.activeBaseChats(get());
+
+    // use optimistic update to avoid the slow waiting
+    const tempId = get().internal_createTmpMessage({
+      content: message,
+      // if message has attached with files, then add files to message and the agent
+      files: fileIdList,
+      role: 'user',
+      sessionId: activeId,
+      // if there is activeTopicId，then add topicId to message
+      topicId: activeTopicId,
+      threadId: activeThreadId,
+    });
+
+    get().internal_toggleMessageLoading(true, tempId);
+    set({ isCreatingMessage: true }, false, 'creatingMessage/start');
+
+    const { model, provider } = agentSelectors.currentAgentConfig(getAgentStoreState());
+
+    const data = await aiChatService.sendMessageInServer({
+      newUserMessage: {
+        content: message,
+        files: fileIdList,
+      },
+      // if there is activeTopicId，then add topicId to message
+      topicId: activeTopicId,
+      threadId: activeThreadId,
+      newTopic: !activeTopicId
+        ? {
+            topicMessageIds: messages.map((m) => m.id),
+            title: t('defaultTitle', { ns: 'topic' }),
+          }
+        : undefined,
+      sessionId: activeId === INBOX_SESSION_ID ? undefined : activeId,
+      newAssistantMessage: { model, provider: provider! },
+    });
+
+    // refresh the total data
+    get().internal_refreshAiChat({
+      messages: data.messages,
+      topics: data.topics,
+      sessionId: activeId,
+      topicId: data.topicId,
+    });
+    get().internal_dispatchMessage({ type: 'deleteMessage', id: tempId });
+
+    if (!activeTopicId) {
+      await get().switchTopic(data.topicId!, true);
+    }
+
+    get().internal_toggleMessageLoading(false, tempId);
+
+    //  update assistant update to make it rerank
+    getSessionStoreState().triggerSessionUpdate(get().activeId);
+
+    // Get the current messages to generate AI response
+    // remove the latest assistant message id
+    const baseMessages = chatSelectors
+      .activeBaseChats(get())
+      .filter((item) => item.id !== data.assistantMessageId);
+
+    try {
+      await internal_execAgentRuntime({
+        messages: baseMessages,
+        userMessageId: data.userMessageId,
+        assistantMessageId: data.assistantMessageId,
+        isWelcomeQuestion,
+        ragQuery: get().internal_shouldUseRAG() ? message : undefined,
+        threadId: activeThreadId,
+      });
+      set({ isCreatingMessage: false }, false, 'creatingMessage/stop');
+
+      const summaryTitle = async () => {
+        // check activeTopic and then auto update topic title
+        if (data.isCreatNewTopic) {
+          await get().summaryTopicTitle(data.topicId, data.messages);
+          return;
+        }
+
+        if (!activeTopicId) return;
+
+        const topic = topicSelectors.getTopicById(activeTopicId)(get());
+
+        if (topic && !topic.title) {
+          const chats = chatSelectors.getBaseChatsByKey(messageMapKey(activeId, topic.id))(get());
+          await get().summaryTopicTitle(topic.id, chats);
+        }
+      };
+      //
+      // // if there is relative files, then add files to agent
+      // // only available in server mode
+      const userFiles = chatSelectors.currentUserFiles(get()).map((f) => f.id);
+      const addFilesToAgent = async () => {
+        await getAgentStoreState().addFilesToAgent(userFiles, false);
+      };
+
+      await Promise.all([summaryTitle(), addFilesToAgent()]);
+    } catch (e) {
+      console.error(e);
+      set({ isCreatingMessage: false }, false, 'creatingMessage/stop');
+    }
+  },
+
+  internal_refreshAiChat: ({ topics, messages, sessionId, topicId }) => {
+    set(
+      {
+        topicMaps: topics ? { ...get().topicMaps, [sessionId]: topics } : get().topicMaps,
+        messagesMap: { ...get().messagesMap, [messageMapKey(sessionId, topicId)]: messages },
+      },
+      false,
+      'refreshAiChat',
+    );
+  },
+
+  internal_execAgentRuntime: async (params) => {
+    const {
+      assistantMessageId: assistantId,
+      userMessageId,
+      ragQuery,
+      messages: originalMessages,
+    } = params;
+    const {
+      internal_fetchAIChatMessage,
+      triggerToolCalls,
+      refreshMessages,
+      internal_updateMessageRAG,
+    } = get();
+
+    // create a new array to avoid the original messages array change
+    const messages = [...originalMessages];
+
+    const agentStoreState = getAgentStoreState();
+    const { model, provider, chatConfig } = agentSelectors.currentAgentConfig(agentStoreState);
+
+    let fileChunks: MessageSemanticSearchChunk[] | undefined;
+    let ragQueryId;
+
+    // go into RAG flow if there is ragQuery flag
+    if (ragQuery) {
+      // 1. get the relative chunks from semantic search
+      const { chunks, queryId, rewriteQuery } = await get().internal_retrieveChunks(
+        userMessageId,
+        ragQuery,
+        // should skip the last content
+        messages.map((m) => m.content).slice(0, messages.length - 1),
+      );
+
+      ragQueryId = queryId;
+
+      const lastMsg = messages.pop() as ChatMessage;
+
+      // 2. build the retrieve context messages
+      const knowledgeBaseQAContext = knowledgeBaseQAPrompts({
+        chunks,
+        userQuery: lastMsg.content,
+        rewriteQuery,
+        knowledge: agentSelectors.currentEnabledKnowledge(agentStoreState),
+      });
+
+      // 3. add the retrieve context messages to the messages history
+      messages.push({
+        ...lastMsg,
+        content: (lastMsg.content + '\n\n' + knowledgeBaseQAContext).trim(),
+      });
+
+      fileChunks = chunks.map((c) => ({ id: c.id, similarity: c.similarity }));
+
+      if (fileChunks.length > 0) {
+        await internal_updateMessageRAG(assistantId, { ragQueryId, fileChunks });
+      }
+    }
+
+    // 3. place a search with the search working model if this model is not support tool use
+    const aiInfraStoreState = getAiInfraStoreState();
+    const isModelSupportToolUse = aiModelSelectors.isModelSupportToolUse(
+      model,
+      provider!,
+    )(aiInfraStoreState);
+    const isProviderHasBuiltinSearch = aiProviderSelectors.isProviderHasBuiltinSearch(provider!)(
+      aiInfraStoreState,
+    );
+    const isModelHasBuiltinSearch = aiModelSelectors.isModelHasBuiltinSearch(
+      model,
+      provider!,
+    )(aiInfraStoreState);
+    const useModelBuiltinSearch = agentChatConfigSelectors.useModelBuiltinSearch(agentStoreState);
+    const useModelSearch =
+      (isProviderHasBuiltinSearch || isModelHasBuiltinSearch) && useModelBuiltinSearch;
+    const isAgentEnableSearch = agentChatConfigSelectors.isAgentEnableSearch(agentStoreState);
+
+    if (isAgentEnableSearch && !useModelSearch && !isModelSupportToolUse) {
+      const { model, provider } = agentChatConfigSelectors.searchFCModel(agentStoreState);
+
+      let isToolsCalling = false;
+      let isError = false;
+
+      const abortController = get().internal_toggleChatLoading(
+        true,
+        assistantId,
+        n('generateMessage(start)', { messageId: assistantId, messages }),
+      );
+
+      get().internal_toggleSearchWorkflow(true, assistantId);
+      await chatService.fetchPresetTaskResult({
+        params: { messages, model, provider, plugins: [WebBrowsingManifest.identifier] },
+        onFinish: async (_, { toolCalls, usage }) => {
+          if (toolCalls && toolCalls.length > 0) {
+            get().internal_toggleToolCallingStreaming(assistantId, undefined);
+            // update tools calling
+            await get().internal_updateMessageContent(assistantId, '', {
+              toolCalls,
+              metadata: usage,
+              model,
+              provider,
+            });
+          }
+        },
+        trace: {
+          traceId: params.traceId,
+          sessionId: get().activeId,
+          topicId: get().activeTopicId,
+          traceName: TraceNameMap.SearchIntentRecognition,
+        },
+        abortController,
+        onMessageHandle: async (chunk) => {
+          if (chunk.type === 'tool_calls') {
+            get().internal_toggleSearchWorkflow(false, assistantId);
+            get().internal_toggleToolCallingStreaming(assistantId, chunk.isAnimationActives);
+            get().internal_dispatchMessage({
+              id: assistantId,
+              type: 'updateMessage',
+              value: { tools: get().internal_transformToolCalls(chunk.tool_calls) },
+            });
+            isToolsCalling = true;
+          }
+
+          if (chunk.type === 'text') {
+            abortController!.abort('not fc');
+          }
+        },
+        onErrorHandle: async (error) => {
+          isError = true;
+          await messageService.updateMessageError(assistantId, error);
+          await refreshMessages();
+        },
+      });
+
+      get().internal_toggleChatLoading(
+        false,
+        assistantId,
+        n('generateMessage(start)', { messageId: assistantId, messages }),
+      );
+      get().internal_toggleSearchWorkflow(false, assistantId);
+
+      // if there is error, then stop
+      if (isError) return;
+
+      // if it's the function call message, trigger the function method
+      if (isToolsCalling) {
+        get().internal_toggleMessageInToolsCalling(true, assistantId);
+        await refreshMessages();
+        await triggerToolCalls(assistantId, {
+          threadId: params?.threadId,
+          inPortalThread: params?.inPortalThread,
+        });
+
+        // then story the workflow
+        return;
+      }
+    }
+
+    // 4. fetch the AI response
+    const { isFunctionCall, content } = await internal_fetchAIChatMessage({
+      messages,
+      messageId: assistantId,
+      params,
+      model,
+      provider: provider!,
+    });
+
+    // 5. if it's the function call message, trigger the function method
+    if (isFunctionCall) {
+      get().internal_toggleMessageInToolsCalling(true, assistantId);
+      await refreshMessages();
+      await triggerToolCalls(assistantId, {
+        threadId: params?.threadId,
+        inPortalThread: params?.inPortalThread,
+      });
+    } else {
+      // 显示桌面通知（仅在桌面端且窗口隐藏时）
+      if (isDesktop) {
+        try {
+          // 动态导入桌面通知服务，避免在非桌面端环境中导入
+          const { desktopNotificationService } = await import(
+            '@/services/electron/desktopNotification'
+          );
+
+          await desktopNotificationService.showNotification({
+            body: content,
+            title: t('notification.finishChatGeneration', { ns: 'electron' }),
+          });
+        } catch (error) {
+          // 静默处理错误，不影响正常流程
+          console.error('Desktop notification error:', error);
+        }
+      }
+    }
+
+    // 6. summary history if context messages is larger than historyCount
+    const historyCount = agentChatConfigSelectors.historyCount(agentStoreState);
+
+    if (
+      agentChatConfigSelectors.enableHistoryCount(agentStoreState) &&
+      chatConfig.enableCompressHistory &&
+      originalMessages.length > historyCount
+    ) {
+      // after generation: [u1,a1,u2,a2,u3,a3]
+      // but the `originalMessages` is still: [u1,a1,u2,a2,u3]
+      // So if historyCount=2, we need to summary [u1,a1,u2,a2]
+      // because user find UI is [u1,a1,u2,a2 | u3,a3]
+      const historyMessages = originalMessages.slice(0, -historyCount + 1);
+
+      await get().internal_summaryHistory(historyMessages);
+    }
+  },
+});

package/src/store/chat/slices/aiChat/actions/index.ts

@@ -3,10 +3,15 @@ import { StateCreator } from 'zustand/vanilla';
 import { ChatStore } from '@/store/chat/store';
 
 import { AIGenerateAction, generateAIChat } from './generateAIChat';
+import { AIGenerateV2Action, generateAIChatV2 } from './generateAIChatV2';
 import { ChatMemoryAction, chatMemory } from './memory';
 import { ChatRAGAction, chatRag } from './rag';
 
-export interface ChatAIChatAction extends ChatRAGAction, ChatMemoryAction, AIGenerateAction {
+export interface ChatAIChatAction
+  extends ChatRAGAction,
+    ChatMemoryAction,
+    AIGenerateAction,
+    AIGenerateV2Action {
   /**/
 }
 
@@ -19,4 +24,5 @@ export const chatAiChat: StateCreator<
   ...chatRag(...params),
   ...generateAIChat(...params),
   ...chatMemory(...params),
+  ...generateAIChatV2(...params),
 });