lmcli/pkg/cli/cmd.go

package cli

import (
	"fmt"
	"os"
	"slices"
	"strings"
	"time"

	"github.com/spf13/cobra"
)

var (
	maxTokens        int
	model            string
	systemPrompt     string
	systemPromptFile string
)

const (
	// Limit number of conversations shown with `ls`, without --all
	LS_LIMIT int = 25
)

func init() {
	inputCmds := []*cobra.Command{newCmd, promptCmd, replyCmd, retryCmd, continueCmd, editCmd}
	for _, cmd := range inputCmds {
		cmd.Flags().IntVar(&maxTokens, "length", *config.OpenAI.DefaultMaxLength, "Maximum response tokens")
		cmd.Flags().StringVar(&model, "model", *config.OpenAI.DefaultModel, "Which model to use model")
		cmd.Flags().StringVar(&systemPrompt, "system-prompt", *config.ModelDefaults.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")
	}

	listCmd.Flags().Bool("all", false, fmt.Sprintf("Show all conversations, by default only the last %d are shown", LS_LIMIT))
	renameCmd.Flags().Bool("generate", false, "Generate a conversation title")
	editCmd.Flags().Int("offset", 1, "Offset from the last reply to edit (Default: edit your last message, assuming there's an assistant reply)")

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

func Execute() error {
	return rootCmd.Execute()
}

func getSystemPrompt() string {
	if systemPromptFile != "" {
		content, err := FileContents(systemPromptFile)
		if err != nil {
			Fatal("Could not read file contents at %s: %v\n", systemPromptFile, err)
		}
		return content
	}
	return systemPrompt
}

// fetchAndShowCompletion prompts the LLM with the given messages and streams
// the response to stdout. Returns all model reply messages.
func fetchAndShowCompletion(messages []Message) ([]Message, error) {
	content := make(chan string) // receives the reponse from LLM
	defer close(content)

	// render all content received over the channel
	go ShowDelayedContent(content)

	var replies []Message
	response, err := CreateChatCompletionStream(model, messages, maxTokens, content, &replies)
	if response != "" {
		// there was some content, so break to a new line after it
		fmt.Println()

		if err != nil {
			Warn("Received partial response. Error: %v\n", err)
			err = nil
		}
	}

	return replies, err
}

// lookupConversation either returns the conversation found by the
// short name or exits the program
func lookupConversation(shortName string) *Conversation {
	c, err := store.ConversationByShortName(shortName)
	if err != nil {
		Fatal("Could not lookup conversation: %v\n", err)
	}
	if c.ID == 0 {
		Fatal("Conversation not found with short name: %s\n", shortName)
	}
	return c
}

func lookupConversationE(shortName string) (*Conversation, error) {
	c, err := store.ConversationByShortName(shortName)
	if err != nil {
		return nil, fmt.Errorf("Could not lookup conversation: %v", err)
	}
	if c.ID == 0 {
		return nil, fmt.Errorf("Conversation not found with short name: %s", shortName)
	}
	return c, nil
}

// handleConversationReply handles sending messages to an existing
// conversation, optionally persisting both the sent replies and responses.
func handleConversationReply(c *Conversation, persist bool, toSend ...Message) {
	existing, err := store.Messages(c)
	if err != nil {
		Fatal("Could not retrieve messages for conversation: %s\n", c.Title)
	}

	if persist {
		for _, message := range toSend {
			err = store.SaveMessage(&message)
			if err != nil {
				Warn("Could not save %s message: %v\n", message.Role, err)
			}
		}
	}

	allMessages := append(existing, toSend...)

	RenderConversation(allMessages, true)

	// render a message header with no contents
	(&Message{Role: MessageRoleAssistant}).RenderTTY()

	replies, err := fetchAndShowCompletion(allMessages)
	if err != nil {
		Fatal("Error fetching LLM response: %v\n", err)
	}

	if persist {
		for _, reply := range replies {
			reply.ConversationID = c.ID

			err = store.SaveMessage(&reply)
			if err != nil {
				Warn("Could not save reply: %v\n", err)
			}
		}
	}
}

// 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 = InputFromEditor(placeholder, "message.*.md", existingMessage)
		if err != nil {
			Fatal("Failed to get input: %v\n", err)
		}
	} else {
		message = strings.Trim(strings.Join(args, " "), " \t\n")
	}
	return
}

