From 7d56726c78d1127420a7acdda0326dd2db620c8e Mon Sep 17 00:00:00 2001
From: Matt Low <matt@mlow.ca>
Date: Tue, 12 Mar 2024 07:43:57 +0000
Subject: [PATCH] Add --model flag completion

---
 pkg/cmd/cmd.go     |  3 +++
 pkg/lmcli/lmcli.go | 10 ++++++++++
 2 files changed, 13 insertions(+)

diff --git a/pkg/cmd/cmd.go b/pkg/cmd/cmd.go
index e3963d5..26b85a6 100644
--- a/pkg/cmd/cmd.go
+++ b/pkg/cmd/cmd.go
@@ -38,6 +38,9 @@ func RootCmd(ctx *lmcli.Context) *cobra.Command {
 	inputCmds := []*cobra.Command{newCmd, promptCmd, replyCmd, retryCmd, continueCmd, editCmd}
 	for _, cmd := range inputCmds {
 		cmd.Flags().StringVar(ctx.Config.Defaults.Model, "model", *ctx.Config.Defaults.Model, "Which model to use")
+		cmd.RegisterFlagCompletionFunc("model", func(*cobra.Command, []string, string) ([]string, cobra.ShellCompDirective) {
+			return ctx.GetModels(), cobra.ShellCompDirectiveDefault
+		})
 		cmd.Flags().IntVar(ctx.Config.Defaults.MaxTokens, "length", *ctx.Config.Defaults.MaxTokens, "Maximum response tokens")
 		cmd.Flags().StringVar(ctx.Config.Defaults.SystemPrompt, "system-prompt", *ctx.Config.Defaults.SystemPrompt, "System prompt")
 		cmd.Flags().StringVar(&systemPromptFile, "system-prompt-file", "", "A path to a file containing the system prompt")
diff --git a/pkg/lmcli/lmcli.go b/pkg/lmcli/lmcli.go
index b2e99d9..4cbb044 100644
--- a/pkg/lmcli/lmcli.go
+++ b/pkg/lmcli/lmcli.go
@@ -42,6 +42,16 @@ func NewContext() (*Context, error) {
 	return &Context{*config, s, hl}, nil
 }
 
+func (c *Context) GetModels() (models []string) {
+	for _, m := range *c.Config.Anthropic.Models {
+		models = append(models, m)
+	}
+	for _, m := range *c.Config.OpenAI.Models {
+		models = append(models, m)
+	}
+	return
+}
+
 func (c *Context) GetCompletionProvider(model string) (provider.ChatCompletionClient, error) {
 	for _, m := range *c.Config.Anthropic.Models {
 		if m == model {