package tui // The terminal UI for lmcli, launched from the `lmcli chat` command // TODO: // - binding to open selected message/input in $EDITOR import ( "context" "fmt" "strings" "git.mlow.ca/mlow/lmcli/pkg/lmcli" models "git.mlow.ca/mlow/lmcli/pkg/lmcli/model" "git.mlow.ca/mlow/lmcli/pkg/lmcli/tools" "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 // 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 ( contentStyle = lipgloss.NewStyle().PaddingLeft(2) userStyle = lipgloss.NewStyle().Bold(true).Foreground(lipgloss.Color("10")) assistantStyle = lipgloss.NewStyle().Bold(true).Foreground(lipgloss.Color("12")) headerStyle = lipgloss.NewStyle(). PaddingLeft(1). Background(lipgloss.Color("0")) 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: c := (*models.Conversation)(msg) cmd = m.loadMessages(c) case msgMessagesLoaded: m.messages = []models.Message(msg) m.updateContent() case msgResponseChunk: chunk := string(msg) last := len(m.messages) - 1 if last >= 0 && m.messages[last].Role == models.MessageRoleAssistant { m.messages[last].Content += chunk } else { m.messages = append(m.messages, 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.messages[last] = reply } else if reply.Role != models.MessageRoleAssistant { m.messages = append(m.messages, 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.messages = append(m.messages, 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) } var toolBag []models.Tool for _, toolName := range *m.ctx.Config.Tools.EnabledTools { tool, ok := tools.AvailableTools[toolName] if ok { toolBag = append(toolBag, tool) } } requestParams := models.RequestParameters{ Model: *m.ctx.Config.Defaults.Model, MaxTokens: *m.ctx.Config.Defaults.MaxTokens, Temperature: *m.ctx.Config.Defaults.Temperature, ToolBag: toolBag, } 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) updateContent() { sb := strings.Builder{} msgCnt := len(m.messages) for i, message := range m.messages { var style lipgloss.Style if message.Role == models.MessageRoleUser { style = userStyle } else { style = assistantStyle } sb.WriteString(fmt.Sprintf("%s:\n\n", style.Render(string(message.Role.FriendlyRole())))) highlighted, _ := m.ctx.Chroma.HighlightS(message.Content) sb.WriteString(contentStyle.Width(m.content.Width - 5).Render(highlighted)) if i < msgCnt-1 { sb.WriteString("\n\n") } } atBottom := m.content.AtBottom() m.content.SetContent(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 }