Merge pull request #863 from trheyi/main

Enhance tool call processing and message handling

neo/assistant/api.go

@@ -620,27 +620,38 @@ func (ast *Assistant) withPrompts(messages []chatMessage.Message) []chatMessage.
 			raw, _ := jsoniter.MarshalToString(ast.Tools.Tools)
 			messages = append(messages, *chatMessage.New().Map(map[string]interface{}{
 				"role":    "system",
+				"name":    "TOOL_CALLS_SCHEMA",
 				"content": raw,
 			}))
 
-			// Add the default system prompts for tool calls
 			messages = append(messages, *chatMessage.New().Map(map[string]interface{}{
 				"role": "system",
-				"content": "## Tool Calls Match Rules:\n" +
+				"name": "TOOL_CALLS",
-					"1. if the user's question is about the tool_calls, just answer one of the tool_calls, do not provide any additional information.\n" +
+				"content": "## Tool Response Format\n" +
-					"2. if the user's question is not about the tool_calls, just answer the user's question directly.\n" +
+					"1. If no matching function exists in TOOL_CALLS_SCHEMA, respond normally without using tool calls\n" +
-					"3. You can only use the functions defined in tool_calls. If none exist, reply directly to the user.",
+					"2. When using tools, wrap function calls in <tool> and </tool> tags\n" +
+					"3. The tool call must be a valid JSON object\n" +
+					"4. Follow the JSON Schema defined in TOOL_CALLS_SCHEMA\n" +
+					"5. One complete tool call per response\n" +
+					"6. Parameter values MUST strictly follow the descriptions and validation rules defined in properties\n" +
+					"7. For each parameter, carefully check and comply with:\n" +
+					"   - Data type requirements\n" +
+					"   - Format restrictions\n" +
+					"   - Value range limitations\n" +
+					"   - Pattern matching rules\n" +
+					"   - Required field validations\n\n" +
+					"Example:\n" +
+					"<tool>\n" + `{"function":"<FunctionName>","arguments":{"<ArgumentName>":"<ArgumentValue>"}}` + "\n</tool>",
 			}))
 			messages = append(messages, *chatMessage.New().Map(map[string]interface{}{
 				"role": "system",
-				"content": "## Tool Calls Response Rules:\n" +
+				"name": "TOOL_CALLS",
-					"1. The response should be a valid JSON object:\n" +
+				"content": "## Tool Usage Guidelines\n" +
-					"  1.1. e.g: <tool>{\"function\":\"function_name\",\"arguments\":{\"arg1\":\"xxxx\"}}</tool>\n" +
+					"1. Use functions defined in TOOL_CALLS_SCHEMA only when they match your needs\n" +
-					"  1.2. strict the example format, do not add any additional information.\n" +
+					"2. If no matching function exists, respond normally as a helpful assistant\n" +
-					"  1.3. The JSON object should be wrapped by <tool> and </tool>.\n" +
+					"3. When using tools, arguments must match the schema definition exactly\n" +
-					"2. The structure of the JSON object is { \"arguments\": {...}, function:\"function_name\"}\n" +
+					"4. All parameter values must strictly adhere to the validation rules specified in properties\n" +
-					"3. The function_name should be the name of the function defined in tool_calls.\n" +
+					"5. Never skip or ignore any validation requirements defined in the schema",
-					"4. The arguments should be the arguments of the function defined in tool_calls.\n",
 			}))
 
 			// Add tool_calls prompts
@@ -747,8 +758,13 @@ func (ast *Assistant) requestMessages(ctx context.Context, messages []chatMessag
 			"content": content,
 		}
 
+		// Keep the name for user messages
 		if name := message.Name; name != "" {
+			if role != "system" {
 				newMessage["name"] = stringHash(name)
+			} else {
+				newMessage["name"] = name
+			}
 		}
 
 		// Special handling for user messages with JSON content last message
