Support custom tools (#223)

Author: janetkuo

cmd/main.go

@@ -86,6 +86,7 @@ type Options struct {
 	ExtraPromptPaths       []string `json:"extraPromptPaths,omitempty"`
 	TracePath              string   `json:"tracePath,omitempty"`
 	RemoveWorkDir          bool     `json:"removeWorkDir,omitempty"`
+	ToolConfigPath         string   `json:"toolConfigPath,omitempty"`
 
 	// UserInterface is the type of user interface to use.
 	UserInterface UserInterface `json:"userInterface,omitempty"`
@@ -236,6 +237,12 @@ func run(ctx context.Context) error {
 		return fmt.Errorf("failed to load config file: %w", err)
 	}
 
+	// Load and register custom tools from config file
+	if err := tools.LoadAndRegisterCustomTools(opt.ToolConfigPath); err != nil {
+		// Log the error but continue execution, as custom tools are optional
+		klog.Warningf("Failed to load or register custom tools (path: %q): %v", opt.ToolConfigPath, err)
+	}
+
 	rootCmd, err := BuildRootCommand(&opt)
 	if err != nil {
 		return err
@@ -267,6 +274,7 @@ func (opt *Options) bindCLIFlags(f *pflag.FlagSet) error {
 	f.StringVar(&opt.ModelID, "model", opt.ModelID, "language model e.g. gemini-2.0-flash-thinking-exp-01-21, gemini-2.0-flash")
 	f.BoolVar(&opt.SkipPermissions, "skip-permissions", opt.SkipPermissions, "(dangerous) skip asking for confirmation before executing kubectl commands that modify resources")
 	f.BoolVar(&opt.MCPServer, "mcp-server", opt.MCPServer, "run in MCP server mode")
+	f.StringVar(&opt.ToolConfigPath, "custom-tools-config", opt.ToolConfigPath, "path to custom tools config file")
 	f.BoolVar(&opt.EnableToolUseShim, "enable-tool-use-shim", opt.EnableToolUseShim, "enable tool use shim")
 	f.BoolVar(&opt.Quiet, "quiet", opt.Quiet, "run in non-interactive mode, requires a query to be provided as a positional argument")
 

pkg/tools/custom_tool.go

@@ -0,0 +1,96 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package tools
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"os/exec"
+	"strings"
+
+	"github.com/GoogleCloudPlatform/kubectl-ai/gollm"
+)
+
+// CustomToolConfig defines the structure for configuring a custom tool.
+type CustomToolConfig struct {
+	Name        string `yaml:"name"`
+	Description string `yaml:"description"`
+	Command     string `yaml:"command"`
+	CommandDesc string `yaml:"command_desc"`
+}
+
+// CustomTool implements the Tool interface for external commands.
+type CustomTool struct {
+	config CustomToolConfig
+}
+
+// NewCustomTool creates a new CustomTool instance.
+func NewCustomTool(config CustomToolConfig) (*CustomTool, error) {
+	if config.Name == "" {
+		return nil, fmt.Errorf("custom tool name cannot be empty")
+	}
+	if len(config.Command) == 0 {
+		return nil, fmt.Errorf("custom tool command cannot be empty for tool %q", config.Name)
+	}
+
+	return &CustomTool{config: config}, nil
+}
+
+// Name returns the tool's name.
+func (t *CustomTool) Name() string {
+	return t.config.Name
+}
+
+// Description returns the tool's description from its function definition.
+func (t *CustomTool) Description() string {
+	return t.config.Description
+}
+
+// FunctionDefinition returns the tool's function definition.
+func (t *CustomTool) FunctionDefinition() *gollm.FunctionDefinition {
+	return &gollm.FunctionDefinition{
+		Name:        t.Name(),
+		Description: t.Description(),
+		Parameters: &gollm.Schema{
+			Type: gollm.TypeObject,
+			Properties: map[string]*gollm.Schema{
+				"command": {
+					Type:        gollm.TypeString,
+					Description: t.config.CommandDesc,
+				},
+			},
+		},
+	}
+}
+
+// Run executes the external command defined for the custom tool.
+func (t *CustomTool) Run(ctx context.Context, args map[string]any) (any, error) {
+	command := strings.Fields(t.config.Command)
+	if len(command) == 0 {
+		return nil, fmt.Errorf("empty command")
+	}
+	cmdArgs := []string{}
+	if len(command) > 1 {
+		cmdArgs = command[1:]
+	}
+	workDir := ctx.Value(WorkDirKey).(string)
+
+	cmd := exec.CommandContext(ctx, command[0], cmdArgs...)
+	cmd.Dir = workDir
+	cmd.Env = os.Environ()
+
+	return executeCommand(cmd)
+}

pkg/tools/tools.go

@@ -19,13 +19,16 @@ import (
 	"encoding/json"
 	"fmt"
 	"maps"
+	"os"
+	"path/filepath"
 	"slices"
 	"sort"
 	"strings"
 	"time"
 
 	"github.com/GoogleCloudPlatform/kubectl-ai/pkg/journal"
 	"github.com/google/uuid"
+	"sigs.k8s.io/yaml"
 )
 
 type ContextKey string
@@ -180,3 +183,51 @@ func ToolResultToMap(result any) (map[string]any, error) {
 	}
 	return m, nil
 }
+
+// LoadAndRegisterCustomTools loads tool configurations from a YAML file
+// and registers them.
+func LoadAndRegisterCustomTools(configPath string) error {
+	if configPath == "" {
+		// Default config path: ~/.config/kubectl-ai/tools.yaml
+		home, err := os.UserHomeDir()
+		if err != nil {
+			return fmt.Errorf("failed to get user home directory for default custom tools config: %w", err)
+		}
+		configPath = filepath.Join(home, ".config", "kubectl-ai", "tools.yaml")
+	}
+
+	yamlFile, err := os.ReadFile(configPath)
+	if os.IsNotExist(err) {
+		return nil
+	} else if err != nil {
+		return fmt.Errorf("failed to read config file %s: %w", configPath, err)
+	}
+
+	var configs []CustomToolConfig
+	err = yaml.Unmarshal(yamlFile, &configs)
+	if err != nil {
+		return fmt.Errorf("failed to parse YAML config file %s: %w", configPath, err)
+	}
+
+	// Register each custom tool
+	var registrationErrors []string
+	for _, config := range configs {
+		tool, err := NewCustomTool(config)
+		if err != nil {
+			registrationErrors = append(registrationErrors, fmt.Sprintf("failed to create tool %q: %v", config.Name, err))
+			continue // Skip registration if creation failed
+		}
+		// Check for duplicate registration attempt
+		if _, exists := allTools.tools[tool.Name()]; exists {
+			registrationErrors = append(registrationErrors, fmt.Sprintf("tool %q already registered (possibly built-in), skipping custom definition", tool.Name()))
+			continue
+		}
+		RegisterTool(tool)
+	}
+
+	if len(registrationErrors) > 0 {
+		return fmt.Errorf("encountered errors during custom tool registration:\n - %s", strings.Join(registrationErrors, "\n - "))
+	}
+
+	return nil
+}