From 168e0cf5d329ded6c9652194d731bbe5a8fa3ca4 Mon Sep 17 00:00:00 2001
From: Matt Low <matt@mlow.ca>
Date: Sun, 5 Nov 2023 18:19:30 +0000
Subject: [PATCH] Parameterize maxTokens

Minor formatting/commet changes
---
 pkg/cli/cmd.go    |  9 ++++++---
 pkg/cli/openai.go | 36 ++++++++++++++----------------------
 pkg/cli/tty.go    |  2 ++
 3 files changed, 22 insertions(+), 25 deletions(-)

diff --git a/pkg/cli/cmd.go b/pkg/cli/cmd.go
index efde7af..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",
@@ -112,7 +115,6 @@ var newCmd = &cobra.Command{
 			ConversationID: conversation.ID,
 			Role:           "assistant",
 		}
-
 		reply.RenderTTY(false)
 
 		receiver := make(chan string)
@@ -120,7 +122,8 @@ var newCmd = &cobra.Command{
 		go func() {
 			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 1251426..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 accepts a slice of Message and returns the response
-// of the Large Language Model.
-func CreateChatCompletion(messages []Message) (string, error) {
+// CreateChatCompletion submits a Chat Completion API request and returns the
+// response.
+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
 	}
@@ -41,18 +37,15 @@ func CreateChatCompletion(messages []Message) (string, error) {
 	return resp.Choices[0].Message.Content, nil
 }
 
-// CreateChatCompletionStream submits an streaming Chat Completion API request
-// and sends the received data to the output channel.
-func CreateChatCompletionStream(messages []Message, output chan 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 {
 	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
 	}
 }
diff --git a/pkg/cli/tty.go b/pkg/cli/tty.go
index e2dab81..159da00 100644
--- a/pkg/cli/tty.go
+++ b/pkg/cli/tty.go
@@ -45,7 +45,9 @@ func HandleDelayedResponse(response chan string) string {
 	firstChunk := true
 	for chunk := range response {
 		if firstChunk {
+			// notify wait animation that we've received data
 			waitSignal <- ""
+			// wait for signal that wait animation has completed
 			<-waitSignal
 			firstChunk = false
 		}