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

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

View File

@@ -19,7 +19,7 @@ const (
/* channel related keys */ /* channel related keys */
ContextKeyChannelId ContextKey = "channel_id" ContextKeyChannelId ContextKey = "channel_id"
ContextKeyChannelName ContextKey = "channel_name" ContextKeyChannelName ContextKey = "channel_name"
ContextKeyChannelCreateTime ContextKey = "channel_create_name" ContextKeyChannelCreateTime ContextKey = "channel_create_time"
ContextKeyChannelBaseUrl ContextKey = "base_url" ContextKeyChannelBaseUrl ContextKey = "base_url"
ContextKeyChannelType ContextKey = "channel_type" ContextKeyChannelType ContextKey = "channel_type"
ContextKeyChannelSetting ContextKey = "channel_setting" ContextKeyChannelSetting ContextKey = "channel_setting"
@@ -29,6 +29,7 @@ const (
ContextKeyChannelModelMapping ContextKey = "model_mapping" ContextKeyChannelModelMapping ContextKey = "model_mapping"
ContextKeyChannelStatusCodeMapping ContextKey = "status_code_mapping" ContextKeyChannelStatusCodeMapping ContextKey = "status_code_mapping"
ContextKeyChannelIsMultiKey ContextKey = "channel_is_multi_key" ContextKeyChannelIsMultiKey ContextKey = "channel_is_multi_key"
ContextKeyChannelMultiKeyIndex ContextKey = "channel_multi_key_index"
ContextKeyChannelKey ContextKey = "channel_key" ContextKeyChannelKey ContextKey = "channel_key"
/* user related keys */ /* user related keys */

View File

@@ -12,6 +12,7 @@ import (
"one-api/model" "one-api/model"
"one-api/service" "one-api/service"
"one-api/setting" "one-api/setting"
"one-api/types"
"strconv" "strconv"
"time" "time"
@@ -415,7 +416,7 @@ func UpdateChannelBalance(c *gin.Context) {
}) })
return return
} }
channel, err := model.GetChannelById(id, true) channel, err := model.CacheGetChannel(id)
if err != nil { if err != nil {
c.JSON(http.StatusOK, gin.H{ c.JSON(http.StatusOK, gin.H{
"success": false, "success": false,
@@ -423,6 +424,13 @@ func UpdateChannelBalance(c *gin.Context) {
}) })
return return
} }
if channel.ChannelInfo.IsMultiKey {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": "多密钥渠道不支持余额查询",
})
return
}
balance, err := updateChannelBalance(channel) balance, err := updateChannelBalance(channel)
if err != nil { if err != nil {
c.JSON(http.StatusOK, gin.H{ c.JSON(http.StatusOK, gin.H{
@@ -436,7 +444,6 @@ func UpdateChannelBalance(c *gin.Context) {
"message": "", "message": "",
"balance": balance, "balance": balance,
}) })
return
} }
func updateAllChannelsBalance() error { func updateAllChannelsBalance() error {
@@ -448,18 +455,21 @@ func updateAllChannelsBalance() error {
if channel.Status != common.ChannelStatusEnabled { if channel.Status != common.ChannelStatusEnabled {
continue continue
} }
if channel.ChannelInfo.IsMultiKey {
continue // skip multi-key channels
}
// TODO: support Azure // TODO: support Azure
//if channel.Type != common.ChannelTypeOpenAI && channel.Type != common.ChannelTypeCustom { //if channel.Type != common.ChannelTypeOpenAI && channel.Type != common.ChannelTypeCustom {
// continue // continue
//} //}
_, err := updateChannelBalance(channel) balance, err := updateChannelBalance(channel)
if err != nil { if err != nil {
continue continue
} else { } else {
// err is nil & balance <= 0 means quota is used up // err is nil & balance <= 0 means quota is used up
//if balance <= 0 { if balance <= 0 {
// service.DisableChannel(channel.Id, channel.Name, "余额不足") service.DisableChannel(*types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey, "", channel.GetAutoBan()), "余额不足")
//} }
} }
time.Sleep(common.RequestInterval) time.Sleep(common.RequestInterval)
} }

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.ContextKeyChannelAutoBan, channel.GetAutoBan())
common.SetContextKey(c, constant.ContextKeyChannelModelMapping, channel.GetModelMapping()) common.SetContextKey(c, constant.ContextKeyChannelModelMapping, channel.GetModelMapping())
common.SetContextKey(c, constant.ContextKeyChannelStatusCodeMapping, channel.GetStatusCodeMapping()) common.SetContextKey(c, constant.ContextKeyChannelStatusCodeMapping, channel.GetStatusCodeMapping())
if channel.ChannelInfo.IsMultiKey {
common.SetContextKey(c, constant.ContextKeyChannelIsMultiKey, true)
} key, index, newAPIError := channel.GetNextEnabledKey()
key, newAPIError := channel.GetNextEnabledKey()
if newAPIError != nil { if newAPIError != nil {
return newAPIError 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)) // c.Request.Header.Set("Authorization", fmt.Sprintf("Bearer %s", key))
common.SetContextKey(c, constant.ContextKeyChannelKey, key) common.SetContextKey(c, constant.ContextKeyChannelKey, key)
common.SetContextKey(c, constant.ContextKeyChannelBaseUrl, channel.GetBaseURL()) common.SetContextKey(c, constant.ContextKeyChannelBaseUrl, channel.GetBaseURL())

View File

@@ -76,17 +76,17 @@ func (channel *Channel) getKeys() []string {
return keys 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 not in multi-key mode, return the original key string directly.
if !channel.ChannelInfo.IsMultiKey { if !channel.ChannelInfo.IsMultiKey {
return channel.Key, nil return channel.Key, 0, nil
} }
// Obtain all keys (split by \n) // Obtain all keys (split by \n)
keys := channel.getKeys() keys := channel.getKeys()
if len(keys) == 0 { if len(keys) == 0 {
// No keys available, return error, should disable the channel // 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 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 no specific status list or none enabled, fall back to first key
if len(enabledIdx) == 0 { if len(enabledIdx) == 0 {
return keys[0], nil return keys[0], 0, nil
} }
switch channel.ChannelInfo.MultiKeyMode { switch channel.ChannelInfo.MultiKeyMode {
case constant.MultiKeyModeRandom: case constant.MultiKeyModeRandom:
// Randomly pick one enabled key // 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: case constant.MultiKeyModePolling:
// Use channel-specific lock to ensure thread-safe polling // Use channel-specific lock to ensure thread-safe polling
lock := getChannelPollingLock(channel.Id) lock := getChannelPollingLock(channel.Id)
@@ -125,7 +126,7 @@ func (channel *Channel) GetNextEnabledKey() (string, *types.NewAPIError) {
channelInfo, err := CacheGetChannelInfo(channel.Id) channelInfo, err := CacheGetChannelInfo(channel.Id)
if err != nil { if err != nil {
return "", types.NewError(err, types.ErrorCodeGetChannelFailed) return "", 0, types.NewError(err, types.ErrorCodeGetChannelFailed)
} }
//println("before polling index:", channel.ChannelInfo.MultiKeyPollingIndex) //println("before polling index:", channel.ChannelInfo.MultiKeyPollingIndex)
defer func() { defer func() {
@@ -148,14 +149,14 @@ func (channel *Channel) GetNextEnabledKey() (string, *types.NewAPIError) {
if getStatus(idx) == common.ChannelStatusEnabled { if getStatus(idx) == common.ChannelStatusEnabled {
// update polling index for next call (point to the next position) // update polling index for next call (point to the next position)
channel.ChannelInfo.MultiKeyPollingIndex = (idx + 1) % len(keys) 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 // Fallback should not happen, but return first enabled key
return keys[enabledIdx[0]], nil return keys[enabledIdx[0]], enabledIdx[0], nil
default: default:
// Unknown mode, default to first enabled key (or original key string) // Unknown mode, default to first enabled key (or original key string)
return keys[enabledIdx[0]], nil return keys[enabledIdx[0]], enabledIdx[0], nil
} }
} }

View File

@@ -1,6 +1,8 @@
package service package service
import ( import (
"one-api/common"
"one-api/constant"
"one-api/dto" "one-api/dto"
relaycommon "one-api/relay/common" relaycommon "one-api/relay/common"
"one-api/relay/helper" "one-api/relay/helper"
@@ -28,6 +30,11 @@ func GenerateTextOtherInfo(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, m
} }
adminInfo := make(map[string]interface{}) adminInfo := make(map[string]interface{})
adminInfo["use_channel"] = ctx.GetStringSlice("use_channel") 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 other["admin_info"] = adminInfo
return other return other
} }

View File

@@ -20,7 +20,7 @@ import {
renderQuota, renderQuota,
stringToColor, stringToColor,
getLogOther, getLogOther,
renderModelTag, renderModelTag
} from '../../helpers'; } from '../../helpers';
import { import {
@@ -356,22 +356,46 @@ const LogsTable = () => {
dataIndex: 'channel', dataIndex: 'channel',
className: isAdmin() ? 'tableShow' : 'tableHiddle', className: isAdmin() ? 'tableShow' : 'tableHiddle',
render: (text, record, index) => { 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 ? ( return isAdminUser ? (
record.type === 0 || record.type === 2 || record.type === 5 ? ( record.type === 0 || record.type === 2 || record.type === 5 ? (
<div> <>
{ {
<Tooltip content={record.channel_name || '[未知]'}> <Tooltip content={record.channel_name || '[未知]'}>
<Tag <Space>
color={colors[parseInt(text) % colors.length]} <Tag
size='large' color={colors[parseInt(text) % colors.length]}
shape='circle' size='large'
> shape='circle'
{' '} >
{text}{' '} {text}
</Tag> </Tag>
{
isMultiKey && (
<Tag
color={'white'}
size='large'
shape='circle'
>
{multiKeyIndex}
</Tag>
)
}
</Space>
</Tooltip> </Tooltip>
} }
</div> </>
) : ( ) : (
<></> <></>
) )