From 142b07b0ce02e385ff9f4e8c9fa0f5c54f3fb0b4 Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 8 Feb 2025 18:09:26 +0800 Subject: [PATCH] Enhance tool call handling and message processing - Add support for native tool call streaming with XML-like tags - Implement stringHash for message name anonymization - Add debug environment variable for printing request messages - Remove unused function and arguments fields in message contents - Modify message type handling for tool calls in streaming chat --- neo/assistant/api.go | 48 ++++++++++++++++++++++++++-- neo/assistant/utils.go | 9 ++++++ neo/message/contents.go | 71 +++-------------------------------------- neo/message/message.go | 37 +++++++-------------- 4 files changed, 71 insertions(+), 94 deletions(-) diff --git a/neo/assistant/api.go b/neo/assistant/api.go index 80d84bce..ca66df6e 100644 --- a/neo/assistant/api.go +++ b/neo/assistant/api.go @@ -4,11 +4,13 @@ import ( "context" "encoding/base64" "fmt" + "os" "strings" "github.com/gin-gonic/gin" jsoniter "github.com/json-iterator/go" "github.com/yaoapp/gou/fs" + "github.com/yaoapp/kun/utils" chatctx "github.com/yaoapp/yao/neo/context" chatMessage "github.com/yaoapp/yao/neo/message" ) @@ -295,6 +297,9 @@ func (ast *Assistant) streamChat( isFirst := true isFirstThink := true isThinking := false + + isFirstTool := true + isTool := false currentMessageID := "" err := ast.Chat(c.Request.Context(), messages, options, func(data []byte) int { select { @@ -350,6 +355,32 @@ func (ast *Assistant) streamChat( contents.ClearToken() } + // for native tool_calls response + if msg.Type == "tool_calls_native" { + if isFirstTool { + msg.Text = "\n" + msg.Text // add the tool_calls begin tag + isFirstTool = false + isTool = true + } + } + + // for tool response + if isTool && msg.Type != "tool_calls_native" { + + if msg.IsDone { + end := chatMessage.New().Map(map[string]interface{}{"text": "}\n\n", "type": "tool", "delta": true}) + end.Write(c.Writer) + end.ID = currentMessageID + end.AppendTo(contents) + contents.UpdateType("tool", map[string]interface{}{"text": contents.Text()}, currentMessageID) + isTool = false + } else { + msg.Text = "\n\n" + msg.Text // add the tool_calls close tag + } + + isTool = false + } + delta := msg.String() // Chunk the delta @@ -404,9 +435,14 @@ func (ast *Assistant) streamChat( // ------------------------------------------------------------------------------ // Write the message to the stream + msgType := msg.Type + if msgType == "tool_calls_native" { + msgType = "tool" + } + output := chatMessage.New().Map(map[string]interface{}{ "text": delta, - "type": msg.Type, + "type": msgType, "done": msg.IsDone, "delta": true, }) @@ -691,7 +727,7 @@ func (ast *Assistant) requestMessages(ctx context.Context, messages []chatMessag } if name := message.Name; name != "" { - newMessage["name"] = name + newMessage["name"] = stringHash(name) } // Special handling for user messages with JSON content last message @@ -726,6 +762,14 @@ func (ast *Assistant) requestMessages(ctx context.Context, messages []chatMessag newMessages = append(newMessages, newMessage) } + + // For debug environment, print the request messages + if os.Getenv("YAO_AGENT_PRINT_REQUEST_MESSAGES") == "true" { + fmt.Println("--------------------------------") + utils.Dump(newMessages) + fmt.Println("--------------------------------") + } + return newMessages, nil } diff --git a/neo/assistant/utils.go b/neo/assistant/utils.go index aa7c6720..c3114269 100644 --- a/neo/assistant/utils.go +++ b/neo/assistant/utils.go @@ -1,6 +1,8 @@ package assistant import ( + "crypto/sha256" + "encoding/hex" "fmt" "strconv" "time" @@ -42,3 +44,10 @@ func timeToMySQLFormat(ts int64) string { } return time.Unix(ts/1e9, ts%1e9).Format("2006-01-02 15:04:05") } + +// stringHash returns the sha256 hash of the string +func stringHash(v string) string { + h := sha256.New() + h.Write([]byte(v)) + return hex.EncodeToString(h.Sum(nil)) +} diff --git a/neo/message/contents.go b/neo/message/contents.go index 05416857..6e0631af 100644 --- a/neo/message/contents.go +++ b/neo/message/contents.go @@ -31,12 +31,10 @@ type Contents struct { // Data the data of the content type Data struct { - Type string `json:"type"` // text, function, error, think, tool - 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,omitempty"` // the function arguments - Props map[string]interface{} `json:"props"` // the props + Type string `json:"type"` // text, function, error, think, tool + ID string `json:"id"` // the id of the content + Bytes []byte `json:"bytes"` // the content bytes + Props map[string]interface{} `json:"props"` // the props } // NewContents create a new contents @@ -117,22 +115,6 @@ func (c *Contents) NewText(bytes []byte, id ...string) *Contents { return c } -// NewTool create a new tool data and append to the contents -func (c *Contents) NewTool(function string, arguments []byte, id ...string) *Contents { - - data := Data{ - Type: "tool", - Function: function, - Arguments: arguments, - } - if len(id) > 0 && id[0] != "" { - data.ID = id[0] - } - c.Data = append(c.Data, data) - c.Current++ - return c -} - // NewType create a new type data and append to the contents func (c *Contents) NewType(typ string, props map[string]interface{}, id ...string) *Contents { @@ -163,15 +145,6 @@ func (c *Contents) UpdateType(typ string, props map[string]interface{}, id ...st return c } -// SetToolID set the id of the current tool content -func (c *Contents) SetToolID(id string) *Contents { - if c.Current == -1 { - c.NewTool("", []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{ @@ -196,16 +169,6 @@ func (c *Contents) AppendText(bytes []byte, id ...string) *Contents { return c } -// AppendTool append the tool to the current content -func (c *Contents) AppendTool(arguments []byte, id ...string) *Contents { - if c.Current == -1 { - c.NewTool("", arguments, id...) - 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 { @@ -254,19 +217,6 @@ func (data *Data) Map() (map[string]interface{}, error) { v["props"] = data.Props } - if data.Arguments != nil && len(data.Arguments) > 0 { - 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 v, nil } @@ -287,18 +237,5 @@ func (data *Data) MarshalJSON() ([]byte, error) { v["props"] = data.Props } - if data.Arguments != nil && len(data.Arguments) > 0 { - 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 3a05d6ea..32bcaa41 100644 --- a/neo/message/message.go +++ b/neo/message/message.go @@ -191,7 +191,8 @@ func NewAny(content interface{}) (*Message, error) { func NewOpenAI(data []byte, isThinking bool) *Message { // For Debug - if os.Getenv("YAO_AGENT_DEBUG") == "true" { + // For debug environment, print the response data + if os.Getenv("YAO_AGENT_PRINT_RESPONSE_DATA") == "true" { fmt.Printf("%s\n", string(data)) } @@ -217,9 +218,14 @@ func NewOpenAI(data []byte, isThinking bool) *Message { msg.Type = "tool_calls_native" if len(toolCalls.Choices) > 0 && len(toolCalls.Choices[0].Delta.ToolCalls) > 0 { - msg.Props["id"] = toolCalls.Choices[0].Delta.ToolCalls[0].ID - msg.Props["function"] = toolCalls.Choices[0].Delta.ToolCalls[0].Function.Name - msg.Text = toolCalls.Choices[0].Delta.ToolCalls[0].Function.Arguments + id := toolCalls.Choices[0].Delta.ToolCalls[0].ID + function := toolCalls.Choices[0].Delta.ToolCalls[0].Function.Name + arguments := toolCalls.Choices[0].Delta.ToolCalls[0].Function.Arguments + text := arguments + if id != "" { + text = fmt.Sprintf(`{"id": "%s", "function": "%s", "arguments": %s`, id, function, arguments) + } + msg.Text = text } case strings.Contains(text, `"delta":{`) && strings.Contains(text, `"content":`): @@ -322,7 +328,7 @@ func (m *Message) String() string { } switch typ { - case "text", "think", "tool": + case "text", "think", "tool", "tool_calls_native": return m.Text case "error": @@ -392,7 +398,7 @@ func (m *Message) AppendTo(contents *Contents) *Message { } switch m.Type { - case "text", "think", "tool": + case "text", "think", "tool", "tool_calls_native": if m.Text != "" { if m.IsNew { contents.NewText([]byte(m.Text), m.ID) @@ -403,25 +409,6 @@ func (m *Message) AppendTo(contents *Contents) *Message { } return m - case "tool_calls_native": - - // Set function name - new := false - if name, ok := m.Props["tool"].(string); ok && name != "" { - contents.NewTool(name, []byte(m.Text)) - new = true - } - - // Set id - if id, ok := m.Props["id"].(string); ok && id != "" { - contents.SetToolID(id) - } - - if !new { - contents.AppendTool([]byte(m.Text)) - } - return m - case "loading", "error", "action": // Ignore loading, action and error messages return m