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
This commit is contained in:
Matt Low 2024-03-17 18:16:10 +00:00
parent 6f737ad19c
commit 6c64f21d9a

View File

@ -2,7 +2,6 @@ package tui
// The terminal UI for lmcli, launched from the `lmcli chat` command // The terminal UI for lmcli, launched from the `lmcli chat` command
// TODO: // TODO:
// - ability to continue an incomplete or missing assistant response
// - conversation list view // - conversation list view
// - change model // - change model
// - rename conversation // - rename conversation
@ -152,7 +151,7 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
switch msg.String() { switch msg.String() {
case "ctrl+c": case "ctrl+c":
if m.waitingForReply { if m.waitingForReply {
m.stopSignal <- "stahp!" m.stopSignal <- ""
} else { } else {
return m, tea.Quit return m, tea.Quit
} }
@ -208,11 +207,16 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
reply := models.Message(msg) reply := models.Message(msg)
last := len(m.messages) - 1 last := len(m.messages) - 1
if last < 0 { 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.setMessageContents(last, strings.TrimSpace(m.messages[last].Content))
m.setMessage(last, reply) if m.messages[last].Role == models.MessageRoleAssistant {
} else if reply.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) m.addMessage(reply)
} }
@ -224,7 +228,7 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
if err != nil { if err != nil {
cmds = append(cmds, wrapError(err)) cmds = append(cmds, wrapError(err))
} else { } 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()) cmds = append(cmds, m.waitForReply())
case msgResponseEnd: case msgResponseEnd:
m.waitingForReply = false 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" m.status = "Press ctrl+s to send"
case msgResponseError: case msgResponseError:
m.waitingForReply = false m.waitingForReply = false
@ -485,6 +495,17 @@ func (m *model) handleMessagesKey(msg tea.KeyMsg) tea.Cmd {
offset := m.messageOffsets[m.selectedMessage] offset := m.messageOffsets[m.selectedMessage]
scrollIntoView(&m.content, offset, 0.1) 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 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 // ensure all messages up to the one we're about to add are
// persistent // persistent
cmd := m.persistRecentMessages() cmd := m.persistConversation()
if cmd != nil { if cmd != nil {
return cmd 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") cmd := openTempfileEditor("message.*.md", m.input.Value(), "# Edit your input below\n")
m.editorTarget = input m.editorTarget = input
return cmd 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 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 { for i, msg := range m.messages {
if msg.ID > 0 { 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 return nil
} }