lmcli/pkg/cli/cmd.go

package cli

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

	"github.com/spf13/cobra"
)

// TODO: allow setting with flag
const MAX_TOKENS = 256

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 showCmd = &cobra.Command{
	Use:   "show [conversation]",
	Short: "Show 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)
		}

		l := len(messages)
		for i, message := range messages {
			message.RenderTTY(i < l-1)
		}
		fmt.Println()
	},
	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",
	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, err := InputFromEditor("# How would you like to reply?\n", "reply.*.md")

		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)
		for _, message := range messages {
			message.RenderTTY(true)
		}

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

		receiver := make(chan string)
		response := make(chan string)
		go func() {
			response <- HandleDelayedResponse(receiver)
		}()

		err = CreateChatCompletionStream(config.OpenAI.DefaultModel, messages, MAX_TOKENS, receiver)
		if err != nil {
			Fatal("%v\n", err)
		}

		assistantReply.OriginalContent = <-response

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

		fmt.Println()
	},
	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",
	Short: "Start a new conversation",
	Long:  `Start a new conversation with the Large Language Model.`,
	Run: func(cmd *cobra.Command, args []string) {
		messageContents, err := InputFromEditor("# What would you like to say?\n", "message.*.md")
		if err != nil {
			Fatal("Failed to get input: %v\n", err)
		}

		if messageContents == "" {
			Fatal("No message was provided.\n")
		}

		// TODO: set title if --title provided, otherwise defer for later(?)
		conversation := Conversation{}
		err = store.SaveConversation(&conversation)
		if err != nil {
			Fatal("Could not save new conversation: %v\n", err)
		}

		const system = "You are a helpful assistant."
		messages := []Message{
			{
				ConversationID:  conversation.ID,
				Role:            "system",
				OriginalContent: system,
			},
			{
				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)
			}
		}

		for _, message := range messages {
			message.RenderTTY(true)
		}

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

		receiver := make(chan string)
		response := make(chan string)
		go func() {
			response <- HandleDelayedResponse(receiver)
		}()

		err = CreateChatCompletionStream(config.OpenAI.DefaultModel, messages, MAX_TOKENS, receiver)
		if err != nil {
			Fatal("%v\n", err)
		}

		reply.OriginalContent = <-response

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

		fmt.Println()

		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",
	Short: "Do a one-shot prompt",
	Long:  `Prompt the Large Language Model and get a response.`,
	Run: func(cmd *cobra.Command, args []string) {
		message := strings.Join(args, " ")
		if len(strings.Trim(message, " \t\n")) == 0 {
			Fatal("No message was provided.\n")
		}

		const system = "You are a helpful assistant."
		messages := []Message{
			{
				Role:            "system",
				OriginalContent: system,
			},
			{
				Role:            "user",
				OriginalContent: message,
			},
		}

		receiver := make(chan string)
		go HandleDelayedResponse(receiver)
		err := CreateChatCompletionStream(config.OpenAI.DefaultModel, messages, MAX_TOKENS, receiver)
		if err != nil {
			Fatal("%v\n", err)
		}

		fmt.Println()
	},
}

func NewRootCmd() *cobra.Command {
	rootCmd.AddCommand(
		lsCmd,
		newCmd,
		promptCmd,
		replyCmd,
		rmCmd,
		showCmd,
	)
	return rootCmd
}