feat: add midjourney price log

Calcium-Ion · Jan 25, 2024 · a3921ea · a3921ea
1 parent 705171e
commit a3921ea
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 4 deletions.
diff --git a/common/model-ratio.go b/common/model-ratio.go
@@ -103,13 +103,15 @@ func UpdateModelPriceByJSONString(jsonStr string) error {
 	return json.Unmarshal([]byte(jsonStr), &ModelPrice)
 }
 
-func GetModelPrice(name string) float64 {
+func GetModelPrice(name string, printErr bool) float64 {
 	if strings.HasPrefix(name, "gpt-4-gizmo") {
 		name = "gpt-4-gizmo-*"
 	}
 	price, ok := ModelPrice[name]
 	if !ok {
-		//SysError("model price not found: " + name)
+		if printErr {
+			SysError("model price not found: " + name)
+		}
 		return -1
 	}
 	return price

diff --git a/controller/relay-mj.go b/controller/relay-mj.go
@@ -392,7 +392,7 @@ func relayMidjourneySubmit(c *gin.Context, relayMode int) *MidjourneyResponse {
 		requestBody = c.Request.Body
 	}
 	mjAction := "mj_" + strings.ToLower(midjRequest.Action)
-	modelPrice := common.GetModelPrice(mjAction)
+	modelPrice := common.GetModelPrice(mjAction, true)
 	// 如果没有配置价格，则使用默认价格
 	if modelPrice == -1 {
 		defaultPrice, ok := DefaultModelPrice[mjAction]

diff --git a/controller/relay-text.go b/controller/relay-text.go
@@ -232,7 +232,7 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
 	case RelayModeModerations:
 		promptTokens = countTokenInput(textRequest.Input, textRequest.Model)
 	}
-	modelPrice := common.GetModelPrice(textRequest.Model)
+	modelPrice := common.GetModelPrice(textRequest.Model, false)
 	groupRatio := common.GetGroupRatio(group)
 
 	var preConsumedQuota int