From de117ac932bdfd6b472a0d50edb353d0ca32366f Mon Sep 17 00:00:00 2001 From: fengsilin Date: Wed, 1 Jul 2026 11:28:26 +0800 Subject: [PATCH] feat: add model display pricing admin api --- .../model_display_pricing_controller.go | 59 ++++++++++++ .../model_display_pricing_controller_test.go | 96 +++++++++++++++++++ model/option.go | 3 + router/api-router.go | 4 + 4 files changed, 162 insertions(+) create mode 100644 controller/model_display_pricing_controller.go create mode 100644 controller/model_display_pricing_controller_test.go diff --git a/controller/model_display_pricing_controller.go b/controller/model_display_pricing_controller.go new file mode 100644 index 0000000..6f40e1e --- /dev/null +++ b/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) +} diff --git a/controller/model_display_pricing_controller_test.go b/controller/model_display_pricing_controller_test.go new file mode 100644 index 0000000..bd9da00 --- /dev/null +++ b/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) +} diff --git a/model/option.go b/model/option.go index 056bfac..a0c2336 100644 --- a/model/option.go +++ b/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": diff --git a/router/api-router.go b/router/api-router.go index edb2a7c..ce023a8 100644 --- a/router/api-router.go +++ b/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)