lmcli/pkg/cli/tty.go
Matt Low 6249fbc8f8 Refactor streamed response handling
Update CreateChangeCompletionStream to return the entire response upon
stream completion. Renamed HandleDelayedResponse to
HandleDelayedContent, which no longer returns the content.

Removes the need wrapping HandleDelayedContent in an immediately invoked
function and the passing of the completed response over a channel. Also
allows us to better handle the case of partial a response.
2023-11-24 03:45:43 +00:00

114 lines
2.9 KiB
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 "system":
roleStyle = color.Style{color.HiRed}
case "user":
roleStyle = color.Style{color.HiGreen}
case "assistant":
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()
}
}