ソースを参照

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

CaIon 7 ヶ月 前
コミット
20607b0b5c

+ 2 - 1
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 */

+ 16 - 6
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)
 	}

+ 5 - 5
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())

+ 10 - 9
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
 	}
 }
 

+ 7 - 0
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
 }

+ 35 - 11
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 ? (
-            <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>
+            </>
           ) : (
             <></>
           )