Merge pull request #397 from trheyi/main

[add] Neo optimizing response data

neo/command/request.go

@@ -5,27 +5,27 @@ import (
 	"time"
 
 	"github.com/google/uuid"
+	"github.com/yaoapp/yao/neo/message"
 )
 
-func output(format string, args ...interface{}) []byte {
+// Run the command
-	content := fmt.Sprintf(format, args...)
+func (req *Request) Run(messages []map[string]interface{}, cb func(msg *message.JSON) int) (interface{}, error) {
-	return []byte(fmt.Sprintf(`{"id":"chatcmpl-7Atx502nGBuYcvoZfIaWU4FREI1mT","object":"chat.completion.chunk","created":1682832715,"model":"gpt-3.5-turbo-0301","choices":[{"delta":{"content":"%s"},"index":0,"finish_reason":null}]}`, content))
+
+	cb(req.msg().Text(fmt.Sprintf("- Command: %s\n", req.Command.Name)))
+	time.Sleep(200 * time.Millisecond)
+
+	cb(req.msg().Text(fmt.Sprintf("- Session: %s\n", req.sid)))
+	time.Sleep(200 * time.Millisecond)
+
+	cb(req.msg().Text(fmt.Sprintf("- Request: %s\n", req.sid)))
+	time.Sleep(200 * time.Millisecond)
+
+	cb(req.msg().Done())
+	return nil, nil
 }
 
-// Run the command
+func (req *Request) msg() *message.JSON {
-func (req *Request) Run(messages []map[string]interface{}, cb func(data []byte) int) (interface{}, error) {
+	return message.New().Command(req.Command.Name, req.Command.ID, req.id)
-
-	cb(output("- Command: %s\\n", req.Command.ID))
-	time.Sleep(200 * time.Millisecond)
-
-	cb(output("- Session: %s\\n", req.sid))
-	time.Sleep(200 * time.Millisecond)
-
-	cb(output("- Request: %s\\n", req.id))
-	time.Sleep(200 * time.Millisecond)
-
-	cb([]byte(`[DONE]`))
-	return nil, nil
 }
 
 // NewRequest create a new request

neo/message/json.go

@@ -0,0 +1,123 @@
+package message
+
+import (
+	"io"
+	"strings"
+
+	jsoniter "github.com/json-iterator/go"
+	"github.com/yaoapp/kun/log"
+	"github.com/yaoapp/yao/openai"
+)
+
+// JSON the JSON message
+type JSON struct{ *Message }
+
+// New create a new JSON message
+func New() *JSON {
+	return &JSON{makeMessage()}
+}
+
+// NewOpenAI create a new JSON message
+func NewOpenAI(data []byte) *JSON {
+
+	if data == nil || len(data) == 0 {
+		return nil
+	}
+
+	msg := makeMessage()
+	data = []byte(strings.TrimPrefix(string(data), "data: "))
+	switch {
+
+	case strings.Contains(string(data), `"delta":{"content"`):
+		var message openai.Message
+		err := jsoniter.Unmarshal(data, &message)
+		if err != nil {
+			msg.Text = err.Error()
+			return &JSON{msg}
+		}
+
+		if len(message.Choices) > 0 {
+			msg.Text = message.Choices[0].Delta.Content
+		}
+		break
+
+	case strings.Contains(string(data), `[DONE]`):
+		msg.Done = true
+		break
+
+	default:
+		return nil
+	}
+
+	return &JSON{msg}
+}
+
+// Text set the text
+func (json *JSON) Text(text string) *JSON {
+	json.Message.Text = text
+	return json
+}
+
+// Done set the done
+func (json *JSON) Done() *JSON {
+	json.Message.Done = true
+	return json
+}
+
+// Confirm set the confirm
+func (json *JSON) Confirm() *JSON {
+	json.Message.Confirm = true
+	return json
+}
+
+// Command set the command
+func (json *JSON) Command(name, id, request string) *JSON {
+	json.Message.Command = &Command{
+		ID:      id,
+		Name:    name,
+		Reqeust: request,
+	}
+	return json
+}
+
+// Action set the action
+func (json *JSON) Action(name string, t string, payload interface{}, next string) *JSON {
+	json.Message.Actions = append(json.Message.Actions, Action{
+		Name:    name,
+		Type:    t,
+		Payload: payload,
+		Next:    next,
+	})
+	return json
+}
+
+// IsDone check if the message is done
+func (json *JSON) IsDone() bool {
+	return json.Message.Done
+}
+
+// Write the message
+func (json *JSON) Write(w io.Writer) bool {
+
+	data, err := jsoniter.Marshal(json.Message)
+	if err != nil {
+		log.Error("%s", err.Error())
+		return false
+	}
+
+	data = append([]byte("data: "), data...)
+	data = append(data, []byte("\n\n")...)
+
+	_, err = w.Write(data)
+	if err != nil {
+		log.Error("%s", err.Error())
+		return false
+	}
+
+	return true
+}
+
+// Append the message
+func (json *JSON) Append(content []byte) []byte {
+	return append(content, []byte(json.Message.Text)...)
+}

neo/message/message.go

@@ -0,0 +1,6 @@
+package message
+
+// makeMessage create a new message
+func makeMessage() *Message {
+	return &Message{Actions: []Action{}}
+}

neo/message/types.go

@@ -0,0 +1,25 @@
+package message
+
+// Message the message
+type Message struct {
+	Text    string   `json:"text,omitempty"`
+	Done    bool     `json:"done,omitempty"`
+	Confirm bool     `json:"confirm,omitempty"`
+	Command *Command `json:"command,omitempty"`
+	Actions []Action `json:"actions,omitempty"`
+}
+
+// Action the action
+type Action struct {
+	Name    string      `json:"name,omitempty"`
+	Type    string      `json:"type"`
+	Payload interface{} `json:"payload,omitempty"`
+	Next    string      `json:"next,omitempty"`
+}
+
+// Command the command
+type Command struct {
+	ID      string `json:"id,omitempty"`
+	Name    string `json:"name,omitempty"`
+	Reqeust string `json:"request,omitempty"`
+}

neo/neo.go

@@ -8,7 +8,6 @@ import (
 
 	"github.com/gin-gonic/gin"
 	"github.com/google/uuid"
-	jsoniter "github.com/json-iterator/go"
 	"github.com/yaoapp/gou/api"
 	"github.com/yaoapp/gou/connector"
 	"github.com/yaoapp/gou/process"
@@ -17,6 +16,7 @@ import (
 	"github.com/yaoapp/yao/neo/command"
 	"github.com/yaoapp/yao/neo/command/query"
 	"github.com/yaoapp/yao/neo/conversation"
+	"github.com/yaoapp/yao/neo/message"
 	"github.com/yaoapp/yao/openai"
 )
 
@@ -84,18 +84,12 @@ func (neo *DSL) API(router *gin.Engine, path string) error {
 // Answer the message
 func (neo *DSL) Answer(ctx command.Context, answer Answer, messages []map[string]interface{}) error {
 
-	chanStream := make(chan []byte, 1)
+	chanStream := make(chan *message.JSON, 1)
 	chanError := make(chan error, 1)
+	content := []byte{}
 
 	// check the command
-	var cmd *command.Command
+	cmd, isCommand := neo.matchCommand(ctx, messages)
-	var isCommand = false
-	input := messages[len(messages)-1]["content"].(string)
-	name, err := command.Match(ctx.Sid, query.Param{Stack: ctx.Stack, Path: ctx.Path}, input)
-	if err == nil && name != "" {
-		cmd, isCommand = command.Commands[name]
-	}
-
 	go func() {
 		defer func() {
 			close(chanStream)
@@ -111,8 +105,8 @@ func (neo *DSL) Answer(ctx command.Context, answer Answer, messages []map[string
 				return
 			}
 
-			_, err = req.Run(messages, func(data []byte) int {
+			_, err = req.Run(messages, func(msg *message.JSON) int {
-				chanStream <- data
+				chanStream <- msg
 				return 1
 			})
 
@@ -125,19 +119,77 @@ func (neo *DSL) Answer(ctx command.Context, answer Answer, messages []map[string
 
 		// chat with AI
 		_, ex := neo.AI.ChatCompletionsWith(ctx, messages, neo.Option, func(data []byte) int {
-			chanStream <- data
+			chanStream <- message.NewOpenAI(data)
 			return 1
 		})
 
 		if ex != nil {
 			chanError <- fmt.Errorf("AI chat error: %s", ex.Message)
 		}
+
+		defer neo.saveHistory(ctx.Sid, content, messages)
+
 	}()
 
-	// save the history
+	answer.Header("Content-Type", "text/event-stream;charset=utf-8")
-	content := []byte{}
+	ok := answer.Stream(func(w io.Writer) bool {
-	defer func() {
+		select {
-		sid := answer.GetString("__sid")
+		case err := <-chanError:
+			if err != nil {
+				message.New().Text(err.Error()).Write(w)
+			}
+
+			message.New().Done().Write(w)
+			return false
+
+		case msg := <-chanStream:
+			if msg == nil {
+				return true
+			}
+			msg.Write(w)
+			content = msg.Append(content)
+			return !msg.IsDone()
+
+		case <-ctx.Done():
+			if err := ctx.Err(); err != nil {
+				message.New().Text(err.Error()).Write(w)
+			}
+			message.New().Done().Write(w)
+			return false
+		}
+	})
+
+	if !ok {
+		answer.Status(500)
+		return nil
+	}
+
+	answer.Status(200)
+	return nil
+}
+
+func (neo *DSL) matchCommand(ctx command.Context, messages []map[string]interface{}) (*command.Command, bool) {
+	if len(messages) < 1 {
+		return nil, false
+	}
+
+	input, ok := messages[len(messages)-1]["content"].(string)
+	if !ok {
+		return nil, false
+	}
+
+	name, err := command.Match(ctx.Sid, query.Param{Stack: ctx.Stack, Path: ctx.Path}, input)
+	if err == nil && name != "" {
+		cmd, isCommand := command.Commands[name]
+		return cmd, isCommand
+	}
+
+	return nil, false
+}
+
+// saveHistory save the history
+func (neo *DSL) saveHistory(sid string, content []byte, messages []map[string]interface{}) {
+
 	if len(content) > 0 && sid != "" && len(messages) > 0 {
 		err := neo.Conversation.SaveHistory(
 			sid,
@@ -151,56 +203,6 @@ func (neo *DSL) Answer(ctx command.Context, answer Answer, messages []map[string
 			log.Error("Save history error: %s", err.Error())
 		}
 	}
-	}()
-
-	answer.Header("Content-Type", "text/event-stream;charset=utf-8")
-	ok := answer.Stream(func(w io.Writer) bool {
-		select {
-		case err := <-chanError:
-			if err != nil {
-				w.Write([]byte(fmt.Sprintf(`data: {"text":"%s"}%s`, err.Error(), "\n\n")))
-			}
-			w.Write([]byte(fmt.Sprintf("data: %s\n\n", `{"done":true}`)))
-			return false
-
-		case msg := <-chanStream:
-			if msg != nil && len(msg) > 0 {
-
-				if strings.Contains(string(msg), `"delta":{"content"`) {
-					msg = []byte(strings.TrimPrefix(string(msg), "data: "))
-					var message openai.Message
-					err := jsoniter.Unmarshal(msg, &message)
-
-					if err != nil {
-						data, _ := jsoniter.Marshal(map[string]interface{}{"text": err.Error()})
-						w.Write([]byte(fmt.Sprintf("data: %s\n\n", data)))
-						return true
-					}
-
-					if len(message.Choices) > 0 {
-						text := message.Choices[0].Delta.Content
-						content = append(content, []byte(text)...)
-						data, _ := jsoniter.Marshal(map[string]interface{}{"text": text})
-						w.Write([]byte(fmt.Sprintf("data: %s\n\n", data)))
-						return true
-					}
-
-				} else if strings.Contains(string(msg), `[DONE]`) {
-					w.Write([]byte(fmt.Sprintf("data: %s\n\n", `{"done":true}`)))
-					return true
-				}
-			}
-			return true
-		}
-	})
-
-	if !ok {
-		answer.Status(500)
-		return nil
-	}
-
-	answer.Status(200)
-	return nil
 }
 
 func (neo *DSL) crossDomain(router *gin.Engine, path string) {

neo/types.go

@@ -29,7 +29,6 @@ type Conversation interface {
 
 // Answer the answer interface
 type Answer interface {
-	GetString(key string) (s string)
 	Stream(func(w io.Writer) bool) bool
 	Status(code int)
 	Header(key, value string)