lmcli/pkg/cli/cmd.go

package cli

import (
	"fmt"
	"strings"

	"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) {
		fmt.Println("Listing conversations...")
		// Example output, asterisk to indicate current converation

		// $ lm 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
	},
}

var viewCmd = &cobra.Command{
	Use:   "view",
	Short: "View messages in a conversation",
	Long:  `Displays all the messages in a coversation.`,
	Run: func(cmd *cobra.Command, args []string) {
		fmt.Println("Displaying conversation messages...")
	},
}

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.`,
	Run: func(cmd *cobra.Command, args []string) {
		fmt.Println("Replying to a conversation...")
	},
}

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(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()
	},
}

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(messages, MAX_TOKENS, receiver)
		if err != nil {
			Fatal("%v\n", err)
		}

		fmt.Println()
	},
}

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