From 35e955a617e7cf86748e93dba803e6959b468751 Mon Sep 17 00:00:00 2001 From: WangLiZhao <1838393649@qq.com> Date: Wed, 1 Jul 2026 17:41:01 +0800 Subject: [PATCH] =?UTF-8?q?feat(billing):=20=E5=AE=9E=E7=8E=B0=E5=A4=9A?= =?UTF-8?q?=E6=9D=A1=E8=AE=A1=E8=B4=B9=E6=95=B0=E6=8D=AE=E6=94=AF=E6=8C=81?= =?UTF-8?q?=E5=B9=B6=E4=BC=98=E5=8C=96=E6=A8=A1=E5=9E=8B=E8=AE=A1=E8=B4=B9?= =?UTF-8?q?=E9=80=BB=E8=BE=91?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- common/util/billing.go | 26 +-- model/entity/model_gateway_task.go | 34 ++-- service/gateway/gateway_http_service.go | 12 +- service/task/task_service.go | 3 +- service/task/worker.go | 222 +++++++++++------------- 5 files changed, 141 insertions(+), 156 deletions(-) diff --git a/common/util/billing.go b/common/util/billing.go index 485f5ce..1e8eb71 100644 --- a/common/util/billing.go +++ b/common/util/billing.go @@ -18,9 +18,9 @@ func CalculateBilling(config map[string]any, billingData map[string]any) map[str return nil } switch config["type"] { - case "inference_tier": + case "inference_tier": //推理模型计费 return calculateInferenceTierBilling(config, billingData) - case "video_resolution": + case "video_resolution": //视频模型计费 return calculateVideoResolutionBilling(config, billingData) } return nil @@ -58,18 +58,20 @@ func calculateInferenceTierBilling(config map[string]any, data map[string]any) m inputCost := float64(promptTokens) * inputPrice / 1000000 outputCost := float64(completionTokens) * outputPrice / 1000000 + // 推理模型 return map[string]any{ - "model_name": data["model_name"], - "has_audio": hasAudio, + "model_name": data["model_name"], + "total_tokens": promptTokens + completionTokens, + "total_fee": inputCost + outputCost, + // 明细 "prompt_tokens": promptTokens, "completion_tokens": completionTokens, - "total_tokens": promptTokens + completionTokens, + "has_audio": hasAudio, "input_tier": fmt.Sprintf("[%v, %v]", matched["input_min"], matched["input_max"]), "input_unit_price": inputPrice, "output_unit_price": outputPrice, "input_cost": inputCost, "output_cost": outputCost, - "total_fee": inputCost + outputCost, } } @@ -94,21 +96,23 @@ func calculateVideoResolutionBilling(config map[string]any, data map[string]any) totalFee += gconv.Float64(data["audio_word_cnt"]) * gconv.Float64(config["audio_unit_price"]) } + // 视频模型 return map[string]any{ - "model_name": data["model_name"], + "model_name": data["model_name"], + "total_tokens": realChargeTokens, + "total_fee": totalFee, + // 明细 + "prompt_tokens": 0, + "completion_tokens": int64(completionTokens), "is_online": data["is_online"], "video_resolution": data["video_resolution"], "input_has_video": data["input_has_video"], "output_duration_sec": data["output_duration_sec"], "aspect_ratio": data["aspect_ratio"], "resolution": data["resolution"], - "prompt_tokens": 0, - "completion_tokens": int64(completionTokens), - "total_tokens": realChargeTokens, "matched_path": pricingPath, "token_unit_price": unitPrice, "effective_min_token": effectiveMinToken, - "total_fee": totalFee, } } diff --git a/model/entity/model_gateway_task.go b/model/entity/model_gateway_task.go index 1bd3d1d..5a5eddc 100644 --- a/model/entity/model_gateway_task.go +++ b/model/entity/model_gateway_task.go @@ -49,23 +49,23 @@ var ModelGatewayTaskCol = modelGatewayTaskCol{ // ModelGatewayTask 模型网关任务 type ModelGatewayTask struct { beans.SQLBaseDO `orm:",inline"` - ModelName string `orm:"model_name" json:"modelName"` - TaskID string `orm:"task_id" json:"taskId"` - BizName string `orm:"biz_name" json:"bizName"` - CallbackURL string `orm:"callback_url" json:"callbackUrl"` - State int `orm:"state" json:"state"` - Phase int `orm:"phase" json:"phase"` - ErrorMsg string `orm:"error_msg" json:"errorMsg"` - ResultFile *ResultFile `orm:"result_file" json:"resultFile"` - TextResult map[string]any `orm:"text_result" json:"text"` - ExpendTokens int64 `orm:"expend_tokens" json:"expendTokens"` - DurationSeconds int64 `orm:"duration_seconds" json:"durationSeconds"` - RetryCount int `orm:"retry_count" json:"retryCount"` - TmpFile string `orm:"tmp_file" json:"tmpFile"` - RequestPayload *RequestPayload `orm:"request_payload" json:"requestPayload"` - EpicycleId int64 `orm:"epicycle_id" json:"epicycleId"` - BuildModelName string `orm:"build_model_name" json:"buildModelName"` - BillingData map[string]any `orm:"billing_data" json:"billingData"` + ModelName string `orm:"model_name" json:"modelName"` + TaskID string `orm:"task_id" json:"taskId"` + BizName string `orm:"biz_name" json:"bizName"` + CallbackURL string `orm:"callback_url" json:"callbackUrl"` + State int `orm:"state" json:"state"` + Phase int `orm:"phase" json:"phase"` + ErrorMsg string `orm:"error_msg" json:"errorMsg"` + ResultFile *ResultFile `orm:"result_file" json:"resultFile"` + TextResult map[string]any `orm:"text_result" json:"text"` + ExpendTokens int64 `orm:"expend_tokens" json:"expendTokens"` + DurationSeconds int64 `orm:"duration_seconds" json:"durationSeconds"` + RetryCount int `orm:"retry_count" json:"retryCount"` + TmpFile string `orm:"tmp_file" json:"tmpFile"` + RequestPayload *RequestPayload `orm:"request_payload" json:"requestPayload"` + EpicycleId int64 `orm:"epicycle_id" json:"epicycleId"` + BuildModelName string `orm:"build_model_name" json:"buildModelName"` + BillingData []map[string]any `orm:"billing_data" json:"billingData"` } // ResultFile OSS 结果文件 diff --git a/service/gateway/gateway_http_service.go b/service/gateway/gateway_http_service.go index 0600dd3..de03280 100644 --- a/service/gateway/gateway_http_service.go +++ b/service/gateway/gateway_http_service.go @@ -77,12 +77,12 @@ func UploadByTask(ctx context.Context, data []byte, fileExt string) (oss *Upload // CallbackPayload 回调请求体 type CallbackPayload struct { - TaskId string `json:"task_id"` - State int `json:"state"` - OssFile string `json:"oss_file"` - FileType string `json:"file_type"` - ErrorMsg string `json:"error_msg"` - BillingDate map[string]any `json:"billing_data"` + TaskId string `json:"task_id"` + State int `json:"state"` + OssFile string `json:"oss_file"` + FileType string `json:"file_type"` + ErrorMsg string `json:"error_msg"` + BillingDate []map[string]any `json:"billing_data"` } // TriggerCallback 任务的回调 diff --git a/service/task/task_service.go b/service/task/task_service.go index 91336d2..3d83461 100644 --- a/service/task/task_service.go +++ b/service/task/task_service.go @@ -109,7 +109,8 @@ func (s *taskService) Create(ctx context.Context, req *dto.CreateTaskReq) (res * // 6) 模型计费 if len(model.BillingConfig) > 0 { - task.BillingData = util.ExtractRequestBilling(ctx, model.BillingConfig, req.RequestPayload) + 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}, BillingData: task.BillingData, diff --git a/service/task/worker.go b/service/task/worker.go index 4568d69..322e1c3 100644 --- a/service/task/worker.go +++ b/service/task/worker.go @@ -6,18 +6,16 @@ import ( "encoding/json" "fmt" "io" - "net/http" - "strings" - "sync" - "time" - "unicode/utf8" - "model-gateway/common/util" "model-gateway/consts/public" "model-gateway/dao" "model-gateway/model/dto" "model-gateway/model/entity" "model-gateway/service/gateway" + "net/http" + "strings" + "sync" + "time" "gitea.redpowerfuture.com/red-future/common/beans" "github.com/gogf/gf/v2/encoding/gjson" @@ -36,7 +34,7 @@ func (w *asyncWorker) handleOne(ctx context.Context, task *entity.ModelGatewayTa body = task.RequestPayload.Body maxRetry = model.RetryTimes startTime = time.Now() - rawBytes []byte + rawData []byte result map[string]any err error ) @@ -52,25 +50,27 @@ func (w *asyncWorker) handleOne(ctx context.Context, task *entity.ModelGatewayTa time.Sleep(time.Duration(attempt) * 5 * time.Second) } + rawData, err = InvokeModel(ctx, model, body) switch { - case model.CallMode != nil && *model.CallMode == public.CallModeStream: // 流式模型 - rawBytes, err = InvokeModel(ctx, model, body) + case model.CallMode != nil && *model.CallMode == public.CallModeStream: if err == nil { - result, err = util.ParseStreamResponse(rawBytes, model.StreamConfig) + result, err = util.ParseStreamResponse(rawData, model.StreamConfig) } - case model.CallMode != nil && *model.CallMode == public.CallModeAsync: // 异步模型 - result, err = w.callModel(ctx, task, model, body) + case model.CallMode != nil && *model.CallMode == public.CallModeAsync: if err == nil { + result = gjson.New(string(rawData)).Map() result, err = util.PullTaskResult(ctx, result, model.QueryConfig, model.HeadMsg) } default: - result, err = w.callModel(ctx, task, model, body) + if err == nil { + result = gjson.New(string(rawData)).Map() + } } if err == nil { break } - + // 模型调用失败 if !strings.Contains(err.Error(), "Timeout") && !strings.Contains(err.Error(), "InternalServiceError") && !strings.Contains(err.Error(), "Invalid video_url") && @@ -86,24 +86,47 @@ func (w *asyncWorker) handleOne(ctx context.Context, task *entity.ModelGatewayTa } // ============================================ - // 2) 解析校验 + 响应映射（可重试） + // 2) 解析返回映射 + 存储 token 相关信息 // ============================================ - result, err = w.parseAndRetry(ctx, result, task, model, req, maxRetry, startTime) + mapped, err := util.MapResponsePayload(model.ResponseMapping, result) if err != nil { - task.TextResult = result w.failTask(ctx, task, startTime, err.Error()) return } + // 计费处理 + if len(model.BillingConfig) > 0 { + responseData := util.ExtractResponseBilling(model.BillingConfig, mapped) + billingResult := util.CalculateBilling(model.BillingConfig, responseData) + if billingResult != nil { + task.BillingData = append(task.BillingData, 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) + return + } // ============================================ - // 3) 上传 OSS（可重试） + // 3) 处理提示词相关数据解析涵盖重试 + // ============================================ + if req.BuildType == public.BuildTypePrompt { + mapped, err = w.parseAndRetry(ctx, mapped, task, model, maxRetry) + if err != nil { + w.failTask(ctx, task, startTime, err.Error()) + return + } + } + + // ============================================ + // 4) 上传 OSS（可重试） // ============================================ var oss *gateway.UploadFileResponse for attempt := 0; attempt <= maxRetry; attempt++ { if attempt > 0 { g.Log().Infof(ctx, "[handleOne] OSS上传重试 第%d/%d次 taskId=%s", attempt, maxRetry, task.TaskID) } - oss, err = gateway.UploadByTask(ctx, gjson.New(result).MustToJson(), "json") + oss, err = gateway.UploadByTask(ctx, gjson.New(mapped).MustToJson(), "json") if err == nil { break } @@ -115,7 +138,7 @@ func (w *asyncWorker) handleOne(ctx context.Context, task *entity.ModelGatewayTa } // ============================================ - // 4) 成功收尾 + // 5) 成功收尾 // ============================================ task.State = public.TaskStatusSuccess task.DurationSeconds = int64(time.Since(startTime).Seconds()) @@ -124,28 +147,7 @@ func (w *asyncWorker) handleOne(ctx context.Context, task *entity.ModelGatewayTa FileType: oss.FileFormat, FileSize: int64(oss.FileSize), } - task.TextResult = result - // 计费处理 - if len(model.BillingConfig) > 0 { - // 补充返回阶段数据 - responseData := util.ExtractResponseBilling(model.BillingConfig, result) - if task.BillingData == nil { - task.BillingData = make(map[string]any) - } - for k, v := range responseData { - task.BillingData[k] = v - } - - // 计算费用 - billingResult := util.CalculateBilling(model.BillingConfig, task.BillingData) - if billingResult != nil { - for k, v := range billingResult { - task.BillingData[k] = v - } - } - } - - // 更新任务表（含计费数据） + task.TextResult = mapped if _, err = dao.ModelGatewayTask.Update(ctx, task); err != nil { g.Log().Errorf(ctx, "[handleOne] 更新DB失败 taskId=%s err=%v", task.TaskID, err) return @@ -171,10 +173,11 @@ var asyncTaskChan = sync.Map{} // taskID → chan asyncResult func (w *asyncWorker) callModelAsync(ctx context.Context, task *entity.ModelGatewayTask, model *entity.ModelGatewayModel, body map[string]any) (map[string]any, error) { // 1. 提交异步任务 - body, err := w.callModel(ctx, task, model, body) + rawData, err := InvokeModel(ctx, model, body) if err != nil { return nil, err } + body = gjson.New(string(rawData)).Map() // 2. 拿到 task_id taskID := gjson.New(body).Get(entity.ResponseBody).String() @@ -212,81 +215,39 @@ func NotifyAsyncResult(taskID string, result map[string]any, err error) { } } -// callModel 调用模型 + 提取文本结果 -func (w *asyncWorker) callModel(ctx context.Context, task *entity.ModelGatewayTask, model *entity.ModelGatewayModel, body map[string]any) (map[string]any, error) { - data, err := InvokeModel(ctx, model, body) - if err != nil { - return nil, err - } - contentType, _ := util.DetectFileType(data) - var textResult string - if utf8.Valid(data) && (strings.HasPrefix(contentType, "text/") || contentType == "application/json") { - textResult = string(data) - } - - if textResult == "" { - return nil, fmt.Errorf("模型返回非文本内容，contentType=%s", contentType) - } - - return gjson.New(textResult).Map(), nil -} +//// callModel 调用模型 + 提取文本结果 +//func (w *asyncWorker) callModel(ctx context.Context, task *entity.ModelGatewayTask, model *entity.ModelGatewayModel, body map[string]any) (map[string]any, error) { +// data, err := InvokeModel(ctx, model, body) +// if err != nil { +// return nil, err +// } +// contentType, _ := util.DetectFileType(data) +// var textResult string +// if utf8.Valid(data) && (strings.HasPrefix(contentType, "text/") || contentType == "application/json") { +// textResult = string(data) +// } +// +// if textResult == "" { +// return nil, fmt.Errorf("模型返回非文本内容，contentType=%s", contentType) +// } +// +// return gjson.New(textResult).Map(), nil +//} // parseAndRetry 解析模型返回结果，并重试 -func (w *asyncWorker) parseAndRetry(ctx context.Context, body map[string]any, task *entity.ModelGatewayTask, model *entity.ModelGatewayModel, req *dto.CreateTaskReq, maxRetry int, startTime time.Time) (map[string]any, error) { - // 0) 如果指定了构建模型，查出校验字段 - var requiredFields []string - if task.BuildModelName != "" { - buildModel, _ := dao.ModelGatewayModels.Get(ctx, &entity.ModelGatewayModel{ - SQLBaseDO: beans.SQLBaseDO{ - TenantId: model.TenantId, - Creator: model.Creator, - }, - ModelName: req.BuildModelName, - }) - if buildModel != nil { - requiredFields = buildModel.RequiredFields - } - } - +func (w *asyncWorker) parseAndRetry(ctx context.Context, body map[string]any, task *entity.ModelGatewayTask, model *entity.ModelGatewayModel, maxRetry int) (map[string]any, error) { var lastErr error for attempt := 0; attempt <= maxRetry; attempt++ { if attempt > 0 { g.Log().Infof(ctx, "[执行任务][重试] JSON解析 第%d/%d次 taskId=%s", attempt, maxRetry, task.TaskID) } - // 1) 响应映射 - mapped, err := util.MapResponsePayload(model.ResponseMapping, body) - if err != nil { - lastErr = err - g.Log().Warningf(ctx, "[执行任务][映射失败] taskId=%s attempt=%d/%d err=%v", task.TaskID, attempt, maxRetry, err) - if attempt == maxRetry { - return nil, fmt.Errorf("响应映射重试耗尽: %w", err) - } - continue - } - // 2) 存 token - if _, ok := mapped[entity.TotalTokens]; ok { - task.ExpendTokens = gconv.Int64(mapped[entity.TotalTokens]) - _, _ = dao.ModelGatewayTask.Update(ctx, &entity.ModelGatewayTask{ - SQLBaseDO: beans.SQLBaseDO{Id: task.Id}, - ExpendTokens: task.ExpendTokens, - }) - } - - // 3) 解析 + 校验（用 buildModel 的 RequiredFields） - var parsed map[string]any - switch req.BuildType { - case public.BuildTypePrompt, public.BuildTypeNode: - parsed, err = util.ParseAndValidate(mapped, requiredFields) - if err == nil { - return parsed, nil - } - lastErr = err - case public.BuildTypeStruct: - return util.ParseStructResult(mapped, entity.ResponseBody), nil - default: - return mapped, nil + // 解析 + 校验 + parsed, err := util.ParseAndValidate(body, model.RequiredFields) + if err == nil { + return parsed, nil } + lastErr = err g.Log().Warningf(ctx, "[执行任务][解析失败] taskId=%s attempt=%d/%d err=%v", task.TaskID, attempt, maxRetry, err) @@ -294,23 +255,45 @@ func (w *asyncWorker) parseAndRetry(ctx context.Context, body map[string]any, ta return nil, fmt.Errorf("JSON解析重试耗尽: %w", lastErr) } - // 4) 拼接错误信息到请求体，重调模型 + // 重试：重新调模型 task.RetryCount++ _, _ = dao.ModelGatewayTask.Update(ctx, task) - body = injectErrorMessage(task.RequestPayload.Body, lastErr) - rawData, callErr := InvokeModel(ctx, model, body) + reqBody := injectErrorMessage(task.RequestPayload.Body, lastErr) + rawData, callErr := InvokeModel(ctx, model, reqBody) if callErr != nil { g.Log().Warningf(ctx, "[执行任务][重调模型失败] taskId=%s attempt=%d/%d err=%v", task.TaskID, attempt, maxRetry, callErr) continue } + // 响应映射 var rawResp map[string]any - if err = json.Unmarshal(rawData, &rawResp); err != nil { + if err := json.Unmarshal(rawData, &rawResp); err != nil { g.Log().Warningf(ctx, "[执行任务][Unmarshal失败] taskId=%s err=%v", task.TaskID, err) continue } - body = rawResp + mapped, mapErr := util.MapResponsePayload(model.ResponseMapping, rawResp) + if mapErr != nil { + g.Log().Warningf(ctx, "[执行任务][映射失败] taskId=%s err=%v", task.TaskID, mapErr) + continue + } + + // 记录重试计费 + if len(model.BillingConfig) > 0 { + responseData := util.ExtractResponseBilling(model.BillingConfig, mapped) + billingResult := util.CalculateBilling(model.BillingConfig, responseData) + if billingResult != nil { + task.BillingData = append(task.BillingData, billingResult) + } + task.ExpendTokens += gconv.Int64(mapped[entity.TotalTokens]) + _, _ = dao.ModelGatewayTask.Update(ctx, &entity.ModelGatewayTask{ + SQLBaseDO: beans.SQLBaseDO{Id: task.Id}, + BillingData: task.BillingData, + ExpendTokens: task.ExpendTokens, + }) + } + + body = mapped } return body, nil @@ -361,13 +344,6 @@ func injectErrorMessage(payload map[string]any, err error) map[string]any { // InvokeModel 调用模型服务，返回二进制结果 // modelKey 用于覆盖/补充模型配置 head_msg（例如每次请求携带不同的 X-API-Key） func InvokeModel(ctx context.Context, model *entity.ModelGatewayModel, body map[string]any) ([]byte, error) { - // 1) 记录模型调用次数 - //_ = dao.ModelGatewayLogsStat.IncRequestCount(ctx, time.Now(), model.TenantId, model.Creator, model.ModelName) - - // 2）请求参数映射：将标准 payload 按模型配置的 requestMapping 转为模型需要的格式 - //—— 请求映射实际处理为提示词构建请求，因为有附加字段及其他字段的拼接。这里不方便做请求映射 - //mappedPayload := util.ReverseMap(model.RequestMapping, payload) - // 3）构建请求 URL 和超时 baseURL := strings.TrimRight(model.BaseURL, "/") timeout := time.Duration(model.TimeoutSeconds) * time.Second @@ -434,6 +410,10 @@ func InvokeModel(ctx context.Context, model *entity.ModelGatewayModel, body map[ msg := string(b) return nil, fmt.Errorf("模型服务返回非2xx: %d, body=%s", resp.StatusCode, msg) } + + // + + // g.Log().Debugf(ctx, "[执行任务][模型调用成功] StatusCode=%v", resp.StatusCode) return b, nil }