| @@ -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 | |||
| } | |||
| @@ -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")) | |||
| } | |||
| @@ -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 | |||
| } | |||