lmcli/pkg/tui/views/chat/update.go

232 lines
6.4 KiB
Go
Raw Normal View History

2024-06-02 16:40:46 -06:00
package chat
import (
"fmt"
"strings"
"time"
models "git.mlow.ca/mlow/lmcli/pkg/lmcli/model"
"git.mlow.ca/mlow/lmcli/pkg/tui/shared"
tuiutil "git.mlow.ca/mlow/lmcli/pkg/tui/util"
"github.com/charmbracelet/bubbles/cursor"
tea "github.com/charmbracelet/bubbletea"
)
func (m *Model) HandleResize(width, height int) {
m.Width, m.Height = width, height
m.content.Width = width
m.input.SetWidth(width - m.input.FocusedStyle.Base.GetHorizontalFrameSize())
if len(m.messages) > 0 {
m.rebuildMessageCache()
m.updateContent()
}
}
func (m *Model) waitForReply() tea.Cmd {
return func() tea.Msg {
return msgAssistantReply(<-m.replyChan)
}
}
func (m *Model) waitForChunk() tea.Cmd {
return func() tea.Msg {
return msgResponseChunk(<-m.replyChunkChan)
}
}
func (m Model) Update(msg tea.Msg) (Model, tea.Cmd) {
var cmds []tea.Cmd
switch msg := msg.(type) {
case shared.MsgViewEnter:
// wake up spinners and cursors
cmds = append(cmds, cursor.Blink, m.spinner.Tick)
if m.State.Values.ConvShortname != "" && m.conversation.ShortName.String != m.State.Values.ConvShortname {
cmds = append(cmds, m.loadConversation(m.State.Values.ConvShortname))
}
m.rebuildMessageCache()
m.updateContent()
case tea.WindowSizeMsg:
m.HandleResize(msg.Width, msg.Height)
case tuiutil.MsgTempfileEditorClosed:
contents := string(msg)
switch m.editorTarget {
case input:
m.input.SetValue(contents)
case selectedMessage:
m.setMessageContents(m.selectedMessage, contents)
if m.persistence && m.messages[m.selectedMessage].ID > 0 {
// update persisted message
err := m.State.Ctx.Store.UpdateMessage(&m.messages[m.selectedMessage])
if err != nil {
cmds = append(cmds, shared.WrapError(fmt.Errorf("Could not save edited message: %v", err)))
}
}
m.updateContent()
}
case msgConversationLoaded:
m.conversation = (*models.Conversation)(msg)
m.rootMessages, _ = m.State.Ctx.Store.RootMessages(m.conversation.ID)
cmds = append(cmds, m.loadMessages(m.conversation))
case msgMessagesLoaded:
m.selectedMessage = len(msg) - 1
m.messages = msg
m.rebuildMessageCache()
m.updateContent()
m.content.GotoBottom()
case msgResponseChunk:
cmds = append(cmds, m.waitForChunk()) // wait for the next chunk
chunk := string(msg)
if chunk == "" {
break
}
last := len(m.messages) - 1
if last >= 0 && m.messages[last].Role.IsAssistant() {
// append chunk to existing message
m.setMessageContents(last, m.messages[last].Content+chunk)
} else {
// use chunk in new message
m.addMessage(models.Message{
Role: models.MessageRoleAssistant,
Content: chunk,
})
}
m.updateContent()
// show cursor and reset blink interval (simulate typing)
m.replyCursor.Blink = false
cmds = append(cmds, m.replyCursor.BlinkCmd())
m.tokenCount++
m.elapsed = time.Now().Sub(m.startTime)
case msgAssistantReply:
cmds = append(cmds, m.waitForReply()) // wait for the next reply
reply := models.Message(msg)
reply.Content = strings.TrimSpace(reply.Content)
last := len(m.messages) - 1
if last < 0 {
panic("Unexpected empty messages handling msgAssistantReply")
}
if reply.Role.IsAssistant() && m.messages[last].Role.IsAssistant() {
// this was a continuation, so replace the previous message with the completed reply
m.setMessage(last, reply)
} else {
m.addMessage(reply)
}
if m.persistence {
err := m.persistConversation()
if err != nil {
cmds = append(cmds, shared.WrapError(err))
}
}
if m.conversation.Title == "" {
cmds = append(cmds, m.generateConversationTitle())
}
m.updateContent()
case msgResponseEnd:
m.waitingForReply = false
last := len(m.messages) - 1
if last < 0 {
panic("Unexpected empty messages handling msgResponseEnd")
}
m.setMessageContents(last, strings.TrimSpace(m.messages[last].Content))
m.updateContent()
m.status = "Press ctrl+s to send"
case msgResponseError:
m.waitingForReply = false
m.status = "Press ctrl+s to send"
m.State.Err = error(msg)
m.updateContent()
case msgConversationTitleChanged:
title := string(msg)
m.conversation.Title = title
if m.persistence {
err := m.State.Ctx.Store.UpdateConversation(m.conversation)
if err != nil {
cmds = append(cmds, shared.WrapError(err))
}
}
case cursor.BlinkMsg:
if m.waitingForReply {
// ensure we show updated "wait for response" cursor blink state
m.updateContent()
}
}
var cmd tea.Cmd
m.spinner, cmd = m.spinner.Update(msg)
if cmd != nil {
cmds = append(cmds, cmd)
}
m.replyCursor, cmd = m.replyCursor.Update(msg)
if cmd != nil {
cmds = append(cmds, cmd)
}
prevInputLineCnt := m.input.LineCount()
inputCaptured := false
m.input, cmd = m.input.Update(msg)
if cmd != nil {
inputCaptured = true
cmds = append(cmds, cmd)
}
if !inputCaptured {
m.content, cmd = m.content.Update(msg)
if cmd != nil {
cmds = append(cmds, cmd)
}
}
// update views once window dimensions are known
if m.Width > 0 {
m.Header = m.headerView()
m.Footer = m.footerView()
m.Error = tuiutil.ErrorBanner(m.Err, m.Width)
fixedHeight := tuiutil.Height(m.Header) + tuiutil.Height(m.Error) + tuiutil.Height(m.Footer)
// calculate clamped input height to accomodate input text
// minimum 4 lines, maximum half of content area
newHeight := max(4, min((m.Height-fixedHeight-1)/2, m.input.LineCount()))
m.input.SetHeight(newHeight)
m.Input = m.input.View()
// remaining height towards content
m.content.Height = m.Height - fixedHeight - tuiutil.Height(m.Input)
m.Content = m.content.View()
}
// this is a pretty nasty hack to ensure the input area viewport doesn't
// scroll below its content, which can happen when the input viewport
// height has grown, or previously entered lines have been deleted
if prevInputLineCnt != m.input.LineCount() {
// dist is the distance we'd need to scroll up from the current cursor
// position to position the last input line at the bottom of the
// viewport. if negative, we're already scrolled above the bottom
dist := m.input.Line() - (m.input.LineCount() - m.input.Height())
if dist > 0 {
for i := 0; i < dist; i++ {
// move cursor up until content reaches the bottom of the viewport
m.input.CursorUp()
}
m.input, cmd = m.input.Update(nil)
for i := 0; i < dist; i++ {
// move cursor back down to its previous position
m.input.CursorDown()
}
m.input, cmd = m.input.Update(nil)
}
}
return m, tea.Batch(cmds...)
}