feat: enhance text protocol conversion and advanced custom routing (#5825)
* refactor: consolidate relay protocol converters * refactor relayconvert text converters * feat: refine relay converters and advanced custom routing * refactor: enhance logging and add thought signature handling for Gemini requests * refactor: enhance channel cache and pricing endpoint handling for advanced custom models * feat: preserve billing usage semantics * feat: add protocol-aware billing usage * Delete useless files * chore: update action versions in workflow files * chore: update Docker action versions in workflow files * fix: harden billing usage settlement and hot-path route matching - estimate Gemini completion tokens locally when billable usageMetadata is prompt-only but output content was received (e.g. client aborts the stream before the final chunk), and rebuild the attached billing_usage as estimated so settlement does not bill zero output tokens - guard NewClaudeMessagesBillingUsage against all-zero ClaudeUsage, matching the OpenAI/Gemini constructors, so a zero billing_usage cannot override a non-zero top-level usage during settlement - cache compiled advanced-custom route model regexes; they run on the request hot path and were recompiled per request - move the effectiveBillingUsage remap to PostTextConsumeQuota only, and document that calculateTextQuotaSummary expects remapped usage - document the updatePricingLock -> channelSyncLock lock ordering that InitChannelCache/CacheUpdateChannel rely on, and the aux-struct pitfall in GeminiChatResponse.UnmarshalJSON
This commit is contained in:
+8
-7
@@ -121,7 +121,7 @@ func GetChannel(group string, model string, retry int, requestPath string) (*Cha
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
abilities = filterAbilitiesByRequestPath(abilities, requestPath)
|
||||
abilities = filterAbilitiesByRequestPathAndModel(abilities, requestPath, model)
|
||||
channel := Channel{}
|
||||
if len(abilities) > 0 {
|
||||
// Randomly choose one
|
||||
@@ -146,11 +146,12 @@ func GetChannel(group string, model string, retry int, requestPath string) (*Cha
|
||||
return &channel, err
|
||||
}
|
||||
|
||||
// filterAbilitiesByRequestPath restricts candidates by request path for the DB
|
||||
// (non-memory-cache) selection path. Only Advanced Custom (type 58) channels are
|
||||
// path-checked: kept only when one of their routes matches requestPath; all other
|
||||
// channel types always pass. When requestPath is empty, filtering is skipped.
|
||||
func filterAbilitiesByRequestPath(abilities []Ability, requestPath string) []Ability {
|
||||
// filterAbilitiesByRequestPathAndModel restricts candidates by request path and
|
||||
// model for the DB (non-memory-cache) selection path. Only Advanced Custom
|
||||
// (type 58) channels are path-checked: kept only when one of their routes matches
|
||||
// requestPath and model; all other channel types always pass. When requestPath is
|
||||
// empty, filtering is skipped.
|
||||
func filterAbilitiesByRequestPathAndModel(abilities []Ability, requestPath string, model string) []Ability {
|
||||
if requestPath == "" || len(abilities) == 0 {
|
||||
return abilities
|
||||
}
|
||||
@@ -185,7 +186,7 @@ func filterAbilitiesByRequestPath(abilities []Ability, requestPath string) []Abi
|
||||
filtered = append(filtered, ability)
|
||||
continue
|
||||
}
|
||||
if config != nil && config.SupportsPath(requestPath) {
|
||||
if config != nil && config.SupportsPathForModel(requestPath, model) {
|
||||
filtered = append(filtered, ability)
|
||||
}
|
||||
}
|
||||
|
||||
+30
-9
@@ -25,6 +25,7 @@ var channelSyncLock sync.RWMutex
|
||||
|
||||
func InitChannelCache() {
|
||||
if !common.MemoryCacheEnabled {
|
||||
InvalidatePricingCache()
|
||||
return
|
||||
}
|
||||
newChannelId2channel := make(map[int]*Channel)
|
||||
@@ -94,6 +95,11 @@ func InitChannelCache() {
|
||||
channelsIDM = newChannelId2channel
|
||||
channel2advancedCustomConfig = newChannel2advancedCustomConfig
|
||||
channelSyncLock.Unlock()
|
||||
// Lock ordering: InvalidatePricingCache acquires updatePricingLock, and
|
||||
// GetPricing (holding updatePricingLock) nests channelSyncLock.RLock via
|
||||
// loadPricingAdvancedCustomConfigs. channelSyncLock MUST be released before
|
||||
// invalidating the pricing cache, otherwise the reversed order deadlocks.
|
||||
InvalidatePricingCache()
|
||||
common.SysLog("channels synced from database")
|
||||
}
|
||||
|
||||
@@ -115,12 +121,12 @@ func GetRandomSatisfiedChannel(group string, model string, retry int, requestPat
|
||||
defer channelSyncLock.RUnlock()
|
||||
|
||||
// First, try to find channels with the exact model name.
|
||||
channels := filterChannelsByRequestPath(group2model2channels[group][model], requestPath)
|
||||
channels := filterChannelsByRequestPathAndModel(group2model2channels[group][model], requestPath, model)
|
||||
|
||||
// If no channels found, try to find channels with the normalized model name.
|
||||
if len(channels) == 0 {
|
||||
normalizedModel := ratio_setting.FormatMatchingModelName(model)
|
||||
channels = filterChannelsByRequestPath(group2model2channels[group][normalizedModel], requestPath)
|
||||
channels = filterChannelsByRequestPathAndModel(group2model2channels[group][normalizedModel], requestPath, model)
|
||||
}
|
||||
|
||||
if len(channels) == 0 {
|
||||
@@ -202,12 +208,12 @@ func GetRandomSatisfiedChannel(group string, model string, retry int, requestPat
|
||||
return nil, errors.New("channel not found")
|
||||
}
|
||||
|
||||
// filterChannelsByRequestPath restricts candidates by request path. Only Advanced
|
||||
// Custom (type 58) channels are path-checked: they are kept only when one of their
|
||||
// configured routes matches requestPath. All other channel types always pass.
|
||||
// When requestPath is empty (non-relay callers) filtering is skipped.
|
||||
// filterChannelsByRequestPathAndModel restricts candidates by request path and
|
||||
// model. Only Advanced Custom (type 58) channels are path-checked: they are kept
|
||||
// only when one of their configured routes matches requestPath and model. All
|
||||
// other channel types always pass. When requestPath is empty, filtering is skipped.
|
||||
// Caller must hold channelSyncLock (read lock). The cached slice is never mutated.
|
||||
func filterChannelsByRequestPath(channels []int, requestPath string) []int {
|
||||
func filterChannelsByRequestPathAndModel(channels []int, requestPath string, model string) []int {
|
||||
if requestPath == "" || len(channels) == 0 {
|
||||
return channels
|
||||
}
|
||||
@@ -223,7 +229,7 @@ func filterChannelsByRequestPath(channels []int, requestPath string) []int {
|
||||
filtered = append(filtered, channelId)
|
||||
continue
|
||||
}
|
||||
if config := channel2advancedCustomConfig[channelId]; config != nil && config.SupportsPath(requestPath) {
|
||||
if config := channel2advancedCustomConfig[channelId]; config != nil && config.SupportsPathForModel(requestPath, model) {
|
||||
filtered = append(filtered, channelId)
|
||||
}
|
||||
}
|
||||
@@ -292,8 +298,8 @@ func CacheUpdateChannel(channel *Channel) {
|
||||
return
|
||||
}
|
||||
channelSyncLock.Lock()
|
||||
defer channelSyncLock.Unlock()
|
||||
if channel == nil {
|
||||
channelSyncLock.Unlock()
|
||||
return
|
||||
}
|
||||
|
||||
@@ -304,5 +310,20 @@ func CacheUpdateChannel(channel *Channel) {
|
||||
logger.LogDebug(nil, "CacheUpdateChannel before: id=%d, name=%s, status=%d, polling_index=%d", channel.Id, channel.Name, channel.Status, oldChannel.ChannelInfo.MultiKeyPollingIndex)
|
||||
}
|
||||
channelsIDM[channel.Id] = channel
|
||||
if channel2advancedCustomConfig == nil {
|
||||
channel2advancedCustomConfig = make(map[int]*dto.AdvancedCustomConfig)
|
||||
}
|
||||
delete(channel2advancedCustomConfig, channel.Id)
|
||||
if channel.Type == constant.ChannelTypeAdvancedCustom {
|
||||
if config := channel.GetOtherSettings().AdvancedCustom; config != nil {
|
||||
channel2advancedCustomConfig[channel.Id] = config
|
||||
}
|
||||
}
|
||||
logger.LogDebug(nil, "CacheUpdateChannel after: id=%d, name=%s, status=%d, polling_index=%d", channel.Id, channel.Name, channel.Status, channel.ChannelInfo.MultiKeyPollingIndex)
|
||||
// Lock ordering: do NOT hold channelSyncLock while calling
|
||||
// InvalidatePricingCache. GetPricing acquires updatePricingLock first and then
|
||||
// channelSyncLock.RLock (via loadPricingAdvancedCustomConfigs); acquiring
|
||||
// updatePricingLock while holding channelSyncLock would be an AB-BA deadlock.
|
||||
channelSyncLock.Unlock()
|
||||
InvalidatePricingCache()
|
||||
}
|
||||
|
||||
+78
-9
@@ -1,7 +1,6 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
@@ -10,6 +9,7 @@ import (
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/setting/billing_setting"
|
||||
"github.com/QuantumNous/new-api/setting/ratio_setting"
|
||||
"github.com/QuantumNous/new-api/types"
|
||||
@@ -107,6 +107,76 @@ func GetModelSupportEndpointTypes(model string) []constant.EndpointType {
|
||||
return make([]constant.EndpointType, 0)
|
||||
}
|
||||
|
||||
func getPricingEndpointTypesForAbility(ability AbilityWithChannel, advancedCustomConfigs map[int]*dto.AdvancedCustomConfig) []constant.EndpointType {
|
||||
if ability.ChannelType != constant.ChannelTypeAdvancedCustom {
|
||||
return common.GetEndpointTypesByChannelType(ability.ChannelType, ability.Model)
|
||||
}
|
||||
if config := advancedCustomConfigs[ability.ChannelId]; config != nil {
|
||||
return config.SupportedEndpointTypesForModel(ability.Model)
|
||||
}
|
||||
return common.GetEndpointTypesByChannelType(ability.ChannelType, ability.Model)
|
||||
}
|
||||
|
||||
// loadPricingAdvancedCustomConfigs runs inside updatePricing while
|
||||
// updatePricingLock is held, and nests channelSyncLock.RLock. This defines the
|
||||
// global lock order updatePricingLock -> channelSyncLock: any code path holding
|
||||
// channelSyncLock must release it before touching the pricing cache (see
|
||||
// InitChannelCache / CacheUpdateChannel), otherwise it deadlocks.
|
||||
// The returned configs are pointers shared with the channel cache; they are
|
||||
// replaced wholesale on update and never mutated in place, so reading them after
|
||||
// RUnlock is safe.
|
||||
func loadPricingAdvancedCustomConfigs(enableAbilities []AbilityWithChannel) map[int]*dto.AdvancedCustomConfig {
|
||||
channelIDs := make([]int, 0)
|
||||
seen := make(map[int]struct{})
|
||||
for _, ability := range enableAbilities {
|
||||
if ability.ChannelType != constant.ChannelTypeAdvancedCustom {
|
||||
continue
|
||||
}
|
||||
if _, exists := seen[ability.ChannelId]; exists {
|
||||
continue
|
||||
}
|
||||
seen[ability.ChannelId] = struct{}{}
|
||||
channelIDs = append(channelIDs, ability.ChannelId)
|
||||
}
|
||||
if len(channelIDs) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
configs := make(map[int]*dto.AdvancedCustomConfig, len(channelIDs))
|
||||
if common.MemoryCacheEnabled {
|
||||
channelSyncLock.RLock()
|
||||
defer channelSyncLock.RUnlock()
|
||||
for _, channelID := range channelIDs {
|
||||
if config := channel2advancedCustomConfig[channelID]; config != nil {
|
||||
configs[channelID] = config
|
||||
}
|
||||
}
|
||||
return configs
|
||||
}
|
||||
|
||||
for _, channelID := range channelIDs {
|
||||
channel, err := CacheGetChannel(channelID)
|
||||
if err != nil {
|
||||
common.SysLog(fmt.Sprintf("load advanced custom channel settings error: channel_id=%d, error=%v", channelID, err))
|
||||
continue
|
||||
}
|
||||
if channel.Type != constant.ChannelTypeAdvancedCustom {
|
||||
continue
|
||||
}
|
||||
if config := channel.GetOtherSettings().AdvancedCustom; config != nil {
|
||||
configs[channelID] = config
|
||||
}
|
||||
}
|
||||
return configs
|
||||
}
|
||||
|
||||
func appendPricingEndpoint(endpoints []string, endpoint string) []string {
|
||||
if endpoint == "" || common.StringsContains(endpoints, endpoint) {
|
||||
return endpoints
|
||||
}
|
||||
return append(endpoints, endpoint)
|
||||
}
|
||||
|
||||
func updatePricing() {
|
||||
//modelRatios := common.GetModelRatios()
|
||||
enableAbilities, err := GetAllEnableAbilityWithChannels()
|
||||
@@ -201,11 +271,12 @@ func updatePricing() {
|
||||
|
||||
//这里使用切片而不是Set,因为一个模型可能支持多个端点类型,并且第一个端点是优先使用端点
|
||||
modelSupportEndpointsStr := make(map[string][]string)
|
||||
advancedCustomConfigs := loadPricingAdvancedCustomConfigs(enableAbilities)
|
||||
|
||||
// 先根据已有能力填充原生端点
|
||||
for _, ability := range enableAbilities {
|
||||
endpoints := modelSupportEndpointsStr[ability.Model]
|
||||
channelTypes := common.GetEndpointTypesByChannelType(ability.ChannelType, ability.Model)
|
||||
channelTypes := getPricingEndpointTypesForAbility(ability, advancedCustomConfigs)
|
||||
for _, channelType := range channelTypes {
|
||||
if !common.StringsContains(endpoints, string(channelType)) {
|
||||
endpoints = append(endpoints, string(channelType))
|
||||
@@ -214,20 +285,18 @@ func updatePricing() {
|
||||
modelSupportEndpointsStr[ability.Model] = endpoints
|
||||
}
|
||||
|
||||
// 再补充模型自定义端点:若配置有效则替换默认端点,不做合并
|
||||
// 再补充模型自定义端点:若配置有效则追加到已有推断,不再裁剪渠道真实能力
|
||||
for modelName, meta := range metaMap {
|
||||
if strings.TrimSpace(meta.Endpoints) == "" {
|
||||
continue
|
||||
}
|
||||
var raw map[string]interface{}
|
||||
if err := json.Unmarshal([]byte(meta.Endpoints), &raw); err == nil {
|
||||
endpoints := make([]string, 0, len(raw))
|
||||
if err := common.Unmarshal([]byte(meta.Endpoints), &raw); err == nil {
|
||||
endpoints := modelSupportEndpointsStr[modelName]
|
||||
for k, v := range raw {
|
||||
switch v.(type) {
|
||||
case string, map[string]interface{}:
|
||||
if !common.StringsContains(endpoints, k) {
|
||||
endpoints = append(endpoints, k)
|
||||
}
|
||||
endpoints = appendPricingEndpoint(endpoints, k)
|
||||
}
|
||||
}
|
||||
if len(endpoints) > 0 {
|
||||
@@ -264,7 +333,7 @@ func updatePricing() {
|
||||
continue
|
||||
}
|
||||
var raw map[string]interface{}
|
||||
if err := json.Unmarshal([]byte(meta.Endpoints), &raw); err == nil {
|
||||
if err := common.Unmarshal([]byte(meta.Endpoints), &raw); err == nil {
|
||||
for k, v := range raw {
|
||||
switch val := v.(type) {
|
||||
case string:
|
||||
|
||||
@@ -0,0 +1,294 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/dto"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func resetPricingEndpointTestTables(t *testing.T) {
|
||||
t.Helper()
|
||||
originalMemoryCacheEnabled := common.MemoryCacheEnabled
|
||||
common.MemoryCacheEnabled = true
|
||||
require.NoError(t, DB.AutoMigrate(&Channel{}, &Ability{}, &Model{}, &Vendor{}))
|
||||
for _, table := range []string{"abilities", "channels", "models", "vendors"} {
|
||||
require.NoError(t, DB.Exec("DELETE FROM "+table).Error)
|
||||
}
|
||||
InitChannelCache()
|
||||
InvalidatePricingCache()
|
||||
t.Cleanup(func() {
|
||||
for _, table := range []string{"abilities", "channels", "models", "vendors"} {
|
||||
require.NoError(t, DB.Exec("DELETE FROM "+table).Error)
|
||||
}
|
||||
InitChannelCache()
|
||||
InvalidatePricingCache()
|
||||
common.MemoryCacheEnabled = originalMemoryCacheEnabled
|
||||
})
|
||||
}
|
||||
|
||||
func insertPricingEndpointChannel(t *testing.T, channelID int, channelType int, settings dto.ChannelOtherSettings) {
|
||||
t.Helper()
|
||||
channel := &Channel{
|
||||
Id: channelID,
|
||||
Type: channelType,
|
||||
Key: fmt.Sprintf("key-%d", channelID),
|
||||
Status: common.ChannelStatusEnabled,
|
||||
Name: fmt.Sprintf("channel-%d", channelID),
|
||||
}
|
||||
if settings.AdvancedCustom != nil {
|
||||
channel.SetOtherSettings(settings)
|
||||
}
|
||||
require.NoError(t, DB.Create(channel).Error)
|
||||
}
|
||||
|
||||
func insertPricingEndpointAbility(t *testing.T, channelID int, modelName string) {
|
||||
t.Helper()
|
||||
require.NoError(t, DB.Create(&Ability{
|
||||
Group: "default",
|
||||
Model: modelName,
|
||||
ChannelId: channelID,
|
||||
Enabled: true,
|
||||
}).Error)
|
||||
}
|
||||
|
||||
func pricingEndpointAdvancedCustomConfig(routes ...dto.AdvancedCustomRoute) dto.ChannelOtherSettings {
|
||||
return dto.ChannelOtherSettings{
|
||||
AdvancedCustom: &dto.AdvancedCustomConfig{
|
||||
Routes: routes,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func pricingEndpointTypesByModel(t *testing.T) map[string][]constant.EndpointType {
|
||||
t.Helper()
|
||||
InitChannelCache()
|
||||
return pricingEndpointTypesFromPricing(GetPricing())
|
||||
}
|
||||
|
||||
func pricingEndpointTypesFromPricing(pricings []Pricing) map[string][]constant.EndpointType {
|
||||
byModel := make(map[string][]constant.EndpointType)
|
||||
for _, pricing := range pricings {
|
||||
byModel[pricing.ModelName] = pricing.SupportedEndpointTypes
|
||||
}
|
||||
return byModel
|
||||
}
|
||||
|
||||
func TestPricingAdvancedCustomUsesConfiguredEndpointTypes(t *testing.T) {
|
||||
resetPricingEndpointTestTables(t)
|
||||
|
||||
insertPricingEndpointChannel(t, 101, constant.ChannelTypeAdvancedCustom, pricingEndpointAdvancedCustomConfig(
|
||||
dto.AdvancedCustomRoute{
|
||||
IncomingPath: "/v1/chat/completions",
|
||||
UpstreamPath: "/v1/chat/completions",
|
||||
},
|
||||
dto.AdvancedCustomRoute{
|
||||
IncomingPath: "/v1/responses",
|
||||
UpstreamPath: "/v1beta/models/{model}:generateContent",
|
||||
Converter: "openai_responses_to_gemini_generate_content",
|
||||
Models: []string{"re:^gemini-"},
|
||||
},
|
||||
))
|
||||
insertPricingEndpointAbility(t, 101, "gemini-2.5-flash")
|
||||
insertPricingEndpointAbility(t, 101, "gpt-4o")
|
||||
|
||||
byModel := pricingEndpointTypesByModel(t)
|
||||
|
||||
assert.Equal(t, []constant.EndpointType{
|
||||
constant.EndpointTypeOpenAI,
|
||||
constant.EndpointTypeOpenAIResponse,
|
||||
}, byModel["gemini-2.5-flash"])
|
||||
assert.Equal(t, []constant.EndpointType{
|
||||
constant.EndpointTypeOpenAI,
|
||||
}, byModel["gpt-4o"])
|
||||
}
|
||||
|
||||
func TestPricingModelMetadataEndpointsMergeWithAdvancedCustomInference(t *testing.T) {
|
||||
resetPricingEndpointTestTables(t)
|
||||
|
||||
insertPricingEndpointChannel(t, 103, constant.ChannelTypeAdvancedCustom, pricingEndpointAdvancedCustomConfig(
|
||||
dto.AdvancedCustomRoute{
|
||||
IncomingPath: "/v1/responses",
|
||||
UpstreamPath: "/v1beta/models/{model}:generateContent",
|
||||
Converter: "openai_responses_to_gemini_generate_content",
|
||||
Models: []string{"re:^gemini-"},
|
||||
},
|
||||
))
|
||||
insertPricingEndpointAbility(t, 103, "gemini-2.5-flash")
|
||||
require.NoError(t, DB.Create(&Model{
|
||||
ModelName: "gemini-2.5-flash",
|
||||
Endpoints: `{
|
||||
"openai": "/v1/chat/completions"
|
||||
}`,
|
||||
Status: 1,
|
||||
NameRule: NameRuleExact,
|
||||
}).Error)
|
||||
|
||||
byModel := pricingEndpointTypesByModel(t)
|
||||
|
||||
assert.Equal(t, []constant.EndpointType{
|
||||
constant.EndpointTypeOpenAIResponse,
|
||||
constant.EndpointTypeOpenAI,
|
||||
}, byModel["gemini-2.5-flash"])
|
||||
}
|
||||
|
||||
func TestPricingModelMetadataEndpointsCanProvideEndpointWithoutChannelInference(t *testing.T) {
|
||||
resetPricingEndpointTestTables(t)
|
||||
|
||||
insertPricingEndpointChannel(t, 104, constant.ChannelTypeAdvancedCustom, pricingEndpointAdvancedCustomConfig(
|
||||
dto.AdvancedCustomRoute{
|
||||
IncomingPath: "/v1/responses",
|
||||
UpstreamPath: "/v1beta/models/{model}:generateContent",
|
||||
Converter: "openai_responses_to_gemini_generate_content",
|
||||
Models: []string{"re:^gemini-"},
|
||||
},
|
||||
))
|
||||
insertPricingEndpointAbility(t, 104, "metadata-only-model")
|
||||
require.NoError(t, DB.Create(&Model{
|
||||
ModelName: "metadata-only-model",
|
||||
Endpoints: `{
|
||||
"openai": "/v1/chat/completions"
|
||||
}`,
|
||||
Status: 1,
|
||||
NameRule: NameRuleExact,
|
||||
}).Error)
|
||||
|
||||
byModel := pricingEndpointTypesByModel(t)
|
||||
|
||||
assert.Equal(t, []constant.EndpointType{constant.EndpointTypeOpenAI}, byModel["metadata-only-model"])
|
||||
}
|
||||
|
||||
func TestPricingAdvancedCustomMissingConfigFallsBackToChannelType(t *testing.T) {
|
||||
resetPricingEndpointTestTables(t)
|
||||
|
||||
insertPricingEndpointChannel(t, 102, constant.ChannelTypeAdvancedCustom, dto.ChannelOtherSettings{})
|
||||
insertPricingEndpointAbility(t, 102, "gpt-4o")
|
||||
|
||||
byModel := pricingEndpointTypesByModel(t)
|
||||
|
||||
assert.Equal(t, []constant.EndpointType{constant.EndpointTypeOpenAI}, byModel["gpt-4o"])
|
||||
}
|
||||
|
||||
func TestPricingNativeChannelEndpointTypesUnchanged(t *testing.T) {
|
||||
resetPricingEndpointTestTables(t)
|
||||
|
||||
insertPricingEndpointChannel(t, 201, constant.ChannelTypeOpenAI, dto.ChannelOtherSettings{})
|
||||
insertPricingEndpointChannel(t, 202, constant.ChannelTypeGemini, dto.ChannelOtherSettings{})
|
||||
insertPricingEndpointChannel(t, 203, constant.ChannelTypeAnthropic, dto.ChannelOtherSettings{})
|
||||
insertPricingEndpointAbility(t, 201, "gpt-4o")
|
||||
insertPricingEndpointAbility(t, 202, "gemini-2.5-flash")
|
||||
insertPricingEndpointAbility(t, 203, "claude-3-5-sonnet")
|
||||
|
||||
byModel := pricingEndpointTypesByModel(t)
|
||||
|
||||
assert.Equal(t, []constant.EndpointType{constant.EndpointTypeOpenAI}, byModel["gpt-4o"])
|
||||
assert.Equal(t, []constant.EndpointType{constant.EndpointTypeGemini, constant.EndpointTypeOpenAI}, byModel["gemini-2.5-flash"])
|
||||
assert.Equal(t, []constant.EndpointType{constant.EndpointTypeAnthropic, constant.EndpointTypeOpenAI}, byModel["claude-3-5-sonnet"])
|
||||
}
|
||||
|
||||
func TestInitChannelCacheInvalidatesPricingCache(t *testing.T) {
|
||||
resetPricingEndpointTestTables(t)
|
||||
|
||||
insertPricingEndpointChannel(t, 301, constant.ChannelTypeAdvancedCustom, pricingEndpointAdvancedCustomConfig(
|
||||
dto.AdvancedCustomRoute{
|
||||
IncomingPath: "/v1/chat/completions",
|
||||
UpstreamPath: "/v1/chat/completions",
|
||||
},
|
||||
))
|
||||
insertPricingEndpointAbility(t, 301, "gemini-3.5-flash")
|
||||
InitChannelCache()
|
||||
|
||||
initial := pricingEndpointTypesByModel(t)
|
||||
require.Equal(t, []constant.EndpointType{constant.EndpointTypeOpenAI}, initial["gemini-3.5-flash"])
|
||||
|
||||
var channel Channel
|
||||
require.NoError(t, DB.First(&channel, "id = ?", 301).Error)
|
||||
channel.SetOtherSettings(pricingEndpointAdvancedCustomConfig(
|
||||
dto.AdvancedCustomRoute{
|
||||
IncomingPath: "/v1/chat/completions",
|
||||
UpstreamPath: "/v1/chat/completions",
|
||||
},
|
||||
dto.AdvancedCustomRoute{
|
||||
IncomingPath: "/v1/responses",
|
||||
UpstreamPath: "/v1beta/models/{model}:generateContent",
|
||||
Converter: "openai_responses_to_gemini_generate_content",
|
||||
Models: []string{"re:^gemini-"},
|
||||
},
|
||||
))
|
||||
require.NoError(t, DB.Model(&Channel{}).Where("id = ?", 301).Update("settings", channel.OtherSettings).Error)
|
||||
InitChannelCache()
|
||||
|
||||
updated := pricingEndpointTypesByModel(t)
|
||||
assert.Equal(t, []constant.EndpointType{
|
||||
constant.EndpointTypeOpenAI,
|
||||
constant.EndpointTypeOpenAIResponse,
|
||||
}, updated["gemini-3.5-flash"])
|
||||
}
|
||||
|
||||
func TestInitChannelCacheInvalidatesStartupPricingBuiltBeforeChannelCache(t *testing.T) {
|
||||
resetPricingEndpointTestTables(t)
|
||||
|
||||
insertPricingEndpointChannel(t, 302, constant.ChannelTypeAdvancedCustom, pricingEndpointAdvancedCustomConfig(
|
||||
dto.AdvancedCustomRoute{
|
||||
IncomingPath: "/v1/chat/completions",
|
||||
UpstreamPath: "/v1/chat/completions",
|
||||
},
|
||||
dto.AdvancedCustomRoute{
|
||||
IncomingPath: "/v1/responses",
|
||||
UpstreamPath: "/v1beta/models/{model}:generateContent",
|
||||
Converter: "openai_responses_to_gemini_generate_content",
|
||||
Models: []string{"re:^gemini-"},
|
||||
},
|
||||
))
|
||||
insertPricingEndpointAbility(t, 302, "gemini-3.5-flash")
|
||||
|
||||
staleByModel := pricingEndpointTypesFromPricing(GetPricing())
|
||||
require.Equal(t, []constant.EndpointType{constant.EndpointTypeOpenAI}, staleByModel["gemini-3.5-flash"])
|
||||
|
||||
InitChannelCache()
|
||||
|
||||
rebuiltByModel := pricingEndpointTypesFromPricing(GetPricing())
|
||||
assert.Equal(t, []constant.EndpointType{
|
||||
constant.EndpointTypeOpenAI,
|
||||
constant.EndpointTypeOpenAIResponse,
|
||||
}, rebuiltByModel["gemini-3.5-flash"])
|
||||
}
|
||||
|
||||
func TestCacheUpdateChannelSyncsAdvancedCustomConfig(t *testing.T) {
|
||||
resetPricingEndpointTestTables(t)
|
||||
|
||||
channel := &Channel{
|
||||
Id: 401,
|
||||
Type: constant.ChannelTypeAdvancedCustom,
|
||||
Key: "key-401",
|
||||
Status: common.ChannelStatusEnabled,
|
||||
Name: "channel-401",
|
||||
}
|
||||
channel.SetOtherSettings(pricingEndpointAdvancedCustomConfig(dto.AdvancedCustomRoute{
|
||||
IncomingPath: "/v1/responses",
|
||||
UpstreamPath: "/v1beta/models/{model}:generateContent",
|
||||
Converter: "openai_responses_to_gemini_generate_content",
|
||||
}))
|
||||
CacheUpdateChannel(channel)
|
||||
|
||||
require.NotNil(t, channel2advancedCustomConfig[401])
|
||||
assert.Equal(t, []constant.EndpointType{constant.EndpointTypeOpenAIResponse}, channel2advancedCustomConfig[401].SupportedEndpointTypesForModel("gemini-3.5-flash"))
|
||||
|
||||
channel.SetOtherSettings(pricingEndpointAdvancedCustomConfig(dto.AdvancedCustomRoute{
|
||||
IncomingPath: "/v1/chat/completions",
|
||||
UpstreamPath: "/v1/chat/completions",
|
||||
}))
|
||||
CacheUpdateChannel(channel)
|
||||
|
||||
require.NotNil(t, channel2advancedCustomConfig[401])
|
||||
assert.Equal(t, []constant.EndpointType{constant.EndpointTypeOpenAI}, channel2advancedCustomConfig[401].SupportedEndpointTypesForModel("gemini-3.5-flash"))
|
||||
|
||||
channel.Type = constant.ChannelTypeOpenAI
|
||||
CacheUpdateChannel(channel)
|
||||
|
||||
assert.Nil(t, channel2advancedCustomConfig[401])
|
||||
}
|
||||
Reference in New Issue
Block a user