Skip to content

Commit

Permalink
Refactor Neo message handling to use gin.ResponseWriter and improve e…
Browse files Browse the repository at this point in the history
…rror management
  • Loading branch information
trheyi committed Nov 7, 2024
1 parent 15fb961 commit 6450fd4
Show file tree
Hide file tree
Showing 3 changed files with 174 additions and 134 deletions.
16 changes: 13 additions & 3 deletions neo/message/json.go
Original file line number Diff line number Diff line change
@@ -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"
Expand Down Expand Up @@ -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
}
Expand Down Expand Up @@ -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 {
Expand All @@ -205,7 +215,7 @@ func (json *JSON) Write(w io.Writer) bool {
log.Error("%s", err.Error())
return false
}

w.Flush()
return true
}

Expand Down
Loading

0 comments on commit 6450fd4

Please sign in to comment.