// AI

galaxy/galaxy/Classes/Modules/AI/AI-Dialogue(AI对话)/C/YHAIChatViewController.swift

@@ -83,6 +83,8 @@ class YHAIChatViewController: YHBaseViewController {
     
     func handleMessage(_ res: YHAIMessage?, _ done: Bool) {
         
+        removeThinkingMessage()
+        
         if let res = res {
             var find = false
             for msg in messages {
@@ -99,21 +101,24 @@ class YHAIChatViewController: YHBaseViewController {
             }
         } else {
             if done { // 思考完成
-                var findIndex = -1
-                for (index, msg) in messages.enumerated() {
-                    if msg.getType() == .thinking {
-                        findIndex = index
-                    }
-                }
-                if findIndex != -1 {
-                    messages.remove(at: findIndex)
-                }
-                self.tableView.reloadData()
+                
             }
         }
         self.scrollToBottom()
     }
     
+    func removeThinkingMessage() {
+        var findIndex = -1
+        for (index, msg) in messages.enumerated() {
+            if msg.getType() == .thinking {
+                findIndex = index
+            }
+        }
+        if findIndex != -1 {
+            messages.remove(at: findIndex)
+        }
+    }
+    
     func scrollToBottom() {
         self.tableView.reloadData {
             let indexPath = IndexPath(row: self.messages.count-1, section: 0)

galaxy/galaxy/Classes/Modules/AI/AI-Dialogue(AI对话)/M/YHAIMessageModel.swift

@@ -53,6 +53,7 @@ class YHAIMessage: CustomStringConvertible {
         body.contentText = text
         question.setFinished()
         question.body = body
+        question.updateBodyToData()
         return question
     }
     
@@ -63,15 +64,21 @@ class YHAIMessage: CustomStringConvertible {
         body.contentType = YHAIMessageType.thinking.rawValue
         message.setFinished()
         message.body = body
+        message.updateBodyToData()
         return message
     }
     
+    func updateBodyToData() {
+        data["contentType"] = body.contentType
+        data["contentText"] = body.contentText
+    }
+    
     // 获取消息类型
     func getType() -> YHAIMessageType {
-        if self.body.contentType == YHAIMessageType.thinking.rawValue {
-            return .thinking
-        }
         if let type = data["contentType"] as? Int {
+            if type == YHAIMessageType.thinking.rawValue {
+                return .thinking
+            }
             if type == YHAIMessageType.text.rawValue {
                 return .text
             }