🐛 fix: 修复重新生成时的逻辑错误

lobehub · Jun 1, 2024 · 3e324e4 · 3e324e4
1 parent 35dfd8b
commit 3e324e4
Show file tree

Hide file tree

Showing 5 changed files with 37 additions and 19 deletions.
diff --git a/src/features/ChatItem/Actions/User.tsx b/src/features/ChatItem/Actions/User.tsx
@@ -1,18 +1,14 @@
-import { ActionIconGroup, useChatListActionsBar } from '@lobehub/ui';
+import { ActionIconGroup } from '@lobehub/ui';
 import { memo } from 'react';
 
 import type { RenderAction } from '@/features/ChatItem/type';
+import { useChatListActionsBar } from '@/hooks/useChatListActionsBar';
 
 const UserActionsBar: RenderAction = ({ onActionClick }) => {
-  const { copy, divider, del, edit, regenerate } = useChatListActionsBar({
-    copy: '复制',
-    delete: '删除',
-    edit: '编辑',
-    regenerate: '重新生成',
-  });
+  const { copy, divider, del, edit, regenerate } = useChatListActionsBar();
   return (
     <ActionIconGroup
-      dropdownMenu={[copy, divider, del]}
+      dropdownMenu={[edit, copy, divider, regenerate, del]}
       items={[regenerate, edit]}
       onActionClick={onActionClick}
       type="ghost"

diff --git a/src/store/session/index.ts b/src/store/session/index.ts
@@ -283,8 +283,30 @@ export const createSessonStore: StateCreator<SessionStore, [['zustand/devtools',
     if (!currentSession) {
       return;
     }
+    const chats = sessionSelectors.currentChats(get());
+    const currentIndex = chats.findIndex((item) => item.id === id);
+    const currentMessage = chats[currentIndex];
 
-    const previousChats = sessionSelectors.previousChats(get(), id);
+    let contextMessages: ChatMessage[] = [];
+
+    switch (currentMessage.role) {
+      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;
+      }
+    }
+
+    const latestMsg = contextMessages.filter((s) => s.role === 'user').at(-1);
+
+    if (!latestMsg) return;
 
     const assistantId = nanoid();
 
@@ -293,12 +315,13 @@ export const createSessonStore: StateCreator<SessionStore, [['zustand/devtools',
       payload: {
         content: LOADING_FLAG,
         id: assistantId,
+        parentId: latestMsg.id,
         role: 'assistant', // 占位符
       },
       type: 'ADD_MESSAGE',
     });
 
-    fetchAIResponse(previousChats, assistantId);
+    fetchAIResponse(contextMessages, assistantId);
   },
   removeSession: (id) => {
     const { sessionList, activeId } = get();
@@ -342,6 +365,7 @@ export const createSessonStore: StateCreator<SessionStore, [['zustand/devtools',
       payload: {
         content: LOADING_FLAG,
         id: assistantId,
+        parentId: userId,
         role: 'assistant', // 占位符
       },
       type: 'ADD_MESSAGE',

diff --git a/src/store/session/reducers/message.ts b/src/store/session/reducers/message.ts
@@ -7,6 +7,7 @@ export interface AddMessageAction {
   payload: {
     content: string;
     id: string;
+    parentId?: string;
     role: LLMRoleType;
   };
   type: 'ADD_MESSAGE';
@@ -34,11 +35,12 @@ export const messageReducer = (state: ChatMessage[], action: MessageActionType):
   switch (action.type) {
     case 'ADD_MESSAGE': {
       return produce(state, (draft) => {
-        const { role, content, id } = action.payload;
+        const { role, content, id, parentId } = action.payload;
         draft.push({
           content,
           createdAt: Date.now(),
           id,
+          parentId,
           meta: {},
           role,
           updatedAt: Date.now(),

diff --git a/src/store/session/selectors.ts b/src/store/session/selectors.ts
@@ -108,13 +108,6 @@ const currentChatIDsWithGreetingMessage = (s: SessionStore): string[] => {
   return currentChats.map((item) => item.id);
 };
 
-const previousChats = (s: SessionStore, id: string): ChatMessage[] => {
-  const chatList = currentChats(s);
-  const index = chatList.findIndex((item) => item.id === id);
-  if (index === -1) return [];
-  return chatList.slice(0, index);
-};
-
 const currentChatsString = (s: SessionStore): string => {
   const session = currentSession(s);
   const agent = currentAgent(s);
@@ -163,6 +156,5 @@ export const sessionSelectors = {
   currentChatsString,
   currentSession,
   currentSystemRole,
-  previousChats,
   sessionListIds,
 };
diff --git a/src/types/chat.ts b/src/types/chat.ts
@@ -61,6 +61,10 @@ export interface ChatMessage {
    * 元数据
    */
   meta: MetaData;
+  /**
+   * 父id
+   */
+  parentId?: string;
   /**
    * 角色
    */