[pull] main from QuantumNous:main

Dockerfile

@@ -2,6 +2,7 @@ FROM oven/bun:latest AS builder
 
 WORKDIR /build
 COPY web/package.json .
+COPY web/bun.lock .
 RUN bun install
 COPY ./web .
 COPY ./VERSION .

common/api_type.go

@@ -63,6 +63,8 @@ func ChannelType2APIType(channelType int) (int, bool) {
 		apiType = constant.APITypeXai
 	case constant.ChannelTypeCoze:
 		apiType = constant.APITypeCoze
+	case constant.ChannelTypeJimeng:
+		apiType = constant.APITypeJimeng
 	}
 	if apiType == -1 {
 		return constant.APITypeOpenAI, false

common/constants.go

@@ -193,3 +193,9 @@ const (
 	ChannelStatusManuallyDisabled = 2 // also don't use 0
 	ChannelStatusAutoDisabled     = 3
 )
+
+const (
+	TopUpStatusPending = "pending"
+	TopUpStatusSuccess = "success"
+	TopUpStatusExpired = "expired"
+)

common/gin.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"github.com/gin-gonic/gin"
 	"io"
+	"net/http"
 	"one-api/constant"
 	"strings"
 	"time"
@@ -32,7 +33,7 @@ func UnmarshalBodyReusable(c *gin.Context, v any) error {
 	}
 	contentType := c.Request.Header.Get("Content-Type")
 	if strings.HasPrefix(contentType, "application/json") {
-		err = UnmarshalJson(requestBody, &v)
+		err = Unmarshal(requestBody, &v)
 	} else {
 		// skip for now
 		// TODO: someday non json request have variant model, we will need to implementation this
@@ -86,3 +87,25 @@ func GetContextKeyType[T any](c *gin.Context, key constant.ContextKey) (T, bool)
 	var t T
 	return t, false
 }
+
+func ApiError(c *gin.Context, err error) {
+	c.JSON(http.StatusOK, gin.H{
+		"success": false,
+		"message": err.Error(),
+	})
+}
+
+func ApiErrorMsg(c *gin.Context, msg string) {
+	c.JSON(http.StatusOK, gin.H{
+		"success": false,
+		"message": msg,
+	})
+}
+
+func ApiSuccess(c *gin.Context, data any) {
+	c.JSON(http.StatusOK, gin.H{
+		"success": true,
+		"message": "",
+		"data":    data,
+	})
+}

common/hash.go

@@ -0,0 +1,34 @@
+package common
+
+import (
+	"crypto/hmac"
+	"crypto/sha1"
+	"crypto/sha256"
+	"encoding/hex"
+)
+
+func Sha256Raw(data []byte) []byte {
+	h := sha256.New()
+	h.Write(data)
+	return h.Sum(nil)
+}
+
+func Sha1Raw(data []byte) []byte {
+	h := sha1.New()
+	h.Write(data)
+	return h.Sum(nil)
+}
+
+func Sha1(data []byte) string {
+	return hex.EncodeToString(Sha1Raw(data))
+}
+
+func HmacSha256Raw(message, key []byte) []byte {
+	h := hmac.New(sha256.New, key)
+	h.Write(message)
+	return h.Sum(nil)
+}
+
+func HmacSha256(message, key string) string {
+	return hex.EncodeToString(HmacSha256Raw([]byte(message), []byte(key)))
+}

common/json.go

@@ -5,7 +5,7 @@ import (
 	"encoding/json"
 )
 
-func UnmarshalJson(data []byte, v any) error {
+func Unmarshal(data []byte, v any) error {
 	return json.Unmarshal(data, v)
 }
 
@@ -17,6 +17,6 @@ func DecodeJson(reader *bytes.Reader, v any) error {
 	return json.NewDecoder(reader).Decode(v)
 }
 
-func EncodeJson(v any) ([]byte, error) {
+func Marshal(v any) ([]byte, error) {
 	return json.Marshal(v)
 }

common/logger.go

@@ -75,6 +75,9 @@ func logHelper(ctx context.Context, level string, msg string) {
 		writer = gin.DefaultWriter
 	}
 	id := ctx.Value(RequestIdKey)
+	if id == nil {
+		id = "SYSTEM"
+	}
 	now := time.Now()
 	_, _ = fmt.Fprintf(writer, "[%s] %v | %s | %s \n", level, now.Format("2006/01/02 - 15:04:05"), id, msg)
 	logCount++ // we don't need accurate count, so no lock here

common/page_info.go

@@ -1,15 +1,14 @@
 package common
 
 import (
-	"github.com/gin-gonic/gin"
 	"strconv"
+
+	"github.com/gin-gonic/gin"
 )
 
 type PageInfo struct {
-	Page           int   `json:"page"`            // page num 页码
-	PageSize       int   `json:"page_size"`       // page size 页大小
-	StartTimestamp int64 `json:"start_timestamp"` // 秒级
-	EndTimestamp   int64 `json:"end_timestamp"`   // 秒级
+	Page     int `json:"page"`      // page num 页码
+	PageSize int `json:"page_size"` // page size 页大小
 
 	Total int `json:"total"` // 总条数，后设置
 	Items any `json:"items"` // 数据，后设置
@@ -39,11 +38,14 @@ func (p *PageInfo) SetItems(items any) {
 	p.Items = items
 }
 
-func GetPageQuery(c *gin.Context) (*PageInfo, error) {
+func GetPageQuery(c *gin.Context) *PageInfo {
 	pageInfo := &PageInfo{}
-	err := c.BindQuery(pageInfo)
-	if err != nil {
-		return nil, err
+	// 手动获取并处理每个参数
+	if page, err := strconv.Atoi(c.Query("page")); err == nil {
+		pageInfo.Page = page
+	}
+	if pageSize, err := strconv.Atoi(c.Query("page_size")); err == nil {
+		pageInfo.PageSize = pageSize
 	}
 	if pageInfo.Page < 1 {
 		// 兼容
@@ -56,7 +58,25 @@ func GetPageQuery(c *gin.Context) (*PageInfo, error) {
 	}
 
 	if pageInfo.PageSize == 0 {
-		pageInfo.PageSize = ItemsPerPage
+		// 兼容
+		pageSize, _ := strconv.Atoi(c.Query("ps"))
+		if pageSize != 0 {
+			pageInfo.PageSize = pageSize
+		}
+		if pageInfo.PageSize == 0 {
+			pageSize, _ = strconv.Atoi(c.Query("size")) // token page
+			if pageSize != 0 {
+				pageInfo.PageSize = pageSize
+			}
+		}
+		if pageInfo.PageSize == 0 {
+			pageInfo.PageSize = ItemsPerPage
+		}
 	}
-	return pageInfo, nil
+
+	if pageInfo.PageSize > 100 {
+		pageInfo.PageSize = 100
+	}
+
+	return pageInfo
 }

common/str.go

@@ -32,16 +32,30 @@ func MapToJsonStr(m map[string]interface{}) string {
 	return string(bytes)
 }
 
-func StrToMap(str string) map[string]interface{} {
+func StrToMap(str string) (map[string]interface{}, error) {
 	m := make(map[string]interface{})
-	err := json.Unmarshal([]byte(str), &m)
+	err := Unmarshal([]byte(str), &m)
 	if err != nil {
-		return nil
+		return nil, err
 	}
-	return m
+	return m, nil
 }
 
-func IsJsonStr(str string) bool {
+func StrToJsonArray(str string) ([]interface{}, error) {
+	var js []interface{}
+	err := json.Unmarshal([]byte(str), &js)
+	if err != nil {
+		return nil, err
+	}
+	return js, nil
+}
+
+func IsJsonArray(str string) bool {
+	var js []interface{}
+	return json.Unmarshal([]byte(str), &js) == nil
+}
+
+func IsJsonObject(str string) bool {
 	var js map[string]interface{}
 	return json.Unmarshal([]byte(str), &js) == nil
 }

constant/api_type.go

@@ -30,5 +30,6 @@ const (
 	APITypeXinference
 	APITypeXai
 	APITypeCoze
+	APITypeJimeng
 	APITypeDummy // this one is only for count, do not add any channel after this
 )

constant/context_key.go

@@ -17,11 +17,20 @@ const (
 	ContextKeyTokenModelLimit        ContextKey = "token_model_limit"
 
 	/* channel related keys */
-	ContextKeyBaseUrl        ContextKey = "base_url"
-	ContextKeyChannelType    ContextKey = "channel_type"
-	ContextKeyChannelId      ContextKey = "channel_id"
-	ContextKeyChannelSetting ContextKey = "channel_setting"
-	ContextKeyParamOverride  ContextKey = "param_override"
+	ContextKeyChannelId                ContextKey = "channel_id"
+	ContextKeyChannelName              ContextKey = "channel_name"
+	ContextKeyChannelCreateTime        ContextKey = "channel_create_time"
+	ContextKeyChannelBaseUrl           ContextKey = "base_url"
+	ContextKeyChannelType              ContextKey = "channel_type"
+	ContextKeyChannelSetting           ContextKey = "channel_setting"
+	ContextKeyChannelParamOverride     ContextKey = "param_override"
+	ContextKeyChannelOrganization      ContextKey = "channel_organization"
+	ContextKeyChannelAutoBan           ContextKey = "auto_ban"
+	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 */
 	ContextKeyUserId      ContextKey = "id"

constant/multi_key_mode.go

@@ -0,0 +1,8 @@
+package constant
+
+type MultiKeyMode string
+
+const (
+	MultiKeyModeRandom  MultiKeyMode = "random"  // 随机
+	MultiKeyModePolling MultiKeyMode = "polling" // 轮询
+)

controller/channel-billing.go

@@ -4,17 +4,19 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
-	"github.com/shopspring/decimal"
 	"io"
 	"net/http"
 	"one-api/common"
 	"one-api/constant"
 	"one-api/model"
 	"one-api/service"
 	"one-api/setting"
+	"one-api/types"
 	"strconv"
 	"time"
 
+	"github.com/shopspring/decimal"
+
 	"github.com/gin-gonic/gin"
 )
 
@@ -409,34 +411,31 @@ func updateChannelBalance(channel *model.Channel) (float64, error) {
 func UpdateChannelBalance(c *gin.Context) {
 	id, err := strconv.Atoi(c.Param("id"))
 	if err != nil {
-		c.JSON(http.StatusOK, gin.H{
-			"success": false,
-			"message": err.Error(),
-		})
+		common.ApiError(c, err)
 		return
 	}
-	channel, err := model.GetChannelById(id, true)
+	channel, err := model.CacheGetChannel(id)
 	if err != nil {
+		common.ApiError(c, err)
+		return
+	}
+	if channel.ChannelInfo.IsMultiKey {
 		c.JSON(http.StatusOK, gin.H{
 			"success": false,
-			"message": err.Error(),
+			"message": "多密钥渠道不支持余额查询",
 		})
 		return
 	}
 	balance, err := updateChannelBalance(channel)
 	if err != nil {
-		c.JSON(http.StatusOK, gin.H{
-			"success": false,
-			"message": err.Error(),
-		})
+		common.ApiError(c, err)
 		return
 	}
 	c.JSON(http.StatusOK, gin.H{
 		"success": true,
 		"message": "",
 		"balance": balance,
 	})
-	return
 }
 
 func updateAllChannelsBalance() error {
@@ -448,6 +447,9 @@ 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
@@ -458,7 +460,7 @@ func updateAllChannelsBalance() error {
 		} else {
 			// err is nil & balance <= 0 means quota is used up
 			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)
@@ -470,10 +472,7 @@ func UpdateAllChannelsBalance(c *gin.Context) {
 	// TODO: make it async
 	err := updateAllChannelsBalance()
 	if err != nil {
-		c.JSON(http.StatusOK, gin.H{
-			"success": false,
-			"message": err.Error(),
-		})
+		common.ApiError(c, err)
 		return
 	}
 	c.JSON(http.StatusOK, gin.H{