The 'stop generate' and 'regenerate response' function is ready

server/chat_handler.go

@@ -3,6 +3,7 @@ package server
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -59,12 +60,13 @@ func (s *Server) ChatHandle(c *gin.Context) {
 				delete(s.ChatClients, sessionId)
 				return
 			}
-
 			logger.Info("Receive a message: ", string(message))
 			//replyMessage(client, "当前 TOKEN 无效，请使用合法的 TOKEN 登录！", false)
 			//replyMessage(client, "![](images/wx.png)", true)
-			// TODO: 当前只保持当前会话的上下文，部保存用户的所有的聊天历史记录，后期要考虑保存所有的历史记录
+			ctx, cancel := context.WithCancel(context.Background())
-			err = s.sendMessage(session, chatRole, string(message), client, false)
+			s.ReqCancelFunc[sessionId] = cancel
+			// 回复消息
+			err = s.sendMessage(ctx, session, chatRole, string(message), client, false)
 			if err != nil {
 				logger.Error(err)
 			}
@@ -73,7 +75,13 @@ func (s *Server) ChatHandle(c *gin.Context) {
 }
 
 // 将消息发送给 ChatGPT 并获取结果，通过 WebSocket 推送到客户端
-func (s *Server) sendMessage(session types.ChatSession, role types.ChatRole, prompt string, ws Client, resetContext bool) error {
+func (s *Server) sendMessage(ctx context.Context, session types.ChatSession, role types.ChatRole, prompt string, ws Client, resetContext bool) error {
+	cancel := s.ReqCancelFunc[session.SessionId]
+	defer func() {
+		cancel()
+		delete(s.ReqCancelFunc, session.SessionId)
+	}()
+
 	user, err := GetUser(session.Username)
 	if err != nil {
 		replyMessage(ws, "当前 TOKEN 无效，请使用合法的 TOKEN 登录！", false)
@@ -98,38 +106,38 @@ func (s *Server) sendMessage(session types.ChatSession, role types.ChatRole, pro
 		replyMessage(ws, "![](images/start.png)", true)
 		return nil
 	}
-	var r = types.ApiRequest{
+	var req = types.ApiRequest{
 		Model:       s.Config.Chat.Model,
 		Temperature: s.Config.Chat.Temperature,
 		MaxTokens:   s.Config.Chat.MaxTokens,
 		Stream:      true,
 	}
-	var context []types.Message
+	var chatCtx []types.Message
 	var ctxKey = fmt.Sprintf("%s-%s", session.SessionId, role.Key)
 	if v, ok := s.ChatContexts[ctxKey]; ok && s.Config.Chat.EnableContext {
-		context = v.Messages
+		chatCtx = v.Messages
 	} else {
-		context = role.Context
+		chatCtx = role.Context
 	}
 
 	if s.DebugMode {
-		logger.Infof("会话上下文：%+v", context)
+		logger.Infof("会话上下文：%+v", chatCtx)
 	}
 
+	req.Messages = append(chatCtx, types.Message{
+		Role:    "user",
+		Content: prompt,
+	})
+
 	// 创建 HttpClient 请求对象
 	var client *http.Client
-	var retryCount = 5
+	var retryCount = 5 // 重试次数
 	var response *http.Response
 	var apiKey string
 	var failedKey = ""
 	var failedProxyURL = ""
 	for retryCount > 0 {
-		r.Messages = append(context, types.Message{
+		requestBody, err := json.Marshal(req)
-			Role:    "user",
-			Content: prompt,
-		})
-
-		requestBody, err := json.Marshal(r)
 		if err != nil {
 			return err
 		}
@@ -139,6 +147,7 @@ func (s *Server) sendMessage(session types.ChatSession, role types.ChatRole, pro
 			return err
 		}
 
+		request = request.WithContext(ctx)
 		request.Header.Add("Content-Type", "application/json")
 
 		proxyURL := s.getProxyURL(failedProxyURL)
@@ -164,7 +173,10 @@ func (s *Server) sendMessage(session types.ChatSession, role types.ChatRole, pro
 		response, err = client.Do(request)
 		if err == nil {
 			break
+		} else if strings.Contains(err.Error(), "context canceled") {
+			return errors.New("用户取消了请求：" + prompt)
 		} else {
+			logger.Error("HTTP API 请求失败：" + err.Error())
 			failedKey = apiKey
 			failedProxyURL = proxyURL
 		}
@@ -205,14 +217,14 @@ func (s *Server) sendMessage(session types.ChatSession, role types.ChatRole, pro
 			_ = utils.SaveConfig(s.Config, s.ConfigPath)
 
 			// 重发当前消息
-			return s.sendMessage(session, role, prompt, ws, false)
+			return s.sendMessage(ctx, session, role, prompt, ws, false)
 
 			// 上下文超出长度了
 		} else if strings.Contains(line, "This model's maximum context length is 4097 tokens") {
 			logger.Infof("会话上下文长度超出限制, Username: %s", user.Name)
 			// 重置上下文，重发当前消息
 			delete(s.ChatContexts, ctxKey)
-			return s.sendMessage(session, role, prompt, ws, true)
+			return s.sendMessage(ctx, session, role, prompt, ws, true)
 		} else if !strings.Contains(line, "data:") {
 			continue
 		}
@@ -246,7 +258,18 @@ func (s *Server) sendMessage(session types.ChatSession, role types.ChatRole, pro
 				IsHelloMsg: false,
 			})
 		}
-	}
+
+		// 监控取消信号
+		select {
+		case <-ctx.Done():
+			// 结束输出
+			replyChunkMessage(ws, types.WsMessage{Type: types.WsEnd, IsHelloMsg: false})
+			_ = response.Body.Close()
+			return errors.New("用户取消了请求：" + prompt)
+		default:
+			continue
+		}
+	} // end for
 	_ = response.Body.Close() // 关闭资源
 
 	// 消息发送成功
@@ -260,27 +283,26 @@ func (s *Server) sendMessage(session types.ChatSession, role types.ChatRole, pro
 		if message.Role == "" {
 			message.Role = "assistant"
 		}
-		// 追加上下文消息
-		useMsg := types.Message{Role: "user", Content: prompt}
-		context = append(context, useMsg)
 		message.Content = strings.Join(contents, "")
+		useMsg := types.Message{Role: "user", Content: prompt}
 
 		// 更新上下文消息
 		if s.Config.Chat.EnableContext {
-			context = append(context, message)
+			chatCtx = append(chatCtx, useMsg)  // 提问消息
+			chatCtx = append(chatCtx, message) // 回复消息
 			s.ChatContexts[ctxKey] = types.ChatContext{
-				Messages:       context,
+				Messages:       chatCtx,
 				LastAccessTime: time.Now().Unix(),
 			}
 		}
 
 		// 追加历史消息
 		if user.EnableHistory {
-			err = AppendChatHistory(user.Name, role.Key, useMsg)
+			err = AppendChatHistory(user.Name, role.Key, useMsg) // 提问消息
 			if err != nil {
 				return err
 			}
-			err = AppendChatHistory(user.Name, role.Key, message)
+			err = AppendChatHistory(user.Name, role.Key, message) // 回复消息
 		}
 	}
 
@@ -431,3 +453,12 @@ func (s *Server) ClearHistoryHandle(c *gin.Context) {
 
 	c.JSON(http.StatusOK, types.BizVo{Code: types.Success})
 }
+
+// StopGenerateHandle 停止生成
+func (s *Server) StopGenerateHandle(c *gin.Context) {
+	sessionId := c.GetHeader(types.TokenName)
+	cancel := s.ReqCancelFunc[sessionId]
+	cancel()
+	delete(s.ReqCancelFunc, sessionId)
+	c.JSON(http.StatusOK, types.BizVo{Code: types.Success})
+}

server/server.go

@@ -1,6 +1,7 @@
 package server
 
 import (
+	"context"
 	"embed"
 	"encoding/json"
 	"github.com/gin-contrib/sessions"
@@ -39,10 +40,11 @@ type Server struct {
 
 	// 保存 Websocket 会话 Username, 每个 Username 只能连接一次
 	// 防止第三方直接连接 socket 调用 OpenAI API
-	ChatSession      map[string]types.ChatSession //map[sessionId]User
+	ChatSession      map[string]types.ChatSession  //map[sessionId]User
-	ApiKeyAccessStat map[string]int64             // 记录每个 API Key 的最后访问之间，保持在 15/min 之内
+	ApiKeyAccessStat map[string]int64              // 记录每个 API Key 的最后访问之间，保持在 15/min 之内
-	ChatClients      map[string]*WsClient         // Websocket 连接集合
+	ChatClients      map[string]*WsClient          // Websocket 连接集合
-	DebugMode        bool                         // 是否开启调试模式
+	ReqCancelFunc    map[string]context.CancelFunc // HttpClient 请求取消 handle function
+	DebugMode        bool                          // 是否开启调试模式
 }
 
 func NewServer(configPath string) (*Server, error) {
@@ -67,6 +69,7 @@ func NewServer(configPath string) (*Server, error) {
 		ChatContexts:     make(map[string]types.ChatContext, 16),
 		ChatSession:      make(map[string]types.ChatSession),
 		ChatClients:      make(map[string]*WsClient),
+		ReqCancelFunc:    make(map[string]context.CancelFunc),
 		ApiKeyAccessStat: make(map[string]int64),
 	}, nil
 }
@@ -83,17 +86,18 @@ func (s *Server) Run(webRoot embed.FS, path string, debug bool) {
 	engine.Use(AuthorizeMiddleware(s))
 	engine.Use(Recover)
 
-	engine.POST("/test", s.TestHandle)
+	engine.POST("test", s.TestHandle)
-	engine.GET("/api/session/get", s.GetSessionHandle)
+	engine.GET("api/session/get", s.GetSessionHandle)
-	engine.POST("/api/login", s.LoginHandle)
+	engine.POST("api/login", s.LoginHandle)
-	engine.POST("/api/logout", s.LogoutHandle)
+	engine.POST("api/logout", s.LogoutHandle)
-	engine.Any("/api/chat", s.ChatHandle)
+	engine.Any("api/chat", s.ChatHandle)
+	engine.POST("api/chat/stop", s.StopGenerateHandle)
 	engine.POST("api/chat/history", s.GetChatHistoryHandle)
 	engine.POST("api/chat/history/clear", s.ClearHistoryHandle)
 
-	engine.POST("/api/config/set", s.ConfigSetHandle)
+	engine.POST("api/config/set", s.ConfigSetHandle)
-	engine.GET("/api/config/chat-roles/get", s.GetChatRoleListHandle)
+	engine.GET("api/config/chat-roles/get", s.GetChatRoleListHandle)
-	engine.GET("/api/config/chat-roles/add", s.AddChatRoleHandle)
+	engine.GET("api/config/chat-roles/add", s.AddChatRoleHandle)
 	engine.POST("api/config/user/add", s.AddUserHandle)
 	engine.POST("api/config/user/batch-add", s.BatchAddUserHandle)
 	engine.POST("api/config/user/set", s.SetUserHandle)

test/test.go

@@ -1,10 +1,73 @@
 package main
 
 import (
+	"context"
 	"fmt"
+	"io"
+	"net/http"
 	"time"
 )
 
 func main() {
+
+	ctx, cancel := context.WithCancel(context.Background())
+
+	http.HandleFunc("/cancel", func(w http.ResponseWriter, r *http.Request) {
+		cancel()
+		_, _ = fmt.Fprintf(w, "请求取消!")
+	})
+
+	go func() {
+		err := http.ListenAndServe(":9999", nil)
+		if err != nil {
+			return
+		}
+	}()
+
+	testHttpClient(ctx)
+}
+
+// Http client 取消操作
+func testHttpClient(ctx context.Context) {
+
+	req, err := http.NewRequest("GET", "http://localhost:2345", nil)
+	if err != nil {
+		fmt.Println(err)
+		return
+	}
+
+	req = req.WithContext(ctx)
+
+	client := &http.Client{}
+
+	resp, err := client.Do(req)
+	if err != nil {
+		fmt.Println(err)
+		return
+	}
+	defer resp.Body.Close()
+	body, err := io.ReadAll(resp.Body)
+	for {
+		time.Sleep(time.Second)
+		fmt.Println(time.Now())
+		select {
+		case <-ctx.Done():
+			fmt.Println("取消退出")
+			return
+		default:
+			continue
+		}
+	}
+
+	if err != nil {
+		fmt.Println(err)
+		return
+	}
+
+	fmt.Println(string(body))
+
+}
+
+func testDate() {
 	fmt.Println(time.Unix(1683336167, 0).Format("2006-01-02 15:04:05"))
 }

web/.env.production

@@ -1,2 +1,2 @@
-VUE_APP_API_HOST=https://ai.r9it.com
+VUE_APP_API_HOST=https://www.chat-plus.net
-VUE_APP_WS_HOST=wss://ai.r9it.com
+VUE_APP_WS_HOST=wss://www.chat-plus.net

web/src/views/ChatPlus.vue

@@ -81,9 +81,26 @@
                             :icon="chat.icon"
                             :content="chat.content"/>
               </div>
-
             </div><!-- end chat box -->
 
+            <div class="re-generate">
+              <div class="btn-box">
+                <el-button type="info" v-if="showStopGenerate" @click="stopGenerate" plain>
+                  <el-icon>
+                    <VideoPause/>
+                  </el-icon>
+                  停止生成
+                </el-button>
+
+                <el-button type="info" v-if="showReGenerate" @click="reGenerate" plain>
+                  <el-icon>
+                    <RefreshRight/>
+                  </el-icon>
+                  重新生成
+                </el-button>
+              </div>
+            </div>
+
             <el-row class="chat-tool-box">
               <el-tooltip
                   class="box-item"
@@ -159,7 +176,17 @@ import ChatPrompt from "@/components/plus/ChatPrompt.vue";
 import ChatReply from "@/components/plus/ChatReply.vue";
 import {isMobile, randString} from "@/utils/libs";
 import {ElMessage, ElMessageBox} from 'element-plus'
-import {Tools, Lock, Delete, Picture, Search, ArrowDown, Monitor} from '@element-plus/icons-vue'
+import {
+  Tools,
+  Lock,
+  Delete,
+  Picture,
+  Search,
+  ArrowDown,
+  Monitor,
+  VideoPause,
+  RefreshRight
+} from '@element-plus/icons-vue'
 import ConfigDialog from '@/components/ConfigDialog.vue'
 import {httpPost, httpGet} from "@/utils/http";
 import {getSessionId, setLoginUser} from "@/utils/storage";
@@ -168,7 +195,20 @@ import 'highlight.js/styles/a11y-dark.css'
 
 export default defineComponent({
   name: "ChatPlus",
-  components: {ArrowDown, Search, ChatPrompt, ChatReply, Tools, Lock, Delete, Picture, Monitor, ConfigDialog},
+  components: {
+    RefreshRight,
+    VideoPause,
+    ArrowDown,
+    Search,
+    ChatPrompt,
+    ChatReply,
+    Tools,
+    Lock,
+    Delete,
+    Picture,
+    Monitor,
+    ConfigDialog
+  },
   data() {
     return {
       title: 'ChatGPT 控制台',
@@ -184,6 +224,11 @@ export default defineComponent({
       replyIcon: 'images/avatar/gpt.png', // 回复信息的头像
       roleName: "", // 搜索角色名称
 
+      showStopGenerate: false, // 停止生成
+      showReGenerate: false, // 重新生成
+      canReGenerate: false, // 是否可以重新生
+      previousText: '', // 上一次提问
+
       lineBuffer: '', // 输出缓冲行
       connectingMessageBox: null, // 保存重连的消息框对象
       errorMessage: null, // 错误信息提示框
@@ -260,8 +305,15 @@ export default defineComponent({
                 content: "",
                 cursor: true
               });
-            } else if (data.type === 'end') {
+              if (data['is_hello_msg'] !== true) {
+                this.canReGenerate = true;
+              }
+            } else if (data.type === 'end') { // 消息接收完毕
               this.sending = false;
+              if (data['is_hello_mgs'] !== true) {
+                this.showReGenerate = true;
+              }
+              this.showStopGenerate = false;
               this.lineBuffer = ''; // 清空缓冲
             } else {
               this.lineBuffer += data.content;
@@ -372,6 +424,7 @@ export default defineComponent({
     inputKeyDown: function (e) {
       if (e.keyCode === 13) {
         if (this.sending) {
+          ElMessage.warning("AI 正在作答中，请稍后...");
           e.preventDefault();
         } else {
           this.sendMessage();
@@ -390,10 +443,8 @@ export default defineComponent({
         target.blur();
       }
 
-      if (this.inputValue.trim().length === 0) {
+      if (this.inputValue.trim().length === 0 || this.sending) {
         return false;
-      } else if (this.sending) {
-        ElMessage.warning("AI 正在作答中请稍后...")
       }
 
       // 追加消息
@@ -405,8 +456,11 @@ export default defineComponent({
       });
 
       this.sending = true;
+      this.showStopGenerate = true;
+      this.showReGenerate = false;
       this.socket.send(this.inputValue);
       this.$refs["text-input"].blur();
+      this.previousText = this.inputValue;
       this.inputValue = '';
       // 等待 textarea 重新调整尺寸之后再自动获取焦点
       setTimeout(() => this.$refs["text-input"].focus(), 100);
@@ -494,6 +548,26 @@ export default defineComponent({
       }).catch(() => {
         ElMessage.error("注销失败");
       })
+    },
+
+    // 停止生成
+    stopGenerate: function () {
+      this.showStopGenerate = false;
+      httpPost("/api/chat/stop").then(() => {
+        console.log("stopped generate.")
+        this.sending = false;
+        if (this.canReGenerate) {
+          this.showReGenerate = true;
+        }
+      })
+    },
+
+    // 重新生成
+    reGenerate: function () {
+      this.sending = true;
+      this.showStopGenerate = true;
+      this.showReGenerate = false;
+      this.socket.send('重新生成上述问题的答案：' + this.previousText);
     }
   },
 
@@ -676,6 +750,24 @@ export default defineComponent({
         }
       }
 
+      .re-generate {
+        position: relative;
+        display: flex;
+        justify-content: center;
+
+        .btn-box {
+          position absolute
+          bottom 10px;
+
+          .el-button {
+            .el-icon {
+              margin-right 5px;
+            }
+          }
+
+        }
+      }
+
       .chat-tool-box {
         padding 10px;
         border-top: 1px solid #2F3032