diff --git a/controller/pricing.go b/controller/pricing.go index de90e0b..12bcaf6 100644 --- a/controller/pricing.go +++ b/controller/pricing.go @@ -138,6 +138,9 @@ func GetUserPricing(c *gin.Context) { 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 { @@ -184,6 +187,9 @@ func respondOriginalPrice(c *gin.Context, pricingData *model.Pricing) { } result["original_price"] = pricingData.ModelPrice } + if len(pricingData.DisplayPricing) > 0 { + result["display_pricing"] = pricingData.DisplayPricing + } c.JSON(200, result) } diff --git a/controller/pricing_user_test.go b/controller/pricing_user_test.go index 31326b8..33e93a1 100644 --- a/controller/pricing_user_test.go +++ b/controller/pricing_user_test.go @@ -9,8 +9,9 @@ import ( "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/setting/ratio_setting" - "github.com/glebarez/sqlite" + "github.com/QuantumNous/new-api/types" "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "gorm.io/gorm" @@ -152,6 +153,36 @@ func TestGetUserPricing_NotLoggedIn_PerCall(t *testing.T) { assert.Equal(t, float64(0.04), resp["original_price"]) } +func TestGetUserPricingIncludesDisplayPricingWithoutChangingComputedPrice(t *testing.T) { + setupPricingTestDB(t) + modelName := "display-pricing-user-model" + router := gin.New() + router.GET("/api/pricing/user/*model", GetUserPricing) + + setPricingCache([]model.Pricing{ + { + ModelName: modelName, + QuotaType: model.QuotaTypeByCall, + ModelPrice: 0.5, + EnableGroup: []string{"default"}, + DisplayPricing: []types.ModelDisplayPricingItem{ + {Specification: "768p-6s", Price: 0.33333, Unit: "second", SortOrder: 1, DiscountRate: 1}, + }, + }, + }) + + w := httptest.NewRecorder() + req, _ := http.NewRequest("GET", "/api/pricing/user/"+modelName, nil) + router.ServeHTTP(w, req) + + var resp map[string]interface{} + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + require.Equal(t, http.StatusOK, w.Code) + require.Contains(t, resp, "display_pricing") + require.Equal(t, float64(0.5), resp["original_price"]) + require.NotEqual(t, float64(0.33333), resp["original_price"]) +} + // TestGetUserPricing_LoggedIn_NoDiscount 已登录但无折扣（分组倍率=1，无个人倍率） func TestGetUserPricing_LoggedIn_NoDiscount(t *testing.T) { db := setupPricingTestDB(t) @@ -226,8 +257,6 @@ func TestGetUserPricing_LoggedIn_GroupDiscount(t *testing.T) { assert.Equal(t, "8折", resp["discount"]) } - - // TestGetUserPricing_PerCall_WithDiscount 按次计费 + 折扣 func TestGetUserPricing_PerCall_WithDiscount(t *testing.T) { db := setupPricingTestDB(t) @@ -280,4 +309,3 @@ func TestFormatDiscount(t *testing.T) { assert.Equal(t, tt.expected, result, "ratio=%.2f", tt.ratio) } } - diff --git a/model/pricing.go b/model/pricing.go index cf53161..eb5db68 100644 --- a/model/pricing.go +++ b/model/pricing.go @@ -14,23 +14,24 @@ import ( ) type Pricing struct { - ModelName string `json:"model_name"` - Description string `json:"description,omitempty"` - Icon string `json:"icon,omitempty"` - Tags string `json:"tags,omitempty"` - VendorID int `json:"vendor_id,omitempty"` - QuotaType int `json:"quota_type"` - ModelRatio float64 `json:"model_ratio"` - ModelPrice float64 `json:"model_price"` - OwnerBy string `json:"owner_by"` - CompletionRatio float64 `json:"completion_ratio"` - CacheRatio float64 `json:"cache_ratio"` - CacheCreationRatio float64 `json:"cache_creation_ratio"` - EnableGroup []string `json:"enable_groups"` - SupportedEndpointTypes []constant.EndpointType `json:"supported_endpoint_types"` - PricingVersion string `json:"pricing_version,omitempty"` - Type int `json:"type"` - PricingConfig *types.PricingConfig `json:"pricing_config,omitempty"` + ModelName string `json:"model_name"` + Description string `json:"description,omitempty"` + Icon string `json:"icon,omitempty"` + Tags string `json:"tags,omitempty"` + VendorID int `json:"vendor_id,omitempty"` + QuotaType int `json:"quota_type"` + ModelRatio float64 `json:"model_ratio"` + ModelPrice float64 `json:"model_price"` + OwnerBy string `json:"owner_by"` + CompletionRatio float64 `json:"completion_ratio"` + CacheRatio float64 `json:"cache_ratio"` + CacheCreationRatio float64 `json:"cache_creation_ratio"` + EnableGroup []string `json:"enable_groups"` + SupportedEndpointTypes []constant.EndpointType `json:"supported_endpoint_types"` + PricingVersion string `json:"pricing_version,omitempty"` + Type int `json:"type"` + PricingConfig *types.PricingConfig `json:"pricing_config,omitempty"` + DisplayPricing []types.ModelDisplayPricingItem `json:"display_pricing,omitempty"` } type PricingVendor struct { @@ -340,6 +341,7 @@ func updatePricing() { } applyGlobalDefault(&pricing, model) + applyDisplayPricing(&pricing, model) pricingMap = append(pricingMap, pricing) } @@ -426,6 +428,13 @@ func applyMatrixPricingConfig(pricing *Pricing, model string) bool { return true } +func applyDisplayPricing(pricing *Pricing, model string) { + items := ratio_setting.GetModelDisplayPricing(model) + if len(items) > 0 { + pricing.DisplayPricing = items + } +} + func minPricingConfigPrice(cfg *types.PricingConfig) float64 { if cfg == nil || len(cfg.Table) == 0 { return 0 diff --git a/model/pricing_test.go b/model/pricing_test.go index 2070c8b..fbe4b5e 100644 --- a/model/pricing_test.go +++ b/model/pricing_test.go @@ -189,6 +189,41 @@ func TestGetPricingByModel_IncludesMatrixPricingConfig(t *testing.T) { require.Equal(t, []constant.EndpointType{constant.EndpointTypeDoubaoVideo}, pricing.SupportedEndpointTypes) } +func TestGetPricingByModelIncludesDisplayPricing(t *testing.T) { + setupPricingTest(t) + + const displayModel = "display-pricing-model" + require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"`+displayModel+`":10}`)) + require.NoError(t, ratio_setting.SetModelDisplayPricing(displayModel, []types.ModelDisplayPricingItem{ + {Specification: "768p-6s", Price: 0.33333, Unit: "second", SortOrder: 1}, + })) + + require.NoError(t, (&Model{ + ModelName: displayModel, + Status: 1, + }).Insert()) + require.NoError(t, (&Channel{ + Id: 1, + Type: constant.ChannelTypeOpenAI, + Key: "k", + Name: "channel-a", + Status: common.ChannelStatusEnabled, + Models: displayModel, + Group: "default", + }).Insert()) + + RefreshPricing() + pricing := GetPricingByModel(displayModel) + + require.NotNil(t, pricing) + require.Len(t, pricing.DisplayPricing, 1) + require.Equal(t, "768p-6s", pricing.DisplayPricing[0].Specification) + require.Equal(t, QuotaTypeByTokens, pricing.QuotaType) + require.Equal(t, 10.0, pricing.ModelRatio) + require.Equal(t, 0.0, pricing.ModelPrice) + require.Nil(t, pricing.PricingConfig) +} + func TestParseModelEndpointTypes_SupportsObjectAndArray(t *testing.T) { objectEndpoints := parseModelEndpointTypes(`{ "anthropic": {"path": "/v1/messages", "method": "POST"},