From 51ce74ad3abc8533d88963ba161f995694d5abc3 Mon Sep 17 00:00:00 2001
From: Matt Low <matt@mlow.ca>
Date: Tue, 9 Jan 2024 18:10:05 +0000
Subject: [PATCH] Add --offset flag to edit command

---
 pkg/cli/cmd.go | 34 ++++++++++++++++++++--------------
 1 file changed, 20 insertions(+), 14 deletions(-)

diff --git a/pkg/cli/cmd.go b/pkg/cli/cmd.go
index 5174c63..507da72 100644
--- a/pkg/cli/cmd.go
+++ b/pkg/cli/cmd.go
@@ -34,6 +34,7 @@ func init() {
 
 	listCmd.Flags().Bool("all", false, fmt.Sprintf("Show all conversations, by default only the last %d are shown", LS_LIMIT))
 	renameCmd.Flags().Bool("generate", false, "Generate a conversation title")
+	editCmd.Flags().Int("offset", 1, "Offset from the last reply to edit (Default: edit your last message, assuming there's an assistant reply)")
 
 	rootCmd.AddCommand(
 		cloneCmd,
@@ -659,34 +660,39 @@ var editCmd = &cobra.Command{
 			Fatal("Could not retrieve messages for conversation: %s\n", conversation.Title)
 		}
 
+		offset, _ := cmd.Flags().GetInt("offset")
+		if offset < 0 {
+			offset = -offset
+		}
+
+		if offset > len(messages) - 1   {
+			Fatal("Offset %d is before the start of the conversation\n", offset)
+		}
+
+		desiredIdx := len(messages) - 1 - offset
+
 		// walk backwards through the conversation deleting messages until and
 		// including the last user message
 		toRemove := []Message{}
-		var lastUserMessage *Message
+		var toEdit *Message
 		for i := len(messages) - 1; i >= 0; i-- {
-			if messages[i].Role == MessageRoleUser {
-				lastUserMessage = &messages[i]
+			if i == desiredIdx {
+				toEdit = &messages[i]
 			}
-
 			toRemove = append(toRemove, messages[i])
 			messages = messages[:i]
-			if lastUserMessage != nil {
+			if toEdit != nil {
 				break
 			}
 		}
 
-		if lastUserMessage == nil {
-			Fatal("No messages left in the conversation, nothing to edit.\n")
-		}
-
-		existingContents := lastUserMessage.OriginalContent
+		existingContents := toEdit.OriginalContent
 
 		newContents := inputFromArgsOrEditor(args[1:], "# Save when finished editing\n", existingContents)
-		if newContents == existingContents {
+		switch newContents {
+		case existingContents:
 			Fatal("No edits were made.\n")
-		}
-
-		if newContents == "" {
+		case "":
 			Fatal("No message was provided.\n")
 		}