Merge pull request #220 from import-ai/feat/msg_ctrl

feat(chat): Support message branch

Author: LucienShui

.env.example

@@ -7,5 +7,5 @@ VITE_API_PATH=http://127.0.0.1:8000
 
 # WebSocket Configuration
 # Set to 'true' to use WebSocket instead of Server-Sent Events (SSE)
-# Default: false (uses SSE)
-VITE_USE_WEBSOCKET=false
+# Default: true (uses WebSocket)
+VITE_USE_WEBSOCKET=true

src/components/copy/index.tsx

@@ -50,7 +50,7 @@ export default function CopyMain(props: IProps) {
         )}
       </TooltipTrigger>
       <TooltipContent>
-        <p>{t('copy.title')}</p>
+        <p>{t('chat.messages.actions.copy')}</p>
       </TooltipContent>
     </Tooltip>
   );

src/i18n/locales/en.json

@@ -317,7 +317,16 @@
     "citations_results": "Search Results",
     "searching": "Searching...",
     "disclaimer": "OmniBox can make mistakes. Check important info.",
-    "save_to_private": "Save to private space",
+    "messages": {
+      "actions": {
+        "regenerate": "Regenerate",
+        "save": "Save to private space",
+        "copy": "Copy Markdown",
+        "edit": "Edit message",
+        "save_edit": "Save",
+        "cancel_edit": "Cancel"
+      }
+    },
     "conversations": {
       "new": "New conversation",
       "new_chat": "New Chat",

src/i18n/locales/zh.json

@@ -317,7 +317,16 @@
     "citations_results": "搜索结果",
     "searching": "搜索中...",
     "disclaimer": "OmniBox 也可能会犯错。请核查重要信息。",
-    "save_to_private": "保存至个人空间",
+    "messages": {
+      "actions": {
+        "regenerate": "重新生成",
+        "save": "保存至个人空间",
+        "copy": "复制 Markdown",
+        "edit": "编辑消息",
+        "save_edit": "保存",
+        "cancel_edit": "取消"
+      }
+    },
     "conversations": {
       "new": "新对话",
       "new_chat": "新建对话",

src/page/chat/components/save.tsx

@@ -55,7 +55,7 @@ export default function SaveMain(props: IProps) {
         </Button>
       </TooltipTrigger>
       <TooltipContent>
-        <p>{t('chat.save_to_private')}</p>
+        <p>{t('chat.messages.actions.save')}</p>
       </TooltipContent>
     </Tooltip>
   );

src/page/chat/conversation/index.tsx

@@ -25,6 +25,9 @@ export default function ChatConversationPage() {
     onContextChange,
     namespaceId,
     conversation,
+    messageOperator,
+    onRegenerate,
+    onEdit,
   } = useContext();
 
   return (
@@ -40,6 +43,9 @@ export default function ChatConversationPage() {
           <Messages
             conversation={conversation}
             messages={normalizeChatData(messages)}
+            messageOperator={messageOperator}
+            onRegenerate={onRegenerate}
+            onEdit={onEdit}
           />
         )}
       </Scrollbar>

src/page/chat/conversation/message-operator.tsx

@@ -4,6 +4,7 @@ import {
   ChatBOSResponse,
   ChatDeltaResponse,
   MessageStatus,
+  OpenAIMessageRole,
 } from '@/page/chat/types/chat-response';
 import {
   ConversationDetail,
@@ -19,9 +20,48 @@ export interface MessageOperator {
   add: (chatResponse: ChatBOSResponse) => string;
   done: (id?: string) => void;
   activate: (id: string) => void;
+  getSiblings: (id: string) => string[];
+  getParent: (id: string) => string;
+}
+
+function getChildren(
+  conversation: ConversationDetail,
+  id: string,
+  targetRole: OpenAIMessageRole
+): string[] {
+  if (targetRole === OpenAIMessageRole.ASSISTANT) {
+    const currentNode = conversation.mapping[id];
+    if (currentNode) {
+      if (
+        currentNode.message.role === OpenAIMessageRole.ASSISTANT &&
+        !currentNode.message.tool_calls
+      ) {
+        return [id];
+      }
+      const targetChildren: string[] = [];
+      for (const childId of currentNode.children || []) {
+        targetChildren.push(...getChildren(conversation, childId, targetRole));
+      }
+      return targetChildren;
+    }
+  } else if (targetRole === OpenAIMessageRole.USER) {
+    const currentNode = conversation.mapping[id];
+    if (currentNode) {
+      return currentNode.children;
+    }
+    const children: string[] = [];
+    for (const node of Object.values(conversation.mapping)) {
+      if (node.parent_id === id) {
+        children.push(node.id);
+      }
+    }
+    return children;
+  }
+  return [];
 }
 
 export function createMessageOperator(
+  conversation: ConversationDetail,
   setConversation: Dispatch<SetStateAction<ConversationDetail>>
 ): MessageOperator {
   return {
@@ -82,11 +122,8 @@ export function createMessageOperator(
 
       setConversation(prev => {
         const newMapping = { ...prev.mapping, [message.id]: message };
-        let currentNode = prev.current_node;
-        if (message.parent_id === currentNode) {
-          currentNode = message.id;
-        }
-        if (message.parent_id) {
+
+        if (message.parent_id && prev.current_node !== undefined) {
           const parentMessage = prev.mapping[message.parent_id];
           if (parentMessage) {
             if (!parentMessage.children.includes(message.id)) {
@@ -101,7 +138,7 @@ export function createMessageOperator(
         return {
           ...prev,
           mapping: newMapping,
-          current_node: currentNode,
+          current_node: message.id,
         };
       });
       return chatResponse.id;
@@ -124,9 +161,50 @@ export function createMessageOperator(
       });
     },
 
+    /**
+     * Get siblings of a message.
+     * @param id
+     */
+    getSiblings: (id: string): string[] => {
+      const currentNode = conversation.mapping[id];
+      if (currentNode.message.tool_calls) {
+        return [id];
+      }
+      if (currentNode) {
+        const currentRole = currentNode.message.role;
+        if (currentNode.message.role === OpenAIMessageRole.USER) {
+          return getChildren(conversation, currentNode.parent_id, currentRole);
+        }
+        let parentNode = currentNode;
+        while (parentNode.message.role !== OpenAIMessageRole.USER) {
+          parentNode = conversation.mapping[parentNode.parent_id];
+        }
+        return getChildren(conversation, parentNode.id, currentRole);
+      }
+      return [];
+    },
+
+    /**
+     * Get non-tool parent of a message.
+     * @param id
+     */
+    getParent: (id: string): string => {
+      let currentNode = conversation.mapping[id];
+      if (currentNode) {
+        const targetRoles =
+          currentNode.message.role === OpenAIMessageRole.ASSISTANT
+            ? [OpenAIMessageRole.USER]
+            : [OpenAIMessageRole.ASSISTANT, OpenAIMessageRole.SYSTEM];
+        while (!targetRoles.includes(currentNode.message.role)) {
+          currentNode = conversation.mapping[currentNode.parent_id];
+        }
+        return currentNode.id;
+      }
+      return '';
+    },
+
     /**
      * When there is multi message, activate one of them.
-     * Designed for future, now enabled for now.
      * @param id
      */
     activate: (id: string) => {

src/page/chat/conversation/types.tsx

@@ -26,7 +26,7 @@ export interface Reasoning extends IChatTool {
   name: ToolType.REASONING;
 }
 
-type ChatTool = WebSearch | PrivateSearch | Reasoning;
+export type ChatTool = WebSearch | PrivateSearch | Reasoning;
 
 export interface ChatRequestBody {
   conversation_id: string;

src/page/chat/conversation/useContext.tsx

@@ -15,7 +15,10 @@ import {
   createMessageOperator,
   MessageOperator,
 } from '@/page/chat/conversation/message-operator';
-import { ask } from '@/page/chat/conversation/utils';
+import {
+  ask,
+  extractOriginalMessageSettings,
+} from '@/page/chat/conversation/utils';
 import {
   ConversationDetail,
   MessageDetail,
@@ -70,8 +73,8 @@ export default function useContext() {
     return result;
   }, [conversation]);
   const messageOperator = useMemo((): MessageOperator => {
-    return createMessageOperator(setConversation);
-  }, [setConversation]);
+    return createMessageOperator(conversation, setConversation);
+  }, [conversation, setConversation]);
   const onAction = async (action?: ChatActionType) => {
     if (action === 'stop') {
       isFunction(askAbortRef.current) && askAbortRef.current();
@@ -96,7 +99,7 @@ export default function useContext() {
         query,
         tools,
         context,
-        messages,
+        messages[messages.length - 1]?.id,
         messageOperator,
         `/api/v1/namespaces/${namespaceId}/wizard/${mode}`,
         getWizardLang(i18n),
@@ -127,6 +130,86 @@ export default function useContext() {
     submit(routeQuery);
   }, []);
 
+  const onRegenerate = async (messageId: string) => {
+    const parentId = messageOperator.getParent(messageId);
+    const parentMessage = conversation.mapping[parentId];
+    if (!parentMessage || !parentMessage.message.content) {
+      console.error('Cannot find parent user message to regenerate from');
+      return;
+    }
+
+    const {
+      originalTools,
+      originalContext,
+      originalLang,
+      originalEnableThinking,
+    } = extractOriginalMessageSettings(parentMessage, {
+      tools,
+      context,
+      lang: getWizardLang(i18n),
+    });
+
+    setLoading(true);
+    try {
+      const askFN = ask(
+        conversationId,
+        parentMessage.message.content,
+        originalTools,
+        originalContext,
+        parentId,
+        messageOperator,
+        `/api/v1/namespaces/${namespaceId}/wizard/${mode}`,
+        originalLang,
+        namespaceId,
+        undefined,
+        undefined,
+        originalEnableThinking
+      );
+      askAbortRef.current = askFN.destroy;
+      await askFN.start();
+    } finally {
+      setLoading(false);
+    }
+  };
+
+  const onEdit = async (messageId: string, newContent: string) => {
+    const parentId = conversation.mapping[messageId].parent_id;
+    const editedMessage = conversation.mapping[messageId];
+
+    const {
+      originalTools,
+      originalContext,
+      originalLang,
+      originalEnableThinking,
+    } = extractOriginalMessageSettings(editedMessage, {
+      tools,
+      context,
+      lang: getWizardLang(i18n),
+    });
+
+    setLoading(true);
+    try {
+      const askFN = ask(
+        conversationId,
+        newContent,
+        originalTools,
+        originalContext,
+        parentId,
+        messageOperator,
+        `/api/v1/namespaces/${namespaceId}/wizard/${mode}`,
+        originalLang,
+        namespaceId,
+        undefined,
+        undefined,
+        originalEnableThinking
+      );
+      askAbortRef.current = askFN.destroy;
+      await askFN.start();
+    } finally {
+      setLoading(false);
+    }
+  };
+
   return {
     mode,
     value,
@@ -141,5 +224,8 @@ export default function useContext() {
     onContextChange,
     namespaceId,
     conversation,
+    messageOperator,
+    onRegenerate,
+    onEdit,
   };
 }