Merge pull request #876 from trheyi/main

Enhance thread safety and tool call handling in streaming components

neo/assistant/api.go

@@ -334,8 +334,7 @@ func (ast *Assistant) streamChat(
 	isFirstThink := true
 	isThinking := false
 
-	isFirstTool := true
+	toolsCount := 0
-	isTool := false
 	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 toolsCount > 1 {
-			if isTool && msg.Type != "tool_calls_native" {
+					msg.Text = "" // clear the text
-
+					msg.Type = "text"
-				if msg.IsDone {
+					msg.IsNew = false
-					end := chatMessage.New().Map(map[string]interface{}{"text": "}\n</tool>\n", "type": "tool", "delta": true})
+					return 1 // continue
-					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
 				}
 
-				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()

neo/message/message.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"os"
 	"strings"
+	"sync"
 
 	"github.com/fatih/color"
 	"github.com/gin-gonic/gin"
@@ -15,6 +16,8 @@ import (
 	"github.com/yaoapp/yao/openai"
 )
 
+var locker = sync.Mutex{}
+
 // Message the message
 type Message struct {
 	ID              string                 `json:"id,omitempty"`               // id for the message
@@ -37,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
@@ -221,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
+			text := ""
-			function := chunk.Choices[0].Delta.ToolCalls[0].Function.Name
+			if len(chunk.Choices[0].Delta.ToolCalls) > 0 {
-			arguments := chunk.Choices[0].Delta.ToolCalls[0].Function.Arguments
+				id := chunk.Choices[0].Delta.ToolCalls[0].ID
-			text := arguments
+				function := chunk.Choices[0].Delta.ToolCalls[0].Function.Name
-			if id != "" {
+				arguments := chunk.Choices[0].Delta.ToolCalls[0].Function.Arguments
-				text = fmt.Sprintf(`{"id": "%s", "function": "%s", "arguments": %s`, id, function, arguments)
+				text = arguments
-				msg.IsNew = true // mark as a new message
+				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
 		}
 
@@ -721,6 +734,11 @@ func (m *Message) Callback(fn interface{}) *Message {
 
 // Write writes the message to response writer
 func (m *Message) Write(w gin.ResponseWriter) bool {
+
+	// Sync write to response writer
+	locker.Lock()
+	defer locker.Unlock()
+
 	defer func() {
 		if r := recover(); r != nil {