From 78f07b0c18b803702c4a6d65025b05b0a396c153 Mon Sep 17 00:00:00 2001 From: Donnie Adams Date: Sat, 4 Jan 2025 21:00:38 -0500 Subject: [PATCH] fix: only count assistant messages when limiting consecutive tool calls Before this change, the number of consecutive tool calls was counted by looking at all messages since the last time a message with role "user" was sent. This was bad logic because the LLM could ask for many parallel tool calls. This change addresses this problem by only considering messages with role "assistant" with tool calls since the last user message. Signed-off-by: Donnie Adams --- pkg/engine/engine.go | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/pkg/engine/engine.go b/pkg/engine/engine.go index 69e16b68..bd7b09ab 100644 --- a/pkg/engine/engine.go +++ b/pkg/engine/engine.go @@ -400,17 +400,24 @@ func (e *Engine) complete(ctx context.Context, state *State) (*Return, error) { } }() - // Limit the number of consecutive tool calls and responses. + // Limit the number of consecutive tool calls. // We don't want the LLM to call tools unrestricted or get stuck in an error loop. var messagesSinceLastUserMessage int for _, msg := range slices.Backward(state.Completion.Messages) { if msg.Role == types.CompletionMessageRoleTypeUser { break + } else if msg.Role == types.CompletionMessageRoleTypeAssistant { + for _, content := range msg.Content { + // If this message is requesting that a tool call be made, then count it towards the limit. + if content.ToolCall != nil { + messagesSinceLastUserMessage++ + break + } + } } - messagesSinceLastUserMessage++ } - // Divide by 2 because tool calls come in pairs: call and response. - if messagesSinceLastUserMessage/2 > maxConsecutiveToolCalls { + + if messagesSinceLastUserMessage > maxConsecutiveToolCalls { msg := fmt.Sprintf("We cannot continue because the number of consecutive tool calls is limited to %d.", maxConsecutiveToolCalls) ret.State.Completion.Messages = append(state.Completion.Messages, types.CompletionMessage{ Role: types.CompletionMessageRoleTypeAssistant,