147 lines
4.0 KiB
Go
147 lines
4.0 KiB
Go
package openapi
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"os"
|
|
"os/exec"
|
|
"path/filepath"
|
|
"strings"
|
|
)
|
|
|
|
// RunCommandArgs defines the arguments for the run_command tool.
|
|
type RunCommandArgs struct {
|
|
Command string `json:"command" desc:"The command to run"`
|
|
}
|
|
|
|
// NewShellRegistry returns a registry with shell-related tools.
|
|
func NewShellRegistry() ToolRegistry {
|
|
r := NewMapRegistry()
|
|
|
|
RegisterTool(r, "run_command", "Execute a shell command",
|
|
func(ctx context.Context, s *Session, params RunCommandArgs) (any, error) {
|
|
cmd := exec.CommandContext(ctx, "sh", "-c", params.Command)
|
|
output, err := cmd.CombinedOutput()
|
|
return map[string]any{
|
|
"output": string(output),
|
|
"error": fmt.Sprintf("%v", err),
|
|
}, nil
|
|
})
|
|
|
|
return r
|
|
}
|
|
|
|
// ReadFileArgs defines the arguments for the read_file tool.
|
|
type ReadFileArgs struct {
|
|
Path string `json:"path" desc:"Path to the file"`
|
|
Start int `json:"start,omitempty" desc:"Start line number"`
|
|
Lines int `json:"lines,omitempty" desc:"number of lines to read, default 200"`
|
|
}
|
|
|
|
// WriteFileArgs defines the arguments for the write_file tool.
|
|
type WriteFileArgs struct {
|
|
Path string `json:"path" desc:"Path to the file"`
|
|
Content string `json:"content" desc:"Content to write"`
|
|
}
|
|
|
|
// ListDirArgs defines the arguments for the list_dir tool.
|
|
type ListDirArgs struct {
|
|
Path string `json:"path" desc:"Path to the directory"`
|
|
}
|
|
|
|
// ReplaceContentArgs defines the arguments for the replace_content tool.
|
|
type ReplaceContentArgs struct {
|
|
Path string `json:"path" desc:"Path to the file"`
|
|
TargetContent string `json:"target_content" desc:"The exact content to replace"`
|
|
NewContent string `json:"new_content" desc:"The new content"`
|
|
}
|
|
|
|
// NewFileRegistry returns a registry with file-related tools.
|
|
func NewFileRegistry(baseDir string) ToolRegistry {
|
|
r := NewMapRegistry()
|
|
|
|
RegisterTool(r, "read_file", "Read the contents of a file",
|
|
func(ctx context.Context, s *Session, params ReadFileArgs) (any, error) {
|
|
fullPath := filepath.Join(baseDir, params.Path)
|
|
content, err := os.ReadFile(fullPath)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
lines := strings.Split(string(content), "\n")
|
|
if params.Start < 1 {
|
|
params.Start = 1
|
|
}
|
|
if params.Lines == 0 {
|
|
params.Lines = 200
|
|
}
|
|
|
|
start := params.Start - 1
|
|
end := start + params.Lines
|
|
if end > len(lines) {
|
|
end = len(lines)
|
|
}
|
|
|
|
if start >= end {
|
|
return nil, fmt.Errorf("invalid start or lines parameters")
|
|
}
|
|
|
|
return strings.Join(lines[start:end], "\n"), nil
|
|
})
|
|
|
|
RegisterTool(r, "write_file", "Write content to a file",
|
|
func(ctx context.Context, s *Session, params WriteFileArgs) (any, error) {
|
|
fullPath := filepath.Join(baseDir, params.Path)
|
|
if err := os.MkdirAll(filepath.Dir(fullPath), 0755); err != nil {
|
|
return nil, err
|
|
}
|
|
if err := os.WriteFile(fullPath, []byte(params.Content), 0644); err != nil {
|
|
return nil, err
|
|
}
|
|
return "File written successfully", nil
|
|
})
|
|
|
|
RegisterTool(r, "list_dir", "List contents of a directory",
|
|
func(ctx context.Context, s *Session, params ListDirArgs) (any, error) {
|
|
fullPath := filepath.Join(baseDir, params.Path)
|
|
entries, err := os.ReadDir(fullPath)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
var result []string
|
|
for _, entry := range entries {
|
|
name := entry.Name()
|
|
if entry.IsDir() {
|
|
name += "/"
|
|
}
|
|
result = append(result, name)
|
|
}
|
|
return result, nil
|
|
})
|
|
|
|
RegisterTool(r, "replace_content", "Replace a block of text in a file",
|
|
func(ctx context.Context, s *Session, params ReplaceContentArgs) (any, error) {
|
|
fullPath := filepath.Join(baseDir, params.Path)
|
|
content, err := os.ReadFile(fullPath)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
oldContent := string(content)
|
|
if !strings.Contains(oldContent, params.TargetContent) {
|
|
return nil, fmt.Errorf("target content not found in file")
|
|
}
|
|
|
|
newContent := strings.Replace(oldContent, params.TargetContent, params.NewContent, 1)
|
|
if err := os.WriteFile(fullPath, []byte(newContent), 0644); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return "Content replaced successfully", nil
|
|
})
|
|
|
|
return r
|
|
}
|
|
|