new-api/controller/model_pricing_controller.go

package controller

import (
	"net/http"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/model"
	"github.com/QuantumNous/new-api/setting/ratio_setting"
	"github.com/QuantumNous/new-api/types"
	"github.com/gin-gonic/gin"
)

func GetModelPricingRules(c *gin.Context) {
	common.ApiSuccess(c, ratio_setting.GetPricingConfigCopy())
}

func GetModelPricingRule(c *gin.Context) {
	modelName := c.Param("model")
	cfg := ratio_setting.GetPricingConfig(modelName)
	if cfg == nil {
		c.JSON(http.StatusNotFound, gin.H{"success": false, "message": "model pricing rule not found"})
		return
	}
	common.ApiSuccess(c, cfg)
}

func UpdateModelPricingRule(c *gin.Context) {
	modelName := c.Param("model")
	var cfg types.PricingConfig
	if err := common.DecodeJson(c.Request.Body, &cfg); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": err.Error()})
		return
	}
	if err := ratio_setting.SetPricingConfig(modelName, &cfg); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": err.Error()})
		return
	}
	if err := model.UpdateOption(ratio_setting.ModelPricingRulesOptionKey, ratio_setting.ModelPricingRules2JSONString()); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"success": false, "message": err.Error()})
		return
	}
	model.RefreshPricing()
	common.ApiSuccess(c, ratio_setting.GetPricingConfig(modelName))
}

func DeleteModelPricingRule(c *gin.Context) {
	modelName := c.Param("model")
	if err := ratio_setting.DeletePricingConfig(modelName); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": err.Error()})
		return
	}
	if err := model.UpdateOption(ratio_setting.ModelPricingRulesOptionKey, ratio_setting.ModelPricingRules2JSONString()); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"success": false, "message": err.Error()})
		return
	}
	model.RefreshPricing()
	common.ApiSuccess(c, nil)
}