diff --git a/neo/api.go b/neo/api.go index c0d80a1a..7f691656 100644 --- a/neo/api.go +++ b/neo/api.go @@ -58,10 +58,7 @@ func (neo *DSL) handleChat(c *gin.Context) { content := c.Query("content") if content == "" { - msg := message.New().Map(map[string]interface{}{ - "error": "content is required", - "done": true, - }) + msg := message.New().Error("content is required").Done() msg.Write(c.Writer) return } diff --git a/neo/message/json.go b/neo/message/json.go index 20fa72f8..767858af 100644 --- a/neo/message/json.go +++ b/neo/message/json.go @@ -7,6 +7,7 @@ import ( "github.com/gin-gonic/gin" jsoniter "github.com/json-iterator/go" "github.com/yaoapp/gou/helper" + "github.com/yaoapp/kun/exception" "github.com/yaoapp/kun/log" "github.com/yaoapp/kun/maps" "github.com/yaoapp/yao/openai" @@ -80,6 +81,16 @@ func (json *JSON) Text(text string) *JSON { return json } +// Error set the error +func (json *JSON) Error(message interface{}) *JSON { + if err, ok := message.(error); ok { + json.Message.Error = err.Error() + } else if msg, ok := message.(string); ok { + json.Message.Error = msg + } + return json +} + // Map set from map func (json *JSON) Map(msg map[string]interface{}) *JSON { if msg == nil { @@ -90,6 +101,14 @@ func (json *JSON) Map(msg map[string]interface{}) *JSON { json.Message.Text = text } + if err, ok := msg["error"].(string); ok { + json.Message.Error = err + } + + if err, ok := msg["error"].(error); ok { + json.Message.Error = err.Error() + } + if done, ok := msg["done"].(bool); ok { json.Message.Done = done } @@ -203,8 +222,8 @@ func (json *JSON) Write(w gin.ResponseWriter) bool { } }() - if json.Error != "" { - json.writeError(w, json.Error) + if json.Message != nil && json.Message.Error != "" { + json.writeError(w, json.Message.Error) return false } @@ -232,7 +251,10 @@ func (json *JSON) Append(content []byte) []byte { } func (json *JSON) writeError(w gin.ResponseWriter, message string) { - data := []byte(`{"text":"` + strings.Trim(message, "\"") + `"}`) + data := []byte(`{"text":"` + strings.Trim(exception.New(message, 500).Message, "\"") + `","type":"error"}`) + if json.Message.Done { + data = []byte(`{"text":"` + strings.Trim(exception.New(message, 500).Message, "\"") + `","type":"error", "done":true}`) + } data = append([]byte("data: "), data...) data = append(data, []byte("\n\n")...) _, err := w.Write(data) diff --git a/neo/neo.go b/neo/neo.go index dc082a1a..a1dd7374 100644 --- a/neo/neo.go +++ b/neo/neo.go @@ -24,20 +24,14 @@ var lock sync.Mutex = sync.Mutex{} func (neo *DSL) Answer(ctx Context, question string, c *gin.Context) error { messages, err := neo.chatMessages(ctx, question) if err != nil { - msg := message.New().Map(map[string]interface{}{ - "error": err.Error(), - "done": true, - }) + msg := message.New().Error(err).Done() msg.Write(c.Writer) return err } err = neo.HookCreate(ctx, messages, c) if err != nil { - msg := message.New().Map(map[string]interface{}{ - "error": err.Error(), - "done": true, - }) + msg := message.New().Error(err).Done() msg.Write(c.Writer) return err } @@ -201,7 +195,7 @@ func (neo *DSL) send(ctx Context, msg *message.JSON, messages []map[string]inter w := c.Writer - if msg.Error != "" { + if msg.Message != nil && msg.Message.Error != "" { msg.Write(w) return nil }