From ae424530f9637c3b06128fad7dfe1a6ebd1d1dcf Mon Sep 17 00:00:00 2001 From: Matt Low Date: Thu, 9 Nov 2023 06:07:52 +0000 Subject: [PATCH] Parameterize the openai model used Add `openai.defaultConfig` to set the default, will allow overriding with CLI flag --- pkg/cli/cmd.go | 4 ++-- pkg/cli/config.go | 1 + pkg/cli/openai.go | 12 ++++++------ 3 files changed, 9 insertions(+), 8 deletions(-) diff --git a/pkg/cli/cmd.go b/pkg/cli/cmd.go index 0562bc8..bdcc946 100644 --- a/pkg/cli/cmd.go +++ b/pkg/cli/cmd.go @@ -123,7 +123,7 @@ var newCmd = &cobra.Command{ response <- HandleDelayedResponse(receiver) }() - err = CreateChatCompletionStream(messages, MAX_TOKENS, receiver) + err = CreateChatCompletionStream(config.OpenAI.DefaultModel, messages, MAX_TOKENS, receiver) if err != nil { Fatal("%v\n", err) } @@ -163,7 +163,7 @@ var promptCmd = &cobra.Command{ receiver := make(chan string) go HandleDelayedResponse(receiver) - err := CreateChatCompletionStream(messages, MAX_TOKENS, receiver) + err := CreateChatCompletionStream(config.OpenAI.DefaultModel, messages, MAX_TOKENS, receiver) if err != nil { Fatal("%v\n", err) } diff --git a/pkg/cli/config.go b/pkg/cli/config.go index 0458a3a..41aa2a7 100644 --- a/pkg/cli/config.go +++ b/pkg/cli/config.go @@ -11,6 +11,7 @@ import ( type Config struct { OpenAI struct { APIKey string `yaml:"apiKey"` + DefaultModel string `yaml:"defaultModel"` } `yaml:"openai"` } diff --git a/pkg/cli/openai.go b/pkg/cli/openai.go index f6d25ec..09435e6 100644 --- a/pkg/cli/openai.go +++ b/pkg/cli/openai.go @@ -8,7 +8,7 @@ import ( openai "github.com/sashabaranov/go-openai" ) -func CreateChatCompletionRequest(messages []Message, maxTokens int) openai.ChatCompletionRequest { +func CreateChatCompletionRequest(model string, messages []Message, maxTokens int) openai.ChatCompletionRequest { chatCompletionMessages := []openai.ChatCompletionMessage{} for _, m := range messages { chatCompletionMessages = append(chatCompletionMessages, openai.ChatCompletionMessage{ @@ -18,7 +18,7 @@ func CreateChatCompletionRequest(messages []Message, maxTokens int) openai.ChatC } return openai.ChatCompletionRequest{ - Model: openai.GPT3Dot5Turbo, + Model: model, Messages: chatCompletionMessages, MaxTokens: maxTokens, } @@ -26,9 +26,9 @@ func CreateChatCompletionRequest(messages []Message, maxTokens int) openai.ChatC // CreateChatCompletion submits a Chat Completion API request and returns the // response. -func CreateChatCompletion(messages []Message, maxTokens int) (string, error) { +func CreateChatCompletion(model string, messages []Message, maxTokens int) (string, error) { client := openai.NewClient(config.OpenAI.APIKey) - req := CreateChatCompletionRequest(messages, maxTokens) + req := CreateChatCompletionRequest(model, messages, maxTokens) resp, err := client.CreateChatCompletion(context.Background(), req) if err != nil { return "", err @@ -39,9 +39,9 @@ func CreateChatCompletion(messages []Message, maxTokens int) (string, error) { // CreateChatCompletionStream submits a streaming Chat Completion API request // and streams the response to the provided output channel. -func CreateChatCompletionStream(messages []Message, maxTokens int, output chan string) error { +func CreateChatCompletionStream(model string, messages []Message, maxTokens int, output chan string) error { client := openai.NewClient(config.OpenAI.APIKey) - req := CreateChatCompletionRequest(messages, maxTokens) + req := CreateChatCompletionRequest(model, messages, maxTokens) defer close(output)