Skip to content

Commit b6f2ada

Browse files
Patel230Patel230
authored
refactor(tool): extend SchemaProperty and migrate WebSearch, MultiAgent to typed schemas (#335)
1 parent 4768694 commit b6f2ada

4 files changed

Lines changed: 203 additions & 70 deletions

File tree

‎internal/tool/agent.go‎

Lines changed: 40 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -210,47 +210,56 @@ func (MultiAgentTool) Description() string {
210210
"Each task is a prompt string (explore by default) or an object with typed spawn fields."
211211
}
212212

213-
func (MultiAgentTool) Parameters() map[string]interface{} {
214-
return map[string]interface{}{
215-
"type": "object",
216-
"properties": map[string]interface{}{
217-
"tasks": map[string]interface{}{
218-
"type": "array",
219-
"items": map[string]interface{}{
220-
"oneOf": []interface{}{
221-
map[string]interface{}{"type": "string"},
222-
map[string]interface{}{
223-
"type": "object",
224-
"properties": map[string]interface{}{
225-
"prompt": map[string]interface{}{"type": "string"},
226-
"subagent_type": map[string]interface{}{"type": "string"},
227-
"capability_mode": map[string]interface{}{"type": "string"},
228-
"isolation": map[string]interface{}{"type": "string"},
229-
"thoroughness": map[string]interface{}{"type": "string"},
230-
"description": map[string]interface{}{"type": "string"},
231-
"model": map[string]interface{}{"type": "string"},
232-
"cwd": map[string]interface{}{"type": "string"},
213+
// MultiAgentInput is the typed input for MultiAgentTool.
214+
type MultiAgentInput struct {
215+
Tasks []json.RawMessage `json:"tasks"`
216+
RunInBackground bool `json:"run_in_background"`
217+
}
218+
219+
// Schema returns the typed input schema. Parameters() delegates to it so the
220+
// two cannot diverge.
221+
func (MultiAgentTool) Schema() ToolSchema {
222+
return ToolSchema{
223+
Type: "object",
224+
Properties: map[string]SchemaProperty{
225+
"tasks": {
226+
Type: "array",
227+
Items: &SchemaProperty{
228+
OneOf: []SchemaProperty{
229+
{Type: "string"},
230+
{
231+
Type: "object",
232+
Properties: map[string]SchemaProperty{
233+
"prompt": {Type: "string"},
234+
"subagent_type": {Type: "string"},
235+
"capability_mode": {Type: "string"},
236+
"isolation": {Type: "string"},
237+
"thoroughness": {Type: "string"},
238+
"description": {Type: "string"},
239+
"model": {Type: "string"},
240+
"cwd": {Type: "string"},
233241
},
234-
"required": []string{"prompt"},
242+
Required: []string{"prompt"},
235243
},
236244
},
237245
},
238246
},
239-
"run_in_background": map[string]interface{}{
240-
"type": "boolean",
241-
"description": "If true, spawn all sub-agents in the background.",
242-
},
247+
"run_in_background": {Type: "boolean", Description: "If true, spawn all sub-agents in the background."},
243248
},
244-
"required": []string{"tasks"},
249+
Required: []string{"tasks"},
245250
}
246251
}
247252

253+
func (MultiAgentTool) Parameters() map[string]interface{} {
254+
return multiAgentSchema.ToJSONSchema()
255+
}
256+
257+
// multiAgentSchema is the single source of truth for MultiAgent's input schema.
258+
var multiAgentSchema = MultiAgentTool{}.Schema()
259+
248260
func (MultiAgentTool) Execute(ctx context.Context, input json.RawMessage) (string, error) {
249-
var p struct {
250-
Tasks []json.RawMessage `json:"tasks"`
251-
RunInBackground bool `json:"run_in_background"`
252-
}
253-
if err := json.Unmarshal(input, &p); err != nil {
261+
p, err := DecodeInput[MultiAgentInput]("MultiAgent", input)
262+
if err != nil {
254263
return "", err
255264
}
256265
if len(p.Tasks) > maxParallelAgentTasks {

‎internal/tool/schema_batch_test.go‎

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1379,3 +1379,94 @@ func TestTerminalKillSchemaProvider(t *testing.T) {
13791379
t.Fatalf("required = %v, want [terminal_id]", TerminalKillTool{}.Parameters()["required"])
13801380
}
13811381
}
1382+
1383+
func TestWebSearchSchemaProvider(t *testing.T) {
1384+
var _ SchemaProvider = WebSearchTool{}
1385+
props := schemaProps(t, WebSearchTool{}.Parameters())
1386+
query := props["query"].(map[string]interface{})
1387+
if query["type"] != "string" {
1388+
t.Fatal("query type wrong")
1389+
}
1390+
if query["maxLength"] != 2000 {
1391+
t.Fatalf("query maxLength = %v, want 2000", query["maxLength"])
1392+
}
1393+
queries := props["queries"].(map[string]interface{})
1394+
if queries["type"] != "array" {
1395+
t.Fatalf("queries type = %v, want array", queries["type"])
1396+
}
1397+
if queries["maxItems"] != 20 {
1398+
t.Fatalf("queries maxItems = %v, want 20", queries["maxItems"])
1399+
}
1400+
items := queries["items"].(map[string]interface{})
1401+
if items["type"] != "string" || items["maxLength"] != 2000 {
1402+
t.Fatalf("queries items = %v, want string+maxLength 2000", items)
1403+
}
1404+
num := props["numResults"].(map[string]interface{})
1405+
if num["default"] != 5 || num["minimum"] != 1 || num["maximum"] != 20 {
1406+
t.Fatalf("numResults = %v, want min/max/default", num)
1407+
}
1408+
st := props["searchType"].(map[string]interface{})
1409+
enum, ok := st["enum"].([]interface{})
1410+
if !ok || len(enum) != 2 || enum[0] != "web" || enum[1] != "news" {
1411+
t.Fatalf("searchType enum = %v, want [web news]", st["enum"])
1412+
}
1413+
if st["default"] != "web" {
1414+
t.Fatalf("searchType default = %v, want web", st["default"])
1415+
}
1416+
}
1417+
1418+
func TestMultiAgentSchemaProvider(t *testing.T) {
1419+
var _ SchemaProvider = MultiAgentTool{}
1420+
props := schemaProps(t, MultiAgentTool{}.Parameters())
1421+
tasks := props["tasks"].(map[string]interface{})
1422+
if tasks["type"] != "array" {
1423+
t.Fatalf("tasks type = %v, want array", tasks["type"])
1424+
}
1425+
if _, hasType := tasks["items"].(map[string]interface{})["type"]; hasType {
1426+
t.Fatal("tasks items must not have a type key (oneOf only)")
1427+
}
1428+
oneOf, ok := tasks["items"].(map[string]interface{})["oneOf"].([]interface{})
1429+
if !ok || len(oneOf) != 2 {
1430+
t.Fatalf("tasks items oneOf = %v, want 2 branches", tasks["items"])
1431+
}
1432+
if oneOf[0].(map[string]interface{})["type"] != "string" {
1433+
t.Fatalf("oneOf[0] = %v, want string", oneOf[0])
1434+
}
1435+
obj := oneOf[1].(map[string]interface{})
1436+
if obj["type"] != "object" {
1437+
t.Fatalf("oneOf[1] type = %v, want object", obj["type"])
1438+
}
1439+
objReq, ok := obj["required"].([]string)
1440+
if !ok || len(objReq) != 1 || objReq[0] != "prompt" {
1441+
t.Fatalf("oneOf[1] required = %v, want [prompt]", obj["required"])
1442+
}
1443+
req, _ := MultiAgentTool{}.Parameters()["required"].([]string)
1444+
if len(req) != 1 || req[0] != "tasks" {
1445+
t.Fatalf("required = %v, want [tasks]", MultiAgentTool{}.Parameters()["required"])
1446+
}
1447+
}
1448+
1449+
func TestSchemaPropertyExtensions(t *testing.T) {
1450+
p := SchemaProperty{Type: "string", MaxLength: 10, MinLength: 1}
1451+
m := p.toMap()
1452+
if m["maxLength"] != 10 || m["minLength"] != 1 {
1453+
t.Fatalf("length bounds = %v", m)
1454+
}
1455+
a := SchemaProperty{Type: "array", Items: &SchemaProperty{Type: "string"}, MaxItems: 5, MinItems: 1}
1456+
am := a.toMap()
1457+
if am["maxItems"] != 5 || am["minItems"] != 1 {
1458+
t.Fatalf("item bounds = %v", am)
1459+
}
1460+
if am["items"].(map[string]interface{})["type"] != "string" {
1461+
t.Fatalf("items = %v", am["items"])
1462+
}
1463+
o := SchemaProperty{OneOf: []SchemaProperty{{Type: "string"}, {Type: "integer"}}}
1464+
om := o.toMap()
1465+
if _, hasType := om["type"]; hasType {
1466+
t.Fatalf("typeless oneOf must omit type key: %v", om)
1467+
}
1468+
branches, ok := om["oneOf"].([]interface{})
1469+
if !ok || len(branches) != 2 || branches[0].(map[string]interface{})["type"] != "string" {
1470+
t.Fatalf("oneOf = %v", om["oneOf"])
1471+
}
1472+
}

‎internal/tool/tool.go‎

Lines changed: 43 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,19 @@ type SchemaProperty struct {
6969
// hand-written schemas exactly.
7070
Minimum interface{} `json:"minimum,omitempty"`
7171
Maximum interface{} `json:"maximum,omitempty"`
72+
// MaxLength/MinLength constrain string fields. MaxItems/MinItems constrain
73+
// array fields. Plain ints (omitted when zero): zero is never a meaningful
74+
// bound, so authors write literals like MaxLength: 2000.
75+
MaxLength int `json:"maxLength,omitempty"`
76+
MinLength int `json:"minLength,omitempty"`
77+
MaxItems int `json:"maxItems,omitempty"`
78+
MinItems int `json:"minItems,omitempty"`
79+
// OneOf/AnyOf list alternative subschemas (e.g. array items that accept a
80+
// string or an object). Each branch is a full SchemaProperty. A property
81+
// that only carries OneOf/AnyOf leaves Type empty, and toMap omits the
82+
// "type" key in that case to match hand-written schemas exactly.
83+
OneOf []SchemaProperty `json:"oneOf,omitempty"`
84+
AnyOf []SchemaProperty `json:"anyOf,omitempty"`
7285
}
7386

7487
// SchemaProvider is an optional interface tools can implement to expose a typed
@@ -98,7 +111,10 @@ func (s ToolSchema) ToJSONSchema() map[string]interface{} {
98111
}
99112

100113
func (p SchemaProperty) toMap() map[string]interface{} {
101-
m := map[string]interface{}{"type": p.Type}
114+
m := map[string]interface{}{}
115+
if p.Type != "" {
116+
m["type"] = p.Type
117+
}
102118
if p.Description != "" {
103119
m["description"] = p.Description
104120
}
@@ -114,6 +130,32 @@ func (p SchemaProperty) toMap() map[string]interface{} {
114130
if p.Maximum != nil {
115131
m["maximum"] = p.Maximum
116132
}
133+
if p.MaxLength != 0 {
134+
m["maxLength"] = p.MaxLength
135+
}
136+
if p.MinLength != 0 {
137+
m["minLength"] = p.MinLength
138+
}
139+
if p.MaxItems != 0 {
140+
m["maxItems"] = p.MaxItems
141+
}
142+
if p.MinItems != 0 {
143+
m["minItems"] = p.MinItems
144+
}
145+
if len(p.OneOf) > 0 {
146+
branches := make([]interface{}, 0, len(p.OneOf))
147+
for _, b := range p.OneOf {
148+
branches = append(branches, b.toMap())
149+
}
150+
m["oneOf"] = branches
151+
}
152+
if len(p.AnyOf) > 0 {
153+
branches := make([]interface{}, 0, len(p.AnyOf))
154+
for _, b := range p.AnyOf {
155+
branches = append(branches, b.toMap())
156+
}
157+
m["anyOf"] = branches
158+
}
117159
if p.Items != nil {
118160
m["items"] = p.Items.toMap()
119161
}

‎internal/tool/web_search.go‎

Lines changed: 29 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -30,55 +30,46 @@ type searchResult struct {
3030

3131
type WebSearchTool struct{}
3232

33-
func (WebSearchTool) Name() string { return "WebSearch" }
33+
func (WebSearchTool) Name() string { return "WebSearch" }
34+
35+
// WebSearchInput is the typed input for WebSearchTool.
36+
type WebSearchInput struct {
37+
Query string `json:"query"`
38+
Queries []string `json:"queries"`
39+
NumResults int `json:"numResults"`
40+
SearchType string `json:"searchType"`
41+
}
42+
3443
func (WebSearchTool) RiskLevel() string { return "low" }
3544
func (WebSearchTool) Aliases() []string { return []string{"web_search"} }
3645
func (WebSearchTool) Description() string {
3746
return "Search the web and return structured results. Supports Brave Search, SearXNG, DeepSeek, Exa, Perplexity, and DuckDuckGo backends."
3847
}
3948

40-
func (WebSearchTool) Parameters() map[string]interface{} {
41-
return map[string]interface{}{
42-
"type": "object",
43-
"properties": map[string]interface{}{
44-
"query": map[string]interface{}{
45-
"type": "string",
46-
"description": "Search query. Provide this OR queries (not both).",
47-
"maxLength": maxWebSearchQueryLength,
48-
},
49-
"queries": map[string]interface{}{
50-
"type": "array",
51-
"items": map[string]interface{}{"type": "string", "maxLength": maxWebSearchQueryLength},
52-
"maxItems": maxWebSearchQueries,
53-
"description": "Multiple search queries to run concurrently in a single call. Use this to research several things at once instead of issuing one WebSearch per query.",
54-
},
55-
"numResults": map[string]interface{}{
56-
"type": "integer",
57-
"description": "Number of results to return (1-20)",
58-
"minimum": 1,
59-
"maximum": 20,
60-
"default": 5,
61-
},
62-
"searchType": map[string]interface{}{
63-
"type": "string",
64-
"description": "Type of search to perform",
65-
"enum": []string{"web", "news"},
66-
"default": "web",
67-
},
49+
// Schema returns the typed input schema. Parameters() delegates to it so the
50+
// two cannot diverge.
51+
func (WebSearchTool) Schema() ToolSchema {
52+
return ToolSchema{
53+
Type: "object",
54+
Properties: map[string]SchemaProperty{
55+
"query": {Type: "string", Description: "Search query. Provide this OR queries (not both).", MaxLength: maxWebSearchQueryLength},
56+
"queries": {Type: "array", Items: &SchemaProperty{Type: "string", MaxLength: maxWebSearchQueryLength}, MaxItems: maxWebSearchQueries, Description: "Multiple search queries to run concurrently in a single call. Use this to research several things at once instead of issuing one WebSearch per query."},
57+
"numResults": {Type: "integer", Description: "Number of results to return (1-20)", Minimum: 1, Maximum: 20, Default: 5},
58+
"searchType": {Type: "string", Description: "Type of search to perform", Enum: []interface{}{"web", "news"}, Default: "web"},
6859
},
69-
// Either query or queries must be supplied; validated in Execute since
70-
// JSON Schema "required" cannot express an exclusive-or cleanly.
7160
}
7261
}
7362

63+
func (WebSearchTool) Parameters() map[string]interface{} {
64+
return webSearchSchema.ToJSONSchema()
65+
}
66+
67+
// webSearchSchema is the single source of truth for WebSearch's input schema.
68+
var webSearchSchema = WebSearchTool{}.Schema()
69+
7470
func (t WebSearchTool) Execute(ctx context.Context, input json.RawMessage) (string, error) {
75-
var p struct {
76-
Query string `json:"query"`
77-
Queries []string `json:"queries"`
78-
NumResults int `json:"numResults"`
79-
SearchType string `json:"searchType"`
80-
}
81-
if err := json.Unmarshal(input, &p); err != nil {
71+
p, err := DecodeInput[WebSearchInput]("WebSearch", input)
72+
if err != nil {
8273
return "", err
8374
}
8475

0 commit comments

Comments
 (0)