new-api/relay/helper/pricing_lookup.go

package helper

import (
	"fmt"
	"net/http"
	"strings"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/dto"
	"github.com/QuantumNous/new-api/types"
)

type pricingLookupResult struct {
	Mode     types.PricingMode
	PriceUSD float64
	Snapshot map[string]any
}

func lookupPricingConfig(cfg *types.PricingConfig, resolved *ResolvedDimensionValues) (*pricingLookupResult, *dto.TaskError) {
	if cfg == nil {
		return nil, localPricingTaskError("pricing config is required", "pricing_config_required")
	}

	var best types.PricingRow
	bestWildcardCount := len(cfg.Dimensions) + 1
	bestCount := 0

	for _, row := range cfg.Table {
		wildcards, ok := rowMatchesDimensions(cfg.Dimensions, row, resolved.Values)
		if !ok {
			continue
		}
		if wildcards < bestWildcardCount {
			best = row
			bestWildcardCount = wildcards
			bestCount = 1
			continue
		}
		if wildcards == bestWildcardCount {
			bestCount++
		}
	}

	if best != nil && bestCount == 1 {
		price, _ := pricingNumberValue(best["price"])
		return &pricingLookupResult{
			Mode:     types.PricingModeMatched,
			PriceUSD: price,
			Snapshot: buildPricingSnapshot(cfg, resolved, price, types.PricingModeMatched, ""),
		}, nil
	}

	if bestCount > 1 {
		return fallbackPricingConfig(cfg, resolved)
	}
	return fallbackPricingConfig(cfg, resolved)
}

func buildMatchedDecision(cfg *types.PricingConfig, result *pricingLookupResult, group types.GroupRatioInfo) *types.PricingDecision {
	return buildMatrixDecision(cfg, result, group)
}

func buildFallbackDecision(cfg *types.PricingConfig, result *pricingLookupResult, group types.GroupRatioInfo) *types.PricingDecision {
	return buildMatrixDecision(cfg, result, group)
}

func buildLegacyDecision(modelPrice float64, group types.GroupRatioInfo, free bool, quota int) *types.PricingDecision {
	return &types.PricingDecision{
		Mode:             types.PricingModeLegacy,
		PriceUSD:         modelPrice,
		Quota:            quota,
		GroupRatioInfo:   group,
		FreeModel:        free,
		LegacyModelPrice: modelPrice,
		UsePrice:         true,
		PerCallBilling:   true,
		Finalized:        true,
	}
}

func buildRemixMatrixDecision(origin *types.OriginPricingSnapshot, group types.GroupRatioInfo) (*types.PricingDecision, *dto.TaskError) {
	if origin == nil {
		return nil, localPricingTaskError("origin pricing snapshot is required", "invalid_pricing_snapshot")
	}
	billingUnit := origin.BillingUnit
	if billingUnit == "" {
		billingUnit = types.BillingUnitPerCall
	}
	result := &pricingLookupResult{
		Mode:     types.PricingModeRemix,
		PriceUSD: origin.PriceUSD,
		Snapshot: types.CloneMapAny(origin.Snapshot),
	}
	cfg := &types.PricingConfig{BillingUnit: billingUnit}
	decision := buildMatrixDecision(cfg, result, group)
	decision.Mode = types.PricingModeRemix
	decision.TokenUnitPriceUSD = origin.TokenUnitPriceUSD
	decision.OtherRatios = types.CloneRatios(origin.OtherRatios)
	decision.PerCallBilling = origin.PerCallBilling
	if billingUnit == types.BillingUnitPer1MTokens {
		decision.PriceUSD = origin.TokenUnitPriceUSD
		decision.Quota = minimumUsagePreconsumeQuota(origin.TokenUnitPriceUSD, group.GroupRatio)
	}
	return decision, nil
}

func buildRemixLegacyDecision(origin *types.OriginPricingSnapshot, group types.GroupRatioInfo) (*types.PricingDecision, *dto.TaskError) {
	if origin == nil {
		return nil, localPricingTaskError("origin legacy pricing snapshot is required", "invalid_legacy_pricing_snapshot")
	}
	quota := int(origin.PriceUSD * common.QuotaPerUnit * group.GroupRatio)
	return &types.PricingDecision{
		Mode:             types.PricingModeRemix,
		PriceUSD:         origin.PriceUSD,
		Quota:            quota,
		GroupRatioInfo:   group,
		LegacyModelPrice: origin.PriceUSD,
		OtherRatios:      types.CloneRatios(origin.OtherRatios),
		UsePrice:         true,
		PerCallBilling:   origin.PerCallBilling,
		Finalized:        true,
	}, nil
}

func buildMatrixDecision(cfg *types.PricingConfig, result *pricingLookupResult, group types.GroupRatioInfo) *types.PricingDecision {
	billingUnit := cfg.BillingUnit
	if billingUnit == "" {
		billingUnit = types.BillingUnitPerCall
	}

	decision := &types.PricingDecision{
		Mode:           result.Mode,
		BillingMode:    types.BillingModeMatrix,
		BillingUnit:    billingUnit,
		GroupRatioInfo: group,
		Snapshot:       types.CloneMapAny(result.Snapshot),
		UsePrice:       true,
		Finalized:      true,
	}

	if billingUnit == types.BillingUnitPer1MTokens {
		decision.TokenUnitPriceUSD = result.PriceUSD
		decision.PriceUSD = result.PriceUSD
		decision.Quota = minimumUsagePreconsumeQuota(result.PriceUSD, group.GroupRatio)
		decision.PerCallBilling = false
		return decision
	}

	decision.PriceUSD = result.PriceUSD
	decision.Quota = int(result.PriceUSD * common.QuotaPerUnit * group.GroupRatio)
	decision.PerCallBilling = true
	if group.GroupRatio == 0 || result.PriceUSD == 0 {
		decision.FreeModel = true
		decision.Quota = 0
	}
	return decision
}

