new-api/controller/pricing.go

package controller

import (
	"fmt"
	"math"
	"sort"
	"strings"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/model"
	"github.com/QuantumNous/new-api/service"
	"github.com/QuantumNous/new-api/setting/ratio_setting"
	"github.com/QuantumNous/new-api/types"

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

// resolveCurrentUser 从 gin.Context 解析当前登录用户信息。
// 未登录或查询失败时 loggedIn=false，groupRatio=1.0。
func resolveCurrentUser(c *gin.Context) (loggedIn bool, userId int, userGroup string, groupRatio float64) {
	groupRatio = 1.0
	raw, exists := c.Get("id")
	if !exists {
		return
	}
	userId = raw.(int)
	user, err := model.GetUserCache(userId)
	if err != nil {
		return
	}
	loggedIn = true
	userGroup = user.Group
	groupRatio = service.GetUserGroupRatio(userGroup, userGroup)
	return
}

func filterPricingByUsableGroups(pricing []model.Pricing, usableGroup map[string]string) []model.Pricing {
	if len(pricing) == 0 {
		return pricing
	}
	if len(usableGroup) == 0 {
		return []model.Pricing{}
	}

	filtered := make([]model.Pricing, 0, len(pricing))
	for _, item := range pricing {
		if common.StringsContains(item.EnableGroup, "all") {
			filtered = append(filtered, item)
			continue
		}
		for _, group := range item.EnableGroup {
			if _, ok := usableGroup[group]; ok {
				filtered = append(filtered, item)
				break
			}
		}
	}
	return filtered
}

func GetPricing(c *gin.Context) {
	pricing := model.GetPricing()
	userId, exists := c.Get("id")
	usableGroup := map[string]string{}
	groupRatio := map[string]float64{}
	for s, f := range ratio_setting.GetGroupRatioCopy() {
		groupRatio[s] = f
	}
	var group string
	if exists {
		user, err := model.GetUserCache(userId.(int))
		if err == nil {
			group = user.Group
			for g := range groupRatio {
				ratio, ok := ratio_setting.GetGroupGroupRatio(group, g)
				if ok {
					groupRatio[g] = ratio
				}
			}
		}
	}

	usableGroup = service.GetUserUsableGroups(group)
	pricing = filterPricingByUsableGroups(pricing, usableGroup)
	// check groupRatio contains usableGroup
	for group := range ratio_setting.GetGroupRatioCopy() {
		if _, ok := usableGroup[group]; !ok {
			delete(groupRatio, group)
		}
	}

	c.JSON(200, gin.H{
		"success":            true,
		"data":               pricing,
		"vendors":            model.GetVendors(),
		"group_ratio":        groupRatio,
		"usable_group":       usableGroup,
		"supported_endpoint": model.GetSupportedEndpointMap(),
		"auto_groups":        service.GetUserAutoGroup(group),
		"_":                  "a42d372ccf0b5dd13ecf71203521f9d2",
	})
}

// GetUserPricing 获取用户对指定模型的价格信息（原价 vs 用户价）
func GetUserPricing(c *gin.Context) {
	modelName := c.Param("model")
	// *model 通配符返回值带前导 /，需要去掉
	modelName = strings.TrimPrefix(modelName, "/")
	if modelName == "" {
		common.ApiErrorMsg(c, "模型名不能为空")
		return
	}

	pricingData := model.GetPricingByModel(modelName)
	if pricingData == nil {
		common.ApiErrorMsg(c, "未找到该模型的定价信息")
		return
	}

	loggedIn, _, userGroup, groupRatio := resolveCurrentUser(c)
	if !loggedIn {
		respondOriginalPrice(c, pricingData)
		return
	}

	totalRatio := groupRatio
	savingsPercent := int(math.Round((1 - totalRatio) * 100))

	result := gin.H{
		"success":         true,
		"model_name":      pricingData.ModelName,
		"quota_type":      pricingData.QuotaType,
		"group":           userGroup,
		"group_ratio":     groupRatio,
		"savings_percent": savingsPercent,
		"logged_in":       true,
	}
	if pricingData.PricingConfig != nil {
		result["pricing_config"] = pricingData.PricingConfig
	}
	if len(pricingData.DisplayPricing) > 0 {
		result["display_pricing"] = pricingData.DisplayPricing
	}

	if pricingData.QuotaType == model.QuotaTypeByTokens {
		if pricingData.PricingConfig != nil {
			originalPrice := representativeMatrixPrice(pricingData.PricingConfig, pricingData.ModelPrice)
			result["original_price"] = originalPrice
			result["user_price"] = originalPrice * totalRatio
		}
		originalInput := pricingData.ModelRatio * 2
		originalOutput := pricingData.ModelRatio * pricingData.CompletionRatio * 2
		result["original_input"] = originalInput
		result["original_output"] = originalOutput
		result["user_input"] = originalInput * totalRatio
		result["user_output"] = originalOutput * totalRatio
	} else {
		result["original_price"] = pricingData.ModelPrice
		result["user_price"] = pricingData.ModelPrice * totalRatio
	}

	if savingsPercent > 0 {
		result["discount"] = formatDiscount(totalRatio)
	}

	c.JSON(200, result)
}

// respondOriginalPrice 未登录或用户查询失败时，只返回原价
func respondOriginalPrice(c *gin.Context, pricingData *model.Pricing) {
	result := gin.H{
		"success":    true,
		"model_name": pricingData.ModelName,
		"quota_type": pricingData.QuotaType,
		"logged_in":  false,
	}
	if pricingData.QuotaType == model.QuotaTypeByTokens {
		if pricingData.PricingConfig != nil {
			result["pricing_config"] = pricingData.PricingConfig
			result["original_price"] = representativeMatrixPrice(pricingData.PricingConfig, pricingData.ModelPrice)
		}
		result["original_input"] = pricingData.ModelRatio * 2
		result["original_output"] = pricingData.ModelRatio * pricingData.CompletionRatio * 2
	} else {
		if pricingData.PricingConfig != nil {
			result["pricing_config"] = pricingData.PricingConfig
		}
		result["original_price"] = pricingData.ModelPrice
	}
	if len(pricingData.DisplayPricing) > 0 {
		result["display_pricing"] = pricingData.DisplayPricing
	}
	c.JSON(200, result)
}

func representativeMatrixPrice(cfg *types.PricingConfig, fallback float64) float64 {
	if cfg == nil || len(cfg.Table) == 0 {
		return fallback
	}
	prices := make([]float64, 0, len(cfg.Table))
	for _, row := range cfg.Table {
		price, ok := pricingRowFloat(row["price"])
		if ok {
			prices = append(prices, price)
		}
	}
	if len(prices) == 0 {
		return fallback
	}
	sort.Float64s(prices)
	mid := len(prices) / 2
	if len(prices)%2 == 1 {
		return prices[mid]
	}
	return (prices[mid-1] + prices[mid]) / 2
}

func pricingRowFloat(value any) (float64, bool) {
	switch n := value.(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
	}
}

// formatDiscount 将倍率转换为中文折扣格式
func formatDiscount(ratio float64) string {
	if ratio <= 0 {
		return "免费"
	}
	rawDiscount := math.Round(ratio*100) / 10
	if rawDiscount >= 10 {
		return ""
	}
	if rawDiscount == math.Trunc(rawDiscount) {
		return fmt.Sprintf("%.0f折", rawDiscount)
	}
	return fmt.Sprintf("%.1f折", rawDiscount)
}

func ResetModelRatio(c *gin.Context) {
	defaultStr := ratio_setting.DefaultModelRatio2JSONString()
	err := model.UpdateOption("ModelRatio", defaultStr)
	if err != nil {
		c.JSON(200, gin.H{
			"success": false,
			"message": err.Error(),
		})
		return
	}
	err = ratio_setting.UpdateModelRatioByJSONString(defaultStr)
	if err != nil {
		c.JSON(200, gin.H{
			"success": false,
			"message": err.Error(),
		})
		return
	}
	c.JSON(200, gin.H{
		"success": true,
		"message": "重置模型倍率成功",
	})
}