diff --git a/cmd/picoclaw/internal/model/add.go b/cmd/picoclaw/internal/model/add.go new file mode 100644 index 000000000..b3ebba340 --- /dev/null +++ b/cmd/picoclaw/internal/model/add.go @@ -0,0 +1,200 @@ +package model + +import ( + "bufio" + "fmt" + "io" + "strconv" + "strings" + + "github.com/spf13/cobra" + + "github.com/sipeed/picoclaw/cmd/picoclaw/internal" + "github.com/sipeed/picoclaw/pkg/config" +) + +const defaultAliasName = "custom-prefer" + +func newAddCommand() *cobra.Command { + var ( + apiBase string + apiKey string + modelID string + alias string + modelType string + ) + + cmd := &cobra.Command{ + Use: "add", + Short: "Add a model from an OpenAI-compatible endpoint", + Long: `Add a model entry by querying an OpenAI-compatible endpoint exposing +GET /models, then setting it as the default model. + +If --model is omitted, the available models are listed and you can pick one +interactively. If --model is provided, the entry is written without contacting +the server. + +Sample interactive session (key shown masked): + + $ picoclaw model add \ + -b https://ark.cn-beijing.volces.com/api/v3 \ + -k 7dff****-****-****-****-********e829 + + 115 model(s) available: + 1) doubao-lite-128k-240428 (doubao-lite-128k) + 2) doubao-pro-128k-240515 (doubao-pro-128k) + ... + 48) deepseek-r1-250120 (deepseek-r1) + 78) kimi-k2-250711 (kimi-k2) + ... + 115) doubao-seed3d-2-0-260328 (doubao-seed3d-2-0) + Pick a model (number or id): 48 + ✓ Saved model 'custom-prefer' (deepseek-r1-250120) and set as default.`, + Example: ` picoclaw model add --api-base https://api.openai.com/v1 --api-key sk-... + picoclaw model add -b http://localhost:8000/v1 -k dummy -m my-model -n local`, + Args: cobra.NoArgs, + RunE: func(cmd *cobra.Command, _ []string) error { + return runAdd(addOptions{ + apiBase: strings.TrimSpace(apiBase), + apiKey: strings.TrimSpace(apiKey), + modelID: strings.TrimSpace(modelID), + alias: strings.TrimSpace(alias), + modelType: strings.TrimSpace(modelType), + stdin: cmd.InOrStdin(), + stdout: cmd.OutOrStdout(), + }) + }, + } + + cmd.Flags().StringVarP(&apiBase, "api-base", "b", "", + "API base URL (required), e.g. https://api.openai.com/v1") + cmd.Flags().StringVarP(&apiKey, "api-key", "k", "", "API key (required)") + cmd.Flags().StringVarP(&modelID, "model", "m", "", + "Model id; when set, skips the interactive picker and the network call") + cmd.Flags().StringVarP(&alias, "name", "n", defaultAliasName, + "Local alias written to model_list and used as the default model name") + cmd.Flags().StringVar(&modelType, "type", "openai-compatible", + "Endpoint type (only 'openai-compatible' is supported today)") + _ = cmd.MarkFlagRequired("api-base") + _ = cmd.MarkFlagRequired("api-key") + + return cmd +} + +type addOptions struct { + apiBase string + apiKey string + modelID string + alias string + modelType string + stdin io.Reader + stdout io.Writer +} + +func runAdd(opt addOptions) error { + if opt.modelType != "" && opt.modelType != "openai-compatible" { + return fmt.Errorf("unsupported --type %q (only 'openai-compatible' is supported)", opt.modelType) + } + if opt.alias == "" { + opt.alias = defaultAliasName + } + + selected := opt.modelID + if selected == "" { + entries, err := fetchOpenAIModels(opt.apiBase, opt.apiKey) + if err != nil { + return fmt.Errorf("fetch models: %w", err) + } + if len(entries) == 0 { + return fmt.Errorf("no models returned by %s", opt.apiBase) + } + selected, err = pickModel(opt.stdin, opt.stdout, entries) + if err != nil { + return err + } + } + + return upsertModelDefault(opt.apiBase, opt.apiKey, opt.alias, selected, opt.stdout) +} + +func pickModel(stdin io.Reader, stdout io.Writer, entries []modelEntry) (string, error) { + fmt.Fprintf(stdout, "\n%d model(s) available:\n", len(entries)) + for i, m := range entries { + line := m.ID + if m.Name != "" && m.Name != m.ID { + line = fmt.Sprintf("%s (%s)", m.ID, m.Name) + } + fmt.Fprintf(stdout, " %3d) %s\n", i+1, line) + } + + scanner := bufio.NewScanner(stdin) + for { + fmt.Fprint(stdout, "Pick a model (number or id): ") + if !scanner.Scan() { + if err := scanner.Err(); err != nil { + return "", fmt.Errorf("read input: %w", err) + } + return "", fmt.Errorf("no selection provided") + } + text := strings.TrimSpace(scanner.Text()) + if text == "" { + continue + } + if idx, err := strconv.Atoi(text); err == nil { + if idx < 1 || idx > len(entries) { + fmt.Fprintf(stdout, "Out of range. Enter 1-%d.\n", len(entries)) + continue + } + return entries[idx-1].ID, nil + } + for _, m := range entries { + if m.ID == text { + return m.ID, nil + } + } + fmt.Fprintln(stdout, "Not a valid number or model id; try again.") + } +} + +func upsertModelDefault(apiBase, apiKey, alias, modelID string, stdout io.Writer) error { + configPath := internal.GetConfigPath() + cfg, err := config.LoadConfig(configPath) + if err != nil { + return fmt.Errorf("failed to load config: %w", err) + } + + secureKeys := config.SimpleSecureStrings(apiKey) + + found := false + for _, m := range cfg.ModelList { + if m == nil { + continue + } + if m.ModelName == alias { + m.Model = modelID + m.APIBase = apiBase + m.APIKeys = secureKeys + m.Enabled = true + found = true + break + } + } + if !found { + cfg.ModelList = append(cfg.ModelList, &config.ModelConfig{ + ModelName: alias, + Model: modelID, + APIBase: apiBase, + APIKeys: secureKeys, + Enabled: true, + }) + } + + cfg.Agents.Defaults.ModelName = alias + + if err := config.SaveConfig(configPath, cfg); err != nil { + return fmt.Errorf("failed to save config: %w", err) + } + + fmt.Fprintf(stdout, "✓ Saved model '%s' (%s) and set as default.\n", alias, modelID) + return nil +} diff --git a/cmd/picoclaw/internal/model/add_test.go b/cmd/picoclaw/internal/model/add_test.go new file mode 100644 index 000000000..5da4d5e7f --- /dev/null +++ b/cmd/picoclaw/internal/model/add_test.go @@ -0,0 +1,257 @@ +package model + +import ( + "bytes" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/sipeed/picoclaw/pkg/config" +) + +func TestFetchOpenAIModels_DataEnvelope(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, "/models", r.URL.Path) + assert.Equal(t, "Bearer secret", r.Header.Get("Authorization")) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"data":[{"id":"gpt-foo","name":"Foo"},{"id":"gpt-bar"}]}`)) + })) + defer srv.Close() + + entries, err := fetchOpenAIModels(srv.URL, "secret") + require.NoError(t, err) + require.Len(t, entries, 2) + assert.Equal(t, "gpt-foo", entries[0].ID) + assert.Equal(t, "Foo", entries[0].Name) + assert.Equal(t, "gpt-bar", entries[1].ID) +} + +func TestFetchOpenAIModels_BareArray(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`[{"id":"a"},{"id":"b"}]`)) + })) + defer srv.Close() + + entries, err := fetchOpenAIModels(srv.URL, "secret") + require.NoError(t, err) + require.Len(t, entries, 2) + assert.Equal(t, "a", entries[0].ID) + assert.Equal(t, "b", entries[1].ID) +} + +func TestFetchOpenAIModels_TrimsTrailingSlash(t *testing.T) { + var gotPath string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotPath = r.URL.Path + _, _ = w.Write([]byte(`{"data":[{"id":"x"}]}`)) + })) + defer srv.Close() + + _, err := fetchOpenAIModels(srv.URL+"/", "k") + require.NoError(t, err) + assert.Equal(t, "/models", gotPath) +} + +func TestFetchOpenAIModels_HTTPError(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + http.Error(w, "nope", http.StatusUnauthorized) + })) + defer srv.Close() + + _, err := fetchOpenAIModels(srv.URL, "bad") + require.Error(t, err) + assert.Contains(t, err.Error(), "HTTP 401") +} + +func TestFetchOpenAIModels_EmptyDataEnvelope(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, _ = w.Write([]byte(`{"data":[]}`)) + })) + defer srv.Close() + + entries, err := fetchOpenAIModels(srv.URL, "k") + require.NoError(t, err) + assert.Empty(t, entries) +} + +func TestFetchOpenAIModels_EmptyBareArray(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, _ = w.Write([]byte(`[]`)) + })) + defer srv.Close() + + entries, err := fetchOpenAIModels(srv.URL, "k") + require.NoError(t, err) + assert.Empty(t, entries) +} + +func TestFetchOpenAIModels_UnrecognizedShape(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, _ = w.Write([]byte(`{"models":"not-supported"}`)) + })) + defer srv.Close() + + _, err := fetchOpenAIModels(srv.URL, "k") + require.Error(t, err) + assert.Contains(t, err.Error(), "unrecognized shape") +} + +func TestFetchOpenAIModels_RequiresInputs(t *testing.T) { + _, err := fetchOpenAIModels("", "k") + require.Error(t, err) + assert.Contains(t, err.Error(), "api base") + + _, err = fetchOpenAIModels("https://example.com", "") + require.Error(t, err) + assert.Contains(t, err.Error(), "api key") +} + +func TestPickModel_ByIndex(t *testing.T) { + entries := []modelEntry{{ID: "a"}, {ID: "b"}, {ID: "c"}} + out := &bytes.Buffer{} + got, err := pickModel(strings.NewReader("2\n"), out, entries) + require.NoError(t, err) + assert.Equal(t, "b", got) + assert.Contains(t, out.String(), "3 model(s) available") +} + +func TestPickModel_ByID(t *testing.T) { + entries := []modelEntry{{ID: "alpha"}, {ID: "beta"}} + out := &bytes.Buffer{} + got, err := pickModel(strings.NewReader("beta\n"), out, entries) + require.NoError(t, err) + assert.Equal(t, "beta", got) +} + +func TestPickModel_RetriesOnInvalid(t *testing.T) { + entries := []modelEntry{{ID: "x"}} + out := &bytes.Buffer{} + got, err := pickModel(strings.NewReader("\n9\nnot-a-model\nx\n"), out, entries) + require.NoError(t, err) + assert.Equal(t, "x", got) + rendered := out.String() + assert.Contains(t, rendered, "Out of range") + assert.Contains(t, rendered, "Not a valid number") +} + +func TestRunAdd_WithExplicitModel_NoNetwork(t *testing.T) { + initTest(t) + + out := &bytes.Buffer{} + err := runAdd(addOptions{ + apiBase: "https://invalid.invalid/v1", + apiKey: "k", + modelID: "explicit-model", + alias: "myalias", + modelType: "openai-compatible", + stdout: out, + }) + require.NoError(t, err) + assert.Contains(t, out.String(), "Saved model 'myalias' (explicit-model)") + + cfg, err := config.LoadConfig(configPath) + require.NoError(t, err) + assert.Equal(t, "myalias", cfg.Agents.Defaults.GetModelName()) + added := findModelByName(cfg, "myalias") + require.NotNil(t, added, "expected model 'myalias' in model_list") + assert.Equal(t, "explicit-model", added.Model) + assert.Equal(t, "https://invalid.invalid/v1", added.APIBase) + assert.True(t, added.Enabled) + require.Len(t, added.APIKeys, 1) + assert.Equal(t, "k", added.APIKeys[0].String()) +} + +func findModelByName(cfg *config.Config, name string) *config.ModelConfig { + for _, m := range cfg.ModelList { + if m != nil && m.ModelName == name { + return m + } + } + return nil +} + +func TestRunAdd_FetchAndPick(t *testing.T) { + initTest(t) + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, "Bearer my-key", r.Header.Get("Authorization")) + _, _ = w.Write([]byte(`{"data":[{"id":"m1"},{"id":"m2"}]}`)) + })) + defer srv.Close() + + out := &bytes.Buffer{} + err := runAdd(addOptions{ + apiBase: srv.URL, + apiKey: "my-key", + alias: defaultAliasName, + modelType: "openai-compatible", + stdin: strings.NewReader("2\n"), + stdout: out, + }) + require.NoError(t, err) + + cfg, err := config.LoadConfig(configPath) + require.NoError(t, err) + assert.Equal(t, defaultAliasName, cfg.Agents.Defaults.GetModelName()) + added := findModelByName(cfg, defaultAliasName) + require.NotNil(t, added) + assert.Equal(t, "m2", added.Model) +} + +func TestRunAdd_UpsertsExistingAlias(t *testing.T) { + initTest(t) + + first := &bytes.Buffer{} + require.NoError(t, runAdd(addOptions{ + apiBase: "https://a.example/v1", + apiKey: "k1", + modelID: "m1", + alias: "shared", + stdout: first, + })) + + second := &bytes.Buffer{} + require.NoError(t, runAdd(addOptions{ + apiBase: "https://b.example/v1", + apiKey: "k2", + modelID: "m2", + alias: "shared", + stdout: second, + })) + + cfg, err := config.LoadConfig(configPath) + require.NoError(t, err) + matches := 0 + for _, m := range cfg.ModelList { + if m != nil && m.ModelName == "shared" { + matches++ + } + } + assert.Equal(t, 1, matches, "alias should be updated, not duplicated") + + updated := findModelByName(cfg, "shared") + require.NotNil(t, updated) + assert.Equal(t, "m2", updated.Model) + assert.Equal(t, "https://b.example/v1", updated.APIBase) + assert.Equal(t, "k2", updated.APIKeys[0].String()) +} + +func TestRunAdd_RejectsUnsupportedType(t *testing.T) { + initTest(t) + + err := runAdd(addOptions{ + apiBase: "https://x/v1", + apiKey: "k", + modelID: "m", + alias: "a", + modelType: "anthropic", + stdout: &bytes.Buffer{}, + }) + require.Error(t, err) + assert.Contains(t, err.Error(), "unsupported --type") +} diff --git a/cmd/picoclaw/internal/model/command.go b/cmd/picoclaw/internal/model/command.go index 330734b82..c412993a0 100644 --- a/cmd/picoclaw/internal/model/command.go +++ b/cmd/picoclaw/internal/model/command.go @@ -21,11 +21,17 @@ func NewModelCommand() *cobra.Command { If no argument is provided, shows the current default model. If a model name is provided, sets it as the default model. +To onboard a model from a custom OpenAI-compatible endpoint (fetch the +available list online and pick one), use the 'add' subcommand: + + picoclaw model add --help + Examples: picoclaw model # Show current default model picoclaw model gpt-5.2 # Set gpt-5.2 as default picoclaw model claude-sonnet-4.6 # Set claude-sonnet-4.6 as default picoclaw model local-model # Set local VLLM server as default + picoclaw model add -b URL -k KEY # Add a model from a custom endpoint Note: 'local-model' is a special value for using a local VLLM server (running at localhost:8000 by default) which does not require an API key.`, @@ -51,6 +57,8 @@ Note: 'local-model' is a special value for using a local VLLM server }, } + cmd.AddCommand(newAddCommand()) + return cmd } @@ -66,6 +74,9 @@ func showCurrentModel(cfg *config.Config) { fmt.Println("\nAvailable models in your config:") listAvailableModels(cfg) } + + fmt.Println("\nTip: 'picoclaw model add -b URL -k KEY' adds a model from a custom") + fmt.Println(" OpenAI-compatible endpoint (see 'picoclaw model add --help').") } func listAvailableModels(cfg *config.Config) { diff --git a/cmd/picoclaw/internal/model/online.go b/cmd/picoclaw/internal/model/online.go new file mode 100644 index 000000000..9b8f7811d --- /dev/null +++ b/cmd/picoclaw/internal/model/online.go @@ -0,0 +1,77 @@ +package model + +import ( + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + "time" +) + +type modelEntry struct { + ID string `json:"id"` + Name string `json:"name"` + Description string `json:"description"` +} + +type modelsAPIResponse struct { + Data []modelEntry `json:"data"` +} + +// fetchOpenAIModels GETs /models with Bearer auth and accepts both the +// {data:[…]} envelope and a bare array shape used by various OpenAI-compatible servers. +func fetchOpenAIModels(baseURL, apiKey string) ([]modelEntry, error) { + if strings.TrimSpace(baseURL) == "" { + return nil, fmt.Errorf("api base is required") + } + if strings.TrimSpace(apiKey) == "" { + return nil, fmt.Errorf("api key is required") + } + + url := strings.TrimRight(baseURL, "/") + "/models" + + client := &http.Client{Timeout: 15 * time.Second} + req, err := http.NewRequest(http.MethodGet, url, nil) + if err != nil { + return nil, fmt.Errorf("build request: %w", err) + } + req.Header.Set("Authorization", "Bearer "+apiKey) + + resp, err := client.Do(req) + if err != nil { + return nil, fmt.Errorf("request failed: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + body, _ := io.ReadAll(io.LimitReader(resp.Body, 512)) + return nil, fmt.Errorf("HTTP %d: %s", resp.StatusCode, strings.TrimSpace(string(body))) + } + + body, err := io.ReadAll(resp.Body) + if err != nil { + return nil, fmt.Errorf("read response: %w", err) + } + + // {"data": [...]} envelope. Distinguish "envelope shape with empty list" + // from "object without a data key" via Data being non-nil after unmarshal: + // json.Unmarshal sets Data to []modelEntry{} for `{"data":[]}` but leaves + // it as nil when "data" is absent or null. + var envelope modelsAPIResponse + if err := json.Unmarshal(body, &envelope); err == nil && envelope.Data != nil { + return envelope.Data, nil + } + + // Bare-array shape, including `[]`. + var arr []modelEntry + if err := json.Unmarshal(body, &arr); err == nil { + return arr, nil + } + + preview := body + if len(preview) > 256 { + preview = preview[:256] + } + return nil, fmt.Errorf("decode response: unrecognized shape: %s", strings.TrimSpace(string(preview))) +}