From 38da63e73cb15f7f428e2b1a39f50828156f0a31 Mon Sep 17 00:00:00 2001
From: Yang Luo <hsluoyz@qq.com>
Date: Tue, 2 May 2023 13:15:54 +0800
Subject: [PATCH] Improve answer text

---
 ai/ai.go            | 19 ++++++++++++++-----
 ai/ai_test.go       |  3 +++
 web/src/ChatBox.js  | 12 +++++++++++-
 web/src/ChatPage.js |  4 ++++
 4 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/ai/ai.go b/ai/ai.go
index fa97d4da..7cf7380e 100644
--- a/ai/ai.go
+++ b/ai/ai.go
@@ -25,7 +25,7 @@ import (
 )
 
 func queryAnswer(authToken string, question string, timeout int) (string, error) {
-	//fmt.Printf("Question: %s\n", question)
+	// fmt.Printf("Question: %s\n", question)
 
 	client := getProxyClientFromToken(authToken)
 
@@ -45,13 +45,12 @@ func queryAnswer(authToken string, question string, timeout int) (string, error)
 		},
 	)
 	if err != nil {
-		//fmt.Printf("%s\n", err.Error())
 		return "", err
 	}
 
 	res := resp.Choices[0].Message.Content
 	res = strings.Trim(res, "\n")
-	//fmt.Printf("Answer: %s\n\n", res)
+	// fmt.Printf("Answer: %s\n\n", res)
 	return res, nil
 }
 
@@ -104,6 +103,7 @@ func QueryAnswerStream(authToken string, question string, writer io.Writer, buil
 	}
 	defer respStream.Close()
 
+	isLeadingReturn := true
 	for {
 		completion, streamErr := respStream.Recv()
 		if streamErr != nil {
@@ -113,13 +113,22 @@ func QueryAnswerStream(authToken string, question string, writer io.Writer, buil
 			return streamErr
 		}
 
+		data := completion.Choices[0].Text
+		if isLeadingReturn && len(data) != 0 {
+			if strings.Count(data, "\n") == len(data) {
+				continue
+			} else {
+				isLeadingReturn = false
+			}
+		}
+
 		// Write the streamed data as Server-Sent Events
-		if _, err := fmt.Fprintf(writer, "data: %s\n\n", completion.Choices[0].Text); err != nil {
+		if _, err = fmt.Fprintf(writer, "data: %s\n\n", data); err != nil {
 			return err
 		}
 
 		// Append the response to the strings.Builder
-		builder.WriteString(completion.Choices[0].Text)
+		builder.WriteString(data)
 	}
 
 	return nil
diff --git a/ai/ai_test.go b/ai/ai_test.go
index b4a094cb..975cdd58 100644
--- a/ai/ai_test.go
+++ b/ai/ai_test.go
@@ -12,6 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+//go:build !skipCi
+// +build !skipCi
+
 package ai
 
 import (
diff --git a/web/src/ChatBox.js b/web/src/ChatBox.js
index edf6e7e2..ed04c46f 100644
--- a/web/src/ChatBox.js
+++ b/web/src/ChatBox.js
@@ -74,6 +74,16 @@ class ChatBox extends React.Component {
     this.setState({inputValue: ""});
   };
 
+  renderText(text) {
+    const lines = text.split("\n").map((line, index) => (
+      <React.Fragment key={index}>
+        {line}
+        <br />
+      </React.Fragment>
+    ));
+    return <div>{lines}</div>;
+  }
+
   renderList() {
     if (this.props.messages === undefined || this.props.messages === null) {
       return (
@@ -110,7 +120,7 @@ class ChatBox extends React.Component {
                     title={
                       <div style={{fontSize: "16px", fontWeight: "normal", lineHeight: "24px", marginTop: "-15px", marginLeft: "5px", marginRight: "80px"}}>
                         {
-                          !item.text.includes("#ERROR#") ? item.text : (
+                          !item.text.includes("#ERROR#") ? this.renderText(item.text) : (
                             <Alert message={item.text.slice("#ERROR#: ".length)} type="error" showIcon />
                           )
                         }
diff --git a/web/src/ChatPage.js b/web/src/ChatPage.js
index 71c643fd..37004186 100644
--- a/web/src/ChatPage.js
+++ b/web/src/ChatPage.js
@@ -89,6 +89,10 @@ class ChatPage extends BaseListPage {
           if (lastMessage.author === "AI" && lastMessage.replyTo !== "" && lastMessage.text === "") {
             let text = "";
             MessageBackend.getMessageAnswer(lastMessage.owner, lastMessage.name, (data) => {
+              if (data === "") {
+                data = "\n";
+              }
+
               const lastMessage2 = Setting.deepCopy(lastMessage);
               text += data;
               lastMessage2.text = text;