@lobehub/chat 1.121.0 → 1.122.0

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),
 });

package/src/store/chat/slices/message/action.ts

@@ -22,6 +22,7 @@ import {
   ModelReasoning,
 } from '@/types/message';
 import { ChatImageItem } from '@/types/message/image';
+import { UpdateMessageRAGParams } from '@/types/message/rag';
 import { GroundingSearch } from '@/types/search';
 import { TraceEventPayloads } from '@/types/trace';
 import { Action, setNamespace } from '@/utils/storeDebug';
@@ -39,6 +40,7 @@ const SWR_USE_FETCH_MESSAGES = 'SWR_USE_FETCH_MESSAGES';
 export interface ChatMessageAction {
   // create
   addAIMessage: () => Promise<void>;
+  addUserMessage: (params: { message: string; fileList?: string[] }) => Promise<void>;
   // delete
   /**
    * clear message on the active session
@@ -59,10 +61,11 @@ export interface ChatMessageAction {
   ) => SWRResponse<ChatMessage[]>;
   copyMessage: (id: string, content: string) => Promise<void>;
   refreshMessages: () => Promise<void>;
-
+  replaceMessages: (messages: ChatMessage[]) => void;
   // =========  ↓ Internal Method ↓  ========== //
   // ========================================== //
   // ========================================== //
+  internal_updateMessageRAG: (id: string, input: UpdateMessageRAGParams) => Promise<void>;
 
   /**
    * update message at the frontend
@@ -213,6 +216,21 @@ export const chatMessage: StateCreator<
 
     updateInputMessage('');
   },
+  addUserMessage: async ({ message, fileList }) => {
+    const { internal_createMessage, updateInputMessage, activeTopicId, activeId } = get();
+    if (!activeId) return;
+
+    await internal_createMessage({
+      content: message,
+      files: fileList,
+      role: 'user',
+      sessionId: activeId,
+      // if there is activeTopicId，then add topicId to message
+      topicId: activeTopicId,
+    });
+
+    updateInputMessage('');
+  },
   copyMessage: async (id, content) => {
     await copyToClipboard(content);
 
@@ -266,6 +284,25 @@ export const chatMessage: StateCreator<
   refreshMessages: async () => {
     await mutate([SWR_USE_FETCH_MESSAGES, get().activeId, get().activeTopicId]);
   },
+  replaceMessages: (messages) => {
+    set(
+      {
+        messagesMap: {
+          ...get().messagesMap,
+          [messageMapKey(get().activeId, get().activeTopicId)]: messages,
+        },
+      },
+      false,
+      'replaceMessages',
+    );
+  },
+
+  internal_updateMessageRAG: async (id, data) => {
+    const { refreshMessages } = get();
+
+    await messageService.updateMessageRAG(id, data);
+    await refreshMessages();
+  },
 
   // the internal process method of the AI message
   internal_dispatchMessage: (payload) => {

package/src/store/chat/slices/message/reducer.ts

@@ -11,6 +11,11 @@ import {
 import { merge } from '@/utils/merge';
 
 interface UpdateMessages {
+  type: 'updateMessages';
+  value: ChatMessage[];
+}
+
+interface UpdateMessage {
   id: string;
   type: 'updateMessage';
   value: Partial<ChatMessage>;
@@ -72,6 +77,7 @@ interface UpdateMessageExtra {
 
 export type MessageDispatch =
   | CreateMessage
+  | UpdateMessage
   | UpdateMessages
   | UpdatePluginState
   | UpdateMessageExtra
@@ -194,6 +200,11 @@ export const messagesReducer = (state: ChatMessage[], payload: MessageDispatch):
         draftState.push({ ...value, createdAt: Date.now(), id, meta: {}, updatedAt: Date.now() });
       });
     }
+
+    case 'updateMessages': {
+      return payload.value;
+    }
+
     case 'deleteMessage': {
       return produce(state, (draft) => {
         const { id } = payload;

package/src/store/chat/slices/topic/reducer.ts

@@ -14,12 +14,21 @@ interface UpdateChatTopicAction {
   value: Partial<ChatTopic>;
 }
 
+interface UpdateTopicsAction {
+  type: 'updateTopics';
+  value: ChatTopic[];
+}
+
 interface DeleteChatTopicAction {
   id: string;
   type: 'deleteTopic';
 }
 
-export type ChatTopicDispatch = AddChatTopicAction | UpdateChatTopicAction | DeleteChatTopicAction;
+export type ChatTopicDispatch =
+  | AddChatTopicAction
+  | UpdateChatTopicAction
+  | DeleteChatTopicAction
+  | UpdateTopicsAction;
 
 export const topicReducer = (state: ChatTopic[] = [], payload: ChatTopicDispatch): ChatTopic[] => {
   switch (payload.type) {
@@ -51,6 +60,10 @@ export const topicReducer = (state: ChatTopic[] = [], payload: ChatTopicDispatch
       });
     }
 
+    case 'updateTopics': {
+      return payload.value;
+    }
+
     case 'deleteTopic': {
       return produce(state, (draftState) => {
         const topicIndex = draftState.findIndex((topic) => topic.id === payload.id);