feat: publish v1.0.0-rc.36-custom.1

model/pricing_usage_schema_test.go

@@ -0,0 +1,141 @@
+package model
+
+import (
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/QuantumNous/new-api/constant"
+	"github.com/QuantumNous/new-api/pkg/jsplugin"
+	"github.com/QuantumNous/new-api/relaykit/dto"
+	"github.com/QuantumNous/new-api/setting/config"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func pricingUsagePluginSource(version, usageSchema string) string {
+	return fmt.Sprintf(`
+export const meta = {
+  apiVersion: 1, key: "pricing-usage-probe", name: "Pricing Usage Probe", version: %q, author: {name: "Test"},
+  models: ["pricing-usage-model"], fetchMode: "per_task", usageSchema: %s
+};
+export function buildSubmitRequest() { return {}; }
+export function parseSubmitResponse() { return {}; }
+export function buildQueryRequest() { return {}; }
+export function parseTaskResult() { return {}; }
+`, version, usageSchema)
+}
+
+func TestPricingCarriesTaskUsageSchemaAndRefreshesWithPluginGeneration(t *testing.T) {
+	resetPricingEndpointTestTables(t)
+	const pluginKey = "pricing-usage-probe"
+	initialSource := pricingUsagePluginSource("1.0.0", `{
+  seconds: {type: "number", unit: "second", description: "Estimated duration."},
+  action: {enum:["video"],enumLabels:{video:{en:"Generate video",zh:"生成视频"}}}
+}`)
+	_, err := jsplugin.DefaultRegistry.Register(initialSource, jsplugin.Options{})
+	require.NoError(t, err)
+	t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister(pluginKey) })
+
+	insertPricingEndpointChannel(t, 901, constant.ChannelTypeTaskPlugin, dto.ChannelOtherSettings{})
+	insertPricingEndpointAbility(t, 901, "pricing-usage-model")
+	insertPricingEndpointAbility(t, 901, "ordinary-model")
+
+	initialPricing := pricingByModel(GetPricing())
+	require.Contains(t, initialPricing, "pricing-usage-model")
+	require.Contains(t, initialPricing, "ordinary-model")
+	assert.Equal(t, "second", initialPricing["pricing-usage-model"].BillingUsageSchema["seconds"].Unit)
+	assert.Equal(t, "Estimated duration.", initialPricing["pricing-usage-model"].BillingUsageSchema["seconds"].Description["en"])
+	assert.Equal(t, "生成视频", initialPricing["pricing-usage-model"].BillingUsageSchema["action"].EnumLabels["video"]["zh"])
+	assert.Nil(t, initialPricing["ordinary-model"].BillingUsageSchema)
+
+	updatedSource := pricingUsagePluginSource("1.1.0", `{
+  seconds: {type: "number", unit: "second", description: "Measured duration."},
+  clips: {type: "number", unit: "count", description: "Generated clip count."}
+}`)
+	_, err = jsplugin.DefaultRegistry.Register(updatedSource, jsplugin.Options{})
+	require.NoError(t, err)
+	lastGetPricingTime = time.Now().Add(-2 * time.Minute)
+
+	refreshedPricing := pricingByModel(GetPricing())
+	require.Len(t, refreshedPricing["pricing-usage-model"].BillingUsageSchema, 2)
+	assert.Equal(t, "Measured duration.", refreshedPricing["pricing-usage-model"].BillingUsageSchema["seconds"].Description["en"])
+	assert.Equal(t, "count", refreshedPricing["pricing-usage-model"].BillingUsageSchema["clips"].Unit)
+}
+
+func TestPricingAliasCarriesPluginUsageSchemaAndTailExpr(t *testing.T) {
+	resetPricingEndpointTestTables(t)
+	const pluginKey = "pricing-usage-probe"
+	source := pricingUsagePluginSource("1.0.0", `{
+  seconds: {type: "number", unit: "second", description: "Estimated duration."}
+}`)
+	_, err := jsplugin.DefaultRegistry.Register(source, jsplugin.Options{})
+	require.NoError(t, err)
+	t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister(pluginKey) })
+
+	mapping := `{"alias-model":"pricing-usage-model"}`
+	channel := &Channel{
+		Id:           910,
+		Type:         constant.ChannelTypeTaskPlugin,
+		Key:          "key-910",
+		Status:       1,
+		Name:         "channel-910",
+		Models:       "alias-model,pricing-usage-model",
+		ModelMapping: &mapping,
+	}
+	require.NoError(t, DB.Create(channel).Error)
+	insertPricingEndpointAbility(t, 910, "alias-model")
+	insertPricingEndpointAbility(t, 910, "pricing-usage-model")
+	InitChannelCache()
+
+	saved := map[string]string{}
+	require.NoError(t, config.GlobalConfig.SaveToDB(func(key, value string) error {
+		saved[key] = value
+		return nil
+	}))
+	t.Cleanup(func() {
+		require.NoError(t, config.GlobalConfig.LoadFromDB(saved))
+	})
+	require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{
+		"billing_setting.billing_mode": `{"pricing-usage-model":"tiered_expr","alias-own-expr":"tiered_expr"}`,
+		"billing_setting.billing_expr": `{"pricing-usage-model":"u(\"seconds\")","alias-own-expr":"u(\"seconds\") * 2"}`,
+	}))
+	InvalidatePricingCache()
+
+	pricing := pricingByModel(GetPricing())
+	require.Contains(t, pricing, "alias-model")
+	require.Contains(t, pricing, "pricing-usage-model")
+	assert.Equal(t, "second", pricing["alias-model"].BillingUsageSchema["seconds"].Unit)
+	assert.Equal(t, "Estimated duration.", pricing["alias-model"].BillingUsageSchema["seconds"].Description["en"])
+	assert.Equal(t, "tiered_expr", pricing["alias-model"].BillingMode)
+	assert.Equal(t, `u("seconds")`, pricing["alias-model"].BillingExpr)
+	assert.Equal(t, "tiered_expr", pricing["pricing-usage-model"].BillingMode)
+	assert.Equal(t, `u("seconds")`, pricing["pricing-usage-model"].BillingExpr)
+
+	ownMapping := `{"alias-own-expr":"pricing-usage-model"}`
+	own := &Channel{
+		Id:           911,
+		Type:         constant.ChannelTypeTaskPlugin,
+		Key:          "key-911",
+		Status:       1,
+		Name:         "channel-911",
+		Models:       "alias-own-expr,pricing-usage-model",
+		ModelMapping: &ownMapping,
+	}
+	require.NoError(t, DB.Create(own).Error)
+	insertPricingEndpointAbility(t, 911, "alias-own-expr")
+	InitChannelCache()
+	InvalidatePricingCache()
+
+	refreshed := pricingByModel(GetPricing())
+	assert.Equal(t, `u("seconds") * 2`, refreshed["alias-own-expr"].BillingExpr)
+	assert.Equal(t, "second", refreshed["alias-own-expr"].BillingUsageSchema["seconds"].Unit)
+}
+
+func pricingByModel(pricings []Pricing) map[string]Pricing {
+	result := make(map[string]Pricing, len(pricings))
+	for _, pricing := range pricings {
+		result[pricing.ModelName] = pricing
+	}
+	return result
+}