From 18a3495968be7443fcae6efb8755cb076e2e613d Mon Sep 17 00:00:00 2001 From: fengsilin Date: Wed, 1 Jul 2026 11:21:00 +0800 Subject: [PATCH] feat: add model display pricing settings --- .../ratio_setting/model_display_pricing.go | 120 ++++++++++++++++++ .../model_display_pricing_test.go | 95 ++++++++++++++ types/model_display_pricing.go | 19 +++ 3 files changed, 234 insertions(+) create mode 100644 setting/ratio_setting/model_display_pricing.go create mode 100644 setting/ratio_setting/model_display_pricing_test.go create mode 100644 types/model_display_pricing.go diff --git a/setting/ratio_setting/model_display_pricing.go b/setting/ratio_setting/model_display_pricing.go new file mode 100644 index 0000000..b53e2e8 --- /dev/null +++ b/setting/ratio_setting/model_display_pricing.go @@ -0,0 +1,120 @@ +package ratio_setting + +import ( + "fmt" + "sort" + "strings" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/types" +) + +const ModelDisplayPricingOptionKey = "ModelDisplayPricing" + +var displayPricingMap = types.NewRWMap[string, []types.ModelDisplayPricingItem]() + +func GetModelDisplayPricing(model string) []types.ModelDisplayPricingItem { + items, ok := displayPricingMap.Get(model) + if !ok { + return nil + } + return types.CloneModelDisplayPricingItems(items) +} + +func GetModelDisplayPricingCopy() map[string][]types.ModelDisplayPricingItem { + src := displayPricingMap.ReadAll() + dst := make(map[string][]types.ModelDisplayPricingItem, len(src)) + for model, items := range src { + dst[model] = types.CloneModelDisplayPricingItems(items) + } + return dst +} + +func SetModelDisplayPricing(model string, items []types.ModelDisplayPricingItem) error { + model = strings.TrimSpace(model) + if model == "" { + return fmt.Errorf("model name is required") + } + normalized, err := normalizeModelDisplayPricingItems(items) + if err != nil { + return err + } + displayPricingMap.Set(model, normalized) + InvalidateExposedDataCache() + return nil +} + +func DeleteModelDisplayPricing(model string) error { + all := displayPricingMap.ReadAll() + delete(all, model) + displayPricingMap.Clear() + displayPricingMap.AddAll(all) + InvalidateExposedDataCache() + return nil +} + +func ModelDisplayPricing2JSONString() string { + bytes, err := common.Marshal(GetModelDisplayPricingCopy()) + if err != nil { + return "{}" + } + return string(bytes) +} + +func UpdateModelDisplayPricingByJSONString(jsonStr string) error { + if strings.TrimSpace(jsonStr) == "" { + jsonStr = "{}" + } + raw := map[string][]types.ModelDisplayPricingItem{} + if err := common.UnmarshalJsonStr(jsonStr, &raw); err != nil { + return err + } + next := make(map[string][]types.ModelDisplayPricingItem, len(raw)) + for model, items := range raw { + model = strings.TrimSpace(model) + if model == "" { + return fmt.Errorf("model name is required") + } + normalized, err := normalizeModelDisplayPricingItems(items) + if err != nil { + return fmt.Errorf("model %s: %w", model, err) + } + next[model] = normalized + } + displayPricingMap.Clear() + displayPricingMap.AddAll(next) + InvalidateExposedDataCache() + return nil +} + +func normalizeModelDisplayPricingItems(items []types.ModelDisplayPricingItem) ([]types.ModelDisplayPricingItem, error) { + normalized := types.CloneModelDisplayPricingItems(items) + for i := range normalized { + item := &normalized[i] + item.Specification = strings.TrimSpace(item.Specification) + item.OfficialSupplierTip = strings.TrimSpace(item.OfficialSupplierTip) + item.Unit = strings.TrimSpace(item.Unit) + if item.Specification == "" { + return nil, fmt.Errorf("row %d: specification is required", i) + } + if item.Price < 0 { + return nil, fmt.Errorf("row %d: price must be >= 0", i) + } + if item.Unit == "" { + return nil, fmt.Errorf("row %d: unit is required", i) + } + if item.DiscountRate == 0 { + item.DiscountRate = 1 + } + if item.DiscountRate <= 0 || item.DiscountRate > 1 { + return nil, fmt.Errorf("row %d: discount_rate must be > 0 and <= 1", i) + } + } + sort.SliceStable(normalized, func(i, j int) bool { + if normalized[i].SortOrder != normalized[j].SortOrder { + return normalized[i].SortOrder < normalized[j].SortOrder + } + return normalized[i].Specification < normalized[j].Specification + }) + return normalized, nil +} diff --git a/setting/ratio_setting/model_display_pricing_test.go b/setting/ratio_setting/model_display_pricing_test.go new file mode 100644 index 0000000..372ce39 --- /dev/null +++ b/setting/ratio_setting/model_display_pricing_test.go @@ -0,0 +1,95 @@ +package ratio_setting + +import ( + "testing" + + "github.com/QuantumNous/new-api/types" + "github.com/stretchr/testify/require" +) + +func TestModelDisplayPricingRoundTripNormalizeAndSort(t *testing.T) { + require.NoError(t, UpdateModelDisplayPricingByJSONString(`{}`)) + + jsonStr := `{ + "kling-video": [ + {"specification":"1080p","official_supplier_tip":"1080P","price":0.58333,"unit":"second","sort_order":3}, + {"specification":"768p-6s","official_supplier_tip":"768P 6s","price":0.33333,"unit":"second","sort_order":1,"discount_rate":0.8}, + {"specification":"768p-10s","official_supplier_tip":"768P 10s","price":0.4,"unit":"second","sort_order":2,"discount_rate":0} + ] + }` + + require.NoError(t, UpdateModelDisplayPricingByJSONString(jsonStr)) + items := GetModelDisplayPricing("kling-video") + require.Len(t, items, 3) + require.Equal(t, "768p-6s", items[0].Specification) + require.Equal(t, 0.8, items[0].DiscountRate) + require.Equal(t, "768p-10s", items[1].Specification) + require.Equal(t, 1.0, items[1].DiscountRate) + require.Equal(t, "1080p", items[2].Specification) + require.Equal(t, 1.0, items[2].DiscountRate) + + items[0].Specification = "mutated" + require.Equal(t, "768p-6s", GetModelDisplayPricing("kling-video")[0].Specification) + + roundTrip := ModelDisplayPricing2JSONString() + require.NoError(t, UpdateModelDisplayPricingByJSONString(roundTrip)) + require.Equal(t, "768p-6s", GetModelDisplayPricing("kling-video")[0].Specification) +} + +func TestSetModelDisplayPricingValidation(t *testing.T) { + tests := []struct { + name string + model string + items []types.ModelDisplayPricingItem + want string + }{ + { + name: "empty_model", + model: "", + items: []types.ModelDisplayPricingItem{{Specification: "768p", Price: 1, Unit: "second"}}, + want: "model name is required", + }, + { + name: "empty_specification", + model: "kling-video", + items: []types.ModelDisplayPricingItem{{Price: 1, Unit: "second"}}, + want: "specification is required", + }, + { + name: "negative_price", + model: "kling-video", + items: []types.ModelDisplayPricingItem{{Specification: "768p", Price: -1, Unit: "second"}}, + want: "price must be >= 0", + }, + { + name: "empty_unit", + model: "kling-video", + items: []types.ModelDisplayPricingItem{{Specification: "768p", Price: 1}}, + want: "unit is required", + }, + { + name: "bad_discount", + model: "kling-video", + items: []types.ModelDisplayPricingItem{{Specification: "768p", Price: 1, Unit: "second", DiscountRate: 1.2}}, + want: "discount_rate must be > 0 and <= 1", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := SetModelDisplayPricing(tt.model, tt.items) + require.Error(t, err) + require.Contains(t, err.Error(), tt.want) + }) + } +} + +func TestDeleteModelDisplayPricing(t *testing.T) { + require.NoError(t, UpdateModelDisplayPricingByJSONString(`{}`)) + require.NoError(t, SetModelDisplayPricing("kling-video", []types.ModelDisplayPricingItem{ + {Specification: "768p", Price: 1, Unit: "second"}, + })) + require.Len(t, GetModelDisplayPricing("kling-video"), 1) + require.NoError(t, DeleteModelDisplayPricing("kling-video")) + require.Empty(t, GetModelDisplayPricing("kling-video")) +} diff --git a/types/model_display_pricing.go b/types/model_display_pricing.go new file mode 100644 index 0000000..3155aea --- /dev/null +++ b/types/model_display_pricing.go @@ -0,0 +1,19 @@ +package types + +type ModelDisplayPricingItem struct { + Specification string `json:"specification"` + OfficialSupplierTip string `json:"official_supplier_tip"` + Price float64 `json:"price"` + Unit string `json:"unit"` + SortOrder int `json:"sort_order"` + DiscountRate float64 `json:"discount_rate"` +} + +func CloneModelDisplayPricingItems(src []ModelDisplayPricingItem) []ModelDisplayPricingItem { + if src == nil { + return nil + } + dst := make([]ModelDisplayPricingItem, len(src)) + copy(dst, src) + return dst +}