From a5e18a1ee280bff2647c59e3573e98b6d610e87a Mon Sep 17 00:00:00 2001
From: Zhang Minghan <zmh13054618081@dingtalk.com>
Date: Sat, 21 Oct 2023 18:15:17 +0800
Subject: [PATCH] update frontend

---
 app/src/assets/chat.less             | 25 +++++++++++++
 app/src/components/Message.tsx       | 30 ++++++++++++---
 app/src/conf.ts                      |  4 +-
 app/src/conversation/conversation.ts | 52 ++++++++++++++++++++------
 app/src/conversation/manager.ts      |  4 +-
 app/src/conversation/types.ts        |  1 +
 app/src/events/connection.ts         | 10 +++++
 app/src/events/sharing.ts            |  2 +-
 app/src/routes/Home.tsx              | 20 ++++++++--
 app/src/routes/Sharing.tsx           |  4 +-
 manager/chat.go                      | 56 ++++++++++++++++++++--------
 manager/conversation/conversation.go |  9 +++++
 manager/manager.go                   |  2 +-
 utils/char.go                        | 20 ++++++++++
 14 files changed, 195 insertions(+), 44 deletions(-)
 create mode 100644 app/src/events/connection.ts

diff --git a/app/src/assets/chat.less b/app/src/assets/chat.less
index 5d725f1c..1e88412a 100644
--- a/app/src/assets/chat.less
+++ b/app/src/assets/chat.less
@@ -55,6 +55,31 @@
     }
   }
 
