From ae424530f9637c3b06128fad7dfe1a6ebd1d1dcf Mon Sep 17 00:00:00 2001
From: Matt Low <matt@mlow.ca>
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)