new-api/relay/helper/price.go

package helper

import (
	"fmt"
	"net/http"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/dto"
	"github.com/QuantumNous/new-api/logger"
	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"

	"github.com/gin-gonic/gin"
)

// HandleGroupRatio checks for "auto_group" in the context and updates the group ratio and relayInfo.UsingGroup if present
func HandleGroupRatio(ctx *gin.Context, relayInfo *relaycommon.RelayInfo) types.GroupRatioInfo {
	groupRatioInfo := types.GroupRatioInfo{
		GroupRatio:        1.0,
		GroupSpecialRatio: -1,
	}

	autoGroup, exists := ctx.Get("auto_group")
	if exists {
		logger.LogDebug(ctx, fmt.Sprintf("final group: %s", autoGroup))
		relayInfo.UsingGroup = autoGroup.(string)
	}

	userGroupRatio, ok := ratio_setting.GetGroupGroupRatio(relayInfo.UserGroup, relayInfo.UsingGroup)
	if ok {
		groupRatioInfo.GroupSpecialRatio = userGroupRatio
		groupRatioInfo.GroupRatio = userGroupRatio
		groupRatioInfo.HasSpecialRatio = true
	} else {
		groupRatioInfo.GroupRatio = ratio_setting.GetGroupRatio(relayInfo.UsingGroup)
	}

	return groupRatioInfo
}

func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens int, meta *types.TokenCountMeta) (types.PriceData, error) {
	var modelPrice float64
	var usePrice bool
	var modelRatio float64
	var completionRatio float64

	modelPrice, usePrice = ratio_setting.GetModelPrice(info.OriginModelName, false)
	groupRatioInfo := HandleGroupRatio(c, info)

	var preConsumedQuota int
	var freeModel bool
	if !usePrice {
		preConsumedTokens := common.Max(promptTokens, common.PreConsumedQuota)
		if meta.MaxTokens != 0 {
			preConsumedTokens += meta.MaxTokens
		}

		var success bool
		var matchName string
		modelRatio, success, matchName = ratio_setting.GetModelRatio(info.OriginModelName)
		if !success {
			acceptUnsetRatio := false
			if info.UserSetting.AcceptUnsetRatioModel {
				acceptUnsetRatio = true
			}
			if !acceptUnsetRatio {
				return types.PriceData{}, fmt.Errorf("模型 %s 倍率或价格未配置，请联系管理员设置或开启自用模式；Model %s ratio or price not set, please set or start self-use mode", matchName, matchName)
			}
		}
		completionRatio = ratio_setting.GetCompletionRatio(info.OriginModelName)

		ratio := modelRatio * groupRatioInfo.GroupRatio
		preConsumedQuota = int(float64(preConsumedTokens) * ratio)
	} else {
		if meta.ImagePriceRatio != 0 {
			modelPrice = modelPrice * meta.ImagePriceRatio
		}
		preConsumedQuota = int(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio)
	}

	cacheRatio, _ := ratio_setting.GetCacheRatio(info.OriginModelName)
	cacheCreationRatio, _ := ratio_setting.GetCreateCacheRatio(info.OriginModelName)
	imageRatio, _ := ratio_setting.GetImageRatio(info.OriginModelName)
	audioRatio := ratio_setting.GetAudioRatio(info.OriginModelName)
	audioCompletionRatio := ratio_setting.GetAudioCompletionRatio(info.OriginModelName)

	if !operation_setting.GetQuotaSetting().EnableFreeModelPreConsume {
		if groupRatioInfo.GroupRatio == 0 {
			preConsumedQuota = 0
			freeModel = true
		} else if usePrice {
			if modelPrice == 0 {
				preConsumedQuota = 0
				freeModel = true
			}
		} else if modelRatio == 0 {
			preConsumedQuota = 0
			freeModel = true
		}
	}

	priceData := types.PriceData{
		FreeModel:            freeModel,
		ModelPrice:           modelPrice,
		ModelRatio:           modelRatio,
		CompletionRatio:      completionRatio,
		GroupRatioInfo:       groupRatioInfo,
		UsePrice:             usePrice,
		QuotaToPreConsume:    preConsumedQuota,
		CacheRatio:           cacheRatio,
		ImageRatio:           imageRatio,
		AudioRatio:           audioRatio,
		AudioCompletionRatio: audioCompletionRatio,
		CacheCreationRatio:   cacheCreationRatio,
		CacheCreation5mRatio: cacheCreationRatio,
		CacheCreation1hRatio: cacheCreationRatio * types.ClaudeCacheCreation1hMultiplier,
	}

	if common.DebugEnabled {
		println(fmt.Sprintf("model_price_helper result: %s", priceData.ToSetting()))
	}
	info.PriceData = priceData
	return priceData, nil
}

