feat: add model display pricing admin api

controller/model_display_pricing_controller.go

@@ -0,0 +1,59 @@
+package controller
+
+import (
+	"net/http"
+	"strings"
+
+	"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 GetModelDisplayPricingRules(c *gin.Context) {
+	common.ApiSuccess(c, ratio_setting.GetModelDisplayPricingCopy())
+}
+
+func GetModelDisplayPricingRule(c *gin.Context) {
+	modelName := strings.TrimPrefix(c.Param("model"), "/")
+	items := ratio_setting.GetModelDisplayPricing(modelName)
+	if len(items) == 0 {
+		c.JSON(http.StatusNotFound, gin.H{"success": false, "message": "model display pricing not found"})
+		return
+	}
+	common.ApiSuccess(c, items)
+}
+
+func UpdateModelDisplayPricingRule(c *gin.Context) {
+	modelName := strings.TrimPrefix(c.Param("model"), "/")
+	var items []types.ModelDisplayPricingItem
+	if err := common.DecodeJson(c.Request.Body, &items); err != nil {
+		c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": err.Error()})
+		return
+	}
+	if err := ratio_setting.SetModelDisplayPricing(modelName, items); err != nil {
+		c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": err.Error()})
+		return
+	}
+	if err := model.UpdateOption(ratio_setting.ModelDisplayPricingOptionKey, ratio_setting.ModelDisplayPricing2JSONString()); err != nil {
+		c.JSON(http.StatusInternalServerError, gin.H{"success": false, "message": err.Error()})
+		return
+	}
+	model.RefreshPricing()
+	common.ApiSuccess(c, ratio_setting.GetModelDisplayPricing(modelName))
+}
+
+func DeleteModelDisplayPricingRule(c *gin.Context) {
+	modelName := strings.TrimPrefix(c.Param("model"), "/")
+	if err := ratio_setting.DeleteModelDisplayPricing(modelName); err != nil {
+		c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": err.Error()})
+		return
+	}
+	if err := model.UpdateOption(ratio_setting.ModelDisplayPricingOptionKey, ratio_setting.ModelDisplayPricing2JSONString()); err != nil {
+		c.JSON(http.StatusInternalServerError, gin.H{"success": false, "message": err.Error()})
+		return
+	}
+	model.RefreshPricing()
+	common.ApiSuccess(c, nil)
+}

controller/model_display_pricing_controller_test.go

@@ -0,0 +1,96 @@
+package controller
+
+import (
+	"net/http"
+	"testing"
+
+	"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"
+	"github.com/glebarez/sqlite"
+	"github.com/stretchr/testify/require"
+	"gorm.io/gorm"
+)
+
+func setupModelDisplayPricingControllerTest(t *testing.T) *gin.Engine {
+	t.Helper()
+	gin.SetMode(gin.TestMode)
+
+	db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
+	require.NoError(t, err)
+	require.NoError(t, db.AutoMigrate(&model.Option{}))
+	sqlDB, err := db.DB()
+	require.NoError(t, err)
+	sqlDB.SetMaxOpenConns(1)
+
+	origDB := model.DB
+	origUsingSQLite := common.UsingSQLite
+	origRedisEnabled := common.RedisEnabled
+	model.DB = db
+	common.UsingSQLite = true
+	common.RedisEnabled = false
+	model.InitOptionMap()
+	require.NoError(t, ratio_setting.UpdateModelDisplayPricingByJSONString(`{}`))
+
+	t.Cleanup(func() {
+		model.DB = origDB
+		common.UsingSQLite = origUsingSQLite
+		common.RedisEnabled = origRedisEnabled
+		require.NoError(t, ratio_setting.UpdateModelDisplayPricingByJSONString(`{}`))
+		require.NoError(t, sqlDB.Close())
+	})
+
+	r := gin.New()
+	r.GET("/api/option/model_display_pricing", GetModelDisplayPricingRules)
+	r.GET("/api/option/model_display_pricing/*model", GetModelDisplayPricingRule)
+	r.PUT("/api/option/model_display_pricing/*model", UpdateModelDisplayPricingRule)
+	r.DELETE("/api/option/model_display_pricing/*model", DeleteModelDisplayPricingRule)
+	return r
+}
+
+func TestModelDisplayPricingControllerCRUDAndWildcard(t *testing.T) {
+	r := setupModelDisplayPricingControllerTest(t)
+	items := []types.ModelDisplayPricingItem{
+		{Specification: "768p-6s", OfficialSupplierTip: "768P 6s", Price: 0.33333, Unit: "second", SortOrder: 1},
+	}
+
+	w := performJSONRequest(t, r, http.MethodPut, "/api/option/model_display_pricing/provider/kling-video", items)
+	require.Equal(t, http.StatusOK, w.Code)
+	require.Len(t, ratio_setting.GetModelDisplayPricing("provider/kling-video"), 1)
+
+	w = performJSONRequest(t, r, http.MethodGet, "/api/option/model_display_pricing/provider/kling-video", nil)
+	require.Equal(t, http.StatusOK, w.Code)
+	require.Contains(t, w.Body.String(), "768p-6s")
+
+	w = performJSONRequest(t, r, http.MethodGet, "/api/option/model_display_pricing", nil)
+	require.Equal(t, http.StatusOK, w.Code)
+	require.Contains(t, w.Body.String(), "provider/kling-video")
+
+	w = performJSONRequest(t, r, http.MethodDelete, "/api/option/model_display_pricing/provider/kling-video", nil)
+	require.Equal(t, http.StatusOK, w.Code)
+	require.Empty(t, ratio_setting.GetModelDisplayPricing("provider/kling-video"))
+}
+
+func TestModelDisplayPricingControllerValidationError(t *testing.T) {
+	r := setupModelDisplayPricingControllerTest(t)
+	items := []types.ModelDisplayPricingItem{{Specification: "", Price: 1, Unit: "second"}}
+
+	w := performJSONRequest(t, r, http.MethodPut, "/api/option/model_display_pricing/kling-video", items)
+	require.Equal(t, http.StatusBadRequest, w.Code)
+	require.Contains(t, w.Body.String(), "specification is required")
+}
+
+func TestModelDisplayPricingOptionMapHydratesCache(t *testing.T) {
+	r := setupModelDisplayPricingControllerTest(t)
+	require.NotNil(t, r)
+
+	require.NoError(t, model.UpdateOption(ratio_setting.ModelDisplayPricingOptionKey, `{
+		"kling-video":[{"specification":"768p","price":1,"unit":"second"}]
+	}`))
+
+	items := ratio_setting.GetModelDisplayPricing("kling-video")
+	require.Len(t, items, 1)
+	require.Equal(t, "768p", items[0].Specification)
+}

