mirror of https://github.com/ohmyzsh/ohmyzsh.git
trying to get flags to work. some movement, but not quite there yet.
This commit is contained in:
sealad886
parent
d94d2bf8d1
commit
1b1e6c9546
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in New Issue