Преглед изворни кода

feat: support free model setting

CaIon пре 2 месеци
родитељ
комит
ede47ef014

+ 7 - 3
controller/relay.go

@@ -140,9 +140,13 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
 
 	// common.SetContextKey(c, constant.ContextKeyTokenCountMeta, meta)
 
-	newAPIError = service.PreConsumeQuota(c, priceData.ShouldPreConsumedQuota, relayInfo)
-	if newAPIError != nil {
-		return
+	if priceData.FreeModel {
+		logger.LogInfo(c, fmt.Sprintf("模型 %s 免费,跳过预扣费", relayInfo.OriginModelName))
+	} else {
+		newAPIError = service.PreConsumeQuota(c, priceData.QuotaToPreConsume, relayInfo)
+		if newAPIError != nil {
+			return
+		}
 	}
 
 	defer func() {

+ 30 - 11
relay/helper/price.go

@@ -5,6 +5,7 @@ import (
 
 	"github.com/QuantumNous/new-api/common"
 	relaycommon "github.com/QuantumNous/new-api/relay/common"
+	"github.com/QuantumNous/new-api/setting/operation_setting"
 	"github.com/QuantumNous/new-api/setting/ratio_setting"
 	"github.com/QuantumNous/new-api/types"
 
@@ -55,6 +56,7 @@ func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens
 	var cacheCreationRatio float64
 	var audioRatio float64
 	var audioCompletionRatio float64
+	var freeModel bool
 	if !usePrice {
 		preConsumedTokens := common.Max(promptTokens, common.PreConsumedQuota)
 		if meta.MaxTokens != 0 {
@@ -87,18 +89,35 @@ func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens
 		preConsumedQuota = int(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio)
 	}
 
+	// check if free model pre-consume is disabled
+	if !operation_setting.GetQuotaSetting().EnableFreeModelPreConsume {
+		// if model price or ratio is 0, do not pre-consume quota
+		if usePrice {
+			if modelPrice == 0 {
+				preConsumedQuota = 0
+				freeModel = true
+			}
+		} else {
+			if modelRatio == 0 {
+				preConsumedQuota = 0
+				freeModel = true
+			}
+		}
+	}
+
 	priceData := types.PriceData{
-		ModelPrice:             modelPrice,
-		ModelRatio:             modelRatio,
-		CompletionRatio:        completionRatio,
-		GroupRatioInfo:         groupRatioInfo,
-		UsePrice:               usePrice,
-		CacheRatio:             cacheRatio,
-		ImageRatio:             imageRatio,
-		AudioRatio:             audioRatio,
-		AudioCompletionRatio:   audioCompletionRatio,
-		CacheCreationRatio:     cacheCreationRatio,
-		ShouldPreConsumedQuota: preConsumedQuota,
+		FreeModel:            freeModel,
+		ModelPrice:           modelPrice,
+		ModelRatio:           modelRatio,
+		CompletionRatio:      completionRatio,
+		GroupRatioInfo:       groupRatioInfo,
+		UsePrice:             usePrice,
+		CacheRatio:           cacheRatio,
+		ImageRatio:           imageRatio,
+		AudioRatio:           audioRatio,
+		AudioCompletionRatio: audioCompletionRatio,
+		CacheCreationRatio:   cacheCreationRatio,
+		QuotaToPreConsume:    preConsumedQuota,
 	}
 
 	if common.DebugEnabled {

+ 21 - 0
setting/operation_setting/quota_setting.go

@@ -0,0 +1,21 @@
+package operation_setting
+
+import "github.com/QuantumNous/new-api/setting/config"
+
+type QuotaSetting struct {
+	EnableFreeModelPreConsume bool `json:"enable_free_model_pre_consume"` // 是否对免费模型启用预消耗
+}
+
+// 默认配置
+var quotaSetting = QuotaSetting{
+	EnableFreeModelPreConsume: true,
+}
+
+func init() {
+	// 注册到全局配置管理器
+	config.GlobalConfig.Register("quota_setting", &quotaSetting)
+}
+
+func GetQuotaSetting() *QuotaSetting {
+	return &quotaSetting
+}

+ 14 - 13
types/price_data.go

@@ -9,18 +9,19 @@ type GroupRatioInfo struct {
 }
 
 type PriceData struct {
-	ModelPrice             float64
-	ModelRatio             float64
-	CompletionRatio        float64
-	CacheRatio             float64
-	CacheCreationRatio     float64
-	ImageRatio             float64
-	AudioRatio             float64
-	AudioCompletionRatio   float64
-	OtherRatios            map[string]float64
-	UsePrice               bool
-	ShouldPreConsumedQuota int
-	GroupRatioInfo         GroupRatioInfo
+	FreeModel            bool
+	ModelPrice           float64
+	ModelRatio           float64
+	CompletionRatio      float64
+	CacheRatio           float64
+	CacheCreationRatio   float64
+	ImageRatio           float64
+	AudioRatio           float64
+	AudioCompletionRatio float64
+	OtherRatios          map[string]float64
+	UsePrice             bool
+	QuotaToPreConsume    int // 预消耗额度
+	GroupRatioInfo       GroupRatioInfo
 }
 
 type PerCallPriceData struct {
@@ -30,5 +31,5 @@ type PerCallPriceData struct {
 }
 
 func (p PriceData) ToSetting() string {
-	return fmt.Sprintf("ModelPrice: %f, ModelRatio: %f, CompletionRatio: %f, CacheRatio: %f, GroupRatio: %f, UsePrice: %t, CacheCreationRatio: %f, ShouldPreConsumedQuota: %d, ImageRatio: %f, AudioRatio: %f, AudioCompletionRatio: %f", p.ModelPrice, p.ModelRatio, p.CompletionRatio, p.CacheRatio, p.GroupRatioInfo.GroupRatio, p.UsePrice, p.CacheCreationRatio, p.ShouldPreConsumedQuota, p.ImageRatio, p.AudioRatio, p.AudioCompletionRatio)
+	return fmt.Sprintf("ModelPrice: %f, ModelRatio: %f, CompletionRatio: %f, CacheRatio: %f, GroupRatio: %f, UsePrice: %t, CacheCreationRatio: %f, QuotaToPreConsume: %d, ImageRatio: %f, AudioRatio: %f, AudioCompletionRatio: %f", p.ModelPrice, p.ModelRatio, p.CompletionRatio, p.CacheRatio, p.GroupRatioInfo.GroupRatio, p.UsePrice, p.CacheCreationRatio, p.QuotaToPreConsume, p.ImageRatio, p.AudioRatio, p.AudioCompletionRatio)
 }

+ 1 - 0
web/src/components/settings/OperationSetting.jsx

@@ -35,6 +35,7 @@ const OperationSetting = () => {
     PreConsumedQuota: 0,
     QuotaForInviter: 0,
     QuotaForInvitee: 0,
+    'quota_setting.enable_free_model_pre_consume': true,
 
     /* 通用设置 */
     TopUpLink: '',

+ 16 - 0
web/src/pages/Setting/Operation/SettingsCreditLimit.jsx

@@ -36,6 +36,7 @@ export default function SettingsCreditLimit(props) {
     PreConsumedQuota: '',
     QuotaForInviter: '',
     QuotaForInvitee: '',
+    'quota_setting.enable_free_model_pre_consume': true,
   });
   const refForm = useRef();
   const [inputsRow, setInputsRow] = useState(inputs);
@@ -166,6 +167,21 @@ export default function SettingsCreditLimit(props) {
                 />
               </Col>
             </Row>
+            <Row>
+              <Col>
+                <Form.Switch
+                  label={t('对免费模型启用预消耗')}
+                  field={'quota_setting.enable_free_model_pre_consume'}
+                  extraText={t('开启后,对免费模型(倍率为0,或者价格为0)的模型也会预消耗额度')}
+                  onChange={(value) =>
+                    setInputs({
+                      ...inputs,
+                      'quota_setting.enable_free_model_pre_consume': value,
+                    })
+                  }
+                />
+              </Col>
+            </Row>
 
             <Row>
               <Button size='default' onClick={onSubmit}>