model/option.go

@@ -134,6 +134,7 @@ func InitOptionMap() {
 	common.OptionMap["ModelRatio"] = ratio_setting.ModelRatio2JSONString()
 	common.OptionMap["ModelPrice"] = ratio_setting.ModelPrice2JSONString()
 	common.OptionMap["ModelPricingRules"] = ratio_setting.ModelPricingRules2JSONString()
+	common.OptionMap["ModelDisplayPricing"] = ratio_setting.ModelDisplayPricing2JSONString()
 	common.OptionMap["CacheRatio"] = ratio_setting.CacheRatio2JSONString()
 	common.OptionMap["CreateCacheRatio"] = ratio_setting.CreateCacheRatio2JSONString()
 	common.OptionMap["GroupRatio"] = ratio_setting.GroupRatio2JSONString()
@@ -514,6 +515,8 @@ func updateOptionMap(key string, value string) (err error) {
 		err = ratio_setting.UpdateModelPriceByJSONString(value)
 	case "ModelPricingRules":
 		err = ratio_setting.UpdateModelPricingRulesByJSONString(value)
+	case "ModelDisplayPricing":
+		err = ratio_setting.UpdateModelDisplayPricingByJSONString(value)
 	case "CacheRatio":
 		err = ratio_setting.UpdateCacheRatioByJSONString(value)
 	case "CreateCacheRatio":

router/api-router.go

@@ -195,6 +195,10 @@ func SetApiRouter(router *gin.Engine) {
 			optionRoute.GET("/model_pricing/:model", controller.GetModelPricingRule)
 			optionRoute.PUT("/model_pricing/:model", controller.UpdateModelPricingRule)
 			optionRoute.DELETE("/model_pricing/:model", controller.DeleteModelPricingRule)
+			optionRoute.GET("/model_display_pricing", controller.GetModelDisplayPricingRules)
+			optionRoute.GET("/model_display_pricing/*model", controller.GetModelDisplayPricingRule)
+			optionRoute.PUT("/model_display_pricing/*model", controller.UpdateModelDisplayPricingRule)
+			optionRoute.DELETE("/model_display_pricing/*model", controller.DeleteModelDisplayPricingRule)
 			optionRoute.GET("/channel_affinity_cache", controller.GetChannelAffinityCacheStats)
 			optionRoute.DELETE("/channel_affinity_cache", controller.ClearChannelAffinityCache)
 			optionRoute.POST("/rest_model_ratio", controller.ResetModelRatio)