From 358f327dee8af80088b26fe520f90fc47896f549 Mon Sep 17 00:00:00 2001
From: Max <shadow.iqka@qq.com>
Date: Mon, 24 Feb 2025 19:29:57 +0800
Subject: [PATCH] Improve native tool calls handling in assistant streaming

- Refactor tool calls processing to handle multiple tool calls more robustly
- Add new message flags (IsBeginTool, IsEndTool) to track tool call stages
- Modify streaming logic to correctly wrap and manage tool call messages
- Ensure only the first tool call is fully processed and displayed
---
 neo/assistant/api.go   | 49 +++++++++++++++++++++---------------------
 neo/message/message.go | 28 ++++++++++++++++--------
 2 files changed, 44 insertions(+), 33 deletions(-)

diff --git a/neo/assistant/api.go b/neo/assistant/api.go
index 79508e36..6222f1ba 100644
--- a/neo/assistant/api.go
+++ b/neo/assistant/api.go
@@ -334,8 +334,7 @@ func (ast *Assistant) streamChat(
 	isFirstThink := true
 	isThinking := false
 
-	isFirstTool := true
-	isTool := false
+	toolsCount := 0
 	currentMessageID := ""
 	err := ast.Chat(c.Request.Context(), messages, options, func(data []byte) int {
 		select {
@@ -401,32 +400,34 @@ func (ast *Assistant) streamChat(
 				contents.ClearToken()
 			}
 
-			// for native tool_calls response
+			// for native tool_calls response, keep the first tool_calls_native message
 			if msg.Type == "tool_calls_native" {
-				if isFirstTool {
-					msg.Text = "\n<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.ID = currentMessageID
-					end.Retry = ctx.Retry
-					end.Silent = ctx.Silent
-					end.Callback(cb).Write(c.Writer)
-					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
+				if toolsCount > 1 {
+					msg.Text = "" // clear the text
+					msg.Type = "text"
+					msg.IsNew = false
+					return 1 // continue
 				}
 
-				isTool = false
+				if msg.IsBeginTool {
+
+					if toolsCount == 1 {
+						msg.IsNew = false
+						msg.Text = "\n</tool>\n" // add the tool_calls close tag
+					}
+
+					if toolsCount == 0 {
+						msg.Text = "\n<tool>\n" + msg.Text // add the tool_calls begin tag
+					}
+
+					toolsCount++
+
+				}
+
+				if msg.IsEndTool {
+					msg.Text = msg.Text + "\n</tool>\n" // add the tool_calls close tag
+				}
 			}
 
 			delta := msg.String()
diff --git a/neo/message/message.go b/neo/message/message.go
index 09832edc..92c7df2c 100644
--- a/neo/message/message.go
+++ b/neo/message/message.go
@@ -40,6 +40,9 @@ type Message struct {
 	Hidden          bool                   `json:"hidden,omitempty"`           // hidden for the message (not show in the UI and history)
 	Retry           bool                   `json:"retry,omitempty"`            // retry for the message
 	Silent          bool                   `json:"silent,omitempty"`           // silent for the message (not show in the UI and history)
+	IsTool          bool                   `json:"-"`                          // is tool for the message for native tool_calls
+	IsBeginTool     bool                   `json:"-"`                          // is new tool for the message for native tool_calls
+	IsEndTool       bool                   `json:"-"`                          // is end tool for the message for native tool_calls
 }
 
 // Mention represents a mention
@@ -224,19 +227,26 @@ func NewOpenAI(data []byte, isThinking bool) *Message {
 		}
 
 		// Tool calls
-		if len(chunk.Choices[0].Delta.ToolCalls) > 0 {
+		if len(chunk.Choices[0].Delta.ToolCalls) > 0 || chunk.Choices[0].FinishReason == "tool_calls" {
 			msg.Type = "tool_calls_native"
-			id := chunk.Choices[0].Delta.ToolCalls[0].ID
-			function := chunk.Choices[0].Delta.ToolCalls[0].Function.Name
-			arguments := chunk.Choices[0].Delta.ToolCalls[0].Function.Arguments
-			text := arguments
-			if id != "" {
-				text = fmt.Sprintf(`{"id": "%s", "function": "%s", "arguments": %s`, id, function, arguments)
-				msg.IsNew = true // mark as a new message
+			text := ""
+			if len(chunk.Choices[0].Delta.ToolCalls) > 0 {
+				id := chunk.Choices[0].Delta.ToolCalls[0].ID
+				function := chunk.Choices[0].Delta.ToolCalls[0].Function.Name
+				arguments := chunk.Choices[0].Delta.ToolCalls[0].Function.Arguments
+				text = arguments
+				if id != "" {
+					msg.IsBeginTool = true
+					msg.IsNew = true // mark as a new message
+					text = fmt.Sprintf(`{"id": "%s", "function": "%s", "arguments": %s`, id, function, arguments)
+				}
+			}
+
+			if chunk.Choices[0].FinishReason == "tool_calls" {
+				msg.IsEndTool = true
 			}
 
 			msg.Text = text
-			msg.IsDone = chunk.Choices[0].FinishReason == "tool_calls" // is done when tool calls are finished
 			return msg
 		}