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/cmd/retry.go

@@ -13,7 +13,7 @@ func RetryCmd(ctx *lmcli.Context) *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "retry <conversation>",
 		Short: "Retry the last user reply in a conversation",
-		Long:  `Re-prompt the conversation up to the last user response. Can be used to regenerate the last assistant reply, or simply generate one if an error occurred.`,
+		Long:  `Prompt the conversation from the last user response.`,
 		Args: func(cmd *cobra.Command, args []string) error {
 			argCount := 1
 			if err := cobra.MinimumNArgs(argCount)(cmd, args); err != nil {
@@ -25,25 +25,36 @@ func RetryCmd(ctx *lmcli.Context) *cobra.Command {
 			shortName := args[0]
 			conversation := cmdutil.LookupConversation(ctx, shortName)
 
-			messages, err := ctx.Store.Messages(conversation)
+			// Load the complete thread from the root message
+			messages, err := ctx.Store.PathToLeaf(conversation.SelectedRoot)
 			if err != nil {
 				return fmt.Errorf("Could not retrieve messages for conversation: %s", conversation.Title)
 			}
 
-			// walk backwards through the conversation and delete messages, break
-			// when we find the latest user response
-			for i := len(messages) - 1; i >= 0; i-- {
-				if messages[i].Role == model.MessageRoleUser {
-					break
-				}
-
-				err = ctx.Store.DeleteMessage(&messages[i])
-				if err != nil {
-					lmcli.Warn("Could not delete previous reply: %v\n", err)
-				}
+			offset, _ := cmd.Flags().GetInt("offset")
+			if offset < 0 {
+				offset = -offset
 			}
 
-			cmdutil.HandleConversationReply(ctx, conversation, true)
+			if offset > len(messages)-1 {
+				return fmt.Errorf("Offset %d is before the start of the conversation.", offset)
+			}
+
+			retryFromIdx := len(messages) - 1 - offset
+
+			// decrease retryFromIdx until we hit a user message
+			for retryFromIdx >= 0 && messages[retryFromIdx].Role != model.MessageRoleUser {
+				retryFromIdx--
+			}
+
+			if messages[retryFromIdx].Role != model.MessageRoleUser {
+				return fmt.Errorf("No user messages to retry")
+			}
+
+			fmt.Printf("Idx: %d Message: %v\n", retryFromIdx, messages[retryFromIdx])
+
+			// Start a new branch at the last user message
+			cmdutil.HandleReply(ctx, &messages[retryFromIdx], true)
 			return nil
 		},
 		ValidArgsFunction: func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
@@ -55,6 +66,8 @@ func RetryCmd(ctx *lmcli.Context) *cobra.Command {
 		},
 	}
 
+	cmd.Flags().Int("offset", 1, "Offset from the last message retry from.")
+
 	applyPromptFlags(ctx, cmd)
 	return cmd
 }