Merge pull request #779 from trheyi/main

Refactor Neo message handling to use gin.ResponseWriter and improve e…

neo/message/json.go

@@ -1,9 +1,9 @@
 package message
 
 import (
-	"io"
 	"strings"
 
+	"github.com/gin-gonic/gin"
 	jsoniter "github.com/json-iterator/go"
 	"github.com/yaoapp/gou/helper"
 	"github.com/yaoapp/kun/log"
@@ -47,6 +47,10 @@ func NewOpenAI(data []byte) *JSON {
 		msg.Done = true
 		break
 
+	case strings.Contains(text, `"finish_reason":"stop"`):
+		msg.Done = true
+		break
+
 	default:
 		msg.Error = text
 	}
@@ -189,7 +193,13 @@ func (json *JSON) IsDone() bool {
 }
 
 // Write the message
-func (json *JSON) Write(w io.Writer) bool {
+func (json *JSON) Write(w gin.ResponseWriter) bool {
+
+	defer func() {
+		if r := recover(); r != nil {
+			log.Error("Write JSON Message Error: %s", r)
+		}
+	}()
 
 	data, err := jsoniter.Marshal(json.Message)
 	if err != nil {
@@ -205,7 +215,7 @@ func (json *JSON) Write(w io.Writer) bool {
 		log.Error("%s", err.Error())
 		return false
 	}
-
+	w.Flush()
 	return true
 }