File size: 3,923 Bytes
d145f59
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
import { useCallback, useState, useRef } from 'react';
import { Toast, Modal } from '@douyinfe/semi-ui';
import { useTranslation } from 'react-i18next';
import { getTextContent, buildApiPayload, createLoadingAssistantMessage } from '../helpers';
import { MESSAGE_ROLES } from '../constants/playground.constants';

export const useMessageEdit = (
  setMessage,
  inputs,
  parameterEnabled,
  sendRequest,
  saveMessages
) => {
  const { t } = useTranslation();
  const [editingMessageId, setEditingMessageId] = useState(null);
  const [editValue, setEditValue] = useState('');
  const editingMessageRef = useRef(null);

  const handleMessageEdit = useCallback((targetMessage) => {
    const editableContent = getTextContent(targetMessage);
    setEditingMessageId(targetMessage.id);
    editingMessageRef.current = targetMessage;
    setEditValue(editableContent);
  }, []);

  const handleEditSave = useCallback(() => {
    if (!editingMessageId || !editValue.trim()) return;

    setMessage(prevMessages => {
      let messageIndex = prevMessages.findIndex(msg => msg === editingMessageRef.current);

      if (messageIndex === -1) {
        messageIndex = prevMessages.findIndex(msg => msg.id === editingMessageId);
      }

      const targetMessage = prevMessages[messageIndex];
      let newContent;

      if (Array.isArray(targetMessage.content)) {
        newContent = targetMessage.content.map(item =>
          item.type === 'text' ? { ...item, text: editValue.trim() } : item
        );
      } else {
        newContent = editValue.trim();
      }

      const updatedMessages = prevMessages.map(msg =>
        msg.id === editingMessageId ? { ...msg, content: newContent } : msg
      );

      // 处理用户消息编辑后的重新生成
      if (targetMessage.role === MESSAGE_ROLES.USER) {
        const hasSubsequentAssistantReply = messageIndex < prevMessages.length - 1 &&
          prevMessages[messageIndex + 1].role === MESSAGE_ROLES.ASSISTANT;

        if (hasSubsequentAssistantReply) {
          Modal.confirm({
            title: t('消息已编辑'),
            content: t('检测到该消息后有AI回复,是否删除后续回复并重新生成?'),
            okText: t('重新生成'),
            cancelText: t('仅保存'),
            onOk: () => {
              const messagesUntilUser = updatedMessages.slice(0, messageIndex + 1);
              setMessage(messagesUntilUser);
              // 编辑后保存(重新生成的情况),传入更新后的消息列表
              setTimeout(() => saveMessages(messagesUntilUser), 0);

              setTimeout(() => {
                const payload = buildApiPayload(messagesUntilUser, null, inputs, parameterEnabled);
                setMessage(prevMsg => [...prevMsg, createLoadingAssistantMessage()]);
                sendRequest(payload, inputs.stream);
              }, 100);
            },
            onCancel: () => {
              setMessage(updatedMessages);
              // 编辑后保存(仅保存的情况),传入更新后的消息列表
              setTimeout(() => saveMessages(updatedMessages), 0);
            }
          });
          return prevMessages;
        }
      }

      // 编辑后保存(普通情况),传入更新后的消息列表
      setTimeout(() => saveMessages(updatedMessages), 0);
      return updatedMessages;
    });

    setEditingMessageId(null);
    editingMessageRef.current = null;
    setEditValue('');
    Toast.success({ content: t('消息已更新'), duration: 2 });
  }, [editingMessageId, editValue, t, inputs, parameterEnabled, sendRequest, setMessage, saveMessages]);

  const handleEditCancel = useCallback(() => {
    setEditingMessageId(null);
    editingMessageRef.current = null;
    setEditValue('');
  }, []);

  return {
    editingMessageId,
    editValue,
    setEditValue,
    handleMessageEdit,
    handleEditSave,
    handleEditCancel
  };
};