oadmin/new-api
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

✨ feat(logs): add multi-key support in LogsTable and enhance log info generation

Browse Source
This commit is contained in:
CaIon
2025-07-12 15:14:55 +08:00
parent 50b76f4466
commit 20607b0b5c
6 changed files with 75 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
constant/context_key.go
View File

@@ -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 */

22
controller/channel-billing.go
View File

@@ -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)
}

10
middleware/distributor.go
View File

@@ -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())

19
model/channel.go
View File

@@ -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
}
}

7
service/log_info_generate.go
View File

@@ -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
}

46
web/src/components/table/LogsTable.js
View File

@@ -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 ? (
<div>
<>
{
<Tooltip content={record.channel_name || '[未知]'}>
<Tag
color={colors[parseInt(text) % colors.length]}
size='large'
shape='circle'
>
{' '}
{text}{' '}
</Tag>
<Space>
<Tag
color={colors[parseInt(text) % colors.length]}
size='large'
shape='circle'
>
{text}
</Tag>
{
isMultiKey && (
<Tag
color={'white'}
size='large'
shape='circle'
>
{multiKeyIndex}
</Tag>
)
}
</Space>
</Tooltip>
}
</div>
</>
) : (
<></>
)