diff --git a/constant/context_key.go b/constant/context_key.go index 36fe39ed..4eaf3d00 100644 --- a/constant/context_key.go +++ b/constant/context_key.go @@ -19,7 +19,7 @@ const ( /* channel related keys */ ContextKeyChannelId ContextKey = "channel_id" ContextKeyChannelName ContextKey = "channel_name" - ContextKeyChannelCreateTime ContextKey = "channel_create_name" + ContextKeyChannelCreateTime ContextKey = "channel_create_time" ContextKeyChannelBaseUrl ContextKey = "base_url" ContextKeyChannelType ContextKey = "channel_type" ContextKeyChannelSetting ContextKey = "channel_setting" @@ -29,6 +29,7 @@ const ( ContextKeyChannelModelMapping ContextKey = "model_mapping" ContextKeyChannelStatusCodeMapping ContextKey = "status_code_mapping" ContextKeyChannelIsMultiKey ContextKey = "channel_is_multi_key" + ContextKeyChannelMultiKeyIndex ContextKey = "channel_multi_key_index" ContextKeyChannelKey ContextKey = "channel_key" /* user related keys */ diff --git a/controller/channel-billing.go b/controller/channel-billing.go index a93f6c64..2c2c25b9 100644 --- a/controller/channel-billing.go +++ b/controller/channel-billing.go @@ -12,6 +12,7 @@ import ( "one-api/model" "one-api/service" "one-api/setting" + "one-api/types" "strconv" "time" @@ -415,7 +416,7 @@ func UpdateChannelBalance(c *gin.Context) { }) return } - channel, err := model.GetChannelById(id, true) + channel, err := model.CacheGetChannel(id) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, @@ -423,6 +424,13 @@ func UpdateChannelBalance(c *gin.Context) { }) return } + if channel.ChannelInfo.IsMultiKey { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": "多密钥渠道不支持余额查询", + }) + return + } balance, err := updateChannelBalance(channel) if err != nil { c.JSON(http.StatusOK, gin.H{ @@ -436,7 +444,6 @@ func UpdateChannelBalance(c *gin.Context) { "message": "", "balance": balance, }) - return } func updateAllChannelsBalance() error { @@ -448,18 +455,21 @@ func updateAllChannelsBalance() error { if channel.Status != common.ChannelStatusEnabled { continue } + if channel.ChannelInfo.IsMultiKey { + continue // skip multi-key channels + } // TODO: support Azure //if channel.Type != common.ChannelTypeOpenAI && channel.Type != common.ChannelTypeCustom { // continue //} - _, err := updateChannelBalance(channel) + balance, err := updateChannelBalance(channel) if err != nil { continue } else { // err is nil & balance <= 0 means quota is used up - //if balance <= 0 { - // service.DisableChannel(channel.Id, channel.Name, "余额不足") - //} + if balance <= 0 { + service.DisableChannel(*types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey, "", channel.GetAutoBan()), "余额不足") + } } time.Sleep(common.RequestInterval) } diff --git a/middleware/distributor.go b/middleware/distributor.go index 7c30daf3..a6889e39 100644 --- a/middleware/distributor.go +++ b/middleware/distributor.go @@ -267,15 +267,15 @@ func SetupContextForSelectedChannel(c *gin.Context, channel *model.Channel, mode common.SetContextKey(c, constant.ContextKeyChannelAutoBan, channel.GetAutoBan()) common.SetContextKey(c, constant.ContextKeyChannelModelMapping, channel.GetModelMapping()) common.SetContextKey(c, constant.ContextKeyChannelStatusCodeMapping, channel.GetStatusCodeMapping()) - if channel.ChannelInfo.IsMultiKey { - common.SetContextKey(c, constant.ContextKeyChannelIsMultiKey, true) - } - - key, newAPIError := channel.GetNextEnabledKey() + key, index, newAPIError := channel.GetNextEnabledKey() if newAPIError != nil { return newAPIError } + if channel.ChannelInfo.IsMultiKey { + common.SetContextKey(c, constant.ContextKeyChannelIsMultiKey, true) + common.SetContextKey(c, constant.ContextKeyChannelMultiKeyIndex, index) + } // c.Request.Header.Set("Authorization", fmt.Sprintf("Bearer %s", key)) common.SetContextKey(c, constant.ContextKeyChannelKey, key) common.SetContextKey(c, constant.ContextKeyChannelBaseUrl, channel.GetBaseURL()) diff --git a/model/channel.go b/model/channel.go index c9f11953..b792c87e 100644 --- a/model/channel.go +++ b/model/channel.go @@ -76,17 +76,17 @@ func (channel *Channel) getKeys() []string { return keys } -func (channel *Channel) GetNextEnabledKey() (string, *types.NewAPIError) { +func (channel *Channel) GetNextEnabledKey() (string, int, *types.NewAPIError) { // If not in multi-key mode, return the original key string directly. if !channel.ChannelInfo.IsMultiKey { - return channel.Key, nil + return channel.Key, 0, nil } // Obtain all keys (split by \n) keys := channel.getKeys() if len(keys) == 0 { // No keys available, return error, should disable the channel - return "", types.NewError(errors.New("no keys available"), types.ErrorCodeChannelNoAvailableKey) + return "", 0, types.NewError(errors.New("no keys available"), types.ErrorCodeChannelNoAvailableKey) } statusList := channel.ChannelInfo.MultiKeyStatusList @@ -110,13 +110,14 @@ func (channel *Channel) GetNextEnabledKey() (string, *types.NewAPIError) { } // If no specific status list or none enabled, fall back to first key if len(enabledIdx) == 0 { - return keys[0], nil + return keys[0], 0, nil } switch channel.ChannelInfo.MultiKeyMode { case constant.MultiKeyModeRandom: // Randomly pick one enabled key - return keys[enabledIdx[rand.Intn(len(enabledIdx))]], nil + selectedIdx := enabledIdx[rand.Intn(len(enabledIdx))] + return keys[selectedIdx], selectedIdx, nil case constant.MultiKeyModePolling: // Use channel-specific lock to ensure thread-safe polling lock := getChannelPollingLock(channel.Id) @@ -125,7 +126,7 @@ func (channel *Channel) GetNextEnabledKey() (string, *types.NewAPIError) { channelInfo, err := CacheGetChannelInfo(channel.Id) if err != nil { - return "", types.NewError(err, types.ErrorCodeGetChannelFailed) + return "", 0, types.NewError(err, types.ErrorCodeGetChannelFailed) } //println("before polling index:", channel.ChannelInfo.MultiKeyPollingIndex) defer func() { @@ -148,14 +149,14 @@ func (channel *Channel) GetNextEnabledKey() (string, *types.NewAPIError) { if getStatus(idx) == common.ChannelStatusEnabled { // update polling index for next call (point to the next position) channel.ChannelInfo.MultiKeyPollingIndex = (idx + 1) % len(keys) - return keys[idx], nil + return keys[idx], idx, nil } } // Fallback – should not happen, but return first enabled key - return keys[enabledIdx[0]], nil + return keys[enabledIdx[0]], enabledIdx[0], nil default: // Unknown mode, default to first enabled key (or original key string) - return keys[enabledIdx[0]], nil + return keys[enabledIdx[0]], enabledIdx[0], nil } } diff --git a/service/log_info_generate.go b/service/log_info_generate.go index affae5fb..020a2ba9 100644 --- a/service/log_info_generate.go +++ b/service/log_info_generate.go @@ -1,6 +1,8 @@ package service import ( + "one-api/common" + "one-api/constant" "one-api/dto" relaycommon "one-api/relay/common" "one-api/relay/helper" @@ -28,6 +30,11 @@ func GenerateTextOtherInfo(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, m } adminInfo := make(map[string]interface{}) adminInfo["use_channel"] = ctx.GetStringSlice("use_channel") + isMultiKey := common.GetContextKeyBool(ctx, constant.ContextKeyChannelIsMultiKey) + if isMultiKey { + adminInfo["is_multi_key"] = true + adminInfo["multi_key_index"] = common.GetContextKeyInt(ctx, constant.ContextKeyChannelMultiKeyIndex) + } other["admin_info"] = adminInfo return other } diff --git a/web/src/components/table/LogsTable.js b/web/src/components/table/LogsTable.js index cc1d2082..a7239357 100644 --- a/web/src/components/table/LogsTable.js +++ b/web/src/components/table/LogsTable.js @@ -20,7 +20,7 @@ import { renderQuota, stringToColor, getLogOther, - renderModelTag, + renderModelTag } from '../../helpers'; import { @@ -356,22 +356,46 @@ const LogsTable = () => { dataIndex: 'channel', className: isAdmin() ? 'tableShow' : 'tableHiddle', render: (text, record, index) => { + let isMultiKey = false + let multiKeyIndex = -1; + let other = getLogOther(record.other); + if (other?.admin_info) { + let adminInfo = other.admin_info; + if (adminInfo?.is_multi_key) { + isMultiKey = true; + multiKeyIndex = adminInfo.multi_key_index; + } + } + return isAdminUser ? ( record.type === 0 || record.type === 2 || record.type === 5 ? ( -
+ <> { - - {' '} - {text}{' '} - + + + {text} + + { + isMultiKey && ( + + {multiKeyIndex} + + ) + } + + } -
+ ) : ( <> )