LocalAI/core/cli/completion.go

package cli

import (
	"fmt"
	"strings"

	"github.com/alecthomas/kong"
	cliContext "github.com/mudler/LocalAI/core/cli/context"
)

type CompletionCMD struct {
	Shell string `arg:"" enum:"bash,zsh,fish" help:"Shell to generate completions for (bash, zsh, fish)"`

	app *kong.Application `kong:"-"`
}

func (c *CompletionCMD) SetApplication(app *kong.Application) {
	c.app = app
}

func (c *CompletionCMD) Run(_ *cliContext.Context) error {
	if c.app == nil {
		return fmt.Errorf("application model not available")
	}

	var script string
	switch c.Shell {
	case "bash":
		script = generateBashCompletion(c.app)
	case "zsh":
		script = generateZshCompletion(c.app)
	case "fish":
		script = generateFishCompletion(c.app)
	default:
		return fmt.Errorf("unsupported shell: %s", c.Shell)
	}

	fmt.Print(script)
	return nil
}

func collectCommands(node *kong.Node, prefix string) []commandInfo {
	var cmds []commandInfo
	for _, child := range node.Children {
		if child.Hidden {
			continue
		}
		name := child.Name
		fullName := name
		if prefix != "" {
			fullName = prefix + " " + name
		}
		help := child.Help
		cmds = append(cmds, commandInfo{
			name:     name,
			fullName: fullName,
			help:     help,
			node:     child,
		})
		cmds = append(cmds, collectCommands(child, fullName)...)
	}
	return cmds
}

type commandInfo struct {
	name     string
	fullName string
	help     string
	node     *kong.Node
}

func collectFlags(node *kong.Node) []flagInfo {
	var flags []flagInfo
	seen := make(map[string]bool)

	// Collect flags from this node and its ancestors
	for n := node; n != nil; n = n.Parent {
		for _, flag := range n.Flags {
			if flag.Hidden || seen[flag.Name] {
				continue
			}
			seen[flag.Name] = true
			flags = append(flags, flagInfo{
				name:  flag.Name,
				short: flag.Short,
				help:  flag.Help,
			})
		}
	}
	return flags
}

type flagInfo struct {
	name  string
	short rune
	help  string
}

