From 1e8ff60c54f1ea554f2f1f852d6099fc7d0b042e Mon Sep 17 00:00:00 2001
From: Matt Low <matt@mlow.ca>
Date: Wed, 29 Nov 2023 15:30:11 +0000
Subject: [PATCH] Add `lmcli rename` to rename conversations

---
 pkg/cli/cmd.go          | 60 ++++++++++++++++++++++++++++++++++++++---
 pkg/cli/conversation.go |  7 +++--
 2 files changed, 59 insertions(+), 8 deletions(-)

diff --git a/pkg/cli/cmd.go b/pkg/cli/cmd.go
index 06f89dc..8312067 100644
--- a/pkg/cli/cmd.go
+++ b/pkg/cli/cmd.go
@@ -26,11 +26,14 @@ func init() {
 		cmd.MarkFlagsMutuallyExclusive("system-prompt", "system-prompt-file")
 	}
 
+	renameCmd.Flags().Bool("generate", false, "Use the LLM to generate the conversation title.")
+
 	rootCmd.AddCommand(
 		continueCmd,
 		lsCmd,
 		newCmd,
 		promptCmd,
+		renameCmd,
 		replyCmd,
 		retryCmd,
 		rmCmd,
@@ -357,8 +360,8 @@ var newCmd = &cobra.Command{
 			Fatal("No message was provided.\n")
 		}
 
-		conversation := Conversation{}
-		err := store.SaveConversation(&conversation)
+		conversation := &Conversation{}
+		err := store.SaveConversation(conversation)
 		if err != nil {
 			Fatal("Could not save new conversation: %v\n", err)
 		}
@@ -387,18 +390,67 @@ var newCmd = &cobra.Command{
 
 		conversation.GenerateAndSaveReplies(messages)
 
-		err = conversation.GenerateTitle()
+		title, err := conversation.GenerateTitle()
 		if err != nil {
 			Warn("Could not generate title for conversation: %v\n", err)
 		}
 
-		err = store.SaveConversation(&conversation)
+		conversation.Title = title
+
+		err = store.SaveConversation(conversation)
 		if err != nil {
 			Warn("Could not save conversation after generating title: %v\n", err)
 		}
 	},
 }
 
+var renameCmd = &cobra.Command{
+	Use:   "rename <conversation> [title]",
+	Short: "Rename a conversation",
+	Long:  `Renames a conversation, either with the provided title or by generating a new name.`,
+	Args: func(cmd *cobra.Command, args []string) error {
+		argCount := 1
+		if err := cobra.MinimumNArgs(argCount)(cmd, args); err != nil {
+			return err
+		}
+		return nil
+	},
+	Run: func(cmd *cobra.Command, args []string) {
+		shortName := args[0]
+		conversation, err := store.ConversationByShortName(shortName)
+		if conversation.ID == 0 {
+			Fatal("Conversation not found with short name: %s\n", shortName)
+		}
+
+		generate, _ := cmd.Flags().GetBool("generate")
+		var title string
+		if generate {
+			title, err = conversation.GenerateTitle()
+			if err != nil {
+				Fatal("Could not generate title for conversation: %v\n", err)
+			}
+		} else {
+			if len(args) < 2 {
+				Fatal("Conversation title not provided.")
+			}
+			title = strings.Join(args[1:], " ")
+		}
+
+		conversation.Title = title
+		err = store.SaveConversation(conversation)
+		if err != nil {
+			Warn("Could not save conversation with new title: %v\n", err)
+		}
+	},
+	ValidArgsFunction: func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+		compMode := cobra.ShellCompDirectiveNoFileComp
+		if len(args) != 0 {
+			return nil, compMode
+		}
+		return store.ConversationShortNameCompletions(toComplete), compMode
+	},
+}
+
 var promptCmd = &cobra.Command{
 	Use:   "prompt [message]",
 	Short: "Do a one-shot prompt",
diff --git a/pkg/cli/conversation.go b/pkg/cli/conversation.go
index 458da44..cac4b73 100644
--- a/pkg/cli/conversation.go
+++ b/pkg/cli/conversation.go
@@ -29,7 +29,7 @@ func (m *Message) FriendlyRole() string {
 	return friendlyRole
 }
 
-func (c *Conversation) GenerateTitle() error {
+func (c *Conversation) GenerateTitle() (string, error) {
 	const header = "Generate a consise 4-5 word title for the conversation below."
 	prompt := fmt.Sprintf("%s\n\n---\n\n%s", header, c.FormatForExternalPrompting(false))
 
@@ -43,11 +43,10 @@ func (c *Conversation) GenerateTitle() error {
 	model := "gpt-3.5-turbo" // use cheap model to generate title
 	response, err := CreateChatCompletion(model, messages, 25, nil)
 	if err != nil {
-		return err
+		return "", err
 	}
 
-	c.Title = response
-	return nil
+	return response, nil
 }
 
 func (c *Conversation) FormatForExternalPrompting(system bool) string {