diff --git a/pkg/cli/config.go b/pkg/cli/config.go index c4108b0..0acc2d1 100644 --- a/pkg/cli/config.go +++ b/pkg/cli/config.go @@ -13,9 +13,10 @@ type Config struct { SystemPrompt *string `yaml:"systemPrompt" default:"You are a helpful assistant."` } `yaml:"modelDefaults"` OpenAI *struct { - APIKey *string `yaml:"apiKey" default:"your_key_here"` - DefaultModel *string `yaml:"defaultModel" default:"gpt-4"` - DefaultMaxLength *int `yaml:"defaultMaxLength" default:"256"` + APIKey *string `yaml:"apiKey" default:"your_key_here"` + DefaultModel *string `yaml:"defaultModel" default:"gpt-4"` + DefaultMaxLength *int `yaml:"defaultMaxLength" default:"256"` + EnabledTools []string `yaml:"enabledTools"` } `yaml:"openai"` Chroma *struct { Style *string `yaml:"style" default:"onedark"` diff --git a/pkg/cli/openai.go b/pkg/cli/openai.go index c4f9843..a308f04 100644 --- a/pkg/cli/openai.go +++ b/pkg/cli/openai.go @@ -34,20 +34,27 @@ func CreateChatCompletionRequest(model string, messages []Message, maxTokens int chatCompletionMessages = append(chatCompletionMessages, message) } - var tools []openai.Tool - for _, t := range AvailableTools { - // TODO: support some way to limit which tools are available per-request - tools = append(tools, t.Tool) + request := openai.ChatCompletionRequest{ + Model: model, + Messages: chatCompletionMessages, + MaxTokens: maxTokens, + N: 1, // limit responses to 1 "choice". we use choices[0] to reference it } - return openai.ChatCompletionRequest{ - Model: model, - Messages: chatCompletionMessages, - MaxTokens: maxTokens, - N: 1, // limit responses to 1 "choice". we use choices[0] to reference it - Tools: tools, - ToolChoice: "auto", // TODO: allow limiting/forcing which function is called? + var tools []openai.Tool + for _, t := range config.OpenAI.EnabledTools { + tool, ok := AvailableTools[t] + if ok { + tools = append(tools, tool.Tool) + } } + + if len(tools) > 0 { + request.Tools = tools + request.ToolChoice = "auto" + } + + return request } // CreateChatCompletion submits a Chat Completion API request and returns the