Skip to content

Commit 7b74cda

Browse files
dgageotdgageot
authored
Merge pull request #520 from dgageot/add-category
Add a category to every builtin tool
2 parents 1635b61 + c660603 commit 7b74cda

18 files changed

Lines changed: 65 additions & 24 deletions

‎examples/golibrary/tool/main.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,7 @@ func run(ctx context.Context) error {
5959

6060
toolAddNumbers := tools.Tool{
6161
Name: "add",
62+
Category: "compute",
6263
Description: "Add two numbers",
6364
Parameters: tools.MustSchemaFor[AddNumbersArgs](),
6465
Handler: addNumbers,

‎pkg/codemode/codemode.go‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -47,9 +47,7 @@ func (c *codeModeTool) Instructions() string {
4747
}
4848

4949
func isExcludedTool(tool tools.Tool) bool {
50-
// TODO(dga): make this more robust. It really a temporary hack to exclude the todo tools.
51-
// See #514
52-
return strings.Contains(tool.Name, "_todo")
50+
return tool.Category == "todo"
5351
}
5452

5553
func (c *codeModeTool) Tools(ctx context.Context) ([]tools.Tool, error) {
@@ -75,6 +73,7 @@ func (c *codeModeTool) Tools(ctx context.Context) ([]tools.Tool, error) {
7573

7674
allTools := []tools.Tool{{
7775
Name: "run_tools_with_javascript",
76+
Category: "code mode",
7877
Description: prompt + strings.Join(functionsDoc, "\n"),
7978
Parameters: tools.MustSchemaFor[RunToolsWithJavascriptArgs](),
8079
Handler: func(ctx context.Context, toolCall tools.ToolCall) (*tools.ToolCallResult, error) {

‎pkg/codemode/codemode_test.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ func TestCodeModeTool_Tools(t *testing.T) {
1818

1919
fetchTool := toolSet[0]
2020
assert.Equal(t, "run_tools_with_javascript", fetchTool.Name)
21+
assert.Equal(t, "code mode", fetchTool.Category)
2122
assert.NotNil(t, fetchTool.Handler)
2223

2324
inputSchema, err := json.Marshal(fetchTool.Parameters)

‎pkg/tools/builtin/fetch.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -309,6 +309,7 @@ func (t *FetchTool) Tools(context.Context) ([]tools.Tool, error) {
309309
return []tools.Tool{
310310
{
311311
Name: "fetch",
312+
Category: "fetch",
312313
Description: "Fetch content from one or more HTTP/HTTPS URLs. Returns the response body and metadata.",
313314
Parameters: map[string]any{
314315
"type": "object",

‎pkg/tools/builtin/fetch_test.go‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -25,13 +25,16 @@ func TestFetchToolWithOptions(t *testing.T) {
2525
func TestFetchTool_Tools(t *testing.T) {
2626
tool := NewFetchTool()
2727

28-
toolSet, err := tool.Tools(t.Context())
28+
allTools, err := tool.Tools(t.Context())
2929
require.NoError(t, err)
30-
require.Len(t, toolSet, 1)
30+
require.Len(t, allTools, 1)
31+
for _, tool := range allTools {
32+
assert.NotNil(t, tool.Handler)
33+
assert.Equal(t, "fetch", tool.Category)
34+
}
3135

32-
fetchTool := toolSet[0]
36+
fetchTool := allTools[0]
3337
assert.Equal(t, "fetch", fetchTool.Name)
34-
assert.NotNil(t, fetchTool.Handler)
3538

3639
schema, err := json.Marshal(fetchTool.Parameters)
3740
require.NoError(t, err)

‎pkg/tools/builtin/filesystem.go‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -153,6 +153,7 @@ func (t *FilesystemTool) Tools(context.Context) ([]tools.Tool, error) {
153153
tls := []tools.Tool{
154154
{
155155
Name: "create_directory",
156+
Category: "filesystem",
156157
Description: "Create a new directory or ensure a directory exists. Can create multiple nested directories in one operation.",
157158
Parameters: tools.MustSchemaFor[CreateDirectoryArgs](),
158159
OutputSchema: tools.MustSchemaFor[string](),
@@ -163,6 +164,7 @@ func (t *FilesystemTool) Tools(context.Context) ([]tools.Tool, error) {
163164
},
164165
{
165166
Name: "directory_tree",
167+
Category: "filesystem",
166168
Description: "Get a recursive tree view of files and directories as a JSON structure.",
167169
Parameters: tools.MustSchemaFor[DirectoryTreeArgs](),
168170
// Manually define the schema here because
@@ -197,6 +199,7 @@ func (t *FilesystemTool) Tools(context.Context) ([]tools.Tool, error) {
197199
},
198200
{
199201
Name: "edit_file",
202+
Category: "filesystem",
200203
Description: "Make line-based edits to a text file. Each edit replaces exact line sequences with new content.",
201204
Parameters: tools.MustSchemaFor[EditFileArgs](),
202205
OutputSchema: tools.MustSchemaFor[string](),
@@ -207,6 +210,7 @@ func (t *FilesystemTool) Tools(context.Context) ([]tools.Tool, error) {
207210
},
208211
{
209212
Name: "get_file_info",
213+
Category: "filesystem",
210214
Description: "Retrieve detailed metadata about a file or directory.",
211215
Parameters: tools.MustSchemaFor[GetFileInfoArgs](),
212216
OutputSchema: tools.MustSchemaFor[FileInfo](),
@@ -218,6 +222,7 @@ func (t *FilesystemTool) Tools(context.Context) ([]tools.Tool, error) {
218222
},
219223
{
220224
Name: "list_allowed_directories",
225+
Category: "filesystem",
221226
Description: "Returns a list of directories that the server has permission to access. Don't call if you access only the current working directory. It's always allowed.",
222227
OutputSchema: tools.MustSchemaFor[string](),
223228
Handler: t.handleListAllowedDirectories,
@@ -228,6 +233,7 @@ func (t *FilesystemTool) Tools(context.Context) ([]tools.Tool, error) {
228233
},
229234
{
230235
Name: "add_allowed_directory",
236+
Category: "filesystem",
231237
Description: "Request to add a new directory to the allowed directories list. This requires explicit user consent for security reasons.",
232238
Parameters: tools.MustSchemaFor[AddAllowedDirectoryArgs](),
233239
OutputSchema: tools.MustSchemaFor[string](),
@@ -238,6 +244,7 @@ func (t *FilesystemTool) Tools(context.Context) ([]tools.Tool, error) {
238244
},
239245
{
240246
Name: "list_directory",
247+
Category: "filesystem",
241248
Description: "Get a detailed listing of all files and directories in a specified path.",
242249
Parameters: tools.MustSchemaFor[ListDirectoryArgs](),
243250
OutputSchema: tools.MustSchemaFor[string](),
@@ -249,6 +256,7 @@ func (t *FilesystemTool) Tools(context.Context) ([]tools.Tool, error) {
249256
},
250257
{
251258
Name: "list_directory_with_sizes",
259+
Category: "filesystem",
252260
Description: "Get a detailed listing of all files and directories in a specified path, including sizes.",
253261
Parameters: tools.MustSchemaFor[ListDirectoryArgs](),
254262
OutputSchema: tools.MustSchemaFor[string](),
@@ -260,6 +268,7 @@ func (t *FilesystemTool) Tools(context.Context) ([]tools.Tool, error) {
260268
},
261269
{
262270
Name: "move_file",
271+
Category: "filesystem",
263272
Description: "Move or rename files and directories.",
264273
Parameters: tools.MustSchemaFor[MoveFileArgs](),
265274
OutputSchema: tools.MustSchemaFor[string](),
@@ -270,6 +279,7 @@ func (t *FilesystemTool) Tools(context.Context) ([]tools.Tool, error) {
270279
},
271280
{
272281
Name: "read_file",
282+
Category: "filesystem",
273283
Description: "Read the complete contents of a file from the file system.",
274284
Parameters: tools.MustSchemaFor[ReadFileArgs](),
275285
OutputSchema: tools.MustSchemaFor[string](),
@@ -281,6 +291,7 @@ func (t *FilesystemTool) Tools(context.Context) ([]tools.Tool, error) {
281291
},
282292
{
283293
Name: "read_multiple_files",
294+
Category: "filesystem",
284295
Description: "Read the contents of multiple files simultaneously.",
285296
Parameters: tools.MustSchemaFor[ReadMultipleFilesArgs](),
286297
// TODO(dga): depends on the json param
@@ -293,6 +304,7 @@ func (t *FilesystemTool) Tools(context.Context) ([]tools.Tool, error) {
293304
},
294305
{
295306
Name: "search_files",
307+
Category: "filesystem",
296308
Description: "Recursively search for files and directories matching a pattern. Prints the full paths of matching files and the total number of files found.",
297309
Parameters: tools.MustSchemaFor[SearchFilesArgs](),
298310
OutputSchema: tools.MustSchemaFor[string](),
@@ -304,6 +316,7 @@ func (t *FilesystemTool) Tools(context.Context) ([]tools.Tool, error) {
304316
},
305317
{
306318
Name: "search_files_content",
319+
Category: "filesystem",
307320
Description: "Searches for text or regex patterns in the content of files matching a GLOB pattern.",
308321
Parameters: tools.MustSchemaFor[SearchFilesContentArgs](),
309322
OutputSchema: tools.MustSchemaFor[string](),
@@ -315,6 +328,7 @@ func (t *FilesystemTool) Tools(context.Context) ([]tools.Tool, error) {
315328
},
316329
{
317330
Name: "write_file",
331+
Category: "filesystem",
318332
Description: "Create a new file or completely overwrite an existing file with new content.",
319333
Parameters: tools.MustSchemaFor[WriteFileArgs](),
320334
OutputSchema: tools.MustSchemaFor[string](),

‎pkg/tools/builtin/filesystem_test.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@ func TestFilesystemTool_Tools(t *testing.T) {
5858
for _, tool := range allTools {
5959
toolNames = append(toolNames, tool.Name)
6060
assert.NotNil(t, tool.Handler)
61+
assert.Equal(t, "filesystem", tool.Category)
6162
}
6263

6364
for _, expected := range expectedTools {

‎pkg/tools/builtin/memory.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ func (t *MemoryTool) Tools(context.Context) ([]tools.Tool, error) {
4747
return []tools.Tool{
4848
{
4949
Name: "add_memory",
50+
Category: "memory",
5051
Description: "Add a new memory to the database",
5152
Parameters: tools.MustSchemaFor[AddMemoryArgs](),
5253
OutputSchema: tools.MustSchemaFor[string](),
@@ -57,6 +58,7 @@ func (t *MemoryTool) Tools(context.Context) ([]tools.Tool, error) {
5758
},
5859
{
5960
Name: "get_memories",
61+
Category: "memory",
6062
Description: "Retrieve all stored memories",
6163
OutputSchema: tools.MustSchemaFor[[]database.UserMemory](),
6264
Handler: t.handleGetMemories,
@@ -67,6 +69,7 @@ func (t *MemoryTool) Tools(context.Context) ([]tools.Tool, error) {
6769
},
6870
{
6971
Name: "delete_memory",
72+
Category: "memory",
7073
Description: "Delete a specific memory by ID",
7174
Parameters: tools.MustSchemaFor[DeleteMemoryArgs](),
7275
OutputSchema: tools.MustSchemaFor[string](),

‎pkg/tools/builtin/memory_test.go‎

Lines changed: 12 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -54,18 +54,21 @@ func TestMemoryTool_Tools(t *testing.T) {
5454
manager := new(MockMemoryManager)
5555
tool := NewMemoryTool(manager)
5656

57-
tls, err := tool.Tools(t.Context())
58-
57+
allTools, err := tool.Tools(t.Context())
5958
require.NoError(t, err)
60-
assert.Len(t, tls, 3)
59+
require.Len(t, allTools, 3)
60+
for _, tool := range allTools {
61+
assert.NotNil(t, tool.Handler)
62+
assert.Equal(t, "memory", tool.Category)
63+
}
6164

6265
// Verify tool functions
63-
assert.Equal(t, "add_memory", tls[0].Name)
64-
assert.Equal(t, "get_memories", tls[1].Name)
65-
assert.Equal(t, "delete_memory", tls[2].Name)
66+
assert.Equal(t, "add_memory", allTools[0].Name)
67+
assert.Equal(t, "get_memories", allTools[1].Name)
68+
assert.Equal(t, "delete_memory", allTools[2].Name)
6669

6770
// Check add_memory parameters
68-
schema, err := json.Marshal(tls[0].Parameters)
71+
schema, err := json.Marshal(allTools[0].Parameters)
6972
require.NoError(t, err)
7073
assert.JSONEq(t, `{
7174
"type": "object",
@@ -82,10 +85,10 @@ func TestMemoryTool_Tools(t *testing.T) {
8285
}`, string(schema))
8386

8487
// Check get_memories parameters
85-
assert.Nil(t, tls[1].Parameters)
88+
assert.Nil(t, allTools[1].Parameters)
8689

8790
// Check delete_memory parameters
88-
schema, err = json.Marshal(tls[2].Parameters)
91+
schema, err = json.Marshal(allTools[2].Parameters)
8992
require.NoError(t, err)
9093
assert.JSONEq(t, `{
9194
"type": "object",

‎pkg/tools/builtin/shell.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -168,6 +168,7 @@ func (t *ShellTool) Tools(context.Context) ([]tools.Tool, error) {
168168
return []tools.Tool{
169169
{
170170
Name: "shell",
171+
Category: "shell",
171172
Description: `Executes the given shell command in the user's default shell.`,
172173
Parameters: tools.MustSchemaFor[RunShellArgs](),
173174
OutputSchema: tools.MustSchemaFor[string](),

0 commit comments

Comments
 (0)