// ModelPriceHelperPerCall 按次计费的 PriceHelper (MJ、Task)
func ModelPriceHelperPerCall(c *gin.Context, info *relaycommon.RelayInfo) (types.PriceData, *dto.TaskError) {
	groupRatioInfo := HandleGroupRatio(c, info)

	if info.PricingDecisionFrozen != nil {
		priceData := priceDataFromDecision(info.PricingDecisionFrozen)
		info.PriceData = priceData
		return priceData, nil
	}

	if info.OriginPricing != nil {
		var (
			decision *types.PricingDecision
			taskErr  *dto.TaskError
		)
		if info.OriginPricing.BillingMode == types.BillingModeMatrix {
			decision, taskErr = buildRemixMatrixDecision(info.OriginPricing, groupRatioInfo)
		} else {
			decision, taskErr = buildRemixLegacyDecision(info.OriginPricing, groupRatioInfo)
		}
		if taskErr != nil {
			return types.PriceData{}, taskErr
		}
		info.PricingDecisionFrozen = decision
		priceData := priceDataFromDecision(decision)
		info.PriceData = priceData
		return priceData, nil
	}

	if info.PricingConfigSnapshotLoaded && info.PricingConfigSnapshot != nil {
		resolver, err := NewDimensionResolver(c)
		if err != nil {
			return types.PriceData{}, pricingDimensionTaskError(err)
		}
		resolved, err := resolver.ResolveDimensions(info.PricingConfigSnapshot.Dimensions)
		if err != nil {
			return types.PriceData{}, pricingDimensionTaskError(err)
		}
		result, taskErr := lookupPricingConfig(info.PricingConfigSnapshot, resolved)
		if taskErr != nil {
			return types.PriceData{}, taskErr
		}
		decision := buildMatchedDecision(info.PricingConfigSnapshot, result, groupRatioInfo)
		if result.Mode != types.PricingModeMatched {
			decision = buildFallbackDecision(info.PricingConfigSnapshot, result, groupRatioInfo)
		}
		info.PricingDecisionFrozen = decision
		priceData := priceDataFromDecision(decision)
		info.PriceData = priceData
		return priceData, nil
	}

	modelPrice, success := ratio_setting.GetModelPrice(info.OriginModelName, true)
	if !success {
		defaultPrice, ok := ratio_setting.GetDefaultModelPriceMap()[info.OriginModelName]
		if !ok {
			modelPrice = 0.1
		} else {
			modelPrice = defaultPrice
		}
	}
	quota := int(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio)

	freeModel := false
	if !operation_setting.GetQuotaSetting().EnableFreeModelPreConsume {
		if groupRatioInfo.GroupRatio == 0 || modelPrice == 0 {
			quota = 0
			freeModel = true
		}
	}

	priceData := types.PriceData{
		FreeModel:      freeModel,
		ModelPrice:     modelPrice,
		Quota:          quota,
		GroupRatioInfo: groupRatioInfo,
	}
	decision := buildLegacyDecision(modelPrice, groupRatioInfo, freeModel, quota)
	priceData.PricingDecision = decision
	info.PricingDecisionFrozen = decision
	info.PriceData = priceData
	return priceData, nil
}

func pricingDimensionTaskError(err error) *dto.TaskError {
	return &dto.TaskError{
		Code:       "pricing_dimension_resolve_failed",
		Message:    err.Error(),
		StatusCode: http.StatusBadRequest,
		LocalError: true,
		Error:      err,
	}
}

func ContainPriceOrRatio(modelName string) bool {
	_, ok := ratio_setting.GetModelPrice(modelName, false)
	if ok {
		return true
	}
	_, ok, _ = ratio_setting.GetModelRatio(modelName)
	if ok {
		return true
	}
	return false
}