diff --git a/plugins/ollama/README.md b/plugins/ollama/README.md index f26f38cf9..9224f9320 100644 --- a/plugins/ollama/README.md +++ b/plugins/ollama/README.md @@ -31,6 +31,13 @@ This plugin enhances your Zsh shell environment by integrating powerful features omz plugin enable ollama ``` + In order to get the most benefit from completions, with helpful usage hints, etc: + ```sh + # ~/.zshrc + # add the following zstyle entry wherever you want + zstyle ':completion:*:*:*:*:descriptions' format '%F{green}%d%f' + ``` + 2. **Restart Your Shell** Apply the changes by reloading Oh-My-Zsh: diff --git a/plugins/ollama/ollama.plugin.zsh b/plugins/ollama/ollama.plugin.zsh index b288bbc18..e57e3b080 100644 --- a/plugins/ollama/ollama.plugin.zsh +++ b/plugins/ollama/ollama.plugin.zsh @@ -1,38 +1,50 @@ -# Function to retrieve available models for completion -_ollama_get_models() { - # Execute 'ollama list' and capture its output, suppressing any error messages - local models_output - models_output="$(ollama list 2>/dev/null)" +# ------------------------------------------------------------------------------ +# ollama.plugin.zsh +# +# Plugin providing Zsh completions for the `ollama` command. +# ------------------------------------------------------------------------------ - # Initialize an array to hold the model suggestions +# ------------------------------------------------------------------------------ +# Function: _ollama_get_models +# Purpose: Retrieves the list of available models for completion. +# Uses `ollama list` with a short timeout and provides candidates. +# ------------------------------------------------------------------------------ +_ollama_get_models() { + local models_output + local timeout=5 # Timeout duration in seconds + + # Attempt to fetch models via `ollama list`; if it fails, show a short message. + models_output="$(timeout $timeout ollama list 2>/dev/null)" || { + _message "Failed to fetch models" + return 1 + } + + # Accumulate parsed model names here local -a models local line + while IFS= read -r line; do + # Skip blank lines and header lines (starting with NAME) + [[ -z "$line" || "$line" =~ ^NAME ]] && continue - # Read the output line by line - while IFS=" " read -r line; do - # Skip blank lines - [[ -z "$line" ]] && continue - - # Skip the header line that starts with 'NAME' - if [[ "$line" =~ ^NAME ]]; then - continue - fi - - # Split the line into words and extract the first word (model name:tag) - set -- $line - local suggestion="${$(echo $1 | cut -d ' ' -f 1)/:/\\:}" # Escape ':' by replacing it with '\:' - models+=( "$suggestion" ) # Add the escaped model name to the array + # Extract the first column and escape any colons for safety + local suggestion="${line%% *}" + suggestion="${suggestion/:/\\:}" + models+=("$suggestion") done <<< "$models_output" - # Use the '_describe' function to provide the model suggestions for completion + # Provide model suggestions using `_describe` _describe -t models 'models' models } -# Main completion function for the 'ollama' command +# ------------------------------------------------------------------------------ +# Function: _ollama +# Purpose: The main completion function for the `ollama` CLI. Determines which +# subcommand is being completed, then sets up the corresponding flags +# and suggestions. +# ------------------------------------------------------------------------------ _ollama() { - # Define an array of available commands with their descriptions - local -a commands - commands=( + # List of top-level commands and their descriptions + local -a commands=( 'serve:Start the Ollama server' 'create:Create a model from a Modelfile' 'show:Display information about a specific model' @@ -47,59 +59,113 @@ _ollama() { 'help:Provide help information for a command' ) - # Initialize context variables for the completion - local context curcontext="$curcontext" state line + # Standard local variables used by _arguments + local curcontext="$curcontext" state line local -A opt_args - # Define the arguments and options for the 'ollama' command + # The main `_arguments` call for handling top-level options (e.g. -h, -v) + # and capturing the first positional argument -> subcommand, then the rest. _arguments -C \ '(-h --help)'{-h,--help}'[Display help information]' \ '(-v --version)'{-v,--version}'[Show version information]' \ '1: :->command' \ '*:: :->args' - # Determine the state of the completion (command or arguments) + # If the user is trying to complete the first argument (the subcommand), + # then we present them the `commands` array above. case $state in command) - # Provide command suggestions _describe -t commands 'ollama commands' commands + return ;; - args) - # Handle argument completion based on the specified command - case $words[1] in - run|rm|stop|show|pull|push) - # For these commands, provide model name suggestions exactly once - if [[ $CURRENT -eq 2 ]]; then - _ollama_get_models - fi - ;; - cp) - if [[ $CURRENT -eq 2 ]]; then - # For the 'cp' command, suggest source model names - _ollama_get_models - elif [[ $CURRENT -eq 3 ]]; then - # For the 'cp' command, prompt for the destination model name - _message 'destination model name' - fi - ;; - create) - # For the 'create' command, suggest Modelfile paths - _arguments \ - '(-f --filename)'{-f,--filename}'[Specify the path to the Modelfile]:Modelfile:_files' - ;; - serve) - # For the 'serve' command, suggest specifying the port number - _arguments \ - '(-p --port)'{-p,--port}'[Specify the port number]:port number:' - ;; - *) - # For any other commands, use the default completion - _default - ;; - esac + esac + + # If the first argument is known, proceed with subcommand-specific completions + case $words[1] in + serve) + _arguments \ + '(-p --port)'{-p,--port}'[Specify the port number]:port number:' + ;; + + create) + # If user typed only `ollama create ` (with no second arg), + # display a short message to remind them to name the new model + if [[ $CURRENT -eq 2 ]]; then + _message 'Specify the new model name' + else + # Otherwise, offer flags for `create` + _arguments \ + '(-f --filename)'{-f,--filename}'[Path to the Modelfile]:Modelfile:_files' \ + '(-q --quantize)'{-q,--quantize}'[Quantization method (e.g. q4_0)]' \ + '--prefix[Set a prefix for the created model]' \ + '(-h --help)--help[Show help for create]' + fi + ;; + + show) + _message 'Usage: ollama show MODEL [flags]' + if [[ $CURRENT -eq 2 ]]; then + _ollama_get_models + else + _arguments \ + '--license[Show the model’s license]' \ + '--modelfile[Show the model’s Modelfile]' \ + '--parameters[Show model parameters]' \ + '--system[Show the system message of the model]' \ + '--template[Show the model’s template]' \ + '(-h --help)--help[Show help for show]' + fi + ;; + + run) + # Display usage message only if there's no argument yet + if [[ $CURRENT -eq 2 ]]; then + _message "Usage: ollama run MODEL [PROMPT] [flags]" + _ollama_get_models + else + # Define flags for the `run` command + local -a _run_flags=( + '--format-string=[Format string for the output (e.g. json)]' + '--insecure[Use an insecure registry]' + '--keepalive=[Time to keep the model loaded (e.g. 5m)]' + '--nowordwrap[Disable word wrapping]' + '--verbose[Show response timings]' + '(-h --help)--help[Show help for run]' + ) + + # Use a mix of `_arguments` and manual handling for freeform input + if [[ $CURRENT -eq 3 ]]; then + # Suggest a freeform prompt (arbitrary input) + _message "Enter a prompt as a string" + else + # Provide flag completions + _arguments -S "${_run_flags[@]}" + fi + fi + ;; + + cp) + # The `cp` command expects `ollama cp SOURCE DEST` + if [[ $CURRENT -eq 2 ]]; then + _ollama_get_models + elif [[ $CURRENT -eq 3 ]]; then + _message 'Specify the destination model name' + fi + ;; + + rm|stop|pull|push) + # All of these commands accept one or more model names + if [[ $CURRENT -eq 2 ]]; then + _ollama_get_models + fi + ;; + + # If the subcommand doesn’t match anything above, fall back to default + *) + _default ;; esac } -# Register the '_ollama' function as the completion handler for the 'ollama' command +# Finally, register the completion function for the `ollama` command compdef _ollama ollama