From 6c64f21d9a807ba1d8790ea76bf485c0c3914d05 Mon Sep 17 00:00:00 2001 From: Matt Low Date: Sun, 17 Mar 2024 18:16:10 +0000 Subject: [PATCH] tui: support for message retry/continue Better handling of persistence, and we now ensure the response we persist is trimmed of whitespace, particularly important when a response is cancelled mid-stream --- pkg/tui/tui.go | 101 ++++++++++++++++++++++++++++++++++++------------- 1 file changed, 75 insertions(+), 26 deletions(-) diff --git a/pkg/tui/tui.go b/pkg/tui/tui.go index 56b043e..2126e92 100644 --- a/pkg/tui/tui.go +++ b/pkg/tui/tui.go @@ -2,7 +2,6 @@ package tui // The terminal UI for lmcli, launched from the `lmcli chat` command // TODO: -// - ability to continue an incomplete or missing assistant response // - conversation list view // - change model // - rename conversation @@ -152,7 +151,7 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { switch msg.String() { case "ctrl+c": if m.waitingForReply { - m.stopSignal <- "stahp!" + m.stopSignal <- "" } else { return m, tea.Quit } @@ -208,11 +207,16 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { reply := models.Message(msg) last := len(m.messages) - 1 if last < 0 { - panic("Unexpected messages length handling msgReply") + panic("Unexpected empty messages handling msgReply") } - if reply.Role == models.MessageRoleToolCall && m.messages[last].Role == models.MessageRoleAssistant { - m.setMessage(last, reply) - } else if reply.Role != models.MessageRoleAssistant { + m.setMessageContents(last, strings.TrimSpace(m.messages[last].Content)) + if m.messages[last].Role == models.MessageRoleAssistant { + // the last message was an assistant message, so this is a continuation + if reply.Role == models.MessageRoleToolCall { + // update last message rrole to tool call + m.messages[last].Role = models.MessageRoleToolCall + } + } else { m.addMessage(reply) } @@ -224,7 +228,7 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { if err != nil { cmds = append(cmds, wrapError(err)) } else { - cmds = append(cmds, m.persistRecentMessages()) + cmds = append(cmds, m.persistConversation()) } } @@ -236,6 +240,12 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { cmds = append(cmds, m.waitForReply()) 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 @@ -485,6 +495,17 @@ func (m *model) handleMessagesKey(msg tea.KeyMsg) tea.Cmd { offset := m.messageOffsets[m.selectedMessage] scrollIntoView(&m.content, offset, 0.1) } + case "ctrl+r": + // resubmit the conversation with all messages up until and including + // the selected message + if len(m.messages) == 0 { + return nil + } + m.messages = m.messages[:m.selectedMessage+1] + m.highlightCache = m.highlightCache[:m.selectedMessage+1] + m.updateContent() + m.content.GotoBottom() + return m.promptLLM() } return nil } @@ -524,7 +545,7 @@ func (m *model) handleInputKey(msg tea.KeyMsg) tea.Cmd { // ensure all messages up to the one we're about to add are // persistent - cmd := m.persistRecentMessages() + cmd := m.persistConversation() if cmd != nil { return cmd } @@ -546,17 +567,6 @@ func (m *model) handleInputKey(msg tea.KeyMsg) tea.Cmd { cmd := openTempfileEditor("message.*.md", m.input.Value(), "# Edit your input below\n") m.editorTarget = input return cmd - case "ctrl+r": - if len(m.messages) == 0 { - return nil - } - // TODO: retry from selected message - if m.messages[len(m.messages)-1].Role != models.MessageRoleUser { - m.messages = m.messages[:len(m.messages)-1] - m.updateContent() - } - m.content.GotoBottom() - return m.promptLLM() } return nil } @@ -653,16 +663,55 @@ func (m *model) promptLLM() tea.Cmd { } } -func (m *model) persistRecentMessages() tea.Cmd { +func (m *model) persistConversation() tea.Cmd { + existingMessages, err := m.ctx.Store.Messages(m.conversation) + if err != nil { + return wrapError(fmt.Errorf("Could not retrieve existing conversation messages while trying to save: %v", err)) + } + + existingById := make(map[uint]*models.Message, len(existingMessages)) + for _, msg := range existingMessages { + existingById[msg.ID] = &msg + } + + currentById := make(map[uint]*models.Message, len(m.messages)) + for _, msg := range m.messages { + currentById[msg.ID] = &msg + } + + for _, msg := range existingMessages { + _, ok := currentById[msg.ID] + if !ok { + err := m.ctx.Store.DeleteMessage(&msg) + if err != nil { + return wrapError(fmt.Errorf("Failed to remove messages: %v", err)) + } + } + } + for i, msg := range m.messages { if msg.ID > 0 { - continue + exist, ok := existingById[msg.ID] + if ok { + if msg.Content == exist.Content { + continue + } + // update message when contents don't match that of store + err := m.ctx.Store.UpdateMessage(&msg) + if err != nil { + return wrapError(err) + } + } else { + // this would be quite odd... and I'm not sure how to handle + // it at the time of writing this + } + } else { + newMessage, err := m.ctx.Store.AddReply(m.conversation, msg) + if err != nil { + return wrapError(err) + } + m.setMessage(i, *newMessage) } - newMessage, err := m.ctx.Store.AddReply(m.conversation, msg) - if err != nil { - return wrapError(err) - } - m.setMessage(i, *newMessage) } return nil }