pref(ChatGPT): advanced streaming and text transformation, support re…

…asoner model

api/openai/openai.go

@@ -2,15 +2,17 @@ package openai
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"github.com/0xJacky/Nginx-UI/internal/chatbot"
 	"github.com/0xJacky/Nginx-UI/settings"
 	"github.com/gin-gonic/gin"
-	"errors"
 	"github.com/sashabaranov/go-openai"
 	"github.com/uozi-tech/cosy"
 	"github.com/uozi-tech/cosy/logger"
 	"io"
+	"strings"
+	"time"
 )
 
 const ChatGPTInitPrompt = `You are a assistant who can help users write and optimise the configurations of Nginx,
@@ -83,31 +85,69 @@ func MakeChatCompletionRequest(c *gin.Context) {
 	msgChan := make(chan string)
 	go func() {
 		defer close(msgChan)
-		for {
-			response, err := stream.Recv()
-			if errors.Is(err, io.EOF) {
-				return
-			}
+		messageCh := make(chan string)
 
-			if err != nil {
-				logger.Errorf("Stream error: %v\n", err)
-				return
+		// 消息接收协程
+		go func() {
+			defer close(messageCh)
+			for {
+				response, err := stream.Recv()
+				if errors.Is(err, io.EOF) {
+					return
+				}
+				if err != nil {
+					messageCh <- fmt.Sprintf("error: %v", err)
+					logger.Errorf("Stream error: %v\n", err)
+					return
+				}
+				messageCh <- response.Choices[0].Delta.Content
 			}
+		}()
+
+		ticker := time.NewTicker(500 * time.Millisecond)
+		defer ticker.Stop()
 
-			message := fmt.Sprintf("%s", response.Choices[0].Delta.Content)
+		var buffer strings.Builder
 
-			msgChan <- message
+		for {
+			select {
+			case msg, ok := <-messageCh:
+				if !ok {
+					if buffer.Len() > 0 {
+						msgChan <- buffer.String()
+					}
+					return
+				}
+				if strings.HasPrefix(msg, "error: ") {
+					msgChan <- msg
+					return
+				}
+				buffer.WriteString(msg)
+			case <-ticker.C:
+				if buffer.Len() > 0 {
+					msgChan <- buffer.String()
+					buffer.Reset()
+				}
+			}
 		}
 	}()
 
 	c.Stream(func(w io.Writer) bool {
-		if m, ok := <-msgChan; ok {
+		m, ok := <-msgChan
+		if !ok {
+			return false
+		}
+		if strings.HasPrefix(m, "error: ") {
 			c.SSEvent("message", gin.H{
-				"type":    "message",
-				"content": m,
+				"type":    "error",
+				"content": strings.TrimPrefix(m, "error: "),
 			})
-			return true
+			return false
 		}
-		return false
+		c.SSEvent("message", gin.H{
+			"type":    "message",
+			"content": m,
+		})
+		return true
 	})
 }