package cmd import ( "fmt" "git.mlow.ca/mlow/lmcli/pkg/api" cmdutil "git.mlow.ca/mlow/lmcli/pkg/cmd/util" "git.mlow.ca/mlow/lmcli/pkg/lmcli" "github.com/spf13/cobra" ) func EditCmd(ctx *lmcli.Context) *cobra.Command { cmd := &cobra.Command{ Use: "edit ", Short: "Edit the last user reply in a conversation", Args: func(cmd *cobra.Command, args []string) error { argCount := 1 if err := cobra.MinimumNArgs(argCount)(cmd, args); err != nil { return err } return nil }, RunE: func(cmd *cobra.Command, args []string) error { shortName := args[0] conversation := cmdutil.LookupConversation(ctx, shortName) messages, err := ctx.Store.PathToLeaf(conversation.SelectedRoot) if err != nil { return fmt.Errorf("Could not retrieve messages for conversation: %s", conversation.Title) } offset, _ := cmd.Flags().GetInt("offset") if offset < 0 { offset = -offset } if offset > len(messages)-1 { return fmt.Errorf("Offset %d is before the start of the conversation.", offset) } desiredIdx := len(messages) - 1 - offset toEdit := messages[desiredIdx] newContents := inputFromArgsOrEditor(args[1:], "# Save when finished editing\n", toEdit.Content) switch newContents { case toEdit.Content: return fmt.Errorf("No edits were made.") case "": return fmt.Errorf("No message was provided.") } toEdit.Content = newContents role, _ := cmd.Flags().GetString("role") if role != "" { if role != string(api.MessageRoleUser) && role != string(api.MessageRoleAssistant) { return fmt.Errorf("Invalid role specified. Please use 'user' or 'assistant'.") } toEdit.Role = api.MessageRole(role) } // Update the message in-place inplace, _ := cmd.Flags().GetBool("in-place") if inplace { return ctx.Store.UpdateMessage(&toEdit) } // Otherwise, create a branch for the edited message message, _, err := ctx.Store.CloneBranch(toEdit) if err != nil { return err } if desiredIdx > 0 { // update selected reply messages[desiredIdx-1].SelectedReply = message err = ctx.Store.UpdateMessage(&messages[desiredIdx-1]) } else { // update selected root conversation.SelectedRoot = message err = ctx.Store.UpdateConversation(conversation) } return err }, ValidArgsFunction: func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) { compMode := cobra.ShellCompDirectiveNoFileComp if len(args) != 0 { return nil, compMode } return ctx.Store.ConversationShortNameCompletions(toComplete), compMode }, } cmd.Flags().BoolP("in-place", "i", true, "Edit the message in-place, rather than creating a branch") cmd.Flags().Int("offset", 1, "Offset from the last message to edit") cmd.Flags().StringP("role", "r", "", "Change the role of the edited message (user or assistant)") return cmd }