当 tokens 超出 OpenAI API 限制大小时自动重置上下文

server/chat_handler.go

@@ -7,7 +7,6 @@ import (
 	"fmt"
 	"github.com/gin-gonic/gin"
 	"github.com/gorilla/websocket"
-	"io"
 	"math/rand"
 	"net/http"
 	"net/url"
@@ -92,6 +91,14 @@ func (s *Server) sendMessage(session types.ChatSession, role types.ChatRole, pro
 		logger.Infof("会话上下文：%+v", context)
 	}
 
+	// 创建 HttpClient 请求对象
+	var client *http.Client
+	var retryCount = 5
+	var response *http.Response
+	var apiKey string
+	var failedKey = ""
+	var failedProxyURL = ""
+	for retryCount > 0 {
 		r.Messages = append(context, types.Message{
 			Role:    "user",
 			Content: prompt,
@@ -102,19 +109,13 @@ func (s *Server) sendMessage(session types.ChatSession, role types.ChatRole, pro
 			return err
 		}
 
-	// 创建 HttpClient 请求对象
-	var client *http.Client
 		request, err := http.NewRequest(http.MethodPost, s.Config.Chat.ApiURL, bytes.NewBuffer(requestBody))
 		if err != nil {
 			return err
 		}
 
 		request.Header.Add("Content-Type", "application/json")
-	var retryCount = 5
-	var response *http.Response
-	var failedKey = ""
-	var failedProxyURL = ""
-	for retryCount > 0 {
+
 		proxyURL := s.getProxyURL(failedProxyURL)
 		if proxyURL == "" {
 			client = &http.Client{}
@@ -127,7 +128,7 @@ func (s *Server) sendMessage(session types.ChatSession, role types.ChatRole, pro
 				},
 			}
 		}
-		apiKey := s.getApiKey(failedKey)
+		apiKey = s.getApiKey(failedKey)
 		if apiKey == "" {
 			logger.Info("Too many requests, all Api Key is not available")
 			time.Sleep(time.Second)
@@ -139,7 +140,12 @@ func (s *Server) sendMessage(session types.ChatSession, role types.ChatRole, pro
 		if err == nil {
 			break
 		} else {
-			logger.Error(err)
+			// 上下文超出长度了
+			if strings.Contains(err.Error(), "This model's maximum context length is 4097 tokens") {
+				logger.Info("会话上下文长度超出限制, Username: %s", user.Name)
+				replyMessage(ws, "温馨提示：会话上下文长度超出限制，已为您重置会话上下文！", false)
+				context = role.Context
+			}
 			failedKey = apiKey
 			failedProxyURL = proxyURL
 		}
@@ -160,16 +166,23 @@ func (s *Server) sendMessage(session types.ChatSession, role types.ChatRole, pro
 	reader := bufio.NewReader(response.Body)
 	for {
 		line, err := reader.ReadString('\n')
-		if err != nil && err != io.EOF {
+		if err != nil {
 			logger.Error(err)
 			break
 		}
-
-		if line == "" {
-			replyChunkMessage(ws, types.WsMessage{Type: types.WsEnd})
-			break
-		} else if len(line) < 20 {
+		if len(line) < 20 {
 			continue
+		} else if strings.Contains(line, "This key is associated with a deactivated account") {
+			logger.Infof("API Key %s is deactivated", apiKey)
+			// 移除当前 API key
+			for i, v := range s.Config.Chat.ApiKeys {
+				if v == apiKey {
+					s.Config.Chat.ApiKeys = append(s.Config.Chat.ApiKeys[:i], s.Config.Chat.ApiKeys[i+1:]...)
+				}
+			}
+
+			// 重发当前消息
+			return s.sendMessage(session, role, prompt, ws)
 		}
 
 		err = json.Unmarshal([]byte(line[6:]), &responseBody)