@@ -786,9 +802,9 @@ func (ast *Assistant) requestMessages(ctx context.Context, messages []chatMessag
 
 	// For debug environment, print the request messages
 	if os.Getenv("YAO_AGENT_PRINT_REQUEST_MESSAGES") == "true" {
-		fmt.Println("--------------------------------")
+		fmt.Println("--- REQUEST_MESSAGES -----------------------------")
 		utils.Dump(newMessages)
-		fmt.Println("--------------------------------")
+		fmt.Println("--- END REQUEST_MESSAGES -----------------------------")
 	}
 
 	return newMessages, nil

neo/assistant/hooks.go

@@ -3,6 +3,7 @@ package assistant
 import (
 	"context"
 	"fmt"
+	"os"
 	"strings"
 	"time"
 
@@ -154,26 +155,28 @@ func (ast *Assistant) HookDone(c *gin.Context, context chatctx.Context, input []
 				props := map[string]interface{}{}
 				if text, ok := data.Props["text"].(string); ok {
 
-					// Format the text keep only the <tool> and </tool> inner text
+					// Extract the content between <tool> and </tool> tags more reliably
-					parts := strings.Split(text, "<tool>")
+					startTag := "<tool>"
-					if len(parts) > 1 {
+					endTag := "</tool>"
-						text = parts[1]
+					startIndex := strings.Index(text, startTag)
+					if startIndex != -1 {
+						// Find the content after <tool>
+						content := text[startIndex+len(startTag):]
+						endIndex := strings.LastIndex(content, endTag)
+						if endIndex != -1 {
+							// Extract the content between tags
+							text = content[:endIndex]
+							text = strings.TrimSpace(text)
+							if os.Getenv("YAO_AGENT_PRINT_TOOL_CALL") == "true" {
+								fmt.Println("---- EXTRACTED TOOL CALL ----")
+								fmt.Println(text)
+								fmt.Println("---- END EXTRACTED TOOL CALL ----")
+							}
 						}
-
-					// Format the text keep only the <tool> and </tool> inner text
-					parts = strings.Split(text, "</tool>")
-					if len(parts) > 1 {
-						text = parts[0]
 					}
-
-					// Escape %7B and %7b to {, %7D and %7d to }
-					text = strings.ReplaceAll(text, "%7B", "{")
-					text = strings.ReplaceAll(text, "%7b", "{")
-					text = strings.ReplaceAll(text, "%7D", "}")
-					text = strings.ReplaceAll(text, "%7d", "}")
 
 					// Parse the text into props
-					err := jsoniter.UnmarshalFromString(text, &props)
+					err := ParseJSON(text, &props)
 					if err != nil {
 						props["error"] = fmt.Sprintf("Can not parse the tool call: %s\n--original--\n%s", err.Error(), text)
 					}

neo/assistant/utils.go

@@ -5,7 +5,10 @@ import (
 	"encoding/hex"
 	"fmt"
 	"strconv"
+	"strings"
 	"time"
+
+	jsoniter "github.com/json-iterator/go"
 )
 
 func getTimestamp(v interface{}) (int64, error) {
@@ -51,3 +54,33 @@ func stringHash(v string) string {
 	h.Write([]byte(v))
 	return hex.EncodeToString(h.Sum(nil))
 }
+
+// ParseJSON attempts to parse a potentially malformed JSON string
+// It tries different approaches:
+// 1. Parse as-is
+// 2. Add a missing closing brace
+// 3. Remove an extra closing brace
+func ParseJSON(jsonStr string, v interface{}) error {
+	// Try parsing as-is first
+	err := jsoniter.UnmarshalFromString(jsonStr, v)
+	if err == nil {
+		return nil
+	}
+	originalErr := err
+
+	// Try adding a closing brace
+	if err := jsoniter.UnmarshalFromString(jsonStr+"}", v); err == nil {
+		return nil
+	}
+
+	// Try removing last closing brace if it exists
+	if strings.HasSuffix(jsonStr, "}") {
+		trimmed := strings.TrimSuffix(jsonStr, "}")
+		if err := jsoniter.UnmarshalFromString(trimmed, v); err == nil {
+			return nil
+		}
+	}
+
+	// If all attempts fail, return the original error
+	return originalErr
+}