lmcli/pkg/cmd/cmd.go

package cmd

import (
	"strings"

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

var (
	systemPromptFile string
)

func RootCmd(ctx *lmcli.Context) *cobra.Command {
	var root = &cobra.Command{
		Use:           "lmcli <command> [flags]",
		Long:          `lmcli - Large Language Model CLI`,
		SilenceErrors: true,
		SilenceUsage:  true,
		Run: func(cmd *cobra.Command, args []string) {
			cmd.Usage()
		},
	}

	chatCmd := ChatCmd(ctx)
	continueCmd := ContinueCmd(ctx)
	cloneCmd := CloneCmd(ctx)
	editCmd := EditCmd(ctx)
	listCmd := ListCmd(ctx)
	newCmd := NewCmd(ctx)
	promptCmd := PromptCmd(ctx)
	renameCmd := RenameCmd(ctx)
	replyCmd := ReplyCmd(ctx)
	retryCmd := RetryCmd(ctx)
	rmCmd := RemoveCmd(ctx)
	viewCmd := ViewCmd(ctx)

	inputCmds := []*cobra.Command{newCmd, promptCmd, replyCmd, retryCmd, continueCmd, editCmd}
	for _, cmd := range inputCmds {
		cmd.Flags().StringVar(ctx.Config.Defaults.Model, "model", *ctx.Config.Defaults.Model, "Which model to use")
		cmd.RegisterFlagCompletionFunc("model", func(*cobra.Command, []string, string) ([]string, cobra.ShellCompDirective) {
			return ctx.GetModels(), cobra.ShellCompDirectiveDefault
		})
		cmd.Flags().IntVar(ctx.Config.Defaults.MaxTokens, "length", *ctx.Config.Defaults.MaxTokens, "Maximum response tokens")
		cmd.Flags().StringVar(ctx.Config.Defaults.SystemPrompt, "system-prompt", *ctx.Config.Defaults.SystemPrompt, "System prompt")
		cmd.Flags().StringVar(&systemPromptFile, "system-prompt-file", "", "A path to a file containing the system prompt")
		cmd.MarkFlagsMutuallyExclusive("system-prompt", "system-prompt-file")
	}

	root.AddCommand(
		chatCmd,
		cloneCmd,
		continueCmd,
		editCmd,
		listCmd,
		newCmd,
		promptCmd,
		renameCmd,
		replyCmd,
		retryCmd,
		rmCmd,
		viewCmd,
	)

	return root
}

func getSystemPrompt(ctx *lmcli.Context) string {
	if systemPromptFile != "" {
		content, err := util.ReadFileContents(systemPromptFile)
		if err != nil {
			lmcli.Fatal("Could not read file contents at %s: %v\n", systemPromptFile, err)
		}
		return content
	}
	return *ctx.Config.Defaults.SystemPrompt
}

// inputFromArgsOrEditor returns either the provided input from the args slice
// (joined with spaces), or if len(args) is 0, opens an editor and returns
// whatever input was provided there. placeholder is a string which populates
// the editor and gets stripped from the final output.
func inputFromArgsOrEditor(args []string, placeholder string, existingMessage string) (message string) {
	var err error
	if len(args) == 0 {
		message, err = util.InputFromEditor(placeholder, "message.*.md", existingMessage)
		if err != nil {
			lmcli.Fatal("Failed to get input: %v\n", err)
		}
	} else {
		message = strings.Join(args, " ")
	}
	message = strings.Trim(message, " \t\n")
	return
}