232 lines
6.4 KiB
Go
232 lines
6.4 KiB
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...)
|
|
}
|