pkg/cli/cmd.go

@@ -7,9 +7,6 @@ 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",
@@ -115,6 +112,7 @@ var newCmd = &cobra.Command{
 			ConversationID: conversation.ID,
 			Role:           "assistant",
 		}
+
 		reply.RenderTTY(false)
 
 		receiver := make(chan string)
@@ -122,8 +120,7 @@ var newCmd = &cobra.Command{
 		go func() {
 			response <- HandleDelayedResponse(receiver)
 		}()
-
-		err = CreateChatCompletionStream(messages, MAX_TOKENS, receiver)
+		err = CreateChatCompletionStream(messages, receiver)
 		if err != nil {
 			Fatal("%v\n", err)
 		}
@@ -163,7 +160,7 @@ var promptCmd = &cobra.Command{
 
 		receiver := make(chan string)
 		go HandleDelayedResponse(receiver)
-		err := CreateChatCompletionStream(messages, MAX_TOKENS, receiver)
+		err := CreateChatCompletionStream(messages, receiver)
 		if err != nil {
 			Fatal("%v\n", err)
 		}

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(messages []Message) *openai.ChatCompletionRequest {
 	chatCompletionMessages := []openai.ChatCompletionMessage{}
 	for _, m := range messages {
 		chatCompletionMessages = append(chatCompletionMessages, openai.ChatCompletionMessage{
@@ -17,19 +17,23 @@ func CreateChatCompletionRequest(messages []Message, maxTokens int) openai.ChatC
 		})
 	}
 
-	return openai.ChatCompletionRequest{
-		Model:     openai.GPT3Dot5Turbo,
+	return &openai.ChatCompletionRequest{
+		Model:     openai.GPT4,
+		MaxTokens: 256,
 		Messages:  chatCompletionMessages,
-		MaxTokens: maxTokens,
+		Stream:    true,
 	}
 }
 
-// CreateChatCompletion submits a Chat Completion API request and returns the
-// response.
-func CreateChatCompletion(messages []Message, maxTokens int) (string, error) {
+// CreateChatCompletion accepts a slice of Message and returns the response
+// of the Large Language Model.
+func CreateChatCompletion(messages []Message) (string, error) {
 	client := openai.NewClient(config.OpenAI.APIKey)
-	req := CreateChatCompletionRequest(messages, maxTokens)
-	resp, err := client.CreateChatCompletion(context.Background(), req)
+	resp, err := client.CreateChatCompletion(
+		context.Background(),
+		*CreateChatCompletionRequest(messages),
+	)
+
 	if err != nil {
 		return "", err
 	}
@@ -37,15 +41,18 @@ func CreateChatCompletion(messages []Message, maxTokens int) (string, error) {
 	return resp.Choices[0].Message.Content, nil
 }
 
-// CreateChatCompletionStream submits a streaming Chat Completion API request
-// and streams the received response to the provided output channel.
-func CreateChatCompletionStream(messages []Message, maxTokens int, output chan string) error {
+// 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 {
 	client := openai.NewClient(config.OpenAI.APIKey)
-	req := CreateChatCompletionRequest(messages, maxTokens)
+	ctx := context.Background()
+
+	req := CreateChatCompletionRequest(messages)
+	req.Stream = true
 
 	defer close(output)
 
-	stream, err := client.CreateChatCompletionStream(context.Background(), req)
+	stream, err := client.CreateChatCompletionStream(ctx, *req)
 	if err != nil {
 		return err
 	}

pkg/cli/tty.go

@@ -45,9 +45,7 @@ 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
 		}