diff --git a/service/task/task_service.go b/service/task/task_service.go index 3d83461..9b89032 100644 --- a/service/task/task_service.go +++ b/service/task/task_service.go @@ -110,6 +110,7 @@ func (s *taskService) Create(ctx context.Context, req *dto.CreateTaskReq) (res * // 6) 模型计费 if len(model.BillingConfig) > 0 { requestData := util.ExtractRequestBilling(ctx, model.BillingConfig, req.RequestPayload) + // 请求数据作为计费记录的基础字段,先存入数组 task.BillingData = append(task.BillingData, requestData) _, _ = dao.ModelGatewayTask.Update(ctx, &entity.ModelGatewayTask{ SQLBaseDO: beans.SQLBaseDO{Id: task.Id}, diff --git a/service/task/worker.go b/service/task/worker.go index 34903e1..580391c 100644 --- a/service/task/worker.go +++ b/service/task/worker.go @@ -93,14 +93,26 @@ func (w *asyncWorker) handleOne(ctx context.Context, task *entity.ModelGatewayTa w.failTask(ctx, task, startTime, err.Error()) return } + // 计费处理 - if len(model.BillingConfig) > 0 { + if len(model.BillingConfig) > 0 && len(task.BillingData) > 0 { + // 取请求阶段数据作为基础 + billingInput := make(map[string]any) + for k, v := range task.BillingData[0] { + billingInput[k] = v + } + // 补充返回数据 responseData := util.ExtractResponseBilling(model.BillingConfig, mapped) - billingResult := util.CalculateBilling(model.BillingConfig, responseData) + for k, v := range responseData { + billingInput[k] = v + } + // 计算费用,替换数组第一个元素 + billingResult := util.CalculateBilling(model.BillingConfig, billingInput) if billingResult != nil { - task.BillingData = append(task.BillingData, billingResult) + task.BillingData[0] = billingResult } } + task.ExpendTokens = gconv.Int64(mapped[entity.TotalTokens]) if _, err = dao.ModelGatewayTask.Update(ctx, task); err != nil { g.Log().Errorf(ctx, "[handleOne] 更新DB失败 taskId=%s err=%v", task.TaskID, err) @@ -279,10 +291,21 @@ func (w *asyncWorker) parseAndRetry(ctx context.Context, body map[string]any, ta continue } - // 记录重试计费 + // parseAndRetry 重试时 if len(model.BillingConfig) > 0 { + requestData := task.BillingData[0] // 请求数据从第一个元素取 + + retryData := make(map[string]any) + for k, v := range requestData { + retryData[k] = v + } + responseData := util.ExtractResponseBilling(model.BillingConfig, mapped) - billingResult := util.CalculateBilling(model.BillingConfig, responseData) + for k, v := range responseData { + retryData[k] = v + } + + billingResult := util.CalculateBilling(model.BillingConfig, retryData) if billingResult != nil { task.BillingData = append(task.BillingData, billingResult) }