Matt Low
94508b1dbf
Syntax highlighting is fairly expensive, and this means we no longer need to do syntax highlighting on the entire conversaion each time a new message chunk is received
449 lines
11 KiB
Go
449 lines
11 KiB
Go
package tui
|
|
|
|
// The terminal UI for lmcli, launched from the `lmcli chat` command
|
|
// TODO:
|
|
// - binding to open selected message/input in $EDITOR
|
|
// - todo ability to continue or retry previous response
|
|
// - save messages to store
|
|
// - conversation list view
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"strings"
|
|
|
|
"git.mlow.ca/mlow/lmcli/pkg/lmcli"
|
|
models "git.mlow.ca/mlow/lmcli/pkg/lmcli/model"
|
|
"github.com/charmbracelet/bubbles/textarea"
|
|
"github.com/charmbracelet/bubbles/viewport"
|
|
tea "github.com/charmbracelet/bubbletea"
|
|
"github.com/charmbracelet/lipgloss"
|
|
)
|
|
|
|
type focusState int
|
|
|
|
const (
|
|
focusInput focusState = iota
|
|
focusMessages
|
|
)
|
|
|
|
type model struct {
|
|
ctx *lmcli.Context
|
|
convShortname string
|
|
|
|
// application state
|
|
conversation *models.Conversation
|
|
messages []models.Message
|
|
waitingForReply bool
|
|
replyChan chan models.Message
|
|
replyChunkChan chan string
|
|
replyCancelFunc context.CancelFunc
|
|
err error
|
|
|
|
// ui state
|
|
focus focusState
|
|
status string // a general status message
|
|
highlightCache []string // a cache of syntax highlighted message content
|
|
|
|
// ui elements
|
|
content viewport.Model
|
|
input textarea.Model
|
|
}
|
|
|
|
type message struct {
|
|
role string
|
|
content string
|
|
}
|
|
|
|
// custom tea.Msg types
|
|
type (
|
|
// sent on each chunk received from LLM
|
|
msgResponseChunk string
|
|
// sent when response is finished being received
|
|
msgResponseEnd string
|
|
// sent on each completed reply
|
|
msgReply models.Message
|
|
// sent when a conversation is (re)loaded
|
|
msgConversationLoaded *models.Conversation
|
|
// send when a conversation's messages are laoded
|
|
msgMessagesLoaded []models.Message
|
|
// sent when an error occurs
|
|
msgError error
|
|
)
|
|
|
|
// styles
|
|
var (
|
|
userStyle = lipgloss.NewStyle().Faint(true).Bold(true).Foreground(lipgloss.Color("10"))
|
|
assistantStyle = lipgloss.NewStyle().Faint(true).Bold(true).Foreground(lipgloss.Color("12"))
|
|
messageStyle = lipgloss.NewStyle().PaddingLeft(1)
|
|
headerStyle = lipgloss.NewStyle().
|
|
PaddingLeft(1).
|
|
Background(lipgloss.Color("0"))
|
|
contentStyle = lipgloss.NewStyle().
|
|
Padding(1)
|
|
footerStyle = lipgloss.NewStyle().
|
|
Faint(true).
|
|
BorderTop(true).
|
|
BorderStyle(lipgloss.NormalBorder())
|
|
)
|
|
|
|
func (m model) Init() tea.Cmd {
|
|
return tea.Batch(
|
|
textarea.Blink,
|
|
m.loadConversation(m.convShortname),
|
|
m.waitForChunk(),
|
|
m.waitForReply(),
|
|
)
|
|
}
|
|
|
|
func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
|
|
var cmd tea.Cmd
|
|
|
|
switch msg := msg.(type) {
|
|
case tea.KeyMsg:
|
|
switch msg.String() {
|
|
case "ctrl+c":
|
|
if m.waitingForReply {
|
|
m.replyCancelFunc()
|
|
} else {
|
|
return m, tea.Quit
|
|
}
|
|
case "q":
|
|
if m.focus != focusInput {
|
|
return m, tea.Quit
|
|
}
|
|
default:
|
|
var inputHandled tea.Cmd
|
|
switch m.focus {
|
|
case focusInput:
|
|
inputHandled = m.handleInputKey(msg)
|
|
case focusMessages:
|
|
inputHandled = m.handleMessagesKey(msg)
|
|
}
|
|
if inputHandled != nil {
|
|
return m, inputHandled
|
|
}
|
|
}
|
|
case tea.WindowSizeMsg:
|
|
m.content.Width = msg.Width
|
|
m.content.Height = msg.Height - m.input.Height() - lipgloss.Height(m.footerView()) - lipgloss.Height(m.headerView())
|
|
m.input.SetWidth(msg.Width - 1)
|
|
m.updateContent()
|
|
case msgConversationLoaded:
|
|
m.conversation = (*models.Conversation)(msg)
|
|
cmd = m.loadMessages(m.conversation)
|
|
case msgMessagesLoaded:
|
|
m.setMessages(msg)
|
|
m.updateContent()
|
|
case msgResponseChunk:
|
|
chunk := string(msg)
|
|
last := len(m.messages) - 1
|
|
if last >= 0 && m.messages[last].Role == models.MessageRoleAssistant {
|
|
m.setMessageContents(last, m.messages[last].Content+chunk)
|
|
} else {
|
|
m.addMessage(models.Message{
|
|
Role: models.MessageRoleAssistant,
|
|
Content: chunk,
|
|
})
|
|
}
|
|
m.updateContent()
|
|
cmd = m.waitForChunk() // wait for the next chunk
|
|
case msgReply:
|
|
// the last reply that was being worked on is finished
|
|
reply := models.Message(msg)
|
|
last := len(m.messages) - 1
|
|
if last < 0 {
|
|
panic("Unexpected messages length handling msgReply")
|
|
}
|
|
if reply.Role == models.MessageRoleToolCall && m.messages[last].Role == models.MessageRoleAssistant {
|
|
m.setMessage(last, reply)
|
|
} else if reply.Role != models.MessageRoleAssistant {
|
|
m.addMessage(reply)
|
|
}
|
|
m.updateContent()
|
|
cmd = m.waitForReply()
|
|
case msgResponseEnd:
|
|
m.replyCancelFunc = nil
|
|
m.waitingForReply = false
|
|
m.status = "Press ctrl+s to send"
|
|
}
|
|
|
|
if cmd != nil {
|
|
return m, cmd
|
|
}
|
|
|
|
m.input, cmd = m.input.Update(msg)
|
|
if cmd != nil {
|
|
return m, cmd
|
|
}
|
|
|
|
m.content, cmd = m.content.Update(msg)
|
|
if cmd != nil {
|
|
return m, cmd
|
|
}
|
|
|
|
return m, cmd
|
|
}
|
|
|
|
func (m model) View() string {
|
|
return lipgloss.JoinVertical(
|
|
lipgloss.Left,
|
|
m.headerView(),
|
|
m.content.View(),
|
|
m.inputView(),
|
|
m.footerView(),
|
|
)
|
|
}
|
|
|
|
func initialModel(ctx *lmcli.Context, convShortname string) model {
|
|
m := model{
|
|
ctx: ctx,
|
|
convShortname: convShortname,
|
|
|
|
replyChan: make(chan models.Message),
|
|
replyChunkChan: make(chan string),
|
|
}
|
|
|
|
m.content = viewport.New(0, 0)
|
|
|
|
m.input = textarea.New()
|
|
m.input.Placeholder = "Enter a message"
|
|
|
|
m.input.FocusedStyle.CursorLine = lipgloss.NewStyle()
|
|
m.input.ShowLineNumbers = false
|
|
m.input.Focus()
|
|
|
|
m.updateContent()
|
|
|
|
m.waitingForReply = false
|
|
m.status = "Press ctrl+s to send"
|
|
return m
|
|
}
|
|
|
|
func (m *model) handleMessagesKey(msg tea.KeyMsg) tea.Cmd {
|
|
switch msg.String() {
|
|
case "tab":
|
|
m.focus = focusInput
|
|
m.input.Focus()
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (m *model) handleInputKey(msg tea.KeyMsg) tea.Cmd {
|
|
switch msg.String() {
|
|
case "esc":
|
|
m.focus = focusMessages
|
|
m.input.Blur()
|
|
case "ctrl+s":
|
|
userInput := strings.TrimSpace(m.input.Value())
|
|
if strings.TrimSpace(userInput) == "" {
|
|
return nil
|
|
}
|
|
m.input.SetValue("")
|
|
m.addMessage(models.Message{
|
|
Role: models.MessageRoleUser,
|
|
Content: userInput,
|
|
})
|
|
|
|
m.updateContent()
|
|
m.content.GotoBottom()
|
|
|
|
m.waitingForReply = true
|
|
m.status = "Waiting for response, press ctrl+c to cancel..."
|
|
return m.promptLLM()
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (m *model) loadConversation(shortname string) tea.Cmd {
|
|
return func() tea.Msg {
|
|
if shortname == "" {
|
|
return nil
|
|
}
|
|
c, err := m.ctx.Store.ConversationByShortName(shortname)
|
|
if err != nil {
|
|
return msgError(fmt.Errorf("Could not lookup conversation: %v\n", err))
|
|
}
|
|
if c.ID == 0 {
|
|
return msgError(fmt.Errorf("Conversation not found with short name: %s\n", shortname))
|
|
}
|
|
return msgConversationLoaded(c)
|
|
}
|
|
}
|
|
|
|
func (m *model) loadMessages(c *models.Conversation) tea.Cmd {
|
|
return func() tea.Msg {
|
|
messages, err := m.ctx.Store.Messages(c)
|
|
if err != nil {
|
|
return msgError(fmt.Errorf("Could not load conversation messages: %v\n", err))
|
|
}
|
|
return msgMessagesLoaded(messages)
|
|
}
|
|
}
|
|
|
|
func (m *model) waitForReply() tea.Cmd {
|
|
return func() tea.Msg {
|
|
return msgReply(<-m.replyChan)
|
|
}
|
|
}
|
|
|
|
func (m *model) waitForChunk() tea.Cmd {
|
|
return func() tea.Msg {
|
|
return msgResponseChunk(<-m.replyChunkChan)
|
|
}
|
|
}
|
|
|
|
func (m *model) promptLLM() tea.Cmd {
|
|
return func() tea.Msg {
|
|
completionProvider, err := m.ctx.GetCompletionProvider(*m.ctx.Config.Defaults.Model)
|
|
if err != nil {
|
|
return msgError(err)
|
|
}
|
|
|
|
requestParams := models.RequestParameters{
|
|
Model: *m.ctx.Config.Defaults.Model,
|
|
MaxTokens: *m.ctx.Config.Defaults.MaxTokens,
|
|
Temperature: *m.ctx.Config.Defaults.Temperature,
|
|
ToolBag: m.ctx.EnabledTools,
|
|
}
|
|
|
|
replyHandler := func(msg models.Message) {
|
|
m.replyChan <- msg
|
|
}
|
|
|
|
ctx, replyCancelFunc := context.WithCancel(context.Background())
|
|
m.replyCancelFunc = replyCancelFunc
|
|
|
|
// TODO: handle error
|
|
resp, _ := completionProvider.CreateChatCompletionStream(
|
|
ctx, requestParams, m.messages, replyHandler, m.replyChunkChan,
|
|
)
|
|
|
|
return msgResponseEnd(resp)
|
|
}
|
|
}
|
|
|
|
func (m *model) setMessages(messages []models.Message) {
|
|
m.messages = messages
|
|
m.highlightCache = make([]string, len(messages))
|
|
for i, msg := range m.messages {
|
|
highlighted, _ := m.ctx.Chroma.HighlightS(msg.Content)
|
|
m.highlightCache[i] = highlighted
|
|
}
|
|
}
|
|
|
|
func (m *model) setMessage(i int, msg models.Message) {
|
|
if i >= len(m.messages) {
|
|
panic("i out of range")
|
|
}
|
|
highlighted, _ := m.ctx.Chroma.HighlightS(msg.Content)
|
|
m.messages[i] = msg
|
|
m.highlightCache[i] = highlighted
|
|
}
|
|
|
|
func (m *model) addMessage(msg models.Message) {
|
|
highlighted, _ := m.ctx.Chroma.HighlightS(msg.Content)
|
|
m.messages = append(m.messages, msg)
|
|
m.highlightCache = append(m.highlightCache, highlighted)
|
|
}
|
|
|
|
func (m *model) setMessageContents(i int, content string) {
|
|
if i >= len(m.messages) {
|
|
panic("i out of range")
|
|
}
|
|
highlighted, _ := m.ctx.Chroma.HighlightS(content)
|
|
m.messages[i].Content = content
|
|
m.highlightCache[i] = highlighted
|
|
}
|
|
|
|
// render the conversation into the main viewport
|
|
func (m *model) updateContent() {
|
|
sb := strings.Builder{}
|
|
msgCnt := len(m.messages)
|
|
for i, message := range m.messages {
|
|
icon := "⚙️"
|
|
friendly := message.Role.FriendlyRole()
|
|
style := lipgloss.NewStyle().Bold(true).Faint(true)
|
|
|
|
switch message.Role {
|
|
case models.MessageRoleUser:
|
|
icon = ""
|
|
style = userStyle
|
|
case models.MessageRoleAssistant:
|
|
icon = ""
|
|
style = assistantStyle
|
|
case models.MessageRoleToolCall:
|
|
case models.MessageRoleToolResult:
|
|
icon = "🔧"
|
|
}
|
|
|
|
// write message heading with space for content
|
|
header := fmt.Sprintf("%s\n\n", style.Render(icon+friendly))
|
|
sb.WriteString(header)
|
|
|
|
// TODO: render something for tool calls/results?
|
|
|
|
// write message contents
|
|
var highlighted string
|
|
if m.highlightCache[i] == "" {
|
|
highlighted = message.Content
|
|
} else {
|
|
highlighted = m.highlightCache[i]
|
|
}
|
|
contents := messageStyle.Width(m.content.Width - 5).Render(highlighted)
|
|
sb.WriteString(contents)
|
|
|
|
if i < msgCnt-1 {
|
|
sb.WriteString("\n\n")
|
|
}
|
|
}
|
|
atBottom := m.content.AtBottom()
|
|
m.content.SetContent(contentStyle.Render(sb.String()))
|
|
if atBottom {
|
|
// if we were at bottom before the update, scroll with the output
|
|
m.content.GotoBottom()
|
|
}
|
|
}
|
|
|
|
func (m *model) headerView() string {
|
|
titleStyle := lipgloss.NewStyle().
|
|
Bold(true)
|
|
var title string
|
|
if m.conversation != nil && m.conversation.Title != "" {
|
|
title = m.conversation.Title
|
|
} else {
|
|
title = "Untitled"
|
|
}
|
|
part := titleStyle.Render(title)
|
|
|
|
return headerStyle.Width(m.content.Width).Render(part)
|
|
}
|
|
|
|
func (m model) inputView() string {
|
|
return m.input.View()
|
|
}
|
|
|
|
func (m *model) footerView() string {
|
|
left := m.status
|
|
right := fmt.Sprintf("Model: %s", *m.ctx.Config.Defaults.Model)
|
|
|
|
totalWidth := lipgloss.Width(left + right)
|
|
var padding string
|
|
if m.content.Width-totalWidth > 0 {
|
|
padding = strings.Repeat(" ", m.content.Width-totalWidth)
|
|
} else {
|
|
padding = ""
|
|
}
|
|
|
|
footer := lipgloss.JoinHorizontal(lipgloss.Center, left, padding, right)
|
|
return footerStyle.Width(m.content.Width).Render(footer)
|
|
}
|
|
|
|
func Launch(ctx *lmcli.Context, convShortname string) error {
|
|
p := tea.NewProgram(initialModel(ctx, convShortname), tea.WithAltScreen())
|
|
if _, err := p.Run(); err != nil {
|
|
return fmt.Errorf("Error running program: %v", err)
|
|
}
|
|
return nil
|
|
}
|