522 lines
14 KiB
Go
522 lines
14 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
|
|
}
|
|
|
|
// LLMRequest prompts the LLM with the given messages, writing the response
|
|
// to stdout. Returns all reply messages added by the LLM, including any
|
|
// function call messages.
|
|
func LLMRequest(messages []Message) ([]Message, 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)
|
|
|
|
var replies []Message
|
|
response, err := CreateChatCompletionStream(model, messages, maxTokens, receiver, &replies)
|
|
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 replies, err
|
|
}
|
|
|
|
func (c *Conversation) GenerateAndSaveReplies(messages []Message) {
|
|
replies, err := LLMRequest(messages)
|
|
if err != nil {
|
|
Fatal("Error fetching LLM response: %v\n", err)
|
|
}
|
|
|
|
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) (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: MessageRoleUser,
|
|
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)
|
|
(&Message{Role: MessageRoleAssistant}).RenderTTY()
|
|
|
|
conversation.GenerateAndSaveReplies(messages)
|
|
},
|
|
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,
|
|
},
|
|
}
|
|
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)
|
|
(&Message{Role: MessageRoleAssistant}).RenderTTY()
|
|
|
|
conversation.GenerateAndSaveReplies(messages)
|
|
|
|
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: MessageRoleSystem,
|
|
OriginalContent: SystemPrompt(),
|
|
},
|
|
{
|
|
Role: MessageRoleUser,
|
|
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
|
|
// walk backwards through conversations to find last user message
|
|
for i := len(messages) - 1; i >= 0; i-- {
|
|
if messages[i].Role == MessageRoleUser {
|
|
lastUserMessageIndex = i
|
|
break
|
|
}
|
|
|
|
if lastUserMessageIndex == 0 {
|
|
// haven't found the the last user message yet, delete this one
|
|
err = store.DeleteMessage(&messages[i])
|
|
if err != nil {
|
|
Warn("Could not delete previous reply: %v\n", err)
|
|
}
|
|
}
|
|
}
|
|
|
|
messages = messages[:lastUserMessageIndex+1]
|
|
|
|
RenderConversation(messages, true)
|
|
(&Message{Role: MessageRoleAssistant}).RenderTTY()
|
|
|
|
conversation.GenerateAndSaveReplies(messages)
|
|
},
|
|
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)
|
|
(&Message{Role: MessageRoleAssistant}).RenderTTY()
|
|
|
|
conversation.GenerateAndSaveReplies(messages)
|
|
},
|
|
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
|
|
},
|
|
}
|