diff --git a/common/util/billing.go b/common/util/billing.go new file mode 100644 index 0000000..5b296c3 --- /dev/null +++ b/common/util/billing.go @@ -0,0 +1,308 @@ +package util + +import ( + "context" + "fmt" + "math" + "model-gateway/service/gateway" + "strings" + + "github.com/gogf/gf/v2/encoding/gjson" + "github.com/gogf/gf/v2/util/gconv" +) + +// ======================== 计费入口 ======================== + +func CalculateBilling(config map[string]any, billingData map[string]any) map[string]any { + if len(config) == 0 { + return nil + } + switch config["type"] { + case "inference_tier": + return calculateInferenceTierBilling(config, billingData) + case "video_resolution": + return calculateVideoResolutionBilling(config, billingData) + } + return nil +} + +// ======================== 推理模型计费 ======================== + +func calculateInferenceTierBilling(config map[string]any, data map[string]any) map[string]any { + promptTokens := gconv.Int64(data["prompt_tokens"]) + completionTokens := gconv.Int64(data["completion_tokens"]) + hasAudio := gconv.Bool(data["has_audio"]) + inputK := promptTokens / 1000 + + tiers := config["pricing"].(map[string]any)["tiers"].([]any) + var matched map[string]any + for _, t := range tiers { + tier := t.(map[string]any) + if inputK >= gconv.Int64(tier["input_min"]) && inputK <= gconv.Int64(tier["input_max"]) { + matched = tier + break + } + } + if matched == nil { + return nil + } + + var inputPrice float64 + if hasAudio && matched["audio_input_price"] != nil { + inputPrice = gconv.Float64(matched["audio_input_price"]) + } else { + inputPrice = gconv.Float64(matched["input_price"]) + } + outputPrice := gconv.Float64(matched["output_price"]) + + inputCost := float64(promptTokens) * inputPrice / 1000000 + outputCost := float64(completionTokens) * outputPrice / 1000000 + + return map[string]any{ + "model_name": data["model_name"], + "has_audio": hasAudio, + "prompt_tokens": promptTokens, + "completion_tokens": completionTokens, + "total_tokens": promptTokens + completionTokens, + "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, + } +} + +// ======================== 视频模型计费 ======================== + +func calculateVideoResolutionBilling(config map[string]any, data map[string]any) map[string]any { + pricingPath := buildPricingPath(config, data) + + pricing := config["pricing"].(map[string]any) + unitPrice := gconv.Float64(pricing[pricingPath]) + + var effectiveMinToken float64 + if gconv.Bool(data["input_has_video"]) { + effectiveMinToken = matchMinToken(config, data) + } + + completionTokens := gconv.Float64(data["actual_tokens"]) + realChargeTokens := int64(math.Max(completionTokens, effectiveMinToken)) + totalFee := float64(realChargeTokens) * unitPrice / 1000000 + + if gconv.Bool(config["enable_audio_charge"]) { + totalFee += gconv.Float64(data["audio_word_cnt"]) * gconv.Float64(config["audio_unit_price"]) + } + + return map[string]any{ + "model_name": data["model_name"], + "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, + } +} + +// ======================== 数据提取 ======================== + +func ExtractRequestBilling(ctx context.Context, config map[string]any, requestPayload map[string]any) map[string]any { + dimensions := config["dimensions"].(map[string]any) + fields := dimensions["request"].(map[string]any) + data := make(map[string]any) + + for key, path := range fields { + data[key] = extractValue(requestPayload, gconv.String(path)) + } + + if defaults, ok := config["defaults"].(map[string]any); ok { + for k, v := range defaults { + if _, exists := data[k]; !exists || data[k] == nil { + data[k] = v + } + } + } + + // 处理 compute 字段 + if compute, ok := config["compute"].(map[string]any); ok { + for targetField, rule := range compute { + r := rule.(map[string]any) + dependsOn := gconv.String(r["depends_on"]) + dependsValue := gconv.Bool(r["depends_value"]) + + if gconv.Bool(data[dependsOn]) != dependsValue { + continue + } + + switch r["service"] { + case "video_duration": + urls := extractVideoUrls(requestPayload) + if len(urls) > 0 { + resp, err := gateway.GetVideoDuration(ctx, urls) + if err == nil { + data[targetField] = resp.Data.TotalDuration + } + } + } + } + } + return data +} + +func ExtractResponseBilling(config map[string]any, response map[string]any) map[string]any { + dimensions := config["dimensions"].(map[string]any) + fields := dimensions["response"].(map[string]any) + data := make(map[string]any) + + for key, path := range fields { + data[key] = gjson.New(response).Get(gconv.String(path)).Val() + } + return data +} + +// ======================== 内部辅助 ======================== + +func buildPricingPath(config map[string]any, data map[string]any) string { + pathConfig := config["pricing_path"].(map[string]any) + segments := pathConfig["segments"].([]any) + mapping := pathConfig["mapping"].(map[string]any) + + var parts []string + for _, seg := range segments { + segName := gconv.String(seg) + val := gconv.String(data[segName]) + if m, ok := mapping[segName].(map[string]any); ok { + if mapped, exists := m[val]; exists { + val = gconv.String(mapped) + } + } + parts = append(parts, val) + } + return strings.Join(parts, ".") +} + +func matchMinToken(config map[string]any, data map[string]any) float64 { + rule, ok := config["min_token_rule"].(map[string]any) + if !ok { + return 0 + } + + conditionDefs := rule["conditions"].([]any) + rows := rule["rows"].([]any) + + for _, row := range rows { + r := row.(map[string]any) + conditions := r["conditions"].([]any) + matched := true + + for i, cond := range conditions { + condStr := gconv.String(cond) + condDef := conditionDefs[i].(map[string]any) + condName := gconv.String(condDef["name"]) + + switch condDef["type"] { + case "range": + if !matchRange(condStr, gconv.Float64(data[condName])) { + matched = false + } + case "enum": + if gconv.String(data[condName]) != condStr { + matched = false + } + } + } + if matched { + return gconv.Float64(r["min_token"]) + } + } + return 0 +} + +func matchRange(rangeStr string, val float64) bool { + rangeStr = strings.Trim(rangeStr, "[]()") + parts := strings.Split(rangeStr, ",") + if len(parts) != 2 { + return false + } + return val >= gconv.Float64(strings.TrimSpace(parts[0])) && val < gconv.Float64(strings.TrimSpace(parts[1])) +} + +func extractValue(source map[string]any, path string) any { + // 数组求和：rounds.#.duration + if strings.HasPrefix(path, "rounds.#.") && !strings.Contains(path, "==") { + field := strings.TrimPrefix(path, "rounds.#.") + rounds := gjson.New(source).Get("rounds").Array() + var sum float64 + for _, r := range rounds { + if strings.Contains(field, "#") { + parts := strings.SplitN(field, ".#.", 2) + arr := gjson.New(r).Get(parts[0]).Array() + for _, item := range arr { + sum += gconv.Float64(gjson.New(item).Get(parts[1]).Val()) + } + } else { + sum += gconv.Float64(gjson.New(r).Get(field).Val()) + } + } + return sum + } + + // 条件判断：content.#.type==xxx + if strings.Count(path, ".#.") == 1 && strings.Contains(path, "==") { + parts := strings.Split(path, "==") + basePath := parts[0] + typ := parts[1] + arr := gjson.New(source).Get(strings.Split(basePath, ".#.")[0]).Array() + for _, item := range arr { + field := strings.Split(basePath, ".#.")[1] + if gconv.String(gjson.New(item).Get(field).Val()) == typ { + return true + } + } + return false + } + + // 条件判断：rounds.#.content.#.type==xxx（嵌套数组） + if strings.Contains(path, "==video_url") || strings.Contains(path, "==input_audio") { + parts := strings.Split(path, "==") + basePath := parts[0] + typ := parts[1] + segments := strings.Split(basePath, ".#.") + topArr := gjson.New(source).Get(segments[0]).Array() + for _, item := range topArr { + subArr := gjson.New(item).Get(segments[1]).Array() + for _, sub := range subArr { + if gconv.String(gjson.New(sub).Get("type").Val()) == typ { + return true + } + } + } + return false + } + + return gjson.New(source).Get(path).Val() +} + +func extractVideoUrls(body map[string]any) []string { + var urls []string + content := gjson.New(body).Get("content").Array() + for _, c := range content { + item := c.(map[string]any) + if item["type"] == "video_url" { + if v, ok := item["video_url"].(map[string]any); ok { + if url, ok := v["url"].(string); ok { + urls = append(urls, url) + } + } + } + } + return urls +} diff --git a/consts/public/table_name.go b/consts/public/table_name.go index a7f89cb..8afde4d 100644 --- a/consts/public/table_name.go +++ b/consts/public/table_name.go @@ -5,8 +5,9 @@ const ( ) const ( - TableNameModel = "model_gateway_models" // 模型表 - TableNameTask = "model_gateway_task" // 任务表 - TableNameOpLog = "model_gateway_logs_op" // 操作日志表 - TableNameStat = "model_gateway_logs_stat" // 按天统计表 + TableNameModel = "model_gateway_models" // 模型表 + TableNameTask = "model_gateway_task" // 任务表 + TableNameUserBilling = "model_user_billing" // 用户账单表 + TableNameOpLog = "model_gateway_logs_op" // 操作日志表 + TableNameStat = "model_gateway_logs_stat" // 按天统计表 ) diff --git a/model/entity/model_gateway_model.go b/model/entity/model_gateway_model.go index afdeb3e..1640aaf 100644 --- a/model/entity/model_gateway_model.go +++ b/model/entity/model_gateway_model.go @@ -30,7 +30,7 @@ type modelGatewayModelCol struct { StreamConfig string FirstFrame string LastFrame string - MaxTokens string + BillingConfig string } var ModelGatewayModelCol = modelGatewayModelCol{ @@ -61,7 +61,7 @@ var ModelGatewayModelCol = modelGatewayModelCol{ StreamConfig: "stream_config", FirstFrame: "first_frame", LastFrame: "last_frame", - MaxTokens: "max_tokens", + BillingConfig: "billing_config", } type ModelGatewayModel struct { @@ -92,7 +92,7 @@ type ModelGatewayModel struct { StreamConfig map[string]any `orm:"stream_config" json:"streamConfig"` FirstFrame string `orm:"first_frame" json:"firstFrame"` LastFrame string `orm:"last_frame" json:"lastFrame"` - MaxTokens int `orm:"max_tokens" json:"maxTokens"` + BillingConfig map[string]any `orm:"billing_config" json:"billingConfig"` } const ( diff --git a/model/entity/model_gateway_task.go b/model/entity/model_gateway_task.go index eabd320..1bd3d1d 100644 --- a/model/entity/model_gateway_task.go +++ b/model/entity/model_gateway_task.go @@ -22,6 +22,7 @@ type modelGatewayTaskCol struct { RequestPayload string EpicycleId string BuildModelName string + BillingData string } var ModelGatewayTaskCol = modelGatewayTaskCol{ @@ -42,6 +43,7 @@ var ModelGatewayTaskCol = modelGatewayTaskCol{ RequestPayload: "request_payload", EpicycleId: "epicycle_id", BuildModelName: "build_model_name", + BillingData: "billing_data", } // ModelGatewayTask 模型网关任务 @@ -63,6 +65,7 @@ type ModelGatewayTask struct { 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 3782545..24a2e9e 100644 --- a/service/gateway/gateway_http_service.go +++ b/service/gateway/gateway_http_service.go @@ -178,6 +178,51 @@ func IsSuperAdmin(ctx context.Context) (res bool, err error) { return r["isSuperAdmin"], err } +// VideoDurationResp 视频时长接口返回 +type VideoDurationResp struct { + Code int `json:"code"` + Message string `json:"message"` + Data struct { + Videos []VideoInfo `json:"videos"` + Count int `json:"count"` + TotalDuration float64 `json:"totalDuration"` // 秒 + TotalDurationStr string `json:"totalDurationStr"` + } `json:"data"` +} + +type VideoInfo struct { + Index int `json:"index"` + VideoUrl string `json:"videoUrl"` + Duration float64 `json:"duration"` + DurationStr string `json:"durationStr"` +} + +// GetVideoDuration 获取视频时长 +func GetVideoDuration(ctx context.Context, urls []string) (VideoDurationResp, error) { + apiURL := "medi/video/duration" + 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 := map[string]any{"video_urls": urls} + jsonData, _ := json.Marshal(body) + + var resp VideoDurationResp + err := commonHttp.Post(ctx, apiURL, headers, &resp, jsonData) + if err != nil { + g.Log().Warningf(ctx, "[视频时长] 获取失败 err=%v", err) + return resp, err + } + + g.Log().Infof(ctx, "[视频时长] 获取成功 count=%d totalDuration=%.2f", resp.Data.Count, resp.Data.TotalDuration) + return resp, nil +} + //// callback 向回调地址 POST 任务结果（与查询接口 GetTaskRes 出参一致） //func (s *audioTaskService) callback(ctx context.Context, taskID, status, errMsg, callbackURL string) { // if callbackURL == "" { diff --git a/service/task/task_service.go b/service/task/task_service.go index 06c095d..91336d2 100644 --- a/service/task/task_service.go +++ b/service/task/task_service.go @@ -107,7 +107,16 @@ func (s *taskService) Create(ctx context.Context, req *dto.CreateTaskReq) (res * ResponsePayload: gdb.Map{"taskId": taskID}, }) - // 6) 异步执行任务 + // 6) 模型计费 + if len(model.BillingConfig) > 0 { + task.BillingData = util.ExtractRequestBilling(ctx, model.BillingConfig, req.RequestPayload) + _, _ = dao.ModelGatewayTask.Update(ctx, &entity.ModelGatewayTask{ + SQLBaseDO: beans.SQLBaseDO{Id: task.Id}, + BillingData: task.BillingData, + }) + } + + // 7) 异步执行任务 go AsyncWorker.handleOne(util.AsyncCtx(ctx), task, model, req) return &dto.CreateTaskRes{TaskID: taskID}, nil diff --git a/service/task/worker.go b/service/task/worker.go index 4737274..6b824f6 100644 --- a/service/task/worker.go +++ b/service/task/worker.go @@ -53,12 +53,12 @@ func (w *asyncWorker) handleOne(ctx context.Context, task *entity.ModelGatewayTa } switch { - case model.CallMode != nil && *model.CallMode == public.CallModeStream: + case model.CallMode != nil && *model.CallMode == public.CallModeStream: // 流式模型 rawBytes, err = InvokeModel(ctx, model, body) if err == nil { result, err = util.ParseStreamResponse(rawBytes, model.StreamConfig) } - case model.CallMode != nil && *model.CallMode == public.CallModeAsync: + case model.CallMode != nil && *model.CallMode == public.CallModeAsync: // 异步模型 result, err = w.callModel(ctx, task, model, body) if err == nil { result, err = util.PullTaskResult(ctx, result, model.QueryConfig, model.HeadMsg) @@ -122,7 +122,27 @@ func (w *asyncWorker) handleOne(ctx context.Context, task *entity.ModelGatewayTa 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 + } + } + } + + // 更新任务表（含计费数据） if _, err = dao.ModelGatewayTask.Update(ctx, task); err != nil { g.Log().Errorf(ctx, "[handleOne] 更新DB失败 taskId=%s err=%v", task.TaskID, err) return @@ -195,7 +215,6 @@ func (w *asyncWorker) callModel(ctx context.Context, task *entity.ModelGatewayTa if err != nil { return nil, err } - contentType, _ := util.DetectFileType(data) var textResult string if utf8.Valid(data) && (strings.HasPrefix(contentType, "text/") || contentType == "application/json") {