func generateBashCompletion(app *kong.Application) string {
	var sb strings.Builder

	cmds := collectCommands(app.Node, "")
	topLevelCmds := []string{}
	for _, cmd := range cmds {
		if !strings.Contains(cmd.fullName, " ") {
			topLevelCmds = append(topLevelCmds, cmd.name)
		}
	}

	globalFlags := collectFlags(app.Node)
	globalFlagNames := []string{}
	for _, f := range globalFlags {
		globalFlagNames = append(globalFlagNames, "--"+f.name)
		if f.short != 0 {
			globalFlagNames = append(globalFlagNames, "-"+string(f.short))
		}
	}

	sb.WriteString(`# bash completion for local-ai
# Generated by local-ai completion bash

_local_ai_completions()
{
    local cur prev words cword
    _init_completion || return

    local commands="` + strings.Join(topLevelCmds, " ") + `"
    local global_flags="` + strings.Join(globalFlagNames, " ") + `"

    # Find the subcommand
    local subcmd=""
    local subcmd_idx=0
    for ((i=1; i < cword; i++)); do
        case "${words[i]}" in
            -*)
                # Skip flags and their values
                ;;
            *)
                if [[ -z "$subcmd" ]]; then
                    subcmd="${words[i]}"
                    subcmd_idx=$i
                fi
                ;;
        esac
    done

    # If completing a flag value, don't suggest anything special
    if [[ "$cur" == -* ]]; then
        case "$subcmd" in
`)

	// Generate flag completions per top-level command
	for _, cmd := range cmds {
		if strings.Contains(cmd.fullName, " ") {
			continue
		}
		flags := collectFlags(cmd.node)
		flagNames := []string{}
		for _, f := range flags {
			flagNames = append(flagNames, "--"+f.name)
			if f.short != 0 {
				flagNames = append(flagNames, "-"+string(f.short))
			}
		}
		sb.WriteString(fmt.Sprintf("            %s)\n", cmd.name))
		sb.WriteString(fmt.Sprintf("                COMPREPLY=($(compgen -W \"%s\" -- \"$cur\"))\n", strings.Join(flagNames, " ")))
		sb.WriteString("                return\n")
		sb.WriteString("                ;;\n")
	}

	sb.WriteString(`            *)
                COMPREPLY=($(compgen -W "$global_flags" -- "$cur"))
                return
                ;;
        esac
    fi

    # Complete subcommands for top-level commands
    case "$subcmd" in
`)

	// Generate subcommand completions
	for _, cmd := range cmds {
		if strings.Contains(cmd.fullName, " ") {
			continue
		}
		subcmds := []string{}
		for _, sub := range cmds {
			parent, child, found := strings.Cut(sub.fullName, " ")
			if found && parent == cmd.name && !strings.Contains(child, " ") {
				subcmds = append(subcmds, child)
			}
		}
		if len(subcmds) > 0 {
			sb.WriteString(fmt.Sprintf("        %s)\n", cmd.name))
			sb.WriteString(fmt.Sprintf("            COMPREPLY=($(compgen -W \"%s\" -- \"$cur\"))\n", strings.Join(subcmds, " ")))
			sb.WriteString("            return\n")
			sb.WriteString("            ;;\n")
		}
	}

	sb.WriteString(`        "")
            COMPREPLY=($(compgen -W "$commands" -- "$cur"))
            return
            ;;
    esac
}

complete -F _local_ai_completions local-ai
`)

	return sb.String()
}

func generateZshCompletion(app *kong.Application) string {
	var sb strings.Builder

	cmds := collectCommands(app.Node, "")
	globalFlags := collectFlags(app.Node)

	sb.WriteString(`#compdef local-ai
# Generated by local-ai completion zsh

_local_ai() {
    local -a commands
    local -a global_flags

    global_flags=(
`)

	for _, f := range globalFlags {
		help := strings.ReplaceAll(f.help, "'", "'\\''")
		help = strings.ReplaceAll(help, "[", "\\[")
		help = strings.ReplaceAll(help, "]", "\\]")
		sb.WriteString(fmt.Sprintf("        '--%s[%s]'\n", f.name, help))
		if f.short != 0 {
			sb.WriteString(fmt.Sprintf("        '-%s[%s]'\n", string(f.short), help))
		}
	}

	sb.WriteString(`    )

    commands=(
`)

	for _, cmd := range cmds {
		if strings.Contains(cmd.fullName, " ") {
			continue
		}
		help := strings.ReplaceAll(cmd.help, "'", "'\\''")
		help = strings.ReplaceAll(help, "[", "\\[")
		help = strings.ReplaceAll(help, "]", "\\]")
		sb.WriteString(fmt.Sprintf("        '%s:%s'\n", cmd.name, help))
	}

	sb.WriteString(`    )

    _arguments -C \
        $global_flags \
        '1:command:->command' \
        '*::arg:->args'

    case $state in
        command)
            _describe -t commands 'local-ai commands' commands
            ;;
        args)
            case $words[1] in
`)

	// Per-command completions
	for _, cmd := range cmds {
		if strings.Contains(cmd.fullName, " ") {
			continue
		}

		sb.WriteString(fmt.Sprintf("                %s)\n", cmd.name))

		// Check for subcommands
		subcmds := []commandInfo{}
		for _, sub := range cmds {
			parent, child, found := strings.Cut(sub.fullName, " ")
			if found && parent == cmd.name && !strings.Contains(child, " ") {
				subcmds = append(subcmds, sub)
			}
		}

		if len(subcmds) > 0 {
			sb.WriteString("                    local -a subcmds\n")
			sb.WriteString("                    subcmds=(\n")
			for _, sub := range subcmds {
				_, child, _ := strings.Cut(sub.fullName, " ")
				help := strings.ReplaceAll(sub.help, "'", "'\\''")
				help = strings.ReplaceAll(help, "[", "\\[")
				help = strings.ReplaceAll(help, "]", "\\]")
				sb.WriteString(fmt.Sprintf("                        '%s:%s'\n", child, help))
			}
			sb.WriteString("                    )\n")
			sb.WriteString("                    _describe -t commands 'subcommands' subcmds\n")
		}

		flags := collectFlags(cmd.node)
		if len(flags) > 0 {
			sb.WriteString("                    _arguments \\\n")
			for i, f := range flags {
				help := strings.ReplaceAll(f.help, "'", "'\\''")
				help = strings.ReplaceAll(help, "[", "\\[")
				help = strings.ReplaceAll(help, "]", "\\]")
				suffix := " \\"
				if i == len(flags)-1 {
					suffix = ""
				}
				sb.WriteString(fmt.Sprintf("                        '--%s[%s]'%s\n", f.name, help, suffix))
			}
		}

		sb.WriteString("                    ;;\n")
	}

	sb.WriteString(`            esac
            ;;
    esac
}

_local_ai "$@"
`)

	return sb.String()
}

