diff --git a/controller/channel.go b/controller/channel.go index a2b5687ab319..fe0865a14995 100644 --- a/controller/channel.go +++ b/controller/channel.go @@ -15,7 +15,6 @@ import ( "github.com/QuantumNous/new-api/i18n" "github.com/QuantumNous/new-api/model" relaychannel "github.com/QuantumNous/new-api/relay/channel" - "github.com/QuantumNous/new-api/relay/channel/gemini" "github.com/QuantumNous/new-api/relay/channel/ollama" "github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/service/authz" @@ -1156,123 +1155,79 @@ func equalStringPtr(a, b *string) bool { return *a == *b } -func FetchModels(c *gin.Context) { - var req struct { - BaseURL string `json:"base_url"` - Type int `json:"type"` - Key string `json:"key"` - } +type fetchModelsRequest struct { + BaseURL string `json:"base_url"` + Type int `json:"type"` + Key string `json:"key"` + Setting string `json:"setting"` + Settings string `json:"settings"` + HeaderOverride string `json:"header_override"` +} - if err := c.ShouldBindJSON(&req); err != nil { - c.JSON(http.StatusBadRequest, gin.H{ - "success": false, - "message": "Invalid request", - }) - return +func (req fetchModelsRequest) channel() model.Channel { + channel := model.Channel{ + Type: req.Type, + Key: req.Key, + OtherSettings: req.Settings, } - - baseURL := req.BaseURL - if baseURL == "" { - baseURL = constant.ChannelBaseURLs[req.Type] + if req.BaseURL != "" { + channel.BaseURL = &req.BaseURL } - - // remove line breaks and extra spaces. - key := strings.TrimSpace(req.Key) - key = strings.Split(key, "\n")[0] - - if req.Type == constant.ChannelTypeOllama { - models, err := ollama.FetchOllamaModels(baseURL, key) - if err != nil { - c.JSON(http.StatusOK, gin.H{ - "success": false, - "message": fmt.Sprintf("获取Ollama模型失败: %s", err.Error()), - }) - return - } - - names := make([]string, 0, len(models)) - for _, modelInfo := range models { - names = append(names, modelInfo.Name) - } - - c.JSON(http.StatusOK, gin.H{ - "success": true, - "data": names, - }) - return + if req.Setting != "" { + channel.Setting = &req.Setting } - - if req.Type == constant.ChannelTypeGemini { - models, err := gemini.FetchGeminiModels(baseURL, key, "") - if err != nil { - c.JSON(http.StatusOK, gin.H{ - "success": false, - "message": fmt.Sprintf("获取Gemini模型失败: %s", err.Error()), - }) - return - } - - c.JSON(http.StatusOK, gin.H{ - "success": true, - "data": models, - }) - return + if req.HeaderOverride != "" { + channel.HeaderOverride = &req.HeaderOverride } + return channel +} - client := &http.Client{} - url := fmt.Sprintf("%s/v1/models", baseURL) - - request, err := http.NewRequest("GET", url, nil) - if err != nil { - c.JSON(http.StatusInternalServerError, gin.H{ - "success": false, - "message": err.Error(), - }) +func normalizeFetchModelsRequestKey(channel *model.Channel) { + key := strings.TrimSpace(channel.Key) + if key == "" || strings.HasPrefix(key, "{") || strings.HasPrefix(key, "[") { + channel.Key = key return } + channel.Key = strings.TrimSpace(strings.Split(key, "\n")[0]) +} + +func fetchModelsRequestRequiresKey(channelType int) bool { + return channelType != constant.ChannelTypeOllama +} - request.Header.Set("Authorization", "Bearer "+key) +func FetchModels(c *gin.Context) { + var req fetchModelsRequest - response, err := client.Do(request) - if err != nil { - c.JSON(http.StatusInternalServerError, gin.H{ + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{ "success": false, - "message": err.Error(), + "message": "Invalid request", }) return } - //check status code - if response.StatusCode != http.StatusOK { - c.JSON(http.StatusInternalServerError, gin.H{ + + channel := req.channel() + normalizeFetchModelsRequestKey(&channel) + if fetchModelsRequestRequiresKey(channel.Type) && channel.Key == "" { + c.JSON(http.StatusOK, gin.H{ "success": false, - "message": "Failed to fetch models", + "message": "Please enter API key first", }) return } - defer response.Body.Close() - - var result struct { - Data []struct { - ID string `json:"id"` - } `json:"data"` - } - if err := json.NewDecoder(response.Body).Decode(&result); err != nil { - c.JSON(http.StatusInternalServerError, gin.H{ + ids, err := fetchChannelUpstreamModelIDs(&channel) + if err != nil { + c.JSON(http.StatusOK, gin.H{ "success": false, - "message": err.Error(), + "message": fmt.Sprintf("获取模型列表失败: %s", err.Error()), }) return } - var models []string - for _, model := range result.Data { - models = append(models, model.ID) - } - c.JSON(http.StatusOK, gin.H{ "success": true, - "data": models, + "data": ids, }) } diff --git a/controller/channel_fetch_models_test.go b/controller/channel_fetch_models_test.go new file mode 100644 index 000000000000..38e45a2b3469 --- /dev/null +++ b/controller/channel_fetch_models_test.go @@ -0,0 +1,63 @@ +package controller + +import ( + "encoding/json" + "testing" + + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" + "github.com/stretchr/testify/require" +) + +func TestNormalizeFetchModelsRequestKeyUsesFirstPlaintextKey(t *testing.T) { + channel := &model.Channel{Key: " sk-first \n sk-second \n"} + + normalizeFetchModelsRequestKey(channel) + + require.Equal(t, "sk-first", channel.Key) +} + +func TestNormalizeFetchModelsRequestKeyKeepsJSONCredential(t *testing.T) { + key := "{\n \"access_token\": \"token\",\n \"account_id\": \"account\"\n}" + channel := &model.Channel{Key: key} + + normalizeFetchModelsRequestKey(channel) + + require.Equal(t, key, channel.Key) +} + +func TestFetchModelsRequestAllowsOllamaWithoutKey(t *testing.T) { + require.False(t, fetchModelsRequestRequiresKey(constant.ChannelTypeOllama)) + require.True(t, fetchModelsRequestRequiresKey(constant.ChannelTypeOpenAI)) +} + +func TestFetchModelsRequestIgnoresPersistedChannelStateFields(t *testing.T) { + body := []byte(`{ + "id": 123, + "type": 1, + "key": "sk-test", + "base_url": "https://api.example.com", + "setting": "{\"proxy\":\"http://proxy.example.com\"}", + "settings": "{\"vertex_key_type\":\"api_key\"}", + "header_override": "{\"X-Test\":\"ok\"}", + "channel_info": { + "is_multi_key": true, + "multi_key_mode": "polling", + "multi_key_status_list": {"0": 1} + } + }`) + + var req fetchModelsRequest + require.NoError(t, json.Unmarshal(body, &req)) + + channel := req.channel() + + require.Zero(t, channel.Id) + require.False(t, channel.ChannelInfo.IsMultiKey) + require.Equal(t, constant.ChannelTypeOpenAI, channel.Type) + require.Equal(t, "sk-test", channel.Key) + require.Equal(t, "https://api.example.com", channel.GetBaseURL()) + require.NotNil(t, channel.Setting) + require.NotNil(t, channel.HeaderOverride) + require.Equal(t, `{"vertex_key_type":"api_key"}`, channel.OtherSettings) +} diff --git a/web/default/src/features/channels/api.ts b/web/default/src/features/channels/api.ts index ba8a67ef2c83..57bad5a80ca0 100644 --- a/web/default/src/features/channels/api.ts +++ b/web/default/src/features/channels/api.ts @@ -529,6 +529,9 @@ export async function fetchModels(data: { base_url: string type: number key: string + setting?: string | null + settings?: string | null + header_override?: string | null }): Promise { const res = await api.post( '/api/channel/fetch_models', diff --git a/web/default/src/features/channels/components/drawers/channel-mutate-drawer.tsx b/web/default/src/features/channels/components/drawers/channel-mutate-drawer.tsx index f21228c34626..fa048d0065c2 100644 --- a/web/default/src/features/channels/components/drawers/channel-mutate-drawer.tsx +++ b/web/default/src/features/channels/components/drawers/channel-mutate-drawer.tsx @@ -159,6 +159,7 @@ import { findMissingModelsInMapping, validateModelMappingJson, hasAdvancedSettingsErrors, + transformFormDataToFetchModelsPayload, } from '../../lib' import { collectInvalidStatusCodeEntries, @@ -1352,15 +1353,13 @@ export function ChannelMutateDrawer({ setFetchModelsDialogOpen(true) }, [isEditing, canEditSensitive, form, t]) - const createModeFetcher = useCallback(async (): Promise => { + const formModeFetcher = useCallback(async (): Promise => { if (!canEditSensitive) { throw new Error(t("You don't have necessary permission")) } - const response = await fetchModels({ - type: form.getValues('type'), - key: form.getValues('key'), - base_url: form.getValues('base_url') || '', - }) + const response = await fetchModels( + transformFormDataToFetchModelsPayload(form.getValues()) + ) if (response.success && response.data) { return response.data } @@ -4552,7 +4551,9 @@ export function ChannelMutateDrawer({ }} redirectModels={redirectModelList} redirectSourceModels={redirectModelKeyList} - customFetcher={!isEditing ? createModeFetcher : undefined} + customFetcher={ + !isEditing || currentKey?.trim() ? formModeFetcher : undefined + } channelName={!isEditing ? currentName?.trim() : undefined} existingModelsOverride={ !isEditing diff --git a/web/default/src/features/channels/lib/channel-form.ts b/web/default/src/features/channels/lib/channel-form.ts index 33a5fb707a92..ff98395d4241 100644 --- a/web/default/src/features/channels/lib/channel-form.ts +++ b/web/default/src/features/channels/lib/channel-form.ts @@ -564,12 +564,15 @@ function buildSettingsJSON(formData: ChannelFormValues): string { settingsObj.allow_inference_geo = formData.allow_inference_geo === true } else { if ('disable_store' in settingsObj) delete settingsObj.disable_store - if ('allow_safety_identifier' in settingsObj) + if ('allow_safety_identifier' in settingsObj) { delete settingsObj.allow_safety_identifier - if ('allow_include_obfuscation' in settingsObj) + } + if ('allow_include_obfuscation' in settingsObj) { delete settingsObj.allow_include_obfuscation - if (formData.type !== 14 && 'allow_inference_geo' in settingsObj) + } + if (formData.type !== 14 && 'allow_inference_geo' in settingsObj) { delete settingsObj.allow_inference_geo + } } // Anthropic (type 14): claude_beta_query, allow_inference_geo, allow_speed @@ -592,14 +595,14 @@ function buildSettingsJSON(formData: ChannelFormValues): string { settingsObj.upstream_model_update_auto_sync_enabled = settingsObj.upstream_model_update_check_enabled === true && formData.upstream_model_update_auto_sync_enabled === true - settingsObj.upstream_model_update_ignored_models = Array.from( - new Set( + settingsObj.upstream_model_update_ignored_models = [ + ...new Set( String(formData.upstream_model_update_ignored_models || '') .split(',') .map((model) => model.trim()) .filter(Boolean) - ) - ) + ), + ] if ( !Array.isArray(settingsObj.upstream_model_update_last_detected_models) || settingsObj.upstream_model_update_check_enabled !== true @@ -739,6 +742,26 @@ export function transformFormDataToUpdatePayload( return payload } +export function transformFormDataToFetchModelsPayload( + formData: ChannelFormValues +): { + base_url: string + type: number + key: string + setting: string + settings: string + header_override: string +} { + return { + base_url: normalizeBaseUrl(formData.base_url) || '', + type: formData.type, + key: formData.key, + setting: buildSettingJSON(formData), + settings: buildSettingsJSON(formData), + header_override: formData.header_override || '', + } +} + // ============================================================================ // Validation Helpers // ============================================================================