lmcli/pkg/cli/functions.go

package cli

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"

	openai "github.com/sashabaranov/go-openai"
)

type FunctionResult struct {
	Message string `json:"message"`
	Result  any    `json:"result,omitempty"`
}

type FunctionParameter struct {
	Type        string   `json:"type"` // "string", "integer", "boolean"
	Description string   `json:"description"`
	Enum        []string `json:"enum,omitempty"`
}

type FunctionParameters struct {
	Type       string                       `json:"type"` // "object"
	Properties map[string]FunctionParameter `json:"properties"`
	Required   []string                     `json:"required,omitempty"` // required function parameter names
}

type AvailableTool struct {
	openai.Tool
	// The tool's implementation. Returns a string, as tool call results
	// are treated as normal messages with string contents.
	Impl func(arguments map[string]interface{}) (string, error)
}

var AvailableTools = map[string]AvailableTool{
	"read_dir": {
		Tool: openai.Tool{Type: "function", Function: openai.FunctionDefinition{
			Name: "read_dir",
			Description: `Return the contents of the CWD (current working directory).

Results are returned as JSON in the following format:
{
	"message": "success", // "success" if successful, or a different message indicating failure
	"result": [
		{"name": "a_file", "type": "file", "length": 123},
		{"name": "a_directory", "type": "dir", "length": 5},
		... // more files or directories
	]
}

For type: file, length represents the size (in bytes) of the file.
For type: dir, length represents the number of entries in that directory.`,
			Parameters: FunctionParameters{
				Type: "object",
				Properties: map[string]FunctionParameter{
					"relative_dir": {
						Type:        "string",
						Description: "If set, read the contents of a directory relative to the current one.",
					},
				},
			},
		}},
		Impl: func(args map[string]interface{}) (string, error) {
			var relativeDir string
			tmp, ok := args["relative_dir"]
			if ok {
				relativeDir, ok = tmp.(string)
				if !ok {
					return "", fmt.Errorf("Invalid relative_dir in function arguments: %v", tmp)
				}
			}
			return ReadDir(relativeDir), nil
		},
	},
}

func resultToJson(result FunctionResult) string {
	if result.Message == "" {
		// When message not supplied, assume success
		result.Message = "success"
	}

	jsonBytes, err := json.Marshal(result)
	if err != nil {
		fmt.Printf("Could not marshal FunctionResult to JSON: %v\n", err)
	}
	return string(jsonBytes)
}

// ExecuteToolCalls handles the execution of all tool_calls provided, and
// returns their results formatted as []Message(s) with role: 'tool' and.
func ExecuteToolCalls(toolCalls []openai.ToolCall) ([]Message, error) {
	var toolResults []Message
	for _, toolCall := range toolCalls {
		if toolCall.Type != "function" {
			// unsupported tool type
			continue
		}

		tool, ok := AvailableTools[toolCall.Function.Name]
		if !ok {
			return nil, fmt.Errorf("Requested tool '%s' does not exist. Hallucination?", toolCall.Function.Name)
		}

		var functionArgs map[string]interface{}
		err := json.Unmarshal([]byte(toolCall.Function.Arguments), &functionArgs)
		if err != nil {
			return nil, fmt.Errorf("Could not unmarshal tool arguments. Malformed JSON? Error: %v", err)
		}

		// TODO: ability to silence this
		fmt.Fprintf(os.Stderr, "INFO: Executing tool '%s' with args %s\n", toolCall.Function.Name, toolCall.Function.Arguments)

		// Execute the tool
		toolResult, err := tool.Impl(functionArgs)
		if err != nil {
			// This can happen if the model missed or supplied invalid tool args
			return nil, fmt.Errorf("Tool '%s' error: %v\n", toolCall.Function.Name, err)
		}

		toolResults = append(toolResults, Message{
			Role:            "tool",
			OriginalContent: toolResult,
			ToolCallID:      sql.NullString{String: toolCall.ID, Valid: true},
			// name is not required since the introduction of ToolCallID
			// hypothesis: by setting it, we inform the model of what a
			// function's purpose was if future requests omit the function
			// definition
		})
	}
	return toolResults, nil
}

func ReadDir(path string) string {
	// TODO: ensure it is not possible to escape to directories above CWD
	// TODO: implement whitelist - list of directories which model is allowed to work in
	targetPath := filepath.Join(".", path)
	files, err := os.ReadDir(targetPath)
	if err != nil {
		return resultToJson(FunctionResult{
			Message: err.Error(),
		})
	}

	var dirContents []map[string]interface{}
	for _, f := range files {
		info, _ := f.Info()

		contentType := "file"
		length := info.Size()

		if info.IsDir() {
			contentType = "dir"
			subdirfiles, _ := os.ReadDir(filepath.Join(".", path, info.Name()))
			length = int64(len(subdirfiles))
		}

		dirContents = append(dirContents, map[string]interface{}{
			"name":   f.Name(),
			"type":   contentType,
			"length": length,
		})
	}

	return resultToJson(FunctionResult{Result: dirContents})
}