red-future/model-gateway
4
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

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

Browse Source 
# 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
This commit is contained in:
WangLiZhao
2026-07-03 09:52:41 +08:00
parent c71823abdb 4d95ff75d2
commit cc23697201
11 changed files with 879 additions and 337 deletions
Download Patch File Download Diff File Expand all files Collapse all files
service/task/task_service.go
+28 -18
View File
@@ -6,9 +6,6 @@ import (
"fmt"
"model-gateway/common/util"
"model-gateway/consts/public"
"model-gateway/service/gateway"
"model-gateway/service/prompt"
"sync"
"time"
"model-gateway/dao"
@@ -18,9 +15,7 @@ import (
"gitea.redpowerfuture.com/red-future/common/beans"
"gitea.redpowerfuture.com/red-future/common/utils"
"github.com/gogf/gf/v2/database/gdb"
"github.com/gogf/gf/v2/encoding/gjson"
"github.com/gogf/gf/v2/frame/g"
"github.com/gogf/gf/v2/os/grpool"
"github.com/gogf/gf/v2/util/gconv"
"github.com/google/uuid"
)
@@ -245,14 +240,17 @@ func (s *taskService) Create(ctx context.Context, req *dto.CreateTaskReq) (res *
// 3) 构建任务实体
task := &entity.ModelGatewayTask{
ModelName: model.ModelName,
TaskID: taskID,
State: public.TaskStatusRunning,
BizName: req.BizName,
CallbackURL: req.CallbackUrl,
RequestPayload: req.RequestPayload,
ModelName: model.ModelName,
TaskID: taskID,
State: public.TaskStatusRunning,
BizName: req.BizName,
CallbackURL: req.CallbackUrl,
RequestPayload: &entity.RequestPayload{
Body: req.RequestPayload,
Headers: util.ParseHeadMsgHeaders(model.HeadMsg),
},
EpicycleId: req.EpicycleId,
BuildType: req.BuildType,
BuildModelName: req.BuildModelName,
}
// 4) 插入任务记录
@@ -285,8 +283,19 @@ func (s *taskService) Create(ctx context.Context, req *dto.CreateTaskReq) (res *
ResponsePayload: gdb.Map{"taskId": taskID},
})
// 6) 异步执行任务
go AsyncWorker.handleOne(util.AsyncCtx(ctx), task, model)
// 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},
BillingData: task.BillingData,
})
}
// 7) 异步执行任务
go AsyncWorker.handleOne(util.AsyncCtx(ctx), task, model, req)
return &dto.CreateTaskRes{TaskID: taskID}, nil
}
@@ -421,7 +430,7 @@ func (s *taskService) GetBatch(ctx context.Context, req *dto.GetTaskBatchReq) (r
if t == nil {
continue
}
if t.State != 2 {
if t.State != public.BuildTypeNode {
continue
}
_ = dao.ModelGatewayTask.MarkDownloadedByID(ctx, t.Id)
@@ -437,9 +446,10 @@ func (s *taskService) GetBatch(ctx context.Context, req *dto.GetTaskBatchReq) (r
continue
}
items = append(items, dto.GetTaskBatchItem{
TaskID: t.TaskID,
State: t.State,
OssFile: t.ResultFile.OssFile,
TaskID: t.TaskID,
State: t.State,
OssFile: t.ResultFile.OssFile,
TextResult: t.TextResult,
})
}
return &dto.GetTaskBatchRes{List: items}, nil
service/task/worker.go
+177 -86
View File
@@ -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
}