fix: shouldn't close c.Request.Body too soon (close #35)

common/constants.go

@@ -48,7 +48,11 @@ var TurnstileSiteKey = ""
 var TurnstileSecretKey = ""
 
 var QuotaForNewUser = 100
-var BytesNumber2Quota = 0.8
+
+// https://platform.openai.com/docs/models/model-endpoint-compatibility
+var RatioGPT3dot5 float64 = 2
+var RatioGPT4 float64 = 30
+var RatioGPT4_32k float64 = 60
 
 const (
 	RoleGuestUser  = 0

controller/relay.go

@@ -6,6 +6,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"github.com/gin-gonic/gin"
+	"github.com/pkoukk/tiktoken-go"
 	"io"
 	"net/http"
 	"one-api/common"
@@ -44,6 +45,13 @@ type StreamResponse struct {
 	} `json:"choices"`
 }
 
+var tokenEncoder, _ = tiktoken.GetEncoding("cl100k_base")
+
+func countToken(text string) int {
+	token := tokenEncoder.Encode(text, nil, nil)
+	return len(token)
+}
+
 func Relay(c *gin.Context) {
 	err := relayHelper(c)
 	if err != nil {
@@ -64,30 +72,28 @@ func relayHelper(c *gin.Context) error {
 	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
+	if consumeQuota {
+		requestBody, err := io.ReadAll(c.Request.Body)
+		if err != nil {
+			return err
+		}
+		err = c.Request.Body.Close()
+		if err != nil {
+			return err
+		}
+		err = json.Unmarshal(requestBody, &textRequest)
+		if err != nil {
+			return err
+		}
+		// Reset request body
+		c.Request.Body = io.NopCloser(bytes.NewBuffer(requestBody))
 	}
-	// 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"))
@@ -101,7 +107,10 @@ func relayHelper(c *gin.Context) error {
 	if err != nil {
 		return err
 	}
-
+	err = c.Request.Body.Close()
+	if err != nil {
+		return err
+	}
 	var textResponse TextResponse
 	isStream := resp.Header.Get("Content-Type") == "text/event-stream"
 	var streamResponseText string
@@ -110,10 +119,24 @@ func relayHelper(c *gin.Context) error {
 		if consumeQuota {
 			quota := 0
 			if isStream {
-				quota = int(float64(len(streamResponseText)) * common.BytesNumber2Quota)
+				var text string
+				for _, message := range textRequest.Messages {
+					text += fmt.Sprintf("%s: %s\n", message.Role, message.Content)
+				}
+				text += fmt.Sprintf("%s: %s\n", "assistant", streamResponseText)
+				quota = countToken(text) + 3
 			} else {
 				quota = textResponse.Usage.TotalTokens
 			}
+			ratio := common.RatioGPT3dot5
+			if strings.HasPrefix(textRequest.Model, "gpt-4-32k") {
+				ratio = common.RatioGPT4_32k
+			} else if strings.HasPrefix(textRequest.Model, "gpt-4") {
+				ratio = common.RatioGPT4
+			} else {
+				ratio = common.RatioGPT3dot5
+			}
+			quota = int(float64(quota) * ratio)
 			err := model.ConsumeTokenQuota(tokenId, quota)
 			if err != nil {
 				common.SysError("Error consuming token remain quota: " + err.Error())
@@ -181,20 +204,22 @@ func relayHelper(c *gin.Context) error {
 		for k, v := range resp.Header {
 			c.Writer.Header().Set(k, v[0])
 		}
-		responseBody, err := io.ReadAll(resp.Body)
-		if err != nil {
-			return err
+		if consumeQuota {
+			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 = 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 {
 			return err

go.mod

@@ -25,6 +25,7 @@ require (
 	github.com/cespare/xxhash/v2 v2.1.2 // indirect
 	github.com/chenzhuoyu/base64x v0.0.0-20221115062448-fe3a3abad311 // indirect
 	github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect
+	github.com/dlclark/regexp2 v1.8.1 // indirect
 	github.com/gin-contrib/sse v0.1.0 // indirect
 	github.com/go-playground/locales v0.14.1 // indirect
 	github.com/go-playground/universal-translator v0.18.1 // indirect
@@ -44,6 +45,7 @@ require (
 	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
 	github.com/modern-go/reflect2 v1.0.2 // indirect
 	github.com/pelletier/go-toml/v2 v2.0.7 // indirect
+	github.com/pkoukk/tiktoken-go v0.1.1 // indirect
 	github.com/twitchyliquid64/golang-asm v0.15.1 // indirect
 	github.com/ugorji/go/codec v1.2.11 // indirect
 	golang.org/x/arch v0.3.0 // indirect

go.sum

@@ -14,6 +14,8 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78=
 github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc=
+github.com/dlclark/regexp2 v1.8.1 h1:6Lcdwya6GjPUNsBct8Lg/yRPwMhABj269AAzdGSiR+0=
+github.com/dlclark/regexp2 v1.8.1/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8=
 github.com/fsnotify/fsnotify v1.4.9 h1:hsms1Qyu0jgnwNXIxa+/V/PDsU6CfLf6CNO8H7IWoS4=
 github.com/gin-contrib/cors v1.4.0 h1:oJ6gwtUl3lqV0WEIwM/LxPF1QZ5qe2lGWdY2+bz7y0g=
 github.com/gin-contrib/cors v1.4.0/go.mod h1:bs9pNM0x/UsmHPBWT2xZz9ROh8xYjYkiURUfmBoMlcs=
@@ -119,6 +121,8 @@ github.com/pelletier/go-toml/v2 v2.0.1/go.mod h1:r9LEWfGN8R5k0VXJ+0BkIe7MYkRdwZO
 github.com/pelletier/go-toml/v2 v2.0.7 h1:muncTPStnKRos5dpVKULv2FVd4bMOhNePj9CjgDb8Us=
 github.com/pelletier/go-toml/v2 v2.0.7/go.mod h1:eumQOmlWiOPt5WriQQqoM5y18pDHwha2N+QD+EUNTek=
 github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA=
+github.com/pkoukk/tiktoken-go v0.1.1 h1:jtkYlIECjyM9OW1w4rjPmTohK4arORP9V25y6TM6nXo=
+github.com/pkoukk/tiktoken-go v0.1.1/go.mod h1:boMWvk9pQCOTx11pgu0DrIdrAKgQzzJKUP6vLXaz7Rw=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/rogpeppe/go-internal v1.6.1/go.mod h1:xXDCJY+GAPziupqXw64V24skbSoqbTEfhy4qGm1nDQc=

model/option.go

@@ -47,7 +47,9 @@ 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["RatioGPT3dot5"] = strconv.FormatFloat(common.RatioGPT3dot5, 'f', -1, 64)
+	common.OptionMap["RatioGPT4"] = strconv.FormatFloat(common.RatioGPT4, 'f', -1, 64)
+	common.OptionMap["RatioGPT4_32k"] = strconv.FormatFloat(common.RatioGPT4_32k, 'f', -1, 64)
 	common.OptionMap["TopUpLink"] = common.TopUpLink
 	common.OptionMapRWMutex.Unlock()
 	options, _ := AllOption()
@@ -136,8 +138,12 @@ 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 "RatioGPT3dot5":
+		common.RatioGPT3dot5, _ = strconv.ParseFloat(value, 64)
+	case "RatioGPT4":
+		common.RatioGPT4, _ = strconv.ParseFloat(value, 64)
+	case "RatioGPT4_32k":
+		common.RatioGPT4_32k, _ = strconv.ParseFloat(value, 64)
 	case "TopUpLink":
 		common.TopUpLink = value
 	}

web/src/components/SystemSetting.js

@@ -25,7 +25,9 @@ const SystemSetting = () => {
     TurnstileSecretKey: '',
     RegisterEnabled: '',
     QuotaForNewUser: 0,
-    BytesNumber2Quota: 0.8,
+    RatioGPT3dot5: 2,
+    RatioGPT4: 30,
+    RatioGPT4_32k: 60,
     TopUpLink: ''
   });
   let originInputs = {};
@@ -91,7 +93,7 @@ const SystemSetting = () => {
       name === 'TurnstileSiteKey' ||
       name === 'TurnstileSecretKey' ||
       name === 'QuotaForNewUser' ||
-      name === 'BytesNumber2Quota' ||
+      name.startsWith('Ratio') ||
       name === 'TopUpLink'
     ) {
       setInputs((inputs) => ({ ...inputs, [name]: value }));
@@ -109,8 +111,14 @@ const SystemSetting = () => {
     if (originInputs['QuotaForNewUser'] !== inputs.QuotaForNewUser) {
       await updateOption('QuotaForNewUser', inputs.QuotaForNewUser);
     }
-    if (originInputs['BytesNumber2Quota'] !== inputs.BytesNumber2Quota) {
-      await updateOption('BytesNumber2Quota', inputs.BytesNumber2Quota);
+    if (originInputs['RatioGPT3dot5'] !== inputs.RatioGPT3dot5) {
+      await updateOption('RatioGPT3dot5', inputs.RatioGPT3dot5);
+    }
+    if (originInputs['RatioGPT4'] !== inputs.RatioGPT4) {
+      await updateOption('RatioGPT4', inputs.RatioGPT4);
+    }
+    if (originInputs['RatioGPT4_32k'] !== inputs.RatioGPT4_32k) {
+      await updateOption('RatioGPT4_32k', inputs.RatioGPT4_32k);
     }
     if (originInputs['TopUpLink'] !== inputs.TopUpLink) {
       await updateOption('TopUpLink', inputs.TopUpLink);
@@ -260,17 +268,6 @@ 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'
@@ -281,6 +278,41 @@ const SystemSetting = () => {
               placeholder='例如发卡网站的购买链接'
             />
           </Form.Group>
+          <Form.Group widths={3}>
+            <Form.Input
+              label='GPT-3.5 系列模型倍率'
+              name='RatioGPT3dot5'
+              onChange={handleInputChange}
+              autoComplete='off'
+              value={inputs.RatioGPT3dot5}
+              type='number'
+              step='0.01'
+              min='0'
+              placeholder='例如：2'
+            />
+            <Form.Input
+              label='GPT-4 系列模型倍率'
+              name='RatioGPT4'
+              onChange={handleInputChange}
+              autoComplete='off'
+              value={inputs.RatioGPT4}
+              type='number'
+              step='0.01'
+              min='0'
+              placeholder='例如：30'
+            />
+            <Form.Input
+              label='GPT-4 32k 系列模型倍率'
+              name='RatioGPT4_32k'
+              onChange={handleInputChange}
+              autoComplete='off'
+              value={inputs.RatioGPT4_32k}
+              type='number'
+              step='0.01'
+              min='0'
+              placeholder='例如：60'
+            />
+          </Form.Group>
           <Form.Button onClick={submitOperationConfig}>保存运营设置</Form.Button>
           <Divider />
           <Header as='h3'>