From 794ccc52ff2adf46a53d9b24e0219f4dedad26fd Mon Sep 17 00:00:00 2001
From: Matt Low <matt@mlow.ca>
Date: Sun, 5 Nov 2023 06:44:06 +0000
Subject: [PATCH] Show waiting animation while waiting for LLM response

---
 pkg/cli/cmd.go    |  9 ++++++---
 pkg/cli/openai.go | 10 ++++++----
 pkg/cli/tty.go    | 51 +++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 63 insertions(+), 7 deletions(-)
 create mode 100644 pkg/cli/tty.go

diff --git a/pkg/cli/cmd.go b/pkg/cli/cmd.go
index 7a61a56..f81562b 100644
--- a/pkg/cli/cmd.go
+++ b/pkg/cli/cmd.go
@@ -2,7 +2,6 @@ package cli
 
 import (
 	"fmt"
-	"os"
 	"strings"
 
 	"github.com/spf13/cobra"
@@ -89,7 +88,9 @@ var newCmd = &cobra.Command{
 			},
 		}
 
-		err = CreateChatCompletionStream("You are a helpful assistant.", messages, os.Stdout)
+		receiver := make(chan string)
+		go HandleDelayedResponse(receiver)
+		err = CreateChatCompletionStream("You are a helpful assistant.", messages, receiver)
 		if err != nil {
 			Fatal("%v\n", err)
 			return
@@ -117,7 +118,9 @@ var promptCmd = &cobra.Command{
 			},
 		}
 
-		err := CreateChatCompletionStream("You are a helpful assistant.", messages, os.Stdout)
+		receiver := make(chan string)
+		go HandleDelayedResponse(receiver)
+		err := CreateChatCompletionStream("You are a helpful assistant.", messages, receiver)
 		if err != nil {
 			Fatal("%v\n", err)
 			return
diff --git a/pkg/cli/openai.go b/pkg/cli/openai.go
index c54e479..060bf63 100644
--- a/pkg/cli/openai.go
+++ b/pkg/cli/openai.go
@@ -3,7 +3,6 @@ package cli
 import (
 	"context"
 	"errors"
-	"fmt"
 	"io"
 
 	openai "github.com/sashabaranov/go-openai"
@@ -48,13 +47,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(system string, messages []Message, output chan string) error {
 	client := openai.NewClient(config.OpenAI.APIKey)
 	ctx := context.Background()
 
 	req := CreateChatCompletionRequest(system, messages)
 	req.Stream = true
 
+	defer close(output)
+
 	stream, err := client.CreateChatCompletionStream(ctx, *req)
 	if err != nil {
 		return err
@@ -69,10 +72,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
 	}
 }
diff --git a/pkg/cli/tty.go b/pkg/cli/tty.go
new file mode 100644
index 0000000..80bf21f
--- /dev/null
+++ b/pkg/cli/tty.go
@@ -0,0 +1,51 @@
+package cli
+
+import (
+	"fmt"
+	"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
+// content received on the response channel to stdout. Blocks until the channel
+// is closed.
+func HandleDelayedResponse(response chan string) {
+	waitSignal := make(chan any)
+	go ShowWaitAnimation(waitSignal)
+
+	firstChunk := true
+	for chunk := range response {
+		if firstChunk {
+			waitSignal <- ""
+			<-waitSignal
+			firstChunk = false
+		}
+		fmt.Print(chunk)
+	}
+}