opencode/yao
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #876 from trheyi/main

Browse source 
Enhance thread safety and tool call handling in streaming components
This commit is contained in:
Max 2025-02-24 19:33:20 +08:00 • committed by GitHub
commit f9f1177e0d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 52 additions and 33 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

49
neo/assistant/api.go
View file

@ -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()

36
neo/message/message.go
View file

@ -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
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
}
@ -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 {