feat: sync deployed v1.0.0-rc.38-custom.67

model/personal_group_ratio.go

@@ -0,0 +1,87 @@
+package model
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"math"
+
+	"github.com/QuantumNous/new-api/setting/ratio_setting"
+	"github.com/QuantumNous/new-api/types"
+	"github.com/gin-gonic/gin"
+	"gorm.io/gorm/clause"
+)
+
+// Stored independently from User and UserSetting to prevent profile mass assignment.
+type PersonalGroupRatio struct {
+	UserID int    `gorm:"primaryKey;autoIncrement:false"`
+	Ratios string `gorm:"type:text;not null"`
+}
+
+const personalRatioContextKey = "personal_group_ratios"
+
+func GetPersonalGroupRatios(ctx context.Context, userID int) (map[string]float64, error) {
+	var rows []PersonalGroupRatio
+	if err := DB.WithContext(ctx).Where("user_id = ?", userID).Find(&rows).Error; err != nil {
+		return nil, err
+	}
+	result := map[string]float64{}
+	if len(rows) != 0 {
+		if err := json.Unmarshal([]byte(rows[0].Ratios), &result); err != nil {
+			return nil, err
+		}
+	}
+	return result, nil
+}
+
+func SavePersonalGroupRatios(ctx context.Context, userID int, ratios map[string]float64) error {
+	for group, ratio := range ratios {
+		if !ratio_setting.ContainsGroupRatio(group) || group == "auto" || ratio < 0 || math.IsNaN(ratio) || math.IsInf(ratio, 0) {
+			return fmt.Errorf("invalid group or ratio: %s", group)
+		}
+	}
+	if ratios == nil {
+		ratios = map[string]float64{}
+	}
+	body, err := json.Marshal(ratios)
+	if err != nil {
+		return err
+	}
+	row := PersonalGroupRatio{UserID: userID, Ratios: string(body)}
+	if err = DB.WithContext(ctx).Clauses(clause.OnConflict{Columns: []clause.Column{{Name: "user_id"}}, DoUpdates: clause.AssignmentColumns([]string{"ratios"})}).Create(&row).Error; err != nil {
+		return err
+	}
+	// Pricing is loaded from the database at request admission, never from a stale
+	// process/Redis cache. The request snapshot survives retries and price updates.
+	return invalidateUserCache(userID)
+}
+
+func LoadPersonalGroupRatios(c *gin.Context, userID int) error {
+	ratios, err := GetPersonalGroupRatios(c.Request.Context(), userID)
+	if err == nil {
+		c.Set(personalRatioContextKey, ratios)
+	}
+	return err
+}
+
+func EffectiveGroupRatio(c *gin.Context, userGroup, group string) types.GroupRatioInfo {
+	result := types.GroupRatioInfo{GroupRatio: ratio_setting.GetGroupRatio(group), GroupSpecialRatio: -1}
+	if ratio, ok := ratio_setting.GetGroupGroupRatio(userGroup, group); ok {
+		result.GroupRatio, result.GroupSpecialRatio, result.HasSpecialRatio = ratio, ratio, true
+	}
+	if value, ok := c.Get(personalRatioContextKey); ok {
+		if ratios, ok := value.(map[string]float64); ok {
+			if ratio, ok := ratios[group]; ok {
+				result.GroupRatio, result.GroupSpecialRatio, result.HasSpecialRatio = ratio, ratio, true
+			}
+		}
+	}
+	return result
+}
+
+func HasPersonalGroupRatio(c *gin.Context, group string) bool {
+	value, _ := c.Get(personalRatioContextKey)
+	ratios, _ := value.(map[string]float64)
+	_, ok := ratios[group]
+	return ok
+}