19904408334
313bb692c7
更新 TenantSurplusResp 结构以匹配实际的嵌套 JSON 响应格式， 并修正 worker 中 billing 逻辑错误引用的租户ID变量。
545 lines
17 KiB
Go
545 lines
17 KiB
Go
package task
|
||
|
||
import (
|
||
"bytes"
|
||
"context"
|
||
"encoding/json"
|
||
"fmt"
|
||
"io"
|
||
"model-gateway/common/util"
|
||
"model-gateway/consts/public"
|
||
"model-gateway/dao"
|
||
"model-gateway/model/dto"
|
||
"model-gateway/model/entity"
|
||
"model-gateway/service/gateway"
|
||
"net/http"
|
||
"strings"
|
||
"sync"
|
||
"time"
|
||
|
||
"gitea.redpowerfuture.com/red-future/common/beans"
|
||
"gitea.redpowerfuture.com/red-future/common/utils"
|
||
"github.com/gogf/gf/v2/encoding/gjson"
|
||
"github.com/gogf/gf/v2/frame/g"
|
||
"github.com/gogf/gf/v2/util/gconv"
|
||
)
|
||
|
||
var AsyncWorker = &asyncWorker{}
|
||
|
||
type asyncWorker struct {
|
||
}
|
||
|
||
// handleOne 执行一次完整的任务
|
||
func (w *asyncWorker) handleOne(ctx context.Context, task *entity.ModelGatewayTask, model *entity.ModelGatewayModel, req *dto.CreateTaskReq) {
|
||
var (
|
||
body = task.RequestPayload.Body
|
||
maxRetry = model.RetryTimes
|
||
startTime = time.Now()
|
||
rawData []byte
|
||
result map[string]any
|
||
err error
|
||
surplus float64
|
||
)
|
||
g.Log().Infof(ctx, "[handleOne] 开始 taskId=%s model=%s", task.TaskID, task.ModelName)
|
||
|
||
// ============================================
|
||
// 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", model.TenantId, surplus)
|
||
|
||
// ============================================
|
||
// 2) 调用模型
|
||
// ============================================
|
||
for attempt := 0; ; attempt++ {
|
||
if attempt > 0 {
|
||
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:
|
||
if err == nil {
|
||
result, err = util.ParseStreamResponse(rawData, model.StreamConfig)
|
||
}
|
||
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:
|
||
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(), "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, "[handleOne] 调模型失败 taskId=%s attempt=%d err=%v", task.TaskID, attempt, err)
|
||
}
|
||
|
||
// ============================================
|
||
// 3) 解析返回映射 + 存储 token 相关信息
|
||
// ============================================
|
||
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), model.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
|
||
}
|
||
|
||
// ============================================
|
||
// 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, "[handleOne] OSS上传重试 第%d/%d次 taskId=%s", attempt, maxRetry, task.TaskID)
|
||
}
|
||
oss, err = gateway.UploadByTask(ctx, gjson.New(mapped).MustToJson(), "json")
|
||
if err == nil {
|
||
break
|
||
}
|
||
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
|
||
}
|
||
}
|
||
|
||
// ============================================
|
||
// 6) 成功收尾
|
||
// ============================================
|
||
task.State = public.TaskStatusSuccess
|
||
task.DurationSeconds = int64(time.Since(startTime).Seconds())
|
||
task.ResultFile = &entity.ResultFile{
|
||
OssFile: oss.FileAddressPrefix + oss.FileURL,
|
||
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)
|
||
}
|
||
|
||
g.Log().Infof(ctx, "[handleOne] 成功 taskId=%s duration=%ds fileType=%s",
|
||
task.TaskID, task.DurationSeconds, oss.FileFormat)
|
||
}
|
||
|
||
// asyncResult 异步任务结果
|
||
type asyncResult struct {
|
||
result map[string]any
|
||
err error
|
||
}
|
||
|
||
// asyncTaskChan 全局异步任务等待通道
|
||
var asyncTaskChan = sync.Map{} // taskID → chan asyncResult
|
||
|
||
func (w *asyncWorker) callModelAsync(ctx context.Context, task *entity.ModelGatewayTask, model *entity.ModelGatewayModel, body map[string]any) (map[string]any, error) {
|
||
// 1. 提交异步任务
|
||
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()
|
||
|
||
// 3. 创建等待通道
|
||
ch := make(chan asyncResult, 1)
|
||
asyncTaskChan.Store(taskID, ch)
|
||
defer func() {
|
||
asyncTaskChan.Delete(taskID)
|
||
close(ch)
|
||
}()
|
||
|
||
// 4. 阻塞等待回调或超时
|
||
timeout := time.Duration(model.TimeoutSeconds) * time.Second
|
||
ctx, cancel := context.WithTimeout(ctx, timeout)
|
||
defer cancel()
|
||
|
||
g.Log().Infof(ctx, "[异步任务] 开始等待结果 taskID=%s timeout=%v", taskID, timeout)
|
||
|
||
select {
|
||
case res, ok := <-ch:
|
||
if !ok {
|
||
return nil, fmt.Errorf("异步任务通道已关闭: taskID=%s", taskID)
|
||
}
|
||
g.Log().Infof(ctx, "[异步任务] 获取结果成功 taskID=%s", taskID)
|
||
return res.result, res.err
|
||
case <-ctx.Done():
|
||
return nil, fmt.Errorf("异步任务超时: taskID=%s", taskID)
|
||
}
|
||
}
|
||
|
||
// NotifyAsyncResult 回调接口调用此方法通知结果
|
||
func NotifyAsyncResult(taskID string, result map[string]any, err error) {
|
||
if ch, ok := asyncTaskChan.Load(taskID); ok {
|
||
ch.(chan asyncResult) <- asyncResult{result: result, err: err}
|
||
}
|
||
}
|
||
|
||
//// 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
|
||
}
|
||
|
||
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)
|
||
}
|
||
|
||
// 解析 + 校验（用构建模型的 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)
|
||
|
||
if attempt == maxRetry {
|
||
return nil, fmt.Errorf("JSON解析重试耗尽: %w", lastErr)
|
||
}
|
||
|
||
// 重试：重新调模型
|
||
task.RetryCount++
|
||
_, _ = dao.ModelGatewayTask.Update(ctx, task)
|
||
|
||
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 {
|
||
g.Log().Warningf(ctx, "[执行任务][Unmarshal失败] taskId=%s err=%v", task.TaskID, err)
|
||
continue
|
||
}
|
||
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
|
||
}
|
||
|
||
// injectErrorMessage 将错误信息插入到最后一个 user 消息之前
|
||
func injectErrorMessage(payload map[string]any, err error) map[string]any {
|
||
if err == nil {
|
||
return payload
|
||
}
|
||
|
||
messages, _ := payload["messages"].([]any)
|
||
if len(messages) == 0 {
|
||
return payload
|
||
}
|
||
|
||
errMsg := fmt.Sprintf("【上一轮输出错误，请修正】%s", err.Error())
|
||
|
||
// 找到最后一个 user 的位置
|
||
lastUserIdx := -1
|
||
for i := len(messages) - 1; i >= 0; i-- {
|
||
msg, ok := messages[i].(map[string]any)
|
||
if !ok {
|
||
continue
|
||
}
|
||
if gconv.String(msg["role"]) == "user" {
|
||
lastUserIdx = i
|
||
break
|
||
}
|
||
}
|
||
|
||
if lastUserIdx == -1 {
|
||
return payload
|
||
}
|
||
|
||
// 在最后一个 user 之前插入错误消息
|
||
errMsgObj := map[string]any{
|
||
"role": "user",
|
||
"content": []map[string]any{{"type": "text", "text": errMsg}},
|
||
}
|
||
|
||
// 切片插入
|
||
messages = append(messages[:lastUserIdx], append([]any{errMsgObj}, messages[lastUserIdx:]...)...)
|
||
payload["messages"] = messages
|
||
return payload
|
||
}
|
||
|
||
// InvokeModel 调用模型服务，返回二进制结果
|
||
// modelKey 用于覆盖/补充模型配置 head_msg（例如每次请求携带不同的 X-API-Key）
|
||
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
|
||
client := &http.Client{Timeout: timeout}
|
||
method := strings.ToUpper(strings.TrimSpace(model.HttpMethod))
|
||
|
||
// 4）构建 HTTP 请求
|
||
var req *http.Request
|
||
switch method {
|
||
case http.MethodGet:
|
||
q, err := util.BodyToQuery(body)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if len(q) > 0 {
|
||
if strings.Contains(baseURL, "?") {
|
||
baseURL = baseURL + "&" + q.Encode()
|
||
} else {
|
||
baseURL = baseURL + "?" + q.Encode()
|
||
}
|
||
}
|
||
// 改用独立超时ctx，隔绝外层截止
|
||
reqCtx, reqCancel := context.WithTimeout(context.Background(), timeout)
|
||
defer reqCancel()
|
||
req, err = http.NewRequestWithContext(reqCtx, http.MethodGet, baseURL, nil)
|
||
//req, err = http.NewRequestWithContext(ctx, http.MethodGet, baseURL, nil)
|
||
default:
|
||
bodyBytes, err := json.Marshal(body)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
reqCtx, reqCancel := context.WithTimeout(context.Background(), timeout)
|
||
defer reqCancel()
|
||
req, err = http.NewRequestWithContext(reqCtx, http.MethodPost, baseURL, bytes.NewReader(bodyBytes))
|
||
//req, err = http.NewRequestWithContext(ctx, http.MethodPost, baseURL, bytes.NewReader(bodyBytes))
|
||
}
|
||
|
||
// 5）注入请求头：先模型静态配置，再动态 modelKey（后者可覆盖前者）
|
||
for hk, hv := range util.ParseHeadMsgHeaders(model.HeadMsg) {
|
||
req.Header.Set(hk, hv)
|
||
}
|
||
if model.ApiKey != "" {
|
||
req.Header.Set("Authorization", "Bearer "+model.ApiKey)
|
||
}
|
||
if method != http.MethodGet {
|
||
req.Header.Set("Content-Type", "application/json")
|
||
}
|
||
|
||
// 6）发送请求
|
||
resp, err := client.Do(req)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer resp.Body.Close()
|
||
|
||
// 7）读取响应体
|
||
b, err := io.ReadAll(resp.Body)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
// 8）检查 HTTP 状态码
|
||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||
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
|
||
}
|
||
|
||
// // InvokeModel 调用模型服务，返回二进制结果
|
||
//
|
||
// func InvokeModel(ctx context.Context, m *entity.AsynchModel, payload any, modelKey string) ([]byte, error) {
|
||
// if m == nil || m.BaseURL == "" {
|
||
// return nil, fmt.Errorf("模型配置不完整")
|
||
// }
|
||
// // 请求参数映射
|
||
// mappedPayload, err := mapRequestPayload(m.RequestMapping, payload)
|
||
// if err != nil {
|
||
// return nil, fmt.Errorf("请求参数映射失败: %w", err)
|
||
// }
|
||
// // 合并请求头
|
||
// headers := util.ForwardHeaders(ctx)
|
||
// for hk, hv := range parseHeadMsgHeaders(m.HeadMsg) {
|
||
// headers[hk] = hv
|
||
// }
|
||
// for hk, hv := range parseHeadMsgHeaders(modelKey) {
|
||
// headers[hk] = hv
|
||
// }
|
||
//
|
||
// // 设置超时
|
||
// timeout := time.Duration(m.TimeoutSeconds) * time.Second
|
||
// if timeout <= 0 {
|
||
// timeout = 600 * time.Second
|
||
// }
|
||
// ctx, cancel := context.WithTimeout(ctx, timeout)
|
||
// defer cancel()
|
||
//
|
||
// invokeUrl := strings.TrimRight(m.BaseURL, "/")
|
||
// method := strings.ToUpper(strings.TrimSpace(m.HttpMethod))
|
||
// if method == "" {
|
||
// method = http.MethodPost
|
||
// }
|
||
//
|
||
// var respBytes []byte
|
||
//
|
||
// switch method {
|
||
// case http.MethodGet:
|
||
// err = commonHttp.Get(ctx, invokeUrl, headers, &respBytes, mappedPayload)
|
||
// default:
|
||
// err = commonHttp.Post(ctx, invokeUrl, headers, &respBytes, mappedPayload)
|
||
// }
|
||
// if err != nil {
|
||
// return nil, err
|
||
// }
|
||
// // 响应参数映射
|
||
// mappedResponse, err := mapResponsePayload(m.ResponseMapping, respBytes)
|
||
// if err != nil {
|
||
// g.Log().Warningf(ctx, "响应参数映射失败: %v，返回原始数据", err)
|
||
// return respBytes, nil
|
||
// }
|
||
// return mappedResponse, nil
|
||
// }
|
||
|
||
// failTask 任务失败统一处理
|
||
func (w *asyncWorker) failTask(ctx context.Context, t *entity.ModelGatewayTask, startTime time.Time, errMsg string) {
|
||
t.State = 3
|
||
t.ErrorMsg = errMsg
|
||
t.DurationSeconds = int64(time.Since(startTime).Seconds())
|
||
_, _ = dao.ModelGatewayTask.Update(ctx, t) // 更新任务状态
|
||
go gateway.TriggerCallback(util.AsyncCtx(ctx), t) // 触发回调
|
||
}
|