diff --git a/cmd/climc/shell/aiproxy/ai_provider.go b/cmd/climc/shell/aiproxy/ai_provider.go index b6eb329220..e1487a429a 100644 --- a/cmd/climc/shell/aiproxy/ai_provider.go +++ b/cmd/climc/shell/aiproxy/ai_provider.go @@ -17,6 +17,7 @@ package aiproxy import ( "yunion.io/x/onecloud/cmd/climc/shell" apmodules "yunion.io/x/onecloud/pkg/mcclient/modules/aiproxy" + "yunion.io/x/onecloud/pkg/mcclient/options" apoptions "yunion.io/x/onecloud/pkg/mcclient/options/aiproxy" ) @@ -28,5 +29,7 @@ func init() { cmd.Update(new(apoptions.AiProviderUpdateOptions)) cmd.Delete(new(apoptions.AiProviderDeleteOptions)) cmd.PerformClass("test-connectivity", new(apoptions.AiProviderTestConnectivityOptions)) + cmd.Perform("test-connectivity", new(options.BaseIdOptions)) + cmd.Perform("set-models", new(apoptions.AiProviderSetModelsOptions)) registerEnableDisable(cmd) } diff --git a/pkg/aiproxy/models/ai_models.go b/pkg/aiproxy/models/ai_models.go index 241d3a9e70..edb485f9a8 100644 --- a/pkg/aiproxy/models/ai_models.go +++ b/pkg/aiproxy/models/ai_models.go @@ -187,6 +187,7 @@ func (manager *SAiModelManager) FetchCustomizeColumns( for i := range objs { rows[i].VirtualResourceDetails = baseRows[i] m := objs[i].(*SAiModel) + rows[i].ContextWindow = CatalogContextWindow(m.ModelKey) providerIds[i] = m.AiProviderId if vid := strings.TrimSpace(m.VisualProviderId); vid != "" { visualProviderIds = append(visualProviderIds, vid) diff --git a/pkg/aiproxy/models/ai_providers.go b/pkg/aiproxy/models/ai_providers.go index 82eb485c2f..41717e0a9d 100644 --- a/pkg/aiproxy/models/ai_providers.go +++ b/pkg/aiproxy/models/ai_providers.go @@ -127,6 +127,26 @@ func (p *SAiProvider) PerformDisable(ctx context.Context, userCred mcclient.Toke return nil, nil } +// PerformSetModels imports selected catalog models. Existing model_key rows are skipped. +func (p *SAiProvider) PerformSetModels( + ctx context.Context, + userCred mcclient.TokenCredential, + query jsonutils.JSONObject, + input api.AiProviderSetModelsInput, +) (jsonutils.JSONObject, error) { + modelKeys, err := normalizeProviderModelKeys(input.ModelKeys) + if err != nil { + return nil, err + } + if len(modelKeys) == 0 { + return nil, errors.Wrap(httperrors.ErrInputParameter, "model_keys is required") + } + if err := createSelectedProviderModels(ctx, userCred, p.GetOwnerId(), p, modelKeys); err != nil { + return nil, err + } + return nil, nil +} + func (manager *SAiProviderManager) FetchCustomizeColumns( ctx context.Context, userCred mcclient.TokenCredential, diff --git a/pkg/aiproxy/models/ai_providers_test.go b/pkg/aiproxy/models/ai_providers_test.go index 4f89537f69..57705263df 100644 --- a/pkg/aiproxy/models/ai_providers_test.go +++ b/pkg/aiproxy/models/ai_providers_test.go @@ -126,3 +126,34 @@ func TestDeleteAiModelsByProviderIdRequiresProviderId(t *testing.T) { t.Fatal("expected error for empty provider id") } } + +func TestResolveProviderSecretForConnectivityEmptyId(t *testing.T) { + _, err := resolveProviderSecretForConnectivity(&SAiProvider{}) + if err == nil { + t.Fatal("expected error for empty provider id") + } +} + +func TestPerformSetModelsRequiresModelKeys(t *testing.T) { + p := &SAiProvider{} + _, err := p.PerformSetModels(context.Background(), nil, nil, api.AiProviderSetModelsInput{}) + if err == nil { + t.Fatal("expected error for empty model_keys") + } +} + +func TestCreateSelectedProviderModelsNilProvider(t *testing.T) { + if err := createSelectedProviderModels(context.Background(), nil, nil, nil, []string{"gpt-4"}); err != nil { + t.Fatalf("nil provider should no-op: %v", err) + } +} + +func TestNormalizeProviderModelKeys(t *testing.T) { + got, err := normalizeProviderModelKeys([]string{" gpt-4 ", "gpt-4", "deepseek-v4-pro"}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(got) != 2 { + t.Fatalf("got %v, want 2 unique keys", got) + } +} diff --git a/pkg/aiproxy/models/catalog_seed_models.go b/pkg/aiproxy/models/catalog_seed_models.go index 8479cb9956..5de6322dbe 100644 --- a/pkg/aiproxy/models/catalog_seed_models.go +++ b/pkg/aiproxy/models/catalog_seed_models.go @@ -20,11 +20,19 @@ import ( api "yunion.io/x/onecloud/pkg/apis/aiproxy" ) +const ( + catalogContextWindow1M = 1_000_000 + catalogContextWindow1050K = 1_050_000 + catalogContextWindow2M = 2_000_000 +) + // catalogSeedModel is a known upstream model id for built-in provider_key values. // ModelKey is the id sent to the upstream API (no "provider/" prefix). +// ContextWindow is 0 when unknown. type catalogSeedModel struct { - ModelKey string - Description string + ModelKey string + Description string + ContextWindow int } // catalogSeedModelsForProvider returns known public model ids for a provider_key. @@ -34,8 +42,8 @@ func catalogSeedModelsForProvider(providerKey string) []catalogSeedModel { switch providerKey { case api.ProviderKeyAnthropic: return []catalogSeedModel{ - {ModelKey: "claude-opus-4-20250514", Description: "Anthropic Claude Opus 4"}, - {ModelKey: "claude-sonnet-4-20250514", Description: "Anthropic Claude Sonnet 4"}, + {ModelKey: "claude-opus-4-20250514", Description: "Anthropic Claude Opus 4", ContextWindow: catalogContextWindow1M}, + {ModelKey: "claude-sonnet-4-20250514", Description: "Anthropic Claude Sonnet 4", ContextWindow: catalogContextWindow1M}, {ModelKey: "claude-3-7-sonnet-20250219", Description: "Anthropic Claude 3.7 Sonnet"}, {ModelKey: "claude-3-5-sonnet-20241022", Description: "Anthropic Claude 3.5 Sonnet"}, {ModelKey: "claude-3-5-haiku-20241022", Description: "Anthropic Claude 3.5 Haiku"}, @@ -85,8 +93,8 @@ func catalogSeedModelsForProvider(providerKey string) []catalogSeedModel { // } case api.ProviderKeyDeepseek: return []catalogSeedModel{ - {ModelKey: "deepseek-v4-flash", Description: "DeepSeek-V4-Flash; 1M context; high concurrency (2500); cost-efficient default"}, - {ModelKey: "deepseek-v4-pro", Description: "DeepSeek-V4-Pro; 1M context; frontier reasoning/coding/agents (500 concurrency)"}, + {ModelKey: "deepseek-v4-flash", Description: "DeepSeek-V4-Flash; 1M context; high concurrency (2500); cost-efficient default", ContextWindow: catalogContextWindow1M}, + {ModelKey: "deepseek-v4-pro", Description: "DeepSeek-V4-Pro; 1M context; frontier reasoning/coding/agents (500 concurrency)", ContextWindow: catalogContextWindow1M}, } // disabled: uncommon provider // case api.ProviderKeyElevenlabs: @@ -106,11 +114,11 @@ func catalogSeedModelsForProvider(providerKey string) []catalogSeedModel { // } case api.ProviderKeyGemini: return []catalogSeedModel{ - {ModelKey: "gemini-2.0-flash", Description: "Google Gemini 2.0 Flash"}, - {ModelKey: "gemini-2.0-flash-lite", Description: "Google Gemini 2.0 Flash-Lite"}, - {ModelKey: "gemini-1.5-pro", Description: "Google Gemini 1.5 Pro"}, - {ModelKey: "gemini-1.5-flash", Description: "Google Gemini 1.5 Flash"}, - {ModelKey: "gemini-1.5-flash-8b", Description: "Google Gemini 1.5 Flash 8B"}, + {ModelKey: "gemini-2.0-flash", Description: "Google Gemini 2.0 Flash", ContextWindow: catalogContextWindow1M}, + {ModelKey: "gemini-2.0-flash-lite", Description: "Google Gemini 2.0 Flash-Lite", ContextWindow: catalogContextWindow1M}, + {ModelKey: "gemini-1.5-pro", Description: "Google Gemini 1.5 Pro", ContextWindow: catalogContextWindow2M}, + {ModelKey: "gemini-1.5-flash", Description: "Google Gemini 1.5 Flash", ContextWindow: catalogContextWindow1M}, + {ModelKey: "gemini-1.5-flash-8b", Description: "Google Gemini 1.5 Flash 8B", ContextWindow: catalogContextWindow1M}, {ModelKey: "gemini-embedding-001", Description: "Google Gemini Embedding 001"}, } case api.ProviderKeyGroq: @@ -175,9 +183,9 @@ func catalogSeedModelsForProvider(providerKey string) []catalogSeedModel { {ModelKey: "gpt-5.2", Description: "OpenAI GPT-5.2"}, {ModelKey: "gpt-5.2-pro", Description: "OpenAI GPT-5.2 pro"}, {ModelKey: "gpt-5.2-codex", Description: "OpenAI GPT-5.2 Codex"}, - {ModelKey: "gpt-4.1", Description: "OpenAI GPT-4.1"}, - {ModelKey: "gpt-4.1-mini", Description: "OpenAI GPT-4.1 mini"}, - {ModelKey: "gpt-4.1-nano", Description: "OpenAI GPT-4.1 nano"}, + {ModelKey: "gpt-4.1", Description: "OpenAI GPT-4.1", ContextWindow: catalogContextWindow1M}, + {ModelKey: "gpt-4.1-mini", Description: "OpenAI GPT-4.1 mini", ContextWindow: catalogContextWindow1M}, + {ModelKey: "gpt-4.1-nano", Description: "OpenAI GPT-4.1 nano", ContextWindow: catalogContextWindow1M}, {ModelKey: "gpt-4o", Description: "OpenAI GPT-4o"}, {ModelKey: "gpt-4o-mini", Description: "OpenAI GPT-4o mini"}, {ModelKey: "chatgpt-4o-latest", Description: "OpenAI ChatGPT-4o latest"}, @@ -200,7 +208,7 @@ func catalogSeedModelsForProvider(providerKey string) []catalogSeedModel { {ModelKey: "openai/gpt-4o-mini", Description: "OpenRouter OpenAI GPT-4o mini"}, {ModelKey: "anthropic/claude-3.5-sonnet", Description: "OpenRouter Claude 3.5 Sonnet"}, {ModelKey: "anthropic/claude-3.5-haiku", Description: "OpenRouter Claude 3.5 Haiku"}, - {ModelKey: "google/gemini-2.0-flash-001", Description: "OpenRouter Gemini 2.0 Flash"}, + {ModelKey: "google/gemini-2.0-flash-001", Description: "OpenRouter Gemini 2.0 Flash", ContextWindow: catalogContextWindow1M}, {ModelKey: "meta-llama/llama-3.3-70b-instruct", Description: "OpenRouter Llama 3.3 70B Instruct"}, {ModelKey: "mistralai/mistral-large", Description: "OpenRouter Mistral Large"}, } @@ -306,9 +314,9 @@ func baiduErnieSeedModels() []catalogSeedModel { func xiaomiMimoSeedModels() []catalogSeedModel { return []catalogSeedModel{ - {ModelKey: "mimo-v2.5-pro", Description: "Xiaomi MiMo 2.5 Pro (flagship text)"}, + {ModelKey: "mimo-v2.5-pro", Description: "Xiaomi MiMo 2.5 Pro (flagship text)", ContextWindow: catalogContextWindow1M}, {ModelKey: "mimo-v2-pro", Description: "Xiaomi MiMo 2 Pro"}, - {ModelKey: "mimo-v2.5", Description: "Xiaomi MiMo 2.5 (multimodal text)"}, + {ModelKey: "mimo-v2.5", Description: "Xiaomi MiMo 2.5 (multimodal text)", ContextWindow: catalogContextWindow1M}, {ModelKey: "mimo-v2-omni", Description: "Xiaomi MiMo 2 Omni (multimodal)"}, {ModelKey: "mimo-v2-flash", Description: "Xiaomi MiMo 2 Flash (fast)"}, } @@ -327,6 +335,7 @@ var catalogSeedProviderKeys = []string{ api.ProviderKeyOpenrouter, api.ProviderKeyXiaomi, api.ProviderKeyMoonshot, + api.ProviderKeyZhipu, } // CatalogSeedDescription returns a known description for modelKey from built-in seed catalogs. @@ -345,6 +354,77 @@ func CatalogSeedDescription(modelKey string) string { return "" } +// extraCatalogContextWindows covers official 1M+ models that are not yet in seed lists. +// Keys are lowercase model ids (and last-path-segment aliases). +var extraCatalogContextWindows = map[string]int{ + "deepseek-chat": catalogContextWindow1M, + "deepseek-reasoner": catalogContextWindow1M, + "glm-5.3": catalogContextWindow1M, + "glm-5.3-flash": catalogContextWindow1M, + "kimi-k3": catalogContextWindow1M, + "gpt-5.4": catalogContextWindow1050K, + "gpt-5.4-pro": catalogContextWindow1050K, + "gpt-5.6": catalogContextWindow1050K, + "gpt-5.6-sol": catalogContextWindow1050K, + "gpt-5.6-terra": catalogContextWindow1050K, + "gpt-5.6-luna": catalogContextWindow1050K, + "gemini-2.5-pro": catalogContextWindow1M, + "gemini-2.5-flash": catalogContextWindow1M, + "gemini-2.5-flash-lite": catalogContextWindow1M, + "gemini-2.0-flash-001": catalogContextWindow1M, + "claude-opus-5": catalogContextWindow1M, + "claude-sonnet-5": catalogContextWindow1M, + "claude-fable-5": catalogContextWindow1M, + "claude-fable-5-1": catalogContextWindow1M, +} + +func catalogModelKeyLeaf(modelKey string) string { + modelKey = strings.TrimSpace(modelKey) + if i := strings.LastIndex(modelKey, "/"); i >= 0 { + return modelKey[i+1:] + } + return modelKey +} + +func extraCatalogContextWindow(modelKey string) int { + if n, ok := extraCatalogContextWindows[strings.ToLower(strings.TrimSpace(modelKey))]; ok { + return n + } + return 0 +} + +// CatalogContextWindow returns the known context window for modelKey from the +// built-in catalog (seed rows plus extra official ids not yet seeded). 0 means unknown. +func CatalogContextWindow(modelKey string) int { + modelKey = strings.TrimSpace(modelKey) + if modelKey == "" { + return 0 + } + if n := extraCatalogContextWindow(modelKey); n > 0 { + return n + } + leaf := catalogModelKeyLeaf(modelKey) + if leaf != modelKey { + if n := extraCatalogContextWindow(leaf); n > 0 { + return n + } + } + for _, providerKey := range catalogSeedProviderKeys { + for _, item := range catalogSeedModelsForProvider(providerKey) { + if item.ContextWindow <= 0 { + continue + } + if strings.EqualFold(item.ModelKey, modelKey) { + return item.ContextWindow + } + if strings.EqualFold(catalogModelKeyLeaf(item.ModelKey), leaf) { + return item.ContextWindow + } + } + } + return 0 +} + func moonshotKimiSeedModels() []catalogSeedModel { return []catalogSeedModel{ {ModelKey: "kimi-k2.7-code", Description: "Moonshot Kimi K2.7 Code"}, @@ -363,7 +443,7 @@ func moonshotKimiSeedModels() []catalogSeedModel { func zhipuGLMSeedModels() []catalogSeedModel { return []catalogSeedModel{ - {ModelKey: "glm-5.2", Description: "Z.AI GLM-5.2 flagship; 1M context; long-horizon agents"}, + {ModelKey: "glm-5.2", Description: "Z.AI GLM-5.2 flagship; 1M context; long-horizon agents", ContextWindow: catalogContextWindow1M}, {ModelKey: "glm-5.1", Description: "Z.AI GLM-5.1; long-horizon tasks"}, {ModelKey: "glm-5-turbo", Description: "Z.AI GLM-5-Turbo; cost-efficient"}, {ModelKey: "glm-4.7", Description: "Z.AI GLM-4.7; general chat and coding"}, diff --git a/pkg/aiproxy/models/catalog_seed_test.go b/pkg/aiproxy/models/catalog_seed_test.go index 467acc8792..d79beb6df8 100644 --- a/pkg/aiproxy/models/catalog_seed_test.go +++ b/pkg/aiproxy/models/catalog_seed_test.go @@ -35,3 +35,40 @@ func TestCatalogSeedModelsSkippedForSelfHostedProviders(t *testing.T) { } } } + +func TestCatalogContextWindow(t *testing.T) { + cases := []struct { + modelKey string + want int + }{ + {modelKey: "deepseek-v4-pro", want: 1_000_000}, + {modelKey: "deepseek-v4-flash", want: 1_000_000}, + {modelKey: "deepseek-chat", want: 1_000_000}, + {modelKey: "glm-5.2", want: 1_000_000}, + {modelKey: "glm-5.3", want: 1_000_000}, + {modelKey: "kimi-k3", want: 1_000_000}, + {modelKey: "mimo-v2.5", want: 1_000_000}, + {modelKey: "gpt-4.1", want: 1_000_000}, + {modelKey: "gpt-5.6", want: 1_050_000}, + {modelKey: "gemini-2.0-flash", want: 1_000_000}, + {modelKey: "gemini-1.5-pro", want: 2_000_000}, + {modelKey: "claude-opus-4-20250514", want: 1_000_000}, + {modelKey: "claude-sonnet-5", want: 1_000_000}, + {modelKey: "glm-5.1", want: 0}, + {modelKey: "gpt-5.2", want: 0}, + {modelKey: "kimi-k2.6", want: 0}, + {modelKey: "", want: 0}, + } + for _, tc := range cases { + got := CatalogContextWindow(tc.modelKey) + if tc.want >= 1_000_000 { + if got != tc.want { + t.Errorf("CatalogContextWindow(%q) = %d, want %d", tc.modelKey, got, tc.want) + } + continue + } + if got >= 1_000_000 { + t.Errorf("CatalogContextWindow(%q) = %d, want 0 or < 1M", tc.modelKey, got) + } + } +} diff --git a/pkg/aiproxy/models/provider_connectivity.go b/pkg/aiproxy/models/provider_connectivity.go index a1c703f677..541c6ac063 100644 --- a/pkg/aiproxy/models/provider_connectivity.go +++ b/pkg/aiproxy/models/provider_connectivity.go @@ -194,6 +194,54 @@ func providerUpstreamModels(modelKeys []string) []api.AiProviderUpstreamModel { return out } +func connectivityTestOutput(modelKeys []string, fromCatalog bool) api.AiProviderTestConnectivityOutput { + out := api.AiProviderTestConnectivityOutput{Ok: true} + if fromCatalog { + out.Message = "connectivity test passed (catalog models)" + out.ModelsSource = api.AiProviderModelsSourceCatalog + } else { + out.Message = "connectivity test passed" + out.ModelsSource = api.AiProviderModelsSourceUpstream + } + out.Models = providerUpstreamModels(modelKeys) + return out +} + +// resolveProviderSecretForConnectivity picks an enabled ai_key secret without model-key filtering. +func resolveProviderSecretForConnectivity(prov *SAiProvider) (string, error) { + if prov == nil { + return "", errors.Wrap(httperrors.ErrInvalidStatus, "ai_provider is nil") + } + pid := strings.TrimSpace(prov.Id) + if pid == "" { + return "", errors.Wrap(httperrors.ErrInvalidStatus, "ai_provider id is empty") + } + keys := make([]SAiKey, 0, 32) + q := AiKeyManager.Query().Equals("ai_provider_id", pid).Equals("enabled", true) + if err := q.All(&keys); err != nil { + return "", errors.Wrap(err, "list ai_key for provider") + } + candidates := make([]*SAiKey, 0, len(keys)) + for i := range keys { + k := &keys[i] + if strings.TrimSpace(k.GetSecret()) == "" { + continue + } + if baseAiKeyWeight(k) <= 0 { + continue + } + candidates = append(candidates, k) + } + if len(candidates) == 0 { + return "", errors.Wrap(httperrors.ErrInvalidStatus, "add an enabled ai_key with secret for this provider") + } + chosen := pickWeightedAiKey(candidates) + if chosen == nil { + return "", errors.Wrap(httperrors.ErrInvalidStatus, "failed to pick ai_key") + } + return strings.TrimSpace(chosen.GetSecret()), nil +} + // PerformTestConnectivity probes upstream list-models without persisting an ai_provider row. func (manager *SAiProviderManager) PerformTestConnectivity( ctx context.Context, @@ -206,14 +254,24 @@ func (manager *SAiProviderManager) PerformTestConnectivity( if err != nil { return out, err } - out.Ok = true - if fromCatalog { - out.Message = "connectivity test passed (catalog models)" - out.ModelsSource = api.AiProviderModelsSourceCatalog - } else { - out.Message = "connectivity test passed" - out.ModelsSource = api.AiProviderModelsSourceUpstream + return connectivityTestOutput(modelKeys, fromCatalog), nil +} + +// PerformTestConnectivity lists upstream models using a stored enabled ai_key. +func (p *SAiProvider) PerformTestConnectivity( + ctx context.Context, + userCred mcclient.TokenCredential, + query jsonutils.JSONObject, + input api.AiProviderInstanceTestConnectivityInput, +) (api.AiProviderTestConnectivityOutput, error) { + out := api.AiProviderTestConnectivityOutput{} + secret, err := resolveProviderSecretForConnectivity(p) + if err != nil { + return out, err } - out.Models = providerUpstreamModels(modelKeys) - return out, nil + modelKeys, fromCatalog, err := listProviderModels(ctx, p.ProviderKey, secret, p.Config, providerTestConnectivityTimeout) + if err != nil { + return out, err + } + return connectivityTestOutput(modelKeys, fromCatalog), nil } diff --git a/pkg/apis/aiproxy/ai_model.go b/pkg/apis/aiproxy/ai_model.go index 5bcd500d18..296c1da8b7 100644 --- a/pkg/apis/aiproxy/ai_model.go +++ b/pkg/apis/aiproxy/ai_model.go @@ -63,6 +63,7 @@ type AiModelDetails struct { VisualModelKey string `json:"visual_model_key"` VisualActive bool `json:"visual_active"` Config *SAiModelConfig `json:"config"` + ContextWindow int `json:"context_window,omitempty"` } // SAiModelConfig stores per-model extension settings. diff --git a/pkg/apis/aiproxy/ai_provider.go b/pkg/apis/aiproxy/ai_provider.go index 1745a07c40..e162609914 100644 --- a/pkg/apis/aiproxy/ai_provider.go +++ b/pkg/apis/aiproxy/ai_provider.go @@ -186,3 +186,12 @@ type AiProviderTestConnectivityOutput struct { ModelsSource string `json:"models_source"` Models []AiProviderUpstreamModel `json:"models"` } + +// AiProviderSetModelsInput imports selected catalog models for an existing provider. +// Existing model_key rows are skipped; nothing is deleted. +type AiProviderSetModelsInput struct { + ModelKeys []string `json:"model_keys"` +} + +// AiProviderInstanceTestConnectivityInput is the empty body for instance test-connectivity. +type AiProviderInstanceTestConnectivityInput struct{} diff --git a/pkg/mcclient/options/aiproxy/resources.go b/pkg/mcclient/options/aiproxy/resources.go index a55f36d1bf..f3d15b3f4c 100644 --- a/pkg/mcclient/options/aiproxy/resources.go +++ b/pkg/mcclient/options/aiproxy/resources.go @@ -86,6 +86,21 @@ func (o *AiProviderTestConnectivityOptions) Params() (jsonutils.JSONObject, erro return params, nil } +type AiProviderSetModelsOptions struct { + options.BaseIdOptions + ModelKeys []string `help:"model keys to import" json:"model_keys"` +} + +func (o *AiProviderSetModelsOptions) Params() (jsonutils.JSONObject, error) { + params := jsonutils.NewDict() + arr := jsonutils.NewArray() + for _, k := range o.ModelKeys { + arr.Add(jsonutils.NewString(k)) + } + params.Set("model_keys", arr) + return params, nil +} + type AiProviderUpdateOptions struct { ID string `help:"ID or name" json:"-"` Name string `json:"name,omitempty"`