Merge branch 'dev未优化' into dev优化中

# Conflicts: # common/util/mapping.go # config.yml # model/entity/model_gateway_model.go # model/entity/model_gateway_task.go # service/gateway/gateway_http_service.go # service/task/task_service.go # service/task/worker.go
2026-07-03 09:52:41 +08:00
parent c71823abdb 4d95ff75d2
commit cc23697201
11 changed files with 879 additions and 337 deletions
@@ -15,6 +15,7 @@ import (
 	"model-gateway/common/util"
 	"model-gateway/consts/public"
 	"model-gateway/dao"
+	"model-gateway/model/dto"
 	"model-gateway/model/entity"
 	"model-gateway/service/gateway"

@@ -30,77 +31,139 @@ type asyncWorker struct {
 }

 // handleOne 执行一次完整的任务
-func (w *asyncWorker) handleOne(ctx context.Context, task *entity.ModelGatewayTask, model *entity.ModelGatewayModel) {
+func (w *asyncWorker) handleOne(ctx context.Context, task *entity.ModelGatewayTask, model *entity.ModelGatewayModel, req *dto.CreateTaskReq) {
 	var (
-		body      = task.RequestPayload
+		body      = task.RequestPayload.Body
 		maxRetry  = model.RetryTimes
 		startTime = time.Now()
-		rawBytes  []byte
+		rawData   []byte
 		result    map[string]any
 		err       error
 	)

-	g.Log().Infof(ctx, "[任务执行] 开始 taskId=%s model=%s", task.TaskID, task.ModelName)
+	g.Log().Infof(ctx, "[handleOne] 开始 taskId=%s model=%s", task.TaskID, task.ModelName)

 	// ============================================
-	// 1) 调用模型
+	// 1) 查询余额
+	// ============================================
+	surplus, _ = gateway.GetTenantSurplus(ctx, model.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 {
-			g.Log().Infof(ctx, "[任务执行] 调用模型重试 第%d次 taskId=%s", attempt, task.TaskID)
-			time.Sleep(time.Duration(attempt) * time.Second)
+			g.Log().Infof(ctx, "[handleOne] 调模型重试 第%d次 taskId=%s", attempt, task.TaskID)
+			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(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(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(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(), "InternalServiceError") &&
+			!strings.Contains(err.Error(), "Invalid video_url") &&
+			!strings.Contains(err.Error(), "Invalid audio track") &&
+			!strings.Contains(err.Error(), "Error while downloading") &&
+			!strings.Contains(err.Error(), "Error while connecting") &&
+			!strings.Contains(err.Error(), "download failed") {
 			w.failTask(ctx, task, startTime, err.Error())
 			return
 		}

-		g.Log().Warningf(ctx, "[任务执行] 调用模型失败 taskId=%s 第%d次 err=%v", task.TaskID, attempt, err)
+		g.Log().Warningf(ctx, "[handleOne] 调模型失败 taskId=%s attempt=%d err=%v", task.TaskID, attempt, err)
 	}

 	// ============================================
-	// 2) 解析校验 + 响应映射（可重试）
+	// 3) 解析返回映射 + 存储 token 相关信息
 	// ============================================
-	result, err = w.parseAndRetry(ctx, result, task, model, maxRetry)
+	mapped, err := util.MapResponsePayload(model.ResponseMapping, result)
 	if err != nil {
 		w.failTask(ctx, task, startTime, err.Error())
 		return
 	}

+	// 计费处理
+	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)
+		for k, v := range responseData {
+			billingInput[k] = v
+		}
+		// 计算费用，替换数组第一个元素
+		billingResult := util.CalculateBilling(model.BillingConfig, billingInput)
+		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])
+	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（可重试）
+	// 4) 处理提示词相关数据解析涵盖重试
+	// ============================================
+	if req.BuildType == public.BuildTypePrompt {
+		mapped, err = w.parseAndRetry(ctx, mapped, model, task, maxRetry)
+		if err != nil {
+			task.TextResult = mapped
+			w.failTask(ctx, task, startTime, err.Error())
+			return
+		}
+	}
+
+	// ============================================
+	// 5) 上传 OSS（可重试）
 	// ============================================
 	var oss *gateway.UploadFileResponse
 	for attempt := 0; attempt <= maxRetry; attempt++ {
 		if attempt > 0 {
-			g.Log().Infof(ctx, "[任务执行] OSS上传重试 第%d/%d次 taskId=%s", attempt, maxRetry, task.TaskID)
+			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
 		}
-		g.Log().Errorf(ctx, "[任务执行] OSS上传失败 taskId=%s 第%d/%d次 err=%v", task.TaskID, attempt, maxRetry, err)
+		g.Log().Errorf(ctx, "[handleOne] OSS上传失败 taskId=%s attempt=%d/%d err=%v", task.TaskID, attempt, maxRetry, err)
 		if attempt == maxRetry {
 			w.failTask(ctx, task, startTime, fmt.Sprintf("OSS上传重试耗尽: %v", err))
 			return
@@ -108,7 +171,7 @@ func (w *asyncWorker) handleOne(ctx context.Context, task *entity.ModelGatewayTa
 	}

 	// ============================================
-	// 4) 成功收尾
+	// 6) 成功收尾
 	// ============================================
 	task.State = public.TaskStatusSuccess
 	task.DurationSeconds = int64(time.Since(startTime).Seconds())
@@ -117,18 +180,18 @@ func (w *asyncWorker) handleOne(ctx context.Context, task *entity.ModelGatewayTa
 		FileType: oss.FileFormat,
 		FileSize: int64(oss.FileSize),
 	}
-
+	task.TextResult = mapped
 	if _, err = dao.ModelGatewayTask.Update(ctx, task); err != nil {
-		g.Log().Errorf(ctx, "[任务执行] 更新数据库失败 taskId=%s err=%v", task.TaskID, err)
+		g.Log().Errorf(ctx, "[handleOne] 更新DB失败 taskId=%s err=%v", task.TaskID, err)
 		return
 	}

 	go gateway.TriggerCallback(util.AsyncCtx(ctx), task)
-	if task.EpicycleId != 0 {
-		go gateway.TriggerPromptsCallback(util.AsyncCtx(ctx), task)
+	if req.EpicycleId != 0 {
+		go gateway.TriggerPromptsCallback(util.AsyncCtx(ctx), task, req.EpicycleId)
 	}

-	g.Log().Infof(ctx, "[任务执行] 成功 taskId=%s 耗时=%ds 文件类型=%s",
+	g.Log().Infof(ctx, "[handleOne] 成功 taskId=%s duration=%ds fileType=%s",
 		task.TaskID, task.DurationSeconds, oss.FileFormat)
 }

@@ -141,12 +204,13 @@ type asyncResult struct {
 // asyncTaskChan 全局异步任务等待通道
 var asyncTaskChan = sync.Map{} // taskID → chan asyncResult

-func (w *asyncWorker) callModelAsync(ctx context.Context, model *entity.ModelGatewayModel, body map[string]any) (map[string]any, error) {
+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(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()

@@ -184,67 +248,52 @@ func NotifyAsyncResult(taskID string, result map[string]any, err error) {
 	}
 }

-// callModel 调用模型 + 提取文本结果
-func (w *asyncWorker) callModel(model *entity.ModelGatewayModel, body map[string]any) (map[string]any, error) {
-	data, err := InvokeModel(model, body)
+//// 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, model *entity.ModelGatewayModel, task *entity.ModelGatewayTask, maxRetry int) (map[string]any, error) {
+	// 获取构建模型的必填字段
+	user, err := utils.GetUserInfo(ctx)
+	if err != nil {
+		return nil, err
+	}
+	buildModel, err := dao.ModelGatewayModels.Get(ctx, &entity.ModelGatewayModel{
+		SQLBaseDO: beans.SQLBaseDO{TenantId: user.TenantId, Creator: user.UserName},
+		ModelName: task.BuildModelName,
+	})
 	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, 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) 解析 + 校验
-		var parsed map[string]any
-		switch task.BuildType {
-		case 1, 3:
-			parsed, err = util.ParseAndValidate(mapped, model.RequiredFields)
-			if err == nil {
-				return parsed, nil
-			}
-			lastErr = err
-		case 2:
-			return util.ParseStructResult(mapped, entity.ResponseBody), nil
-		default:
-			return mapped, nil
+		// 解析 + 校验（用构建模型的 RequiredFields）
+		parsed, err := util.ParseAndValidate(body, buildModel.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)

@@ -252,23 +301,57 @@ 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, lastErr)
-		rawData, callErr := InvokeModel(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 && len(task.BillingData) > 0 {
+			requestData := task.BillingData[0]
+			retryData := make(map[string]any)
+			for k, v := range requestData {
+				retryData[k] = v
+			}
+			responseData := util.ExtractResponseBilling(model.BillingConfig, mapped)
+			for k, v := range responseData {
+				retryData[k] = v
+			}
+			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},
+				BillingData:  task.BillingData,
+				ExpendTokens: task.ExpendTokens,
+			})
+		}
+
+		body = mapped
 	}

 	return body, nil
@@ -318,7 +401,12 @@ func injectErrorMessage(payload map[string]any, err error) map[string]any {

 // InvokeModel 调用模型服务，返回二进制结果
 // modelKey 用于覆盖/补充模型配置 head_msg（例如每次请求携带不同的 X-API-Key）
-func InvokeModel(model *entity.ModelGatewayModel, body map[string]any) ([]byte, error) {
+func InvokeModel(ctx context.Context, model *entity.ModelGatewayModel, body map[string]any) ([]byte, error) {
+	//surplus, _ := gateway.GetTenantSurplus(ctx, model.TenantId)
+	//if surplus <= 0 {
+	//	return nil, fmt.Errorf("租户余额不足")
+	//}
+
 	// 3）构建请求 URL 和超时
 	baseURL := strings.TrimRight(model.BaseURL, "/")
 	timeout := time.Duration(model.TimeoutSeconds) * time.Second
@@ -385,6 +473,9 @@ func InvokeModel(model *entity.ModelGatewayModel, body map[string]any) ([]byte,
 		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
 }