@lobehub/lobehub 2.0.0-next.35 → 2.0.0-next.36

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

@@ -0,0 +1,157 @@
+/* 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 { MESSAGE_CANCEL_FLAT } from '@lobechat/const';
+import { produce } from 'immer';
+import { StateCreator } from 'zustand/vanilla';
+
+import { ChatStore } from '@/store/chat/store';
+import { setNamespace } from '@/utils/storeDebug';
+
+import { messageMapKey } from '../../../utils/messageMapKey';
+import { MainSendMessageOperation } from '../initialState';
+
+const n = setNamespace('ai');
+
+/**
+ * Actions for controlling conversation operations like cancellation and error handling
+ */
+export interface ConversationControlAction {
+  /**
+   * Interrupts the ongoing ai message generation process
+   */
+  stopGenerateMessage: () => void;
+  /**
+   * Cancels sendMessage operation for a specific topic/session
+   */
+  cancelSendMessageInServer: (topicId?: string) => void;
+  /**
+   * Clears any error messages from the send message operation
+   */
+  clearSendMessageError: () => void;
+  /**
+   * Switches to a different branch of a message
+   */
+  switchMessageBranch: (messageId: string, branchIndex: number) => Promise<void>;
+  /**
+   * Toggle sendMessage operation state
+   */
+  internal_toggleSendMessageOperation: (
+    key: string | { sessionId: string; topicId?: string | null },
+    loading: boolean,
+    cancelReason?: string,
+  ) => AbortController | undefined;
+  /**
+   * Update sendMessage operation metadata
+   */
+  internal_updateSendMessageOperation: (
+    key: string | { sessionId: string; topicId?: string | null },
+    value: Partial<MainSendMessageOperation> | null,
+    actionName?: any,
+  ) => void;
+}
+
+export const conversationControl: StateCreator<
+  ChatStore,
+  [['zustand/devtools', never]],
+  [],
+  ConversationControlAction
+> = (set, get) => ({
+  stopGenerateMessage: () => {
+    const { chatLoadingIdsAbortController, internal_toggleChatLoading } = get();
+
+    if (!chatLoadingIdsAbortController) return;
+
+    chatLoadingIdsAbortController.abort(MESSAGE_CANCEL_FLAT);
+
+    internal_toggleChatLoading(false, undefined, n('stopGenerateMessage') as string);
+  },
+
+  cancelSendMessageInServer: (topicId?: string) => {
+    const { activeId, activeTopicId } = get();
+
+    // Determine which operation to cancel
+    const targetTopicId = topicId ?? activeTopicId;
+    const operationKey = messageMapKey(activeId, targetTopicId);
+
+    // Cancel the specific operation
+    get().internal_toggleSendMessageOperation(
+      operationKey,
+      false,
+      'User cancelled sendMessage operation',
+    );
+
+    // Only clear creating message state if it's the active session
+    if (operationKey === messageMapKey(activeId, activeTopicId)) {
+      const editorTempState = get().mainSendMessageOperations[operationKey]?.inputEditorTempState;
+
+      if (editorTempState) get().mainInputEditor?.setJSONState(editorTempState);
+    }
+  },
+
+  clearSendMessageError: () => {
+    get().internal_updateSendMessageOperation(
+      { sessionId: get().activeId, topicId: get().activeTopicId },
+      null,
+      'clearSendMessageError',
+    );
+  },
+
+  switchMessageBranch: async (messageId, branchIndex) => {
+    await get().optimisticUpdateMessageMetadata(messageId, { activeBranchIndex: branchIndex });
+  },
+
+  internal_toggleSendMessageOperation: (key, loading: boolean, cancelReason?: string) => {
+    if (loading) {
+      const abortController = new AbortController();
+
+      get().internal_updateSendMessageOperation(
+        key,
+        { isLoading: true, abortController },
+        n('toggleSendMessageOperation(start)', { key }),
+      );
+
+      return abortController;
+    } else {
+      const operationKey =
+        typeof key === 'string' ? key : messageMapKey(key.sessionId, key.topicId);
+
+      const operation = get().mainSendMessageOperations[operationKey];
+
+      // If cancelReason is provided, abort the operation first
+      if (cancelReason && operation?.isLoading) {
+        operation.abortController?.abort(cancelReason);
+      }
+
+      get().internal_updateSendMessageOperation(
+        key,
+        { isLoading: false, abortController: null },
+        n('toggleSendMessageOperation(stop)', { key, cancelReason }),
+      );
+
+      return undefined;
+    }
+  },
+
+  internal_updateSendMessageOperation: (key, value, actionName) => {
+    const operationKey = typeof key === 'string' ? key : messageMapKey(key.sessionId, key.topicId);
+
+    set(
+      produce((draft) => {
+        if (!draft.mainSendMessageOperations[operationKey])
+          draft.mainSendMessageOperations[operationKey] = value;
+        else {
+          if (value === null) {
+            delete draft.mainSendMessageOperations[operationKey];
+          } else {
+            draft.mainSendMessageOperations[operationKey] = {
+              ...draft.mainSendMessageOperations[operationKey],
+              ...value,
+            };
+          }
+        }
+      }),
+      false,
+      actionName ?? n('updateSendMessageOperation', { operationKey, value }),
+    );
+  },
+});

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

