From e8800415b82a1cc3a5974da231515450ff71ef9b Mon Sep 17 00:00:00 2001
From: CaIon <1808837298@qq.com>
Date: Sun, 12 May 2024 20:15:56 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E6=94=AF=E6=8C=81=E8=87=AA=E5=AE=9A?=
 =?UTF-8?q?=E4=B9=89=E7=89=B9=E6=AE=8A=E6=A8=A1=E5=9E=8B=E8=A1=A5=E5=85=A8?=
 =?UTF-8?q?=E5=80=8D=E7=8E=87?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 common/model-ratio.go                  | 33 ++++++++++++++++++++++++--
 model/option.go                        |  3 +++
 web/src/components/OperationSetting.js | 20 ++++++++++++++++
 3 files changed, 54 insertions(+), 2 deletions(-)

diff --git a/common/model-ratio.go b/common/model-ratio.go
index 9ddbe12..f470317 100644
--- a/common/model-ratio.go
+++ b/common/model-ratio.go
@@ -138,6 +138,12 @@ var DefaultModelPrice = map[string]float64{
 var modelPrice map[string]float64 = nil
 var modelRatio map[string]float64 = nil
 
+var CompletionRatio map[string]float64 = nil
+var DefaultCompletionRatio = map[string]float64{
+	"gpt-4-gizmo-*": 2,
+	"gpt-4-all":     2,
+}
+
 func ModelPrice2JSONString() string {
 	if modelPrice == nil {
 		modelPrice = DefaultModelPrice
@@ -202,6 +208,22 @@ func GetModelRatio(name string) float64 {
 	return ratio
 }
 
+func CompletionRatio2JSONString() string {
+	if CompletionRatio == nil {
+		CompletionRatio = DefaultCompletionRatio
+	}
+	jsonBytes, err := json.Marshal(CompletionRatio)
+	if err != nil {
+		SysError("error marshalling completion ratio: " + err.Error())
+	}
+	return string(jsonBytes)
+}
+
+func UpdateCompletionRatioByJSONString(jsonStr string) error {
+	CompletionRatio = make(map[string]float64)
+	return json.Unmarshal([]byte(jsonStr), &CompletionRatio)
+}
+
 func GetCompletionRatio(name string) float64 {
 	if strings.HasPrefix(name, "gpt-3.5") {
 		if name == "gpt-3.5-turbo" || strings.HasSuffix(name, "0125") {
@@ -214,7 +236,7 @@ func GetCompletionRatio(name string) float64 {
 		}
 		return 4.0 / 3.0
 	}
-	if strings.HasPrefix(name, "gpt-4") {
+	if strings.HasPrefix(name, "gpt-4") && name != "gpt-4-all" && !strings.HasPrefix(name, "gpt-4-gizmo") {
 		if strings.HasPrefix(name, "gpt-4-turbo") || strings.HasSuffix(name, "preview") {
 			return 3
 		}
@@ -248,7 +270,14 @@ func GetCompletionRatio(name string) float64 {
 	}
 	switch name {
 	case "llama2-70b-4096":
-		return 0.8 / 0.7
+		return 0.8 / 0.64
+	case "llama3-8b-8192":
+		return 2
+	case "llama3-70b-8192":
+		return 0.79 / 0.59
+	}
+	if ratio, ok := CompletionRatio[name]; ok {
+		return ratio
 	}
 	return 1
 }
diff --git a/model/option.go b/model/option.go
index 1adc84c..bfa7ddc 100644
--- a/model/option.go
+++ b/model/option.go
@@ -83,6 +83,7 @@ func InitOptionMap() {
 	common.OptionMap["ModelRatio"] = common.ModelRatio2JSONString()
 	common.OptionMap["ModelPrice"] = common.ModelPrice2JSONString()
 	common.OptionMap["GroupRatio"] = common.GroupRatio2JSONString()
+	common.OptionMap["CompletionRatio"] = common.CompletionRatio2JSONString()
 	common.OptionMap["TopUpLink"] = common.TopUpLink
 	common.OptionMap["ChatLink"] = common.ChatLink
 	common.OptionMap["ChatLink2"] = common.ChatLink2
@@ -290,6 +291,8 @@ func updateOptionMap(key string, value string) (err error) {
 		err = common.UpdateModelRatioByJSONString(value)
 	case "GroupRatio":
 		err = common.UpdateGroupRatioByJSONString(value)
+	case "CompletionRatio":
+		err = common.UpdateCompletionRatioByJSONString(value)
 	case "ModelPrice":
 		err = common.UpdateModelPriceByJSONString(value)
 	case "TopUpLink":
diff --git a/web/src/components/OperationSetting.js b/web/src/components/OperationSetting.js
index 7566faa..41573af 100644
--- a/web/src/components/OperationSetting.js
+++ b/web/src/components/OperationSetting.js
@@ -20,6 +20,7 @@ const OperationSetting = () => {
     PreConsumedQuota: 0,
     StreamCacheQueueLength: 0,
     ModelRatio: '',
+    CompletionRatio: '',
     ModelPrice: '',
     GroupRatio: '',
     TopUpLink: '',
@@ -68,6 +69,7 @@ const OperationSetting = () => {
         if (
           item.key === 'ModelRatio' ||
           item.key === 'GroupRatio' ||
+          item.key === 'CompletionRatio' ||
           item.key === 'ModelPrice'
         ) {
           item.value = JSON.stringify(JSON.parse(item.value), null, 2);
@@ -157,6 +159,13 @@ const OperationSetting = () => {
           }
           await updateOption('ModelRatio', inputs.ModelRatio);
         }
+        if (originInputs['CompletionRatio'] !== inputs.CompletionRatio) {
+          if (!verifyJSON(inputs.CompletionRatio)) {
+            showError('模型补全倍率不是合法的 JSON 字符串');
+            return;
+          }
+          await updateOption('CompletionRatio', inputs.CompletionRatio);
+        }
         if (originInputs['GroupRatio'] !== inputs.GroupRatio) {
           if (!verifyJSON(inputs.GroupRatio)) {
             showError('分组倍率不是合法的 JSON 字符串');
@@ -593,6 +602,17 @@ const OperationSetting = () => {
               placeholder='为一个 JSON 文本，键为模型名称，值为倍率'
             />
           </Form.Group>
+          <Form.Group widths='equal'>
+            <Form.TextArea
+              label='模型补全倍率（仅对自定义模型有效）'
+              name='CompletionRatio'
+              onChange={handleInputChange}
+              style={{ minHeight: 250, fontFamily: 'JetBrains Mono, Consolas' }}
+              autoComplete='new-password'
+              value={inputs.CompletionRatio}
+              placeholder='为一个 JSON 文本，键为分组名称，值为倍率'
+            />
+          </Form.Group>
           <Form.Group widths='equal'>
             <Form.TextArea
               label='分组倍率'