diff --git a/service/gateway/gateway_http_service.go b/service/gateway/gateway_http_service.go index fca6761..0ad2635 100644 --- a/service/gateway/gateway_http_service.go +++ b/service/gateway/gateway_http_service.go @@ -221,6 +221,66 @@ func GetVideoDuration(ctx context.Context, urls []string) (VideoDurationResp, er return resp, nil } +// DeductBalanceReq 扣减余额请求 +type DeductBalanceReq struct { + Id uint64 `json:"id"` + Surplus float64 `json:"surplus"` +} + +// DeductBalance 扣减租户余额 +func DeductBalance(ctx context.Context, tenantId uint64, amount float64) error { + apiURL := "admin-go/api/v1/system/tenant/edit" + headers := make(map[string]string) + if r := g.RequestFromCtx(ctx); r != nil { + for k, v := range r.Request.Header { + if len(v) > 0 { + headers[k] = v[0] + } + } + } + + body := DeductBalanceReq{ + Id: tenantId, + Surplus: amount, + } + jsonData, _ := json.Marshal(body) + + var resp struct{} + err := commonHttp.Put(ctx, apiURL, headers, &resp, jsonData) + if err != nil { + g.Log().Warningf(ctx, "[扣减余额] 失败 tenantId=%d amount=%.6f err=%v", tenantId, amount, err) + return err + } + g.Log().Infof(ctx, "[扣减余额] 成功 tenantId=%d amount=%.6f", tenantId, amount) + return nil +} + +// TenantSurplusResp 租户余额返回 +type TenantSurplusResp struct { + Surplus float64 `json:"surplus"` +} + +// GetTenantSurplus 获取租户余额 +func GetTenantSurplus(ctx context.Context, tenantId uint64) (float64, error) { + apiURL := fmt.Sprintf("admin-go/api/v1/system/tenant/getTenantDetails?tenantId=%d", tenantId) + headers := make(map[string]string) + if r := g.RequestFromCtx(ctx); r != nil { + for k, v := range r.Request.Header { + if len(v) > 0 { + headers[k] = v[0] + } + } + } + + var resp TenantSurplusResp + err := commonHttp.Get(ctx, apiURL, headers, &resp, nil) + if err != nil { + g.Log().Warningf(ctx, "[获取余额] 失败 tenantId=%d err=%v", tenantId, err) + return 0, err + } + return resp.Surplus, nil +} + //// callback 向回调地址 POST 任务结果(与查询接口 GetTaskRes 出参一致) //func (s *audioTaskService) callback(ctx context.Context, taskID, status, errMsg, callbackURL string) { // if callbackURL == "" { diff --git a/service/task/worker.go b/service/task/worker.go index 036b9cd..f1229c0 100644 --- a/service/task/worker.go +++ b/service/task/worker.go @@ -39,11 +39,23 @@ func (w *asyncWorker) handleOne(ctx context.Context, task *entity.ModelGatewayTa result map[string]any err error ) - g.Log().Infof(ctx, "[handleOne] 开始 taskId=%s model=%s", task.TaskID, task.ModelName) // ============================================ - // 1) 调用模型 + // 1) 查询余额 + // ============================================ + var surplus float64 + if len(model.BillingConfig) > 0 { + surplus, _ = gateway.GetTenantSurplus(ctx, task.TenantId) + if surplus <= 0 { + w.failTask(ctx, task, startTime, "租户余额不足") + return + } + g.Log().Infof(ctx, "[handleOne] 当前余额 tenantId=%d surplus=%.2f", task.TenantId, surplus) + } + + // ============================================ + // 2) 调用模型 // ============================================ for attempt := 0; ; attempt++ { if attempt > 0 { @@ -87,7 +99,7 @@ func (w *asyncWorker) handleOne(ctx context.Context, task *entity.ModelGatewayTa } // ============================================ - // 2) 解析返回映射 + 存储 token 相关信息 + // 3) 解析返回映射 + 存储 token 相关信息 // ============================================ mapped, err := util.MapResponsePayload(model.ResponseMapping, result) if err != nil { @@ -112,6 +124,14 @@ func (w *asyncWorker) handleOne(ctx context.Context, task *entity.ModelGatewayTa if billingResult != nil { task.BillingData[0] = billingResult } + + if billingResult != nil { + task.BillingData[0] = billingResult + totalFee := gconv.Float64(billingResult["total_fee"]) + if totalFee > 0 { + _ = gateway.DeductBalance(util.AsyncCtx(ctx), task.TenantId, -totalFee) + } + } } task.ExpendTokens = gconv.Int64(mapped[entity.TotalTokens]) @@ -121,7 +141,7 @@ func (w *asyncWorker) handleOne(ctx context.Context, task *entity.ModelGatewayTa } // ============================================ - // 3) 处理提示词相关数据解析涵盖重试 + // 4) 处理提示词相关数据解析涵盖重试 // ============================================ if req.BuildType == public.BuildTypePrompt { mapped, err = w.parseAndRetry(ctx, mapped, model, task, maxRetry) @@ -133,7 +153,7 @@ func (w *asyncWorker) handleOne(ctx context.Context, task *entity.ModelGatewayTa } // ============================================ - // 4) 上传 OSS(可重试) + // 5) 上传 OSS(可重试) // ============================================ var oss *gateway.UploadFileResponse for attempt := 0; attempt <= maxRetry; attempt++ { @@ -152,7 +172,7 @@ func (w *asyncWorker) handleOne(ctx context.Context, task *entity.ModelGatewayTa } // ============================================ - // 5) 成功收尾 + // 6) 成功收尾 // ============================================ task.State = public.TaskStatusSuccess task.DurationSeconds = int64(time.Since(startTime).Seconds()) @@ -318,7 +338,12 @@ func (w *asyncWorker) parseAndRetry(ctx context.Context, body map[string]any, mo billingResult := util.CalculateBilling(model.BillingConfig, retryData) if billingResult != nil { task.BillingData = append(task.BillingData, billingResult) + totalFee := gconv.Float64(billingResult["total_fee"]) + if totalFee > 0 { + _ = gateway.DeductBalance(util.AsyncCtx(ctx), task.TenantId, -totalFee) + } } + task.ExpendTokens += gconv.Int64(mapped[entity.TotalTokens]) _, _ = dao.ModelGatewayTask.Update(ctx, &entity.ModelGatewayTask{ SQLBaseDO: beans.SQLBaseDO{Id: task.Id},