1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
|
// Copyright 2026 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.
package aflow
import (
"net/http"
"strings"
"testing"
"github.com/stretchr/testify/assert"
"google.golang.org/genai"
)
func TestLLMTool(t *testing.T) {
type inputs struct {
Input int
}
type outputs struct {
Reply string
}
type toolArgs struct {
Something string `jsonschema:"something"`
}
testFlow[inputs, outputs](t, map[string]any{"Input": 42}, map[string]any{"Reply": "YES"},
Pipeline(
&LLMAgent{
Name: "smarty",
Model: "model",
TaskType: FormalReasoningTask,
Reply: "Reply",
Instruction: "Do something!",
Prompt: "Prompt",
Tools: []Tool{
&LLMTool{
Name: "researcher",
Model: "sub-agent-model",
TaskType: FormalReasoningTask,
Description: "researcher description",
Instruction: "researcher instruction",
Tools: []Tool{
NewFuncTool("researcher-tool", func(ctx *Context, state inputs, args toolArgs) (struct{}, error) {
// State passed all the way from the workflow inputs.
assert.Equal(t, state.Input, 42)
assert.True(t, strings.HasPrefix(args.Something, "subtool input"),
"args.Something=%q", args.Something)
return struct{}{}, nil
}, "researcher-tool description"),
},
},
},
},
),
[]any{
// Main agent calls the tool sub-agent.
&genai.Part{
FunctionCall: &genai.FunctionCall{
ID: "id0",
Name: "researcher",
Args: map[string]any{
"Question": "What do you think?",
},
},
},
// Sub-agent calls own tool.
&genai.Part{
FunctionCall: &genai.FunctionCall{
ID: "id1",
Name: "researcher-tool",
Args: map[string]any{
"Something": "subtool input 1",
},
},
},
// Sub-agent returns result.
genai.NewPartFromText("Nothing."),
// Repeat the same one more time.
&genai.Part{
FunctionCall: &genai.FunctionCall{
ID: "id2",
Name: "researcher",
Args: map[string]any{
"Question": "But really?",
},
},
},
&genai.Part{
FunctionCall: &genai.FunctionCall{
ID: "id3",
Name: "researcher-tool",
Args: map[string]any{
"Something": "subtool input 2",
},
},
},
// Now model input token overflow.
&genai.Part{
FunctionCall: &genai.FunctionCall{
ID: "id4",
Name: "researcher-tool",
Args: map[string]any{
"Something": "subtool input 3",
},
},
},
genai.APIError{
Code: http.StatusBadRequest,
Message: "The input token count exceeds the maximum number of tokens allowed 1048576.",
},
genai.NewPartFromText("Still nothing."),
// Main returns result.
genai.NewPartFromText("YES"),
},
nil,
)
}
|
