设置会话超时,避免出现cpu 100%

This commit is contained in:
cookeem
2023-02-13 10:34:02 +08:00
parent 1b9d4b5e65
commit be04f67050

View File

@@ -94,7 +94,9 @@ func (api *Api) wsPingMsg(conn *websocket.Conn, chClose, chIsCloseSet chan int)
} }
func (api *Api) GetChatMessage(conn *websocket.Conn, cli *gogpt.Client, mutex *sync.Mutex, requestMsg string) { func (api *Api) GetChatMessage(conn *websocket.Conn, cli *gogpt.Client, mutex *sync.Mutex, requestMsg string) {
var err error
var strResp string var strResp string
var end bool
req := gogpt.CompletionRequest{ req := gogpt.CompletionRequest{
Model: gogpt.GPT3TextDavinci003, Model: gogpt.GPT3TextDavinci003,
MaxTokens: api.Config.MaxLength, MaxTokens: api.Config.MaxLength,
@@ -107,8 +109,23 @@ func (api *Api) GetChatMessage(conn *websocket.Conn, cli *gogpt.Client, mutex *s
PresencePenalty: 0.1, PresencePenalty: 0.1,
} }
ctx, cancel := context.WithTimeout(context.Background(), time.Duration(time.Second*time.Duration(api.Config.TimeoutSeconds))) ctx, cancel := context.WithTimeout(context.Background(), time.Second*time.Duration(api.Config.TimeoutSeconds))
defer cancel() defer func() {
if !end {
err = fmt.Errorf("[ERROR] context timeout")
chatMsg := Message{
Kind: "error",
Msg: err.Error(),
MsgId: uuid.New().String(),
CreateTime: time.Now().Format("2006-01-02 15:04:05"),
}
mutex.Lock()
_ = conn.WriteJSON(chatMsg)
mutex.Unlock()
api.Logger.LogError(err.Error())
}
cancel()
}()
stream, err := cli.CreateCompletionStream(ctx, req) stream, err := cli.CreateCompletionStream(ctx, req)
if err != nil { if err != nil {
@@ -132,6 +149,7 @@ func (api *Api) GetChatMessage(conn *websocket.Conn, cli *gogpt.Client, mutex *s
for { for {
response, err := stream.Recv() response, err := stream.Recv()
if errors.Is(err, io.EOF) { if errors.Is(err, io.EOF) {
end = true
var s string var s string
var kind string var kind string
if i == 0 { if i == 0 {
@@ -155,6 +173,7 @@ func (api *Api) GetChatMessage(conn *websocket.Conn, cli *gogpt.Client, mutex *s
} }
break break
} else if err != nil { } else if err != nil {
end = true
err = fmt.Errorf("[ERROR] receive chatGPT stream error: %s", err.Error()) err = fmt.Errorf("[ERROR] receive chatGPT stream error: %s", err.Error())
chatMsg := Message{ chatMsg := Message{
Kind: "error", Kind: "error",
@@ -294,7 +313,6 @@ func (api *Api) WsChat(c *gin.Context) {
mutex.Unlock() mutex.Unlock()
api.Logger.LogError(err.Error()) api.Logger.LogError(err.Error())
} else { } else {
chatMsg := Message{ chatMsg := Message{
Kind: "receive", Kind: "receive",
Msg: requestMsg, Msg: requestMsg,