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

Improve native tool calls handling in assistant streaming

Browse source 
- 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
This commit is contained in:
Max 2025-02-24 19:29:57 +08:00
parent 101927e343
commit 358f327dee
2 changed files with 44 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()

28
neo/message/message.go
View file

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