feat: 增加重置模型倍率功能 (close #62)

common/model-ratio.go

@@ -20,9 +20,11 @@ const (
 // 1 === $0.002 / 1K tokens
 // 1 === ￥0.014 / 1k tokens
 
-var DefaultModelRatio = map[string]float64{
+var defaultModelRatio = map[string]float64{
 	//"midjourney":                50,
 	"gpt-4-gizmo-*": 15,
+	"gpt-4-all":     15,
+	"gpt-4o-all":    15,
 	"gpt-4":         15,
 	//"gpt-4-0314":                   15, //deprecated
 	"gpt-4-0613": 15,
@@ -149,7 +151,7 @@ var DefaultModelRatio = map[string]float64{
 	"llama-3-sonar-large-32k-online": 1 / 1000 * USD,
 }
 
-var DefaultModelPrice = map[string]float64{
+var defaultModelPrice = map[string]float64{
 	"dall-e-3":          0.04,
 	"gpt-4-gizmo-*":     0.1,
 	"mj_imagine":        0.1,
@@ -173,14 +175,14 @@ var modelPrice map[string]float64 = nil
 var modelRatio map[string]float64 = nil
 
 var CompletionRatio map[string]float64 = nil
-var DefaultCompletionRatio = map[string]float64{
+var defaultCompletionRatio = map[string]float64{
 	"gpt-4-gizmo-*": 2,
 	"gpt-4-all":     2,
 }
 
 func ModelPrice2JSONString() string {
 	if modelPrice == nil {
-		modelPrice = DefaultModelPrice
+		modelPrice = defaultModelPrice
 	}
 	jsonBytes, err := json.Marshal(modelPrice)
 	if err != nil {
@@ -197,7 +199,7 @@ func UpdateModelPriceByJSONString(jsonStr string) error {
 // GetModelPrice 返回模型的价格，如果模型不存在则返回-1，false
 func GetModelPrice(name string, printErr bool) (float64, bool) {
 	if modelPrice == nil {
-		modelPrice = DefaultModelPrice
+		modelPrice = defaultModelPrice
 	}
 	if strings.HasPrefix(name, "gpt-4-gizmo") {
 		name = "gpt-4-gizmo-*"
@@ -212,16 +214,16 @@ func GetModelPrice(name string, printErr bool) (float64, bool) {
 	return price, true
 }
 
-func GetModelPrices() map[string]float64 {
+func GetModelPriceMap() map[string]float64 {
 	if modelPrice == nil {
-		modelPrice = DefaultModelPrice
+		modelPrice = defaultModelPrice
 	}
 	return modelPrice
 }
 
 func ModelRatio2JSONString() string {
 	if modelRatio == nil {
-		modelRatio = DefaultModelRatio
+		modelRatio = defaultModelRatio
 	}
 	jsonBytes, err := json.Marshal(modelRatio)
 	if err != nil {
@@ -237,7 +239,7 @@ func UpdateModelRatioByJSONString(jsonStr string) error {
 
 func GetModelRatio(name string) float64 {
 	if modelRatio == nil {
-		modelRatio = DefaultModelRatio
+		modelRatio = defaultModelRatio
 	}
 	if strings.HasPrefix(name, "gpt-4-gizmo") {
 		name = "gpt-4-gizmo-*"
@@ -250,16 +252,21 @@ func GetModelRatio(name string) float64 {
 	return ratio
 }
 
-func GetModelRatios() map[string]float64 {
+func DefaultModelRatio2JSONString() string {
-	if modelRatio == nil {
+	jsonBytes, err := json.Marshal(defaultModelRatio)
-		modelRatio = DefaultModelRatio
+	if err != nil {
+		SysError("error marshalling model ratio: " + err.Error())
 	}
-	return modelRatio
+	return string(jsonBytes)
+}
+
+func GetDefaultModelRatioMap() map[string]float64 {
+	return defaultModelRatio
 }
 
 func CompletionRatio2JSONString() string {
 	if CompletionRatio == nil {
-		CompletionRatio = DefaultCompletionRatio
+		CompletionRatio = defaultCompletionRatio
 	}
 	jsonBytes, err := json.Marshal(CompletionRatio)
 	if err != nil {
@@ -343,9 +350,9 @@ func GetCompletionRatio(name string) float64 {
 	return 1
 }
 
-func GetCompletionRatios() map[string]float64 {
+func GetCompletionRatioMap() map[string]float64 {
 	if CompletionRatio == nil {
-		CompletionRatio = DefaultCompletionRatio
+		CompletionRatio = defaultCompletionRatio
 	}
 	return CompletionRatio
 }

controller/model.go

@@ -200,19 +200,3 @@ func RetrieveModel(c *gin.Context) {
 		})
 	}
 }
-
-func GetPricing(c *gin.Context) {
-	userId := c.GetInt("id")
-	// if no login, get default group ratio
-	groupRatio := common.GetGroupRatio("default")
-	group, err := model.CacheGetUserGroup(userId)
-	if err == nil {
-		groupRatio = common.GetGroupRatio(group)
-	}
-	pricing := model.GetPricing(group)
-	c.JSON(200, gin.H{
-		"success":     true,
-		"data":        pricing,
-		"group_ratio": groupRatio,
-	})
-}

controller/pricing.go

@@ -0,0 +1,47 @@
+package controller
+
+import (
+	"github.com/gin-gonic/gin"
+	"one-api/common"
+	"one-api/model"
+)
+
+func GetPricing(c *gin.Context) {
+	userId := c.GetInt("id")
+	// if no login, get default group ratio
+	groupRatio := common.GetGroupRatio("default")
+	group, err := model.CacheGetUserGroup(userId)
+	if err == nil {
+		groupRatio = common.GetGroupRatio(group)
+	}
+	pricing := model.GetPricing(group)
+	c.JSON(200, gin.H{
+		"success":     true,
+		"data":        pricing,
+		"group_ratio": groupRatio,
+	})
+}
+
+func ResetModelRatio(c *gin.Context) {
+	defaultStr := common.DefaultModelRatio2JSONString()
+	err := model.UpdateOption("ModelRatio", defaultStr)
+	if err != nil {
+		c.JSON(200, gin.H{
+			"success": false,
+			"message": err.Error(),
+		})
+		return
+	}
+	err = common.UpdateModelRatioByJSONString(defaultStr)
+	if err != nil {
+		c.JSON(200, gin.H{
+			"success": false,
+			"message": err.Error(),
+		})
+		return
+	}
+	c.JSON(200, gin.H{
+		"success": true,
+		"message": "重置模型倍率成功",
+	})
+}

relay/relay-mj.go

@@ -158,7 +158,7 @@ func RelaySwapFace(c *gin.Context) *dto.MidjourneyResponse {
 	modelPrice, success := common.GetModelPrice(modelName, true)
 	// 如果没有配置价格，则使用默认价格
 	if !success {
-		defaultPrice, ok := common.DefaultModelPrice[modelName]
+		defaultPrice, ok := common.GetDefaultModelRatioMap()[modelName]
 		if !ok {
 			modelPrice = 0.1
 		} else {
@@ -457,7 +457,7 @@ func RelayMidjourneySubmit(c *gin.Context, relayMode int) *dto.MidjourneyRespons
 	modelPrice, success := common.GetModelPrice(modelName, true)
 	// 如果没有配置价格，则使用默认价格
 	if !success {
-		defaultPrice, ok := common.DefaultModelPrice[modelName]
+		defaultPrice, ok := common.GetDefaultModelRatioMap()[modelName]
 		if !ok {
 			modelPrice = 0.1
 		} else {

router/api-router.go

@@ -72,6 +72,7 @@ func SetApiRouter(router *gin.Engine) {
 		{
 			optionRoute.GET("/", controller.GetOptions)
 			optionRoute.PUT("/", controller.UpdateOption)
+			optionRoute.POST("/rest_model_ratio", controller.ResetModelRatio)
 		}
 		channelRoute := apiRouter.Group("/channel")
 		channelRoute.Use(middleware.AdminAuth())

service/token_counter.go

@@ -30,7 +30,7 @@ func InitTokenEncoders() {
 	if err != nil {
 		common.FatalLog(fmt.Sprintf("failed to get gpt-4 token encoder: %s", err.Error()))
 	}
-	for model, _ := range common.DefaultModelRatio {
+	for model, _ := range common.GetDefaultModelRatioMap() {
 		if strings.HasPrefix(model, "gpt-3.5") {
 			tokenEncoderMap[model] = gpt35TokenEncoder
 		} else if strings.HasPrefix(model, "gpt-4") {

web/src/helpers/utils.js

@@ -212,6 +212,16 @@ export const verifyJSON = (str) => {
   return true;
 };
 
+export function verifyJSONPromise(value) {
+  try {
+    JSON.parse(value);
+    return Promise.resolve();
+  } catch (e) {
+    return Promise.reject('不是合法的 JSON 字符串');
+  }
+}
+
+
 export function shouldShowPrompt(id) {
   let prompt = localStorage.getItem(`prompt-${id}`);
   return !prompt;

web/src/pages/Setting/Operation/SettingsMagnification.js

@@ -1,5 +1,5 @@
 import React, { useEffect, useState, useRef } from 'react';
-import { Button, Col, Form, Row, Spin } from '@douyinfe/semi-ui';
+import { Button, Col, Form, Popconfirm, Row, Space, Spin } from '@douyinfe/semi-ui';
 import {
   compareObjects,
   API,
@@ -7,6 +7,7 @@ import {
   showSuccess,
   showWarning,
   verifyJSON,
+  verifyJSONPromise
 } from '../../../helpers';
 
 export default function SettingsMagnification(props) {
@@ -15,14 +16,16 @@ export default function SettingsMagnification(props) {
     ModelPrice: '',
     ModelRatio: '',
     CompletionRatio: '',
-    GroupRatio: '',
+    GroupRatio: ''
   });
   const refForm = useRef();
   const [inputsRow, setInputsRow] = useState(inputs);
 
   async function onSubmit() {
     try {
-      await refForm.current.validate();
+      console.log('Starting validation...');
+      await refForm.current.validate().then(() => {
+        console.log('Validation passed');
         const updateArray = compareObjects(inputs, inputsRow);
         if (!updateArray.length) return showWarning('你似乎并没有修改什么');
         const requestQueue = updateArray.map((item) => {
@@ -34,7 +37,7 @@ export default function SettingsMagnification(props) {
           }
           return API.put('/api/option/', {
             key: item.key,
-          value,
+            value
           });
         });
         setLoading(true);
@@ -55,10 +58,28 @@ export default function SettingsMagnification(props) {
           .finally(() => {
             setLoading(false);
           });
+      }).catch((error) => {
+        console.error('Validation failed:', error);
+        showError('请检查输入');
+      });
     } catch (error) {
       showError('请检查输入');
       console.error(error);
-    } finally {
+    }
+  }
+
+  async function resetModelRatio() {
+    try {
+      let res = await API.post(`/api/option/rest_model_ratio`);
+      // return {success, message}
+      if (res.data.success) {
+        showSuccess(res.data.message);
+        props.refresh();
+      } else {
+        showError(res.data.message);
+      }
+    } catch (error) {
+      showError(error);
     }
   }
 
@@ -73,8 +94,8 @@ export default function SettingsMagnification(props) {
     setInputsRow(structuredClone(currentInputs));
     refForm.current.setValues(currentInputs);
   }, [props.options]);
+
   return (
-    <>
     <Spin spinning={loading}>
       <Form
         values={inputs}
@@ -93,16 +114,19 @@ export default function SettingsMagnification(props) {
                 field={'ModelPrice'}
                 autosize={{ minRows: 6, maxRows: 12 }}
                 trigger='blur'
+                stopValidateWithError
                 rules={[
                   {
-                      validator: (rule, value) => verifyJSON(value),
+                    validator: (rule, value) => {
-                      message: '不是合法的 JSON 字符串',
+                      return verifyJSON(value);
                     },
+                    message: '不是合法的 JSON 字符串'
+                  }
                 ]}
                 onChange={(value) =>
                   setInputs({
                     ...inputs,
-                      ModelPrice: value,
+                    ModelPrice: value
                   })
                 }
               />
@@ -117,16 +141,19 @@ export default function SettingsMagnification(props) {
                 field={'ModelRatio'}
                 autosize={{ minRows: 6, maxRows: 12 }}
                 trigger='blur'
+                stopValidateWithError
                 rules={[
                   {
-                      validator: (rule, value) => verifyJSON(value),
+                    validator: (rule, value) => {
-                      message: '不是合法的 JSON 字符串',
+                      return verifyJSON(value);
                     },
+                    message: '不是合法的 JSON 字符串'
+                  }
                 ]}
                 onChange={(value) =>
                   setInputs({
                     ...inputs,
-                      ModelRatio: value,
+                    ModelRatio: value
                   })
                 }
               />
@@ -141,16 +168,19 @@ export default function SettingsMagnification(props) {
                 field={'CompletionRatio'}
                 autosize={{ minRows: 6, maxRows: 12 }}
                 trigger='blur'
+                stopValidateWithError
                 rules={[
                   {
-                      validator: (rule, value) => verifyJSON(value),
+                    validator: (rule, value) => {
-                      message: '不是合法的 JSON 字符串',
+                      return verifyJSON(value);
                     },
+                    message: '不是合法的 JSON 字符串'
+                  }
                 ]}
                 onChange={(value) =>
                   setInputs({
                     ...inputs,
-                      CompletionRatio: value,
+                    CompletionRatio: value
                   })
                 }
               />
@@ -165,30 +195,44 @@ export default function SettingsMagnification(props) {
                 field={'GroupRatio'}
                 autosize={{ minRows: 6, maxRows: 12 }}
                 trigger='blur'
+                stopValidateWithError
                 rules={[
                   {
-                      validator: (rule, value) => verifyJSON(value),
+                    validator: (rule, value) => {
-                      message: '不是合法的 JSON 字符串',
+                      return verifyJSON(value);
                     },
+                    message: '不是合法的 JSON 字符串'
+                  }
                 ]}
                 onChange={(value) =>
                   setInputs({
                     ...inputs,
-                      GroupRatio: value,
+                    GroupRatio: value
                   })
                 }
               />
             </Col>
           </Row>
-
-            <Row>
-              <Button size='large' onClick={onSubmit}>
-                保存倍率设置
-              </Button>
-            </Row>
         </Form.Section>
       </Form>
+      <Space>
+        <Button onClick={onSubmit}>
+          保存倍率设置
+        </Button>
+        <Popconfirm
+          title='确定重置模型倍率吗？'
+          content='此修改将不可逆'
+          okType={'danger'}
+          position={'top'}
+          onConfirm={() => {
+            resetModelRatio();
+          }}
+        >
+          <Button type={'danger'}>
+            重置模型倍率
+          </Button>
+        </Popconfirm>
+      </Space>
     </Spin>
-    </>
   );
 }