lmcli/pkg/cmd/edit.go

package cmd

import (
	"fmt"

	cmdutil "git.mlow.ca/mlow/lmcli/pkg/cmd/util"
	"git.mlow.ca/mlow/lmcli/pkg/lmcli"
	"git.mlow.ca/mlow/lmcli/pkg/lmcli/model"
	"github.com/spf13/cobra"
)

func EditCmd(ctx *lmcli.Context) *cobra.Command {
	cmd := &cobra.Command{
		Use:   "edit <conversation>",
		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(model.MessageRoleUser) && role != string(model.MessageRoleAssistant) {
					return fmt.Errorf("Invalid role specified. Please use 'user' or 'assistant'.")
				}
				toEdit.Role = model.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
}