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
)

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")
	}

	renameCmd.Flags().Bool("generate", false, "Generate a conversation title")

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

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

func SystemPrompt() string {
	if systemPromptFile != "" {
		content, err := FileContents(systemPromptFile)
		if err != nil {
			Fatal("Could not read file contents at %s: %v", 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
}

// lookupConversationByShortname either returns the conversation found by the
// short name or exits the program
func lookupConversationByShortname(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
}

// handleConversationReply handles sending messages to an existing
// conversation, optionally persisting them. It displays the entire
// conversation before
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",
	Long: `lmcli - command-line interface with Large Language Models.`,
	Run: func(cmd *cobra.Command, args []string) {
		// execute `lm ls` by default
	},
}

var lsCmd = &cobra.Command{
	Use:   "ls",
	Short: "List existing conversations",
	Long:  `List all existing conversations in descending order of recent activity.`,
	Run: func(cmd *cobra.Command, args []string) {
		conversations, err := store.Conversations()
		if err != nil {
			fmt.Println("Could not fetch conversations.")
			return
		}

		// Example output
		// $ lmcli ls
		// last hour:
		//   98sg - 12 minutes ago - Project discussion
		// last day:
		//   tj3l - 10 hours ago - Deep learning concepts
		// last week:
		//   bwfm - 2 days ago - Machine learning study
		//   8n3h - 3 days ago - Weekend plans
		//   f3n7 - 6 days ago - CLI development
		// last month:
		//   5hn2 - 8 days ago - Book club discussion
		//   b7ze - 20 days ago - Gardening tips and tricks
		// last 6 months:
		//   3jn2 - 30 days ago - Web development best practices
		//   43jk - 2 months ago - Longboard maintenance
		//   g8d9 - 3 months ago - History book club
		//   4lk3 - 4 months ago - Local events and meetups
		//   43jn - 6 months ago - Mobile photography techniques

		type ConversationLine struct {
			timeSinceReply time.Duration
			formatted      string
		}

		now := time.Now()
		categories := []string{
			"recent",
			"last hour",
			"last 6 hours",
			"last day",
			"last week",
			"last month",
			"last 6 months",
			"older",
		}
		categorized := map[string][]ConversationLine{}

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

			var category string
			switch {
			case messageAge <= 10*time.Minute:
				category = "recent"
			case messageAge <= time.Hour:
				category = "last hour"
			case messageAge <= 6*time.Hour:
				category = "last 6 hours"
			case messageAge <= 24*time.Hour:
				category = "last day"
			case messageAge <= 7*24*time.Hour:
				category = "last week"
			case messageAge <= 30*24*time.Hour:
				category = "last month"
			case messageAge <= 6*30*24*time.Hour: // Approximate as 6 months
				category = "last 6 months"
			default:
				category = "older"
			}

			formatted := fmt.Sprintf(
				"%s - %s - %s",
				conversation.ShortName.String,
				humanTimeElapsedSince(messageAge),
				conversation.Title,
			)
			categorized[category] = append(
				categorized[category],
				ConversationLine{messageAge, formatted},
			)
		}

		for _, category := range categories {
			conversations, ok := categorized[category]
			if !ok {
				continue
			}
			slices.SortFunc(conversations, func(a, b ConversationLine) int {
				return int(a.timeSinceReply - b.timeSinceReply)
			})
			fmt.Printf("%s:\n", category)
			for _, conv := range conversations {
				fmt.Printf("  %s\n", conv.formatted)
			}
		}
	},
}

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 := lookupConversationByShortname(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 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 := lookupConversationByShortname(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 := lookupConversationByShortname(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.")
			}
			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: "Send a 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 := lookupConversationByShortname(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: SystemPrompt(),
			},
			{
				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: SystemPrompt(),
			},
			{
				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: "Retries the last conversation prompt.",
	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 := lookupConversationByShortname(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: "Continues where the previous prompt left off.",
	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 := lookupConversationByShortname(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 message 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 := lookupConversationByShortname(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 deleting messages until and
		// including the last user message
		toRemove := []Message{}
		var lastUserMessage *Message
		for i := len(messages) - 1; i >= 0; i-- {
			if messages[i].Role == MessageRoleUser {
				lastUserMessage = &messages[i]
			}

			toRemove = append(toRemove, messages[i])
			messages = messages[:i]
			if lastUserMessage != nil {
				break
			}
		}

		if lastUserMessage == nil {
			Fatal("No messages left in the conversation, nothing to edit.\n")
		}

		existingContents := lastUserMessage.OriginalContent

		newContents := InputFromArgsOrEditor(args[1:], "# Save when finished editing\n", existingContents)
		if newContents == existingContents {
			Fatal("No edits were made.\n")
		}

		if newContents == "" {
			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
	},
}