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

refactor(util): 重构映射工具函数并优化异步任务轮询逻辑

Browse Source
This commit is contained in:
王利钊
2026-06-03 13:30:39 +08:00
parent 2c7838807b
commit bcfcc7ed47
6 changed files with 99 additions and 131 deletions
Download Patch File Download Diff File Expand all files Collapse all files

185
common/util/mapping.go
View File

@@ -52,21 +52,24 @@ func ValidatePromptResult(raw map[string]any, model *entity.AsynchModel) error {
return nil
}
// ReverseMap 映射 payload 到 mapping
func ReverseMap(mapping map[string]any, payload map[string]any) map[string]any {
jsonObj := gjson.New("{}")
for path, defaultValue := range mapping {
// 从 payload 取对应路径的值
val := gjson.New(payload).Get(path)
if !val.IsNil() {
// payload 有值，用它
_ = jsonObj.Set(path, val.Val())
} else if !g.IsEmpty(defaultValue) {
// payload 没值，用默认值
_ = jsonObj.Set(path, defaultValue)
// ParseHeadMsgHeaders 从 head_msg JSON 中提取请求头
// head_msg 格式示例：
//
// {
// "Authorization": "Bearer xxx",
// "Content-Type": "application/json",
// "X-Api-App-Id": "5147401364",
// "X-Api-Access-Key": "VCqRX7..."
// }
func ParseHeadMsgHeaders(headMsg map[string]any) map[string]string {
if len(headMsg) == 0 {
return nil
}
out := make(map[string]string, len(headMsg))
for k, v := range headMsg {
out[k] = gconv.String(v)
}
return jsonObj.Map()
return out
}
// MapResponsePayload 映射模型响应为标准格式
@@ -106,26 +109,6 @@ func MapResponsePayload(mapping map[string]any, result map[string]any) (map[stri
return mapped, nil
}
// ParseHeadMsgHeaders 从 head_msg JSON 中提取请求头
// head_msg 格式示例：
//
// {
// "Authorization": "Bearer xxx",
// "Content-Type": "application/json",
// "X-Api-App-Id": "5147401364",
// "X-Api-Access-Key": "VCqRX7..."
// }
func ParseHeadMsgHeaders(headMsg map[string]any) map[string]string {
if len(headMsg) == 0 {
return nil
}
out := make(map[string]string, len(headMsg))
for k, v := range headMsg {
out[k] = gconv.String(v)
}
return out
}
// GetModelBody 获取数据库中保存的模型信息
func GetModelBody(v map[string]any) map[string]any {
if v == nil {
@@ -149,32 +132,44 @@ func BodyToQuery(payload map[string]any) (url.Values, error) {
return q, nil
}
// PullTaskResult 轮询查询任务结果直到完成
func PullTaskResult(ctx context.Context, taskID string, queryConfig map[string]any) (map[string]any, error) {
// 1. 解析配置
url := gconv.String(queryConfig["url"])
// PullTaskResult 轮询查询异步任务结果直到完成
func PullTaskResult(ctx context.Context, body map[string]any, queryConfig map[string]any, headMsg map[string]any) (map[string]any, error) {
// 1) 解析配置
// 1.1 提取 taskID
taskIDPath := gconv.String(queryConfig["task_id"])
taskID := gconv.String(gjson.New(body).Get(taskIDPath).Val())
if taskID == "" {
return nil, fmt.Errorf("无法从路径 %s 提取 taskID", taskIDPath)
}
g.Log().Infof(ctx, "[PullTaskResult] taskID=%s", taskID)
// 1.2 请求地址，替换 {id}
queryUrl := gconv.String(queryConfig["url"])
queryUrl = replaceURLParams(queryUrl, map[string]any{"id": taskID})
// 1.3 请求方式
method := gconv.String(queryConfig["method"])
headers, _ := queryConfig["headers"].(map[string]any)
if method == "" {
method = "GET"
}
// 1.4 状态判断配置
statusPath := gconv.String(queryConfig["status_path"])
statusValues, _ := queryConfig["status_values"].(map[string]any)
if statusPath == "" {
statusPath = "status"
}
// 1.5 轮询间隔
interval := gconv.Int(queryConfig["interval_seconds"])
if interval <= 0 {
interval = 2
}
if method == "" {
method = "GET"
}
// 1.6 请求体
reqBodyMap := map[string]any{"task_id": taskID}
// 2. 构建参数
params := map[string]any{"id": taskID}
// 3. 替换 URL 中的 {id}
finalURL := replaceURLParams(url, params)
// 4. 构建请求体
bodyCfg, _ := queryConfig["body"].(map[string]any)
body := buildParams(bodyCfg, params)
// 5. 轮询
// 2) 轮询请求
for {
select {
case <-ctx.Done():
@@ -183,21 +178,19 @@ func PullTaskResult(ctx context.Context, taskID string, queryConfig map[string]a
}
var reqBody io.Reader
if method == "POST" && body != nil {
bs, _ := json.Marshal(body)
if method == "POST" {
bs, _ := json.Marshal(reqBodyMap)
reqBody = bytes.NewReader(bs)
}
req, err := http.NewRequestWithContext(ctx, method, finalURL, reqBody)
req, err := http.NewRequestWithContext(ctx, method, queryUrl, reqBody)
if err != nil {
return nil, fmt.Errorf("创建请求失败: %w", err)
}
for k, v := range headers {
req.Header.Set(k, gconv.String(v))
}
if req.Header.Get("Content-Type") == "" && reqBody != nil {
req.Header.Set("Content-Type", "application/json")
// 统一用 headMsg 注入请求头
for hk, hv := range ParseHeadMsgHeaders(headMsg) {
req.Header.Set(hk, hv)
}
client := &http.Client{Timeout: 30 * time.Second}
@@ -208,56 +201,54 @@ func PullTaskResult(ctx context.Context, taskID string, queryConfig map[string]a
continue
}
raw, _ := io.ReadAll(resp.Body)
_ = resp.Body.Close()
g.Log().Infof(ctx, "[PullTaskResult] taskID=%s statusCode=%d body=%s", taskID, resp.StatusCode, string(raw))
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
all, _ := io.ReadAll(resp.Body)
resp.Body.Close()
g.Log().Warningf(ctx, "[PullTaskResult] 请求异常 taskID=%s status=%d body=%s", taskID, resp.StatusCode, string(all))
time.Sleep(time.Duration(interval) * time.Second)
continue
}
var result map[string]any
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
resp.Body.Close()
g.Log().Warningf(ctx, "[PullTaskResult] 解析失败 taskID=%s err=%v", taskID, err)
time.Sleep(time.Duration(interval) * time.Second)
continue
}
resp.Body.Close()
_ = json.Unmarshal(raw, &result)
status := gconv.String(result["status"])
g.Log().Infof(ctx, "[PullTaskResult] 轮询 taskID=%s status=%s", taskID, status)
statusVal := gjson.New(result).Get(statusPath).Val()
statusStr := gconv.String(statusVal)
g.Log().Infof(ctx, "[PullTaskResult] 状态 taskID=%s status=%v", taskID, statusVal)
switch status {
case "succeeded":
return result, nil
case "failed", "expired":
return result, fmt.Errorf("任务失败: status=%s", status)
case "queued", "running":
time.Sleep(time.Duration(interval) * time.Second)
continue
default:
// 兼容没有 status 字段的情况，直接返回
if matchStatus(statusStr, statusValues["succeeded"]) {
g.Log().Infof(ctx, "[PullTaskResult] 任务成功 taskID=%s", taskID)
return result, nil
}
if matchStatus(statusStr, statusValues["failed"]) {
g.Log().Errorf(ctx, "[PullTaskResult] 任务失败 taskID=%s", taskID)
return result, fmt.Errorf("任务失败")
}
time.Sleep(time.Duration(interval) * time.Second)
}
}
// buildParams 构建请求参数，用 params 覆盖 bodyCfg 中对应 key
func buildParams(bodyCfg map[string]any, params map[string]any) map[string]any {
result := make(map[string]any, len(bodyCfg)+len(params))
for k, v := range bodyCfg {
result[k] = v
func matchStatus(actual string, expected any) bool {
switch v := expected.(type) {
case string:
return actual == v
case []any:
for _, item := range v {
if actual == gconv.String(item) {
return true
}
for k, v := range params {
result[k] = v
}
return result
}
return false
}
// replaceURLParams 替换 URL 中的 {key}
func replaceURLParams(url string, params map[string]any) string {
re := regexp.MustCompile(`\{([^}]+)\}`)
re := regexp.MustCompile(`\{([^}]+)}`)
return re.ReplaceAllStringFunc(url, func(s string) string {
key := strings.Trim(s, "{}")
if val, ok := params[key]; ok {
@@ -267,18 +258,6 @@ func replaceURLParams(url string, params map[string]any) string {
})
}
// replaceBodyParams 用 params 覆盖 body 中对应 key
func replaceBodyParams(bodyCfg map[string]any, params map[string]any) map[string]any {
result := make(map[string]any)
for k, v := range bodyCfg {
result[k] = v
}
for k, v := range params {
result[k] = v
}
return result
}
// InjectCallbackURL 将回调地址注入到请求体中
func InjectCallbackURL(ctx context.Context, payload map[string]any, callbackURL string) map[string]any {
if callbackURL == "" {

12
dao/model_dao.go
View File

@@ -60,18 +60,6 @@ func (d *modelDao) Delete(ctx context.Context, req *entity.AsynchModel) (rows in
// Get 按ID获取（带租户隔离，只查当前租户）
func (d *modelDao) Get(ctx context.Context, req *entity.AsynchModel, fields ...string) (m *entity.AsynchModel, err error) {
//r, err := gfdb.DB(ctx, public.DbNameModelGateway).Model(ctx, public.TableNameModel).
// OmitEmpty().
// Where(entity.AsynchModelCol.Id, req.Id).
// Where(entity.AsynchModelCol.Creator, req.Creator).
// Where(entity.AsynchModelCol.IsChatModel, req.IsChatModel).
// Where(entity.AsynchModelCol.ModelName, req.ModelName).
// Fields(fields).One()
//if err != nil {
// return
//}
//err = r.Struct(&m)
var whereCondition strings.Builder
var queryParams []interface{}
if !g.IsEmpty(req.Id) {

1
model/dto/model_dto.go
View File

@@ -98,6 +98,7 @@ type GetModelReq struct {
g.Meta `path:"/getModel" method:"get" tags:"模型管理" summary:"获取模型配置" dc:"根据模型ID获取配置详情"`
ID int64 `p:"id" json:"id,string" dc:"配置ID"`
Creator string `p:"creator" json:"creator" dc:"创建人"`
IsChatModel *int `p:"isChatModel" json:"isChatModel" dc:"是否为聊天模型"`
ModelName string `p:"modelName" json:"modelName" dc:"模型名称（唯一标识）"`
}

14
service/model/model_service.go
View File

@@ -91,12 +91,14 @@ func (s *modelService) Get(ctx context.Context, req *dto.GetModelReq) (*dto.GetM
if g.IsEmpty(req.ID) {
req.Creator = user.UserName
}
modelReq := new(entity.AsynchModel)
err = gconv.Struct(req, modelReq)
if err != nil {
return nil, err
}
model, err := dao.Model.Get(ctx, modelReq)
model, err := dao.Model.Get(ctx, &entity.AsynchModel{
SQLBaseDO: beans.SQLBaseDO{
Id: req.ID,
Creator: user.UserName,
},
ModelName: req.ModelName,
IsChatModel: req.IsChatModel,
})
if err != nil {
return nil, err
}

2
service/task/task_service.go
View File

@@ -179,7 +179,7 @@ func (s *taskService) QueryPendingTasks(ctx context.Context, req *dto.QueryPendi
if err != nil || model == nil || model.QueryConfig == nil {
continue
}
result, err := util.PullTaskResult(ctx, t.TaskID, model.QueryConfig)
result, err := util.PullTaskResult(ctx, nil, model.QueryConfig, model.HeadMsg)
if err != nil {
g.Log().Warningf(ctx, "[轮询] 查询失败 taskID=%s err=%v", t.TaskID, err)
continue

4
service/task/worker.go
View File

@@ -122,9 +122,7 @@ func (w *asyncWorker) handleOne(ctx context.Context, task *entity.AsynchTask, mo
w.failTask(ctx, task, err.Error())
return
}
// 拿到 task_id，启动轮询
taskID := gjson.New(body).Get(model.ResponseBody).String()
body, err = util.PullTaskResult(ctx, taskID, model.QueryConfig)
body, err = util.PullTaskResult(ctx, body, model.QueryConfig, model.HeadMsg)
if err != nil {
w.failTask(ctx, task, err.Error())
return