var rootCmd = &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()
	},
}

var listCmd = &cobra.Command{
	Use:   "list",
	Aliases: []string{"ls"},
	Short: "List conversations",
	Long:  `List conversations in order of recent activity`,
	Run: func(cmd *cobra.Command, args []string) {
		conversations, err := store.Conversations()
		if err != nil {
			Fatal("Could not fetch conversations.\n")
			return
		}

		type Category struct {
			name   string
			cutoff time.Duration
		}

		type ConversationLine struct {
			timeSinceReply time.Duration
			formatted      string
		}

		now := time.Now()

		midnight := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, now.Location())
		monthStart := time.Date(now.Year(), now.Month(), 1, 0, 0, 0, 0, now.Location())
		dayOfWeek := int(now.Weekday())
		categories := []Category{
			{"today", now.Sub(midnight)},
			{"yesterday", now.Sub(midnight.AddDate(0, 0, -1))},
			{"this week", now.Sub(midnight.AddDate(0, 0, -dayOfWeek))},
			{"last week", now.Sub(midnight.AddDate(0, 0, -(dayOfWeek + 7)))},
			{"this month", now.Sub(monthStart)},
			{"last month", now.Sub(monthStart.AddDate(0, -1, 0))},
			{"2 months ago", now.Sub(monthStart.AddDate(0, -2, 0))},
			{"3 months ago", now.Sub(monthStart.AddDate(0, -3, 0))},
			{"4 months ago", now.Sub(monthStart.AddDate(0, -4, 0))},
			{"5 months ago", now.Sub(monthStart.AddDate(0, -5, 0))},
			{"older", now.Sub(time.Time{})},
		}
		categorized := map[string][]ConversationLine{}

		all, _ := cmd.Flags().GetBool("all")

		for _, conversation := range conversations {
			lastMessage, err := store.LastMessage(&conversation)
			if lastMessage == nil || err != nil {
				continue
			}

			messageAge := now.Sub(lastMessage.CreatedAt)

			var category string
			for _, c := range categories {
				if messageAge < c.cutoff {
					category = c.name
					break
				}
			}

			formatted := fmt.Sprintf(
				"%s - %s - %s",
				conversation.ShortName.String,
				humanTimeElapsedSince(messageAge),
				conversation.Title,
			)

			categorized[category] = append(
				categorized[category],
				ConversationLine{messageAge, formatted},
			)
		}

		var conversationsPrinted int
	outer:
		for _, category := range categories {
			conversations, ok := categorized[category.name]
			if !ok {
				continue
			}

			slices.SortFunc(conversations, func(a, b ConversationLine) int {
				return int(a.timeSinceReply - b.timeSinceReply)
			})

			fmt.Printf("%s:\n", category.name)
			for _, conv := range conversations {
				if conversationsPrinted >= LS_LIMIT && !all {
					fmt.Printf("%d remaining message(s), use --all to view.\n", len(conversations)-conversationsPrinted)
					break outer
				}

				fmt.Printf("  %s\n", conv.formatted)
				conversationsPrinted++
			}
		}
	},
}

