Add read_file and write_file tools

Also improve `read_dir` description, and make it skip hidden files
This commit is contained in:
Matt Low 2023-11-26 07:32:28 +00:00
parent 4ae5c5e717
commit 07cc8306c1

View File

@ -44,6 +44,7 @@ var AvailableTools = map[string]AvailableTool{
Results are returned as JSON in the following format: Results are returned as JSON in the following format:
{ {
"message": "success", // if successful, or a different message indicating failure "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": [ "result": [
{"name": "a_file", "type": "file", "size": 123}, {"name": "a_file", "type": "file", "size": 123},
{"name": "a_directory/", "type": "dir", "size": 11}, {"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 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 { func resultToJson(result FunctionResult) string {
@ -143,27 +221,36 @@ func ExecuteToolCalls(toolCalls []openai.ToolCall) ([]Message, error) {
// VM/container. // VM/container.
func isPathContained(directory string, path string) (bool, error) { func isPathContained(directory string, path string) (bool, error) {
// Clean and resolve symlinks for both paths // Clean and resolve symlinks for both paths
absPath, err := filepath.Abs(path) path, err := filepath.Abs(path)
if err != nil {
return false, err
}
realPath, err := filepath.EvalSymlinks(absPath)
if err != nil { if err != nil {
return false, err 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 { if err != nil {
return false, err return false, err
} }
realDirectory, err := filepath.EvalSymlinks(absDirectory) directory, err = filepath.EvalSymlinks(directory)
if err != nil { if err != nil {
return false, err return false, err
} }
// Case insensitive checks // Case insensitive checks
if !strings.EqualFold(realPath, realDirectory) && if !strings.EqualFold(path, directory) &&
!strings.HasPrefix(strings.ToLower(realPath), strings.ToLower(realDirectory)+string(os.PathSeparator)) { !strings.HasPrefix(strings.ToLower(path), strings.ToLower(directory)+string(os.PathSeparator)) {
return false, nil return false, nil
} }
@ -206,6 +293,11 @@ func ReadDir(path string) string {
info, _ := f.Info() info, _ := f.Info()
name := f.Name() name := f.Name()
if strings.HasPrefix(name, ".") {
// skip hidden files
continue
}
entryType := "file" entryType := "file"
size := info.Size() size := info.Size()
@ -225,3 +317,29 @@ func ReadDir(path string) string {
return resultToJson(FunctionResult{Result: dirContents}) 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{})
}