func generateFishCompletion(app *kong.Application) string {
	var sb strings.Builder

	cmds := collectCommands(app.Node, "")
	globalFlags := collectFlags(app.Node)

	sb.WriteString("# fish completion for local-ai\n")
	sb.WriteString("# Generated by local-ai completion fish\n\n")

	// Disable file completions by default
	sb.WriteString("complete -c local-ai -f\n\n")

	// Global flags
	for _, f := range globalFlags {
		help := strings.ReplaceAll(f.help, "'", "\\'")
		args := fmt.Sprintf("complete -c local-ai -l %s", f.name)
		if f.short != 0 {
			args += fmt.Sprintf(" -s %s", string(f.short))
		}
		args += fmt.Sprintf(" -d '%s'", help)
		sb.WriteString(args + "\n")
	}
	sb.WriteString("\n")

	// Top-level commands (no condition means they show when no subcommand is given)
	topLevelCmds := []string{}
	for _, cmd := range cmds {
		if strings.Contains(cmd.fullName, " ") {
			continue
		}
		topLevelCmds = append(topLevelCmds, cmd.name)
		help := strings.ReplaceAll(cmd.help, "'", "\\'")
		sb.WriteString(fmt.Sprintf("complete -c local-ai -n '__fish_use_subcommand' -a %s -d '%s'\n", cmd.name, help))
	}
	sb.WriteString("\n")

	// Subcommands and per-command flags
	for _, cmd := range cmds {
		if strings.Contains(cmd.fullName, " ") {
			continue
		}

		// Subcommands
		for _, sub := range cmds {
			parent, child, found := strings.Cut(sub.fullName, " ")
			if found && parent == cmd.name && !strings.Contains(child, " ") {
				help := strings.ReplaceAll(sub.help, "'", "\\'")
				sb.WriteString(fmt.Sprintf("complete -c local-ai -n '__fish_seen_subcommand_from %s' -a %s -d '%s'\n", cmd.name, child, help))
			}
		}

		// Per-command flags
		flags := collectFlags(cmd.node)
		for _, f := range flags {
			help := strings.ReplaceAll(f.help, "'", "\\'")
			args := fmt.Sprintf("complete -c local-ai -n '__fish_seen_subcommand_from %s' -l %s", cmd.name, f.name)
			if f.short != 0 {
				args += fmt.Sprintf(" -s %s", string(f.short))
			}
			args += fmt.Sprintf(" -d '%s'", help)
			sb.WriteString(args + "\n")
		}
	}

	return sb.String()
}