+  .content-wrapper {
+    display: flex;
+    flex-direction: row;
+
+    .message-toolbar {
+      display: flex;
+      flex-direction: column;
+      padding: 0 4px;
+      user-select: none;
+      height: max-content;
+      margin-top: auto;
+      gap: 4px;
+
+      svg {
+        cursor: pointer;
+        color: hsl(var(--text-secondary));
+        transition: 0.25s;
+
+        &:hover {
+          color: hsl(var(--text));
+        }
+      }
+    }
+  }
+
   .message-quota {
     display: flex;
     flex-direction: row;
diff --git a/app/src/components/Message.tsx b/app/src/components/Message.tsx
index aa6105ca..75518209 100644
--- a/app/src/components/Message.tsx
+++ b/app/src/components/Message.tsx
@@ -6,7 +6,7 @@ import {
   Copy,
   File,
   Loader2,
-  MousePointerSquare,
+  MousePointerSquare, Power, RotateCcw,
 } from "lucide-react";
 import {
   ContextMenu,
@@ -30,16 +30,19 @@ import {
 
 type MessageProps = {
   message: Message;
+  end?: boolean;
+  onEvent?: (event: string) => void;
 };
 
-function MessageSegment({ message }: MessageProps) {
+function MessageSegment(props: MessageProps) {
   const { t } = useTranslation();
+  const { message } = props;
 
   return (
     <ContextMenu>
       <ContextMenuTrigger asChild>
         <div className={`message ${message.role}`}>
-          <MessageContent message={message} />
+          <MessageContent {...props} />
           {message.quota && message.quota !== 0 ? (
             <TooltipProvider>
               <Tooltip>
@@ -86,9 +89,9 @@ function MessageSegment({ message }: MessageProps) {
   );
 }
 
-function MessageContent({ message }: MessageProps) {
+function MessageContent({ message, end, onEvent }: MessageProps) {
   return (
-    <>
+    <div className={`content-wrapper`}>
       <div className={`message-content`}>
         {message.keyword && message.keyword.length ? (
           <div className={`bing`}>
@@ -162,7 +165,22 @@ function MessageContent({ message }: MessageProps) {
           <Loader2 className={`h-5 w-5 m-1 animate-spin`} />
         )}
       </div>
-    </>
+      {
+        (message.role === "assistant" && end === true) && (
+          <div className={`message-toolbar`}>
+            {
+              (message.end !== false) ?
+                <RotateCcw className={`h-4 w-4 m-0.5`} onClick={() => (
+                  onEvent && onEvent("restart")
+                )} /> :
+                <Power className={`h-4 w-4 m-0.5`} onClick={() => (
+                  onEvent && onEvent("stop")
+                )} />
+            }
+          </div>
+        )
+      }
+    </div>
   );
 }
 
diff --git a/app/src/conf.ts b/app/src/conf.ts
index 691731fe..b267b41a 100644
--- a/app/src/conf.ts
+++ b/app/src/conf.ts
@@ -1,7 +1,7 @@
 import axios from "axios";
 
-export const version = "3.4.4";
-export const deploy: boolean = false;
+export const version = "3.4.5";
+export const deploy: boolean = true;
 export let rest_api: string = "http://localhost:8094";
 export let ws_api: string = "ws://localhost:8094";
 
diff --git a/app/src/conversation/conversation.ts b/app/src/conversation/conversation.ts
index dfe9744b..edf55a78 100644
--- a/app/src/conversation/conversation.ts
+++ b/app/src/conversation/conversation.ts
@@ -1,6 +1,7 @@
 import { ChatProps, Connection, StreamMessage } from "./connection.ts";
 import { Message } from "./types.ts";
-import { event } from "../events/sharing.ts";
+import { sharingEvent } from "../events/sharing.ts";
+import {connectionEvent} from "../events/connection.ts";
 
 type ConversationCallback = (idx: number, message: Message[]) => void;
 
@@ -11,7 +12,6 @@ export class Conversation {
   public id: number;
   public data: Message[];
   public end: boolean;
-  public refer: string;
 
   public constructor(id: number, callback?: ConversationCallback) {
     if (callback) this.setCallback(callback);
@@ -20,23 +20,50 @@ export class Conversation {
     this.id = id;
     this.end = true;
     this.connection = new Connection(this.id);
-    this.refer = "";
 
     if (id === -1 && this.idx === -1) {
-      event.bind(({ refer, data }) => {
+      sharingEvent.bind(({ refer, data }) => {
         console.log(
           `[conversation] load from sharing event (ref: ${refer}, length: ${data.length})`,
         );
-        this.refer = refer;
-        this.load(data);
 
-        this.connection?.sendWithRetry(null, {
-          type: "share",
-          message: this.refer,
-          model: "gpt-3.5-turbo",
-        });
+        this.load(data);
+        this.sendEvent("share", refer);
       });
     }
+
+    connectionEvent.addEventListener((ev) => {
+      if (ev.id === this.id) {
+        console.debug(`[conversation] connection event (id: ${this.id}, event: ${ev.event})`);
+
+        switch (ev.event) {
+          case "stop":
+            this.end = true;
+            this.data[this.data.length - 1].end = true;
+            this.sendEvent("stop");
+            this.triggerCallback();
+            break;
+
+          case "restart":
+            this.end = false;
+            delete this.data[this.data.length - 1];
+            this.connection?.setCallback(this.useMessage());
+            this.sendEvent("restart");
+            break;
+
+          default:
+            console.debug(`[conversation] unknown event: ${ev.event} (from: ${ev.id})`);
+        }
+      }
+    })
+  }
+
+  protected sendEvent(event: string, data?: string) {
+    this.connection?.sendWithRetry(null, {
+      type: event,
+      message: data || "",
+      model: "event",
+    });
   }
 
   public setId(id: number): void {
@@ -98,10 +125,12 @@ export class Conversation {
     message: string,
     keyword?: string,
     quota?: number,
+    end?: boolean,
   ) {
     this.data[idx].content += message;
     if (keyword) this.data[idx].keyword = keyword;
     if (quota) this.data[idx].quota = quota;
+    this.data[idx].end = end;
     this.triggerCallback();
   }
 
@@ -117,6 +146,7 @@ export class Conversation {
         message.message,
         message.keyword,
         message.quota,
+        message.end,
       );
       if (message.end) {
         this.end = true;
diff --git a/app/src/conversation/manager.ts b/app/src/conversation/manager.ts
index 047271b8..b83e73e1 100644
--- a/app/src/conversation/manager.ts
+++ b/app/src/conversation/manager.ts
@@ -11,7 +11,7 @@ import { useShared } from "../utils.ts";
 import { ChatProps } from "./connection.ts";
 import { supportModelConvertor } from "../conf.ts";
 import { AppDispatch } from "../store";
-import { event } from "../events/sharing.ts";
+import { sharingEvent } from "../events/sharing.ts";
 
 export class Manager {
   conversations: Record<number, Conversation>;
@@ -23,7 +23,7 @@ export class Manager {
     this.conversations[-1] = this.createConversation(-1);
     this.current = -1;
 
-    event.addEventListener(async (data) => {
+    sharingEvent.addEventListener(async (data) => {
       console.debug(`[manager] accept sharing event (refer: ${data.refer})`);
 
       const interval = setInterval(() => {
diff --git a/app/src/conversation/types.ts b/app/src/conversation/types.ts
index fec02f94..b8e52be9 100644
--- a/app/src/conversation/types.ts
+++ b/app/src/conversation/types.ts
@@ -5,6 +5,7 @@ export type Message = {
   content: string;
   keyword?: string;
   quota?: number;
+  end?: boolean;
 };
 
 export type Id = number;
diff --git a/app/src/events/connection.ts b/app/src/events/connection.ts
new file mode 100644
index 00000000..95d1fc2e
--- /dev/null
+++ b/app/src/events/connection.ts
@@ -0,0 +1,10 @@
+import {EventCommitter} from "./struct.ts";
+
+export type ConnectionEvent = {
+  id: number;
+  event: string;
+};
+
+export const connectionEvent = new EventCommitter<ConnectionEvent>({
+  name: "connection",
+});
diff --git a/app/src/events/sharing.ts b/app/src/events/sharing.ts
index 3a7aca53..ab546d00 100644
--- a/app/src/events/sharing.ts
+++ b/app/src/events/sharing.ts
@@ -6,7 +6,7 @@ export type SharingEvent = {
   data: Message[];
 };
 
-export const event = new EventCommitter<SharingEvent>({
+export const sharingEvent = new EventCommitter<SharingEvent>({
   name: "sharing",
   destroyedAfterTrigger: true,
 });
diff --git a/app/src/routes/Home.tsx b/app/src/routes/Home.tsx
index 158bd46c..414c4e47 100644
--- a/app/src/routes/Home.tsx
+++ b/app/src/routes/Home.tsx
@@ -69,6 +69,7 @@ import router from "../router.ts";
 import SelectGroup from "../components/SelectGroup.tsx";
 import EditorProvider from "../components/EditorProvider.tsx";
 import ConversationSegment from "../components/home/ConversationSegment.tsx";
+import {connectionEvent} from "../events/connection.ts";
 
 function SideBar() {
   const { t } = useTranslation();
@@ -307,6 +308,7 @@ function ChatInterface() {
   const ref = useRef(null);
   const [scroll, setScroll] = useState(false);
   const messages: Message[] = useSelector(selectMessages);
+  const current: number = useSelector(selectCurrent);
 
   function listenScrolling() {
     if (!ref.current) return;
@@ -351,9 +353,21 @@ function ChatInterface() {
           </Button>
         </div>
 
-        {messages.map((message, i) => (
-          <MessageSegment message={message} key={i} />
-        ))}
+        {
+          messages.map((message, i) =>
+            <MessageSegment
+              message={message}
+              end={i === messages.length - 1}
+              onEvent={(e: string) => {
+                connectionEvent.emit({
+                  id: current,
+                  event: e,
+                });
+              }}
+              key={i}
+            />
+          )
+        }
       </div>
     </>
   );
diff --git a/app/src/routes/Sharing.tsx b/app/src/routes/Sharing.tsx
index a666dd0a..59dd6d10 100644
--- a/app/src/routes/Sharing.tsx
+++ b/app/src/routes/Sharing.tsx
@@ -13,7 +13,7 @@ import MessageSegment from "../components/Message.tsx";
 import { Button } from "../components/ui/button.tsx";
 import router from "../router.ts";
 import { useToast } from "../components/ui/use-toast.ts";
-import { event } from "../events/sharing.ts";
+import { sharingEvent } from "../events/sharing.ts";
 import { Message } from "../conversation/types.ts";
 
 type SharingFormProps = {
@@ -72,7 +72,7 @@ function SharingForm({ refer, data }: SharingFormProps) {
         <Button
           variant={`outline`}
           onClick={async () => {
-            event.emit({
+            sharingEvent.emit({
               refer: refer as string,
               data: data?.messages as Message[],
             });
diff --git a/manager/chat.go b/manager/chat.go
index 48b31e40..eb211951 100644
--- a/manager/chat.go
+++ b/manager/chat.go
@@ -10,6 +10,7 @@ import (
 	"fmt"
 	"github.com/gin-gonic/gin"
 	"strings"
+	"time"
 )
 
 const defaultMessage = "Sorry, I don't understand. Please try again."
@@ -53,6 +54,27 @@ func ImageHandler(conn *Connection, user *auth.User, instance *conversation.Conv
 	}
 }
 
+func MockStreamSender(conn *Connection, message string) {
+	for _, line := range utils.SplitLangItems(message) {
+		time.Sleep(100 * time.Millisecond)
+		conn.Send(globals.ChatSegmentResponse{
+			Message: line + " ",
+			End:     false,
+			Quota:   0,
+		})
+
+		if signal := conn.PeekWithType(StopType); signal != nil {
+			// stop signal from client
+			break
+		}
+	}
+
+	conn.Send(globals.ChatSegmentResponse{
+		End:   true,
+		Quota: 0,
+	})
+}
+
 func ChatHandler(conn *Connection, user *auth.User, instance *conversation.Conversation) string {
 	defer func() {
 		if err := recover(); err != nil {
@@ -84,16 +106,12 @@ func ChatHandler(conn *Connection, user *auth.User, instance *conversation.Conve
 		Model:      model,
 		Reversible: reversible,
 	}); form != nil {
-		conn.Send(globals.ChatSegmentResponse{
-			Message: form.Message,
-			Quota:   0,
-			End:     true,
-		})
+		MockStreamSender(conn, form.Message)
 		return form.Message
 	}
 
 	buffer := utils.NewBuffer(model, segment)
-	if err := adapter.NewChatRequest(&adapter.ChatProps{
+	err := adapter.NewChatRequest(&adapter.ChatProps{
 		Model:      model,
 		Message:    segment,
 		Reversible: reversible && globals.IsGPT4Model(model),
@@ -107,7 +125,9 @@ func ChatHandler(conn *Connection, user *auth.User, instance *conversation.Conve
 			Quota:   buffer.GetQuota(),
 			End:     false,
 		})
-	}); err != nil && err.Error() != "signal" {
+	})
+
+	if err != nil && err.Error() != "signal" {
 		CollectQuota(conn.GetCtx(), user, buffer.GetQuota(), reversible)
 		conn.Send(globals.ChatSegmentResponse{
 			Message: err.Error(),
@@ -130,14 +150,18 @@ func ChatHandler(conn *Connection, user *auth.User, instance *conversation.Conve
 
 	conn.Send(globals.ChatSegmentResponse{End: true, Quota: buffer.GetQuota()})
 
-	SaveCacheData(conn.GetCtx(), &CacheProps{
-		Message:    segment,
-		Model:      model,
-		Reversible: reversible,
-	}, &CacheData{
-		Keyword: keyword,
-		Message: buffer.ReadWithDefault(defaultMessage),
-	})
+	result := buffer.ReadWithDefault(defaultMessage)
+
+	if err == nil && result != defaultMessage {
+		SaveCacheData(conn.GetCtx(), &CacheProps{
+			Message:    segment,
+			Model:      model,
+			Reversible: reversible,
+		}, &CacheData{
+			Keyword: keyword,
+			Message: result,
+		})
+	}
 
-	return buffer.ReadWithDefault(defaultMessage)
+	return result
 }
diff --git a/manager/conversation/conversation.go b/manager/conversation/conversation.go
index 7de6f60d..60556e07 100644
--- a/manager/conversation/conversation.go
+++ b/manager/conversation/conversation.go
@@ -74,6 +74,9 @@ func ExtractConversation(db *sql.DB, user *auth.User, id int64, ref string) *Con
 }
 
 func (c *Conversation) GetModel() string {
+	if len(c.Model) == 0 {
+		return globals.GPT3Turbo
+	}
 	return c.Model
 }
 
@@ -82,6 +85,9 @@ func (c *Conversation) IsEnableWeb() bool {
 }
 
 func (c *Conversation) SetModel(model string) {
+	if len(model) == 0 {
+		model = globals.GPT3Turbo
+	}
 	c.Model = model
 }
 
@@ -246,6 +252,9 @@ func (c *Conversation) SaveResponse(db *sql.DB, message string) {
 }
 
 func (c *Conversation) RemoveMessage(index int) globals.Message {
+	if index < 0 || index >= len(c.Message) {
+		return globals.Message{}
+	}
 	message := c.Message[index]
 	c.Message = append(c.Message[:index], c.Message[index+1:]...)
 	return message
diff --git a/manager/manager.go b/manager/manager.go
index 10fb2ee4..41b79b2d 100644
--- a/manager/manager.go
+++ b/manager/manager.go
@@ -62,7 +62,7 @@ func ChatAPI(c *gin.Context) {
 		case ShareType:
 			instance.LoadSharing(db, form.Message)
 		case RestartType:
-			if message := instance.RemoveLatestMessage(); message.Role != "user" {
+			if message := instance.RemoveLatestMessage(); message.Role != "assistant" {
 				return fmt.Errorf("message type error")
 			}
 			response := EventHandler(buf, instance, user)
diff --git a/utils/char.go b/utils/char.go
index 9f31cbf4..6774f7b6 100644
--- a/utils/char.go
+++ b/utils/char.go
@@ -103,6 +103,26 @@ func SplitItem(data string, sep string) []string {
 	return result
 }
 
+func SplitItems(data string, seps []string) []string {
+	if len(seps) == 0 {
+		return []string{}
+	}
+
+	result := []string{data}
+	for _, sep := range seps {
+		var temp []string
+		for _, item := range result {
+			temp = append(temp, SplitItem(item, sep)...)
+		}
+		result = temp
+	}
+	return result
+}
+
+func SplitLangItems(data string) []string {
+	return SplitItems(data, []string{",", "，", " ", "\n"})
+}
+
 func Extract(data string, length int, flow string) string {
 	value := []rune(data)
 	if len(value) > length {