@@ -0,0 +1,329 @@
+/* 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 { DEFAULT_AGENT_CHAT_CONFIG, INBOX_SESSION_ID } from '@lobechat/const';
+import {
+  ChatImageItem,
+  ChatVideoItem,
+  SendMessageParams,
+  SendMessageServerResponse,
+  TraceEventType,
+} from '@lobechat/types';
+import { TRPCClientError } from '@trpc/client';
+import { t } from 'i18next';
+import { StateCreator } from 'zustand/vanilla';
+
+import { aiChatService } from '@/services/aiChat';
+import { getAgentStoreState } from '@/store/agent';
+import { agentChatConfigSelectors, agentSelectors } from '@/store/agent/slices/chat';
+import { ChatStore } from '@/store/chat/store';
+import { getFileStoreState } from '@/store/file/store';
+import { getSessionStoreState } from '@/store/session';
+
+import {
+  dbMessageSelectors,
+  displayMessageSelectors,
+  messageStateSelectors,
+  topicSelectors,
+} from '../../../selectors';
+import { messageMapKey } from '../../../utils/messageMapKey';
+
+/**
+ * Actions managing the complete lifecycle of conversations including sending,
+ * regenerating, and resending messages
+ */
+export interface ConversationLifecycleAction {
+  /**
+   * Sends a new message to the AI chat system
+   */
+  sendMessage: (params: SendMessageParams) => Promise<void>;
+  regenerateUserMessage: (
+    id: string,
+    params?: { skipTrace?: boolean; traceId?: string },
+  ) => Promise<void>;
+  regenerateAssistantMessage: (
+    id: string,
+    params?: { skipTrace?: boolean; traceId?: string },
+  ) => Promise<void>;
+  /**
+   * Deletes an existing message and generates a new one in its place
+   */
+  delAndRegenerateMessage: (id: string) => Promise<void>;
+}
+
+export const conversationLifecycle: StateCreator<
+  ChatStore,
+  [['zustand/devtools', never]],
+  [],
+  ConversationLifecycleAction
+> = (set, get) => ({
+  sendMessage: async ({ message, files, onlyAddUserMessage }) => {
+    const { activeTopicId, activeId, activeThreadId, internal_execAgentRuntime, mainInputEditor } =
+      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 = displayMessageSelectors.activeDisplayMessages(get());
+    const parentId = displayMessageSelectors.lastDisplayMessageId(get());
+
+    const chatConfig = agentChatConfigSelectors.currentChatConfig(getAgentStoreState());
+    const autoCreateThreshold =
+      chatConfig.autoCreateTopicThreshold ?? DEFAULT_AGENT_CHAT_CONFIG.autoCreateTopicThreshold;
+    const shouldCreateNewTopic =
+      !activeTopicId &&
+      !!chatConfig.enableAutoCreateTopic &&
+      messages.length + 2 >= autoCreateThreshold;
+
+    // 构造服务端模式临时消息的本地媒体预览（优先使用 S3 URL）
+    const filesInStore = getFileStoreState().chatUploadFileList;
+    const tempImages: ChatImageItem[] = filesInStore
+      .filter((f) => f.file?.type?.startsWith('image'))
+      .map((f) => ({
+        id: f.id,
+        url: f.fileUrl || f.base64Url || f.previewUrl || '',
+        alt: f.file?.name || f.id,
+      }));
+    const tempVideos: ChatVideoItem[] = filesInStore
+      .filter((f) => f.file?.type?.startsWith('video'))
+      .map((f) => ({
+        id: f.id,
+        url: f.fileUrl || f.base64Url || f.previewUrl || '',
+        alt: f.file?.name || f.id,
+      }));
+
+    // use optimistic update to avoid the slow waiting
+    const tempId = get().optimisticCreateTmpMessage({
+      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,
+      imageList: tempImages.length > 0 ? tempImages : undefined,
+      videoList: tempVideos.length > 0 ? tempVideos : undefined,
+    });
+    get().internal_toggleMessageLoading(true, tempId);
+
+    const operationKey = messageMapKey(activeId, activeTopicId);
+
+    // Start tracking sendMessage operation with AbortController
+    const abortController = get().internal_toggleSendMessageOperation(operationKey, true)!;
+
+    const jsonState = mainInputEditor?.getJSONState();
+    get().internal_updateSendMessageOperation(
+      operationKey,
+      { inputSendErrorMsg: undefined, inputEditorTempState: jsonState },
+      'creatingMessage/start',
+    );
+
+    let data: SendMessageServerResponse | undefined;
+    try {
+      const { model, provider } = agentSelectors.currentAgentConfig(getAgentStoreState());
+
+      data = await aiChatService.sendMessageInServer(
+        {
+          newUserMessage: { content: message, files: fileIdList, parentId },
+          // if there is activeTopicId，then add topicId to message
+          topicId: activeTopicId,
+          threadId: activeThreadId,
+          newTopic: shouldCreateNewTopic
+            ? {
+                topicMessageIds: messages.map((m) => m.id),
+                title: t('defaultTitle', { ns: 'topic' }),
+              }
+            : undefined,
+          sessionId: activeId === INBOX_SESSION_ID ? undefined : activeId,
+          newAssistantMessage: { model, provider: provider! },
+        },
+        abortController,
+      );
+      let topicId = activeTopicId;
+      // refresh the total data
+      if (data?.topics) {
+        get().internal_dispatchTopic({ type: 'updateTopics', value: data.topics });
+        topicId = data.topicId;
+      }
+
+      get().replaceMessages(data.messages, { sessionId: activeId, topicId: topicId });
+
+      if (data.isCreateNewTopic && data.topicId) {
+        await get().switchTopic(data.topicId, true);
+      }
+    } catch (e) {
+      if (e instanceof TRPCClientError) {
+        const isAbort = e.message.includes('aborted') || e.name === 'AbortError';
+        // Check if error is due to cancellation
+        if (!isAbort) {
+          get().internal_updateSendMessageOperation(operationKey, { inputSendErrorMsg: e.message });
+          get().mainInputEditor?.setJSONState(jsonState);
+        }
+      }
+    } finally {
+      // Stop tracking sendMessage operation
+      get().internal_toggleSendMessageOperation(operationKey, false);
+    }
+
+    // remove temporally message
+    if (data?.isCreateNewTopic) {
+      get().internal_dispatchMessage(
+        { type: 'deleteMessage', id: tempId },
+        { topicId: activeTopicId, sessionId: activeId },
+      );
+    }
+
+    get().internal_toggleMessageLoading(false, tempId);
+    get().internal_updateSendMessageOperation(
+      operationKey,
+      { inputEditorTempState: null },
+      'creatingMessage/finished',
+    );
+
+    if (!data) return;
+
+    //  update assistant update to make it rerank
+    getSessionStoreState().triggerSessionUpdate(get().activeId);
+
+    if (data.topicId) get().internal_updateTopicLoading(data.topicId, true);
+
+    const summaryTitle = async () => {
+      // check activeTopic and then auto update topic title
+      if (data.isCreateNewTopic) {
+        await get().summaryTopicTitle(data.topicId, data.messages);
+        return;
+      }
+
+      if (!data.topicId) return;
+
+      const topic = topicSelectors.getTopicById(data.topicId)(get());
+
+      if (topic && !topic.title) {
+        const chats = displayMessageSelectors
+          .getDisplayMessagesByKey(messageMapKey(activeId, topic.id))(get())
+          .filter((item) => item.id !== data.assistantMessageId);
+
+        await get().summaryTopicTitle(topic.id, chats);
+      }
+    };
+
+    summaryTitle().catch(console.error);
+
+    // Get the current messages to generate AI response
+    const displayMessages = displayMessageSelectors.activeDisplayMessages(get());
+
+    try {
+      await internal_execAgentRuntime({
+        messages: displayMessages,
+        parentMessageId: data.assistantMessageId,
+        parentMessageType: 'assistant',
+        ragQuery: get().internal_shouldUseRAG() ? message : undefined,
+        threadId: activeThreadId,
+      });
+
+      //
+      // // if there is relative files, then add files to agent
+      // // only available in server mode
+      const userFiles = dbMessageSelectors
+        .dbUserFiles(get())
+        .map((f) => f?.id)
+        .filter(Boolean) as string[];
+
+      await getAgentStoreState().addFilesToAgent(userFiles, false);
+    } catch (e) {
+      console.error(e);
+    } finally {
+      if (data.topicId) get().internal_updateTopicLoading(data.topicId, false);
+    }
+  },
+
+  regenerateUserMessage: async (id, params) => {
+    const isRegenerating = messageStateSelectors.isMessageRegenerating(id)(get());
+    if (isRegenerating) return;
+
+    const item = displayMessageSelectors.getDisplayMessageById(id)(get());
+    if (!item) return;
+
+    const chats = displayMessageSelectors.mainAIChats(get());
+
+    const currentIndex = chats.findIndex((c) => c.id === id);
+    const contextMessages = chats.slice(0, currentIndex + 1);
+
+    if (contextMessages.length <= 0) return;
+
+    try {
+      const { internal_execAgentRuntime, activeThreadId } = get();
+
+      // Mark message as regenerating
+      set(
+        { regeneratingIds: [...get().regeneratingIds, id] },
+        false,
+        'regenerateUserMessage/start',
+      );
+
+      const traceId = params?.traceId ?? dbMessageSelectors.getTraceIdByDbMessageId(id)(get());
+
+      // 切一个新的激活分支
+      await get().switchMessageBranch(id, item.branch ? item.branch.count : 1);
+
+      await internal_execAgentRuntime({
+        messages: contextMessages,
+        parentMessageId: id,
+        parentMessageType: 'user',
+        traceId,
+        ragQuery: get().internal_shouldUseRAG() ? item.content : undefined,
+        threadId: activeThreadId,
+      });
+
+      // trace the regenerate message
+      if (!params?.skipTrace)
+        get().internal_traceMessage(id, { eventType: TraceEventType.RegenerateMessage });
+    } finally {
+      // Remove message from regenerating state
+      set(
+        { regeneratingIds: get().regeneratingIds.filter((msgId) => msgId !== id) },
+        false,
+        'regenerateUserMessage/end',
+      );
+    }
+  },
+
+  regenerateAssistantMessage: async (id, params) => {
+    const isRegenerating = messageStateSelectors.isMessageRegenerating(id)(get());
+    if (isRegenerating) return;
+
+    const chats = displayMessageSelectors.mainAIChats(get());
+    const currentIndex = chats.findIndex((c) => c.id === id);
+    const currentMessage = chats[currentIndex];
+
+    // 消息是 AI 发出的因此需要找到它的 user 消息
+    const userId = currentMessage.parentId;
+    const userIndex = chats.findIndex((c) => c.id === userId);
+    // 如果消息没有 parentId，那么同 user 模式
+    const contextMessages = chats.slice(0, userIndex < 0 ? currentIndex + 1 : userIndex + 1);
+
+    if (contextMessages.length <= 0 || !userId) return;
+
+    await get().regenerateUserMessage(userId, params);
+  },
+
+  delAndRegenerateMessage: async (id) => {
+    const traceId = dbMessageSelectors.getTraceIdByDbMessageId(id)(get());
+    get().regenerateAssistantMessage(id, { skipTrace: true, traceId });
+    get().deleteMessage(id);
+
+    // trace the delete and regenerate message
+    get().internal_traceMessage(id, { eventType: TraceEventType.DeleteAndRegenerateMessage });
+  },
+});

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

