lmcli/pkg/cli/cmd.go

package cli

import (
	"fmt"
	"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}
	for _, cmd := range inputCmds {
		cmd.Flags().IntVar(&maxTokens, "length", *config.OpenAI.DefaultMaxLength, "Max response length in tokens")
		cmd.Flags().StringVar(&model, "model", *config.OpenAI.DefaultModel, "The language model to use")
		cmd.Flags().StringVar(&systemPrompt, "system-prompt", *config.ModelDefaults.SystemPrompt, "The system prompt to use.")
		cmd.Flags().StringVar(&systemPromptFile, "system-prompt-file", "", "A path to a file whose contents are used as the system prompt.")
		cmd.MarkFlagsMutuallyExclusive("system-prompt", "system-prompt-file")
	}

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

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
}

// LLMRequest prompts the LLM with the given Message, writes the (partial)
// response to stdout, and returns the (partial) response or any errors.
func LLMRequest(messages []Message) (string, error) {
	// receiver receives the reponse from LLM
	receiver := make(chan string)
	defer close(receiver)

	// start HandleDelayedContent goroutine to print received data to stdout
	go HandleDelayedContent(receiver)

	response, err := CreateChatCompletionStream(model, messages, maxTokens, receiver)
	if response != "" {
		if err != nil {
			Warn("Received partial response. Error: %v\n", err)
			err = nil
		}
		// there was some content, so break to a new line after it
		fmt.Println()
	}

	return response, 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) (message string) {
	var err error
	if len(args) == 0 {
		message, err = InputFromEditor(placeholder, "message.*.md")
		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",
	Short: "Interact with Large Language Models",
	Long:  `lmcli is a CLI tool to interact 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 a conversation",
	Long:  `Removes a conversation by its short 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, err := store.ConversationByShortName(shortName)
		if err != nil {
			Fatal("Could not search for conversation: %v\n", err)
		}
		if conversation.ID == 0 {
			Fatal("Conversation not found with short name: %s\n", shortName)
		}
		err = store.DeleteConversation(conversation)
		if err != nil {
			Fatal("Could not delete conversation: %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 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, err := store.ConversationByShortName(shortName)
		if conversation.ID == 0 {
			Fatal("Conversation not found with short name: %s\n", 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 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, err := store.ConversationByShortName(shortName)
		if conversation.ID == 0 {
			Fatal("Conversation not found with short name: %s\n", shortName)
		}

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

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

		userReply := Message{
			ConversationID:  conversation.ID,
			Role:            "user",
			OriginalContent: messageContents,
		}

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

		messages = append(messages, userReply)

		RenderConversation(messages, true)
		assistantReply := Message{
			ConversationID: conversation.ID,
			Role:           "assistant",
		}
		assistantReply.RenderTTY()

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

		assistantReply.OriginalContent = response

		err = store.SaveMessage(&assistantReply)
		if err != nil {
			Fatal("Could not save assistant reply: %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 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:            "system",
				OriginalContent: SystemPrompt(),
			},
			{
				ConversationID:  conversation.ID,
				Role:            "user",
				OriginalContent: messageContents,
			},
		}
		for _, message := range messages {
			err = store.SaveMessage(&message)
			if err != nil {
				Warn("Could not save %s message: %v\n", message.Role, err)
			}
		}

		RenderConversation(messages, true)
		reply := Message{
			ConversationID: conversation.ID,
			Role:           "assistant",
		}
		reply.RenderTTY()

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

		reply.OriginalContent = response

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

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

		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:            "system",
				OriginalContent: SystemPrompt(),
			},
			{
				Role:            "user",
				OriginalContent: message,
			},
		}

		_, err := LLMRequest(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, err := store.ConversationByShortName(shortName)
		if conversation.ID == 0 {
			Fatal("Conversation not found with short name: %s\n", shortName)
		}

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

		var lastUserMessageIndex int
		for i := len(messages) - 1; i >= 0; i-- {
			if messages[i].Role == "user" {
				lastUserMessageIndex = i
				break
			}
		}

		messages = messages[:lastUserMessageIndex+1]

		RenderConversation(messages, true)
		assistantReply := Message{
			ConversationID: conversation.ID,
			Role:           "assistant",
		}
		assistantReply.RenderTTY()

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

		assistantReply.OriginalContent = response

		err = store.SaveMessage(&assistantReply)
		if err != nil {
			Fatal("Could not save assistant reply: %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 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, err := store.ConversationByShortName(shortName)
		if conversation.ID == 0 {
			Fatal("Conversation not found with short name: %s\n", shortName)
		}

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

		RenderConversation(messages, true)
		assistantReply := Message{
			ConversationID: conversation.ID,
			Role:           "assistant",
		}
		assistantReply.RenderTTY()

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

		assistantReply.OriginalContent = response

		err = store.SaveMessage(&assistantReply)
		if err != nil {
			Fatal("Could not save assistant reply: %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
	},
}