diff --git a/examples/extensions/README.md b/examples/extensions/README.md index 84d5ee91..84db8d5b 100644 --- a/examples/extensions/README.md +++ b/examples/extensions/README.md @@ -74,6 +74,7 @@ kit install github.com/mark3labs/kit/examples/extensions --local | `summarize.go` | Summarize conversation | Custom tool with parameters | | `confirm-destructive.go` | Confirm destructive commands | `OnToolCall` blocking | | `lsp-diagnostics.go` | LSP integration | Complex extension, external process | +| `preferred-models.go` | Switch between the models that you list in an option | `RegisterCommand` with `Complete`, `RegisterOption`, `PromptSelect`, `SetModel`, `GetCurrentProvider` | ### Subagents & Background Tasks diff --git a/examples/extensions/preferred-models.go b/examples/extensions/preferred-models.go new file mode 100644 index 00000000..a36c32e5 --- /dev/null +++ b/examples/extensions/preferred-models.go @@ -0,0 +1,120 @@ +//go:build ignore + +package main + +import ( + "fmt" + "strconv" + "strings" + + "kit/ext" +) + +// preferredModels reads the option preferred-models: provider/model strings, +// separated by commas or spaces. +func preferredModels(ctx ext.Context) []string { + return strings.Fields(strings.ReplaceAll(ctx.GetOption("preferred-models"), ",", " ")) +} + +// currentModel returns the model in use as provider/model. ctx.Model is not +// enough: until the first model switch it is only the model ID, and a model +// ID can contain "/", as in openrouter/anthropic/claude-opus-5. +func currentModel(ctx ext.Context) string { + return ctx.GetCurrentProvider() + "/" + ctx.GetCurrentModelID() +} + +// pickModel returns the model that arg selects. arg is the number of the +// model in the list, starting at 1, or a part of a model name. A part of a +// name selects the first model that contains it. pickModel returns "" if +// nothing matches. +func pickModel(list []string, arg string) string { + if n, err := strconv.Atoi(arg); err == nil { + if n >= 1 && n <= len(list) { + return list[n-1] + } + return "" + } + for _, m := range list { + if strings.Contains(strings.ToLower(m), strings.ToLower(arg)) { + return m + } + } + return "" +} + +// switchTo switches to model and returns the message for the transcript. +func switchTo(ctx ext.Context, model string) (string, error) { + if model == currentModel(ctx) { + return "Already using " + model, nil + } + if err := ctx.SetModel(model); err != nil { + return "", err + } + return "Switched to " + model, nil +} + +// Init adds the command /preferred. It switches between the models in the +// option preferred-models, which lists provider/model strings separated by +// commas or spaces. Write the list as one string; a YAML list is not read. +// +// options: +// preferred-models: anthropic/claude-opus-5-5, amazon-bedrock/eu.anthropic.claude-opus-5-5 +// +// or for one run: KIT_OPT_PREFERRED_MODELS="ollama/llama3 ollama/qwen3" kit +// +// /preferred select a model from the list +// /preferred 2 switch to the second model in the list +// /preferred bedrock switch to the first model whose name contains "bedrock" +// +// Unlike /model, the list contains only the models that you put in it. +// +// Usage: kit -e examples/extensions/preferred-models.go +func Init(api ext.API) { + api.RegisterOption(ext.OptionDef{ + Name: "preferred-models", + Description: "Models for /preferred, as provider/model, separated by commas", + }) + + api.RegisterCommand(ext.CommandDef{ + Name: "preferred", + Description: "Switch to one of your preferred models", + Execute: func(args string, ctx ext.Context) (string, error) { + list := preferredModels(ctx) + if len(list) == 0 { + return "No preferred models. List them in the option preferred-models, for example:\n options:\n preferred-models: anthropic/claude-opus-5-5, anthropic/claude-sonnet-5-5", nil + } + if arg := strings.TrimSpace(args); arg != "" { + model := pickModel(list, arg) + if model == "" { + return "", fmt.Errorf("no preferred model matches %q", arg) + } + return switchTo(ctx, model) + } + current := currentModel(ctx) + options := make([]string, len(list)) + for i, m := range list { + options[i] = fmt.Sprintf("%d %s", i+1, m) + if m == current { + options[i] += " (current)" + } + } + res := ctx.PromptSelect(ext.PromptSelectConfig{ + Message: "Switch to a preferred model", + Options: options, + }) + if res.Cancelled || res.Index < 0 || res.Index >= len(list) { + return "", nil + } + return switchTo(ctx, list[res.Index]) + }, + Complete: func(prefix string, ctx ext.Context) []string { + var out []string + for _, m := range preferredModels(ctx) { + if strings.HasPrefix(m, prefix) { + out = append(out, m) + } + } + return out + }, + }) +} diff --git a/pkg/extensions/test/preferred_models_yaegi_test.go b/pkg/extensions/test/preferred_models_yaegi_test.go new file mode 100644 index 00000000..8c2a1d65 --- /dev/null +++ b/pkg/extensions/test/preferred_models_yaegi_test.go @@ -0,0 +1,209 @@ +package test + +import ( + "reflect" + "strings" + "testing" + + "github.com/mark3labs/kit/pkg/extensions" +) + +// preferredList is the option value for the preferred-models tests. The +// model ID of the first model is also the name of the second model. +const preferredList = "openrouter/anthropic/claude-opus-5, anthropic/claude-opus-5 ollama/qwen3" + +// loadPreferred loads the preferred-models example and returns its command +// /preferred. +func loadPreferred(t *testing.T) (*Harness, extensions.CommandDef) { + t.Helper() + harness := New(t) + harness.LoadFile("../../../examples/extensions/preferred-models.go") + for _, cmd := range harness.RegisteredCommands() { + if cmd.Name == "preferred" && cmd.Execute != nil && cmd.Complete != nil { + return harness, cmd + } + } + t.Fatal("the example does not register /preferred with Execute and Complete") + return nil, extensions.CommandDef{} +} + +// preferredKit gives /preferred a context as Kit does. active is the model +// in use. GetCurrentProvider and GetCurrentModelID report it, and SetModel +// switches it. ctx.Model is only the model ID until the first switch, and +// provider/model after it. switched records the models given to SetModel. +type preferredKit struct { + ctx extensions.Context + active string + switched []string +} + +// newPreferredKit returns a preferredKit with active in use. model is the +// value of ctx.Model. +func newPreferredKit(harness *Harness, active, model string) *preferredKit { + k := &preferredKit{ctx: harness.Context().ToContext(), active: active} + k.ctx.Model = model + k.ctx.GetOption = func(name string) string { + if name == "preferred-models" { + return preferredList + } + return "" + } + k.ctx.GetCurrentProvider = func() string { + provider, _, _ := strings.Cut(k.active, "/") + return provider + } + k.ctx.GetCurrentModelID = func() string { + _, id, _ := strings.Cut(k.active, "/") + return id + } + k.ctx.SetModel = func(model string) error { + k.switched = append(k.switched, model) + k.active = model + k.ctx.Model = model + return nil + } + return k +} + +// picker opens the picker with /preferred, selects the option pick (-1 for +// Esc), and returns the options and the output. +func (k *preferredKit) picker(t *testing.T, cmd extensions.CommandDef, pick int) ([]string, string) { + t.Helper() + var shown []string + k.ctx.PromptSelect = func(cfg extensions.PromptSelectConfig) extensions.PromptSelectResult { + shown = cfg.Options + if pick < 0 { + return extensions.PromptSelectResult{Cancelled: true} + } + return extensions.PromptSelectResult{Value: cfg.Options[pick], Index: pick} + } + out, err := cmd.Execute("", k.ctx) + if err != nil { + t.Fatalf("/preferred: %v", err) + } + return shown, out +} + +var ( + // markFirst and markSecond are the options of the picker with the + // first or the second model in use. + markFirst = []string{ + "1 openrouter/anthropic/claude-opus-5 (current)", + "2 anthropic/claude-opus-5", + "3 ollama/qwen3", + } + markSecond = []string{ + "1 openrouter/anthropic/claude-opus-5", + "2 anthropic/claude-opus-5 (current)", + "3 ollama/qwen3", + } +) + +// TestPreferredModelsExample runs /preferred from the start. The configured +// model is openrouter/anthropic/claude-opus-5, so ctx.Model starts as +// anthropic/claude-opus-5, the name of the second model. +func TestPreferredModelsExample(t *testing.T) { + harness, cmd := loadPreferred(t) + k := newPreferredKit(harness, "openrouter/anthropic/claude-opus-5", "anthropic/claude-opus-5") + + if shown, out := k.picker(t, cmd, -1); out != "" || !reflect.DeepEqual(shown, markFirst) { + t.Fatalf("picker at the start = %q, output %q; want %q and no output", shown, out, markFirst) + } + for _, c := range []struct{ args, want string }{ + {"2", "Switched to anthropic/claude-opus-5"}, + {"1", "Switched to openrouter/anthropic/claude-opus-5"}, + {"2", "Switched to anthropic/claude-opus-5"}, + {"2", "Already using anthropic/claude-opus-5"}, + } { + if out, err := cmd.Execute(c.args, k.ctx); out != c.want || err != nil { + t.Fatalf("/preferred %s = %q, %v; want %q", c.args, out, err, c.want) + } + } + if shown, _ := k.picker(t, cmd, -1); !reflect.DeepEqual(shown, markSecond) { + t.Fatalf("picker after a switch = %q, want %q", shown, markSecond) + } + if _, out := k.picker(t, cmd, 2); out != "Switched to ollama/qwen3" { + t.Fatalf("picker, third option = %q, want the switch", out) + } + + // A part of a name selects the first model that contains it. The match + // ignores case. + for _, c := range []struct{ args, want string }{ + {"OPENROUTER", "Switched to openrouter/anthropic/claude-opus-5"}, + {"openrouter", "Already using openrouter/anthropic/claude-opus-5"}, + } { + if out, err := cmd.Execute(c.args, k.ctx); out != c.want || err != nil { + t.Errorf("/preferred %s = %q, %v; want %q", c.args, out, err, c.want) + } + } + for _, args := range []string{"0", "4", "gpt"} { + if _, err := cmd.Execute(args, k.ctx); err == nil || !strings.Contains(err.Error(), "no preferred model matches") { + t.Errorf("/preferred %s gave error %v, want no match", args, err) + } + } + want := []string{ + "anthropic/claude-opus-5", + "openrouter/anthropic/claude-opus-5", + "anthropic/claude-opus-5", + "ollama/qwen3", + "openrouter/anthropic/claude-opus-5", + } + if !reflect.DeepEqual(k.switched, want) { + t.Errorf("SetModel calls = %q, want %q", k.switched, want) + } + + if got := cmd.Complete("anthropic/", k.ctx); !reflect.DeepEqual(got, []string{"anthropic/claude-opus-5"}) { + t.Errorf("Complete(anthropic/) = %q, want [anthropic/claude-opus-5]", got) + } + if got := cmd.Complete("", k.ctx); len(got) != 3 { + t.Errorf("Complete() = %q, want all three models", got) + } +} + +// TestPreferredModelsExampleAfterReload loads the example after a switch to +// anthropic/claude-opus-5, as /reload-ext does. ctx.Model is then the same +// string as at the start with openrouter/anthropic/claude-opus-5. +func TestPreferredModelsExampleAfterReload(t *testing.T) { + harness, cmd := loadPreferred(t) + k := newPreferredKit(harness, "anthropic/claude-opus-5", "anthropic/claude-opus-5") + + if shown, _ := k.picker(t, cmd, -1); !reflect.DeepEqual(shown, markSecond) { + t.Fatalf("picker after the reload = %q, want %q", shown, markSecond) + } + if out, err := cmd.Execute("1", k.ctx); out != "Switched to openrouter/anthropic/claude-opus-5" || err != nil { + t.Fatalf("/preferred 1 after the reload = %q, %v; want the switch", out, err) + } +} + +// TestPreferredModelsExampleIgnoresInvalidPickerIndex checks that an index +// outside the list switches nothing. +func TestPreferredModelsExampleIgnoresInvalidPickerIndex(t *testing.T) { + harness, cmd := loadPreferred(t) + k := newPreferredKit(harness, "ollama/qwen3", "ollama/qwen3") + for _, pick := range []int{-2, 3, 5} { + k.ctx.PromptSelect = func(extensions.PromptSelectConfig) extensions.PromptSelectResult { + return extensions.PromptSelectResult{Index: pick} + } + if out, err := cmd.Execute("", k.ctx); out != "" || err != nil { + t.Errorf("picker index %d = %q, %v; want no output", pick, out, err) + } + } + if len(k.switched) != 0 { + t.Errorf("SetModel calls = %q, want none", k.switched) + } +} + +// TestPreferredModelsExampleWithoutList checks that /preferred explains the +// option when it is not set, and opens no picker. +func TestPreferredModelsExampleWithoutList(t *testing.T) { + harness, cmd := loadPreferred(t) + ctx := harness.Context().ToContext() + ctx.PromptSelect = func(extensions.PromptSelectConfig) extensions.PromptSelectResult { + t.Error("the picker opened without a list") + return extensions.PromptSelectResult{Cancelled: true} + } + out, err := cmd.Execute("", ctx) + if err != nil || !strings.Contains(out, "No preferred models") || !strings.Contains(out, "preferred-models:") { + t.Errorf("/preferred without the option = %q, %v; want a hint with the option", out, err) + } +} diff --git a/www/pages/extensions/examples.md b/www/pages/extensions/examples.md index 123c4ccb..6f245573 100644 --- a/www/pages/extensions/examples.md +++ b/www/pages/extensions/examples.md @@ -56,6 +56,7 @@ environment restrictions when input is untrusted. | [`inline-bash.go`](https://github.com/mark3labs/kit/blob/master/examples/extensions/inline-bash.go) | Inline bash execution | | [`interactive-shell.go`](https://github.com/mark3labs/kit/blob/master/examples/extensions/interactive-shell.go) | Interactive shell integration | | [`notify.go`](https://github.com/mark3labs/kit/blob/master/examples/extensions/notify.go) | Desktop notifications | +| [`preferred-models.go`](https://github.com/mark3labs/kit/blob/master/examples/extensions/preferred-models.go) | `/preferred` switches between the models that you list in the option `preferred-models` | ## Agent and context