From 617893a57995f5a200c5dfd049d1f31764d407d9 Mon Sep 17 00:00:00 2001
From: Max <shadow.iqka@qq.com>
Date: Mon, 10 Feb 2025 11:13:59 +0800
Subject: [PATCH] Improve tool call parsing and formatting

- Refine tool call text extraction in HookDone method
- Add URL decoding for JSON-encoded tool call arguments
- Enhance error handling for tool call parsing
- Modify tool call tag handling in streaming chat
---
 neo/assistant/api.go   |  2 +-
 neo/assistant/hooks.go | 22 ++++++++++++++++++----
 neo/message/message.go |  1 +
 3 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/neo/assistant/api.go b/neo/assistant/api.go
index ed0e9266..fc731c4f 100644
--- a/neo/assistant/api.go
+++ b/neo/assistant/api.go
@@ -379,7 +379,7 @@ func (ast *Assistant) streamChat(
 			// 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
+					msg.Text = "\n<tool>\n" + msg.Text // add the tool_calls begin tag
 					isFirstTool = false
 					isTool = true
 				}
diff --git a/neo/assistant/hooks.go b/neo/assistant/hooks.go
index 8d448788..e6dc7ff5 100644
--- a/neo/assistant/hooks.go
+++ b/neo/assistant/hooks.go
@@ -154,14 +154,28 @@ func (ast *Assistant) HookDone(c *gin.Context, context chatctx.Context, input []
 				props := map[string]interface{}{}
 				if text, ok := data.Props["text"].(string); ok {
 
-					// Remove <tool> and </tool> tags
-					text = strings.ReplaceAll(text, "<tool>", "")
-					text = strings.ReplaceAll(text, "</tool>", "")
+					// Format the text keep only the <tool> and </tool> inner text
+					parts := strings.Split(text, "<tool>")
+					if len(parts) > 1 {
+						text = parts[1]
+					}
+
+					// 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)
 					if err != nil {
-						props["error"] = err.Error()
+						props["error"] = fmt.Sprintf("Can not parse the tool call: %s\n--original--\n%s", err.Error(), text)
 					}
 				}
 
diff --git a/neo/message/message.go b/neo/message/message.go
index 6b6d81bf..19e8fbe0 100644
--- a/neo/message/message.go
+++ b/neo/message/message.go
@@ -228,6 +228,7 @@ func NewOpenAI(data []byte, isThinking bool) *Message {
 			text := arguments
 			if id != "" {
 				text = fmt.Sprintf(`{"id": "%s", "function": "%s", "arguments": %s`, id, function, arguments)
+				msg.IsNew = true // mark as a new message
 			}
 
 			msg.Text = text