mirror of
https://github.com/ohmyzsh/ohmyzsh.git
synced 2025-12-05 01:46:46 +01:00
105 lines
3.4 KiB
Bash
105 lines
3.4 KiB
Bash
# 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)"
|
|
|
|
# Initialize an array to hold the model suggestions
|
|
local -a models
|
|
local line
|
|
|
|
# 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
|
|
done <<< "$models_output"
|
|
|
|
# Use the '_describe' function to provide the model suggestions for completion
|
|
_describe -t models 'models' models
|
|
}
|
|
|
|
# Main completion function for the 'ollama' command
|
|
_ollama() {
|
|
# Define an array of available commands with their descriptions
|
|
local -a commands
|
|
commands=(
|
|
'serve:Start the Ollama server'
|
|
'create:Create a model from a Modelfile'
|
|
'show:Display information about a specific model'
|
|
'run:Execute a model with a given prompt'
|
|
'stop:Terminate a running model'
|
|
'pull:Download a model from the registry'
|
|
'push:Upload a model to the registry'
|
|
'list:Display all available models'
|
|
'ps:Show currently running models'
|
|
'cp:Duplicate an existing model'
|
|
'rm:Delete a model from the local system'
|
|
'help:Provide help information for a command'
|
|
)
|
|
|
|
# Initialize context variables for the completion
|
|
local context curcontext="$curcontext" state line
|
|
local -A opt_args
|
|
|
|
# Define the arguments and options for the 'ollama' command
|
|
_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)
|
|
case $state in
|
|
command)
|
|
# Provide command suggestions
|
|
_describe -t commands 'ollama commands' commands
|
|
;;
|
|
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
|
|
}
|
|
|
|
# Register the '_ollama' function as the completion handler for the 'ollama' command
|
|
compdef _ollama ollama
|