var rmCmd = &cobra.Command{
	Use:   "rm <conversation>...",
	Short: "Remove conversations",
	Long:  `Remove conversations by their short names.`,
	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) {
		var toRemove []*Conversation
		for _, shortName := range args {
			conversation := lookupConversation(shortName)
			toRemove = append(toRemove, conversation)
		}
		var errors []error
		for _, c := range toRemove {
			err := store.DeleteConversation(c)
			if err != nil {
				errors = append(errors, fmt.Errorf("Could not remove conversation %s: %v", c.ShortName.String, err))
			}
		}
		for _, err := range errors {
			fmt.Fprintln(os.Stderr, err.Error())
		}
		if len(errors) > 0 {
			os.Exit(1)
		}
	},
	ValidArgsFunction: func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
		compMode := cobra.ShellCompDirectiveNoFileComp
		var completions []string
	outer:
		for _, completion := range store.ConversationShortNameCompletions(toComplete) {
			parts := strings.Split(completion, "\t")
			for _, arg := range args {
				if parts[0] == arg {
					continue outer
				}
			}
			completions = append(completions, completion)
		}
		return completions, compMode
	},
}

var cloneCmd = &cobra.Command{
	Use:   "clone <conversation>",
	Short: "Clone conversations",
	Long:  `Clones the provided 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]
		toClone, err := lookupConversationE(shortName)
		if err != nil {
			return err
		}

		messagesToCopy, err := store.Messages(toClone)
		if err != nil {
			return fmt.Errorf("Could not retrieve messages for conversation: %s", toClone.ShortName.String)
		}

		clone := &Conversation{
			Title: toClone.Title + " - Clone",
		}
		if err := store.SaveConversation(clone); err != nil {
			return fmt.Errorf("Cloud not create clone: %s", err)
		}

		var errors []error
		messageCnt := 0
		for _, message := range messagesToCopy {
			newMessage := message
			newMessage.ConversationID = clone.ID
			newMessage.ID = 0
			if err := store.SaveMessage(&newMessage); err != nil {
				errors = append(errors, err)
			} else {
				messageCnt++
			}
		}

		if len(errors) > 0 {
			return fmt.Errorf("Messages failed to be cloned: %v", errors)
		}

		fmt.Printf("Cloned %d messages to: %s\n", messageCnt, clone.Title)
		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 store.ConversationShortNameCompletions(toComplete), compMode
	},
}

var viewCmd = &cobra.Command{
	Use:   "view <conversation>",
	Short: "View messages in a conversation",
	Long:  `Finds a conversation by its short name and displays its contents.`,
	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 := lookupConversation(shortName)

		messages, err := store.Messages(conversation)
		if err != nil {
			Fatal("Could not retrieve messages for conversation: %s\n", conversation.Title)
		}

		RenderConversation(messages, false)
	},
	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 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 := lookupConversation(shortName)
		var err error

		generate, _ := cmd.Flags().GetBool("generate")
		var title string
		if generate {
			title, err = conversation.GenerateTitle()
			if err != nil {
				Fatal("Could not generate conversation title: %v\n", err)
			}
		} else {
			if len(args) < 2 {
				Fatal("Conversation title not provided.\n")
			}
			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 replyCmd = &cobra.Command{
	Use:   "reply <conversation> [message]",
	Short: "Reply to a conversation",
	Long:  `Sends a reply to conversation and writes the response to stdout.`,
	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 := lookupConversation(shortName)

		reply := inputFromArgsOrEditor(args[1:], "# How would you like to reply?\n", "")
		if reply == "" {
			Fatal("No reply was provided.\n")
		}

		handleConversationReply(conversation, true, Message{
			ConversationID:  conversation.ID,
			Role:            MessageRoleUser,
			OriginalContent: reply,
		})
	},
	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 newCmd = &cobra.Command{
	Use:   "new [message]",
	Short: "Start a new conversation",
	Long:  `Start a new conversation with the Large Language Model.`,
	Run: func(cmd *cobra.Command, args []string) {
		messageContents := inputFromArgsOrEditor(args, "# What would you like to say?\n", "")
		if messageContents == "" {
			Fatal("No message was provided.\n")
		}

		conversation := &Conversation{}
		err := store.SaveConversation(conversation)
		if err != nil {
			Fatal("Could not save new conversation: %v\n", err)
		}

		messages := []Message{
			{
				ConversationID:  conversation.ID,
				Role:            MessageRoleSystem,
				OriginalContent: getSystemPrompt(),
			},
			{
				ConversationID:  conversation.ID,
				Role:            MessageRoleUser,
				OriginalContent: messageContents,
			},
		}

		handleConversationReply(conversation, true, messages...)

		title, err := conversation.GenerateTitle()
		if err != nil {
			Warn("Could not generate title for conversation: %v\n", err)
		}

		conversation.Title = title

		err = store.SaveConversation(conversation)
		if err != nil {
			Warn("Could not save conversation after generating title: %v\n", err)
		}
	},
}

var promptCmd = &cobra.Command{
	Use:   "prompt [message]",
	Short: "Do a one-shot prompt",
	Long:  `Prompt the Large Language Model and get a response.`,
	Run: func(cmd *cobra.Command, args []string) {
		message := inputFromArgsOrEditor(args, "# What would you like to say?\n", "")
		if message == "" {
			Fatal("No message was provided.\n")
		}

		messages := []Message{
			{
				Role:            MessageRoleSystem,
				OriginalContent: getSystemPrompt(),
			},
			{
				Role:            MessageRoleUser,
				OriginalContent: message,
			},
		}

		_, err := fetchAndShowCompletion(messages)
		if err != nil {
			Fatal("Error fetching LLM response: %v\n", err)
		}
	},
}

var retryCmd = &cobra.Command{
	Use:   "retry <conversation>",
	Short: "Retry the last user reply in a conversation",
	Long:  `Re-prompt the conversation up to the last user response. Can be used to regenerate the last assistant reply, or simply generate one if an error occurred.`,
	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 := lookupConversation(shortName)

		messages, err := store.Messages(conversation)
		if err != nil {
			Fatal("Could not retrieve messages for conversation: %s\n", conversation.Title)
		}

		// walk backwards through the conversation and delete messages, break
		// when we find the latest user response
		for i := len(messages) - 1; i >= 0; i-- {
			if messages[i].Role == MessageRoleUser {
				break
			}

			err = store.DeleteMessage(&messages[i])
			if err != nil {
				Warn("Could not delete previous reply: %v\n", err)
			}
		}

		handleConversationReply(conversation, true)
	},
	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 continueCmd = &cobra.Command{
	Use:   "continue <conversation>",
	Short: "Continue a conversation from the last message",
	Long:  `Re-prompt the conversation with all existing prompts. Useful if a reply was cut short.`,
	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 := lookupConversation(shortName)
		handleConversationReply(conversation, true)
	},
	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 editCmd = &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
	},
	Run: func(cmd *cobra.Command, args []string) {
		shortName := args[0]
		conversation := lookupConversation(shortName)

		messages, err := store.Messages(conversation)
		if err != nil {
			Fatal("Could not retrieve messages for conversation: %s\n", conversation.Title)
		}

		offset, _ := cmd.Flags().GetInt("offset")
		if offset < 0 {
			offset = -offset
		}

		if offset > len(messages) - 1   {
			Fatal("Offset %d is before the start of the conversation\n", offset)
		}

		desiredIdx := len(messages) - 1 - offset

		// walk backwards through the conversation deleting messages until and
		// including the last user message
		toRemove := []Message{}
		var toEdit *Message
		for i := len(messages) - 1; i >= 0; i-- {
			if i == desiredIdx {
				toEdit = &messages[i]
			}
			toRemove = append(toRemove, messages[i])
			messages = messages[:i]
			if toEdit != nil {
				break
			}
		}

		existingContents := toEdit.OriginalContent

		newContents := inputFromArgsOrEditor(args[1:], "# Save when finished editing\n", existingContents)
		switch newContents {
		case existingContents:
			Fatal("No edits were made.\n")
		case "":
			Fatal("No message was provided.\n")
		}

		for _, message := range toRemove {
			err = store.DeleteMessage(&message)
			if err != nil {
				Warn("Could not delete message: %v\n", err)
			}
		}

		handleConversationReply(conversation, true, Message{
			ConversationID:  conversation.ID,
			Role:            MessageRoleUser,
			OriginalContent: newContents,
		})
	},
	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
	},
}