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

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