feat: support configuring ratio when estimating token number in stream mode

common/constants.go

@@ -48,6 +48,7 @@ var TurnstileSiteKey = ""
 var TurnstileSecretKey = ""
 
 var QuotaForNewUser = 100
+var BytesNumber2Quota = 0.8
 
 const (
 	RoleGuestUser  = 0

controller/relay.go

@@ -2,6 +2,8 @@ package controller
 
 import (
 	"bufio"
+	"bytes"
+	"encoding/json"
 	"fmt"
 	"github.com/gin-gonic/gin"
 	"io"
@@ -11,16 +13,39 @@ import (
 	"strings"
 )
 
+type Message struct {
+	Role    string `json:"role"`
+	Content string `json:"content"`
+}
+
+type TextRequest struct {
+	Model    string    `json:"model"`
+	Messages []Message `json:"messages"`
+	Prompt   string    `json:"prompt"`
+	//Stream   bool      `json:"stream"`
+}
+
+type Usage struct {
+	PromptTokens     int `json:"prompt_tokens"`
+	CompletionTokens int `json:"completion_tokens"`
+	TotalTokens      int `json:"total_tokens"`
+}
+
+type TextResponse struct {
+	Usage `json:"usage"`
+}
+
+type StreamResponse struct {
+	Choices []struct {
+		Delta struct {
+			Content string `json:"content"`
+		} `json:"delta"`
+		FinishReason string `json:"finish_reason"`
+	} `json:"choices"`
+}
+
 func Relay(c *gin.Context) {
-	channelType := c.GetInt("channel")
-	tokenId := c.GetInt("token_id")
-	isUnlimitedTimes := c.GetBool("unlimited_times")
-	baseURL := common.ChannelBaseURLs[channelType]
-	if channelType == common.ChannelTypeCustom {
-		baseURL = c.GetString("base_url")
-	}
-	requestURL := c.Request.URL.String()
-	req, err := http.NewRequest(c.Request.Method, fmt.Sprintf("%s%s", baseURL, requestURL), c.Request.Body)
+	err := relayHelper(c)
 	if err != nil {
 		c.JSON(http.StatusOK, gin.H{
 			"error": gin.H{
@@ -28,42 +53,74 @@ func Relay(c *gin.Context) {
 				"type":    "one_api_error",
 			},
 		})
-		return
 	}
-	//req.Header = c.Request.Header.Clone()
-	// Fix HTTP Decompression failed
-	// https://github.com/stoplightio/prism/issues/1064#issuecomment-824682360
-	//req.Header.Del("Accept-Encoding")
+}
+
+func relayHelper(c *gin.Context) error {
+	channelType := c.GetInt("channel")
+	tokenId := c.GetInt("token_id")
+	consumeQuota := c.GetBool("consume_quota")
+	baseURL := common.ChannelBaseURLs[channelType]
+	if channelType == common.ChannelTypeCustom {
+		baseURL = c.GetString("base_url")
+	}
+	requestBody, err := io.ReadAll(c.Request.Body)
+	if err != nil {
+		return err
+	}
+	err = c.Request.Body.Close()
+	if err != nil {
+		return err
+	}
+	var textRequest TextRequest
+	err = json.Unmarshal(requestBody, &textRequest)
+	if err != nil {
+		return err
+	}
+	// Reset request body
+	c.Request.Body = io.NopCloser(bytes.NewBuffer(requestBody))
+	requestURL := c.Request.URL.String()
+	req, err := http.NewRequest(c.Request.Method, fmt.Sprintf("%s%s", baseURL, requestURL), c.Request.Body)
+	if err != nil {
+		return err
+	}
+	err = c.Request.Body.Close()
+	if err != nil {
+		return err
+	}
 	req.Header.Set("Authorization", c.Request.Header.Get("Authorization"))
 	req.Header.Set("Content-Type", c.Request.Header.Get("Content-Type"))
 	req.Header.Set("Accept", c.Request.Header.Get("Accept"))
 	req.Header.Set("Connection", c.Request.Header.Get("Connection"))
 	client := &http.Client{}
-
 	resp, err := client.Do(req)
 	if err != nil {
-		c.JSON(http.StatusOK, gin.H{
-			"error": gin.H{
-				"message": err.Error(),
-				"type":    "one_api_error",
-			},
-		})
-		return
+		return err
+	}
+	err = req.Body.Close()
+	if err != nil {
+		return err
 	}
 
+	var textResponse TextResponse
+	isStream := resp.Header.Get("Content-Type") == "text/event-stream"
+	var streamResponseText string
+
 	defer func() {
-		err := req.Body.Close()
-		if err != nil {
-			common.SysError("Error closing request body: " + err.Error())
-		}
-		if !isUnlimitedTimes && requestURL == "/v1/chat/completions" {
-			err := model.DecreaseTokenRemainTimesById(tokenId)
+		if consumeQuota {
+			quota := 0
+			if isStream {
+				quota = int(float64(len(streamResponseText)) * common.BytesNumber2Quota)
+			} else {
+				quota = textResponse.Usage.TotalTokens
+			}
+			err := model.ConsumeTokenQuota(tokenId, quota)
 			if err != nil {
-				common.SysError("Error decreasing token remain times: " + err.Error())
+				common.SysError("Error consuming token remain quota: " + err.Error())
 			}
 		}
 	}()
-	isStream := resp.Header.Get("Content-Type") == "text/event-stream"
+
 	if isStream {
 		scanner := bufio.NewScanner(resp.Body)
 		scanner.Split(func(data []byte, atEOF bool) (advance int, token []byte, err error) {
@@ -87,6 +144,18 @@ func Relay(c *gin.Context) {
 			for scanner.Scan() {
 				data := scanner.Text()
 				dataChan <- data
+				data = data[6:]
+				if data != "[DONE]" {
+					var streamResponse StreamResponse
+					err = json.Unmarshal([]byte(data), &streamResponse)
+					if err != nil {
+						common.SysError("Error unmarshalling stream response: " + err.Error())
+						return
+					}
+					for _, choice := range streamResponse.Choices {
+						streamResponseText += choice.Delta.Content
+					}
+				}
 			}
 			stopChan <- true
 		}()
@@ -103,20 +172,46 @@ func Relay(c *gin.Context) {
 				return false
 			}
 		})
-		return
+		err = resp.Body.Close()
+		if err != nil {
+			return err
+		}
+		return nil
 	} else {
 		for k, v := range resp.Header {
 			c.Writer.Header().Set(k, v[0])
 		}
+		responseBody, err := io.ReadAll(resp.Body)
+		if err != nil {
+			return err
+		}
+		err = resp.Body.Close()
+		if err != nil {
+			return err
+		}
+		err = json.Unmarshal(responseBody, &textResponse)
+		if err != nil {
+			return err
+		}
+		// Reset response body
+		resp.Body = io.NopCloser(bytes.NewBuffer(responseBody))
 		_, err = io.Copy(c.Writer, resp.Body)
 		if err != nil {
-			c.JSON(http.StatusOK, gin.H{
-				"error": gin.H{
-					"message": err.Error(),
-					"type":    "one_api_error",
-				},
-			})
-			return
+			return err
 		}
+		err = resp.Body.Close()
+		if err != nil {
+			return err
+		}
+		return nil
 	}
 }
+
+func RelayNotImplemented(c *gin.Context) {
+	c.JSON(http.StatusOK, gin.H{
+		"error": gin.H{
+			"message": "Not Implemented",
+			"type":    "one_api_error",
+		},
+	})
+}

controller/token.go

@@ -102,8 +102,8 @@ func AddToken(c *gin.Context) {
 		ExpiredTime:  token.ExpiredTime,
 	}
 	if isAdmin {
-		cleanToken.RemainTimes = token.RemainTimes
-		cleanToken.UnlimitedTimes = token.UnlimitedTimes
+		cleanToken.RemainQuota = token.RemainQuota
+		cleanToken.UnlimitedQuota = token.UnlimitedQuota
 	} else {
 		userId := c.GetInt("id")
 		quota, err := model.GetUserQuota(userId)
@@ -115,7 +115,7 @@ func AddToken(c *gin.Context) {
 			return
 		}
 		if quota > 0 {
-			cleanToken.RemainTimes = quota
+			cleanToken.RemainQuota = quota
 		}
 	}
 	err = cleanToken.Insert()
@@ -128,7 +128,7 @@ func AddToken(c *gin.Context) {
 	}
 	if !isAdmin {
 		// update user quota
-		err = model.DecreaseUserQuota(c.GetInt("id"), cleanToken.RemainTimes)
+		err = model.DecreaseUserQuota(c.GetInt("id"), cleanToken.RemainQuota)
 	}
 	c.JSON(http.StatusOK, gin.H{
 		"success": true,
@@ -184,7 +184,7 @@ func UpdateToken(c *gin.Context) {
 			})
 			return
 		}
-		if cleanToken.Status == common.TokenStatusExhausted && cleanToken.RemainTimes <= 0 && !cleanToken.UnlimitedTimes {
+		if cleanToken.Status == common.TokenStatusExhausted && cleanToken.RemainQuota <= 0 && !cleanToken.UnlimitedQuota {
 			c.JSON(http.StatusOK, gin.H{
 				"success": false,
 				"message": "令牌可用次数已用尽，无法启用，请先修改令牌剩余次数，或者设置为无限次数",
@@ -199,8 +199,8 @@ func UpdateToken(c *gin.Context) {
 		cleanToken.Name = token.Name
 		cleanToken.ExpiredTime = token.ExpiredTime
 		if isAdmin {
-			cleanToken.RemainTimes = token.RemainTimes
-			cleanToken.UnlimitedTimes = token.UnlimitedTimes
+			cleanToken.RemainQuota = token.RemainQuota
+			cleanToken.UnlimitedQuota = token.UnlimitedQuota
 		}
 	}
 	err = cleanToken.Update()

middleware/auth.go

@@ -110,7 +110,17 @@ func TokenAuth() func(c *gin.Context) {
 		}
 		c.Set("id", token.UserId)
 		c.Set("token_id", token.Id)
-		c.Set("unlimited_times", token.UnlimitedTimes)
+		requestURL := c.Request.URL.String()
+		consumeQuota := false
+		switch requestURL {
+		case "/v1/chat/completions":
+			consumeQuota = !token.UnlimitedQuota
+		case "/v1/completions":
+			consumeQuota = !token.UnlimitedQuota
+		case "/v1/edits":
+			consumeQuota = !token.UnlimitedQuota
+		}
+		c.Set("consume_quota", consumeQuota)
 		if len(parts) > 1 {
 			if model.IsAdmin(token.UserId) {
 				c.Set("channelId", parts[1])

model/option.go

@@ -47,6 +47,7 @@ func InitOptionMap() {
 	common.OptionMap["TurnstileSiteKey"] = ""
 	common.OptionMap["TurnstileSecretKey"] = ""
 	common.OptionMap["QuotaForNewUser"] = strconv.Itoa(common.QuotaForNewUser)
+	common.OptionMap["BytesNumber2Quota"] = strconv.FormatFloat(common.BytesNumber2Quota, 'f', -1, 64)
 	common.OptionMap["TopUpLink"] = common.TopUpLink
 	common.OptionMapRWMutex.Unlock()
 	options, _ := AllOption()
@@ -135,6 +136,8 @@ func updateOptionMap(key string, value string) {
 		common.TurnstileSecretKey = value
 	case "QuotaForNewUser":
 		common.QuotaForNewUser, _ = strconv.Atoi(value)
+	case "BytesNumber2Quota":
+		common.BytesNumber2Quota, _ = strconv.ParseFloat(value, 64)
 	case "TopUpLink":
 		common.TopUpLink = value
 	}

model/redemption.go

@@ -55,7 +55,7 @@ func Redeem(key string, tokenId int) (quota int, err error) {
 	if redemption.Status != common.RedemptionCodeStatusEnabled {
 		return 0, errors.New("该兑换码已被使用")
 	}
-	err = TopUpToken(tokenId, redemption.Quota)
+	err = TopUpTokenQuota(tokenId, redemption.Quota)
 	if err != nil {
 		return 0, err
 	}

model/token.go

@@ -17,8 +17,8 @@ type Token struct {
 	CreatedTime    int64  `json:"created_time" gorm:"bigint"`
 	AccessedTime   int64  `json:"accessed_time" gorm:"bigint"`
 	ExpiredTime    int64  `json:"expired_time" gorm:"bigint;default:-1"` // -1 means never expired
-	RemainTimes    int    `json:"remain_times" gorm:"default:0"`
-	UnlimitedTimes bool   `json:"unlimited_times" gorm:"default:false"`
+	RemainQuota    int    `json:"remain_quota" gorm:"default:0"`
+	UnlimitedQuota bool   `json:"unlimited_quota" gorm:"default:false"`
 }
 
 func GetAllUserTokens(userId int, startIdx int, num int) ([]*Token, error) {
@@ -52,13 +52,13 @@ func ValidateUserToken(key string) (token *Token, err error) {
 			}
 			return nil, errors.New("该 token 已过期")
 		}
-		if !token.UnlimitedTimes && token.RemainTimes <= 0 {
+		if !token.UnlimitedQuota && token.RemainQuota <= 0 {
 			token.Status = common.TokenStatusExhausted
 			err := token.SelectUpdate()
 			if err != nil {
 				common.SysError("更新 token 状态失败：" + err.Error())
 			}
-			return nil, errors.New("该 token 可用次数已用尽")
+			return nil, errors.New("该 token 额度已用尽")
 		}
 		go func() {
 			token.AccessedTime = common.GetTimestamp()
@@ -91,7 +91,7 @@ func (token *Token) Insert() error {
 // Update Make sure your token's fields is completed, because this will update non-zero values
 func (token *Token) Update() error {
 	var err error
-	err = DB.Model(token).Select("name", "status", "expired_time", "remain_times", "unlimited_times").Updates(token).Error
+	err = DB.Model(token).Select("name", "status", "expired_time", "remain_quota", "unlimited_quota").Updates(token).Error
 	return err
 }
 
@@ -119,12 +119,12 @@ func DeleteTokenById(id int, userId int) (err error) {
 	return token.Delete()
 }
 
-func DecreaseTokenRemainTimesById(id int) (err error) {
-	err = DB.Model(&Token{}).Where("id = ?", id).Update("remain_times", gorm.Expr("remain_times - ?", 1)).Error
+func ConsumeTokenQuota(id int, quota int) (err error) {
+	err = DB.Model(&Token{}).Where("id = ?", id).Update("remain_quota", gorm.Expr("remain_quota - ?", quota)).Error
 	return err
 }
 
-func TopUpToken(id int, times int) (err error) {
-	err = DB.Model(&Token{}).Where("id = ?", id).Update("remain_times", gorm.Expr("remain_times + ?", times)).Error
+func TopUpTokenQuota(id int, quota int) (err error) {
+	err = DB.Model(&Token{}).Where("id = ?", id).Update("remain_quota", gorm.Expr("remain_quota + ?", quota)).Error
 	return err
 }

router/relay-router.go

@@ -7,12 +7,35 @@ import (
 )
 
 func SetRelayRouter(router *gin.Engine) {
+	// https://platform.openai.com/docs/api-reference/introduction
 	relayV1Router := router.Group("/v1")
 	relayV1Router.Use(middleware.TokenAuth(), middleware.Distribute())
 	{
-		relayV1Router.Any("/*path", controller.Relay)
+		relayV1Router.GET("/models", controller.Relay)
+		relayV1Router.GET("/models/:model", controller.Relay)
+		relayV1Router.POST("/completions", controller.RelayNotImplemented)
+		relayV1Router.POST("/chat/completions", controller.Relay)
+		relayV1Router.POST("/edits", controller.RelayNotImplemented)
+		relayV1Router.POST("/images/generations", controller.RelayNotImplemented)
+		relayV1Router.POST("/images/edits", controller.RelayNotImplemented)
+		relayV1Router.POST("/images/variations", controller.RelayNotImplemented)
+		relayV1Router.POST("/embeddings", controller.RelayNotImplemented)
+		relayV1Router.POST("/audio/transcriptions", controller.RelayNotImplemented)
+		relayV1Router.POST("/audio/translations", controller.RelayNotImplemented)
+		relayV1Router.GET("/files", controller.RelayNotImplemented)
+		relayV1Router.POST("/files", controller.RelayNotImplemented)
+		relayV1Router.DELETE("/files/:id", controller.RelayNotImplemented)
+		relayV1Router.GET("/files/:id", controller.RelayNotImplemented)
+		relayV1Router.GET("/files/:id/content", controller.RelayNotImplemented)
+		relayV1Router.POST("/fine-tunes", controller.RelayNotImplemented)
+		relayV1Router.GET("/fine-tunes", controller.RelayNotImplemented)
+		relayV1Router.GET("/fine-tunes/:id", controller.RelayNotImplemented)
+		relayV1Router.POST("/fine-tunes/:id/cancel", controller.RelayNotImplemented)
+		relayV1Router.GET("/fine-tunes/:id/events", controller.RelayNotImplemented)
+		relayV1Router.DELETE("/models/:model", controller.RelayNotImplemented)
+		relayV1Router.POST("/moderations", controller.RelayNotImplemented)
 	}
-	relayDashboardRouter := router.Group("/dashboard")
+	relayDashboardRouter := router.Group("/dashboard") // TODO: return system's own token info
 	relayDashboardRouter.Use(middleware.TokenAuth(), middleware.Distribute())
 	{
 		relayDashboardRouter.Any("/*path", controller.Relay)

web/src/components/SystemSetting.js

@@ -25,6 +25,7 @@ const SystemSetting = () => {
     TurnstileSecretKey: '',
     RegisterEnabled: '',
     QuotaForNewUser: 0,
+    BytesNumber2Quota: 0.8,
     TopUpLink: ''
   });
   let originInputs = {};
@@ -90,6 +91,7 @@ const SystemSetting = () => {
       name === 'TurnstileSiteKey' ||
       name === 'TurnstileSecretKey' ||
       name === 'QuotaForNewUser' ||
+      name === 'BytesNumber2Quota' ||
       name === 'TopUpLink'
     ) {
       setInputs((inputs) => ({ ...inputs, [name]: value }));
@@ -107,6 +109,9 @@ const SystemSetting = () => {
     if (originInputs['QuotaForNewUser'] !== inputs.QuotaForNewUser) {
       await updateOption('QuotaForNewUser', inputs.QuotaForNewUser);
     }
+    if (originInputs['BytesNumber2Quota'] !== inputs.BytesNumber2Quota) {
+      await updateOption('BytesNumber2Quota', inputs.BytesNumber2Quota);
+    }
     if (originInputs['TopUpLink'] !== inputs.TopUpLink) {
       await updateOption('TopUpLink', inputs.TopUpLink);
     }
@@ -255,6 +260,17 @@ const SystemSetting = () => {
               min='0'
               placeholder='例如：100'
             />
+            <Form.Input
+              label='Stream 模式下估算 token 时所使用的倍率'
+              name='BytesNumber2Quota'
+              onChange={handleInputChange}
+              autoComplete='off'
+              value={inputs.BytesNumber2Quota}
+              type='number'
+              step='0.01'
+              min='0'
+              placeholder='例如：0.8'
+            />
             <Form.Input
               label='充值链接'
               name='TopUpLink'

web/src/components/TokensTable.js

@@ -164,7 +164,7 @@ const TokensTable = () => {
       showSuccess('充值成功！');
       let newTokens = [...tokens];
       let realIdx = (activePage - 1) * ITEMS_PER_PAGE + targetTokenIdx;
-      newTokens[realIdx].remain_times += data;
+      newTokens[realIdx].remain_quota += data;
       setTokens(newTokens);
       setRedemptionCode('');
       setShowTopUpModal(false);
@@ -217,10 +217,10 @@ const TokensTable = () => {
             <Table.HeaderCell
               style={{ cursor: 'pointer' }}
               onClick={() => {
-                sortToken('remain_times');
+                sortToken('remain_quota');
               }}
             >
-              剩余次数
+              额度
             </Table.HeaderCell>
             <Table.HeaderCell
               style={{ cursor: 'pointer' }}
@@ -255,7 +255,7 @@ const TokensTable = () => {
                   <Table.Cell>{token.id}</Table.Cell>
                   <Table.Cell>{token.name ? token.name : '无'}</Table.Cell>
                   <Table.Cell>{renderStatus(token.status)}</Table.Cell>
-                  <Table.Cell>{token.unlimited_times ? '无限制' : token.remain_times}</Table.Cell>
+                  <Table.Cell>{token.unlimited_quota ? '无限制' : token.remain_quota}</Table.Cell>
                   <Table.Cell>{renderTimestamp(token.created_time)}</Table.Cell>
                   <Table.Cell>{token.expired_time === -1 ? '永不过期' : renderTimestamp(token.expired_time)}</Table.Cell>
                   <Table.Cell>

web/src/pages/Token/EditToken.js

@@ -10,13 +10,13 @@ const EditToken = () => {
   const [loading, setLoading] = useState(isEdit);
   const originInputs = {
     name: '',
-    remain_times: 0,
+    remain_quota: 0,
     expired_time: -1,
-    unlimited_times: false
+    unlimited_quota: false
   };
   const isAdminUser = isAdmin();
   const [inputs, setInputs] = useState(originInputs);
-  const { name, remain_times, expired_time, unlimited_times } = inputs;
+  const { name, remain_quota, expired_time, unlimited_quota } = inputs;
 
   const handleInputChange = (e, { name, value }) => {
     setInputs((inputs) => ({ ...inputs, [name]: value }));
@@ -37,8 +37,8 @@ const EditToken = () => {
     }
   };
 
-  const setUnlimitedTimes = () => {
-    setInputs({ ...inputs, unlimited_times: !unlimited_times });
+  const setUnlimitedQuota = () => {
+    setInputs({ ...inputs, unlimited_quota: !unlimited_quota });
   };
 
   const loadToken = async () => {
@@ -63,7 +63,7 @@ const EditToken = () => {
   const submit = async () => {
     if (!isEdit && inputs.name === '') return;
     let localInputs = inputs;
-    localInputs.remain_times = parseInt(localInputs.remain_times);
+    localInputs.remain_quota = parseInt(localInputs.remain_quota);
     if (localInputs.expired_time !== -1) {
       let time = Date.parse(localInputs.expired_time);
       if (isNaN(time)) {
@@ -111,19 +111,19 @@ const EditToken = () => {
             isAdminUser && <>
               <Form.Field>
                 <Form.Input
-                  label='剩余次数'
-                  name='remain_times'
-                  placeholder={'请输入剩余次数'}
+                  label='额度'
+                  name='remain_quota'
+                  placeholder={'请输入额度'}
                   onChange={handleInputChange}
-                  value={remain_times}
+                  value={remain_quota}
                   autoComplete='off'
                   type='number'
-                  disabled={unlimited_times}
+                  disabled={unlimited_quota}
                 />
               </Form.Field>
               <Button type={'button'} style={{marginBottom: '14px'}} onClick={() => {
-                setUnlimitedTimes();
-              }}>{unlimited_times ? '取消无限次' : '设置为无限次'}</Button>
+                setUnlimitedQuota();
+              }}>{unlimited_quota ? '取消无限额度' : '设置为无限额度'}</Button>
             </>
           }
           <Form.Field>