Skip to content

Commit ad1d0af

Browse files
YellowDusk04YellowDusk04
authored
chore: ensure that tool calls match tool results and correct comments (#1091)
* chore: ensure that tool calls match tool results and correct comments * chore: adjust condition
1 parent 13e1a25 commit ad1d0af

3 files changed

Lines changed: 177 additions & 16 deletions

File tree

‎adk/middlewares/reduction/reduction.go‎

Lines changed: 61 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -97,15 +97,22 @@ type TypedConfig[M adk.MessageType] struct {
9797

9898
// TokenCounter is used to count the number of tokens in the conversation messages.
9999
// It is used to determine when to trigger clearing based on token usage, and token usage after clearing.
100-
// Required.
100+
// Optional. If not provided, a default token counter will be used, which estimates tokens by counting 1 token per 4 characters.
101101
TokenCounter func(ctx context.Context, msg []M, tools []*schema.ToolInfo) (int64, error)
102102

103103
// MaxTokensForClear is the maximum number of tokens allowed in the conversation before clearing is attempted.
104104
// Required. Default is 160000.
105105
MaxTokensForClear int64
106106

107-
// ClearRetentionSuffixLimit is the number of most recent messages to retain without clearing.
108-
// This ensures the model has some immediate context.
107+
// ClearRetentionSuffixLimit is the number of most recent tool-call rounds to retain without clearing.
108+
// A round consists of one assistant message (which may contain multiple tool calls) and its corresponding tool-result messages.
109+
// This ensures the model has immediate context to process pending tool results.
110+
//
111+
// Example with ClearRetentionSuffixLimit = 2, the retained suffix looks like:
112+
// Round 2: assistant [call_A, call_B] → result_A, result_B
113+
// Round 1: assistant [call_C] → result_C
114+
// (Round 1 is the most recent, closest to the end of the message list.)
115+
//
109116
// Optional. Default is 1.
110117
ClearRetentionSuffixLimit int
111118

@@ -681,16 +688,14 @@ func (t *typedToolReductionMiddleware[M]) beforeModelRewriteStateGeneric(ctx con
681688
toolCallMsg := editTarget[toolCallMsgIndex]
682689
toolCalls := getToolCallsGeneric(toolCallMsg)
683690
if isAssistantMsg(toolCallMsg) && len(toolCalls) > 0 {
684-
toolMsgIndex := toolCallMsgIndex
685691
for _, tc := range toolCalls {
686-
toolMsgIndex++
687-
if toolMsgIndex >= end {
688-
break
689-
}
690-
resultMsg := editTarget[toolMsgIndex]
691-
if !isToolResultMsg(resultMsg) { // unexpected
692-
break
692+
// Find the corresponding tool-result message by callID
693+
resultMsgIndex, found := findToolResultByCallID(editTarget, toolCallMsgIndex+1, toolCallMsgIndex+1+len(toolCalls), tc.CallID)
694+
if !found {
695+
continue // No corresponding tool result found, skip
693696
}
697+
resultMsg := editTarget[resultMsgIndex]
698+
694699
if _, found := t.excludeClearTools[tc.Name]; found {
695700
continue
696701
}
@@ -849,6 +854,51 @@ func (t *typedToolReductionMiddleware[M]) applyClearRewriteGeneric(ctx context.C
849854
return editTarget, end, nil
850855
}
851856

857+
func agenticResultCallID(block *schema.ContentBlock) (string, bool) {
858+
if block == nil {
859+
return "", false
860+
}
861+
if block.Type == schema.ContentBlockTypeFunctionToolResult && block.FunctionToolResult != nil {
862+
return block.FunctionToolResult.CallID, true
863+
}
864+
if block.Type == schema.ContentBlockTypeToolSearchResult && block.ToolSearchFunctionToolResult != nil {
865+
return block.ToolSearchFunctionToolResult.CallID, true
866+
}
867+
return "", false
868+
}
869+
870+
func findToolResultByCallID[M adk.MessageType](messages []M, startIndex, endIndex int, callID string) (int, bool) {
871+
for i := startIndex; i < endIndex; i++ {
872+
msg := messages[i]
873+
if isToolResultMsg(msg) {
874+
if getToolResultCallID(msg) == callID {
875+
return i, true
876+
}
877+
} else {
878+
return -1, false
879+
}
880+
}
881+
return -1, false
882+
}
883+
884+
func getToolResultCallID[M adk.MessageType](msg M) string {
885+
switch m := any(msg).(type) {
886+
case *schema.Message:
887+
if m.Role == schema.Tool {
888+
return m.ToolCallID
889+
}
890+
case *schema.AgenticMessage:
891+
if m.Role == schema.AgenticRoleTypeUser {
892+
for _, block := range m.ContentBlocks {
893+
if callID, ok := agenticResultCallID(block); ok {
894+
return callID
895+
}
896+
}
897+
}
898+
}
899+
return ""
900+
}
901+
852902
type offloadStashItem struct {
853903
config *ToolReductionConfig
854904
offloadInfo *ClearResult

‎adk/middlewares/reduction/reduction_generic_test.go‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -296,6 +296,49 @@ func testHelperFunctions[M adk.MessageType](t *testing.T) {
296296
copiedTCs := getToolCallsGeneric(copied[0])
297297
assert.Equal(t, `{"modified":"true"}`, copiedTCs[0].Arguments)
298298
})
299+
300+
t.Run("getToolResultCallID", func(t *testing.T) {
301+
// Tool result message with matching callID
302+
tr := makeToolResultMsgG[M]("result content", "call_123", "my_tool")
303+
assert.Equal(t, "call_123", getToolResultCallID(tr))
304+
305+
// Non-tool-result message should return empty string
306+
user := makeUserMsgG[M]("hello")
307+
assert.Equal(t, "", getToolResultCallID(user))
308+
})
309+
310+
t.Run("findToolResultByCallID", func(t *testing.T) {
311+
messages := []M{
312+
makeAssistantMsgWithToolCallsG[M]([]testToolCall{
313+
{ID: "call_1", Name: "tool1", Arguments: `{}`},
314+
}),
315+
makeToolResultMsgG[M]("result 1", "call_1", "tool1"),
316+
makeToolResultMsgG[M]("result 2", "call_2", "tool2"),
317+
}
318+
319+
// Find existing tool result
320+
idx, found := findToolResultByCallID(messages, 1, 3, "call_2")
321+
assert.True(t, found)
322+
assert.Equal(t, 2, idx)
323+
324+
// Find non-existent tool result
325+
idx, found = findToolResultByCallID(messages, 1, 3, "call_999")
326+
assert.False(t, found)
327+
assert.Equal(t, -1, idx)
328+
329+
// Stop searching when encountering a non-tool-result message
330+
messages2 := []M{
331+
makeAssistantMsgWithToolCallsG[M]([]testToolCall{
332+
{ID: "call_3", Name: "tool3", Arguments: `{}`},
333+
}),
334+
makeToolResultMsgG[M]("result 3", "call_3", "tool3"),
335+
makeUserMsgG[M]("not a tool result"),
336+
makeToolResultMsgG[M]("result 4", "call_4", "tool4"),
337+
}
338+
idx, found = findToolResultByCallID(messages2, 1, 4, "call_4")
339+
assert.False(t, found)
340+
assert.Equal(t, -1, idx)
341+
})
299342
}
300343

301344
// ---------------------------------------------------------------------------

‎adk/middlewares/reduction/reduction_test.go‎

Lines changed: 73 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -353,7 +353,7 @@ func TestReductionMiddlewareClear(t *testing.T) {
353353
Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London, UK", "unit": "c"}`},
354354
},
355355
}),
356-
schema.ToolMessage("Sunny", "call_123456789"),
356+
schema.ToolMessage("Sunny", "call_987654321"),
357357
schema.AssistantMessage("", []schema.ToolCall{
358358
{
359359
ID: "call_123456789",
@@ -460,7 +460,7 @@ func TestReductionMiddlewareClear(t *testing.T) {
460460
Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London, UK", "unit": "c"}`},
461461
},
462462
}),
463-
schema.ToolMessage("Sunny", "call_123456789"),
463+
schema.ToolMessage("Sunny", "call_987654321"),
464464
schema.AssistantMessage("", []schema.ToolCall{
465465
{
466466
ID: "call_123456789",
@@ -491,6 +491,74 @@ func TestReductionMiddlewareClear(t *testing.T) {
491491
assert.Equal(t, "[Old tool result content cleared]", s.Messages[3].Content)
492492
})
493493

494+
t.Run("test clear with reordered tool results", func(t *testing.T) {
495+
backend := filesystem.NewInMemoryBackend()
496+
config := &Config{
497+
SkipTruncation: true,
498+
TokenCounter: defaultTokenCounter,
499+
MaxTokensForClear: 20,
500+
ClearRetentionSuffixLimit: 0,
501+
ToolConfig: map[string]*ToolReductionConfig{
502+
"get_weather": {
503+
Backend: backend,
504+
SkipClear: false,
505+
ClearHandler: defaultClearHandler(testClearOffloadPath("/tmp"), true, "read_file"),
506+
},
507+
},
508+
}
509+
510+
mw, err := New(ctx, config)
511+
assert.NoError(t, err)
512+
_, s, err := mw.BeforeModelRewriteState(ctx, &adk.ChatModelAgentState{
513+
Messages: []adk.Message{
514+
schema.SystemMessage("you are a helpful assistant"),
515+
schema.UserMessage("get weather for two cities"),
516+
// Assistant message with two tool calls
517+
schema.AssistantMessage("", []schema.ToolCall{
518+
{
519+
ID: "two_call_one",
520+
Type: "function",
521+
Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London"}`},
522+
},
523+
{
524+
ID: "two_call_two",
525+
Type: "function",
526+
Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "Paris"}`},
527+
},
528+
}),
529+
// Intentionally put tool result two BEFORE tool result one
530+
schema.ToolMessage("Cloudy in Paris", "two_call_two"),
531+
schema.ToolMessage("Sunny in London", "two_call_one"),
532+
schema.AssistantMessage("", []schema.ToolCall{{ID: "dummy", Type: "function", Function: schema.FunctionCall{Name: "dummy_tool"}}}),
533+
},
534+
}, &adk.ModelContext{
535+
Tools: toolsInfo,
536+
})
537+
assert.NoError(t, err)
538+
539+
// Verify tool call arguments are preserved (not offloaded since offload is true)
540+
assert.Equal(t, `{"location": "London"}`, s.Messages[2].ToolCalls[0].Function.Arguments)
541+
assert.Equal(t, `{"location": "Paris"}`, s.Messages[2].ToolCalls[1].Function.Arguments)
542+
543+
// Verify tool result two is correctly matched and offloaded (it's at index 3)
544+
assert.Equal(t, "<persisted-output>Tool result saved to: /tmp/clear/two_call_two\nUse read_file to view</persisted-output>", s.Messages[3].Content)
545+
fileContent2, err := backend.Read(ctx, &filesystem.ReadRequest{
546+
FilePath: "/tmp/clear/two_call_two",
547+
})
548+
assert.NoError(t, err)
549+
fileContentStr2 := strings.TrimPrefix(strings.TrimSpace(fileContent2.Content), "1\t")
550+
assert.Equal(t, "Cloudy in Paris", fileContentStr2)
551+
552+
// Verify tool result one is correctly matched and offloaded (it's at index 4)
553+
assert.Equal(t, "<persisted-output>Tool result saved to: /tmp/clear/two_call_one\nUse read_file to view</persisted-output>", s.Messages[4].Content)
554+
fileContent1, err := backend.Read(ctx, &filesystem.ReadRequest{
555+
FilePath: "/tmp/clear/two_call_one",
556+
})
557+
assert.NoError(t, err)
558+
fileContentStr1 := strings.TrimPrefix(strings.TrimSpace(fileContent1.Content), "1\t")
559+
assert.Equal(t, "Sunny in London", fileContentStr1)
560+
})
561+
494562
t.Run("test clear", func(t *testing.T) {
495563
backend := filesystem.NewInMemoryBackend()
496564
handler := func(ctx context.Context, detail *ToolDetail) (*ClearResult, error) {
@@ -547,7 +615,7 @@ func TestReductionMiddlewareClear(t *testing.T) {
547615
Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London, UK", "unit": "c"}`},
548616
},
549617
}),
550-
schema.ToolMessage("Sunny", "call_123456789"),
618+
schema.ToolMessage("Sunny", "call_987654321"),
551619
schema.AssistantMessage("", []schema.ToolCall{
552620
{
553621
ID: "call_123456789",
@@ -621,7 +689,7 @@ func TestReductionMiddlewareClear(t *testing.T) {
621689
Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London, UK", "unit": "c"}`},
622690
},
623691
}),
624-
schema.ToolMessage("Sunny", "call_123456789"),
692+
schema.ToolMessage("Sunny", "call_987654321"),
625693
schema.AssistantMessage("", []schema.ToolCall{
626694
{
627695
ID: "call_123456789",
@@ -872,7 +940,7 @@ func TestReductionMiddlewareClear(t *testing.T) {
872940
Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London, UK", "unit": "c"}`},
873941
},
874942
}),
875-
schema.ToolMessage("Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny", "call_123456789"),
943+
schema.ToolMessage("Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny", "call_987654321"),
876944
schema.AssistantMessage("", []schema.ToolCall{
877945
{
878946
ID: "call_123456789",

0 commit comments

Comments
 (0)