Selaa lähdekoodia

feat(ratio): add functions to check for audio ratios and clean up unused code

CaIon 16 tuntia sitten
vanhempi
sitoutus
d06915c30d
2 muutettua tiedostoa jossa 21 lisäystä ja 21 poistoa
  1. 5 1
      relay/compatible_handler.go
  2. 16 20
      setting/ratio_setting/model_ratio.go

+ 5 - 1
relay/compatible_handler.go

@@ -18,6 +18,7 @@ import (
 	"github.com/QuantumNous/new-api/service"
 	"github.com/QuantumNous/new-api/setting/model_setting"
 	"github.com/QuantumNous/new-api/setting/operation_setting"
+	"github.com/QuantumNous/new-api/setting/ratio_setting"
 	"github.com/QuantumNous/new-api/types"
 
 	"github.com/shopspring/decimal"
@@ -181,7 +182,10 @@ func TextHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types
 		return newApiErr
 	}
 
-	if usage.(*dto.Usage).CompletionTokenDetails.AudioTokens > 0 || usage.(*dto.Usage).PromptTokensDetails.AudioTokens > 0 {
+	var containAudioTokens = usage.(*dto.Usage).CompletionTokenDetails.AudioTokens > 0 || usage.(*dto.Usage).PromptTokensDetails.AudioTokens > 0
+	var containsAudioRatios = ratio_setting.ContainsAudioRatio(info.OriginModelName) || ratio_setting.ContainsAudioCompletionRatio(info.OriginModelName)
+
+	if containAudioTokens && containsAudioRatios {
 		service.PostAudioConsumeQuota(c, info, usage.(*dto.Usage), "")
 	} else {
 		postConsumeQuota(c, info, usage.(*dto.Usage))

+ 16 - 20
setting/ratio_setting/model_ratio.go

@@ -674,6 +674,22 @@ func GetAudioCompletionRatio(name string) float64 {
 	return 1
 }
 
+func ContainsAudioRatio(name string) bool {
+	audioRatioMapMutex.RLock()
+	defer audioRatioMapMutex.RUnlock()
+	name = FormatMatchingModelName(name)
+	_, ok := audioRatioMap[name]
+	return ok
+}
+
+func ContainsAudioCompletionRatio(name string) bool {
+	audioCompletionRatioMapMutex.RLock()
+	defer audioCompletionRatioMapMutex.RUnlock()
+	name = FormatMatchingModelName(name)
+	_, ok := audioCompletionRatioMap[name]
+	return ok
+}
+
 func ModelRatio2JSONString() string {
 	modelRatioMapMutex.RLock()
 	defer modelRatioMapMutex.RUnlock()
@@ -749,16 +765,6 @@ func UpdateAudioRatioByJSONString(jsonStr string) error {
 	return nil
 }
 
-func GetAudioRatioCopy() map[string]float64 {
-	audioRatioMapMutex.RLock()
-	defer audioRatioMapMutex.RUnlock()
-	copyMap := make(map[string]float64, len(audioRatioMap))
-	for k, v := range audioRatioMap {
-		copyMap[k] = v
-	}
-	return copyMap
-}
-
 func AudioCompletionRatio2JSONString() string {
 	audioCompletionRatioMapMutex.RLock()
 	defer audioCompletionRatioMapMutex.RUnlock()
@@ -781,16 +787,6 @@ func UpdateAudioCompletionRatioByJSONString(jsonStr string) error {
 	return nil
 }
 
-func GetAudioCompletionRatioCopy() map[string]float64 {
-	audioCompletionRatioMapMutex.RLock()
-	defer audioCompletionRatioMapMutex.RUnlock()
-	copyMap := make(map[string]float64, len(audioCompletionRatioMap))
-	for k, v := range audioCompletionRatioMap {
-		copyMap[k] = v
-	}
-	return copyMap
-}
-
 func GetModelRatioCopy() map[string]float64 {
 	modelRatioMapMutex.RLock()
 	defer modelRatioMapMutex.RUnlock()