lmcli/pkg/cli/tty.go

package cli

import (
	"fmt"
	"os"
	"time"

	"github.com/alecthomas/chroma/v2/quick"
	"github.com/gookit/color"
)

// 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)
		}
	}
}

// HandleDelayedContent displays a waiting animation to stdout while waiting
// for content to be received on the provided channel. As soon as any (possibly
// chunked) content is received on the channel, the waiting animation is
// replaced by the content.
// Blocks until the channel is closed.
func HandleDelayedContent(content <-chan string) {
	waitSignal := make(chan any)
	go ShowWaitAnimation(waitSignal)

	firstChunk := true
	for chunk := range content {
		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)
	}
}

// RenderConversation renders the given messages to TTY, with optional space
// for a subsequent message. spaceForResponse controls how many '\n' characters
// are printed immediately after the final message (1 if false, 2 if true)
func RenderConversation(messages []Message, spaceForResponse bool) {
	l := len(messages)
	for i, message := range messages {
		message.RenderTTY()
		if i < l-1 || spaceForResponse {
			// print an additional space before the next message
			fmt.Println()
		}
	}
}

// HighlightMarkdown applies syntax highlighting to the provided markdown text
// and writes it to stdout.
func HighlightMarkdown(markdownText string) error {
	return quick.Highlight(os.Stdout, markdownText, "md", *config.Chroma.Formatter, *config.Chroma.Style)
}

func (m *Message) RenderTTY() {
	var messageAge string
	if m.CreatedAt.IsZero() {
		messageAge = "now"
	} else {
		now := time.Now()
		messageAge = humanTimeElapsedSince(now.Sub(m.CreatedAt))
	}

	var roleStyle color.Style
	switch m.Role {
	case MessageRoleSystem:
		roleStyle = color.Style{color.HiRed}
	case MessageRoleUser:
		roleStyle = color.Style{color.HiGreen}
	case MessageRoleAssistant:
		roleStyle = color.Style{color.HiBlue}
	default:
		roleStyle = color.Style{color.FgWhite}
	}
	roleStyle.Add(color.Bold)

	headerColor := color.FgYellow
	separator := headerColor.Sprint("===")
	timestamp := headerColor.Sprint(messageAge)
	role := roleStyle.Sprint(m.FriendlyRole())

	fmt.Printf("%s %s - %s %s\n\n", separator, role, timestamp, separator)
	if m.OriginalContent != "" {
		HighlightMarkdown(m.OriginalContent)
		fmt.Println()
	}
}