lmcli/pkg/cli/cmd.go

575 lines
16 KiB
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
}
// HandlePartialResponse accepts a response and an err. If err is nil, it does
// nothing and returns nil. If response != "" and err != nil, it prints a
// warning and returns nil. If response == "" and err != nil, it returns the
// error.
func HandlePartialResponse(response string, err error) (e error) {
if err != nil {
if response != "" {
Warn("Received partial response. Error: %v\n", err)
} else {
e = err
}
}
return
}
// 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()
receiver := make(chan string)
go HandleDelayedContent(receiver)
response, err := CreateChatCompletionStream(model, messages, maxTokens, receiver)
err = HandlePartialResponse(response, err)
if err != nil {
Fatal("Error while receiving response: %v\n", err)
}
assistantReply.OriginalContent = response
fmt.Println()
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()
receiver := make(chan string)
go HandleDelayedContent(receiver)
response, err := CreateChatCompletionStream(model, messages, maxTokens, receiver)
err = HandlePartialResponse(response, err)
if err != nil {
Fatal("Error while receiving response: %v\n", err)
}
fmt.Println()
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,
},
}
receiver := make(chan string)
go HandleDelayedContent(receiver)
response, err := CreateChatCompletionStream(model, messages, maxTokens, receiver)
err = HandlePartialResponse(response, err)
if err != nil {
Fatal("%v\n", err)
}
fmt.Println()
},
}
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()
receiver := make(chan string)
go HandleDelayedContent(receiver)
response, err := CreateChatCompletionStream(model, messages, maxTokens, receiver)
err = HandlePartialResponse(response, err)
if err != nil {
Fatal("Error while receiving response: %v\n", err)
}
fmt.Println()
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 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()
receiver := make(chan string)
go HandleDelayedContent(receiver)
response, err := CreateChatCompletionStream(model, messages, maxTokens, receiver)
err = HandlePartialResponse(response, err)
if err != nil {
Fatal("Error while receiving response: %v\n", err)
}
fmt.Println()
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
},
}