@@ -271,7 +271,7 @@ export const chatAiGroupChat: StateCreator<
   return {
     sendGroupMessage: async ({ groupId, message, files, onlyAddUserMessage, targetMemberId }) => {
       const {
-        internal_createMessage,
+        optimisticCreateMessage,
         internal_triggerSupervisorDecisionDebounced,
         internal_setActiveGroup,
         activeTopicId,
@@ -294,7 +294,7 @@ export const chatAiGroupChat: StateCreator<
           targetId: targetMemberId,
         };
 
-        const result = await internal_createMessage(userMessage);
+        const result = await optimisticCreateMessage(userMessage);
 
         // if only add user message, then stop
         if (onlyAddUserMessage) {
@@ -379,7 +379,7 @@ export const chatAiGroupChat: StateCreator<
       const {
         messagesMap,
         internal_toggleSupervisorLoading,
-        internal_createMessage,
+        optimisticCreateMessage,
         supervisorTodos,
       } = get();
 
@@ -398,8 +398,8 @@ export const chatAiGroupChat: StateCreator<
         const content = formatSupervisorTodoContent(todoList);
         const supervisorMessage: CreateMessageParams = {
           content,
-          fromModel: groupConfig.orchestratorModel,
-          fromProvider: groupConfig.orchestratorProvider,
+          model: groupConfig.orchestratorModel,
+          provider: groupConfig.orchestratorProvider,
           groupId,
           role: 'supervisor',
           sessionId,
@@ -408,7 +408,7 @@ export const chatAiGroupChat: StateCreator<
 
         console.log('Creating supervisor todo message:', supervisorMessage);
 
-        await internal_createMessage(supervisorMessage);
+        await optimisticCreateMessage(supervisorMessage);
       };
 
       const messages = messagesMap[messageMapKey(groupId, currentTopicId)] || [];
@@ -597,7 +597,7 @@ export const chatAiGroupChat: StateCreator<
       });
       const {
         messagesMap,
-        internal_createMessage,
+        optimisticCreateMessage,
         internal_fetchAIChatMessage,
         refreshMessages,
         activeTopicId,
@@ -659,10 +659,10 @@ export const chatAiGroupChat: StateCreator<
         // Create agent message using real agent config
         const agentMessage: CreateMessageParams = {
           role: 'assistant',
-          fromModel: agentModel,
+          model: agentModel,
           groupId,
           content: LOADING_FLAT,
-          fromProvider: agentProvider,
+          provider: agentProvider,
           agentId,
           sessionId: useSessionStore.getState().activeId,
           topicId: activeTopicId,
@@ -671,7 +671,7 @@ export const chatAiGroupChat: StateCreator<
 
         console.log('DEBUG: Creating agent message with:', agentMessage);
 
-        const result = await internal_createMessage(agentMessage);
+        const result = await optimisticCreateMessage(agentMessage);
         if (!result) return;
         const assistantId = result.id;
 
@@ -936,7 +936,7 @@ export const chatAiGroupChat: StateCreator<
     },
 
     internal_createSupervisorErrorMessage: async (groupId: string, error: Error | string) => {
-      const { internal_createTmpMessage, activeTopicId } = get();
+      const { optimisticCreateTmpMessage, activeTopicId } = get();
 
       try {
         const errorMessage = error instanceof Error ? error.message : error;
@@ -944,8 +944,8 @@ export const chatAiGroupChat: StateCreator<
 
         const supervisorMessage: CreateMessageParams = {
           role: 'supervisor',
-          fromModel: groupConfig.orchestratorModel,
-          fromProvider: groupConfig.orchestratorProvider,
+          model: groupConfig.orchestratorModel,
+          provider: groupConfig.orchestratorProvider,
           groupId,
           sessionId: useSessionStore.getState().activeId || groupId,
           topicId: activeTopicId,
@@ -957,7 +957,7 @@ export const chatAiGroupChat: StateCreator<
         };
 
         // Create a temporary message that only exists in UI state, no API call
-        internal_createTmpMessage(supervisorMessage);
+        optimisticCreateTmpMessage(supervisorMessage);
       } catch (createError) {
         console.error('Failed to create supervisor error message:', createError);
       }

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

@@ -2,16 +2,20 @@ import { StateCreator } from 'zustand/vanilla';
 
 import { ChatStore } from '@/store/chat/store';
 
-import { AIGenerateAction, generateAIChat } from './generateAIChat';
-import { AIGenerateV2Action, generateAIChatV2 } from './generateAIChatV2';
+import { ConversationControlAction, conversationControl } from './conversationControl';
+import { ConversationLifecycleAction, conversationLifecycle } from './conversationLifecycle';
 import { ChatMemoryAction, chatMemory } from './memory';
 import { ChatRAGAction, chatRag } from './rag';
+import { StreamingExecutorAction, streamingExecutor } from './streamingExecutor';
+import { StreamingStatesAction, streamingStates } from './streamingStates';
 
 export interface ChatAIChatAction
   extends ChatRAGAction,
     ChatMemoryAction,
-    AIGenerateAction,
-    AIGenerateV2Action {
+    ConversationLifecycleAction,
+    ConversationControlAction,
+    StreamingExecutorAction,
+    StreamingStatesAction {
   /**/
 }
 
@@ -22,7 +26,9 @@ export const chatAiChat: StateCreator<
   ChatAIChatAction
 > = (...params) => ({
   ...chatRag(...params),
-  ...generateAIChat(...params),
   ...chatMemory(...params),
-  ...generateAIChatV2(...params),
+  ...conversationLifecycle(...params),
+  ...conversationControl(...params),
+  ...streamingExecutor(...params),
+  ...streamingStates(...params),
 });

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

@@ -6,7 +6,7 @@ import { ragService } from '@/services/rag';
 import { useAgentStore } from '@/store/agent';
 import { agentSelectors } from '@/store/agent/selectors';
 import { ChatStore } from '@/store/chat';
-import { chatSelectors } from '@/store/chat/selectors';
+import { dbMessageSelectors, displayMessageSelectors } from '@/store/chat/selectors';
 import { toggleBooleanList } from '@/store/chat/utils';
 import { useUserStore } from '@/store/user';
 import { systemAgentSelectors } from '@/store/user/selectors';
@@ -43,7 +43,7 @@ export const chatRag: StateCreator<ChatStore, [['zustand/devtools', never]], [],
   get,
 ) => ({
   deleteUserMessageRagQuery: async (id) => {
-    const message = chatSelectors.getMessageById(id)(get());
+    const message = dbMessageSelectors.getDbMessageById(id)(get());
 
     if (!message || !message.ragQueryId) return;
 
@@ -61,7 +61,7 @@ export const chatRag: StateCreator<ChatStore, [['zustand/devtools', never]], [],
   internal_retrieveChunks: async (id, userQuery, messages) => {
     get().internal_toggleMessageRAGLoading(true, id);
 
-    const message = chatSelectors.getMessageById(id)(get());
+    const message = dbMessageSelectors.getDbMessageById(id)(get());
 
     // 1. get the rewrite query
     let rewriteQuery = message?.ragQuery as string | undefined;
@@ -73,7 +73,10 @@ export const chatRag: StateCreator<ChatStore, [['zustand/devtools', never]], [],
     }
 
     // 2. retrieve chunks from semantic search
-    const files = chatSelectors.currentUserFiles(get()).map((f) => f.id);
+    const files = dbMessageSelectors
+      .dbUserFiles(get())
+      .map((f) => f?.id)
+      .filter(Boolean) as string[];
     try {
       const { chunks, queryId } = await ragService.semanticSearchForChat({
         fileIds: knowledgeIds().fileIds.concat(files),
@@ -145,13 +148,13 @@ export const chatRag: StateCreator<ChatStore, [['zustand/devtools', never]], [],
   },
 
   rewriteQuery: async (id) => {
-    const message = chatSelectors.getMessageById(id)(get());
+    const message = dbMessageSelectors.getDbMessageById(id)(get());
     if (!message) return;
 
     // delete the current ragQuery
     await get().deleteUserMessageRagQuery(id);
 
-    const chats = chatSelectors.mainAIChatsWithHistoryConfig(get());
+    const chats = displayMessageSelectors.mainAIChatsWithHistoryConfig(get());
 
     await get().internal_rewriteQuery(
       id,