Fix expend cost.

yiwen-ai · Aug 27, 2023 · 7bd32ff · 7bd32ff
1 parent 279c5ca
commit 7bd32ff
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 8 deletions.
diff --git a/src/api/publication.go b/src/api/publication.go
@@ -144,6 +144,7 @@ func (a *Publication) Create(ctx *gear.Context) error {
 				CID:      src.CID,
 				Language: *input.ToLanguage,
 				Version:  src.Version,
+				Model:    input.Model,
 				Price:    bll.Pricing(input.Model),
 				Tokens:   teOutput.Tokens,
 			}

diff --git a/src/bll/walletbase.go b/src/bll/walletbase.go
@@ -28,8 +28,8 @@ func Pricing(model string) float64 {
 	return pricing[model]
 }
 
-func CostWEN(model string, tokens uint32) int64 {
-	f := pricing[model] * float64(tokens) / 1000
+func CostWEN(price float64, tokens uint32) int64 {
+	f := price * float64(tokens) / 1000
 	c := int64(f)
 	if f > float64(c) {
 		c += 1
@@ -106,7 +106,7 @@ func (b *Walletbase) Expend(ctx context.Context, uid util.ID, input *ExpendPaylo
 
 	ex := ExpendInput{
 		UID:         uid,
-		Amount:      CostWEN(input.Model, input.Tokens),
+		Amount:      CostWEN(input.Price, input.Tokens),
 		Description: "Create Publication",
 		Payload:     data,
 	}

diff --git a/src/bll/walletbase_test.go b/src/bll/walletbase_test.go
@@ -17,9 +17,11 @@ func TestModel(t *testing.T) {
 	assert.Equal(float64(10.0), Pricing("gpt4"))
 	assert.Equal(float64(0), Pricing("davinci"))
 
-	assert.Equal(int64(1), CostWEN(DefaultModel, 100))
-	assert.Equal(int64(2), CostWEN(DefaultModel, 1600))
-	assert.Equal(int64(3), CostWEN(DefaultModel, 2100))
-	assert.Equal(int64(1), CostWEN("gpt4", 100))
-	assert.Equal(int64(15), CostWEN("gpt4", 1500))
+	assert.Equal(int64(1), CostWEN(Pricing(DefaultModel), 100))
+	assert.Equal(int64(2), CostWEN(Pricing(DefaultModel), 1600))
+	assert.Equal(int64(3), CostWEN(Pricing(DefaultModel), 2100))
+	assert.Equal(int64(3), CostWEN(Pricing(DefaultModel), 2597))
+	assert.Equal(int64(1), CostWEN(Pricing("gpt4"), 100))
+	assert.Equal(int64(15), CostWEN(Pricing("gpt4"), 1500))
+	assert.Equal(int64(26), CostWEN(Pricing("gpt4"), 2597))
 }