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 RetryCmd(ctx *lmcli.Context) *cobra.Command { cmd := &cobra.Command{ Use: "retry ", Short: "Retry the last user reply in a conversation", 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 { return err } return nil }, RunE: func(cmd *cobra.Command, args []string) error { shortName := args[0] conversation := cmdutil.LookupConversation(ctx, shortName) // 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) } 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) } 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) { compMode := cobra.ShellCompDirectiveNoFileComp if len(args) != 0 { return nil, compMode } return ctx.Store.ConversationShortNameCompletions(toComplete), compMode }, } cmd.Flags().Int("offset", 0, "Offset from the last message to retry from.") applyPromptFlags(ctx, cmd) return cmd }