Pull message rendering into its own method

Removed the system parameter on ChatCopmletion functions, and persist it
in conversations as well.

pkg/cli/cli.go

@@ -12,3 +12,7 @@ func Fatal(format string, args ...any) {
 	fmt.Fprintf(os.Stderr, format, args...)
 	os.Exit(1)
 }
+
+func Warn(format string, args ...any) {
+	fmt.Fprintf(os.Stderr, format, args...)
+}

pkg/cli/cmd.go

@@ -2,7 +2,6 @@ package cli
 
 import (
 	"fmt"
-	"os"
 	"strings"
 
 	"github.com/spf13/cobra"
@@ -72,27 +71,65 @@ var newCmd = &cobra.Command{
 		messageContents, err := InputFromEditor("# What would you like to say?\n", "message.*.md")
 		if err != nil {
 			Fatal("Failed to get input: %v\n", err)
-			return
 		}
 
 		if messageContents == "" {
 			Fatal("No message was provided.\n")
-			return
 		}
 
-		fmt.Printf("> %s\n", messageContents)
+		// TODO: set title if --title provided, otherwise defer for later(?)
+		conversation := Conversation{}
+		err = store.SaveConversation(&conversation)
+		if err != nil {
+			Fatal("Could not save new conversation: %v\n", err)
+		}
 
+		const system = "You are a helpful assistant."
 		messages := []Message{
 			{
+				ConversationID:  conversation.ID,
+				Role:            "system",
+				OriginalContent: system,
+			},
+			{
+				ConversationID:  conversation.ID,
 				Role:            "user",
 				OriginalContent: messageContents,
 			},
 		}
+		for _, message := range messages {
+			err = store.SaveMessage(&message)
+			if err != nil {
+				Warn("Could not save %s message: %v\n", message.Role, err)
+			}
+		}
 
-		err = CreateChatCompletionStream("You are a helpful assistant.", messages, os.Stdout)
+		for _, message := range messages {
+			message.RenderTTY(true)
+		}
+
+		reply := Message{
+			ConversationID: conversation.ID,
+			Role:           "assistant",
+		}
+
+		reply.RenderTTY(false)
+
+		receiver := make(chan string)
+		response := make(chan string)
+		go func() {
+			response <- HandleDelayedResponse(receiver)
+		}()
+		err = CreateChatCompletionStream(messages, receiver)
 		if err != nil {
 			Fatal("%v\n", err)
-			return
+		}
+
+		reply.OriginalContent = <-response
+
+		err = store.SaveMessage(&reply)
+		if err != nil {
+			Fatal("Could not save reply: %v\n", err)
 		}
 
 		fmt.Println()
@@ -107,20 +144,25 @@ var promptCmd = &cobra.Command{
 		message := strings.Join(args, " ")
 		if len(strings.Trim(message, " \t\n")) == 0 {
 			Fatal("No message was provided.\n")
-			return
 		}
 
+		const system = "You are a helpful assistant."
 		messages := []Message{
+			{
+				Role:            "system",
+				OriginalContent: system,
+			},
 			{
 				Role:            "user",
 				OriginalContent: message,
 			},
 		}
 
-		err := CreateChatCompletionStream("You are a helpful assistant.", messages, os.Stdout)
+		receiver := make(chan string)
+		go HandleDelayedResponse(receiver)
+		err := CreateChatCompletionStream(messages, receiver)
 		if err != nil {
 			Fatal("%v\n", err)
-			return
 		}
 
 		fmt.Println()
@@ -128,6 +170,9 @@ var promptCmd = &cobra.Command{
 }
 
 func NewRootCmd() *cobra.Command {
-	rootCmd.AddCommand(newCmd, promptCmd)
+	rootCmd.AddCommand(
+		newCmd,
+		promptCmd,
+	)
 	return rootCmd
 }

pkg/cli/config.go

@@ -49,7 +49,7 @@ func InitializeConfig() *Config {
 
 		_, err = file.Write(bytes)
 		if err != nil {
-			Fatal("Could not save default configuratoin: %v", err)
+			Fatal("Could not save default configuration: %v", err)
 			return nil
 		}
 	} else if err != nil {

pkg/cli/openai.go

@@ -3,20 +3,13 @@ package cli
 import (
 	"context"
 	"errors"
-	"fmt"
 	"io"
 
 	openai "github.com/sashabaranov/go-openai"
 )
 
-func CreateChatCompletionRequest(system string, messages []Message) *openai.ChatCompletionRequest {
-	chatCompletionMessages := []openai.ChatCompletionMessage{
-		{
-			Role:    "system",
-			Content: system,
-		},
-	}
-
+func CreateChatCompletionRequest(messages []Message) *openai.ChatCompletionRequest {
+	chatCompletionMessages := []openai.ChatCompletionMessage{}
 	for _, m := range messages {
 		chatCompletionMessages = append(chatCompletionMessages, openai.ChatCompletionMessage{
 			Role:    m.Role,
@@ -34,11 +27,11 @@ func CreateChatCompletionRequest(system string, messages []Message) *openai.Chat
 
 // CreateChatCompletion accepts a slice of Message and returns the response
 // of the Large Language Model.
-func CreateChatCompletion(system string, messages []Message) (string, error) {
+func CreateChatCompletion(messages []Message) (string, error) {
 	client := openai.NewClient(config.OpenAI.APIKey)
 	resp, err := client.CreateChatCompletion(
 		context.Background(),
-		*CreateChatCompletionRequest(system, messages),
+		*CreateChatCompletionRequest(messages),
 	)
 
 	if err != nil {
@@ -48,13 +41,17 @@ func CreateChatCompletion(system string, messages []Message) (string, error) {
 	return resp.Choices[0].Message.Content, nil
 }
 
-func CreateChatCompletionStream(system string, messages []Message, output io.Writer) 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)
 	ctx := context.Background()
 
-	req := CreateChatCompletionRequest(system, messages)
+	req := CreateChatCompletionRequest(messages)
 	req.Stream = true
 
+	defer close(output)
+
 	stream, err := client.CreateChatCompletionStream(ctx, *req)
 	if err != nil {
 		return err
@@ -69,10 +66,9 @@ func CreateChatCompletionStream(system string, messages []Message, output io.Wri
 		}
 
 		if err != nil {
-			//fmt.Printf("\nStream error: %v\n", err)
 			return err
 		}
 
-		fmt.Fprint(output, response.Choices[0].Delta.Content)
+		output <- response.Choices[0].Delta.Content
 	}
 }

pkg/cli/tty.go

@@ -0,0 +1,76 @@
+package cli
+
+import (
+	"fmt"
+	"strings"
+	"time"
+)
+
+// ShowWaitAnimation "draws" an animated ellipses to stdout until something is
+// received on the signal channel. An empty string sent to the channel to
+// noftify the caller that the animation has completed (carriage returned).
+func ShowWaitAnimation(signal chan any) {
+	animationStep := 0
+	for {
+		select {
+		case _ = <-signal:
+			fmt.Print("\r")
+			signal <- ""
+			return
+		default:
+			modSix := animationStep % 6
+			if modSix == 3 || modSix == 0 {
+				fmt.Print("\r")
+			}
+			if modSix < 3 {
+				fmt.Print(".")
+			} else {
+				fmt.Print(" ")
+			}
+			animationStep++
+			time.Sleep(250 * time.Millisecond)
+		}
+	}
+}
+
+// HandledDelayedResponse writes a waiting animation (abusing \r) and the
+// (possibly chunked) content received on the response channel to stdout.
+// Blocks until the channel is closed.
+func HandleDelayedResponse(response chan string) string {
+	waitSignal := make(chan any)
+	go ShowWaitAnimation(waitSignal)
+
+	sb := strings.Builder{}
+
+	firstChunk := true
+	for chunk := range response {
+		if firstChunk {
+			waitSignal <- ""
+			<-waitSignal
+			firstChunk = false
+		}
+		fmt.Print(chunk)
+		sb.WriteString(chunk)
+	}
+
+	return sb.String()
+}
+
+func (m *Message) RenderTTY(paddingDown bool) {
+	var friendlyRole string
+	switch m.Role {
+		case "user":
+			friendlyRole = "You"
+		case "system":
+			friendlyRole = "System"
+		case "assistant":
+			friendlyRole = "Assistant"
+	}
+	fmt.Printf("<%s>\n\n", friendlyRole)
+	if m.OriginalContent != "" {
+		fmt.Print(m.OriginalContent)
+	}
+	if paddingDown {
+		fmt.Print("\n\n")
+	}
+}

pkg/cli/util.go

@@ -46,5 +46,5 @@ func InputFromEditor(placeholder string, pattern string) (string, error) {
 		}
 	}
 
-	return content, nil
+	return strings.Trim(content, "\n \t"), nil
 }