From 5d5fa4c2703356e6d6e80ac2e4589d2ea5aa8b67 Mon Sep 17 00:00:00 2001 From: Max Date: Sun, 19 Jan 2025 12:23:23 +0800 Subject: [PATCH] Refactor Neo API assistant message handling and content structure - Replaced the Content struct with a new Contents struct to improve message handling and organization. - Updated methods across the Assistant and Message structs to utilize the new Contents structure, enhancing the clarity and efficiency of message processing. - Improved error handling by ensuring proper message appending and content management during chat streaming. - Removed the deprecated Content struct, streamlining the codebase and reducing complexity. These changes enhance the maintainability and robustness of the Neo API assistant, paving the way for improved message handling and assistant functionalities. --- neo/assistant/api.go | 48 ++++++------- neo/assistant/hooks.go | 22 ++++-- neo/assistant/load.go | 4 ++ neo/message/content.go | 100 --------------------------- neo/message/contents.go | 150 ++++++++++++++++++++++++++++++++++++++++ neo/message/message.go | 41 +++++++++-- neo/neo.go | 6 +- 7 files changed, 230 insertions(+), 141 deletions(-) delete mode 100644 neo/message/content.go create mode 100644 neo/message/contents.go diff --git a/neo/assistant/api.go b/neo/assistant/api.go index 41f2b0d4..2b5611d6 100644 --- a/neo/assistant/api.go +++ b/neo/assistant/api.go @@ -56,6 +56,11 @@ func (ast *Assistant) Execute(c *gin.Context, ctx chatctx.Context, input string, // Run init hook res, err := ast.HookInit(c, ctx, messages, options) if err != nil { + chatMessage.New(). + Assistant(ast.ID, ast.Name, ast.Avatar). + Error(err). + Done(). + Write(c.Writer) return err } @@ -63,6 +68,11 @@ func (ast *Assistant) Execute(c *gin.Context, ctx chatctx.Context, input string, if res != nil && res.AssistantID != ctx.AssistantID { newAst, err := Get(res.AssistantID) if err != nil { + chatMessage.New(). + Assistant(ast.ID, ast.Name, ast.Avatar). + Error(err). + Done(). + Write(c.Writer) return err } *ast = *newAst @@ -163,16 +173,16 @@ func (next *NextAction) Execute(c *gin.Context, ctx chatctx.Context) error { func (ast *Assistant) handleChatStream(c *gin.Context, ctx chatctx.Context, messages []chatMessage.Message, options map[string]interface{}) error { clientBreak := make(chan bool, 1) done := make(chan bool, 1) - content := chatMessage.NewContent("text") + contents := chatMessage.NewContents() // Chat with AI in background go func() { - err := ast.streamChat(c, ctx, messages, options, clientBreak, done, content) + err := ast.streamChat(c, ctx, messages, options, clientBreak, done, contents) if err != nil { chatMessage.New().Error(err).Done().Write(c.Writer) } - ast.saveChatHistory(ctx, messages, content) + ast.saveChatHistory(ctx, messages, contents) done <- true }() @@ -194,7 +204,7 @@ func (ast *Assistant) streamChat( options map[string]interface{}, clientBreak chan bool, done chan bool, - content *chatMessage.Content) error { + contents *chatMessage.Contents) error { return ast.Chat(c.Request.Context(), messages, options, func(data []byte) int { select { @@ -210,7 +220,7 @@ func (ast *Assistant) streamChat( // Handle error if msg.Type == "error" { value := msg.String() - res, hookErr := ast.HookFail(c, ctx, messages, content.String(), fmt.Errorf("%s", value)) + res, hookErr := ast.HookFail(c, ctx, messages, contents.JSON(), fmt.Errorf("%s", value)) if hookErr == nil && res != nil && (res.Output != "" || res.Error != "") { value = res.Output if res.Error != "" { @@ -221,26 +231,12 @@ func (ast *Assistant) streamChat( return 0 // break } - // Handle tool call - if msg.Type == "tool_calls" { - content.SetType("function") // Set type to function - // Set id - if id, ok := msg.Props["id"].(string); ok && id != "" { - content.SetID(id) - } - - // Set name - if name, ok := msg.Props["name"].(string); ok && name != "" { - content.SetName(name) - } - } - // Append content and send message + msg.Append(contents) value := msg.String() - content.Append(value) if value != "" { // Handle stream - res, err := ast.HookStream(c, ctx, messages, content.String(), content.Type == "function") + res, err := ast.HookStream(c, ctx, messages, contents.JSON()) if err == nil && res != nil { if res.Output != "" { value = res.Output @@ -278,9 +274,7 @@ func (ast *Assistant) streamChat( msg.Write(c.Writer) } - // Call HookDone - content.SetStatus(chatMessage.ContentStatusDone) - res, hookErr := ast.HookDone(c, ctx, messages, content.String(), content.Type == "function") + res, hookErr := ast.HookDone(c, ctx, messages, contents.JSON()) if hookErr == nil && res != nil { if res.Output != "" { chatMessage.New(). @@ -322,8 +316,8 @@ func (ast *Assistant) streamChat( } // saveChatHistory saves the chat history if storage is available -func (ast *Assistant) saveChatHistory(ctx chatctx.Context, messages []chatMessage.Message, content *chatMessage.Content) { - if len(content.Bytes) > 0 && ctx.Sid != "" && len(messages) > 0 { +func (ast *Assistant) saveChatHistory(ctx chatctx.Context, messages []chatMessage.Message, contents *chatMessage.Contents) { + if len(contents.Data) > 0 && ctx.Sid != "" && len(messages) > 0 { userMessage := messages[len(messages)-1] data := []map[string]interface{}{ { @@ -333,7 +327,7 @@ func (ast *Assistant) saveChatHistory(ctx chatctx.Context, messages []chatMessag }, { "role": "assistant", - "content": content.String(), + "content": contents.JSON(), "name": ctx.Sid, "assistant_id": ast.ID, "assistant_name": ast.Name, diff --git a/neo/assistant/hooks.go b/neo/assistant/hooks.go index bf21fa25..f422b60b 100644 --- a/neo/assistant/hooks.go +++ b/neo/assistant/hooks.go @@ -6,6 +6,7 @@ import ( "time" "github.com/gin-gonic/gin" + jsoniter "github.com/json-iterator/go" chatctx "github.com/yaoapp/yao/neo/context" "github.com/yaoapp/yao/neo/message" ) @@ -34,6 +35,17 @@ func (ast *Assistant) HookInit(c *gin.Context, context chatctx.Context, input [] response.ChatID = res } + // input + if input, has := v["input"]; has { + raw, _ := jsoniter.MarshalToString(input) + vv := []message.Message{} + err := jsoniter.UnmarshalFromString(raw, &vv) + if err != nil { + return nil, err + } + response.Input = vv + } + if res, ok := v["next"].(map[string]interface{}); ok { response.Next = &NextAction{} if name, ok := res["action"].(string); ok { @@ -57,13 +69,13 @@ func (ast *Assistant) HookInit(c *gin.Context, context chatctx.Context, input [] } // HookStream Handle streaming response from LLM -func (ast *Assistant) HookStream(c *gin.Context, context chatctx.Context, input []message.Message, output string, toolcall bool) (*ResHookStream, error) { +func (ast *Assistant) HookStream(c *gin.Context, context chatctx.Context, input []message.Message, output string) (*ResHookStream, error) { // Create timeout context ctx, cancel := ast.createTimeoutContext(c) defer cancel() - v, err := ast.call(ctx, "Stream", context, input, output, toolcall, c.Writer) + v, err := ast.call(ctx, "Stream", context, input, output, c.Writer) if err != nil { if err.Error() == HookErrorMethodNotFound { return nil, nil @@ -100,12 +112,12 @@ func (ast *Assistant) HookStream(c *gin.Context, context chatctx.Context, input } // HookDone Handle completion of assistant response -func (ast *Assistant) HookDone(c *gin.Context, context chatctx.Context, input []message.Message, output string, toolcall bool) (*ResHookDone, error) { +func (ast *Assistant) HookDone(c *gin.Context, context chatctx.Context, input []message.Message, output string) (*ResHookDone, error) { // Create timeout context ctx, cancel := ast.createTimeoutContext(c) defer cancel() - v, err := ast.call(ctx, "Done", context, input, output, toolcall, c.Writer) + v, err := ast.call(ctx, "Done", context, input, output, c.Writer) if err != nil { if err.Error() == HookErrorMethodNotFound { return nil, nil @@ -185,7 +197,7 @@ func (ast *Assistant) HookFail(c *gin.Context, context chatctx.Context, input [] // createTimeoutContext creates a timeout context with 5 seconds timeout func (ast *Assistant) createTimeoutContext(c *gin.Context) (context.Context, context.CancelFunc) { - ctx, cancel := context.WithTimeout(c.Request.Context(), 5*time.Second) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) return ctx, cancel } diff --git a/neo/assistant/load.go b/neo/assistant/load.go index 9cbd7dff..13563acb 100644 --- a/neo/assistant/load.go +++ b/neo/assistant/load.go @@ -28,6 +28,10 @@ var defaultConnector string = "" // default connector // LoadBuiltIn load the built-in assistants func LoadBuiltIn() error { + + // Clear the cache + loaded.Clear() + root := `/assistants` app, err := fs.Get("app") if err != nil { diff --git a/neo/message/content.go b/neo/message/content.go deleted file mode 100644 index 2d2466c5..00000000 --- a/neo/message/content.go +++ /dev/null @@ -1,100 +0,0 @@ -package message - -import ( - jsoniter "github.com/json-iterator/go" -) - -const ( - // ContentStatusPending the content status pending - ContentStatusPending = iota - // ContentStatusDone the content status done - ContentStatusDone - // ContentStatusError the content status error - ContentStatusError -) - -// Content the content -type Content struct { - ID string `json:"id"` - Name string `json:"name"` - Bytes []byte `json:"bytes"` - Type string `json:"type"` // text, function, error - Status uint8 `json:"status"` // 0: pending, 1: done -} - -// NewContent create a new content -func NewContent(typ string) *Content { - if typ == "" { - typ = "text" - } - - return &Content{ - Bytes: []byte{}, - Type: typ, - Status: ContentStatusPending, - } -} - -// String the content string -func (c *Content) String() string { - if c.Type == "function" { - - var arguments interface{} = string(c.Bytes) - if c.Status == ContentStatusDone { - var vv interface{} = nil - err := jsoniter.Unmarshal(c.Bytes, &vv) - if err != nil { - return "" - } - arguments = vv - } - - data := map[string]interface{}{ - "id": c.ID, - "type": "function", - "text": c.Name, - "props": map[string]interface{}{ - "id": c.ID, - "name": c.Name, - "arguments": arguments, - }, - "function": map[string]interface{}{ - "name": c.Name, - "arguments": arguments, - }, - } - - raw, err := jsoniter.MarshalToString(data) - if err != nil { - return "" - } - - return raw - } - return string(c.Bytes) -} - -// SetID set the content id -func (c *Content) SetID(id string) { - c.ID = id -} - -// SetName set the content name -func (c *Content) SetName(name string) { - c.Name = name -} - -// SetType set the content type -func (c *Content) SetType(typ string) { - c.Type = typ -} - -// Append append the content -func (c *Content) Append(data string) { - c.Bytes = append(c.Bytes, []byte(data)...) -} - -// SetStatus set the content status -func (c *Content) SetStatus(status uint8) { - c.Status = status -} diff --git a/neo/message/contents.go b/neo/message/contents.go new file mode 100644 index 00000000..62ed6851 --- /dev/null +++ b/neo/message/contents.go @@ -0,0 +1,150 @@ +package message + +import ( + jsoniter "github.com/json-iterator/go" +) + +const ( + // ContentStatusPending the content status pending + ContentStatusPending = iota + // ContentStatusDone the content status done + ContentStatusDone + // ContentStatusError the content status error + ContentStatusError +) + +// Contents the contents +type Contents struct { + Current int `json:"current"` // the current content index + Data []Data `json:"data"` // the data +} + +// Data the data of the content +type Data struct { + Type string `json:"type"` // text, function, error, ... + ID string `json:"id"` // the id of the content + Function string `json:"function"` // the function name + Bytes []byte `json:"bytes"` // the content bytes + Arguments []byte `json:"arguments"` // the function arguments +} + +// NewContents create a new contents +func NewContents() *Contents { + return &Contents{ + Current: -1, + Data: []Data{}, + } +} + +// NewText create a new text data and append to the contents +func (c *Contents) NewText(bytes []byte) *Contents { + c.Data = append(c.Data, Data{ + Type: "text", + Bytes: bytes, + }) + c.Current++ + return c +} + +// NewFunction create a new function data and append to the contents +func (c *Contents) NewFunction(function string, arguments []byte) *Contents { + c.Data = append(c.Data, Data{ + Type: "function", + Function: function, + Arguments: arguments, + }) + c.Current++ + return c +} + +// SetFunctionID set the id of the current function content +func (c *Contents) SetFunctionID(id string) *Contents { + if c.Current == -1 { + c.NewFunction("", []byte{}) + } + c.Data[c.Current].ID = id + return c +} + +// NewError create a new error data and append to the contents +func (c *Contents) NewError(err []byte) *Contents { + c.Data = append(c.Data, Data{ + Type: "error", + Bytes: err, + }) + c.Current++ + return c +} + +// AppendText append the text to the current content +func (c *Contents) AppendText(bytes []byte) *Contents { + if c.Current == -1 { + c.NewText(bytes) + return c + } + c.Data[c.Current].Bytes = append(c.Data[c.Current].Bytes, bytes...) + return c +} + +// AppendFunction append the function to the current content +func (c *Contents) AppendFunction(arguments []byte) *Contents { + if c.Current == -1 { + c.NewFunction("", arguments) + return c + } + c.Data[c.Current].Arguments = append(c.Data[c.Current].Arguments, arguments...) + return c +} + +// AppendError append the error to the current content +func (c *Contents) AppendError(err []byte) *Contents { + if c.Current == -1 { + c.NewError(err) + return c + } + c.Data[c.Current].Bytes = append(c.Data[c.Current].Bytes, err...) + return c +} + +// JSON returns the json representation +func (c *Contents) JSON() string { + raw, _ := jsoniter.MarshalToString(c.Data) + return raw +} + +// Text returns the text of the current content +func (c *Contents) Text() string { + if c.Current == -1 { + return "" + } + return string(c.Data[c.Current].Bytes) +} + +// MarshalJSON returns the json representation +func (data *Data) MarshalJSON() ([]byte, error) { + + v := map[string]interface{}{"type": data.Type} + + if data.ID != "" { + v["id"] = data.ID + } + + if data.Bytes != nil { + v["bytes"] = string(data.Bytes) + } + + if data.Arguments != nil { + var vv interface{} = nil + err := jsoniter.Unmarshal(data.Arguments, &vv) + if err != nil { + return nil, err + } + v["arguments"] = vv + } + + if data.Function != "" { + v["function"] = data.Function + } + + return jsoniter.Marshal(v) +} diff --git a/neo/message/message.go b/neo/message/message.go index 39551102..d96f82e4 100644 --- a/neo/message/message.go +++ b/neo/message/message.go @@ -99,7 +99,7 @@ func NewOpenAI(data []byte) *Message { msg.Type = "tool_calls" if len(toolCalls.Choices) > 0 && len(toolCalls.Choices[0].Delta.ToolCalls) > 0 { msg.Props["id"] = toolCalls.Choices[0].Delta.ToolCalls[0].ID - msg.Props["name"] = toolCalls.Choices[0].Delta.ToolCalls[0].Function.Name + msg.Props["function"] = toolCalls.Choices[0].Delta.ToolCalls[0].Function.Name msg.Text = toolCalls.Choices[0].Delta.ToolCalls[0].Function.Arguments } @@ -184,6 +184,32 @@ func (m *Message) SetContent(content string) *Message { return m } +// Append append the contents +func (m *Message) Append(contents *Contents) *Message { + + switch m.Type { + case "text": + if m.Text != "" { + contents.AppendText([]byte(m.Text)) + } + + case "tool_calls": + + // Set function name + if name, ok := m.Props["function"].(string); ok && name != "" { + contents.NewFunction(name, []byte(m.Text)) + } + + // Set id + if id, ok := m.Props["id"].(string); ok && id != "" { + contents.SetFunctionID(id) + } + + contents.AppendFunction([]byte(m.Text)) + } + return m +} + // Content get the content func (m *Message) Content() string { content := map[string]interface{}{"text": m.Text} @@ -291,6 +317,14 @@ func (m *Message) Done() *Message { return m } +// Assistant set the assistant +func (m *Message) Assistant(id string, name string, avatar string) *Message { + m.AssistantID = id + m.AssistantName = name + m.AssistantAvatar = avatar + return m +} + // Action add an action func (m *Message) Action(name string, t string, payload interface{}, next string) *Message { if m.Data != nil { @@ -339,11 +373,6 @@ func (m *Message) Write(w gin.ResponseWriter) bool { return true } -// Append appends content to the byte slice -func (m *Message) Append(content []byte) []byte { - return append(content, []byte(m.Text)...) -} - // WriteError writes an error message to response writer func (m *Message) WriteError(w gin.ResponseWriter, message string) { errMsg := strings.Trim(exception.New(message, 500).Message, "\"") diff --git a/neo/neo.go b/neo/neo.go index c9a45f98..e34a55b7 100644 --- a/neo/neo.go +++ b/neo/neo.go @@ -100,7 +100,7 @@ func (neo *DSL) GenerateWithAI(ctx chatctx.Context, input string, messageType st clientBreak := make(chan bool, 1) done := make(chan bool, 1) fail := make(chan error, 1) - content := []byte{} + contents := message.NewContents() // Chat with AI in background go func() { @@ -126,7 +126,7 @@ func (neo *DSL) GenerateWithAI(ctx chatctx.Context, input string, messageType st } // Append content and send message - content = msg.Append(content) + msg.Append(contents) if !silent { value := msg.String() if value != "" { @@ -166,7 +166,7 @@ func (neo *DSL) GenerateWithAI(ctx chatctx.Context, input string, messageType st // Wait for completion or client disconnect select { case <-done: - return string(content), nil + return contents.Text(), nil case err := <-fail: return "", err case <-c.Writer.CloseNotify():