Kaynağa Gözat

feat: add model display pricing settings

master
fengsilin 1 hafta önce
ebeveyn
işleme
18a3495968
3 değiştirilmiş dosya ile 234 ekleme ve 0 silme
  1. +120
    -0
      setting/ratio_setting/model_display_pricing.go
  2. +95
    -0
      setting/ratio_setting/model_display_pricing_test.go
  3. +19
    -0
      types/model_display_pricing.go

+ 120
- 0
setting/ratio_setting/model_display_pricing.go Dosyayı Görüntüle

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

+ 95
- 0
setting/ratio_setting/model_display_pricing_test.go Dosyayı Görüntüle

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

+ 19
- 0
types/model_display_pricing.go Dosyayı Görüntüle

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

Yükleniyor…
İptal
Kaydet