Merge pull request #856 from trheyi/main

Enhance tool call handling and message processing

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 = "<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
 }
 

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

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
+	Type  string                 `json:"type"`  // text, function, error, think, tool
-	ID        string                 `json:"id"`                  // the id of the content
+	ID    string                 `json:"id"`    // the id of the content
-	Function  string                 `json:"function"`            // the function name
+	Bytes []byte                 `json:"bytes"` // the content bytes
-	Bytes     []byte                 `json:"bytes"`               // the content bytes
+	Props map[string]interface{} `json:"props"` // the props
-	Arguments []byte                 `json:"arguments,omitempty"` // the function arguments
-	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)
 }

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
+			id := toolCalls.Choices[0].Delta.ToolCalls[0].ID
-			msg.Props["function"] = toolCalls.Choices[0].Delta.ToolCalls[0].Function.Name
+			function := toolCalls.Choices[0].Delta.ToolCalls[0].Function.Name
-			msg.Text = toolCalls.Choices[0].Delta.ToolCalls[0].Function.Arguments
+			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