From 0b3ebceffa0d542bfb876017ea1882b6782b12c9 Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 4 May 2023 04:59:17 +0800 Subject: [PATCH] [add] Neo optimizing response data --- neo/command/request.go | 34 ++++++------ neo/message/json.go | 123 +++++++++++++++++++++++++++++++++++++++++ neo/message/message.go | 6 ++ neo/message/types.go | 25 +++++++++ neo/neo.go | 118 ++++++++++++++++++++------------------- neo/types.go | 1 - 6 files changed, 231 insertions(+), 76 deletions(-) create mode 100644 neo/message/json.go create mode 100644 neo/message/message.go create mode 100644 neo/message/types.go diff --git a/neo/command/request.go b/neo/command/request.go index aad045af..2357e7d2 100644 --- a/neo/command/request.go +++ b/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 { - content := fmt.Sprintf(format, args...) - 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)) +// Run the command +func (req *Request) Run(messages []map[string]interface{}, cb func(msg *message.JSON) int) (interface{}, error) { + + 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) Run(messages []map[string]interface{}, cb func(data []byte) int) (interface{}, error) { - - 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 +func (req *Request) msg() *message.JSON { + return message.New().Command(req.Command.Name, req.Command.ID, req.id) } // NewRequest create a new request diff --git a/neo/message/json.go b/neo/message/json.go new file mode 100644 index 00000000..cc084830 --- /dev/null +++ b/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)...) +} diff --git a/neo/message/message.go b/neo/message/message.go new file mode 100644 index 00000000..7726dd26 --- /dev/null +++ b/neo/message/message.go @@ -0,0 +1,6 @@ +package message + +// makeMessage create a new message +func makeMessage() *Message { + return &Message{Actions: []Action{}} +} diff --git a/neo/message/types.go b/neo/message/types.go new file mode 100644 index 00000000..cd8323ec --- /dev/null +++ b/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"` +} diff --git a/neo/neo.go b/neo/neo.go index 84064219..c17970af 100644 --- a/neo/neo.go +++ b/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 - 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] - } - + cmd, isCommand := neo.matchCommand(ctx, messages) 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 { - chanStream <- data + _, err = req.Run(messages, func(msg *message.JSON) int { + chanStream <- msg return 1 }) @@ -125,32 +119,16 @@ 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) } - }() - // save the history - content := []byte{} - defer func() { - sid := answer.GetString("__sid") - if len(content) > 0 && sid != "" && len(messages) > 0 { - err := neo.Conversation.SaveHistory( - sid, - []map[string]interface{}{ - {"role": "user", "content": messages[len(messages)-1]["content"], "name": sid}, - {"role": "assistant", "content": string(content), "name": sid}, - }, - ) + defer neo.saveHistory(ctx.Sid, content, messages) - if err != nil { - log.Error("Save history error: %s", err.Error()) - } - } }() answer.Header("Content-Type", "text/event-stream;charset=utf-8") @@ -158,39 +136,26 @@ func (neo *DSL) Answer(ctx command.Context, answer Answer, messages []map[string select { case err := <-chanError: if err != nil { - w.Write([]byte(fmt.Sprintf(`data: {"text":"%s"}%s`, err.Error(), "\n\n"))) + message.New().Text(err.Error()).Write(w) } - w.Write([]byte(fmt.Sprintf("data: %s\n\n", `{"done":true}`))) + + message.New().Done().Write(w) 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 - } + if msg == nil { + return true } - 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 } }) @@ -203,6 +168,43 @@ func (neo *DSL) Answer(ctx command.Context, answer Answer, messages []map[string 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, + []map[string]interface{}{ + {"role": "user", "content": messages[len(messages)-1]["content"], "name": sid}, + {"role": "assistant", "content": string(content), "name": sid}, + }, + ) + + if err != nil { + log.Error("Save history error: %s", err.Error()) + } + } +} + func (neo *DSL) crossDomain(router *gin.Engine, path string) { if len(neo.Allows) == 0 { diff --git a/neo/types.go b/neo/types.go index 4ab97d74..9494c8f0 100644 --- a/neo/types.go +++ b/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)