Add message branching

Updated the behaviour of commands:

- `lmcli edit`
  - by default create a new branch/message branch with the edited contents
  - add --in-place to avoid creating a branch
  - no longer delete messages after the edited message
  - only do the edit, don't fetch a new response
- `lmcli retry`
  - create a new branch rather than replacing old messages
  - add --offset to change where to retry from

pkg/tui/chat.go

@@ -307,14 +307,9 @@ func (m chatModel) Update(msg tea.Msg) (chatModel, tea.Cmd) {
 		}
 
 		if m.persistence {
-			var err error
-			if m.conversation.ID == 0 {
-				err = m.ctx.Store.SaveConversation(m.conversation)
-			}
+			err := m.persistConversation()
 			if err != nil {
 				cmds = append(cmds, wrapError(err))
-			} else {
-				cmds = append(cmds, m.persistConversation())
 			}
 		}
 
@@ -341,7 +336,7 @@ func (m chatModel) Update(msg tea.Msg) (chatModel, tea.Cmd) {
 		title := string(msg)
 		m.conversation.Title = title
 		if m.persistence {
-			err := m.ctx.Store.SaveConversation(m.conversation)
+			err := m.ctx.Store.UpdateConversation(m.conversation)
 			if err != nil {
 				cmds = append(cmds, wrapError(err))
 			}
@@ -469,8 +464,8 @@ func (m *chatModel) handleInputKey(msg tea.KeyMsg) (bool, tea.Cmd) {
 		m.input.Blur()
 		return true, nil
 	case "ctrl+s":
-		userInput := strings.TrimSpace(m.input.Value())
-		if strings.TrimSpace(userInput) == "" {
+		input := strings.TrimSpace(m.input.Value())
+		if input == "" {
 			return true, nil
 		}
 
@@ -478,35 +473,19 @@ func (m *chatModel) handleInputKey(msg tea.KeyMsg) (bool, tea.Cmd) {
 			return true, wrapError(fmt.Errorf("Can't reply to a user message"))
 		}
 
-		reply := models.Message{
+		m.addMessage(models.Message{
 			Role:    models.MessageRoleUser,
-			Content: userInput,
-		}
-
-		if m.persistence {
-			var err error
-			if m.conversation.ID == 0 {
-				err = m.ctx.Store.SaveConversation(m.conversation)
-			}
-			if err != nil {
-				return true, wrapError(err)
-			}
-
-			// ensure all messages up to the one we're about to add are persisted
-			cmd := m.persistConversation()
-			if cmd != nil {
-				return true, cmd
-			}
-
-			savedReply, err := m.ctx.Store.AddReply(m.conversation, reply)
-			if err != nil {
-				return true, wrapError(err)
-			}
-			reply = *savedReply
-		}
+			Content: input,
+		})
 
 		m.input.SetValue("")
-		m.addMessage(reply)
+
+		if m.persistence {
+			err := m.persistConversation()
+			if err != nil {
+				return true, wrapError(err)
+			}
+		}
 
 		m.updateContent()
 		m.content.GotoBottom()
@@ -783,7 +762,7 @@ func (m *chatModel) loadConversation(shortname string) tea.Cmd {
 
 func (m *chatModel) loadMessages(c *models.Conversation) tea.Cmd {
 	return func() tea.Msg {
-		messages, err := m.ctx.Store.Messages(c)
+		messages, err := m.ctx.Store.PathToLeaf(c.SelectedRoot)
 		if err != nil {
 			return msgError(fmt.Errorf("Could not load conversation messages: %v\n", err))
 		}
@@ -791,62 +770,48 @@ func (m *chatModel) loadMessages(c *models.Conversation) tea.Cmd {
 	}
 }
 
-func (m *chatModel) 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))
-			}
+func (m *chatModel) persistConversation() error {
+	if m.conversation.ID == 0 {
+		// Start a new conversation with all messages so far
+		c, messages, err := m.ctx.Store.StartConversation(m.messages...)
+		if err != nil {
+			return err
 		}
+		m.conversation = c
+		m.messages = messages
+
+		return nil
 	}
 
-	for i, msg := range m.messages {
-		if msg.ID > 0 {
-			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, we'll handle updating an existing conversation's messages
+	for i := 0; i < len(m.messages); i++ {
+		if m.messages[i].ID > 0 {
+			// message has an ID, update its contents
+			// TODO: check for content/tool equality before updating?
+			err := m.ctx.Store.UpdateMessage(&m.messages[i])
+			if err != nil {
+				return err
 			}
+		} else if i > 0  {
+			// messages is new, so add it as a reply to previous message
+			saved, err := m.ctx.Store.Reply(&m.messages[i-1], m.messages[i])
+			if err != nil {
+				return err
+			}
+			m.messages[i] = saved[0]
 		} else {
-			newMessage, err := m.ctx.Store.AddReply(m.conversation, msg)
-			if err != nil {
-				return wrapError(err)
-			}
-			m.setMessage(i, *newMessage)
+			// message has no id and no previous messages to add it to
+			// this shouldn't happen?
+			return fmt.Errorf("Error: no messages to reply to")
 		}
 	}
+
 	return nil
 }
 
 func (m *chatModel) generateConversationTitle() tea.Cmd {
 	return func() tea.Msg {
-		title, err := cmdutil.GenerateTitle(m.ctx, m.conversation)
+		title, err := cmdutil.GenerateTitle(m.ctx, m.messages)
 		if err != nil {
 			return msgError(err)
 		}

pkg/tui/conversations.go

@@ -2,7 +2,6 @@ package tui
 
 import (
 	"fmt"
-	"slices"
 	"strings"
 	"time"
 
@@ -145,25 +144,17 @@ func (m conversationsModel) Update(msg tea.Msg) (conversationsModel, tea.Cmd) {
 
 func (m *conversationsModel) loadConversations() tea.Cmd {
 	return func() tea.Msg {
-		conversations, err := m.ctx.Store.Conversations()
+		messages, err := m.ctx.Store.LatestConversationMessages()
 		if err != nil {
 			return msgError(fmt.Errorf("Could not load conversations: %v", err))
 		}
 
-		loaded := make([]loadedConversation, len(conversations))
-		for i, c := range conversations {
-			lastMessage, err := m.ctx.Store.LastMessage(&c)
-			if err != nil {
-				return msgError(err)
-			}
-			loaded[i].conv = c
-			loaded[i].lastReply = *lastMessage
+		loaded := make([]loadedConversation, len(messages))
+		for i, m := range messages {
+			loaded[i].lastReply = m
+			loaded[i].conv = m.Conversation
 		}
 
-		slices.SortFunc(loaded, func(a, b loadedConversation) int {
-			return b.lastReply.CreatedAt.Compare(a.lastReply.CreatedAt)
-		})
-
 		return msgConversationsLoaded(loaded)
 	}
 }