diff --git a/controller/channel_upstream_update.go b/controller/channel_upstream_update.go index 71ab0e53fafe..a734be09f6ae 100644 --- a/controller/channel_upstream_update.go +++ b/controller/channel_upstream_update.go @@ -516,6 +516,7 @@ func checkAndPersistChannelUpstreamModelUpdates( if err = channel.UpdateAbilities(nil); err != nil { return true, autoAdded, err } + _ = model.EnsureChannelVendorAndModels(channel, nil) } return modelsChanged, autoAdded, nil } @@ -962,6 +963,7 @@ func applyChannelUpstreamModelUpdates( if err := channel.UpdateAbilities(nil); err != nil { return addModels, removeModels, remainingModels, remainingRemoveModels, true, err } + _ = model.EnsureChannelVendorAndModels(channel, nil) } return addModels, removeModels, remainingModels, remainingRemoveModels, modelsChanged, nil } diff --git a/model/channel.go b/model/channel.go index 2cd7c3115ff6..019e6da6c789 100644 --- a/model/channel.go +++ b/model/channel.go @@ -447,6 +447,7 @@ func BatchInsertChannels(channels []Channel) error { tx.Rollback() return err } + _ = EnsureChannelVendorAndModels(&channel_, tx) } } return tx.Commit().Error @@ -526,6 +527,9 @@ func (channel *Channel) Insert() error { return err } err = channel.AddAbilities(nil) + if err == nil { + _ = EnsureChannelVendorAndModels(channel, nil) + } return err } @@ -575,6 +579,9 @@ func (channel *Channel) Update() error { } DB.Model(channel).First(channel, "id = ?", channel.Id) err = channel.UpdateAbilities(nil) + if err == nil { + _ = EnsureChannelVendorAndModels(channel, nil) + } return err } diff --git a/model/vendor_channel_auto_test.go b/model/vendor_channel_auto_test.go new file mode 100644 index 000000000000..c4ee2e887141 --- /dev/null +++ b/model/vendor_channel_auto_test.go @@ -0,0 +1,48 @@ +package model + +import ( + "fmt" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestEnsureVendorAndAutoBindModelsForNewChannel(t *testing.T) { + require.NoError(t, DB.AutoMigrate(&Channel{}, &Ability{}, &Vendor{}, &Model{})) + + uniqueChannelName := fmt.Sprintf("AutoVendorChan_%d", common.GetTimestamp()) + testModel1 := fmt.Sprintf("auto-test-model-1-%d", common.GetTimestamp()) + testModel2 := fmt.Sprintf("auto-test-model-2-%d", common.GetTimestamp()) + + channel := &Channel{ + Type: constant.ChannelTypeAdvancedCustom, + Name: uniqueChannelName, + Key: "sk-test-key", + Models: fmt.Sprintf("%s,%s", testModel1, testModel2), + Status: common.ChannelStatusEnabled, + } + + err := channel.Insert() + require.NoError(t, err) + + // 1. Verify Vendor was auto-created + var vendor Vendor + err = DB.Where("name = ? AND deleted_at IS NULL", uniqueChannelName).First(&vendor).Error + require.NoError(t, err) + assert.Equal(t, uniqueChannelName, vendor.Name) + assert.Equal(t, 1, vendor.Status) + + // 2. Verify model metadata was auto-created and bound to vendor.Id + var meta1 Model + err = DB.Where("model_name = ? AND deleted_at IS NULL", testModel1).First(&meta1).Error + require.NoError(t, err) + assert.Equal(t, vendor.Id, meta1.VendorID) + + var meta2 Model + err = DB.Where("model_name = ? AND deleted_at IS NULL", testModel2).First(&meta2).Error + require.NoError(t, err) + assert.Equal(t, vendor.Id, meta2.VendorID) +} diff --git a/model/vendor_meta.go b/model/vendor_meta.go index 2bb357f82a1f..b749de3a10b1 100644 --- a/model/vendor_meta.go +++ b/model/vendor_meta.go @@ -1,7 +1,12 @@ package model import ( + "errors" + "fmt" + "strings" + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" "gorm.io/gorm" ) @@ -86,3 +91,123 @@ func SearchVendors(keyword string, offset int, limit int) ([]*Vendor, int64, err } return vendors, total, nil } + +func resolveVendorNameAndIcon(channel *Channel) (string, string) { + if channel == nil { + return "Custom", "Globe" + } + + name := strings.TrimSpace(channel.Name) + switch channel.Type { + case constant.ChannelTypeAdvancedCustom, constant.ChannelTypeNewAPI, constant.ChannelTypeSub2API, constant.ChannelTypeCustom: + if name != "" { + return name, "Globe" + } + return constant.GetChannelTypeName(channel.Type), "Globe" + default: + typeName := constant.GetChannelTypeName(channel.Type) + if typeName != "" && typeName != "Unknown" { + return typeName, "Globe" + } + if name != "" { + return name, "Globe" + } + return "Custom", "Globe" + } +} + +// EnsureVendorForChannel 检查并自动创建渠道对应的供应商记录,返回 Vendor ID +func EnsureVendorForChannel(channel *Channel, tx *gorm.DB) (int, error) { + if channel == nil { + return 0, nil + } + + name, icon := resolveVendorNameAndIcon(channel) + if name == "" { + return 0, nil + } + + useDB := DB + if tx != nil { + useDB = tx + } + + var vendor Vendor + err := useDB.Where("name = ? AND deleted_at IS NULL", name).First(&vendor).Error + if err == nil { + return vendor.Id, nil + } + + if errors.Is(err, gorm.ErrRecordNotFound) { + newVendor := Vendor{ + Name: name, + Description: fmt.Sprintf("%s 渠道自动创建供应商", name), + Icon: icon, + Status: 1, + CreatedTime: common.GetTimestamp(), + UpdatedTime: common.GetTimestamp(), + } + if err := useDB.Create(&newVendor).Error; err != nil { + return 0, err + } + return newVendor.Id, nil + } + + return 0, err +} + +// AutoBindChannelModelsToVendor 为渠道拥有的模型自动补充元数据并绑定 Vendor ID +func AutoBindChannelModelsToVendor(channel *Channel, vendorID int, tx *gorm.DB) error { + if channel == nil || channel.Models == "" || vendorID <= 0 { + return nil + } + + useDB := DB + if tx != nil { + useDB = tx + } + + models := strings.Split(channel.Models, ",") + now := common.GetTimestamp() + + for _, modelName := range models { + modelName = strings.TrimSpace(modelName) + if modelName == "" { + continue + } + + var existing Model + err := useDB.Where("model_name = ? AND deleted_at IS NULL", modelName).First(&existing).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + newModel := Model{ + ModelName: modelName, + VendorID: vendorID, + Status: 1, + SyncOfficial: 0, + Endpoints: `["openai"]`, + CreatedTime: now, + UpdatedTime: now, + } + if err := useDB.Create(&newModel).Error; err != nil { + common.SysError(fmt.Sprintf("auto bind model failed: %s, err: %v", modelName, err)) + } + } else if err == nil && existing.VendorID == 0 { + useDB.Model(&existing).Update("vendor_id", vendorID) + } + } + return nil +} + +// EnsureChannelVendorAndModels 自动确保渠道对应的供应商和模型元数据已关联 +func EnsureChannelVendorAndModels(channel *Channel, tx *gorm.DB) error { + if channel == nil { + return nil + } + vendorID, err := EnsureVendorForChannel(channel, tx) + if err != nil { + common.SysError(fmt.Sprintf("EnsureVendorForChannel failed for channel %d (%s): %v", channel.Id, channel.Name, err)) + return err + } + return AutoBindChannelModelsToVendor(channel, vendorID, tx) +} + diff --git a/relay/channel/gemini/constant.go b/relay/channel/gemini/constant.go index 08667b9b548f..7e868da2336f 100644 --- a/relay/channel/gemini/constant.go +++ b/relay/channel/gemini/constant.go @@ -4,7 +4,8 @@ var ModelList = []string{ // stable version "gemini-2.5-flash", "gemini-2.5-pro", "gemini-2.0-flash", "gemini-2.0-flash-001", "gemini-2.0-flash-lite-001", "gemini-2.0-flash-lite", - "gemini-2.5-flash-lite", + "gemini-2.5-flash-lite", "gemini-3.1-flash-lite", "gemini-3.5-flash-lite", + "gemini-3.6", "gemini-3.6-flash", "gemini-3-pro-image", "gemini-3.1-flash-image", // latest version "gemini-flash-latest", "gemini-flash-lite-latest", "gemini-pro-latest", diff --git a/setting/ratio_setting/cache_ratio.go b/setting/ratio_setting/cache_ratio.go index 6e874b5bc5f9..e4bcb6d59196 100644 --- a/setting/ratio_setting/cache_ratio.go +++ b/setting/ratio_setting/cache_ratio.go @@ -8,6 +8,10 @@ var defaultCacheRatio = map[string]float64{ "gemini-3-flash-preview": 0.1, "gemini-3-pro-preview": 0.1, "gemini-3.1-pro-preview": 0.1, + "gemini-3.1-flash-lite": 0.1, + "gemini-3.5-flash-lite": 0.1, + "gemini-3.6": 0.1, + "gemini-3.6-flash": 0.1, "gpt-4": 0.5, "o1": 0.5, "o1-2024-12-17": 0.5, diff --git a/setting/ratio_setting/model_ratio.go b/setting/ratio_setting/model_ratio.go index 829e0794a157..b2d4776ad801 100644 --- a/setting/ratio_setting/model_ratio.go +++ b/setting/ratio_setting/model_ratio.go @@ -190,6 +190,10 @@ var defaultModelRatio = map[string]float64{ "gemini-2.5-flash-lite-preview-thinking-*": 0.05, "gemini-2.5-flash-lite-preview-06-17": 0.05, "gemini-2.5-flash": 0.15, + "gemini-3.1-flash-lite": 0.125, + "gemini-3.5-flash-lite": 0.15, + "gemini-3.6": 0.75, + "gemini-3.6-flash": 0.75, "gemini-robotics-er-1.5-preview": 0.15, "gemini-embedding-001": 0.075, "text-embedding-004": 0.001, @@ -326,10 +330,14 @@ var modelRatioMap = types.NewRWMap[string, float64]() var completionRatioMap = types.NewRWMap[string, float64]() var defaultCompletionRatio = map[string]float64{ - "gpt-4-gizmo-*": 2, - "gpt-4o-gizmo-*": 3, - "gpt-4-all": 2, - "gpt-image-1": 8, + "gpt-4-gizmo-*": 2, + "gpt-4o-gizmo-*": 3, + "gpt-4-all": 2, + "gpt-image-1": 8, + "gemini-3.1-flash-lite": 6, + "gemini-3.5-flash-lite": 2.5 / 0.3, + "gemini-3.6": 5, + "gemini-3.6-flash": 5, } // InitRatioSettings initializes all model related settings maps @@ -577,6 +585,12 @@ func getHardcodedCompletionModelRatio(name string) (float64, bool) { return 2.5 / 0.3, false } else if strings.HasPrefix(name, "gemini-robotics-er-1.5") { return 2.5 / 0.3, false + } else if strings.HasPrefix(name, "gemini-3.1-flash-lite") { + return 6, false + } else if strings.HasPrefix(name, "gemini-3.5-flash-lite") { + return 2.5 / 0.3, false + } else if strings.HasPrefix(name, "gemini-3.6") { + return 5, false } else if strings.HasPrefix(name, "gemini-3-pro") { if strings.HasPrefix(name, "gemini-3-pro-image") { return 60, false diff --git a/setting/ratio_setting/model_ratio_test.go b/setting/ratio_setting/model_ratio_test.go new file mode 100644 index 000000000000..c228f912117c --- /dev/null +++ b/setting/ratio_setting/model_ratio_test.go @@ -0,0 +1,52 @@ +package ratio_setting_test + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/QuantumNous/new-api/setting/ratio_setting" +) + +func TestGemini3ModelRatios(t *testing.T) { + ratio_setting.InitRatioSettings() + + tests := []struct { + model string + expectedModelRatio float64 + expectedCompRatio float64 + }{ + { + model: "gemini-3.1-flash-lite", + expectedModelRatio: 0.125, + expectedCompRatio: 6.0, + }, + { + model: "gemini-3.5-flash-lite", + expectedModelRatio: 0.15, + expectedCompRatio: 2.5 / 0.3, + }, + { + model: "gemini-3.6", + expectedModelRatio: 0.75, + expectedCompRatio: 5.0, + }, + { + model: "gemini-3.6-flash", + expectedModelRatio: 0.75, + expectedCompRatio: 5.0, + }, + } + + for _, tt := range tests { + t.Run(tt.model, func(t *testing.T) { + ratio, ok, _ := ratio_setting.GetModelRatio(tt.model) + require.True(t, ok, "model ratio for %s should exist", tt.model) + assert.InDelta(t, tt.expectedModelRatio, ratio, 0.0001, "model ratio mismatch for %s", tt.model) + + compRatio := ratio_setting.GetCompletionRatio(tt.model) + assert.InDelta(t, tt.expectedCompRatio, compRatio, 0.0001, "completion ratio mismatch for %s", tt.model) + }) + } +}