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

@@ -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
 }