From 963b29567e0a37761c8094892e9c420365b9c479 Mon Sep 17 00:00:00 2001
From: Matt Low <matt@mlow.ca>
Date: Sun, 5 Nov 2023 18:29:57 +0000
Subject: [PATCH] Parameterize maxTokens

---
 pkg/cli/cmd.go    |  7 +++++--
 pkg/cli/openai.go | 30 +++++++++++-------------------
 2 files changed, 16 insertions(+), 21 deletions(-)

diff --git a/pkg/cli/cmd.go b/pkg/cli/cmd.go
index 6aff35f..0562bc8 100644
--- a/pkg/cli/cmd.go
+++ b/pkg/cli/cmd.go
@@ -7,6 +7,9 @@ import (
 	"github.com/spf13/cobra"
 )
 
+// TODO: allow setting with flag
+const MAX_TOKENS = 256
+
 var rootCmd = &cobra.Command{
 	Use:   "lmcli",
 	Short: "Interact with Large Language Models",
@@ -120,7 +123,7 @@ var newCmd = &cobra.Command{
 			response <- HandleDelayedResponse(receiver)
 		}()
 
-		err = CreateChatCompletionStream(messages, receiver)
+		err = CreateChatCompletionStream(messages, MAX_TOKENS, receiver)
 		if err != nil {
 			Fatal("%v\n", err)
 		}
@@ -160,7 +163,7 @@ var promptCmd = &cobra.Command{
 
 		receiver := make(chan string)
 		go HandleDelayedResponse(receiver)
-		err := CreateChatCompletionStream(messages, receiver)
+		err := CreateChatCompletionStream(messages, MAX_TOKENS, receiver)
 		if err != nil {
 			Fatal("%v\n", err)
 		}
diff --git a/pkg/cli/openai.go b/pkg/cli/openai.go
index 21d77af..f6d25ec 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) *openai.ChatCompletionRequest {
+func CreateChatCompletionRequest(messages []Message, maxTokens int) openai.ChatCompletionRequest {
 	chatCompletionMessages := []openai.ChatCompletionMessage{}
 	for _, m := range messages {
 		chatCompletionMessages = append(chatCompletionMessages, openai.ChatCompletionMessage{
@@ -17,23 +17,19 @@ func CreateChatCompletionRequest(messages []Message) *openai.ChatCompletionReque
 		})
 	}
 
-	return &openai.ChatCompletionRequest{
-		Model:     openai.GPT4,
-		MaxTokens: 256,
+	return openai.ChatCompletionRequest{
+		Model:     openai.GPT3Dot5Turbo,
 		Messages:  chatCompletionMessages,
-		Stream:    true,
+		MaxTokens: maxTokens,
 	}
 }
 
 // CreateChatCompletion submits a Chat Completion API request and returns the
 // response.
-func CreateChatCompletion(messages []Message) (string, error) {
+func CreateChatCompletion(messages []Message, maxTokens int) (string, error) {
 	client := openai.NewClient(config.OpenAI.APIKey)
-	resp, err := client.CreateChatCompletion(
-		context.Background(),
-		*CreateChatCompletionRequest(messages),
-	)
-
+	req := CreateChatCompletionRequest(messages, maxTokens)
+	resp, err := client.CreateChatCompletion(context.Background(), req)
 	if err != nil {
 		return "", err
 	}
@@ -42,17 +38,14 @@ func CreateChatCompletion(messages []Message) (string, error) {
 }
 
 // CreateChatCompletionStream submits a streaming Chat Completion API request
-// and streams the received response to the provided output channel.
-func CreateChatCompletionStream(messages []Message, output chan string) error {
+// and streams the response to the provided output channel.
+func CreateChatCompletionStream(messages []Message, maxTokens int, output chan string) error {
 	client := openai.NewClient(config.OpenAI.APIKey)
-	ctx := context.Background()
-
-	req := CreateChatCompletionRequest(messages)
-	req.Stream = true
+	req := CreateChatCompletionRequest(messages, maxTokens)
 
 	defer close(output)
 
-	stream, err := client.CreateChatCompletionStream(ctx, *req)
+	stream, err := client.CreateChatCompletionStream(context.Background(), req)
 	if err != nil {
 		return err
 	}
@@ -68,7 +61,6 @@ func CreateChatCompletionStream(messages []Message, output chan string) error {
 		if err != nil {
 			return err
 		}
-
 		output <- response.Choices[0].Delta.Content
 	}
 }