@lobehub/chat 1.11.2 → 1.11.3

package/src/libs/langchain/loaders/code/__tests__/long.txt

@@ -0,0 +1,885 @@
+import { copyToClipboard } from '@lobehub/ui';
+import isEqual from 'fast-deep-equal';
+import { produce } from 'immer';
+import { template } from 'lodash-es';
+import { SWRResponse, mutate } from 'swr';
+import { StateCreator } from 'zustand/vanilla';
+
+import { chainAnswerWithContext } from '@/chains/answerWithContext';
+import { LOADING_FLAT } from '@/const/message';
+import { TraceEventType, TraceNameMap } from '@/const/trace';
+import { isServerMode } from '@/const/version';
+import { useClientDataSWR } from '@/libs/swr';
+import { chatService } from '@/services/chat';
+import { messageService } from '@/services/message';
+import { topicService } from '@/services/topic';
+import { traceService } from '@/services/trace';
+import { useAgentStore } from '@/store/agent';
+import { agentSelectors } from '@/store/agent/selectors';
+import { chatHelpers } from '@/store/chat/helpers';
+import { messageMapKey } from '@/store/chat/slices/message/utils';
+import { ChatStore } from '@/store/chat/store';
+import { useSessionStore } from '@/store/session';
+import { UploadFileItem } from '@/types/files/upload';
+import {
+  ChatMessage,
+  ChatMessageError,
+  CreateMessageParams,
+  MessageToolCall,
+} from '@/types/message';
+import { TraceEventPayloads } from '@/types/trace';
+import { setNamespace } from '@/utils/storeDebug';
+import { nanoid } from '@/utils/uuid';
+
+import type { ChatStoreState } from '../../initialState';
+import { chatSelectors, topicSelectors } from '../../selectors';
+import { preventLeavingFn, toggleBooleanList } from '../../utils';
+import { ChatRAGAction, chatRag } from './actions/rag';
+import { MessageDispatch, messagesReducer } from './reducer';
+
+const n = setNamespace('m');
+
+const SWR_USE_FETCH_MESSAGES = 'SWR_USE_FETCH_MESSAGES';
+
+export interface SendMessageParams {
+  message: string;
+  files?: UploadFileItem[];
+  onlyAddUserMessage?: boolean;
+  /**
+   *
+   * https://github.com/lobehub/lobe-chat/pull/2086
+   */
+  isWelcomeQuestion?: boolean;
+}
+
+interface ProcessMessageParams {
+  traceId?: string;
+  isWelcomeQuestion?: boolean;
+  /**
+   * the RAG query content, should be embedding and used in the semantic search
+   */
+  ragQuery?: string;
+}
+
+export interface ChatMessageAction extends ChatRAGAction {
+  // create
+  sendMessage: (params: SendMessageParams) => Promise<void>;
+  addAIMessage: () => Promise<void>;
+  /**
+   * regenerate message
+   * trace enabled
+   * @param id
+   */
+  regenerateMessage: (id: string) => Promise<void>;
+
+  // delete
+  /**
+   * clear message on the active session
+   */
+  clearMessage: () => Promise<void>;
+  deleteMessage: (id: string) => Promise<void>;
+  deleteToolMessage: (id: string) => Promise<void>;
+  delAndRegenerateMessage: (id: string) => Promise<void>;
+  clearAllMessages: () => Promise<void>;
+  // update
+  updateInputMessage: (message: string) => void;
+  modifyMessageContent: (id: string, content: string) => Promise<void>;
+  // query
+  useFetchMessages: (sessionId: string, topicId?: string) => SWRResponse<ChatMessage[]>;
+  stopGenerateMessage: () => void;
+  copyMessage: (id: string, content: string) => Promise<void>;
+  refreshMessages: () => Promise<void>;
+  toggleMessageEditing: (id: string, editing: boolean) => void;
+
+  // =========  ↓ Internal Method ↓  ========== //
+  // ========================================== //
+  // ========================================== //
+
+  /**
+   * update message at the frontend point
+   * this method will not update messages to database
+   */
+  internal_dispatchMessage: (payload: MessageDispatch) => void;
+  /**
+   * core process of the AI message (include preprocess and postprocess)
+   */
+  internal_coreProcessMessage: (
+    messages: ChatMessage[],
+    parentId: string,
+    params?: ProcessMessageParams,
+  ) => Promise<void>;
+  /**
+   * the method to fetch the AI message
+   */
+  internal_fetchAIChatMessage: (
+    messages: ChatMessage[],
+    assistantMessageId: string,
+    params?: ProcessMessageParams,
+  ) => Promise<{
+    isFunctionCall: boolean;
+    traceId?: string;
+  }>;
+
+  /**
+   * update the message content with optimistic update
+   * a method used by other action
+   */
+  internal_updateMessageContent: (
+    id: string,
+    content: string,
+    toolCalls?: MessageToolCall[],
+  ) => Promise<void>;
+  /**
+   * update the message error with optimistic update
+   */
+  internal_updateMessageError: (id: string, error: ChatMessageError | null) => Promise<void>;
+  /**
+   * create a message with optimistic update
+   */
+  internal_createMessage: (
+    params: CreateMessageParams,
+    context?: { tempMessageId?: string; skipRefresh?: boolean },
+  ) => Promise<string>;
+  /**
+   * create a temp message for optimistic update
+   * otherwise the message will be too slow to show
+   */
+  internal_createTmpMessage: (params: CreateMessageParams) => string;
+  /**
+   * delete the message content with optimistic update
+   */
+  internal_deleteMessage: (id: string) => Promise<void>;
+  internal_resendMessage: (id: string, traceId?: string) => Promise<void>;
+
+  internal_fetchMessages: () => Promise<void>;
+  internal_traceMessage: (id: string, payload: TraceEventPayloads) => Promise<void>;
+
+  /**
+   * method to toggle message create loading state
+   * the AI message status is creating -> generating
+   * other message role like user and tool , only this method need to be called
+   */
+  internal_toggleMessageLoading: (loading: boolean, id: string) => void;
+  /**
+   * method to toggle ai message generating loading
+   */
+  internal_toggleChatLoading: (
+    loading: boolean,
+    id?: string,
+    action?: string,
+  ) => AbortController | undefined;
+  /**
+   * method to toggle the tool calling loading state
+   */
+  internal_toggleToolCallingStreaming: (id: string, streaming: boolean[] | undefined) => void;
+  /**
+   * helper to toggle the loading state of the array,used by these three toggleXXXLoading
+   */
+  internal_toggleLoadingArrays: (
+    key: keyof ChatStoreState,
+    loading: boolean,
+    id?: string,
+    action?: string,
+  ) => AbortController | undefined;
+}
+
+const getAgentConfig = () => agentSelectors.currentAgentConfig(useAgentStore.getState());
+const getAgentChatConfig = () => agentSelectors.currentAgentChatConfig(useAgentStore.getState());
+const hasEnabledKnowledge = () => agentSelectors.hasEnabledKnowledge(useAgentStore.getState());
+
+export const chatMessage: StateCreator<
+  ChatStore,
+  [['zustand/devtools', never]],
+  [],
+  ChatMessageAction
+> = (set, get, ...rest) => ({
+  ...chatRag(set, get, ...rest),
+
+  deleteMessage: async (id) => {
+    const message = chatSelectors.getMessageById(id)(get());
+    if (!message) return;
+
+    let ids = [message.id];
+
+    // if the message is a tool calls, then delete all the related messages
+    if (message.tools) {
+      const toolMessageIds = message.tools.flatMap((tool) => {
+        const messages = chatSelectors
+          .currentChats(get())
+          .filter((m) => m.tool_call_id === tool.id);
+
+        return messages.map((m) => m.id);
+      });
+      ids = ids.concat(toolMessageIds);
+    }
+
+    get().internal_dispatchMessage({ type: 'deleteMessages', ids });
+    await messageService.removeMessages(ids);
+    await get().refreshMessages();
+  },
+
+  deleteToolMessage: async (id) => {
+    const message = chatSelectors.getMessageById(id)(get());
+    if (!message || message.role !== 'tool') return;
+
+    const removeToolInAssistantMessage = async () => {
+      if (!message.parentId) return;
+      await get().internal_removeToolToAssistantMessage(message.parentId, message.tool_call_id);
+    };
+
+    await Promise.all([
+      // 1. remove tool message
+      get().internal_deleteMessage(id),
+      // 2. remove the tool item in the assistant tools
+      removeToolInAssistantMessage(),
+    ]);
+  },
+
+  delAndRegenerateMessage: async (id) => {
+    const traceId = chatSelectors.getTraceIdByMessageId(id)(get());
+    get().internal_resendMessage(id, traceId);
+    get().deleteMessage(id);
+
+    // trace the delete and regenerate message
+    get().internal_traceMessage(id, { eventType: TraceEventType.DeleteAndRegenerateMessage });
+  },
+  regenerateMessage: async (id: string) => {
+    const traceId = chatSelectors.getTraceIdByMessageId(id)(get());
+    await get().internal_resendMessage(id, traceId);
+
+    // trace the delete and regenerate message
+    get().internal_traceMessage(id, { eventType: TraceEventType.RegenerateMessage });
+  },
+  clearMessage: async () => {
+    const { activeId, activeTopicId, refreshMessages, refreshTopic, switchTopic } = get();
+
+    await messageService.removeMessagesByAssistant(activeId, activeTopicId);
+
+    if (activeTopicId) {
+      await topicService.removeTopic(activeTopicId);
+    }
+    await refreshTopic();
+    await refreshMessages();
+
+    // after remove topic , go back to default topic
+    switchTopic();
+  },
+  clearAllMessages: async () => {
+    const { refreshMessages } = get();
+    await messageService.removeAllMessages();
+    await refreshMessages();
+  },
+  sendMessage: async ({ message, files, onlyAddUserMessage, isWelcomeQuestion }) => {
+    const { internal_coreProcessMessage, activeTopicId, activeId } = 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;
+
+    set({ isCreatingMessage: true }, false, 'creatingMessage/start');
+
+    const newMessage: CreateMessageParams = {
+      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,
+    };
+
+    const agentConfig = getAgentChatConfig();
+
+    let tempMessageId: string | undefined = undefined;
+    let newTopicId: string | undefined = undefined;
+
+    // it should be the default topic, then
+    // if autoCreateTopic is enabled, check to whether we need to create a topic
+    if (!onlyAddUserMessage && !activeTopicId && agentConfig.enableAutoCreateTopic) {
+      // check activeTopic and then auto create topic
+      const chats = chatSelectors.currentChats(get());
+
+      // we will add two messages (user and assistant), so the finial length should +2
+      const featureLength = chats.length + 2;
+
+      // if there is no activeTopicId and the feature length is greater than the threshold
+      // then create a new topic and active it
+      if (!get().activeTopicId && featureLength >= agentConfig.autoCreateTopicThreshold) {
+        // we need to create a temp message for optimistic update
+        tempMessageId = get().internal_createTmpMessage(newMessage);
+        get().internal_toggleMessageLoading(true, tempMessageId);
+
+        const topicId = await get().createTopic();
+
+        if (topicId) {
+          newTopicId = topicId;
+          newMessage.topicId = topicId;
+
+          // we need to copy the messages to the new topic or the message will disappear
+          const mapKey = chatSelectors.currentChatKey(get());
+          const newMaps = {
+            ...get().messagesMap,
+            [messageMapKey(activeId, topicId)]: get().messagesMap[mapKey],
+          };
+          set({ messagesMap: newMaps }, false, 'internal_copyMessages');
+
+          // get().internal_dispatchMessage({ type: 'deleteMessage', id: tempMessageId });
+          get().internal_toggleMessageLoading(false, tempMessageId);
+
+          // make the topic loading
+          get().internal_updateTopicLoading(topicId, true);
+        }
+      }
+    }
+    //  update assistant update to make it rerank
+    useSessionStore.getState().triggerSessionUpdate(get().activeId);
+
+    const id = await get().internal_createMessage(newMessage, {
+      tempMessageId,
+      skipRefresh: !onlyAddUserMessage,
+    });
+
+    // switch to the new topic if create the new topic
+    if (!!newTopicId) {
+      await get().switchTopic(newTopicId, true);
+      await get().internal_fetchMessages();
+
+      // delete previous messages
+      // remove the temp message map
+      const newMaps = { ...get().messagesMap, [messageMapKey(activeId, null)]: [] };
+      set({ messagesMap: newMaps }, false, 'internal_copyMessages');
+    }
+
+    // if only add user message, then stop
+    if (onlyAddUserMessage) {
+      set({ isCreatingMessage: false }, false, 'creatingMessage/start');
+      return;
+    }
+
+    // Get the current messages to generate AI response
+    const messages = chatSelectors.currentChats(get());
+    const userFiles = chatSelectors.currentUserFiles(get()).map((f) => f.id);
+
+    await internal_coreProcessMessage(messages, id, {
+      isWelcomeQuestion,
+      // if there is relative files or enabled knowledge, try with ragQuery
+      ragQuery: hasEnabledKnowledge() || userFiles.length > 0 ? message : undefined,
+    });
+
+    set({ isCreatingMessage: false }, false, 'creatingMessage/stop');
+
+    const summaryTitle = async () => {
+      // if autoCreateTopic is false, then stop
+      if (!agentConfig.enableAutoCreateTopic) return;
+
+      // check activeTopic and then auto update topic title
+      if (newTopicId) {
+        const chats = chatSelectors.currentChats(get());
+        await get().summaryTopicTitle(newTopicId, chats);
+        return;
+      }
+
+      const topic = topicSelectors.currentActiveTopic(get());
+
+      if (topic && !topic.title) {
+        const chats = chatSelectors.currentChats(get());
+        await get().summaryTopicTitle(topic.id, chats);
+      }
+    };
+
+    // if there is relative files, then add files to agent
+    // only available in server mode
+    const addFilesToAgent = async () => {
+      if (userFiles.length === 0 || !isServerMode) return;
+
+      await useAgentStore.getState().addFilesToAgent(userFiles, false);
+    };
+
+    await Promise.all([summaryTitle(), addFilesToAgent()]);
+  },
+  addAIMessage: async () => {
+    const { internal_createMessage, updateInputMessage, activeTopicId, activeId, inputMessage } =
+      get();
+    if (!activeId) return;
+
+    await internal_createMessage({
+      content: inputMessage,
+      role: 'assistant',
+      sessionId: activeId,
+      // if there is activeTopicId，then add topicId to message
+      topicId: activeTopicId,
+    });
+
+    updateInputMessage('');
+  },
+  copyMessage: async (id, content) => {
+    await copyToClipboard(content);
+
+    get().internal_traceMessage(id, { eventType: TraceEventType.CopyMessage });
+  },
+  toggleMessageEditing: (id, editing) => {
+    set(
+      { messageEditingIds: toggleBooleanList(get().messageEditingIds, id, editing) },
+      false,
+      'toggleMessageEditing',
+    );
+  },
+  stopGenerateMessage: () => {
+    const { abortController, internal_toggleChatLoading } = get();
+    if (!abortController) return;
+
+    abortController.abort('canceled');
+
+    internal_toggleChatLoading(false, undefined, n('stopGenerateMessage') as string);
+  },
+
+  updateInputMessage: (message) => {
+    if (isEqual(message, get().inputMessage)) return;
+
+    set({ inputMessage: message }, false, n('updateInputMessage', message));
+  },
+  modifyMessageContent: async (id, content) => {
+    // tracing the diff of update
+    // due to message content will change, so we need send trace before update,or will get wrong data
+    get().internal_traceMessage(id, {
+      eventType: TraceEventType.ModifyMessage,
+      nextContent: content,
+    });
+
+    await get().internal_updateMessageContent(id, content);
+  },
+  useFetchMessages: (sessionId, activeTopicId) =>
+    useClientDataSWR<ChatMessage[]>(
+      [SWR_USE_FETCH_MESSAGES, sessionId, activeTopicId],
+      async ([, sessionId, topicId]: [string, string, string | undefined]) =>
+        messageService.getMessages(sessionId, topicId),
+      {
+        onSuccess: (messages, key) => {
+          const nextMap = {
+            ...get().messagesMap,
+            [messageMapKey(sessionId, activeTopicId)]: messages,
+          };
+          // no need to update map if the messages have been init and the map is the same
+          if (get().messagesInit && isEqual(nextMap, get().messagesMap)) return;
+
+          set(
+            { messagesInit: true, messagesMap: nextMap },
+            false,
+            n('useFetchMessages', { messages, queryKey: key }),
+          );
+        },
+      },
+    ),
+  refreshMessages: async () => {
+    await mutate([SWR_USE_FETCH_MESSAGES, get().activeId, get().activeTopicId]);
+  },
+
+  // the internal process method of the AI message
+  internal_coreProcessMessage: async (originalMessages, userMessageId, params) => {
+    const { internal_fetchAIChatMessage, triggerToolCalls, refreshMessages, activeTopicId } = get();
+
+    // create a new array to avoid the original messages array change
+    const messages = [...originalMessages];
+
+    const { model, provider } = getAgentConfig();
+
+    let fileChunkIds: string[] | undefined;
+
+    // go into RAG flow if there is ragQuery flag
+    if (params?.ragQuery) {
+      // 1. get the relative chunks from semantic search
+      const chunks = await get().internal_retrieveChunks(
+        userMessageId,
+        params?.ragQuery,
+        // should skip the last content
+        messages.map((m) => m.content).slice(0, messages.length - 1),
+      );
+      console.log('召回 chunks', chunks);
+
+      // 2. build the retrieve context messages
+      const retrieveContext = chainAnswerWithContext(
+        params?.ragQuery,
+        chunks.map((c) => c.text as string),
+      );
+
+      // 3. add the retrieve context messages to the messages history
+      if (retrieveContext.messages && retrieveContext.messages?.length > 0) {
+        // remove the last message due to the query is in the retrieveContext
+        messages.pop();
+        retrieveContext.messages?.forEach((m) => messages.push(m as ChatMessage));
+      }
+
+      fileChunkIds = chunks.map((c) => c.id);
+    }
+
+    // 2. Add an empty message to place the AI response
+    const assistantMessage: CreateMessageParams = {
+      role: 'assistant',
+      content: LOADING_FLAT,
+      fromModel: model,
+      fromProvider: provider,
+
+      parentId: userMessageId,
+      sessionId: get().activeId,
+      topicId: activeTopicId, // if there is activeTopicId，then add it to topicId
+      fileChunkIds,
+      ragQueryId: userMessageId,
+    };
+
+    const assistantId = await get().internal_createMessage(assistantMessage);
+
+    // 3. fetch the AI response
+    const { isFunctionCall } = await internal_fetchAIChatMessage(messages, assistantId, params);
+
+    // 4. if it's the function call message, trigger the function method
+    if (isFunctionCall) {
+      await refreshMessages();
+      await triggerToolCalls(assistantId);
+    }
+  },
+  internal_dispatchMessage: (payload) => {
+    const { activeId } = get();
+
+    if (!activeId) return;
+
+    const messages = messagesReducer(chatSelectors.currentChats(get()), payload);
+
+    const nextMap = { ...get().messagesMap, [chatSelectors.currentChatKey(get())]: messages };
+
+    if (isEqual(nextMap, get().messagesMap)) return;
+
+    set({ messagesMap: nextMap }, false, { type: `dispatchMessage/${payload.type}`, payload });
+  },
+  internal_fetchAIChatMessage: async (messages, assistantId, params) => {
+    const {
+      internal_toggleChatLoading,
+      refreshMessages,
+      internal_updateMessageContent,
+      internal_dispatchMessage,
+      internal_toggleToolCallingStreaming,
+    } = get();
+
+    const abortController = internal_toggleChatLoading(
+      true,
+      assistantId,
+      n('generateMessage(start)', { assistantId, messages }) as string,
+    );
+
+    const agentConfig = getAgentConfig();
+    const chatConfig = agentConfig.chatConfig;
+
+    const compiler = template(chatConfig.inputTemplate, { interpolate: /{{([\S\s]+?)}}/g });
+
+    // ================================== //
+    //   messages uniformly preprocess    //
+    // ================================== //
+
+    // 1. slice messages with config
+    let preprocessMsgs = chatHelpers.getSlicedMessagesWithConfig(messages, chatConfig);
+
+    // 2. replace inputMessage template
+    preprocessMsgs = !chatConfig.inputTemplate
+      ? preprocessMsgs
+      : preprocessMsgs.map((m) => {
+          if (m.role === 'user') {
+            try {
+              return { ...m, content: compiler({ text: m.content }) };
+            } catch (error) {
+              console.error(error);
+
+              return m;
+            }
+          }
+
+          return m;
+        });
+
+    // 3. add systemRole
+    if (agentConfig.systemRole) {
+      preprocessMsgs.unshift({ content: agentConfig.systemRole, role: 'system' } as ChatMessage);
+    }
+
+    // 4. handle max_tokens
+    agentConfig.params.max_tokens = chatConfig.enableMaxTokens
+      ? agentConfig.params.max_tokens
+      : undefined;
+
+    // 5. handle config for the vision model
+    // Due to the gpt-4-vision-preview model's default max_tokens is very small
+    // we need to set the max_tokens a larger one.
+    if (agentConfig.model === 'gpt-4-vision-preview') {
+      /* eslint-disable unicorn/no-lonely-if */
+      if (!agentConfig.params.max_tokens)
+        // refs: https://github.com/lobehub/lobe-chat/issues/837
+        agentConfig.params.max_tokens = 2048;
+    }
+
+    let isFunctionCall = false;
+    let msgTraceId: string | undefined;
+    let output = '';
+
+    await chatService.createAssistantMessageStream({
+      abortController,
+      params: {
+        messages: preprocessMsgs,
+        model: agentConfig.model,
+        provider: agentConfig.provider,
+        ...agentConfig.params,
+        plugins: agentConfig.plugins,
+      },
+      trace: {
+        traceId: params?.traceId,
+        sessionId: get().activeId,
+        topicId: get().activeTopicId,
+        traceName: TraceNameMap.Conversation,
+      },
+      isWelcomeQuestion: params?.isWelcomeQuestion,
+      onErrorHandle: async (error) => {
+        await messageService.updateMessageError(assistantId, error);
+        await refreshMessages();
+      },
+      onFinish: async (content, { traceId, observationId, toolCalls }) => {
+        // if there is traceId, update it
+        if (traceId) {
+          msgTraceId = traceId;
+          await messageService.updateMessage(assistantId, {
+            traceId,
+            observationId: observationId ?? undefined,
+          });
+        }
+
+        if (toolCalls && toolCalls.length > 0) {
+          internal_toggleToolCallingStreaming(assistantId, undefined);
+        }
+
+        // update the content after fetch result
+        await internal_updateMessageContent(assistantId, content, toolCalls);
+      },
+      onMessageHandle: async (chunk) => {
+        switch (chunk.type) {
+          case 'text': {
+            output += chunk.text;
+            internal_dispatchMessage({
+              id: assistantId,
+              type: 'updateMessage',
+              value: { content: output },
+            });
+            break;
+          }
+
+          // is this message is just a tool call
+          case 'tool_calls': {
+            internal_toggleToolCallingStreaming(assistantId, chunk.isAnimationActives);
+            internal_dispatchMessage({
+              id: assistantId,
+              type: 'updateMessage',
+              value: { tools: get().internal_transformToolCalls(chunk.tool_calls) },
+            });
+            isFunctionCall = true;
+          }
+        }
+      },
+    });
+
+    internal_toggleChatLoading(false, assistantId, n('generateMessage(end)') as string);
+
+    return {
+      isFunctionCall,
+      traceId: msgTraceId,
+    };
+  },
+
+  internal_resendMessage: async (messageId, traceId) => {
+    // 1. 构造所有相关的历史记录
+    const chats = chatSelectors.currentChats(get());
+
+    const currentIndex = chats.findIndex((c) => c.id === messageId);
+    if (currentIndex < 0) return;
+
+    const currentMessage = chats[currentIndex];
+
+    let contextMessages: ChatMessage[] = [];
+
+    switch (currentMessage.role) {
+      case 'tool':
+      case 'user': {
+        contextMessages = chats.slice(0, currentIndex + 1);
+        break;
+      }
+      case 'assistant': {
+        // 消息是 AI 发出的因此需要找到它的 user 消息
+        const userId = currentMessage.parentId;
+        const userIndex = chats.findIndex((c) => c.id === userId);
+        // 如果消息没有 parentId，那么同 user/function 模式
+        contextMessages = chats.slice(0, userIndex < 0 ? currentIndex + 1 : userIndex + 1);
+        break;
+      }
+    }
+
+    if (contextMessages.length <= 0) return;
+
+    const { internal_coreProcessMessage } = get();
+
+    const latestMsg = contextMessages.findLast((s) => s.role === 'user');
+
+    if (!latestMsg) return;
+
+    await internal_coreProcessMessage(contextMessages, latestMsg.id, {
+      traceId,
+      ragQuery: currentMessage.content,
+    });
+  },
+
+  internal_updateMessageError: async (id, error) => {
+    get().internal_dispatchMessage({ id, type: 'updateMessage', value: { error } });
+    await messageService.updateMessage(id, { error });
+    await get().refreshMessages();
+  },
+  internal_updateMessageContent: async (id, content, toolCalls) => {
+    const { internal_dispatchMessage, refreshMessages, internal_transformToolCalls } = get();
+
+    // Due to the async update method and refresh need about 100ms
+    // we need to update the message content at the frontend to avoid the update flick
+    // refs: https://medium.com/@kyledeguzmanx/what-are-optimistic-updates-483662c3e171
+    if (toolCalls) {
+      internal_dispatchMessage({
+        id,
+        type: 'updateMessage',
+        value: { tools: internal_transformToolCalls(toolCalls) },
+      });
+    } else {
+      internal_dispatchMessage({ id, type: 'updateMessage', value: { content } });
+    }
+
+    await messageService.updateMessage(id, {
+      content,
+      tools: toolCalls ? internal_transformToolCalls(toolCalls) : undefined,
+    });
+    await refreshMessages();
+  },
+
+  internal_createMessage: async (message, context) => {
+    const { internal_createTmpMessage, refreshMessages, internal_toggleMessageLoading } = get();
+    let tempId = context?.tempMessageId;
+    if (!tempId) {
+      // use optimistic update to avoid the slow waiting
+      tempId = internal_createTmpMessage(message);
+
+      internal_toggleMessageLoading(true, tempId);
+    }
+
+    const id = await messageService.createMessage(message);
+    if (!context?.skipRefresh) {
+      await refreshMessages();
+    }
+
+    internal_toggleMessageLoading(false, tempId);
+    return id;
+  },
+
+  internal_fetchMessages: async () => {
+    const messages = await messageService.getMessages(get().activeId, get().activeTopicId);
+    const nextMap = { ...get().messagesMap, [chatSelectors.currentChatKey(get())]: messages };
+    // no need to update map if the messages have been init and the map is the same
+    if (get().messagesInit && isEqual(nextMap, get().messagesMap)) return;
+
+    set(
+      { messagesInit: true, messagesMap: nextMap },
+      false,
+      n('internal_fetchMessages', { messages }),
+    );
+  },
+  internal_createTmpMessage: (message) => {
+    const { internal_dispatchMessage } = get();
+
+    // use optimistic update to avoid the slow waiting
+    const tempId = 'tmp_' + nanoid();
+    internal_dispatchMessage({ type: 'createMessage', id: tempId, value: message });
+
+    return tempId;
+  },
+  internal_deleteMessage: async (id: string) => {
+    get().internal_dispatchMessage({ type: 'deleteMessage', id });
+    await messageService.removeMessage(id);
+    await get().refreshMessages();
+  },
+  internal_traceMessage: async (id, payload) => {
+    // tracing the diff of update
+    const message = chatSelectors.getMessageById(id)(get());
+    if (!message) return;
+
+    const traceId = message?.traceId;
+    const observationId = message?.observationId;
+
+    if (traceId && message?.role === 'assistant') {
+      traceService
+        .traceEvent({ traceId, observationId, content: message.content, ...payload })
+        .catch();
+    }
+  },
+
+  // ----- Loading ------- //
+  internal_toggleMessageLoading: (loading, id) => {
+    set(
+      {
+        messageLoadingIds: toggleBooleanList(get().messageLoadingIds, id, loading),
+      },
+      false,
+      'internal_toggleMessageLoading',
+    );
+  },
+  internal_toggleChatLoading: (loading, id, action) => {
+    return get().internal_toggleLoadingArrays('chatLoadingIds', loading, id, action);
+  },
+  internal_toggleToolCallingStreaming: (id, streaming) => {
+    set(
+      {
+        toolCallingStreamIds: produce(get().toolCallingStreamIds, (draft) => {
+          if (!!streaming) {
+            draft[id] = streaming;
+          } else {
+            delete draft[id];
+          }
+        }),
+      },
+
+      false,
+      'toggleToolCallingStreaming',
+    );
+  },
+  internal_toggleLoadingArrays: (key, loading, id, action) => {
+    if (loading) {
+      window.addEventListener('beforeunload', preventLeavingFn);
+
+      const abortController = new AbortController();
+      set(
+        {
+          abortController,
+          [key]: toggleBooleanList(get()[key] as string[], id!, loading),
+        },
+        false,
+        action,
+      );
+
+      return abortController;
+    } else {
+      if (!id) {
+        set({ abortController: undefined, [key]: [] }, false, action);
+      } else
+        set(
+          {
+            abortController: undefined,
+            [key]: toggleBooleanList(get()[key] as string[], id, loading),
+          },
+          false,
+          action,
+        );
+
+      window.removeEventListener('beforeunload', preventLeavingFn);
+    }
+  },
+});