chore: add model option to collection translating input

yiwen-ai · Nov 26, 2023 · 46cf161 · 46cf161
1 parent deb3d2a
commit 46cf161
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 10 deletions.
diff --git a/src/api/collection.go b/src/api/collection.go
@@ -372,6 +372,11 @@ func (a *Collection) TranslateInfo(ctx *gear.Context) error {
 		return err
 	}
 
+	model := bll.DefaultModel
+	if input.Model != nil {
+		model = bll.GetAIModel(*input.Model)
+	}
+
 	languages := strings.Join(input.Languages, ",")
 	msg, err := a.blls.Writing.GetCollectionInfo(ctx, &bll.QueryGidID{
 		ID: input.ID, GID: input.GID,
@@ -417,7 +422,7 @@ func (a *Collection) TranslateInfo(ctx *gear.Context) error {
 	}
 
 	tokens := a.blls.Jarvis.EstimateTranslatingTokens(trans, *msg.Language, input.Languages[0])
-	estimate_cost := bll.DefaultModel.CostWEN(tokens) * int64(len(input.Languages))
+	estimate_cost := model.CostWEN(tokens) * int64(len(input.Languages))
 	if b := wallet.Balance(); b < estimate_cost {
 		return gear.ErrPaymentRequired.WithMsgf("insufficient balance, expected %d, got %d", estimate_cost, b)
 	}
@@ -474,7 +479,7 @@ func (a *Collection) TranslateInfo(ctx *gear.Context) error {
 					Version:      input.Version,
 					FromLanguage: msg.Language,
 					Context:      msg.Context,
-					Model:        util.Ptr(bll.DefaultModel.ID),
+					Model:        util.Ptr(model.ID),
 					Content:      util.Ptr(util.Bytes(teData)),
 				})
 			}
@@ -509,8 +514,8 @@ func (a *Collection) TranslateInfo(ctx *gear.Context) error {
 				Action:   bll.LogActionMessageUpdate,
 				Language: languages,
 				Version:  input.Version,
-				Model:    bll.DefaultModel.ID,
-				Price:    bll.DefaultModel.Price,
+				Model:    model.ID,
+				Price:    model.Price,
 				Tokens:   usedTokens,
 			}
 
@@ -550,7 +555,7 @@ func (a *Collection) TranslateInfo(ctx *gear.Context) error {
 			log["error"] = err.Error()
 		} else {
 			auditLog.Status = 1
-			log["cost"] = bll.DefaultModel.CostWEN(*auditLog.Tokens)
+			log["cost"] = model.CostWEN(*auditLog.Tokens)
 		}
 
 		go a.blls.Logbase.Update(gctx, auditLog)

diff --git a/src/api/message.go b/src/api/message.go
@@ -106,6 +106,11 @@ func (a *Message) UpdateI18n(ctx *gear.Context) error {
 		return err
 	}
 
+	model := bll.DefaultModel
+	if input.Model != nil {
+		model = bll.GetAIModel(*input.Model)
+	}
+
 	lang := *input.Language
 	msg, err := a.blls.Writing.GetMessage(ctx, &bll.QueryID{
 		ID: input.ID, Fields: "version,language,attach_to,context,message," + lang,
@@ -177,7 +182,7 @@ func (a *Message) UpdateI18n(ctx *gear.Context) error {
 	}
 
 	tokens := a.blls.Jarvis.EstimateTranslatingTokens(trans, *msg.Language, *input.Language)
-	estimate_cost := bll.DefaultModel.CostWEN(tokens)
+	estimate_cost := model.CostWEN(tokens)
 	if b := wallet.Balance(); b < estimate_cost {
 		return gear.ErrPaymentRequired.WithMsgf("insufficient balance, expected %d, got %d", estimate_cost, b)
 	}
@@ -219,7 +224,7 @@ func (a *Message) UpdateI18n(ctx *gear.Context) error {
 			Version:      *msg.Version,
 			FromLanguage: msg.Language,
 			Context:      msg.Context,
-			Model:        util.Ptr(bll.DefaultModel.ID),
+			Model:        util.Ptr(model.ID),
 			Content:      util.Ptr(util.Bytes(teData)),
 		})
 
@@ -236,8 +241,8 @@ func (a *Message) UpdateI18n(ctx *gear.Context) error {
 				Action:   bll.LogActionMessageUpdate,
 				Language: *input.Language,
 				Version:  *msg.Version,
-				Model:    bll.DefaultModel.ID,
-				Price:    bll.DefaultModel.Price,
+				Model:    model.ID,
+				Price:    model.Price,
 				Tokens:   tmOutput.Tokens,
 			}
 
@@ -283,7 +288,7 @@ func (a *Message) UpdateI18n(ctx *gear.Context) error {
 			log["error"] = err.Error()
 		} else {
 			auditLog.Status = 1
-			log["cost"] = bll.DefaultModel.CostWEN(*auditLog.Tokens)
+			log["cost"] = model.CostWEN(*auditLog.Tokens)
 		}
 
 		go a.blls.Logbase.Update(gctx, auditLog)

diff --git a/src/bll/writing_collection.go b/src/bll/writing_collection.go
@@ -196,6 +196,7 @@ type TranslateCollectionInfoInput struct {
 	GID       util.ID  `json:"gid" cbor:"gid" validate:"required"`
 	Version   uint16   `json:"version" cbor:"version" validate:"gte=0,lte=32767"`
 	Languages []string `json:"languages" cbor:"languages" validate:"gte=1,lte=100"`
+	Model     *string  `json:"model,omitempty" cbor:"model,omitempty" validate:"omitempty,gte=2,lte=16"`
 }
 
 func (i *TranslateCollectionInfoInput) Validate() error {

diff --git a/src/bll/writing_message.go b/src/bll/writing_message.go
@@ -221,6 +221,7 @@ type UpdateMessageInput struct {
 	Version  uint16      `json:"version" cbor:"version" validate:"gte=1,lte=32767"`
 	Context  *string     `json:"context,omitempty" cbor:"context,omitempty" validate:"omitempty,gte=4,lte=1024"`
 	Language *string     `json:"language,omitempty" cbor:"language,omitempty"`
+	Model    *string     `json:"model,omitempty" cbor:"model,omitempty" validate:"omitempty,gte=2,lte=16"`
 	Message  *util.Bytes `json:"message,omitempty" cbor:"message,omitempty"`
 	NewlyAdd *bool       `json:"newly_add,omitempty" cbor:"newly_add,omitempty"` // default true
 }