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

refactor(task): 重构任务处理逻辑并移除文件类型检测功能

Browse Source
This commit is contained in:
WangLiZhao
2026-07-03 15:55:43 +08:00
parent cc23697201
commit 071ebe2adc
9 changed files with 81 additions and 258 deletions
Download Patch File Download Diff File Expand all files Collapse all files
common/util/files.go
-53
View File
@@ -1,64 +1,11 @@
package util
import (
"net/http"
"path/filepath"
"regexp"
"strings"
)
// DetectFileType 根据返回的二进制内容推断 contentType + 扩展名
func DetectFileType(data []byte) (contentType string, ext string) {
if len(data) == 0 {
return "application/octet-stream", ".bin"
}
ct := http.DetectContentType(data)
if idx := strings.Index(ct, ";"); idx > 0 {
ct = strings.TrimSpace(ct[:idx])
}
switch ct {
case "audio/mpeg":
return ct, ".mp3"
case "audio/wave", "audio/wav", "audio/x-wav":
return ct, ".wav"
case "audio/mp4", "audio/x-m4a":
return ct, ".m4a"
case "video/mp4":
return ct, ".mp4"
case "video/webm":
return ct, ".webm"
case "image/png":
return ct, ".png"
case "image/jpeg":
return ct, ".jpg"
case "image/gif":
return ct, ".gif"
case "image/webp":
return ct, ".webp"
case "application/pdf":
return ct, ".pdf"
case "text/plain":
return ct, ".txt"
case "application/json":
return ct, ".json"
case "application/zip":
return ct, ".zip"
case "application/octet-stream":
return ct, ".bin"
default:
if parts := strings.Split(ct, "/"); len(parts) == 2 {
sub := parts[1]
if idx := strings.Index(sub, ";"); idx > 0 {
sub = strings.TrimSpace(sub[:idx])
}
return ct, "." + sub
}
return ct, ".bin"
}
}
var (
// AllowedMIMEPrefixes 允许的文本类 MIME 类型前缀
AllowedMIMEPrefixes = []string{
common/util/headers.go
-50
View File
@@ -29,56 +29,6 @@ func AsyncCtx(ctx context.Context) context.Context {
return asyncCtx
}
// ForwardHeaders 透传调用链路的头信息，优先使用 ctx 中的固化值
func ForwardHeaders(ctx context.Context) map[string]string {
headers := make(map[string]string)
SetHeaderFromContext(headers, ctx, "Authorization", "token")
SetHeaderFromContext(headers, ctx, "X-User-Info", "xUserInfo")
FallbackToRequestHeaders(headers, ctx)
return headers
}
// SetHeaderFromContext 从上下文中设置 header
func SetHeaderFromContext(headers map[string]string, ctx context.Context, headerKey, ctxKey string) {
if value, ok := ctx.Value(ctxKey).(string); ok && value != "" {
headers[headerKey] = value
}
}
// FallbackToRequestHeaders 从请求头中获取作为兜底
func FallbackToRequestHeaders(headers map[string]string, ctx context.Context) {
r := g.RequestFromCtx(ctx)
if r == nil {
return
}
if headers["Authorization"] == "" {
if token := r.Header.Get("Authorization"); token != "" {
headers["Authorization"] = token
}
}
if headers["X-User-Info"] == "" {
if userInfo := r.Header.Get("X-User-Info"); userInfo != "" {
headers["X-User-Info"] = userInfo
}
}
}
// SetTaskHeadersToCtx 把任务入库时保存的 header 信息注入 ctx，给 worker 调 OSS 用
func SetTaskHeadersToCtx(ctx context.Context, headers map[string]string) context.Context {
if headers == nil {
return ctx
}
if v := gconv.String(headers["Authorization"]); v != "" {
ctx = context.WithValue(ctx, "token", v)
}
if v := gconv.String(headers["X-User-Info"]); v != "" {
ctx = context.WithValue(ctx, "xUserInfo", v)
}
return ctx
}
// ======================== 请求工具 ========================
// ParseHeadMsgHeaders 从 head_msg 中提取 HTTP 请求头
common/util/mapping.go
-74
View File
@@ -1,20 +1,12 @@
package util
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"model-gateway/model/entity"
"net/http"
"net/url"
"regexp"
"strings"
"time"
"github.com/gogf/gf/v2/encoding/gjson"
"github.com/gogf/gf/v2/frame/g"
"github.com/gogf/gf/v2/util/gconv"
tgjson "github.com/tidwall/gjson"
)
@@ -49,56 +41,6 @@ func ParseAndValidate(raw map[string]any, requiredFields []string) (map[string]a
return map[string]any{"total_rounds": len(arr), "rounds": arr}, nil
}
// ParseStructResult 解析结构结果
func ParseStructResult(raw map[string]any, responseBody string) map[string]any {
contentStr := gconv.String(raw[responseBody])
if contentStr == "" || contentStr == "0" {
return map[string]any{
"total_rounds": 1,
"rounds": []map[string]any{{responseBody: raw}},
}
}
if arr := tryParseArray(contentStr); arr != nil {
return map[string]any{
"total_rounds": 1,
"rounds": []map[string]any{{responseBody: arr}},
}
}
if parsed := tryParseAny(contentStr); parsed != nil {
return map[string]any{
"total_rounds": 1,
"rounds": []map[string]any{{responseBody: parsed}},
}
}
return map[string]any{
"total_rounds": 1,
"rounds": []map[string]any{{responseBody: contentStr}},
}
}
// 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
}
// MapResponsePayload 映射模型响应为标准格式
func MapResponsePayload(mapping map[string]any, result map[string]any) (map[string]any, error) {
if len(mapping) == 0 {
@@ -146,19 +88,3 @@ func cleanControlChars(s string) string {
return r
}, s)
}
func tryParseArray(s string) []any {
var arr []any
if err := json.Unmarshal([]byte(s), &arr); err == nil && len(arr) > 0 {
return arr
}
return nil
}
func tryParseAny(s string) any {
var parsed any
if err := json.Unmarshal([]byte(s), &parsed); err == nil {
return parsed
}
return nil
}
common/util/pull_task.go
+1
View File
@@ -16,6 +16,7 @@ import (
"github.com/gogf/gf/v2/util/gconv"
)
// PullTaskResult 拉取任务结果
func PullTaskResult(ctx context.Context, body map[string]any, queryConfig map[string]any, headMsg map[string]any) (map[string]any, error) {
taskID, err := extractTaskID(body, queryConfig)
if err != nil {
model/dto/model_gateway_task_dto.go
-1
View File
@@ -38,7 +38,6 @@ type CreateTaskReq struct {
CallbackUrl string `p:"callbackUrl" json:"callbackUrl" dc:"回调地址（可选，用于后续业务通知）"`
RequestPayload map[string]any `p:"requestPayload" json:"requestPayload" dc:"请求负载（透传给模型服务）"`
EpicycleId int64 `json:"epicycleId" dc:"轮次ID"`
BuildType int64 `json:"buildType" dc:"构建类型：1-提示词构建 2-节点构建"`
BuildModelName string `json:"buildModelName" json:"buildModelName" dc:"构建模型名称"`
}
model/entity/model_gateway_task.go
+9 -27
View File
@@ -11,18 +11,14 @@ type modelGatewayTaskCol struct {
BizName string
CallbackURL string
State string
Phase string
RetryCount string
ErrorMsg string
ResultFile string
TextResult string
ExpendTokens string
DurationSeconds string
RetryCount string
TmpFile string
RequestPayload string
DurationSeconds string
EpicycleId string
BuildModelName string
BillingData string
BuildModelName string
}
var ModelGatewayTaskCol = modelGatewayTaskCol{
@@ -32,18 +28,14 @@ var ModelGatewayTaskCol = modelGatewayTaskCol{
BizName: "biz_name",
CallbackURL: "callback_url",
State: "state",
Phase: "phase",
RetryCount: "retry_count",
ErrorMsg: "error_msg",
ResultFile: "result_file",
TextResult: "text_result",
ExpendTokens: "expend_tokens",
DurationSeconds: "duration_seconds",
RetryCount: "retry_count",
TmpFile: "tmp_file",
RequestPayload: "request_payload",
DurationSeconds: "duration_seconds",
EpicycleId: "epicycle_id",
BuildModelName: "build_model_name",
BillingData: "billing_data",
BuildModelName: "build_model_name",
}
// ModelGatewayTask 模型网关任务
@@ -54,18 +46,14 @@ type ModelGatewayTask struct {
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"`
RetryCount int `orm:"retry_count" json:"retryCount"`
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"`
RequestPayload map[string]any `orm:"request_payload" json:"requestPayload"`
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"`
BuildModelName string `orm:"build_model_name" json:"buildModelName"`
}
// ResultFile OSS 结果文件
@@ -74,9 +62,3 @@ type ResultFile struct {
FileType string `json:"fileType"`
FileSize int64 `json:"fileSize"`
}
// RequestPayload 请求参数结构体
type RequestPayload struct {
Headers map[string]string `json:"headers"`
Body map[string]any `json:"body"`
}
service/gateway/gateway_http_service.go
+36 -2
View File
@@ -12,6 +12,7 @@ import (
commonHttp "gitea.redpowerfuture.com/red-future/common/http"
"github.com/gogf/gf/v2/frame/g"
"github.com/gogf/gf/v2/os/gtime"
"github.com/gogf/gf/v2/util/guid"
)
@@ -147,8 +148,7 @@ func TriggerPromptsCallback(ctx context.Context, t *entity.ModelGatewayTask) {
}
var resp struct{}
payload := PromptsCallbackPayload{
EpicycleId: epicycleId,
Messages: t.TextResult,
EpicycleId: t.EpicycleId,
}
jsonData, err := json.Marshal(payload)
if err != nil {
@@ -166,6 +166,40 @@ func TriggerPromptsCallback(ctx context.Context, t *entity.ModelGatewayTask) {
g.Log().Infof(ctx, "[提示词回调] 发送成功 epicycleId=%d 回调地址=%s 消息体大小=%d字节", t.EpicycleId, callbackURL, len(jsonData))
}
// BuildCallbackPayload 构建回调请求体
type BuildCallbackPayload struct {
TaskId string `json:"taskId"`
Status int `json:"status"`
Messages any `json:"messages"`
ErrorMsg string `json:"errorMsg"`
}
// CallbackBuildResult 回调构建结果
func CallbackBuildResult(ctx context.Context, record *entity.ModelGatewayBuildRecord) {
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]
}
}
}
payload := BuildCallbackPayload{
TaskId: record.TaskID,
Status: record.Status,
Messages: record.ResultMessages,
ErrorMsg: record.ErrorMsg,
}
jsonData, _ := json.Marshal(payload)
var resp struct{}
if err := commonHttp.Post(ctx, record.CallbackURL, headers, &resp, jsonData); err != nil {
g.Log().Warningf(ctx, "[构建回调] 发送失败 taskId=%s err=%v", record.TaskID, err)
return
}
g.Log().Infof(ctx, "[构建回调] 发送成功 taskId=%s", record.TaskID)
}
// IsSuperAdmin 调用admin-go服务检查是否是超级管理员
func IsSuperAdmin(ctx context.Context) (res bool, err error) {
//headers := util.ForwardHeaders(ctx)
service/task/task_service.go
+18 -21
View File
@@ -6,6 +6,9 @@ import (
"fmt"
"model-gateway/common/util"
"model-gateway/consts/public"
"model-gateway/service/gateway"
"model-gateway/service/prompt"
"sync"
"time"
"model-gateway/dao"
@@ -15,7 +18,9 @@ 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"
)
@@ -142,7 +147,6 @@ func (s *taskService) buildResult(ctx context.Context, req *dto.BuildMessagesReq
State: public.TaskStatusRunning,
BizName: "model-gateway",
RequestPayload: reqBody,
BuildType: req.BuildType,
}
id, err := dao.ModelGatewayTask.Insert(ctx, task)
if err != nil {
@@ -150,7 +154,7 @@ func (s *taskService) buildResult(ctx context.Context, req *dto.BuildMessagesReq
}
task.Id = id
rawData, err := AsyncWorker.callModel(chatModel, reqBody)
rawData, err := InvokeModel(ctx, chatModel, reqBody)
if err != nil {
task.State = public.TaskStatusFailed
task.ErrorMsg = err.Error()
@@ -158,13 +162,10 @@ func (s *taskService) buildResult(ctx context.Context, req *dto.BuildMessagesReq
return nil, err
}
mapped, err := util.MapResponsePayload(chatModel.ResponseMapping, rawData)
mapped, err := util.MapResponsePayload(chatModel.ResponseMapping, gjson.New(string(rawData)).Map())
if err != nil {
return nil, err
}
if _, ok := mapped[entity.TotalTokens]; ok {
task.ExpendTokens = gconv.Int64(mapped[entity.TotalTokens])
}
var rounds []map[string]any
contentStr := gjson.New(mapped).Get(entity.ResponseBody).String()
@@ -240,15 +241,12 @@ 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: &entity.RequestPayload{
Body: req.RequestPayload,
Headers: util.ParseHeadMsgHeaders(model.HeadMsg),
},
ModelName: model.ModelName,
TaskID: taskID,
State: public.TaskStatusRunning,
BizName: req.BizName,
CallbackURL: req.CallbackUrl,
RequestPayload: req.RequestPayload,
EpicycleId: req.EpicycleId,
BuildModelName: req.BuildModelName,
}
@@ -295,7 +293,7 @@ func (s *taskService) Create(ctx context.Context, req *dto.CreateTaskReq) (res *
}
// 7) 异步执行任务
go AsyncWorker.handleOne(util.AsyncCtx(ctx), task, model, req)
go AsyncWorker.handleOne(util.AsyncCtx(ctx), task, model)
return &dto.CreateTaskRes{TaskID: taskID}, nil
}
@@ -430,7 +428,7 @@ func (s *taskService) GetBatch(ctx context.Context, req *dto.GetTaskBatchReq) (r
if t == nil {
continue
}
if t.State != public.BuildTypeNode {
if t.State != 2 {
continue
}
_ = dao.ModelGatewayTask.MarkDownloadedByID(ctx, t.Id)
@@ -446,10 +444,9 @@ 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,
TextResult: t.TextResult,
TaskID: t.TaskID,
State: t.State,
OssFile: t.ResultFile.OssFile,
})
}
return &dto.GetTaskBatchRes{List: items}, nil
service/task/worker.go
+17 -30
View File
@@ -6,18 +6,15 @@ import (
"encoding/json"
"fmt"
"io"
"model-gateway/common/util"
"model-gateway/consts/public"
"model-gateway/dao"
"model-gateway/model/entity"
"model-gateway/service/gateway"
"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"
"gitea.redpowerfuture.com/red-future/common/beans"
"github.com/gogf/gf/v2/encoding/gjson"
@@ -31,9 +28,9 @@ type asyncWorker struct {
}
// handleOne 执行一次完整的任务
func (w *asyncWorker) handleOne(ctx context.Context, task *entity.ModelGatewayTask, model *entity.ModelGatewayModel, req *dto.CreateTaskReq) {
func (w *asyncWorker) handleOne(ctx context.Context, task *entity.ModelGatewayTask, model *entity.ModelGatewayModel) {
var (
body = task.RequestPayload.Body
body = task.RequestPayload
maxRetry = model.RetryTimes
startTime = time.Now()
rawData []byte
@@ -46,7 +43,7 @@ func (w *asyncWorker) handleOne(ctx context.Context, task *entity.ModelGatewayTa
// ============================================
// 1) 查询余额
// ============================================
surplus, _ = gateway.GetTenantSurplus(ctx, model.TenantId)
surplus, _ := gateway.GetTenantSurplus(ctx, model.TenantId)
if surplus <= 0 {
w.failTask(ctx, task, startTime, "租户余额不足")
return
@@ -64,11 +61,11 @@ func (w *asyncWorker) handleOne(ctx context.Context, task *entity.ModelGatewayTa
rawData, err = InvokeModel(ctx, model, body)
switch {
case model.CallMode != nil && *model.CallMode == public.CallModeStream:
case model.CallMode != nil && *model.CallMode == public.CallModeStream: // 流式
if err == nil {
result, err = util.ParseStreamResponse(rawData, model.StreamConfig)
}
case model.CallMode != nil && *model.CallMode == public.CallModeAsync:
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)
@@ -133,7 +130,6 @@ func (w *asyncWorker) handleOne(ctx context.Context, task *entity.ModelGatewayTa
}
}
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
@@ -142,10 +138,9 @@ func (w *asyncWorker) handleOne(ctx context.Context, task *entity.ModelGatewayTa
// ============================================
// 4) 处理提示词相关数据解析涵盖重试
// ============================================
if req.BuildType == public.BuildTypePrompt {
if task.BizName == "prompts-core" {
mapped, err = w.parseAndRetry(ctx, mapped, model, task, maxRetry)
if err != nil {
task.TextResult = mapped
w.failTask(ctx, task, startTime, err.Error())
return
}
@@ -180,15 +175,14 @@ 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, "[handleOne] 更新DB失败 taskId=%s err=%v", task.TaskID, err)
return
}
go gateway.TriggerCallback(util.AsyncCtx(ctx), task)
if req.EpicycleId != 0 {
go gateway.TriggerPromptsCallback(util.AsyncCtx(ctx), task, req.EpicycleId)
if task.EpicycleId != 0 {
go gateway.TriggerPromptsCallback(util.AsyncCtx(ctx), task)
}
g.Log().Infof(ctx, "[handleOne] 成功 taskId=%s duration=%ds fileType=%s",
@@ -269,13 +263,8 @@ func NotifyAsyncResult(taskID string, result map[string]any, err error) {
// 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},
SQLBaseDO: beans.SQLBaseDO{TenantId: model.TenantId, Creator: model.Creator},
ModelName: task.BuildModelName,
})
if err != nil {
@@ -305,7 +294,7 @@ func (w *asyncWorker) parseAndRetry(ctx context.Context, body map[string]any, mo
task.RetryCount++
_, _ = dao.ModelGatewayTask.Update(ctx, task)
reqBody := injectErrorMessage(task.RequestPayload.Body, lastErr)
reqBody := injectErrorMessage(task.RequestPayload, 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)
@@ -343,11 +332,9 @@ func (w *asyncWorker) parseAndRetry(ctx context.Context, body map[string]any, mo
}
}
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,
SQLBaseDO: beans.SQLBaseDO{Id: task.Id},
BillingData: task.BillingData,
})
}