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
This commit is contained in:
Max 2025-02-08 18:09:26 +08:00
parent a278c4c63a
commit 142b07b0ce
4 changed files with 71 additions and 94 deletions

View file

@ -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 = "<tool>\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</tool>\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</tool>\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
}

View file

@ -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))
}

View file

@ -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)
}

View file

@ -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