func fallbackPricingConfig(cfg *types.PricingConfig, resolved *ResolvedDimensionValues) (*pricingLookupResult, *dto.TaskError) {
	switch cfg.Fallback.Strategy {
	case types.PricingFallbackReject:
		return nil, localPricingTaskError("pricing dimensions did not match any row", "pricing_no_match")
	case types.PricingFallbackMax:
		price := maxPricingTablePrice(cfg.Table)
		return &pricingLookupResult{
			Mode:     types.PricingModeFallbackMax,
			PriceUSD: price,
			Snapshot: buildPricingSnapshot(cfg, resolved, price, types.PricingModeFallbackMax, cfg.Fallback.Strategy),
		}, nil
	case types.PricingFallbackDefault:
		price := cfg.Fallback.DefaultPrice
		return &pricingLookupResult{
			Mode:     types.PricingModeFallbackDefault,
			PriceUSD: price,
			Snapshot: buildPricingSnapshot(cfg, resolved, price, types.PricingModeFallbackDefault, cfg.Fallback.Strategy),
		}, nil
	default:
		return nil, localPricingTaskError(fmt.Sprintf("unsupported fallback strategy %q", cfg.Fallback.Strategy), "invalid_pricing_config")
	}
}

func buildPricingSnapshot(cfg *types.PricingConfig, resolved *ResolvedDimensionValues, price float64, mode types.PricingMode, fallbackStrategy string) map[string]any {
	snapshot := make(map[string]any, len(cfg.Dimensions)+4)
	for _, dim := range cfg.Dimensions {
		snapshot[dim.Key] = resolved.Values[dim.Key]
	}
	snapshot["price"] = price
	snapshot["billing_unit"] = cfg.BillingUnit
	snapshot["pricing_mode"] = string(mode)
	if fallbackStrategy != "" {
		snapshot["fallback_strategy"] = fallbackStrategy
	}
	return snapshot
}

func rowMatchesDimensions(dimensions []types.PricingDimension, row types.PricingRow, values map[string]any) (int, bool) {
	wildcards := 0
	for _, dim := range dimensions {
		rv := row[dim.Key]
		if s, ok := rv.(string); ok && s == "*" {
			wildcards++
			continue
		}
		if !samePricingLiteral(rv, values[dim.Key], dim.Type) {
			return 0, false
		}
	}
	return wildcards, true
}

func maxPricingTablePrice(rows []types.PricingRow) float64 {
	var max float64
	for i, row := range rows {
		price, ok := pricingNumberValue(row["price"])
		if !ok {
			continue
		}
		if i == 0 || price > max {
			max = price
		}
	}
	return max
}

func minimumUsagePreconsumeQuota(tokenUnitPriceUSD float64, groupRatio float64) int {
	if tokenUnitPriceUSD == 0 || groupRatio == 0 {
		return 0
	}
	if common.PreConsumedQuota <= 0 {
		return 0
	}
	quota := int(float64(common.PreConsumedQuota) * tokenUnitPriceUSD / 1_000_000 * common.QuotaPerUnit * groupRatio)
	if quota <= 0 {
		return 1
	}
	return quota
}

func priceDataFromDecision(decision *types.PricingDecision) types.PriceData {
	if decision == nil {
		return types.PriceData{}
	}
	return types.PriceData{
		FreeModel:       decision.FreeModel,
		ModelPrice:      decision.PriceUSD,
		Quota:           decision.Quota,
		GroupRatioInfo:  decision.GroupRatioInfo,
		OtherRatios:     types.CloneRatios(decision.OtherRatios),
		UsePrice:        decision.UsePrice,
		PricingDecision: decision,
	}
}

func pricingNumberValue(v any) (float64, bool) {
	switch n := v.(type) {
	case float64:
		return n, true
	case float32:
		return float64(n), true
	case int:
		return float64(n), true
	case int64:
		return float64(n), true
	case int32:
		return float64(n), true
	default:
		return 0, false
	}
}

func samePricingLiteral(a, b any, dimensionType string) bool {
	af, aok := pricingNumberValue(a)
	bf, bok := pricingNumberValue(b)
	if aok && bok {
		return af == bf
	}
	if dimensionType == "string" {
		return strings.EqualFold(strings.TrimSpace(fmt.Sprintf("%v", a)), strings.TrimSpace(fmt.Sprintf("%v", b)))
	}
	return fmt.Sprintf("%v", a) == fmt.Sprintf("%v", b)
}

func localPricingTaskError(message string, code string) *dto.TaskError {
	err := fmt.Errorf("%s", message)
	return &dto.TaskError{
		Code:       code,
		Message:    message,
		StatusCode: http.StatusBadRequest,
		LocalError: true,
		Error:      err,
	}
}