Skip to content

Commit 4aaf4dc

Browse files
dgageotdgageot
committed
Move code from displaynames.go to annotations
Signed-off-by: David Gageot <david.gageot@docker.com>
1 parent a3d7fd6 commit 4aaf4dc

17 files changed

Lines changed: 84 additions & 75 deletions

File tree

‎internal/app/app.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,10 @@ func (a *App) FirstMessage() *string {
3131
return a.firstMessage
3232
}
3333

34+
func (a *App) ToolDisplayName(ctx context.Context, toolName string) string {
35+
return a.runtime.CurrentAgent().ToolDisplayName(ctx, toolName)
36+
}
37+
3438
// Run one agent loop
3539
func (a *App) Run(ctx context.Context, message string) {
3640
go func() {

‎internal/tui/components/tool/tool.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package tool
22

33
import (
4+
"context"
45
"fmt"
56
"strings"
67

@@ -13,7 +14,6 @@ import (
1314
"github.com/docker/cagent/internal/tui/core/layout"
1415
"github.com/docker/cagent/internal/tui/styles"
1516
"github.com/docker/cagent/internal/tui/types"
16-
"github.com/docker/cagent/pkg/tools"
1717
)
1818

1919
// Model represents a view that can render a message
@@ -129,7 +129,7 @@ func (mv *toolModel) Render(width int) string {
129129
spinnerText = " " + mv.spinner.View()
130130
}
131131

132-
displayName := tools.DisplayName(msg.ToolName)
132+
displayName := mv.app.ToolDisplayName(context.TODO(), msg.ToolName)
133133
content := fmt.Sprintf("%s %s%s", icon, styles.HighlightStyle.Render(displayName), spinnerText)
134134

135135
if msg.Arguments != "" {

‎pkg/agent/agent.go‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package agent
33
import (
44
"context"
55
"fmt"
6+
"log/slog"
67
"math/rand"
78
"sync/atomic"
89

@@ -102,6 +103,22 @@ func (a *Agent) Tools(ctx context.Context) ([]tools.Tool, error) {
102103
return agentTools, nil
103104
}
104105

106+
func (a *Agent) ToolDisplayName(ctx context.Context, toolName string) string {
107+
allTools, err := a.Tools(ctx)
108+
if err != nil {
109+
slog.Error("Failed to get tools for display name", "agent", a.Name(), "error", err)
110+
return toolName
111+
}
112+
113+
for _, tool := range allTools {
114+
if tool.Function.Name == toolName {
115+
return tool.DisplayName()
116+
}
117+
}
118+
119+
return toolName
120+
}
121+
105122
func (a *Agent) ToolSets() []tools.ToolSet {
106123
return a.toolsets
107124
}

‎pkg/runtime/runtime.go‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -101,17 +101,17 @@ func (r *Runtime) Team() *team.Team {
101101
return r.team
102102
}
103103

104+
func (r *Runtime) CurrentAgent() *agent.Agent {
105+
return r.team.Agent(r.currentAgent)
106+
}
107+
104108
// registerDefaultTools registers the default tool handlers
105109
func (r *Runtime) registerDefaultTools() {
106110
slog.Debug("Registering default tools")
107111
r.toolMap["transfer_task"] = r.handleTaskTransfer
108112
slog.Debug("Registered default tools", "count", len(r.toolMap))
109113
}
110114

111-
func (r *Runtime) CurrentAgent() *agent.Agent {
112-
return r.team.Agent(r.currentAgent)
113-
}
114-
115115
func (r *Runtime) finalizeEventChannel(ctx context.Context, sess *session.Session, events chan Event) {
116116
defer close(events)
117117

‎pkg/tools/builtin/filesystem.go‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,9 @@ func (t *FilesystemTool) Tools(context.Context) ([]tools.Tool, error) {
7171
Function: &tools.FunctionDefinition{
7272
Name: "create_directory",
7373
Description: "Create a new directory or ensure a directory exists. Can create multiple nested directories in one operation.",
74+
Annotations: tools.ToolAnnotation{
75+
Title: "Create Directory",
76+
},
7477
Parameters: tools.FunctionParamaters{
7578
Type: "object",
7679
Properties: map[string]any{
@@ -90,6 +93,7 @@ func (t *FilesystemTool) Tools(context.Context) ([]tools.Tool, error) {
9093
Description: "Get a recursive tree view of files and directories as a JSON structure.",
9194
Annotations: tools.ToolAnnotation{
9295
ReadOnlyHint: &[]bool{true}[0],
96+
Title: "Directory Tree",
9397
},
9498
Parameters: tools.FunctionParamaters{
9599
Type: "object",
@@ -112,6 +116,9 @@ func (t *FilesystemTool) Tools(context.Context) ([]tools.Tool, error) {
112116
Function: &tools.FunctionDefinition{
113117
Name: "edit_file",
114118
Description: "Make line-based edits to a text file. Each edit replaces exact line sequences with new content.",
119+
Annotations: tools.ToolAnnotation{
120+
Title: "Edit File",
121+
},
115122
Parameters: tools.FunctionParamaters{
116123
Type: "object",
117124
Properties: map[string]any{
@@ -153,6 +160,7 @@ func (t *FilesystemTool) Tools(context.Context) ([]tools.Tool, error) {
153160
Description: "Retrieve detailed metadata about a file or directory.",
154161
Annotations: tools.ToolAnnotation{
155162
ReadOnlyHint: &[]bool{true}[0],
163+
Title: "Get File Info",
156164
},
157165
Parameters: tools.FunctionParamaters{
158166
Type: "object",
@@ -173,6 +181,7 @@ func (t *FilesystemTool) Tools(context.Context) ([]tools.Tool, error) {
173181
Description: "Returns a list of directories that the server has permission to access.",
174182
Annotations: tools.ToolAnnotation{
175183
ReadOnlyHint: &[]bool{true}[0],
184+
Title: "List Allowed Directories",
176185
},
177186
Parameters: tools.FunctionParamaters{
178187
Type: "object",
@@ -185,6 +194,9 @@ func (t *FilesystemTool) Tools(context.Context) ([]tools.Tool, error) {
185194
Function: &tools.FunctionDefinition{
186195
Name: "add_allowed_directory",
187196
Description: "Request to add a new directory to the allowed directories list. This requires explicit user consent for security reasons.",
197+
Annotations: tools.ToolAnnotation{
198+
Title: "Add Allowed Directory",
199+
},
188200
Parameters: tools.FunctionParamaters{
189201
Type: "object",
190202
Properties: map[string]any{
@@ -212,6 +224,7 @@ func (t *FilesystemTool) Tools(context.Context) ([]tools.Tool, error) {
212224
Description: "Get a detailed listing of all files and directories in a specified path.",
213225
Annotations: tools.ToolAnnotation{
214226
ReadOnlyHint: &[]bool{true}[0],
227+
Title: "List Directory",
215228
},
216229
Parameters: tools.FunctionParamaters{
217230
Type: "object",
@@ -232,6 +245,7 @@ func (t *FilesystemTool) Tools(context.Context) ([]tools.Tool, error) {
232245
Description: "Get a detailed listing of all files and directories in a specified path, including sizes.",
233246
Annotations: tools.ToolAnnotation{
234247
ReadOnlyHint: &[]bool{true}[0],
248+
Title: "List Directory With Sizes",
235249
},
236250
Parameters: tools.FunctionParamaters{
237251
Type: "object",
@@ -250,6 +264,9 @@ func (t *FilesystemTool) Tools(context.Context) ([]tools.Tool, error) {
250264
Function: &tools.FunctionDefinition{
251265
Name: "move_file",
252266
Description: "Move or rename files and directories.",
267+
Annotations: tools.ToolAnnotation{
268+
Title: "Move File",
269+
},
253270
Parameters: tools.FunctionParamaters{
254271
Type: "object",
255272
Properties: map[string]any{
@@ -273,6 +290,7 @@ func (t *FilesystemTool) Tools(context.Context) ([]tools.Tool, error) {
273290
Description: "Read the complete contents of a file from the file system.",
274291
Annotations: tools.ToolAnnotation{
275292
ReadOnlyHint: &[]bool{true}[0],
293+
Title: "Read File",
276294
},
277295
Parameters: tools.FunctionParamaters{
278296
Type: "object",
@@ -293,6 +311,7 @@ func (t *FilesystemTool) Tools(context.Context) ([]tools.Tool, error) {
293311
Description: "Read the contents of multiple files simultaneously.",
294312
Annotations: tools.ToolAnnotation{
295313
ReadOnlyHint: &[]bool{true}[0],
314+
Title: "Read Multiple Files",
296315
},
297316
Parameters: tools.FunctionParamaters{
298317
Type: "object",
@@ -316,6 +335,7 @@ func (t *FilesystemTool) Tools(context.Context) ([]tools.Tool, error) {
316335
Description: "Recursively search for files and directories matching a pattern.",
317336
Annotations: tools.ToolAnnotation{
318337
ReadOnlyHint: &[]bool{true}[0],
338+
Title: "Search Files",
319339
},
320340
Parameters: tools.FunctionParamaters{
321341
Type: "object",
@@ -347,6 +367,7 @@ func (t *FilesystemTool) Tools(context.Context) ([]tools.Tool, error) {
347367
Description: "Searches for text or regex patterns in the content of files matching a GLOB pattern.",
348368
Annotations: tools.ToolAnnotation{
349369
ReadOnlyHint: &[]bool{true}[0],
370+
Title: "Search Files Content",
350371
},
351372
Parameters: tools.FunctionParamaters{
352373
Type: "object",
@@ -380,6 +401,9 @@ func (t *FilesystemTool) Tools(context.Context) ([]tools.Tool, error) {
380401
Function: &tools.FunctionDefinition{
381402
Name: "write_file",
382403
Description: "Create a new file or completely overwrite an existing file with new content.",
404+
Annotations: tools.ToolAnnotation{
405+
Title: "Write File",
406+
},
383407
Parameters: tools.FunctionParamaters{
384408
Type: "object",
385409
Properties: map[string]any{

‎pkg/tools/builtin/filesystem_test.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,8 @@ func TestFilesystemTool_DisplayNames(t *testing.T) {
7272
require.NoError(t, err)
7373

7474
for _, tool := range all {
75-
assert.NotEqual(t, tool.Function.Name, tools.DisplayName(tool.Function.Name))
75+
assert.NotEmpty(t, tool.DisplayName())
76+
assert.NotEqual(t, tool.Function.Name, tool.DisplayName())
7677
}
7778
}
7879

‎pkg/tools/builtin/memory.go‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,9 @@ func (t *MemoryTool) Tools(context.Context) ([]tools.Tool, error) {
4040
Function: &tools.FunctionDefinition{
4141
Name: "add_memory",
4242
Description: "Add a new memory to the database",
43+
Annotations: tools.ToolAnnotation{
44+
Title: "Add Memory",
45+
},
4346
Parameters: tools.FunctionParamaters{
4447
Type: "object",
4548
Properties: map[string]any{
@@ -59,6 +62,7 @@ func (t *MemoryTool) Tools(context.Context) ([]tools.Tool, error) {
5962
Description: "Retrieve all stored memories",
6063
Annotations: tools.ToolAnnotation{
6164
ReadOnlyHint: &[]bool{true}[0],
65+
Title: "Get Memories",
6266
},
6367
Parameters: tools.FunctionParamaters{
6468
Type: "object",
@@ -71,6 +75,9 @@ func (t *MemoryTool) Tools(context.Context) ([]tools.Tool, error) {
7175
Function: &tools.FunctionDefinition{
7276
Name: "delete_memory",
7377
Description: "Delete a specific memory by ID",
78+
Annotations: tools.ToolAnnotation{
79+
Title: "Delete Memory",
80+
},
7481
Parameters: tools.FunctionParamaters{
7582
Type: "object",
7683
Properties: map[string]any{

‎pkg/tools/builtin/memory_test.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,8 @@ func TestMemoryTool_DisplayNames(t *testing.T) {
7373
require.NoError(t, err)
7474

7575
for _, tool := range all {
76-
assert.NotEqual(t, tool.Function.Name, tools.DisplayName(tool.Function.Name))
76+
assert.NotEmpty(t, tool.DisplayName())
77+
assert.NotEqual(t, tool.Function.Name, tool.DisplayName())
7778
}
7879
}
7980

‎pkg/tools/builtin/shell.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -162,6 +162,9 @@ func (t *ShellTool) Tools(context.Context) ([]tools.Tool, error) {
162162
Function: &tools.FunctionDefinition{
163163
Name: "shell",
164164
Description: `Executes the given shell command in the user's default shell.`,
165+
Annotations: tools.ToolAnnotation{
166+
Title: "Run Shell Command",
167+
},
165168
Parameters: tools.FunctionParamaters{
166169
Type: "object",
167170
Properties: map[string]any{

‎pkg/tools/builtin/shell_test.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,8 @@ func TestShellTool_DisplayNames(t *testing.T) {
6666
require.NoError(t, err)
6767

6868
for _, tool := range all {
69-
assert.NotEqual(t, tool.Function.Name, tools.DisplayName(tool.Function.Name))
69+
assert.NotEmpty(t, tool.DisplayName())
70+
assert.NotEqual(t, tool.Function.Name, tool.DisplayName())
7071
}
7172
}
7273

0 commit comments

Comments
 (0)