Merge 4ea010a520cb3bd9e21c641b38d06b7abc8a401a into d7eb05b9361febead29a74e71ddffc2ebeff5302

.gitignore

@@ -1,5 +1,6 @@
 .DS_Store
 .vscode
+.idea
 .env
 .venv
 .swp

README.md

@@ -28,6 +28,18 @@ curl -fsSL https://ollama.com/install.sh | sh
 
 The official [Ollama Docker image](https://hub.docker.com/r/ollama/ollama) `ollama/ollama` is available on Docker Hub.
 
+### Autocompletion
+To enable autocompletion generate completion script for your shell (`bash`, `zsh`, `fish`):
+```
+echo "source <(./ollama completion bash)" >> ~/.bashrc # Required bash-completion package installed
+```
+```
+echo '[[ $commands[ollama] ]] && source <(ollama completion zsh)' >> ~/.zshrc
+```
+```
+ollama completion fish > ~/.config/fish/completions/ollama.fish
+```
+
 ### Libraries
 
 - [ollama-python](https://github.com/ollama/ollama-python)

cmd/cmd.go

@@ -1279,6 +1279,52 @@ func checkServerHeartbeat(cmd *cobra.Command, _ []string) error {
 	return nil
 }
 
+func completionHandler(cmd *cobra.Command, args []string) error {
+	var err error
+	switch args[0] {
+	case "bash":
+		err = cmd.Root().GenBashCompletion(os.Stdout)
+	case "zsh":
+		err = cmd.Root().GenZshCompletion(os.Stdout)
+	case "fish":
+		err = cmd.Root().GenFishCompletion(os.Stdout, true)
+	default:
+		err = errors.New("unsupported shell. Supported shells: zsh, fish, bash")
+	}
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "error: %s\n", err)
+	}
+	return nil
+}
+
+func autocompleteModelName(_ *cobra.Command, _ []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+	log.Printf("autocomplete: %s", toComplete)
+	client, err := api.ClientFromEnvironment()
+	if err != nil {
+		return nil, cobra.ShellCompDirectiveError
+	}
+
+	models, err := client.List(context.Background())
+	if err != nil {
+		return nil, cobra.ShellCompDirectiveError
+	}
+
+	var data []string
+
+	for _, m := range models.Models {
+		log.Printf("model: %s", m.Name)
+		if strings.HasPrefix(m.Name, toComplete) {
+			data = append(data, m.Name[:strings.IndexByte(m.Name, ':')])
+		}
+	}
+
+	return data, cobra.ShellCompDirectiveNoFileComp
+}
+
+func doNotAutocomplete(_ *cobra.Command, _ []string, _ string) ([]string, cobra.ShellCompDirective) {
+	return []string{}, cobra.ShellCompDirectiveNoFileComp
+}
+
 func versionHandler(cmd *cobra.Command, _ []string) {
 	client, err := api.ClientFromEnvironment()
 	if err != nil {
@@ -1348,6 +1394,7 @@ func NewCLI() *cobra.Command {
 		Args:              cobra.ExactArgs(1),
 		PreRunE:           checkServerHeartbeat,
 		RunE:              CreateHandler,
+		ValidArgsFunction: doNotAutocomplete,
 	}
 
 	createCmd.Flags().StringP("file", "f", "", "Name of the Modelfile (default \"Modelfile\"")
@@ -1359,6 +1406,7 @@ func NewCLI() *cobra.Command {
 		Args:              cobra.ExactArgs(1),
 		PreRunE:           checkServerHeartbeat,
 		RunE:              ShowHandler,
+		ValidArgsFunction: autocompleteModelName,
 	}
 
 	showCmd.Flags().Bool("license", false, "Show license of a model")
@@ -1373,6 +1421,7 @@ func NewCLI() *cobra.Command {
 		Args:              cobra.MinimumNArgs(1),
 		PreRunE:           checkServerHeartbeat,
 		RunE:              RunHandler,
+		ValidArgsFunction: autocompleteModelName,
 	}
 
 	runCmd.Flags().String("keepalive", "", "Duration to keep a model loaded (e.g. 5m)")
@@ -1395,6 +1444,7 @@ func NewCLI() *cobra.Command {
 		Short:             "Start ollama",
 		Args:              cobra.ExactArgs(0),
 		RunE:              RunServer,
+		ValidArgsFunction: doNotAutocomplete,
 	}
 
 	pullCmd := &cobra.Command{
@@ -1403,6 +1453,7 @@ func NewCLI() *cobra.Command {
 		Args:              cobra.ExactArgs(1),
 		PreRunE:           checkServerHeartbeat,
 		RunE:              PullHandler,
+		ValidArgsFunction: doNotAutocomplete,
 	}
 
 	pullCmd.Flags().Bool("insecure", false, "Use an insecure registry")
@@ -1413,6 +1464,7 @@ func NewCLI() *cobra.Command {
 		Args:              cobra.ExactArgs(1),
 		PreRunE:           checkServerHeartbeat,
 		RunE:              PushHandler,
+		ValidArgsFunction: autocompleteModelName,
 	}
 
 	pushCmd.Flags().Bool("insecure", false, "Use an insecure registry")
@@ -1423,6 +1475,7 @@ func NewCLI() *cobra.Command {
 		Short:             "List models",
 		PreRunE:           checkServerHeartbeat,
 		RunE:              ListHandler,
+		ValidArgsFunction: doNotAutocomplete,
 	}
 
 	psCmd := &cobra.Command{
@@ -1438,6 +1491,7 @@ func NewCLI() *cobra.Command {
 		Args:              cobra.ExactArgs(2),
 		PreRunE:           checkServerHeartbeat,
 		RunE:              CopyHandler,
+		ValidArgsFunction: autocompleteModelName,
 	}
 
 	deleteCmd := &cobra.Command{
@@ -1446,6 +1500,19 @@ func NewCLI() *cobra.Command {
 		Args:              cobra.MinimumNArgs(1),
 		PreRunE:           checkServerHeartbeat,
 		RunE:              DeleteHandler,
+		ValidArgsFunction: autocompleteModelName,
+	}
+
+	completionCmd := &cobra.Command{
+		Use:                   "completion [bash|zsh|fish]",
+		Short:                 "Generate completion scripts",
+		DisableFlagsInUseLine: true,
+		Hidden:                true,
+		Args:                  cobra.ExactArgs(1),
+		RunE:                  completionHandler,
+		ValidArgsFunction: func(_ *cobra.Command, _ []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+			return []string{"bash", "zsh", "fish"}, cobra.ShellCompDirectiveNoFileComp
+		},
 	}
 
 	envVars := envconfig.AsMap()
@@ -1503,6 +1570,7 @@ func NewCLI() *cobra.Command {
 		psCmd,
 		copyCmd,
 		deleteCmd,
+		completionCmd,
 	)
 
 	return rootCmd