From d41931097ecf9c60f8994024204a880608438cca Mon Sep 17 00:00:00 2001
From: Zhang Minghan <zmh13054618081@dingtalk.com>
Date: Sun, 13 Aug 2023 08:32:34 +0800
Subject: [PATCH] update conversation

---
 api/chat.go                  | 11 +++++++++++
 api/utils.go                 |  6 ++++++
 conversation/conversation.go |  5 +++++
 3 files changed, 22 insertions(+)

diff --git a/api/chat.go b/api/chat.go
index bfc7e98..c1dcc32 100644
--- a/api/chat.go
+++ b/api/chat.go
@@ -73,12 +73,23 @@ func ChatAPI(c *gin.Context) {
 			keyword, segment := ChatWithWeb(instance.GetMessageSegment(12), true)
 			SendSegmentMessage(conn, types.ChatGPTSegmentResponse{Keyword: keyword, End: false})
 
+			msg := ""
 			StreamRequest("gpt-3.5-turbo-16k-0613", segment, 2000, func(resp string) {
+				msg += resp
 				SendSegmentMessage(conn, types.ChatGPTSegmentResponse{
 					Message: resp,
 					End:     false,
 				})
 			})
+			if msg == "" {
+				msg = "There was something wrong... Please try again later."
+				SendSegmentMessage(conn, types.ChatGPTSegmentResponse{
+					Message: msg,
+					End:     false,
+				})
+			}
+
+			instance.SaveResponse(db, msg)
 			SendSegmentMessage(conn, types.ChatGPTSegmentResponse{End: true})
 		}
 	}
diff --git a/api/utils.go b/api/utils.go
index 5ff884a..a700fa4 100644
--- a/api/utils.go
+++ b/api/utils.go
@@ -70,6 +70,12 @@ func SearchWeb(message []types.ChatGPTMessage) string {
 	}, {
 		Role:    "assistant",
 		Content: "{\"keyword\":\"\"}",
+	}, {
+		Role:    "user",
+		Content: "?",
+	}, {
+		Role:    "assistant",
+		Content: "{\"keyword\":\"\"}",
 	}, {
 		Role:    "user",
 		Content: message[len(message)-1].Content,
diff --git a/conversation/conversation.go b/conversation/conversation.go
index 95459aa..c177d69 100644
--- a/conversation/conversation.go
+++ b/conversation/conversation.go
@@ -124,3 +124,8 @@ func (c *Conversation) HandleMessage(db *sql.DB, data []byte) bool {
 	c.SaveConversation(db)
 	return true
 }
+
+func (c *Conversation) SaveResponse(db *sql.DB, message string) {
+	c.AddMessageFromAssistant(message)
+	c.SaveConversation(db)
+}