lmcli/pkg/cli/tty.go

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 {
			// notify wait animation that we've received data
			waitSignal <- ""
			// wait for signal that wait animation has completed
			<-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")
	}
}