From 07cc8306c1778d42dae6e23edc2108a2309616d7 Mon Sep 17 00:00:00 2001 From: Matt Low Date: Sun, 26 Nov 2023 07:32:28 +0000 Subject: [PATCH] Add `read_file` and `write_file` tools Also improve `read_dir` description, and make it skip hidden files --- pkg/cli/functions.go | 136 ++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 127 insertions(+), 9 deletions(-) diff --git a/pkg/cli/functions.go b/pkg/cli/functions.go index ba12bbc..72af6fa 100644 --- a/pkg/cli/functions.go +++ b/pkg/cli/functions.go @@ -44,6 +44,7 @@ var AvailableTools = map[string]AvailableTool{ Results are returned as JSON in the following format: { "message": "success", // if successful, or a different message indicating failure + // result may be an empty array if there are no files in the directory "result": [ {"name": "a_file", "type": "file", "size": 123}, {"name": "a_directory/", "type": "dir", "size": 11}, @@ -75,6 +76,83 @@ For directories, size represents the number of entries in that directory.`, return ReadDir(relativeDir), nil }, }, + "read_file": { + Tool: openai.Tool{Type: "function", Function: openai.FunctionDefinition{ + Name: "read_file", + Description: `Read the contents of a file relative to the current working directory. + +Result is returned as JSON in the following format: +{ + "message": "success", // if successful, or a different message indicating failure + "result": "the contents\nof the file\n" +}`, + Parameters: FunctionParameters{ + Type: "object", + Properties: map[string]FunctionParameter{ + "path": { + Type: "string", + Description: "Path to a file within the current working directory to read.", + }, + }, + Required: []string{"file_path"}, + }, + }}, + Impl: func(args map[string]interface{}) (string, error) { + tmp, ok := args["path"] + if !ok { + return "", fmt.Errorf("Path parameter to read_file was not included.") + } + path, ok := tmp.(string) + if !ok { + return "", fmt.Errorf("Invalid path in function arguments: %v", tmp) + } + return ReadFile(path), nil + }, + }, + "write_file": { + Tool: openai.Tool{Type: "function", Function: openai.FunctionDefinition{ + Name: "write_file", + Description: `Write the provided contents to a file relative to the current working directory. + +Result is returned as JSON in the following format: +{ + "message": "success", // if successful, or a different message indicating failure +}`, + Parameters: FunctionParameters{ + Type: "object", + Properties: map[string]FunctionParameter{ + "path": { + Type: "string", + Description: "Path to a file within the current working directory to write to.", + }, + "content": { + Type: "string", + Description: "The content to write to the file. Overwrites any existing content!", + }, + }, + Required: []string{"file_path"}, + }, + }}, + Impl: func(args map[string]interface{}) (string, error) { + tmp, ok := args["path"] + if !ok { + return "", fmt.Errorf("Path parameter to write_file was not included.") + } + path, ok := tmp.(string) + if !ok { + return "", fmt.Errorf("Invalid path in function arguments: %v", tmp) + } + tmp, ok = args["content"] + if !ok { + return "", fmt.Errorf("Content parameter to write_file was not included.") + } + content, ok := tmp.(string) + if !ok { + return "", fmt.Errorf("Invalido content in function arguments: %v", tmp) + } + return WriteFile(path, content), nil + }, + }, } func resultToJson(result FunctionResult) string { @@ -143,27 +221,36 @@ func ExecuteToolCalls(toolCalls []openai.ToolCall) ([]Message, error) { // VM/container. func isPathContained(directory string, path string) (bool, error) { // Clean and resolve symlinks for both paths - absPath, err := filepath.Abs(path) - if err != nil { - return false, err - } - realPath, err := filepath.EvalSymlinks(absPath) + path, err := filepath.Abs(path) if err != nil { return false, err } - absDirectory, err := filepath.Abs(directory) + // check if path exists + _, err = os.Stat(path) + if err != nil { + if !os.IsNotExist(err) { + return false, fmt.Errorf("Could not stat path: %v", err) + } + } else { + path, err = filepath.EvalSymlinks(path) + if err != nil { + return false, err + } + } + + directory, err = filepath.Abs(directory) if err != nil { return false, err } - realDirectory, err := filepath.EvalSymlinks(absDirectory) + directory, err = filepath.EvalSymlinks(directory) if err != nil { return false, err } // Case insensitive checks - if !strings.EqualFold(realPath, realDirectory) && - !strings.HasPrefix(strings.ToLower(realPath), strings.ToLower(realDirectory)+string(os.PathSeparator)) { + if !strings.EqualFold(path, directory) && + !strings.HasPrefix(strings.ToLower(path), strings.ToLower(directory)+string(os.PathSeparator)) { return false, nil } @@ -206,6 +293,11 @@ func ReadDir(path string) string { info, _ := f.Info() name := f.Name() + if strings.HasPrefix(name, ".") { + // skip hidden files + continue + } + entryType := "file" size := info.Size() @@ -225,3 +317,29 @@ func ReadDir(path string) string { return resultToJson(FunctionResult{Result: dirContents}) } + +func ReadFile(path string) string { + ok, res := isPathWithinCWD(path) + if !ok { + return resultToJson(*res) + } + data, err := os.ReadFile(path) + if err != nil { + return resultToJson(FunctionResult{Message: fmt.Sprintf("Could not read path: %s", err.Error())}) + } + return resultToJson(FunctionResult{ + Result: string(data), + }) +} + +func WriteFile(path string, content string) string { + ok, res := isPathWithinCWD(path) + if !ok { + return resultToJson(*res) + } + err := os.WriteFile(path, []byte(content), 0644) + if err != nil { + return resultToJson(FunctionResult{Message: fmt.Sprintf("Could not write to path: %s", err.Error())}) + } + return resultToJson(FunctionResult{}) +}