Skip to content

phsl / new-api

Go to a project
Commit 0f09dbda by Seefs
Options

Merge branch 'upstream-main' into feature/improve-param-override 

# Conflicts:
#	relay/channel/api_request_test.go
#	relay/common/override_test.go
#	web/src/components/table/channels/modals/EditChannelModal.jsx
parents 71886f4e 8f14687d
Showing with 2163 additions and 937 deletions
README.fr.md
......@@ -30,8 +30,8 @@
</p>
<p align="center">
<a href="https://trendshift.io/repositories/8227" target="_blank">
<img src="https://trendshift.io/api/badge/repositories/8227" alt="Calcium-Ion%2Fnew-api | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/>
<a href="https://trendshift.io/repositories/20180" target="_blank">
<img src="https://trendshift.io/api/badge/repositories/20180" alt="QuantumNous%2Fnew-api | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/>
</a>
<br>
<a href="https://hellogithub.com/repository/QuantumNous/new-api" target="_blank">
......
README.ja.md
......@@ -30,8 +30,8 @@
</p>
<p align="center">
<a href="https://trendshift.io/repositories/8227" target="_blank">
<img src="https://trendshift.io/api/badge/repositories/8227" alt="Calcium-Ion%2Fnew-api | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/>
<a href="https://trendshift.io/repositories/20180" target="_blank">
<img src="https://trendshift.io/api/badge/repositories/20180" alt="QuantumNous%2Fnew-api | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/>
</a>
<br>
<a href="https://hellogithub.com/repository/QuantumNous/new-api" target="_blank">
......
README.md
......@@ -30,8 +30,8 @@
</p>
<p align="center">
<a href="https://trendshift.io/repositories/8227" target="_blank">
<img src="https://trendshift.io/api/badge/repositories/8227" alt="Calcium-Ion%2Fnew-api | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/>
<a href="https://trendshift.io/repositories/20180" target="_blank">
<img src="https://trendshift.io/api/badge/repositories/20180" alt="QuantumNous%2Fnew-api | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/>
</a>
<br>
<a href="https://hellogithub.com/repository/QuantumNous/new-api" target="_blank">
......
README.zh_CN.md
......@@ -30,8 +30,8 @@
</p>
<p align="center">
<a href="https://trendshift.io/repositories/8227" target="_blank">
<img src="https://trendshift.io/api/badge/repositories/8227" alt="Calcium-Ion%2Fnew-api | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/>
<a href="https://trendshift.io/repositories/20180" target="_blank">
<img src="https://trendshift.io/api/badge/repositories/20180" alt="QuantumNous%2Fnew-api | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/>
</a>
<br>
<a href="https://hellogithub.com/repository/QuantumNous/new-api" target="_blank">
......
README.zh_TW.md
......@@ -30,8 +30,8 @@
</p>
<p align="center">
<a href="https://trendshift.io/repositories/8227" target="_blank">
<img src="https://trendshift.io/api/badge/repositories/8227" alt="Calcium-Ion%2Fnew-api | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/>
<a href="https://trendshift.io/repositories/20180" target="_blank">
<img src="https://trendshift.io/api/badge/repositories/20180" alt="QuantumNous%2Fnew-api | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/>
</a>
<br>
<a href="https://hellogithub.com/repository/QuantumNous/new-api" target="_blank">
......
common/gin.go
......@@ -243,7 +243,15 @@ func ParseMultipartFormReusable(c *gin.Context) (*multipart.Form, error) {
return nil, err
}
contentType := c.Request.Header.Get("Content-Type")
// Use the original Content-Type saved on first call to avoid boundary
// mismatch when callers overwrite c.Request.Header after multipart rebuild.
var contentType string
if saved, ok := c.Get("_original_multipart_ct"); ok {
contentType = saved.(string)
} else {
contentType = c.Request.Header.Get("Content-Type")
c.Set("_original_multipart_ct", contentType)
}
boundary, err := parseBoundary(contentType)
if err != nil {
return nil, err
......@@ -295,7 +303,13 @@ func parseFormData(data []byte, v any) error {
}
func parseMultipartFormData(c *gin.Context, data []byte, v any) error {
contentType := c.Request.Header.Get("Content-Type")
var contentType string
if saved, ok := c.Get("_original_multipart_ct"); ok {
contentType = saved.(string)
} else {
contentType = c.Request.Header.Get("Content-Type")
c.Set("_original_multipart_ct", contentType)
}
boundary, err := parseBoundary(contentType)
if err != nil {
if errors.Is(err, errBoundaryNotFound) {
......
common/init.go
......@@ -145,6 +145,8 @@ func initConstantEnv() {
constant.ErrorLogEnabled = GetEnvOrDefaultBool("ERROR_LOG_ENABLED", false)
// 任务轮询时查询的最大数量
constant.TaskQueryLimit = GetEnvOrDefault("TASK_QUERY_LIMIT", 1000)
// 异步任务超时时间（分钟），超过此时间未完成的任务将被标记为失败并退款。0 表示禁用。
constant.TaskTimeoutMinutes = GetEnvOrDefault("TASK_TIMEOUT_MINUTES", 1440)
soraPatchStr := GetEnvOrDefaultString("TASK_PRICE_PATCH", "")
if soraPatchStr != "" {
......
constant/env.go
......@@ -16,6 +16,7 @@ var NotificationLimitDurationMinute int
var GenerateDefaultToken bool
var ErrorLogEnabled bool
var TaskQueryLimit int
var TaskTimeoutMinutes int
// temporary variable for sora patch, will be removed in future
var TaskPricePatches []string
......
controller/codex_oauth.go
......@@ -145,6 +145,7 @@ func completeCodexOAuthWithChannelID(c *gin.Context, channelID int) {
return
}
channelProxy := ""
if channelID > 0 {
ch, err := model.GetChannelById(channelID, false)
if err != nil {
......@@ -159,6 +160,7 @@ func completeCodexOAuthWithChannelID(c *gin.Context, channelID int) {
c.JSON(http.StatusOK, gin.H{"success": false, "message": "channel type is not Codex"})
return
}
channelProxy = ch.GetSetting().Proxy
}
session := sessions.Default(c)
......@@ -176,7 +178,7 @@ func completeCodexOAuthWithChannelID(c *gin.Context, channelID int) {
ctx, cancel := context.WithTimeout(c.Request.Context(), 15*time.Second)
defer cancel()
tokenRes, err := service.ExchangeCodexAuthorizationCode(ctx, code, verifier)
tokenRes, err := service.ExchangeCodexAuthorizationCodeWithProxy(ctx, code, verifier, channelProxy)
if err != nil {
common.SysError("failed to exchange codex authorization code: " + err.Error())
c.JSON(http.StatusOK, gin.H{"success": false, "message": "授权码交换失败，请重试"})
......
controller/codex_usage.go
......@@ -2,7 +2,6 @@ package controller
import (
"context"
"encoding/json"
"fmt"
"net/http"
"strconv"
......@@ -80,7 +79,7 @@ func GetCodexChannelUsage(c *gin.Context) {
refreshCtx, refreshCancel := context.WithTimeout(c.Request.Context(), 10*time.Second)
defer refreshCancel()
res, refreshErr := service.RefreshCodexOAuthToken(refreshCtx, oauthKey.RefreshToken)
res, refreshErr := service.RefreshCodexOAuthTokenWithProxy(refreshCtx, oauthKey.RefreshToken, ch.GetSetting().Proxy)
if refreshErr == nil {
oauthKey.AccessToken = res.AccessToken
oauthKey.RefreshToken = res.RefreshToken
......@@ -109,7 +108,7 @@ func GetCodexChannelUsage(c *gin.Context) {
}
var payload any
if json.Unmarshal(body, &payload) != nil {
if common.Unmarshal(body, &payload) != nil {
payload = string(body)
}
......
controller/midjourney.go
......@@ -105,13 +105,13 @@ func UpdateMidjourneyTaskBulk() {
}
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
logger.LogError(ctx, fmt.Sprintf("Get Task parse body error: %v", err))
logger.LogError(ctx, fmt.Sprintf("Get Mjp Task parse body error: %v", err))
continue
}
var responseItems []dto.MidjourneyDto
err = json.Unmarshal(responseBody, &responseItems)
if err != nil {
logger.LogError(ctx, fmt.Sprintf("Get Task parse body error2: %v, body: %s", err, string(responseBody)))
logger.LogError(ctx, fmt.Sprintf("Get Mjp Task parse body error2: %v, body: %s", err, string(responseBody)))
continue
}
resp.Body.Close()
......@@ -130,6 +130,7 @@ func UpdateMidjourneyTaskBulk() {
if !checkMjTaskNeedUpdate(task, responseItem) {
continue
}
preStatus := task.Status
task.Code = 1
task.Progress = responseItem.Progress
task.PromptEn = responseItem.PromptEn
......@@ -172,18 +173,26 @@ func UpdateMidjourneyTaskBulk() {
shouldReturnQuota = true
}
}
err = task.Update()
won, err := task.UpdateWithStatus(preStatus)
if err != nil {
logger.LogError(ctx, "UpdateMidjourneyTask task error: "+err.Error())
} else {
if shouldReturnQuota {
err = model.IncreaseUserQuota(task.UserId, task.Quota, false)
if err != nil {
logger.LogError(ctx, "fail to increase user quota: "+err.Error())
}
logContent := fmt.Sprintf("构图失败 %s，补偿 %s", task.MjId, logger.LogQuota(task.Quota))
model.RecordLog(task.UserId, model.LogTypeSystem, logContent)
} else if won && shouldReturnQuota {
err = model.IncreaseUserQuota(task.UserId, task.Quota, false)
if err != nil {
logger.LogError(ctx, "fail to increase user quota: "+err.Error())
}
model.RecordTaskBillingLog(model.RecordTaskBillingLogParams{
UserId: task.UserId,
LogType: model.LogTypeRefund,
Content: "",
ChannelId: task.ChannelId,
ModelName: service.CovertMjpActionToModelName(task.Action),
Quota: task.Quota,
Other: map[string]interface{}{
"task_id": task.MjId,
"reason": "构图失败",
},
})
}
}
}
......
controller/misc.go
......@@ -134,8 +134,10 @@ func GetStatus(c *gin.Context) {
customProviders := oauth.GetEnabledCustomProviders()
if len(customProviders) > 0 {
type CustomOAuthInfo struct {
Id int `json:"id"`
Name string `json:"name"`
Slug string `json:"slug"`
Icon string `json:"icon"`
ClientId string `json:"client_id"`
AuthorizationEndpoint string `json:"authorization_endpoint"`
Scopes string `json:"scopes"`
......@@ -144,8 +146,10 @@ func GetStatus(c *gin.Context) {
for _, p := range customProviders {
config := p.GetConfig()
providersInfo = append(providersInfo, CustomOAuthInfo{
Id: config.Id,
Name: config.Name,
Slug: config.Slug,
Icon: config.Icon,
ClientId: config.ClientId,
AuthorizationEndpoint: config.AuthorizationEndpoint,
Scopes: config.Scopes,
......
controller/oauth.go
......@@ -237,6 +237,16 @@ func findOrCreateOAuthUser(c *gin.Context, provider oauth.Provider, oauthUser *o
// Set up new user
user.Username = provider.GetProviderPrefix() + strconv.Itoa(model.GetMaxUserId()+1)
if oauthUser.Username != "" {
if exists, err := model.CheckUserExistOrDeleted(oauthUser.Username, ""); err == nil && !exists {
// 防止索引退化
if len(oauthUser.Username) <= model.UserNameMaxLength {
user.Username = oauthUser.Username
}
}
}
if oauthUser.DisplayName != "" {
user.DisplayName = oauthUser.DisplayName
} else if oauthUser.Username != "" {
......@@ -295,12 +305,12 @@ func findOrCreateOAuthUser(c *gin.Context, provider oauth.Provider, oauthUser *o
// Set the provider user ID on the user model and update
provider.SetProviderUserID(user, oauthUser.ProviderUserID)
if err := tx.Model(user).Updates(map[string]interface{}{
"github_id": user.GitHubId,
"discord_id": user.DiscordId,
"oidc_id": user.OidcId,
"linux_do_id": user.LinuxDOId,
"wechat_id": user.WeChatId,
"telegram_id": user.TelegramId,
"github_id": user.GitHubId,
"discord_id": user.DiscordId,
"oidc_id": user.OidcId,
"linux_do_id": user.LinuxDOId,
"wechat_id": user.WeChatId,
"telegram_id": user.TelegramId,
}).Error; err != nil {
return err
}
......@@ -340,6 +350,8 @@ func handleOAuthError(c *gin.Context, err error) {
} else {
common.ApiErrorI18n(c, e.MsgKey)
}
case *oauth.AccessDeniedError:
common.ApiErrorMsg(c, e.Message)
case *oauth.TrustLevelError:
common.ApiErrorI18n(c, i18n.MsgOAuthTrustLevelLow)
default:
......
controller/relay.go
......@@ -455,72 +455,147 @@ func RelayNotFound(c *gin.Context) {
})
}
func RelayTaskFetch(c *gin.Context) {
relayInfo, err := relaycommon.GenRelayInfo(c, types.RelayFormatTask, nil, nil)
if err != nil {
c.JSON(http.StatusInternalServerError, &dto.TaskError{
Code: "gen_relay_info_failed",
Message: err.Error(),
StatusCode: http.StatusInternalServerError,
})
return
}
if taskErr := relay.RelayTaskFetch(c, relayInfo.RelayMode); taskErr != nil {
respondTaskError(c, taskErr)
}
}
func RelayTask(c *gin.Context) {
retryTimes := common.RetryTimes
channelId := c.GetInt("channel_id")
c.Set("use_channel", []string{fmt.Sprintf("%d", channelId)})
relayInfo, err := relaycommon.GenRelayInfo(c, types.RelayFormatTask, nil, nil)
if err != nil {
c.JSON(http.StatusInternalServerError, &dto.TaskError{
Code: "gen_relay_info_failed",
Message: err.Error(),
StatusCode: http.StatusInternalServerError,
})
return
}
taskErr := taskRelayHandler(c, relayInfo)
if taskErr == nil {
retryTimes = 0
if taskErr := relay.ResolveOriginTask(c, relayInfo); taskErr != nil {
respondTaskError(c, taskErr)
return
}
var result *relay.TaskSubmitResult
var taskErr *dto.TaskError
defer func() {
if taskErr != nil && relayInfo.Billing != nil {
relayInfo.Billing.Refund(c)
}
}()
retryParam := &service.RetryParam{
Ctx: c,
TokenGroup: relayInfo.TokenGroup,
ModelName: relayInfo.OriginModelName,
Retry: common.GetPointer(0),
}
for ; shouldRetryTaskRelay(c, channelId, taskErr, retryTimes) && retryParam.GetRetry() < retryTimes; retryParam.IncreaseRetry() {
channel, newAPIError := getChannel(c, relayInfo, retryParam)
if newAPIError != nil {
logger.LogError(c, fmt.Sprintf("CacheGetRandomSatisfiedChannel failed: %s", newAPIError.Error()))
taskErr = service.TaskErrorWrapperLocal(newAPIError.Err, "get_channel_failed", http.StatusInternalServerError)
break
for ; retryParam.GetRetry() <= common.RetryTimes; retryParam.IncreaseRetry() {
var channel *model.Channel
if lockedCh, ok := relayInfo.LockedChannel.(*model.Channel); ok && lockedCh != nil {
channel = lockedCh
if retryParam.GetRetry() > 0 {
if setupErr := middleware.SetupContextForSelectedChannel(c, channel, relayInfo.OriginModelName); setupErr != nil {
taskErr = service.TaskErrorWrapperLocal(setupErr.Err, "setup_locked_channel_failed", http.StatusInternalServerError)
break
}
}
} else {
var channelErr *types.NewAPIError
channel, channelErr = getChannel(c, relayInfo, retryParam)
if channelErr != nil {
logger.LogError(c, channelErr.Error())
taskErr = service.TaskErrorWrapperLocal(channelErr.Err, "get_channel_failed", http.StatusInternalServerError)
break
}
}
channelId = channel.Id
useChannel := c.GetStringSlice("use_channel")
useChannel = append(useChannel, fmt.Sprintf("%d", channelId))
c.Set("use_channel", useChannel)
logger.LogInfo(c, fmt.Sprintf("using channel #%d to retry (remain times %d)", channel.Id, retryParam.GetRetry()))
//middleware.SetupContextForSelectedChannel(c, channel, originalModel)
bodyStorage, err := common.GetBodyStorage(c)
if err != nil {
if common.IsRequestBodyTooLargeError(err) || errors.Is(err, common.ErrRequestBodyTooLarge) {
taskErr = service.TaskErrorWrapperLocal(err, "read_request_body_failed", http.StatusRequestEntityTooLarge)
addUsedChannel(c, channel.Id)
bodyStorage, bodyErr := common.GetBodyStorage(c)
if bodyErr != nil {
if common.IsRequestBodyTooLargeError(bodyErr) || errors.Is(bodyErr, common.ErrRequestBodyTooLarge) {
taskErr = service.TaskErrorWrapperLocal(bodyErr, "read_request_body_failed", http.StatusRequestEntityTooLarge)
} else {
taskErr = service.TaskErrorWrapperLocal(err, "read_request_body_failed", http.StatusBadRequest)
taskErr = service.TaskErrorWrapperLocal(bodyErr, "read_request_body_failed", http.StatusBadRequest)
}
break
}
c.Request.Body = io.NopCloser(bodyStorage)
taskErr = taskRelayHandler(c, relayInfo)
result, taskErr = relay.RelayTaskSubmit(c, relayInfo)
if taskErr == nil {
break
}
if !taskErr.LocalError {
processChannelError(c,
*types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey,
common.GetContextKeyString(c, constant.ContextKeyChannelKey), channel.GetAutoBan()),
types.NewOpenAIError(taskErr.Error, types.ErrorCodeBadResponseStatusCode, taskErr.StatusCode))
}
if !shouldRetryTaskRelay(c, channel.Id, taskErr, common.RetryTimes-retryParam.GetRetry()) {
break
}
}
useChannel := c.GetStringSlice("use_channel")
if len(useChannel) > 1 {
retryLogStr := fmt.Sprintf("重试：%s", strings.Trim(strings.Join(strings.Fields(fmt.Sprint(useChannel)), "->"), "[]"))
logger.LogInfo(c, retryLogStr)
}
if taskErr != nil {
if taskErr.StatusCode == http.StatusTooManyRequests {
taskErr.Message = "当前分组上游负载已饱和，请稍后再试"
// ── 成功：结算 + 日志 + 插入任务 ──
if taskErr == nil {
if settleErr := service.SettleBilling(c, relayInfo, result.Quota); settleErr != nil {
common.SysError("settle task billing error: " + settleErr.Error())
}
c.JSON(taskErr.StatusCode, taskErr)
service.LogTaskConsumption(c, relayInfo)
task := model.InitTask(result.Platform, relayInfo)
task.PrivateData.UpstreamTaskID = result.UpstreamTaskID
task.PrivateData.BillingSource = relayInfo.BillingSource
task.PrivateData.SubscriptionId = relayInfo.SubscriptionId
task.PrivateData.TokenId = relayInfo.TokenId
task.PrivateData.BillingContext = &model.TaskBillingContext{
ModelPrice: relayInfo.PriceData.ModelPrice,
GroupRatio: relayInfo.PriceData.GroupRatioInfo.GroupRatio,
ModelRatio: relayInfo.PriceData.ModelRatio,
OtherRatios: relayInfo.PriceData.OtherRatios,
OriginModelName: relayInfo.OriginModelName,
PerCallBilling: common.StringsContains(constant.TaskPricePatches, relayInfo.OriginModelName),
}
task.Quota = result.Quota
task.Data = result.TaskData
task.Action = relayInfo.Action
if insertErr := task.Insert(); insertErr != nil {
common.SysError("insert task error: " + insertErr.Error())
}
}
if taskErr != nil {
respondTaskError(c, taskErr)
}
}
func taskRelayHandler(c *gin.Context, relayInfo *relaycommon.RelayInfo) *dto.TaskError {
var err *dto.TaskError
switch relayInfo.RelayMode {
case relayconstant.RelayModeSunoFetch, relayconstant.RelayModeSunoFetchByID, relayconstant.RelayModeVideoFetchByID:
err = relay.RelayTaskFetch(c, relayInfo.RelayMode)
default:
err = relay.RelayTaskSubmit(c, relayInfo)
// respondTaskError 统一输出 Task 错误响应（含 429 限流提示改写）
func respondTaskError(c *gin.Context, taskErr *dto.TaskError) {
if taskErr.StatusCode == http.StatusTooManyRequests {
taskErr.Message = "当前分组上游负载已饱和，请稍后再试"
}
return err
c.JSON(taskErr.StatusCode, taskErr)
}
func shouldRetryTaskRelay(c *gin.Context, channelId int, taskErr *dto.TaskError, retryTimes int) bool {
......@@ -544,7 +619,7 @@ func shouldRetryTaskRelay(c *gin.Context, channelId int, taskErr *dto.TaskError,
}
if taskErr.StatusCode/100 == 5 {
// 超时不重试
if taskErr.StatusCode == 504 || taskErr.StatusCode == 524 {
if operation_setting.IsAlwaysSkipRetryStatusCode(taskErr.StatusCode) {
return false
}
return true
......
controller/task.go
package controller
import (
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"sort"
"strconv"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/relay"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/types"
"github.com/gin-gonic/gin"
"github.com/samber/lo"
)
// UpdateTaskBulk 薄入口，实际轮询逻辑在 service 层
func UpdateTaskBulk() {
//revocer
//imageModel := "midjourney"
for {
time.Sleep(time.Duration(15) * time.Second)
common.SysLog("任务进度轮询开始")
ctx := context.TODO()
allTasks := model.GetAllUnFinishSyncTasks(constant.TaskQueryLimit)
platformTask := make(map[constant.TaskPlatform][]*model.Task)
for _, t := range allTasks {
platformTask[t.Platform] = append(platformTask[t.Platform], t)
}
for platform, tasks := range platformTask {
if len(tasks) == 0 {
continue
}
taskChannelM := make(map[int][]string)
taskM := make(map[string]*model.Task)
nullTaskIds := make([]int64, 0)
for _, task := range tasks {
if task.TaskID == "" {
// 统计失败的未完成任务
nullTaskIds = append(nullTaskIds, task.ID)
continue
}
taskM[task.TaskID] = task
taskChannelM[task.ChannelId] = append(taskChannelM[task.ChannelId], task.TaskID)
}
if len(nullTaskIds) > 0 {
err := model.TaskBulkUpdateByID(nullTaskIds, map[string]any{
"status": "FAILURE",
"progress": "100%",
})
if err != nil {
logger.LogError(ctx, fmt.Sprintf("Fix null task_id task error: %v", err))
} else {
logger.LogInfo(ctx, fmt.Sprintf("Fix null task_id task success: %v", nullTaskIds))
}
}
if len(taskChannelM) == 0 {
continue
}
UpdateTaskByPlatform(platform, taskChannelM, taskM)
}
common.SysLog("任务进度轮询完成")
}
}
func UpdateTaskByPlatform(platform constant.TaskPlatform, taskChannelM map[int][]string, taskM map[string]*model.Task) {
switch platform {
case constant.TaskPlatformMidjourney:
//_ = UpdateMidjourneyTaskAll(context.Background(), tasks)
case constant.TaskPlatformSuno:
_ = UpdateSunoTaskAll(context.Background(), taskChannelM, taskM)
default:
if err := UpdateVideoTaskAll(context.Background(), platform, taskChannelM, taskM); err != nil {
common.SysLog(fmt.Sprintf("UpdateVideoTaskAll fail: %s", err))
}
}
}
func UpdateSunoTaskAll(ctx context.Context, taskChannelM map[int][]string, taskM map[string]*model.Task) error {
for channelId, taskIds := range taskChannelM {
err := updateSunoTaskAll(ctx, channelId, taskIds, taskM)
if err != nil {
logger.LogError(ctx, fmt.Sprintf("渠道 #%d 更新异步任务失败: %s", channelId, err.Error()))
}
}
return nil
}
func updateSunoTaskAll(ctx context.Context, channelId int, taskIds []string, taskM map[string]*model.Task) error {
logger.LogInfo(ctx, fmt.Sprintf("渠道 #%d 未完成的任务有: %d", channelId, len(taskIds)))
if len(taskIds) == 0 {
return nil
}
channel, err := model.CacheGetChannel(channelId)
if err != nil {
common.SysLog(fmt.Sprintf("CacheGetChannel: %v", err))
err = model.TaskBulkUpdate(taskIds, map[string]any{
"fail_reason": fmt.Sprintf("获取渠道信息失败，请联系管理员，渠道ID：%d", channelId),
"status": "FAILURE",
"progress": "100%",
})
if err != nil {
common.SysLog(fmt.Sprintf("UpdateMidjourneyTask error2: %v", err))
}
return err
}
adaptor := relay.GetTaskAdaptor(constant.TaskPlatformSuno)
if adaptor == nil {
return errors.New("adaptor not found")
}
proxy := channel.GetSetting().Proxy
resp, err := adaptor.FetchTask(*channel.BaseURL, channel.Key, map[string]any{
"ids": taskIds,
}, proxy)
if err != nil {
common.SysLog(fmt.Sprintf("Get Task Do req error: %v", err))
return err
}
if resp.StatusCode != http.StatusOK {
logger.LogError(ctx, fmt.Sprintf("Get Task status code: %d", resp.StatusCode))
return errors.New(fmt.Sprintf("Get Task status code: %d", resp.StatusCode))
}
defer resp.Body.Close()
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
common.SysLog(fmt.Sprintf("Get Task parse body error: %v", err))
return err
}
var responseItems dto.TaskResponse[[]dto.SunoDataResponse]
err = json.Unmarshal(responseBody, &responseItems)
if err != nil {
logger.LogError(ctx, fmt.Sprintf("Get Task parse body error2: %v, body: %s", err, string(responseBody)))
return err
}
if !responseItems.IsSuccess() {
common.SysLog(fmt.Sprintf("渠道 #%d 未完成的任务有: %d, 成功获取到任务数: %s", channelId, len(taskIds), string(responseBody)))
return err
}
for _, responseItem := range responseItems.Data {
task := taskM[responseItem.TaskID]
if !checkTaskNeedUpdate(task, responseItem) {
continue
}
task.Status = lo.If(model.TaskStatus(responseItem.Status) != "", model.TaskStatus(responseItem.Status)).Else(task.Status)
task.FailReason = lo.If(responseItem.FailReason != "", responseItem.FailReason).Else(task.FailReason)
task.SubmitTime = lo.If(responseItem.SubmitTime != 0, responseItem.SubmitTime).Else(task.SubmitTime)
task.StartTime = lo.If(responseItem.StartTime != 0, responseItem.StartTime).Else(task.StartTime)
task.FinishTime = lo.If(responseItem.FinishTime != 0, responseItem.FinishTime).Else(task.FinishTime)
if responseItem.FailReason != "" || task.Status == model.TaskStatusFailure {
logger.LogInfo(ctx, task.TaskID+" 构建失败，"+task.FailReason)
task.Progress = "100%"
//err = model.CacheUpdateUserQuota(task.UserId) ?
if err != nil {
logger.LogError(ctx, "error update user quota cache: "+err.Error())
} else {
quota := task.Quota
if quota != 0 {
err = model.IncreaseUserQuota(task.UserId, quota, false)
if err != nil {
logger.LogError(ctx, "fail to increase user quota: "+err.Error())
}
logContent := fmt.Sprintf("异步任务执行失败 %s，补偿 %s", task.TaskID, logger.LogQuota(quota))
model.RecordLog(task.UserId, model.LogTypeSystem, logContent)
}
}
}
if responseItem.Status == model.TaskStatusSuccess {
task.Progress = "100%"
}
task.Data = responseItem.Data
err = task.Update()
if err != nil {
common.SysLog("UpdateMidjourneyTask task error: " + err.Error())
}
}
return nil
}
func checkTaskNeedUpdate(oldTask *model.Task, newTask dto.SunoDataResponse) bool {
if oldTask.SubmitTime != newTask.SubmitTime {
return true
}
if oldTask.StartTime != newTask.StartTime {
return true
}
if oldTask.FinishTime != newTask.FinishTime {
return true
}
if string(oldTask.Status) != newTask.Status {
return true
}
if oldTask.FailReason != newTask.FailReason {
return true
}
if oldTask.FinishTime != newTask.FinishTime {
return true
}
if (oldTask.Status == model.TaskStatusFailure || oldTask.Status == model.TaskStatusSuccess) && oldTask.Progress != "100%" {
return true
}
oldData, _ := json.Marshal(oldTask.Data)
newData, _ := json.Marshal(newTask.Data)
sort.Slice(oldData, func(i, j int) bool {
return oldData[i] < oldData[j]
})
sort.Slice(newData, func(i, j int) bool {
return newData[i] < newData[j]
})
if string(oldData) != string(newData) {
return true
}
return false
service.TaskPollingLoop()
}
func GetAllTask(c *gin.Context) {
......@@ -247,7 +38,7 @@ func GetAllTask(c *gin.Context) {
items := model.TaskGetAllTasks(pageInfo.GetStartIdx(), pageInfo.GetPageSize(), queryParams)
total := model.TaskCountAllTasks(queryParams)
pageInfo.SetTotal(int(total))
pageInfo.SetItems(items)
pageInfo.SetItems(tasksToDto(items, true))
common.ApiSuccess(c, pageInfo)
}
......@@ -271,6 +62,33 @@ func GetUserTask(c *gin.Context) {
items := model.TaskGetAllUserTask(userId, pageInfo.GetStartIdx(), pageInfo.GetPageSize(), queryParams)
total := model.TaskCountAllUserTask(userId, queryParams)
pageInfo.SetTotal(int(total))
pageInfo.SetItems(items)
pageInfo.SetItems(tasksToDto(items, false))
common.ApiSuccess(c, pageInfo)
}
func tasksToDto(tasks []*model.Task, fillUser bool) []*dto.TaskDto {
var userIdMap map[int]*model.UserBase
if fillUser {
userIdMap = make(map[int]*model.UserBase)
userIds := types.NewSet[int]()
for _, task := range tasks {
userIds.Add(task.UserId)
}
for _, userId := range userIds.Items() {
cacheUser, err := model.GetUserCache(userId)
if err == nil {
userIdMap[userId] = cacheUser
}
}
}
result := make([]*dto.TaskDto, len(tasks))
for i, task := range tasks {
if fillUser {
if user, ok := userIdMap[task.UserId]; ok {
task.Username = user.Username
}
}
result[i] = relay.TaskModel2Dto(task)
}
return result
}
controller/user.go
......@@ -582,6 +582,44 @@ func UpdateUser(c *gin.Context) {
return
}
func AdminClearUserBinding(c *gin.Context) {
id, err := strconv.Atoi(c.Param("id"))
if err != nil {
common.ApiErrorI18n(c, i18n.MsgInvalidParams)
return
}
bindingType := strings.ToLower(strings.TrimSpace(c.Param("binding_type")))
if bindingType == "" {
common.ApiErrorI18n(c, i18n.MsgInvalidParams)
return
}
user, err := model.GetUserById(id, false)
if err != nil {
common.ApiError(c, err)
return
}
myRole := c.GetInt("role")
if myRole <= user.Role && myRole != common.RoleRootUser {
common.ApiErrorI18n(c, i18n.MsgUserNoPermissionSameLevel)
return
}
if err := user.ClearBinding(bindingType); err != nil {
common.ApiError(c, err)
return
}
model.RecordLog(user.Id, model.LogTypeManage, fmt.Sprintf("admin cleared %s binding for user %s", bindingType, user.Username))
c.JSON(http.StatusOK, gin.H{
"success": true,
"message": "success",
})
}
func UpdateSelf(c *gin.Context) {
var requestData map[string]interface{}
err := json.NewDecoder(c.Request.Body).Decode(&requestData)
......
controller/video_proxy.go
......@@ -16,59 +16,44 @@ import (
"github.com/gin-gonic/gin"
)
// videoProxyError returns a standardized OpenAI-style error response.
func videoProxyError(c *gin.Context, status int, errType, message string) {
c.JSON(status, gin.H{
"error": gin.H{
"message": message,
"type": errType,
},
})
}
func VideoProxy(c *gin.Context) {
taskID := c.Param("task_id")
if taskID == "" {
c.JSON(http.StatusBadRequest, gin.H{
"error": gin.H{
"message": "task_id is required",
"type": "invalid_request_error",
},
})
videoProxyError(c, http.StatusBadRequest, "invalid_request_error", "task_id is required")
return
}
task, exists, err := model.GetByOnlyTaskId(taskID)
if err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to query task %s: %s", taskID, err.Error()))
c.JSON(http.StatusInternalServerError, gin.H{
"error": gin.H{
"message": "Failed to query task",
"type": "server_error",
},
})
videoProxyError(c, http.StatusInternalServerError, "server_error", "Failed to query task")
return
}
if !exists || task == nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to get task %s: %v", taskID, err))
c.JSON(http.StatusNotFound, gin.H{
"error": gin.H{
"message": "Task not found",
"type": "invalid_request_error",
},
})
videoProxyError(c, http.StatusNotFound, "invalid_request_error", "Task not found")
return
}
if task.Status != model.TaskStatusSuccess {
c.JSON(http.StatusBadRequest, gin.H{
"error": gin.H{
"message": fmt.Sprintf("Task is not completed yet, current status: %s", task.Status),
"type": "invalid_request_error",
},
})
videoProxyError(c, http.StatusBadRequest, "invalid_request_error",
fmt.Sprintf("Task is not completed yet, current status: %s", task.Status))
return
}
channel, err := model.CacheGetChannel(task.ChannelId)
if err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to get task %s: not found", taskID))
c.JSON(http.StatusInternalServerError, gin.H{
"error": gin.H{
"message": "Failed to retrieve channel information",
"type": "server_error",
},
})
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to get channel for task %s: %s", taskID, err.Error()))
videoProxyError(c, http.StatusInternalServerError, "server_error", "Failed to retrieve channel information")
return
}
baseURL := channel.GetBaseURL()
......@@ -81,12 +66,7 @@ func VideoProxy(c *gin.Context) {
client, err := service.GetHttpClientWithProxy(proxy)
if err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to create proxy client for task %s: %s", taskID, err.Error()))
c.JSON(http.StatusInternalServerError, gin.H{
"error": gin.H{
"message": "Failed to create proxy client",
"type": "server_error",
},
})
videoProxyError(c, http.StatusInternalServerError, "server_error", "Failed to create proxy client")
return
}
......@@ -95,12 +75,7 @@ func VideoProxy(c *gin.Context) {
req, err := http.NewRequestWithContext(ctx, http.MethodGet, "", nil)
if err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to create request: %s", err.Error()))
c.JSON(http.StatusInternalServerError, gin.H{
"error": gin.H{
"message": "Failed to create proxy request",
"type": "server_error",
},
})
videoProxyError(c, http.StatusInternalServerError, "server_error", "Failed to create proxy request")
return
}
......@@ -109,68 +84,43 @@ func VideoProxy(c *gin.Context) {
apiKey := task.PrivateData.Key
if apiKey == "" {
logger.LogError(c.Request.Context(), fmt.Sprintf("Missing stored API key for Gemini task %s", taskID))
c.JSON(http.StatusInternalServerError, gin.H{
"error": gin.H{
"message": "API key not stored for task",
"type": "server_error",
},
})
videoProxyError(c, http.StatusInternalServerError, "server_error", "API key not stored for task")
return
}
videoURL, err = getGeminiVideoURL(channel, task, apiKey)
if err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to resolve Gemini video URL for task %s: %s", taskID, err.Error()))
c.JSON(http.StatusBadGateway, gin.H{
"error": gin.H{
"message": "Failed to resolve Gemini video URL",
"type": "server_error",
},
})
videoProxyError(c, http.StatusBadGateway, "server_error", "Failed to resolve Gemini video URL")
return
}
req.Header.Set("x-goog-api-key", apiKey)
case constant.ChannelTypeOpenAI, constant.ChannelTypeSora:
videoURL = fmt.Sprintf("%s/v1/videos/%s/content", baseURL, task.TaskID)
videoURL = fmt.Sprintf("%s/v1/videos/%s/content", baseURL, task.GetUpstreamTaskID())
req.Header.Set("Authorization", "Bearer "+channel.Key)
default:
// Video URL is directly in task.FailReason
videoURL = task.FailReason
// Video URL is stored in PrivateData.ResultURL (fallback to FailReason for old data)
videoURL = task.GetResultURL()
}
req.URL, err = url.Parse(videoURL)
if err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to parse URL %s: %s", videoURL, err.Error()))
c.JSON(http.StatusInternalServerError, gin.H{
"error": gin.H{
"message": "Failed to create proxy request",
"type": "server_error",
},
})
videoProxyError(c, http.StatusInternalServerError, "server_error", "Failed to create proxy request")
return
}
resp, err := client.Do(req)
if err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to fetch video from %s: %s", videoURL, err.Error()))
c.JSON(http.StatusBadGateway, gin.H{
"error": gin.H{
"message": "Failed to fetch video content",
"type": "server_error",
},
})
videoProxyError(c, http.StatusBadGateway, "server_error", "Failed to fetch video content")
return
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
logger.LogError(c.Request.Context(), fmt.Sprintf("Upstream returned status %d for %s", resp.StatusCode, videoURL))
c.JSON(http.StatusBadGateway, gin.H{
"error": gin.H{
"message": fmt.Sprintf("Upstream service returned status %d", resp.StatusCode),
"type": "server_error",
},
})
videoProxyError(c, http.StatusBadGateway, "server_error",
fmt.Sprintf("Upstream service returned status %d", resp.StatusCode))
return
}
......@@ -180,10 +130,9 @@ func VideoProxy(c *gin.Context) {
}
}
c.Writer.Header().Set("Cache-Control", "public, max-age=86400") // Cache for 24 hours
c.Writer.Header().Set("Cache-Control", "public, max-age=86400")
c.Writer.WriteHeader(resp.StatusCode)
_, err = io.Copy(c.Writer, resp.Body)
if err != nil {
if _, err = io.Copy(c.Writer, resp.Body); err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to stream video content: %s", err.Error()))
}
}
controller/video_proxy_gemini.go
package controller
import (
"encoding/json"
"fmt"
"io"
"strconv"
"strings"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/relay"
......@@ -37,7 +37,7 @@ func getGeminiVideoURL(channel *model.Channel, task *model.Task, apiKey string)
proxy := channel.GetSetting().Proxy
resp, err := adaptor.FetchTask(baseURL, apiKey, map[string]any{
"task_id": task.TaskID,
"task_id": task.GetUpstreamTaskID(),
"action": task.Action,
}, proxy)
if err != nil {
......@@ -71,7 +71,7 @@ func extractGeminiVideoURLFromTaskData(task *model.Task) string {
return ""
}
var payload map[string]any
if err := json.Unmarshal(task.Data, &payload); err != nil {
if err := common.Unmarshal(task.Data, &payload); err != nil {
return ""
}
return extractGeminiVideoURLFromMap(payload)
......@@ -79,7 +79,7 @@ func extractGeminiVideoURLFromTaskData(task *model.Task) string {
func extractGeminiVideoURLFromPayload(body []byte) string {
var payload map[string]any
if err := json.Unmarshal(body, &payload); err != nil {
if err := common.Unmarshal(body, &payload); err != nil {
return ""
}
return extractGeminiVideoURLFromMap(payload)
......
dto/channel_settings.go
......@@ -24,14 +24,16 @@ const (
)
type ChannelOtherSettings struct {
AzureResponsesVersion string `json:"azure_responses_version,omitempty"`
VertexKeyType VertexKeyType `json:"vertex_key_type,omitempty"` // "json" or "api_key"
OpenRouterEnterprise *bool `json:"openrouter_enterprise,omitempty"`
ClaudeBetaQuery bool `json:"claude_beta_query,omitempty"` // Claude 渠道是否强制追加 ?beta=true
AllowServiceTier bool `json:"allow_service_tier,omitempty"` // 是否允许 service_tier 透传（默认过滤以避免额外计费）
DisableStore bool `json:"disable_store,omitempty"` // 是否禁用 store 透传（默认允许透传，禁用后可能导致 Codex 无法使用）
AllowSafetyIdentifier bool `json:"allow_safety_identifier,omitempty"` // 是否允许 safety_identifier 透传（默认过滤以保护用户隐私）
AwsKeyType AwsKeyType `json:"aws_key_type,omitempty"`
AzureResponsesVersion string `json:"azure_responses_version,omitempty"`
VertexKeyType VertexKeyType `json:"vertex_key_type,omitempty"` // "json" or "api_key"
OpenRouterEnterprise *bool `json:"openrouter_enterprise,omitempty"`
ClaudeBetaQuery bool `json:"claude_beta_query,omitempty"` // Claude 渠道是否强制追加 ?beta=true
AllowServiceTier bool `json:"allow_service_tier,omitempty"` // 是否允许 service_tier 透传（默认过滤以避免额外计费）
AllowInferenceGeo bool `json:"allow_inference_geo,omitempty"` // 是否允许 inference_geo 透传（仅 Claude，默认过滤以满足数据驻留合规）
DisableStore bool `json:"disable_store,omitempty"` // 是否禁用 store 透传（默认允许透传，禁用后可能导致 Codex 无法使用）
AllowSafetyIdentifier bool `json:"allow_safety_identifier,omitempty"` // 是否允许 safety_identifier 透传（默认过滤以保护用户隐私）
AllowIncludeObfuscation bool `json:"allow_include_obfuscation,omitempty"` // 是否允许 stream_options.include_obfuscation 透传（默认过滤以避免关闭流混淆保护）
AwsKeyType AwsKeyType `json:"aws_key_type,omitempty"`
}
func (s *ChannelOtherSettings) IsOpenRouterEnterprise() bool {
......
dto/claude.go
......@@ -190,10 +190,13 @@ type ClaudeToolChoice struct {
}
type ClaudeRequest struct {
Model string `json:"model"`
Prompt string `json:"prompt,omitempty"`
System any `json:"system,omitempty"`
Messages []ClaudeMessage `json:"messages,omitempty"`
Model string `json:"model"`
Prompt string `json:"prompt,omitempty"`
System any `json:"system,omitempty"`
Messages []ClaudeMessage `json:"messages,omitempty"`
// InferenceGeo controls Claude data residency region.
// This field is filtered by default and can be enabled via channel setting allow_inference_geo.
InferenceGeo string `json:"inference_geo,omitempty"`
MaxTokens uint `json:"max_tokens,omitempty"`
MaxTokensToSample uint `json:"max_tokens_to_sample,omitempty"`
StopSequences []string `json:"stop_sequences,omitempty"`
......@@ -210,7 +213,8 @@ type ClaudeRequest struct {
Thinking *Thinking `json:"thinking,omitempty"`
McpServers json.RawMessage `json:"mcp_servers,omitempty"`
Metadata json.RawMessage `json:"metadata,omitempty"`
// 服务层级字段，用于指定 API 服务等级。允许透传可能导致实际计费高于预期，默认应过滤
// ServiceTier specifies upstream service level and may affect billing.
// This field is filtered by default and can be enabled via channel setting allow_service_tier.
ServiceTier string `json:"service_tier,omitempty"`
}
......
dto/gemini.go
......@@ -324,25 +324,26 @@ type GeminiChatTool struct {
}
type GeminiChatGenerationConfig struct {
Temperature *float64 `json:"temperature,omitempty"`
TopP float64 `json:"topP,omitempty"`
TopK float64 `json:"topK,omitempty"`
MaxOutputTokens uint `json:"maxOutputTokens,omitempty"`
CandidateCount int `json:"candidateCount,omitempty"`
StopSequences []string `json:"stopSequences,omitempty"`
ResponseMimeType string `json:"responseMimeType,omitempty"`
ResponseSchema any `json:"responseSchema,omitempty"`
ResponseJsonSchema json.RawMessage `json:"responseJsonSchema,omitempty"`
PresencePenalty *float32 `json:"presencePenalty,omitempty"`
FrequencyPenalty *float32 `json:"frequencyPenalty,omitempty"`
ResponseLogprobs bool `json:"responseLogprobs,omitempty"`
Logprobs *int32 `json:"logprobs,omitempty"`
MediaResolution MediaResolution `json:"mediaResolution,omitempty"`
Seed int64 `json:"seed,omitempty"`
ResponseModalities []string `json:"responseModalities,omitempty"`
ThinkingConfig *GeminiThinkingConfig `json:"thinkingConfig,omitempty"`
SpeechConfig json.RawMessage `json:"speechConfig,omitempty"` // RawMessage to allow flexible speech config
ImageConfig json.RawMessage `json:"imageConfig,omitempty"` // RawMessage to allow flexible image config
Temperature *float64 `json:"temperature,omitempty"`
TopP float64 `json:"topP,omitempty"`
TopK float64 `json:"topK,omitempty"`
MaxOutputTokens uint `json:"maxOutputTokens,omitempty"`
CandidateCount int `json:"candidateCount,omitempty"`
StopSequences []string `json:"stopSequences,omitempty"`
ResponseMimeType string `json:"responseMimeType,omitempty"`
ResponseSchema any `json:"responseSchema,omitempty"`
ResponseJsonSchema json.RawMessage `json:"responseJsonSchema,omitempty"`
PresencePenalty *float32 `json:"presencePenalty,omitempty"`
FrequencyPenalty *float32 `json:"frequencyPenalty,omitempty"`
ResponseLogprobs bool `json:"responseLogprobs,omitempty"`
Logprobs *int32 `json:"logprobs,omitempty"`
EnableEnhancedCivicAnswers *bool `json:"enableEnhancedCivicAnswers,omitempty"`
MediaResolution MediaResolution `json:"mediaResolution,omitempty"`
Seed int64 `json:"seed,omitempty"`
ResponseModalities []string `json:"responseModalities,omitempty"`
ThinkingConfig *GeminiThinkingConfig `json:"thinkingConfig,omitempty"`
SpeechConfig json.RawMessage `json:"speechConfig,omitempty"` // RawMessage to allow flexible speech config
ImageConfig json.RawMessage `json:"imageConfig,omitempty"` // RawMessage to allow flexible image config
}
// UnmarshalJSON allows GeminiChatGenerationConfig to accept both snake_case and camelCase fields.
......@@ -350,22 +351,23 @@ func (c *GeminiChatGenerationConfig) UnmarshalJSON(data []byte) error {
type Alias GeminiChatGenerationConfig
var aux struct {
Alias
TopPSnake float64 `json:"top_p,omitempty"`
TopKSnake float64 `json:"top_k,omitempty"`
MaxOutputTokensSnake uint `json:"max_output_tokens,omitempty"`
CandidateCountSnake int `json:"candidate_count,omitempty"`
StopSequencesSnake []string `json:"stop_sequences,omitempty"`
ResponseMimeTypeSnake string `json:"response_mime_type,omitempty"`
ResponseSchemaSnake any `json:"response_schema,omitempty"`
ResponseJsonSchemaSnake json.RawMessage `json:"response_json_schema,omitempty"`
PresencePenaltySnake *float32 `json:"presence_penalty,omitempty"`
FrequencyPenaltySnake *float32 `json:"frequency_penalty,omitempty"`
ResponseLogprobsSnake bool `json:"response_logprobs,omitempty"`
MediaResolutionSnake MediaResolution `json:"media_resolution,omitempty"`
ResponseModalitiesSnake []string `json:"response_modalities,omitempty"`
ThinkingConfigSnake *GeminiThinkingConfig `json:"thinking_config,omitempty"`
SpeechConfigSnake json.RawMessage `json:"speech_config,omitempty"`
ImageConfigSnake json.RawMessage `json:"image_config,omitempty"`
TopPSnake float64 `json:"top_p,omitempty"`
TopKSnake float64 `json:"top_k,omitempty"`
MaxOutputTokensSnake uint `json:"max_output_tokens,omitempty"`
CandidateCountSnake int `json:"candidate_count,omitempty"`
StopSequencesSnake []string `json:"stop_sequences,omitempty"`
ResponseMimeTypeSnake string `json:"response_mime_type,omitempty"`
ResponseSchemaSnake any `json:"response_schema,omitempty"`
ResponseJsonSchemaSnake json.RawMessage `json:"response_json_schema,omitempty"`
PresencePenaltySnake *float32 `json:"presence_penalty,omitempty"`
FrequencyPenaltySnake *float32 `json:"frequency_penalty,omitempty"`
ResponseLogprobsSnake bool `json:"response_logprobs,omitempty"`
EnableEnhancedCivicAnswersSnake *bool `json:"enable_enhanced_civic_answers,omitempty"`
MediaResolutionSnake MediaResolution `json:"media_resolution,omitempty"`
ResponseModalitiesSnake []string `json:"response_modalities,omitempty"`
ThinkingConfigSnake *GeminiThinkingConfig `json:"thinking_config,omitempty"`
SpeechConfigSnake json.RawMessage `json:"speech_config,omitempty"`
ImageConfigSnake json.RawMessage `json:"image_config,omitempty"`
}
if err := common.Unmarshal(data, &aux); err != nil {
......@@ -408,6 +410,9 @@ func (c *GeminiChatGenerationConfig) UnmarshalJSON(data []byte) error {
if aux.ResponseLogprobsSnake {
c.ResponseLogprobs = aux.ResponseLogprobsSnake
}
if aux.EnableEnhancedCivicAnswersSnake != nil {
c.EnableEnhancedCivicAnswers = aux.EnableEnhancedCivicAnswersSnake
}
if aux.MediaResolutionSnake != "" {
c.MediaResolution = aux.MediaResolutionSnake
}
......@@ -453,12 +458,14 @@ type GeminiChatResponse struct {
}
type GeminiUsageMetadata struct {
PromptTokenCount int `json:"promptTokenCount"`
CandidatesTokenCount int `json:"candidatesTokenCount"`
TotalTokenCount int `json:"totalTokenCount"`
ThoughtsTokenCount int `json:"thoughtsTokenCount"`
CachedContentTokenCount int `json:"cachedContentTokenCount"`
PromptTokensDetails []GeminiPromptTokensDetails `json:"promptTokensDetails"`
PromptTokenCount int `json:"promptTokenCount"`
ToolUsePromptTokenCount int `json:"toolUsePromptTokenCount"`
CandidatesTokenCount int `json:"candidatesTokenCount"`
TotalTokenCount int `json:"totalTokenCount"`
ThoughtsTokenCount int `json:"thoughtsTokenCount"`
CachedContentTokenCount int `json:"cachedContentTokenCount"`
PromptTokensDetails []GeminiPromptTokensDetails `json:"promptTokensDetails"`
ToolUsePromptTokensDetails []GeminiPromptTokensDetails `json:"toolUsePromptTokensDetails"`
}
type GeminiPromptTokensDetails struct {
......
dto/openai_request.go
......@@ -54,18 +54,22 @@ type GeneralOpenAIRequest struct {
ParallelTooCalls *bool `json:"parallel_tool_calls,omitempty"`
Tools []ToolCallRequest `json:"tools,omitempty"`
ToolChoice any `json:"tool_choice,omitempty"`
FunctionCall json.RawMessage `json:"function_call,omitempty"`
User string `json:"user,omitempty"`
LogProbs bool `json:"logprobs,omitempty"`
TopLogProbs int `json:"top_logprobs,omitempty"`
Dimensions int `json:"dimensions,omitempty"`
Modalities json.RawMessage `json:"modalities,omitempty"`
Audio json.RawMessage `json:"audio,omitempty"`
// ServiceTier specifies upstream service level and may affect billing.
// This field is filtered by default and can be enabled via channel setting allow_service_tier.
ServiceTier string `json:"service_tier,omitempty"`
LogProbs bool `json:"logprobs,omitempty"`
TopLogProbs int `json:"top_logprobs,omitempty"`
Dimensions int `json:"dimensions,omitempty"`
Modalities json.RawMessage `json:"modalities,omitempty"`
Audio json.RawMessage `json:"audio,omitempty"`
// 安全标识符，用于帮助 OpenAI 检测可能违反使用政策的应用程序用户
// 注意：此字段会向 OpenAI 发送用户标识信息，默认过滤以保护用户隐私
// 注意：此字段会向 OpenAI 发送用户标识信息，默认过滤，可通过 allow_safety_identifier 开启
SafetyIdentifier string `json:"safety_identifier,omitempty"`
// Whether or not to store the output of this chat completion request for use in our model distillation or evals products.
// 是否存储此次请求数据供 OpenAI 用于评估和优化产品
// 注意：默认过滤此字段以保护用户隐私，但过滤后可能导致 Codex 无法正常使用
// 注意：默认允许透传，可通过 disable_store 禁用；禁用后可能导致 Codex 无法正常使用
Store json.RawMessage `json:"store,omitempty"`
// Used by OpenAI to cache responses for similar requests to optimize your cache hit rates. Replaces the user field
PromptCacheKey string `json:"prompt_cache_key,omitempty"`
......@@ -261,6 +265,9 @@ type FunctionRequest struct {
type StreamOptions struct {
IncludeUsage bool `json:"include_usage,omitempty"`
// IncludeObfuscation is only for /v1/responses stream payload.
// This field is filtered by default and can be enabled via channel setting allow_include_obfuscation.
IncludeObfuscation bool `json:"include_obfuscation,omitempty"`
}
func (r *GeneralOpenAIRequest) GetMaxTokens() uint {
......@@ -799,30 +806,42 @@ type WebSearchOptions struct {
// https://platform.openai.com/docs/api-reference/responses/create
type OpenAIResponsesRequest struct {
Model string `json:"model"`
Input json.RawMessage `json:"input,omitempty"`
Include json.RawMessage `json:"include,omitempty"`
Model string `json:"model"`
Input json.RawMessage `json:"input,omitempty"`
Include json.RawMessage `json:"include,omitempty"`
// 在后台运行推理，暂时还不支持依赖的接口
// Background json.RawMessage `json:"background,omitempty"`
Conversation json.RawMessage `json:"conversation,omitempty"`
ContextManagement json.RawMessage `json:"context_management,omitempty"`
Instructions json.RawMessage `json:"instructions,omitempty"`
MaxOutputTokens uint `json:"max_output_tokens,omitempty"`
TopLogProbs *int `json:"top_logprobs,omitempty"`
Metadata json.RawMessage `json:"metadata,omitempty"`
ParallelToolCalls json.RawMessage `json:"parallel_tool_calls,omitempty"`
PreviousResponseID string `json:"previous_response_id,omitempty"`
Reasoning *Reasoning `json:"reasoning,omitempty"`
// 服务层级字段，用于指定 API 服务等级。允许透传可能导致实际计费高于预期，默认应过滤
ServiceTier string `json:"service_tier,omitempty"`
// ServiceTier specifies upstream service level and may affect billing.
// This field is filtered by default and can be enabled via channel setting allow_service_tier.
ServiceTier string `json:"service_tier,omitempty"`
// Store controls whether upstream may store request/response data.
// This field is allowed by default and can be disabled via channel setting disable_store.
Store json.RawMessage `json:"store,omitempty"`
PromptCacheKey json.RawMessage `json:"prompt_cache_key,omitempty"`
PromptCacheRetention json.RawMessage `json:"prompt_cache_retention,omitempty"`
Stream bool `json:"stream,omitempty"`
Temperature *float64 `json:"temperature,omitempty"`
Text json.RawMessage `json:"text,omitempty"`
ToolChoice json.RawMessage `json:"tool_choice,omitempty"`
Tools json.RawMessage `json:"tools,omitempty"` // 需要处理的参数很少，MCP 参数太多不确定，所以用 map
TopP *float64 `json:"top_p,omitempty"`
Truncation string `json:"truncation,omitempty"`
User string `json:"user,omitempty"`
MaxToolCalls uint `json:"max_tool_calls,omitempty"`
Prompt json.RawMessage `json:"prompt,omitempty"`
// SafetyIdentifier carries client identity for policy abuse detection.
// This field is filtered by default and can be enabled via channel setting allow_safety_identifier.
SafetyIdentifier string `json:"safety_identifier,omitempty"`
Stream bool `json:"stream,omitempty"`
StreamOptions *StreamOptions `json:"stream_options,omitempty"`
Temperature *float64 `json:"temperature,omitempty"`
Text json.RawMessage `json:"text,omitempty"`
ToolChoice json.RawMessage `json:"tool_choice,omitempty"`
Tools json.RawMessage `json:"tools,omitempty"` // 需要处理的参数很少，MCP 参数太多不确定，所以用 map
TopP *float64 `json:"top_p,omitempty"`
Truncation string `json:"truncation,omitempty"`
User string `json:"user,omitempty"`
MaxToolCalls uint `json:"max_tool_calls,omitempty"`
Prompt json.RawMessage `json:"prompt,omitempty"`
// qwen
EnableThinking json.RawMessage `json:"enable_thinking,omitempty"`
// perplexity
......
dto/suno.go
......@@ -4,10 +4,6 @@ import (
"encoding/json"
)
type TaskData interface {
SunoDataResponse | []SunoDataResponse | string | any
}
type SunoSubmitReq struct {
GptDescriptionPrompt string `json:"gpt_description_prompt,omitempty"`
Prompt string `json:"prompt,omitempty"`
......@@ -20,10 +16,6 @@ type SunoSubmitReq struct {
MakeInstrumental bool `json:"make_instrumental"`
}
type FetchReq struct {
IDs []string `json:"ids"`
}
type SunoDataResponse struct {
TaskID string `json:"task_id" gorm:"type:varchar(50);index"`
Action string `json:"action" gorm:"type:varchar(40);index"` // 任务类型, song, lyrics, description-mode
......@@ -66,30 +58,6 @@ type SunoLyrics struct {
Text string `json:"text"`
}
const TaskSuccessCode = "success"
type TaskResponse[T TaskData] struct {
Code string `json:"code"`
Message string `json:"message"`
Data T `json:"data"`
}
func (t *TaskResponse[T]) IsSuccess() bool {
return t.Code == TaskSuccessCode
}
type TaskDto struct {
TaskID string `json:"task_id"` // 第三方id，不一定有/ song id\ Task id
Action string `json:"action"` // 任务类型, song, lyrics, description-mode
Status string `json:"status"` // 任务状态, submitted, queueing, processing, success, failed
FailReason string `json:"fail_reason"`
SubmitTime int64 `json:"submit_time"`
StartTime int64 `json:"start_time"`
FinishTime int64 `json:"finish_time"`
Progress string `json:"progress"`
Data json.RawMessage `json:"data"`
}
type SunoGoAPISubmitReq struct {
CustomMode bool `json:"custom_mode"`
......
dto/task.go
package dto
import (
"encoding/json"
)
type TaskError struct {
Code string `json:"code"`
Message string `json:"message"`
......@@ -8,3 +12,46 @@ type TaskError struct {
LocalError bool `json:"-"`
Error error `json:"-"`
}
type TaskData interface {
SunoDataResponse | []SunoDataResponse | string | any
}
const TaskSuccessCode = "success"
type TaskResponse[T TaskData] struct {
Code string `json:"code"`
Message string `json:"message"`
Data T `json:"data"`
}
func (t *TaskResponse[T]) IsSuccess() bool {
return t.Code == TaskSuccessCode
}
type TaskDto struct {
ID int64 `json:"id"`
CreatedAt int64 `json:"created_at"`
UpdatedAt int64 `json:"updated_at"`
TaskID string `json:"task_id"`
Platform string `json:"platform"`
UserId int `json:"user_id"`
Group string `json:"group"`
ChannelId int `json:"channel_id"`
Quota int `json:"quota"`
Action string `json:"action"`
Status string `json:"status"`
FailReason string `json:"fail_reason"`
ResultURL string `json:"result_url,omitempty"` // 任务结果 URL（视频地址等）
SubmitTime int64 `json:"submit_time"`
StartTime int64 `json:"start_time"`
FinishTime int64 `json:"finish_time"`
Progress string `json:"progress"`
Properties any `json:"properties"`
Username string `json:"username,omitempty"`
Data json.RawMessage `json:"data"`
}
type FetchReq struct {
IDs []string `json:"ids"`
}
logger/logger.go
......@@ -2,7 +2,6 @@ package logger
import (
"context"
"encoding/json"
"fmt"
"io"
"log"
......@@ -151,7 +150,7 @@ func FormatQuota(quota int) string {
// LogJson 仅供测试使用 only for test
func LogJson(ctx context.Context, msg string, obj any) {
jsonStr, err := json.Marshal(obj)
jsonStr, err := common.Marshal(obj)
if err != nil {
LogError(ctx, fmt.Sprintf("json marshal failed: %s", err.Error()))
return
......
main.go
......@@ -19,6 +19,7 @@ import (
"github.com/QuantumNous/new-api/middleware"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/oauth"
"github.com/QuantumNous/new-api/relay"
"github.com/QuantumNous/new-api/router"
"github.com/QuantumNous/new-api/service"
_ "github.com/QuantumNous/new-api/setting/performance_setting"
......@@ -111,6 +112,15 @@ func main() {
// Subscription quota reset task (daily/weekly/monthly/custom)
service.StartSubscriptionQuotaResetTask()
// Wire task polling adaptor factory (breaks service -> relay import cycle)
service.GetTaskAdaptorFunc = func(platform constant.TaskPlatform) service.TaskPollingAdaptor {
a := relay.GetTaskAdaptor(platform)
if a == nil {
return nil
}
return a
}
if common.IsMasterNode && constant.UpdateTask {
gopool.Go(func() {
controller.UpdateMidjourneyTaskBulk()
......
middleware/auth.go
......@@ -170,6 +170,24 @@ func WssAuth(c *gin.Context) {
}
// TokenOrUserAuth allows either session-based user auth or API token auth.
// Used for endpoints that need to be accessible from both the dashboard and API clients.
func TokenOrUserAuth() func(c *gin.Context) {
return func(c *gin.Context) {
// Try session auth first (dashboard users)
session := sessions.Default(c)
if id := session.Get("id"); id != nil {
if status, ok := session.Get("status").(int); ok && status == common.UserStatusEnabled {
c.Set("id", id)
c.Next()
return
}
}
// Fall back to token auth (API clients)
TokenAuth()(c)
}
}
// TokenAuthReadOnly 宽松版本的令牌认证中间件，用于只读查询接口。
// 只验证令牌 key 是否存在，不检查令牌状态、过期时间和额度。
// 即使令牌已过期、已耗尽或已禁用，也允许访问。
......
middleware/logger.go
......@@ -7,14 +7,28 @@ import (
"github.com/gin-gonic/gin"
)
const RouteTagKey = "route_tag"
func RouteTag(tag string) gin.HandlerFunc {
return func(c *gin.Context) {
c.Set(RouteTagKey, tag)
c.Next()
}
}
func SetUpLogger(server *gin.Engine) {
server.Use(gin.LoggerWithFormatter(func(param gin.LogFormatterParams) string {
var requestID string
if param.Keys != nil {
requestID = param.Keys[common.RequestIdKey].(string)
requestID, _ = param.Keys[common.RequestIdKey].(string)
}
tag, _ := param.Keys[RouteTagKey].(string)
if tag == "" {
tag = "web"
}
return fmt.Sprintf("[GIN] %s | %s | %3d | %13v | %15s | %7s %s\n",
return fmt.Sprintf("[GIN] %s | %s | %s | %3d | %13v | %15s | %7s %s\n",
param.TimeStamp.Format("2006/01/02 - 15:04:05"),
tag,
requestID,
param.StatusCode,
param.Latency,
......
model/custom_oauth_provider.go
......@@ -2,32 +2,65 @@ package model
import (
"errors"
"fmt"
"strings"
"time"
"github.com/QuantumNous/new-api/common"
)
type accessPolicyPayload struct {
Logic string `json:"logic"`
Conditions []accessConditionItem `json:"conditions"`
Groups []accessPolicyPayload `json:"groups"`
}
type accessConditionItem struct {
Field string `json:"field"`
Op string `json:"op"`
Value any `json:"value"`
}
var supportedAccessPolicyOps = map[string]struct{}{
"eq": {},
"ne": {},
"gt": {},
"gte": {},
"lt": {},
"lte": {},
"in": {},
"not_in": {},
"contains": {},
"not_contains": {},
"exists": {},
"not_exists": {},
}
// CustomOAuthProvider stores configuration for custom OAuth providers
type CustomOAuthProvider struct {
Id int `json:"id" gorm:"primaryKey"`
Name string `json:"name" gorm:"type:varchar(64);not null"` // Display name, e.g., "GitHub Enterprise"
Slug string `json:"slug" gorm:"type:varchar(64);uniqueIndex;not null"` // URL identifier, e.g., "github-enterprise"
Enabled bool `json:"enabled" gorm:"default:false"` // Whether this provider is enabled
ClientId string `json:"client_id" gorm:"type:varchar(256)"` // OAuth client ID
ClientSecret string `json:"-" gorm:"type:varchar(512)"` // OAuth client secret (not returned to frontend)
AuthorizationEndpoint string `json:"authorization_endpoint" gorm:"type:varchar(512)"` // Authorization URL
TokenEndpoint string `json:"token_endpoint" gorm:"type:varchar(512)"` // Token exchange URL
UserInfoEndpoint string `json:"user_info_endpoint" gorm:"type:varchar(512)"` // User info URL
Scopes string `json:"scopes" gorm:"type:varchar(256);default:'openid profile email'"` // OAuth scopes
Id int `json:"id" gorm:"primaryKey"`
Name string `json:"name" gorm:"type:varchar(64);not null"` // Display name, e.g., "GitHub Enterprise"
Slug string `json:"slug" gorm:"type:varchar(64);uniqueIndex;not null"` // URL identifier, e.g., "github-enterprise"
Icon string `json:"icon" gorm:"type:varchar(128);default:''"` // Icon name from @lobehub/icons
Enabled bool `json:"enabled" gorm:"default:false"` // Whether this provider is enabled
ClientId string `json:"client_id" gorm:"type:varchar(256)"` // OAuth client ID
ClientSecret string `json:"-" gorm:"type:varchar(512)"` // OAuth client secret (not returned to frontend)
AuthorizationEndpoint string `json:"authorization_endpoint" gorm:"type:varchar(512)"` // Authorization URL
TokenEndpoint string `json:"token_endpoint" gorm:"type:varchar(512)"` // Token exchange URL
UserInfoEndpoint string `json:"user_info_endpoint" gorm:"type:varchar(512)"` // User info URL
Scopes string `json:"scopes" gorm:"type:varchar(256);default:'openid profile email'"` // OAuth scopes
// Field mapping configuration (supports JSONPath via gjson)
UserIdField string `json:"user_id_field" gorm:"type:varchar(128);default:'sub'"` // User ID field path, e.g., "sub", "id", "data.user.id"
UsernameField string `json:"username_field" gorm:"type:varchar(128);default:'preferred_username'"` // Username field path
DisplayNameField string `json:"display_name_field" gorm:"type:varchar(128);default:'name'"` // Display name field path
EmailField string `json:"email_field" gorm:"type:varchar(128);default:'email'"` // Email field path
UserIdField string `json:"user_id_field" gorm:"type:varchar(128);default:'sub'"` // User ID field path, e.g., "sub", "id", "data.user.id"
UsernameField string `json:"username_field" gorm:"type:varchar(128);default:'preferred_username'"` // Username field path
DisplayNameField string `json:"display_name_field" gorm:"type:varchar(128);default:'name'"` // Display name field path
EmailField string `json:"email_field" gorm:"type:varchar(128);default:'email'"` // Email field path
// Advanced options
WellKnown string `json:"well_known" gorm:"type:varchar(512)"` // OIDC discovery endpoint (optional)
AuthStyle int `json:"auth_style" gorm:"default:0"` // 0=auto, 1=params, 2=header (Basic Auth)
WellKnown string `json:"well_known" gorm:"type:varchar(512)"` // OIDC discovery endpoint (optional)
AuthStyle int `json:"auth_style" gorm:"default:0"` // 0=auto, 1=params, 2=header (Basic Auth)
AccessPolicy string `json:"access_policy" gorm:"type:text"` // JSON policy for access control based on user info
AccessDeniedMessage string `json:"access_denied_message" gorm:"type:varchar(512)"` // Custom error message template when access is denied
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
......@@ -158,6 +191,57 @@ func validateCustomOAuthProvider(provider *CustomOAuthProvider) error {
if provider.Scopes == "" {
provider.Scopes = "openid profile email"
}
if strings.TrimSpace(provider.AccessPolicy) != "" {
var policy accessPolicyPayload
if err := common.UnmarshalJsonStr(provider.AccessPolicy, &policy); err != nil {
return errors.New("access_policy must be valid JSON")
}
if err := validateAccessPolicyPayload(&policy); err != nil {
return fmt.Errorf("access_policy is invalid: %w", err)
}
}
return nil
}
func validateAccessPolicyPayload(policy *accessPolicyPayload) error {
if policy == nil {
return errors.New("policy is nil")
}
logic := strings.ToLower(strings.TrimSpace(policy.Logic))
if logic == "" {
logic = "and"
}
if logic != "and" && logic != "or" {
return fmt.Errorf("unsupported logic: %s", logic)
}
if len(policy.Conditions) == 0 && len(policy.Groups) == 0 {
return errors.New("policy requires at least one condition or group")
}
for index, condition := range policy.Conditions {
field := strings.TrimSpace(condition.Field)
if field == "" {
return fmt.Errorf("condition[%d].field is required", index)
}
op := strings.ToLower(strings.TrimSpace(condition.Op))
if _, ok := supportedAccessPolicyOps[op]; !ok {
return fmt.Errorf("condition[%d].op is unsupported: %s", index, op)
}
if op == "in" || op == "not_in" {
if _, ok := condition.Value.([]any); !ok {
return fmt.Errorf("condition[%d].value must be an array for op %s", index, op)
}
}
}
for index := range policy.Groups {
if err := validateAccessPolicyPayload(&policy.Groups[index]); err != nil {
return fmt.Errorf("group[%d]: %w", index, err)
}
}
return nil
}
model/log.go
......@@ -199,6 +199,49 @@ func RecordConsumeLog(c *gin.Context, userId int, params RecordConsumeLogParams)
}
}
type RecordTaskBillingLogParams struct {
UserId int
LogType int
Content string
ChannelId int
ModelName string
Quota int
TokenId int
Group string
Other map[string]interface{}
}
func RecordTaskBillingLog(params RecordTaskBillingLogParams) {
if params.LogType == LogTypeConsume && !common.LogConsumeEnabled {
return
}
username, _ := GetUsernameById(params.UserId, false)
tokenName := ""
if params.TokenId > 0 {
if token, err := GetTokenById(params.TokenId); err == nil {
tokenName = token.Name
}
}
log := &Log{
UserId: params.UserId,
Username: username,
CreatedAt: common.GetTimestamp(),
Type: params.LogType,
Content: params.Content,
TokenName: tokenName,
ModelName: params.ModelName,
Quota: params.Quota,
ChannelId: params.ChannelId,
TokenId: params.TokenId,
Group: params.Group,
Other: common.MapToJsonStr(params.Other),
}
err := LOG_DB.Create(log).Error
if err != nil {
common.SysLog("failed to record task billing log: " + err.Error())
}
}
func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string, startIdx int, num int, channel int, group string, requestId string) (logs []*Log, total int64, err error) {
var tx *gorm.DB
if logType == LogTypeUnknown {
......@@ -252,8 +295,24 @@ func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName
Id int `gorm:"column:id"`
Name string `gorm:"column:name"`
}
if err = DB.Table("channels").Select("id, name").Where("id IN ?", channelIds.Items()).Find(&channels).Error; err != nil {
return logs, total, err
if common.MemoryCacheEnabled {
// Cache get channel
for _, channelId := range channelIds.Items() {
if cacheChannel, err := CacheGetChannel(channelId); err == nil {
channels = append(channels, struct {
Id int `gorm:"column:id"`
Name string `gorm:"column:name"`
}{
Id: channelId,
Name: cacheChannel.Name,
})
}
}
} else {
// Bulk query channels from DB
if err = DB.Table("channels").Select("id, name").Where("id IN ?", channelIds.Items()).Find(&channels).Error; err != nil {
return logs, total, err
}
}
channelMap := make(map[int]string, len(channels))
for _, channel := range channels {
......
model/midjourney.go
......@@ -157,6 +157,19 @@ func (midjourney *Midjourney) Update() error {
return err
}
// UpdateWithStatus performs a conditional UPDATE guarded by fromStatus (CAS).
// Returns (true, nil) if this caller won the update, (false, nil) if
// another process already moved the task out of fromStatus.
// UpdateWithStatus performs a conditional UPDATE guarded by fromStatus (CAS).
// Uses Model().Select("*").Updates() to avoid GORM Save()'s INSERT fallback.
func (midjourney *Midjourney) UpdateWithStatus(fromStatus string) (bool, error) {
result := DB.Model(midjourney).Where("status = ?", fromStatus).Select("*").Updates(midjourney)
if result.Error != nil {
return false, result.Error
}
return result.RowsAffected > 0, nil
}
func MjBulkUpdate(mjIds []string, params map[string]any) error {
return DB.Model(&Midjourney{}).
Where("mj_id in (?)", mjIds).
......
model/task.go
package model
import (
"bytes"
"database/sql/driver"
"encoding/json"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto"
commonRelay "github.com/QuantumNous/new-api/relay/common"
......@@ -64,13 +66,12 @@ type Task struct {
}
func (t *Task) SetData(data any) {
b, _ := json.Marshal(data)
b, _ := common.Marshal(data)
t.Data = json.RawMessage(b)
}
func (t *Task) GetData(v any) error {
err := json.Unmarshal(t.Data, &v)
return err
return common.Unmarshal(t.Data, &v)
}
type Properties struct {
......@@ -85,18 +86,59 @@ func (m *Properties) Scan(val interface{}) error {
*m = Properties{}
return nil
}
return json.Unmarshal(bytesValue, m)
return common.Unmarshal(bytesValue, m)
}
func (m Properties) Value() (driver.Value, error) {
if m == (Properties{}) {
return nil, nil
}
return json.Marshal(m)
return common.Marshal(m)
}
type TaskPrivateData struct {
Key string `json:"key,omitempty"`
Key string `json:"key,omitempty"`
UpstreamTaskID string `json:"upstream_task_id,omitempty"` // 上游真实 task ID
ResultURL string `json:"result_url,omitempty"` // 任务成功后的结果 URL（视频地址等）
// 计费上下文：用于异步退款/差额结算（轮询阶段读取）
BillingSource string `json:"billing_source,omitempty"` // "wallet" 或 "subscription"
SubscriptionId int `json:"subscription_id,omitempty"` // 订阅 ID，用于订阅退款
TokenId int `json:"token_id,omitempty"` // 令牌 ID，用于令牌额度退款
BillingContext *TaskBillingContext `json:"billing_context,omitempty"` // 计费参数快照（用于轮询阶段重新计算）
}
// TaskBillingContext 记录任务提交时的计费参数，以便轮询阶段可以重新计算额度。
type TaskBillingContext struct {
ModelPrice float64 `json:"model_price,omitempty"` // 模型单价
GroupRatio float64 `json:"group_ratio,omitempty"` // 分组倍率
ModelRatio float64 `json:"model_ratio,omitempty"` // 模型倍率
OtherRatios map[string]float64 `json:"other_ratios,omitempty"` // 附加倍率（时长、分辨率等）
OriginModelName string `json:"origin_model_name,omitempty"` // 模型名称，必须为OriginModelName
PerCallBilling bool `json:"per_call_billing,omitempty"` // 按次计费：跳过轮询阶段的差额结算
}
// GetUpstreamTaskID 获取上游真实 task ID（用于与 provider 通信）
// 旧数据没有 UpstreamTaskID 时，TaskID 本身就是上游 ID
func (t *Task) GetUpstreamTaskID() string {
if t.PrivateData.UpstreamTaskID != "" {
return t.PrivateData.UpstreamTaskID
}
return t.TaskID
}
// GetResultURL 获取任务结果 URL（视频地址等）
// 新数据存在 PrivateData.ResultURL 中；旧数据回退到 FailReason（历史兼容）
func (t *Task) GetResultURL() string {
if t.PrivateData.ResultURL != "" {
return t.PrivateData.ResultURL
}
return t.FailReason
}
// GenerateTaskID 生成对外暴露的 task_xxxx 格式 ID
func GenerateTaskID() string {
key, _ := common.GenerateRandomCharsKey(32)
return "task_" + key
}
func (p *TaskPrivateData) Scan(val interface{}) error {
......@@ -104,14 +146,14 @@ func (p *TaskPrivateData) Scan(val interface{}) error {
if len(bytesValue) == 0 {
return nil
}
return json.Unmarshal(bytesValue, p)
return common.Unmarshal(bytesValue, p)
}
func (p TaskPrivateData) Value() (driver.Value, error) {
if (p == TaskPrivateData{}) {
return nil, nil
}
return json.Marshal(p)
return common.Marshal(p)
}
// SyncTaskQueryParams 用于包含所有搜索条件的结构体，可以根据需求添加更多字段
......@@ -142,7 +184,16 @@ func InitTask(platform constant.TaskPlatform, relayInfo *commonRelay.RelayInfo)
}
}
// 使用预生成的公开 ID（如果有），否则新生成
taskID := ""
if relayInfo.TaskRelayInfo != nil && relayInfo.TaskRelayInfo.PublicTaskID != "" {
taskID = relayInfo.TaskRelayInfo.PublicTaskID
} else {
taskID = GenerateTaskID()
}
t := &Task{
TaskID: taskID,
UserId: relayInfo.UserId,
Group: relayInfo.UsingGroup,
SubmitTime: time.Now().Unix(),
......@@ -237,6 +288,20 @@ func TaskGetAllTasks(startIdx int, num int, queryParams SyncTaskQueryParams) []*
return tasks
}
func GetTimedOutUnfinishedTasks(cutoffUnix int64, limit int) []*Task {
var tasks []*Task
err := DB.Where("progress != ?", "100%").
Where("status NOT IN ?", []string{TaskStatusFailure, TaskStatusSuccess}).
Where("submit_time < ?", cutoffUnix).
Order("submit_time").
Limit(limit).
Find(&tasks).Error
if err != nil {
return nil
}
return tasks
}
func GetAllUnFinishSyncTasks(limit int) []*Task {
var tasks []*Task
var err error
......@@ -291,40 +356,70 @@ func GetByTaskIds(userId int, taskIds []any) ([]*Task, error) {
return task, nil
}
func TaskUpdateProgress(id int64, progress string) error {
return DB.Model(&Task{}).Where("id = ?", id).Update("progress", progress).Error
}
func (Task *Task) Insert() error {
var err error
err = DB.Create(Task).Error
return err
}
type taskSnapshot struct {
Status TaskStatus
Progress string
StartTime int64
FinishTime int64
FailReason string
ResultURL string
Data json.RawMessage
}
func (s taskSnapshot) Equal(other taskSnapshot) bool {
return s.Status == other.Status &&
s.Progress == other.Progress &&
s.StartTime == other.StartTime &&
s.FinishTime == other.FinishTime &&
s.FailReason == other.FailReason &&
s.ResultURL == other.ResultURL &&
bytes.Equal(s.Data, other.Data)
}
func (t *Task) Snapshot() taskSnapshot {
return taskSnapshot{
Status: t.Status,
Progress: t.Progress,
StartTime: t.StartTime,
FinishTime: t.FinishTime,
FailReason: t.FailReason,
ResultURL: t.PrivateData.ResultURL,
Data: t.Data,
}
}
func (Task *Task) Update() error {
var err error
err = DB.Save(Task).Error
return err
}
func TaskBulkUpdate(TaskIds []string, params map[string]any) error {
if len(TaskIds) == 0 {
return nil
}
return DB.Model(&Task{}).
Where("task_id in (?)", TaskIds).
Updates(params).Error
}
func TaskBulkUpdateByTaskIds(taskIDs []int64, params map[string]any) error {
if len(taskIDs) == 0 {
return nil
}
return DB.Model(&Task{}).
Where("id in (?)", taskIDs).
Updates(params).Error
// UpdateWithStatus performs a conditional UPDATE guarded by fromStatus (CAS).
// Returns (true, nil) if this caller won the update, (false, nil) if
// another process already moved the task out of fromStatus.
//
// Uses Model().Select("*").Updates() instead of Save() because GORM's Save
// falls back to INSERT ON CONFLICT when the WHERE-guarded UPDATE matches
// zero rows, which silently bypasses the CAS guard.
func (t *Task) UpdateWithStatus(fromStatus TaskStatus) (bool, error) {
result := DB.Model(t).Where("status = ?", fromStatus).Select("*").Updates(t)
if result.Error != nil {
return false, result.Error
}
return result.RowsAffected > 0, nil
}
// TaskBulkUpdateByID performs an unconditional bulk UPDATE by primary key IDs.
// WARNING: This function has NO CAS (Compare-And-Swap) guard — it will overwrite
// any concurrent status changes. DO NOT use in billing/quota lifecycle flows
// (e.g., timeout, success, failure transitions that trigger refunds or settlements).
// For status transitions that involve billing, use Task.UpdateWithStatus() instead.
func TaskBulkUpdateByID(ids []int64, params map[string]any) error {
if len(ids) == 0 {
return nil
......@@ -339,37 +434,6 @@ type TaskQuotaUsage struct {
Count float64 `json:"count"`
}
func SumUsedTaskQuota(queryParams SyncTaskQueryParams) (stat []TaskQuotaUsage, err error) {
query := DB.Model(Task{})
// 添加过滤条件
if queryParams.ChannelID != "" {
query = query.Where("channel_id = ?", queryParams.ChannelID)
}
if queryParams.UserID != "" {
query = query.Where("user_id = ?", queryParams.UserID)
}
if len(queryParams.UserIDs) != 0 {
query = query.Where("user_id in (?)", queryParams.UserIDs)
}
if queryParams.TaskID != "" {
query = query.Where("task_id = ?", queryParams.TaskID)
}
if queryParams.Action != "" {
query = query.Where("action = ?", queryParams.Action)
}
if queryParams.Status != "" {
query = query.Where("status = ?", queryParams.Status)
}
if queryParams.StartTimestamp != 0 {
query = query.Where("submit_time >= ?", queryParams.StartTimestamp)
}
if queryParams.EndTimestamp != 0 {
query = query.Where("submit_time <= ?", queryParams.EndTimestamp)
}
err = query.Select("mode, sum(quota) as count").Group("mode").Find(&stat).Error
return stat, err
}
// TaskCountAllTasks returns total tasks that match the given query params (admin usage)
func TaskCountAllTasks(queryParams SyncTaskQueryParams) int64 {
var total int64
......@@ -438,6 +502,6 @@ func (t *Task) ToOpenAIVideo() *dto.OpenAIVideo {
openAIVideo.SetProgressStr(t.Progress)
openAIVideo.CreatedAt = t.CreatedAt
openAIVideo.CompletedAt = t.UpdatedAt
openAIVideo.SetMetadata("url", t.FailReason)
openAIVideo.SetMetadata("url", t.GetResultURL())
return openAIVideo
}
model/task_cas_test.go 0 → 100644
package model
import (
"encoding/json"
"os"
"sync"
"testing"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
)
func TestMain(m *testing.M) {
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
if err != nil {
panic("failed to open test db: " + err.Error())
}
DB = db
LOG_DB = db
common.UsingSQLite = true
common.RedisEnabled = false
common.BatchUpdateEnabled = false
common.LogConsumeEnabled = true
sqlDB, err := db.DB()
if err != nil {
panic("failed to get sql.DB: " + err.Error())
}
sqlDB.SetMaxOpenConns(1)
if err := db.AutoMigrate(&Task{}, &User{}, &Token{}, &Log{}, &Channel{}); err != nil {
panic("failed to migrate: " + err.Error())
}
os.Exit(m.Run())
}
func truncateTables(t *testing.T) {
t.Helper()
t.Cleanup(func() {
DB.Exec("DELETE FROM tasks")
DB.Exec("DELETE FROM users")
DB.Exec("DELETE FROM tokens")
DB.Exec("DELETE FROM logs")
DB.Exec("DELETE FROM channels")
})
}
func insertTask(t *testing.T, task *Task) {
t.Helper()
task.CreatedAt = time.Now().Unix()
task.UpdatedAt = time.Now().Unix()
require.NoError(t, DB.Create(task).Error)
}
// ---------------------------------------------------------------------------
// Snapshot / Equal — pure logic tests (no DB)
// ---------------------------------------------------------------------------
func TestSnapshotEqual_Same(t *testing.T) {
s := taskSnapshot{
Status: TaskStatusInProgress,
Progress: "50%",
StartTime: 1000,
FinishTime: 0,
FailReason: "",
ResultURL: "",
Data: json.RawMessage(`{"key":"value"}`),
}
assert.True(t, s.Equal(s))
}
func TestSnapshotEqual_DifferentStatus(t *testing.T) {
a := taskSnapshot{Status: TaskStatusInProgress, Data: json.RawMessage(`{}`)}
b := taskSnapshot{Status: TaskStatusSuccess, Data: json.RawMessage(`{}`)}
assert.False(t, a.Equal(b))
}
func TestSnapshotEqual_DifferentProgress(t *testing.T) {
a := taskSnapshot{Status: TaskStatusInProgress, Progress: "30%", Data: json.RawMessage(`{}`)}
b := taskSnapshot{Status: TaskStatusInProgress, Progress: "60%", Data: json.RawMessage(`{}`)}
assert.False(t, a.Equal(b))
}
func TestSnapshotEqual_DifferentData(t *testing.T) {
a := taskSnapshot{Status: TaskStatusInProgress, Data: json.RawMessage(`{"a":1}`)}
b := taskSnapshot{Status: TaskStatusInProgress, Data: json.RawMessage(`{"a":2}`)}
assert.False(t, a.Equal(b))
}
func TestSnapshotEqual_NilVsEmpty(t *testing.T) {
a := taskSnapshot{Status: TaskStatusInProgress, Data: nil}
b := taskSnapshot{Status: TaskStatusInProgress, Data: json.RawMessage{}}
// bytes.Equal(nil, []byte{}) == true
assert.True(t, a.Equal(b))
}
func TestSnapshot_Roundtrip(t *testing.T) {
task := &Task{
Status: TaskStatusInProgress,
Progress: "42%",
StartTime: 1234,
FinishTime: 5678,
FailReason: "timeout",
PrivateData: TaskPrivateData{
ResultURL: "https://example.com/result.mp4",
},
Data: json.RawMessage(`{"model":"test-model"}`),
}
snap := task.Snapshot()
assert.Equal(t, task.Status, snap.Status)
assert.Equal(t, task.Progress, snap.Progress)
assert.Equal(t, task.StartTime, snap.StartTime)
assert.Equal(t, task.FinishTime, snap.FinishTime)
assert.Equal(t, task.FailReason, snap.FailReason)
assert.Equal(t, task.PrivateData.ResultURL, snap.ResultURL)
assert.JSONEq(t, string(task.Data), string(snap.Data))
}
// ---------------------------------------------------------------------------
// UpdateWithStatus CAS — DB integration tests
// ---------------------------------------------------------------------------
func TestUpdateWithStatus_Win(t *testing.T) {
truncateTables(t)
task := &Task{
TaskID: "task_cas_win",
Status: TaskStatusInProgress,
Progress: "50%",
Data: json.RawMessage(`{}`),
}
insertTask(t, task)
task.Status = TaskStatusSuccess
task.Progress = "100%"
won, err := task.UpdateWithStatus(TaskStatusInProgress)
require.NoError(t, err)
assert.True(t, won)
var reloaded Task
require.NoError(t, DB.First(&reloaded, task.ID).Error)
assert.EqualValues(t, TaskStatusSuccess, reloaded.Status)
assert.Equal(t, "100%", reloaded.Progress)
}
func TestUpdateWithStatus_Lose(t *testing.T) {
truncateTables(t)
task := &Task{
TaskID: "task_cas_lose",
Status: TaskStatusFailure,
Data: json.RawMessage(`{}`),
}
insertTask(t, task)
task.Status = TaskStatusSuccess
won, err := task.UpdateWithStatus(TaskStatusInProgress) // wrong fromStatus
require.NoError(t, err)
assert.False(t, won)
var reloaded Task
require.NoError(t, DB.First(&reloaded, task.ID).Error)
assert.EqualValues(t, TaskStatusFailure, reloaded.Status) // unchanged
}
func TestUpdateWithStatus_ConcurrentWinner(t *testing.T) {
truncateTables(t)
task := &Task{
TaskID: "task_cas_race",
Status: TaskStatusInProgress,
Quota: 1000,
Data: json.RawMessage(`{}`),
}
insertTask(t, task)
const goroutines = 5
wins := make([]bool, goroutines)
var wg sync.WaitGroup
wg.Add(goroutines)
for i := 0; i < goroutines; i++ {
go func(idx int) {
defer wg.Done()
t := &Task{}
*t = Task{
ID: task.ID,
TaskID: task.TaskID,
Status: TaskStatusSuccess,
Progress: "100%",
Quota: task.Quota,
Data: json.RawMessage(`{}`),
}
t.CreatedAt = task.CreatedAt
t.UpdatedAt = time.Now().Unix()
won, err := t.UpdateWithStatus(TaskStatusInProgress)
if err == nil {
wins[idx] = won
}
}(i)
}
wg.Wait()
winCount := 0
for _, w := range wins {
if w {
winCount++
}
}
assert.Equal(t, 1, winCount, "exactly one goroutine should win the CAS")
}
model/token.go
......@@ -360,7 +360,7 @@ func DeleteTokenById(id int, userId int) (err error) {
return token.Delete()
}
func IncreaseTokenQuota(id int, key string, quota int) (err error) {
func IncreaseTokenQuota(tokenId int, key string, quota int) (err error) {
if quota < 0 {
return errors.New("quota 不能为负数！")
}
......@@ -373,10 +373,10 @@ func IncreaseTokenQuota(id int, key string, quota int) (err error) {
})
}
if common.BatchUpdateEnabled {
addNewRecord(BatchUpdateTypeTokenQuota, id, quota)
addNewRecord(BatchUpdateTypeTokenQuota, tokenId, quota)
return nil
}
return increaseTokenQuota(id, quota)
return increaseTokenQuota(tokenId, quota)
}
func increaseTokenQuota(id int, quota int) (err error) {
......
model/user.go
package model
import (
"database/sql"
"encoding/json"
"errors"
"fmt"
......@@ -15,6 +16,8 @@ import (
"gorm.io/gorm"
)
const UserNameMaxLength = 20
// User if you add sensitive fields, don't forget to clean them in setupLogin function.
// Otherwise, the sensitive information will be saved on local storage in plain text!
type User struct {
......@@ -536,6 +539,37 @@ func (user *User) Edit(updatePassword bool) error {
return updateUserCache(*user)
}
func (user *User) ClearBinding(bindingType string) error {
if user.Id == 0 {
return errors.New("user id is empty")
}
bindingColumnMap := map[string]string{
"email": "email",
"github": "github_id",
"discord": "discord_id",
"oidc": "oidc_id",
"wechat": "wechat_id",
"telegram": "telegram_id",
"linuxdo": "linux_do_id",
}
column, ok := bindingColumnMap[bindingType]
if !ok {
return errors.New("invalid binding type")
}
if err := DB.Model(&User{}).Where("id = ?", user.Id).Update(column, "").Error; err != nil {
return err
}
if err := DB.Where("id = ?", user.Id).First(user).Error; err != nil {
return err
}
return updateUserCache(*user)
}
func (user *User) Delete() error {
if user.Id == 0 {
return errors.New("id 为空！")
......@@ -820,10 +854,17 @@ func GetUserSetting(id int, fromDB bool) (settingMap dto.UserSetting, err error)
// Don't return error - fall through to DB
}
fromDB = true
err = DB.Model(&User{}).Where("id = ?", id).Select("setting").Find(&setting).Error
// can be nil setting
var safeSetting sql.NullString
err = DB.Model(&User{}).Where("id = ?", id).Select("setting").Find(&safeSetting).Error
if err != nil {
return settingMap, err
}
if safeSetting.Valid {
setting = safeSetting.String
} else {
setting = ""
}
userBase := &UserBase{
Setting: setting,
}
......
oauth/types.go
......@@ -57,3 +57,12 @@ func NewOAuthErrorWithRaw(msgKey string, params map[string]any, rawError string)
RawError: rawError,
}
}
// AccessDeniedError is a direct user-facing access denial message.
type AccessDeniedError struct {
Message string
}
func (e *AccessDeniedError) Error() string {
return e.Message
}
relay/channel/adapter.go
......@@ -36,6 +36,32 @@ type TaskAdaptor interface {
ValidateRequestAndSetAction(c *gin.Context, info *relaycommon.RelayInfo) *dto.TaskError
// ── Billing ──────────────────────────────────────────────────────
// EstimateBilling returns OtherRatios for pre-charge based on user request.
// Called after ValidateRequestAndSetAction, before price calculation.
// Adaptors should extract duration, resolution, etc. from the parsed request
// and return them as ratio multipliers (e.g. {"seconds": 5, "size": 1.666}).
// Return nil to use the base model price without extra ratios.
EstimateBilling(c *gin.Context, info *relaycommon.RelayInfo) map[string]float64
// AdjustBillingOnSubmit returns adjusted OtherRatios from the upstream
// submit response. Called after a successful DoResponse.
// If the upstream returned actual parameters that differ from the estimate
// (e.g. actual seconds), return updated ratios so the caller can recalculate
// the quota and settle the delta with the pre-charge.
// Return nil if no adjustment is needed.
AdjustBillingOnSubmit(info *relaycommon.RelayInfo, taskData []byte) map[string]float64
// AdjustBillingOnComplete returns the actual quota when a task reaches a
// terminal state (success/failure) during polling.
// Called by the polling loop after ParseTaskResult.
// Return a positive value to trigger delta settlement (supplement / refund).
// Return 0 to keep the pre-charged amount unchanged.
AdjustBillingOnComplete(task *model.Task, taskResult *relaycommon.TaskInfo) int
// ── Request / Response ───────────────────────────────────────────
BuildRequestURL(info *relaycommon.RelayInfo) (string, error)
BuildRequestHeader(c *gin.Context, req *http.Request, info *relaycommon.RelayInfo) error
BuildRequestBody(c *gin.Context, info *relaycommon.RelayInfo) (io.Reader, error)
......@@ -46,9 +72,9 @@ type TaskAdaptor interface {
GetModelList() []string
GetChannelName() string
// FetchTask
FetchTask(baseUrl, key string, body map[string]any, proxy string) (*http.Response, error)
// ── Polling ──────────────────────────────────────────────────────
FetchTask(baseUrl, key string, body map[string]any, proxy string) (*http.Response, error)
ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error)
}
......
relay/channel/api_request.go
......@@ -61,8 +61,9 @@ var passthroughSkipHeaderNamesLower = map[string]struct{}{
"cookie": {},
// Additional headers that should not be forwarded by name-matching passthrough rules.
"host": {},
"content-length": {},
"host": {},
"content-length": {},
"accept-encoding": {},
// Do not passthrough credentials by wildcard/regex.
"authorization": {},
......
relay/channel/api_request_test.go
......@@ -110,3 +110,30 @@ func TestProcessHeaderOverride_RuntimeOverrideHasPriority(t *testing.T) {
_, ok := headers["X-Legacy"]
require.False(t, ok)
}
func TestProcessHeaderOverride_PassthroughSkipsAcceptEncoding(t *testing.T) {
t.Parallel()
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(recorder)
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
ctx.Request.Header.Set("X-Trace-Id", "trace-123")
ctx.Request.Header.Set("Accept-Encoding", "gzip")
info := &relaycommon.RelayInfo{
IsChannelTest: false,
ChannelMeta: &relaycommon.ChannelMeta{
HeadersOverride: map[string]any{
"*": "",
},
},
}
headers, err := processHeaderOverride(info, ctx)
require.NoError(t, err)
require.Equal(t, "trace-123", headers["X-Trace-Id"])
_, hasAcceptEncoding := headers["Accept-Encoding"]
require.False(t, hasAcceptEncoding)
}
relay/channel/gemini/relay-gemini-native.go
......@@ -42,22 +42,7 @@ func GeminiTextGenerationHandler(c *gin.Context, info *relaycommon.RelayInfo, re
}
// 计算使用量（基于 UsageMetadata）
usage := dto.Usage{
PromptTokens: geminiResponse.UsageMetadata.PromptTokenCount,
CompletionTokens: geminiResponse.UsageMetadata.CandidatesTokenCount + geminiResponse.UsageMetadata.ThoughtsTokenCount,
TotalTokens: geminiResponse.UsageMetadata.TotalTokenCount,
}
usage.CompletionTokenDetails.ReasoningTokens = geminiResponse.UsageMetadata.ThoughtsTokenCount
usage.PromptTokensDetails.CachedTokens = geminiResponse.UsageMetadata.CachedContentTokenCount
for _, detail := range geminiResponse.UsageMetadata.PromptTokensDetails {
if detail.Modality == "AUDIO" {
usage.PromptTokensDetails.AudioTokens = detail.TokenCount
} else if detail.Modality == "TEXT" {
usage.PromptTokensDetails.TextTokens = detail.TokenCount
}
}
usage := buildUsageFromGeminiMetadata(geminiResponse.UsageMetadata, info.GetEstimatePromptTokens())
service.IOCopyBytesGracefully(c, resp, responseBody)
......
relay/channel/gemini/relay-gemini.go
......@@ -1032,6 +1032,46 @@ func getResponseToolCall(item *dto.GeminiPart) *dto.ToolCallResponse {
}
}
func buildUsageFromGeminiMetadata(metadata dto.GeminiUsageMetadata, fallbackPromptTokens int) dto.Usage {
promptTokens := metadata.PromptTokenCount + metadata.ToolUsePromptTokenCount
if promptTokens <= 0 && fallbackPromptTokens > 0 {
promptTokens = fallbackPromptTokens
}
usage := dto.Usage{
PromptTokens: promptTokens,
CompletionTokens: metadata.CandidatesTokenCount + metadata.ThoughtsTokenCount,
TotalTokens: metadata.TotalTokenCount,
}
usage.CompletionTokenDetails.ReasoningTokens = metadata.ThoughtsTokenCount
usage.PromptTokensDetails.CachedTokens = metadata.CachedContentTokenCount
for _, detail := range metadata.PromptTokensDetails {
if detail.Modality == "AUDIO" {
usage.PromptTokensDetails.AudioTokens += detail.TokenCount
} else if detail.Modality == "TEXT" {
usage.PromptTokensDetails.TextTokens += detail.TokenCount
}
}
for _, detail := range metadata.ToolUsePromptTokensDetails {
if detail.Modality == "AUDIO" {
usage.PromptTokensDetails.AudioTokens += detail.TokenCount
} else if detail.Modality == "TEXT" {
usage.PromptTokensDetails.TextTokens += detail.TokenCount
}
}
if usage.TotalTokens > 0 && usage.CompletionTokens <= 0 {
usage.CompletionTokens = usage.TotalTokens - usage.PromptTokens
}
if usage.PromptTokens > 0 && usage.PromptTokensDetails.TextTokens == 0 && usage.PromptTokensDetails.AudioTokens == 0 {
usage.PromptTokensDetails.TextTokens = usage.PromptTokens
}
return usage
}
func responseGeminiChat2OpenAI(c *gin.Context, response *dto.GeminiChatResponse) *dto.OpenAITextResponse {
fullTextResponse := dto.OpenAITextResponse{
Id: helper.GetResponseID(c),
......@@ -1272,18 +1312,8 @@ func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http
// 更新使用量统计
if geminiResponse.UsageMetadata.TotalTokenCount != 0 {
usage.PromptTokens = geminiResponse.UsageMetadata.PromptTokenCount
usage.CompletionTokens = geminiResponse.UsageMetadata.CandidatesTokenCount + geminiResponse.UsageMetadata.ThoughtsTokenCount
usage.CompletionTokenDetails.ReasoningTokens = geminiResponse.UsageMetadata.ThoughtsTokenCount
usage.TotalTokens = geminiResponse.UsageMetadata.TotalTokenCount
usage.PromptTokensDetails.CachedTokens = geminiResponse.UsageMetadata.CachedContentTokenCount
for _, detail := range geminiResponse.UsageMetadata.PromptTokensDetails {
if detail.Modality == "AUDIO" {
usage.PromptTokensDetails.AudioTokens = detail.TokenCount
} else if detail.Modality == "TEXT" {
usage.PromptTokensDetails.TextTokens = detail.TokenCount
}
}
mappedUsage := buildUsageFromGeminiMetadata(geminiResponse.UsageMetadata, info.GetEstimatePromptTokens())
*usage = mappedUsage
}
return callback(data, &geminiResponse)
......@@ -1295,11 +1325,6 @@ func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http
}
}
usage.PromptTokensDetails.TextTokens = usage.PromptTokens
if usage.TotalTokens > 0 {
usage.CompletionTokens = usage.TotalTokens - usage.PromptTokens
}
if usage.CompletionTokens <= 0 {
if info.ReceivedResponseCount > 0 {
usage = service.ResponseText2Usage(c, responseText.String(), info.UpstreamModelName, info.GetEstimatePromptTokens())
......@@ -1416,21 +1441,7 @@ func GeminiChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.R
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
}
if len(geminiResponse.Candidates) == 0 {
usage := dto.Usage{
PromptTokens: geminiResponse.UsageMetadata.PromptTokenCount,
}
usage.CompletionTokenDetails.ReasoningTokens = geminiResponse.UsageMetadata.ThoughtsTokenCount
usage.PromptTokensDetails.CachedTokens = geminiResponse.UsageMetadata.CachedContentTokenCount
for _, detail := range geminiResponse.UsageMetadata.PromptTokensDetails {
if detail.Modality == "AUDIO" {
usage.PromptTokensDetails.AudioTokens = detail.TokenCount
} else if detail.Modality == "TEXT" {
usage.PromptTokensDetails.TextTokens = detail.TokenCount
}
}
if usage.PromptTokens <= 0 {
usage.PromptTokens = info.GetEstimatePromptTokens()
}
usage := buildUsageFromGeminiMetadata(geminiResponse.UsageMetadata, info.GetEstimatePromptTokens())
var newAPIError *types.NewAPIError
if geminiResponse.PromptFeedback != nil && geminiResponse.PromptFeedback.BlockReason != nil {
......@@ -1466,23 +1477,7 @@ func GeminiChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.R
}
fullTextResponse := responseGeminiChat2OpenAI(c, &geminiResponse)
fullTextResponse.Model = info.UpstreamModelName
usage := dto.Usage{
PromptTokens: geminiResponse.UsageMetadata.PromptTokenCount,
CompletionTokens: geminiResponse.UsageMetadata.CandidatesTokenCount,
TotalTokens: geminiResponse.UsageMetadata.TotalTokenCount,
}
usage.CompletionTokenDetails.ReasoningTokens = geminiResponse.UsageMetadata.ThoughtsTokenCount
usage.PromptTokensDetails.CachedTokens = geminiResponse.UsageMetadata.CachedContentTokenCount
usage.CompletionTokens = usage.TotalTokens - usage.PromptTokens
for _, detail := range geminiResponse.UsageMetadata.PromptTokensDetails {
if detail.Modality == "AUDIO" {
usage.PromptTokensDetails.AudioTokens = detail.TokenCount
} else if detail.Modality == "TEXT" {
usage.PromptTokensDetails.TextTokens = detail.TokenCount
}
}
usage := buildUsageFromGeminiMetadata(geminiResponse.UsageMetadata, info.GetEstimatePromptTokens())
fullTextResponse.Usage = usage
......
relay/channel/gemini/relay_gemini_usage_test.go 0 → 100644
package gemini
import (
"bytes"
"io"
"net/http"
"net/http/httptest"
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/types"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func TestGeminiChatHandlerCompletionTokensExcludeToolUsePromptTokens(t *testing.T) {
t.Parallel()
gin.SetMode(gin.TestMode)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
info := &relaycommon.RelayInfo{
RelayFormat: types.RelayFormatGemini,
OriginModelName: "gemini-3-flash-preview",
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: "gemini-3-flash-preview",
},
}
payload := dto.GeminiChatResponse{
Candidates: []dto.GeminiChatCandidate{
{
Content: dto.GeminiChatContent{
Role: "model",
Parts: []dto.GeminiPart{
{Text: "ok"},
},
},
},
},
UsageMetadata: dto.GeminiUsageMetadata{
PromptTokenCount: 151,
ToolUsePromptTokenCount: 18329,
CandidatesTokenCount: 1089,
ThoughtsTokenCount: 1120,
TotalTokenCount: 20689,
},
}
body, err := common.Marshal(payload)
require.NoError(t, err)
resp := &http.Response{
Body: io.NopCloser(bytes.NewReader(body)),
}
usage, newAPIError := GeminiChatHandler(c, info, resp)
require.Nil(t, newAPIError)
require.NotNil(t, usage)
require.Equal(t, 18480, usage.PromptTokens)
require.Equal(t, 2209, usage.CompletionTokens)
require.Equal(t, 20689, usage.TotalTokens)
require.Equal(t, 1120, usage.CompletionTokenDetails.ReasoningTokens)
}
func TestGeminiStreamHandlerCompletionTokensExcludeToolUsePromptTokens(t *testing.T) {
gin.SetMode(gin.TestMode)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
oldStreamingTimeout := constant.StreamingTimeout
constant.StreamingTimeout = 300
t.Cleanup(func() {
constant.StreamingTimeout = oldStreamingTimeout
})
info := &relaycommon.RelayInfo{
OriginModelName: "gemini-3-flash-preview",
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: "gemini-3-flash-preview",
},
}
chunk := dto.GeminiChatResponse{
Candidates: []dto.GeminiChatCandidate{
{
Content: dto.GeminiChatContent{
Role: "model",
Parts: []dto.GeminiPart{
{Text: "partial"},
},
},
},
},
UsageMetadata: dto.GeminiUsageMetadata{
PromptTokenCount: 151,
ToolUsePromptTokenCount: 18329,
CandidatesTokenCount: 1089,
ThoughtsTokenCount: 1120,
TotalTokenCount: 20689,
},
}
chunkData, err := common.Marshal(chunk)
require.NoError(t, err)
streamBody := []byte("data: " + string(chunkData) + "\n" + "data: [DONE]\n")
resp := &http.Response{
Body: io.NopCloser(bytes.NewReader(streamBody)),
}
usage, newAPIError := geminiStreamHandler(c, info, resp, func(_ string, _ *dto.GeminiChatResponse) bool {
return true
})
require.Nil(t, newAPIError)
require.NotNil(t, usage)
require.Equal(t, 18480, usage.PromptTokens)
require.Equal(t, 2209, usage.CompletionTokens)
require.Equal(t, 20689, usage.TotalTokens)
require.Equal(t, 1120, usage.CompletionTokenDetails.ReasoningTokens)
}
func TestGeminiTextGenerationHandlerPromptTokensIncludeToolUsePromptTokens(t *testing.T) {
t.Parallel()
gin.SetMode(gin.TestMode)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1beta/models/gemini-3-flash-preview:generateContent", nil)
info := &relaycommon.RelayInfo{
OriginModelName: "gemini-3-flash-preview",
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: "gemini-3-flash-preview",
},
}
payload := dto.GeminiChatResponse{
Candidates: []dto.GeminiChatCandidate{
{
Content: dto.GeminiChatContent{
Role: "model",
Parts: []dto.GeminiPart{
{Text: "ok"},
},
},
},
},
UsageMetadata: dto.GeminiUsageMetadata{
PromptTokenCount: 151,
ToolUsePromptTokenCount: 18329,
CandidatesTokenCount: 1089,
ThoughtsTokenCount: 1120,
TotalTokenCount: 20689,
},
}
body, err := common.Marshal(payload)
require.NoError(t, err)
resp := &http.Response{
Body: io.NopCloser(bytes.NewReader(body)),
}
usage, newAPIError := GeminiTextGenerationHandler(c, info, resp)
require.Nil(t, newAPIError)
require.NotNil(t, usage)
require.Equal(t, 18480, usage.PromptTokens)
require.Equal(t, 2209, usage.CompletionTokens)
require.Equal(t, 20689, usage.TotalTokens)
require.Equal(t, 1120, usage.CompletionTokenDetails.ReasoningTokens)
}
func TestGeminiChatHandlerUsesEstimatedPromptTokensWhenUsagePromptMissing(t *testing.T) {
t.Parallel()
gin.SetMode(gin.TestMode)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
info := &relaycommon.RelayInfo{
RelayFormat: types.RelayFormatGemini,
OriginModelName: "gemini-3-flash-preview",
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: "gemini-3-flash-preview",
},
}
info.SetEstimatePromptTokens(20)
payload := dto.GeminiChatResponse{
Candidates: []dto.GeminiChatCandidate{
{
Content: dto.GeminiChatContent{
Role: "model",
Parts: []dto.GeminiPart{
{Text: "ok"},
},
},
},
},
UsageMetadata: dto.GeminiUsageMetadata{
PromptTokenCount: 0,
ToolUsePromptTokenCount: 0,
CandidatesTokenCount: 90,
ThoughtsTokenCount: 10,
TotalTokenCount: 110,
},
}
body, err := common.Marshal(payload)
require.NoError(t, err)
resp := &http.Response{
Body: io.NopCloser(bytes.NewReader(body)),
}
usage, newAPIError := GeminiChatHandler(c, info, resp)
require.Nil(t, newAPIError)
require.NotNil(t, usage)
require.Equal(t, 20, usage.PromptTokens)
require.Equal(t, 100, usage.CompletionTokens)
require.Equal(t, 110, usage.TotalTokens)
}
func TestGeminiStreamHandlerUsesEstimatedPromptTokensWhenUsagePromptMissing(t *testing.T) {
gin.SetMode(gin.TestMode)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
oldStreamingTimeout := constant.StreamingTimeout
constant.StreamingTimeout = 300
t.Cleanup(func() {
constant.StreamingTimeout = oldStreamingTimeout
})
info := &relaycommon.RelayInfo{
OriginModelName: "gemini-3-flash-preview",
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: "gemini-3-flash-preview",
},
}
info.SetEstimatePromptTokens(20)
chunk := dto.GeminiChatResponse{
Candidates: []dto.GeminiChatCandidate{
{
Content: dto.GeminiChatContent{
Role: "model",
Parts: []dto.GeminiPart{
{Text: "partial"},
},
},
},
},
UsageMetadata: dto.GeminiUsageMetadata{
PromptTokenCount: 0,
ToolUsePromptTokenCount: 0,
CandidatesTokenCount: 90,
ThoughtsTokenCount: 10,
TotalTokenCount: 110,
},
}
chunkData, err := common.Marshal(chunk)
require.NoError(t, err)
streamBody := []byte("data: " + string(chunkData) + "\n" + "data: [DONE]\n")
resp := &http.Response{
Body: io.NopCloser(bytes.NewReader(streamBody)),
}
usage, newAPIError := geminiStreamHandler(c, info, resp, func(_ string, _ *dto.GeminiChatResponse) bool {
return true
})
require.Nil(t, newAPIError)
require.NotNil(t, usage)
require.Equal(t, 20, usage.PromptTokens)
require.Equal(t, 100, usage.CompletionTokens)
require.Equal(t, 110, usage.TotalTokens)
}
func TestGeminiTextGenerationHandlerUsesEstimatedPromptTokensWhenUsagePromptMissing(t *testing.T) {
t.Parallel()
gin.SetMode(gin.TestMode)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1beta/models/gemini-3-flash-preview:generateContent", nil)
info := &relaycommon.RelayInfo{
OriginModelName: "gemini-3-flash-preview",
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: "gemini-3-flash-preview",
},
}
info.SetEstimatePromptTokens(20)
payload := dto.GeminiChatResponse{
Candidates: []dto.GeminiChatCandidate{
{
Content: dto.GeminiChatContent{
Role: "model",
Parts: []dto.GeminiPart{
{Text: "ok"},
},
},
},
},
UsageMetadata: dto.GeminiUsageMetadata{
PromptTokenCount: 0,
ToolUsePromptTokenCount: 0,
CandidatesTokenCount: 90,
ThoughtsTokenCount: 10,
TotalTokenCount: 110,
},
}
body, err := common.Marshal(payload)
require.NoError(t, err)
resp := &http.Response{
Body: io.NopCloser(bytes.NewReader(body)),
}
usage, newAPIError := GeminiTextGenerationHandler(c, info, resp)
require.Nil(t, newAPIError)
require.NotNil(t, usage)
require.Equal(t, 20, usage.PromptTokens)
require.Equal(t, 100, usage.CompletionTokens)
require.Equal(t, 110, usage.TotalTokens)
}
relay/channel/minimax/adaptor.go
......@@ -10,6 +10,7 @@ import (
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/relay/channel"
"github.com/QuantumNous/new-api/relay/channel/claude"
"github.com/QuantumNous/new-api/relay/channel/openai"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/constant"
......@@ -26,7 +27,8 @@ func (a *Adaptor) ConvertGeminiRequest(*gin.Context, *relaycommon.RelayInfo, *dt
}
func (a *Adaptor) ConvertClaudeRequest(c *gin.Context, info *relaycommon.RelayInfo, req *dto.ClaudeRequest) (any, error) {
return nil, errors.New("not implemented")
adaptor := claude.Adaptor{}
return adaptor.ConvertClaudeRequest(c, info, req)
}
func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.AudioRequest) (io.Reader, error) {
......@@ -119,8 +121,14 @@ func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycom
return handleTTSResponse(c, resp, info)
}
adaptor := openai.Adaptor{}
return adaptor.DoResponse(c, resp, info)
switch info.RelayFormat {
case types.RelayFormatClaude:
adaptor := claude.Adaptor{}
return adaptor.DoResponse(c, resp, info)
default:
adaptor := openai.Adaptor{}
return adaptor.DoResponse(c, resp, info)
}
}
func (a *Adaptor) GetModelList() []string {
......
relay/channel/minimax/relay-minimax.go
......@@ -6,6 +6,7 @@ import (
channelconstant "github.com/QuantumNous/new-api/constant"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/types"
)
func GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
......@@ -13,13 +14,17 @@ func GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
if baseUrl == "" {
baseUrl = channelconstant.ChannelBaseURLs[channelconstant.ChannelTypeMiniMax]
}
switch info.RelayMode {
case constant.RelayModeChatCompletions:
return fmt.Sprintf("%s/v1/text/chatcompletion_v2", baseUrl), nil
case constant.RelayModeAudioSpeech:
return fmt.Sprintf("%s/v1/t2a_v2", baseUrl), nil
switch info.RelayFormat {
case types.RelayFormatClaude:
return fmt.Sprintf("%s/anthropic/v1/messages", info.ChannelBaseUrl), nil
default:
return "", fmt.Errorf("unsupported relay mode: %d", info.RelayMode)
switch info.RelayMode {
case constant.RelayModeChatCompletions:
return fmt.Sprintf("%s/v1/text/chatcompletion_v2", baseUrl), nil
case constant.RelayModeAudioSpeech:
return fmt.Sprintf("%s/v1/t2a_v2", baseUrl), nil
default:
return "", fmt.Errorf("unsupported relay mode: %d", info.RelayMode)
}
}
}
relay/channel/task/ali/adaptor.go
......@@ -13,6 +13,7 @@ import (
"github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/relay/channel"
"github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/service"
"github.com/samber/lo"
......@@ -108,10 +109,10 @@ type AliMetadata struct {
// ============================
type TaskAdaptor struct {
taskcommon.BaseBilling
ChannelType int
apiKey string
baseURL string
aliReq *AliVideoRequest
}
func (a *TaskAdaptor) Init(info *relaycommon.RelayInfo) {
......@@ -121,17 +122,7 @@ func (a *TaskAdaptor) Init(info *relaycommon.RelayInfo) {
}
func (a *TaskAdaptor) ValidateRequestAndSetAction(c *gin.Context, info *relaycommon.RelayInfo) (taskErr *dto.TaskError) {
// 阿里通义万相支持 JSON 格式，不使用 multipart
var taskReq relaycommon.TaskSubmitReq
if err := common.UnmarshalBodyReusable(c, &taskReq); err != nil {
return service.TaskErrorWrapper(err, "unmarshal_task_request_failed", http.StatusBadRequest)
}
aliReq, err := a.convertToAliRequest(info, taskReq)
if err != nil {
return service.TaskErrorWrapper(err, "convert_to_ali_request_failed", http.StatusInternalServerError)
}
a.aliReq = aliReq
logger.LogJson(c, "ali video request body", aliReq)
// ValidateMultipartDirect 负责解析并将原始 TaskSubmitReq 存入 context
return relaycommon.ValidateMultipartDirect(c, info)
}
......@@ -148,11 +139,21 @@ func (a *TaskAdaptor) BuildRequestHeader(c *gin.Context, req *http.Request, info
}
func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayInfo) (io.Reader, error) {
bodyBytes, err := common.Marshal(a.aliReq)
taskReq, err := relaycommon.GetTaskRequest(c)
if err != nil {
return nil, errors.Wrap(err, "marshal_ali_request_failed")
return nil, errors.Wrap(err, "get_task_request_failed")
}
aliReq, err := a.convertToAliRequest(info, taskReq)
if err != nil {
return nil, errors.Wrap(err, "convert_to_ali_request_failed")
}
logger.LogJson(c, "ali video request body", aliReq)
bodyBytes, err := common.Marshal(aliReq)
if err != nil {
return nil, errors.Wrap(err, "marshal_ali_request_failed")
}
return bytes.NewReader(bodyBytes), nil
}
......@@ -252,8 +253,12 @@ func ProcessAliOtherRatios(aliReq *AliVideoRequest) (map[string]float64, error)
}
func (a *TaskAdaptor) convertToAliRequest(info *relaycommon.RelayInfo, req relaycommon.TaskSubmitReq) (*AliVideoRequest, error) {
upstreamModel := req.Model
if info.IsModelMapped {
upstreamModel = info.UpstreamModelName
}
aliReq := &AliVideoRequest{
Model: req.Model,
Model: upstreamModel,
Input: AliVideoInput{
Prompt: req.Prompt,
ImgURL: req.InputReference,
......@@ -331,23 +336,37 @@ func (a *TaskAdaptor) convertToAliRequest(info *relaycommon.RelayInfo, req relay
}
}
if aliReq.Model != req.Model {
if aliReq.Model != upstreamModel {
return nil, errors.New("can't change model with metadata")
}
info.PriceData.OtherRatios = map[string]float64{
"seconds": float64(aliReq.Parameters.Duration),
return aliReq, nil
}
// EstimateBilling 根据用户请求参数计算 OtherRatios（时长、分辨率等）。
// 在 ValidateRequestAndSetAction 之后、价格计算之前调用。
func (a *TaskAdaptor) EstimateBilling(c *gin.Context, info *relaycommon.RelayInfo) map[string]float64 {
taskReq, err := relaycommon.GetTaskRequest(c)
if err != nil {
return nil
}
aliReq, err := a.convertToAliRequest(info, taskReq)
if err != nil {
return nil
}
otherRatios := map[string]float64{
"seconds": float64(aliReq.Parameters.Duration),
}
ratios, err := ProcessAliOtherRatios(aliReq)
if err != nil {
return nil, err
return otherRatios
}
for s, f := range ratios {
info.PriceData.OtherRatios[s] = f
for k, v := range ratios {
otherRatios[k] = v
}
return aliReq, nil
return otherRatios
}
// DoRequest delegates to common helper
......@@ -384,7 +403,8 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
// 转换为 OpenAI 格式响应
openAIResp := dto.NewOpenAIVideo()
openAIResp.ID = aliResp.Output.TaskID
openAIResp.ID = info.PublicTaskID
openAIResp.TaskID = info.PublicTaskID
openAIResp.Model = c.GetString("model")
if openAIResp.Model == "" && info != nil {
openAIResp.Model = info.OriginModelName
......
relay/channel/task/doubao/adaptor.go
......@@ -2,7 +2,6 @@ package doubao
import (
"bytes"
"encoding/json"
"fmt"
"io"
"net/http"
......@@ -14,6 +13,7 @@ import (
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/relay/channel"
taskcommon "github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/service"
......@@ -89,6 +89,7 @@ type responseTask struct {
// ============================
type TaskAdaptor struct {
taskcommon.BaseBilling
ChannelType int
apiKey string
baseURL string
......@@ -130,8 +131,12 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
if err != nil {
return nil, errors.Wrap(err, "convert request payload failed")
}
info.UpstreamModelName = body.Model
data, err := json.Marshal(body)
if info.IsModelMapped {
body.Model = info.UpstreamModelName
} else {
info.UpstreamModelName = body.Model
}
data, err := common.Marshal(body)
if err != nil {
return nil, err
}
......@@ -154,7 +159,7 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
// Parse Doubao response
var dResp responsePayload
if err := json.Unmarshal(responseBody, &dResp); err != nil {
if err := common.Unmarshal(responseBody, &dResp); err != nil {
taskErr = service.TaskErrorWrapper(errors.Wrapf(err, "body: %s", responseBody), "unmarshal_response_body_failed", http.StatusInternalServerError)
return
}
......@@ -165,8 +170,8 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
}
ov := dto.NewOpenAIVideo()
ov.ID = dResp.ID
ov.TaskID = dResp.ID
ov.ID = info.PublicTaskID
ov.TaskID = info.PublicTaskID
ov.CreatedAt = time.Now().Unix()
ov.Model = info.OriginModelName
......@@ -234,12 +239,7 @@ func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*
}
metadata := req.Metadata
medaBytes, err := json.Marshal(metadata)
if err != nil {
return nil, errors.Wrap(err, "metadata marshal metadata failed")
}
err = json.Unmarshal(medaBytes, &r)
if err != nil {
if err := taskcommon.UnmarshalMetadata(metadata, &r); err != nil {
return nil, errors.Wrap(err, "unmarshal metadata failed")
}
......@@ -248,7 +248,7 @@ func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*
func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) {
resTask := responseTask{}
if err := json.Unmarshal(respBody, &resTask); err != nil {
if err := common.Unmarshal(respBody, &resTask); err != nil {
return nil, errors.Wrap(err, "unmarshal task result failed")
}
......@@ -286,7 +286,7 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, error) {
var dResp responseTask
if err := json.Unmarshal(originTask.Data, &dResp); err != nil {
if err := common.Unmarshal(originTask.Data, &dResp); err != nil {
return nil, errors.Wrap(err, "unmarshal doubao task data failed")
}
......@@ -307,6 +307,5 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, erro
}
}
jsonData, _ := common.Marshal(openAIVideo)
return jsonData, nil
return common.Marshal(openAIVideo)
}
relay/channel/task/gemini/adaptor.go
......@@ -2,8 +2,6 @@ package gemini
import (
"bytes"
"encoding/base64"
"encoding/json"
"fmt"
"io"
"net/http"
......@@ -16,10 +14,10 @@ import (
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/relay/channel"
taskcommon "github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting/model_setting"
"github.com/QuantumNous/new-api/setting/system_setting"
"github.com/gin-gonic/gin"
"github.com/pkg/errors"
)
......@@ -87,6 +85,7 @@ type operationResponse struct {
// ============================
type TaskAdaptor struct {
taskcommon.BaseBilling
ChannelType int
apiKey string
baseURL string
......@@ -106,7 +105,7 @@ func (a *TaskAdaptor) ValidateRequestAndSetAction(c *gin.Context, info *relaycom
// BuildRequestURL constructs the upstream URL.
func (a *TaskAdaptor) BuildRequestURL(info *relaycommon.RelayInfo) (string, error) {
modelName := info.OriginModelName
modelName := info.UpstreamModelName
version := model_setting.GetGeminiVersionSetting(modelName)
return fmt.Sprintf(
......@@ -145,16 +144,11 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
}
metadata := req.Metadata
medaBytes, err := json.Marshal(metadata)
if err != nil {
return nil, errors.Wrap(err, "metadata marshal metadata failed")
}
err = json.Unmarshal(medaBytes, &body.Parameters)
if err != nil {
if err := taskcommon.UnmarshalMetadata(metadata, &body.Parameters); err != nil {
return nil, errors.Wrap(err, "unmarshal metadata failed")
}
data, err := json.Marshal(body)
data, err := common.Marshal(body)
if err != nil {
return nil, err
}
......@@ -175,16 +169,16 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
_ = resp.Body.Close()
var s submitResponse
if err := json.Unmarshal(responseBody, &s); err != nil {
if err := common.Unmarshal(responseBody, &s); err != nil {
return "", nil, service.TaskErrorWrapper(err, "unmarshal_response_failed", http.StatusInternalServerError)
}
if strings.TrimSpace(s.Name) == "" {
return "", nil, service.TaskErrorWrapper(fmt.Errorf("missing operation name"), "invalid_response", http.StatusInternalServerError)
}
taskID = encodeLocalTaskID(s.Name)
taskID = taskcommon.EncodeLocalTaskID(s.Name)
ov := dto.NewOpenAIVideo()
ov.ID = taskID
ov.TaskID = taskID
ov.ID = info.PublicTaskID
ov.TaskID = info.PublicTaskID
ov.CreatedAt = time.Now().Unix()
ov.Model = info.OriginModelName
c.JSON(http.StatusOK, ov)
......@@ -206,7 +200,7 @@ func (a *TaskAdaptor) FetchTask(baseUrl, key string, body map[string]any, proxy
return nil, fmt.Errorf("invalid task_id")
}
upstreamName, err := decodeLocalTaskID(taskID)
upstreamName, err := taskcommon.DecodeLocalTaskID(taskID)
if err != nil {
return nil, fmt.Errorf("decode task_id failed: %w", err)
}
......@@ -232,7 +226,7 @@ func (a *TaskAdaptor) FetchTask(baseUrl, key string, body map[string]any, proxy
func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) {
var op operationResponse
if err := json.Unmarshal(respBody, &op); err != nil {
if err := common.Unmarshal(respBody, &op); err != nil {
return nil, fmt.Errorf("unmarshal operation response failed: %w", err)
}
......@@ -254,9 +248,8 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
ti.Status = model.TaskStatusSuccess
ti.Progress = "100%"
taskID := encodeLocalTaskID(op.Name)
ti.TaskID = taskID
ti.Url = fmt.Sprintf("%s/v1/videos/%s/content", system_setting.ServerAddress, taskID)
ti.TaskID = taskcommon.EncodeLocalTaskID(op.Name)
// Url intentionally left empty — the caller constructs the proxy URL using the public task ID
// Extract URL from generateVideoResponse if available
if len(op.Response.GenerateVideoResponse.GeneratedSamples) > 0 {
......@@ -269,7 +262,10 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
}
func (a *TaskAdaptor) ConvertToOpenAIVideo(task *model.Task) ([]byte, error) {
upstreamName, err := decodeLocalTaskID(task.TaskID)
// Use GetUpstreamTaskID() to get the real upstream operation name for model extraction.
// task.TaskID is now a public task_xxxx ID, no longer a base64-encoded upstream name.
upstreamTaskID := task.GetUpstreamTaskID()
upstreamName, err := taskcommon.DecodeLocalTaskID(upstreamTaskID)
if err != nil {
upstreamName = ""
}
......@@ -297,18 +293,6 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(task *model.Task) ([]byte, error) {
// helpers
// ============================
func encodeLocalTaskID(name string) string {
return base64.RawURLEncoding.EncodeToString([]byte(name))
}
func decodeLocalTaskID(local string) (string, error) {
b, err := base64.RawURLEncoding.DecodeString(local)
if err != nil {
return "", err
}
return string(b), nil
}
var modelRe = regexp.MustCompile(`models/([^/]+)/operations/`)
func extractModelFromOperationName(name string) string {
......
relay/channel/task/hailuo/adaptor.go
......@@ -2,7 +2,6 @@ package hailuo
import (
"bytes"
"encoding/json"
"fmt"
"io"
"net/http"
......@@ -18,12 +17,14 @@ import (
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/relay/channel"
taskcommon "github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/service"
)
// https://platform.minimaxi.com/docs/api-reference/video-generation-intro
type TaskAdaptor struct {
taskcommon.BaseBilling
ChannelType int
apiKey string
baseURL string
......@@ -60,12 +61,12 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
return nil, fmt.Errorf("invalid request type in context")
}
body, err := a.convertToRequestPayload(&req)
body, err := a.convertToRequestPayload(&req, info)
if err != nil {
return nil, errors.Wrap(err, "convert request payload failed")
}
data, err := json.Marshal(body)
data, err := common.Marshal(body)
if err != nil {
return nil, err
}
......@@ -86,7 +87,7 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
_ = resp.Body.Close()
var hResp VideoResponse
if err := json.Unmarshal(responseBody, &hResp); err != nil {
if err := common.Unmarshal(responseBody, &hResp); err != nil {
taskErr = service.TaskErrorWrapper(errors.Wrapf(err, "body: %s", responseBody), "unmarshal_response_body_failed", http.StatusInternalServerError)
return
}
......@@ -101,8 +102,8 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
}
ov := dto.NewOpenAIVideo()
ov.ID = hResp.TaskID
ov.TaskID = hResp.TaskID
ov.ID = info.PublicTaskID
ov.TaskID = info.PublicTaskID
ov.CreatedAt = time.Now().Unix()
ov.Model = info.OriginModelName
......@@ -141,8 +142,8 @@ func (a *TaskAdaptor) GetChannelName() string {
return ChannelName
}
func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*VideoRequest, error) {
modelConfig := GetModelConfig(req.Model)
func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq, info *relaycommon.RelayInfo) (*VideoRequest, error) {
modelConfig := GetModelConfig(info.UpstreamModelName)
duration := DefaultDuration
if req.Duration > 0 {
duration = req.Duration
......@@ -153,7 +154,7 @@ func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*
}
videoRequest := &VideoRequest{
Model: req.Model,
Model: info.UpstreamModelName,
Prompt: req.Prompt,
Duration: &duration,
Resolution: resolution,
......@@ -182,7 +183,7 @@ func (a *TaskAdaptor) parseResolutionFromSize(size string, modelConfig ModelConf
func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) {
resTask := QueryTaskResponse{}
if err := json.Unmarshal(respBody, &resTask); err != nil {
if err := common.Unmarshal(respBody, &resTask); err != nil {
return nil, errors.Wrap(err, "unmarshal task result failed")
}
......@@ -224,7 +225,7 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, error) {
var hailuoResp QueryTaskResponse
if err := json.Unmarshal(originTask.Data, &hailuoResp); err != nil {
if err := common.Unmarshal(originTask.Data, &hailuoResp); err != nil {
return nil, errors.Wrap(err, "unmarshal hailuo task data failed")
}
......@@ -271,7 +272,7 @@ func (a *TaskAdaptor) buildVideoURL(_, fileID string) string {
}
var retrieveResp RetrieveFileResponse
if err := json.Unmarshal(responseBody, &retrieveResp); err != nil {
if err := common.Unmarshal(responseBody, &retrieveResp); err != nil {
return ""
}
......
relay/channel/task/jimeng/adaptor.go
......@@ -6,7 +6,6 @@ import (
"crypto/sha256"
"encoding/base64"
"encoding/hex"
"encoding/json"
"fmt"
"io"
"net/http"
......@@ -25,6 +24,7 @@ import (
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/relay/channel"
taskcommon "github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/service"
)
......@@ -77,6 +77,7 @@ const (
// ============================
type TaskAdaptor struct {
taskcommon.BaseBilling
ChannelType int
accessKey string
secretKey string
......@@ -164,11 +165,11 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
}
}
body, err := a.convertToRequestPayload(&req)
body, err := a.convertToRequestPayload(&req, info)
if err != nil {
return nil, errors.Wrap(err, "convert request payload failed")
}
data, err := json.Marshal(body)
data, err := common.Marshal(body)
if err != nil {
return nil, err
}
......@@ -191,7 +192,7 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
// Parse Jimeng response
var jResp responsePayload
if err := json.Unmarshal(responseBody, &jResp); err != nil {
if err := common.Unmarshal(responseBody, &jResp); err != nil {
taskErr = service.TaskErrorWrapper(errors.Wrapf(err, "body: %s", responseBody), "unmarshal_response_body_failed", http.StatusInternalServerError)
return
}
......@@ -202,8 +203,8 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
}
ov := dto.NewOpenAIVideo()
ov.ID = jResp.Data.TaskID
ov.TaskID = jResp.Data.TaskID
ov.ID = info.PublicTaskID
ov.TaskID = info.PublicTaskID
ov.CreatedAt = time.Now().Unix()
ov.Model = info.OriginModelName
c.JSON(http.StatusOK, ov)
......@@ -225,7 +226,7 @@ func (a *TaskAdaptor) FetchTask(baseUrl, key string, body map[string]any, proxy
"req_key": "jimeng_vgfm_t2v_l20", // This is fixed value from doc: https://www.volcengine.com/docs/85621/1544774
"task_id": taskID,
}
payloadBytes, err := json.Marshal(payload)
payloadBytes, err := common.Marshal(payload)
if err != nil {
return nil, errors.Wrap(err, "marshal fetch task payload failed")
}
......@@ -377,9 +378,9 @@ func hmacSHA256(key []byte, data []byte) []byte {
return h.Sum(nil)
}
func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*requestPayload, error) {
func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq, info *relaycommon.RelayInfo) (*requestPayload, error) {
r := requestPayload{
ReqKey: req.Model,
ReqKey: info.UpstreamModelName,
Prompt: req.Prompt,
}
......@@ -398,13 +399,7 @@ func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*
r.BinaryDataBase64 = req.Images
}
}
metadata := req.Metadata
medaBytes, err := json.Marshal(metadata)
if err != nil {
return nil, errors.Wrap(err, "metadata marshal metadata failed")
}
err = json.Unmarshal(medaBytes, &r)
if err != nil {
if err := taskcommon.UnmarshalMetadata(req.Metadata, &r); err != nil {
return nil, errors.Wrap(err, "unmarshal metadata failed")
}
......@@ -432,7 +427,7 @@ func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*
func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) {
resTask := responseTask{}
if err := json.Unmarshal(respBody, &resTask); err != nil {
if err := common.Unmarshal(respBody, &resTask); err != nil {
return nil, errors.Wrap(err, "unmarshal task result failed")
}
taskResult := relaycommon.TaskInfo{}
......@@ -458,7 +453,7 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, error) {
var jimengResp responseTask
if err := json.Unmarshal(originTask.Data, &jimengResp); err != nil {
if err := common.Unmarshal(originTask.Data, &jimengResp); err != nil {
return nil, errors.Wrap(err, "unmarshal jimeng task data failed")
}
......@@ -477,8 +472,7 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, erro
}
}
jsonData, _ := common.Marshal(openAIVideo)
return jsonData, nil
return common.Marshal(openAIVideo)
}
func isNewAPIRelay(apiKey string) bool {
......
relay/channel/task/kling/adaptor.go
......@@ -2,7 +2,6 @@ package kling
import (
"bytes"
"encoding/json"
"fmt"
"io"
"net/http"
......@@ -21,6 +20,7 @@ import (
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/relay/channel"
taskcommon "github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/service"
)
......@@ -97,6 +97,7 @@ type responsePayload struct {
// ============================
type TaskAdaptor struct {
taskcommon.BaseBilling
ChannelType int
apiKey string
baseURL string
......@@ -149,14 +150,14 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
}
req := v.(relaycommon.TaskSubmitReq)
body, err := a.convertToRequestPayload(&req)
body, err := a.convertToRequestPayload(&req, info)
if err != nil {
return nil, err
}
if body.Image == "" && body.ImageTail == "" {
c.Set("action", constant.TaskActionTextGenerate)
}
data, err := json.Marshal(body)
data, err := common.Marshal(body)
if err != nil {
return nil, err
}
......@@ -180,7 +181,7 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
}
var kResp responsePayload
err = json.Unmarshal(responseBody, &kResp)
err = common.Unmarshal(responseBody, &kResp)
if err != nil {
taskErr = service.TaskErrorWrapper(err, "unmarshal_response_failed", http.StatusInternalServerError)
return
......@@ -190,8 +191,8 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
return
}
ov := dto.NewOpenAIVideo()
ov.ID = kResp.Data.TaskId
ov.TaskID = kResp.Data.TaskId
ov.ID = info.PublicTaskID
ov.TaskID = info.PublicTaskID
ov.CreatedAt = time.Now().Unix()
ov.Model = info.OriginModelName
c.JSON(http.StatusOK, ov)
......@@ -247,15 +248,15 @@ func (a *TaskAdaptor) GetChannelName() string {
// helpers
// ============================
func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*requestPayload, error) {
func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq, info *relaycommon.RelayInfo) (*requestPayload, error) {
r := requestPayload{
Prompt: req.Prompt,
Image: req.Image,
Mode: defaultString(req.Mode, "std"),
Duration: fmt.Sprintf("%d", defaultInt(req.Duration, 5)),
Mode: taskcommon.DefaultString(req.Mode, "std"),
Duration: fmt.Sprintf("%d", taskcommon.DefaultInt(req.Duration, 5)),
AspectRatio: a.getAspectRatio(req.Size),
ModelName: req.Model,
Model: req.Model, // Keep consistent with model_name, double writing improves compatibility
ModelName: info.UpstreamModelName,
Model: info.UpstreamModelName,
CfgScale: 0.5,
StaticMask: "",
DynamicMasks: []DynamicMask{},
......@@ -265,14 +266,9 @@ func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*
}
if r.ModelName == "" {
r.ModelName = "kling-v1"
r.Model = "kling-v1"
}
metadata := req.Metadata
medaBytes, err := json.Marshal(metadata)
if err != nil {
return nil, errors.Wrap(err, "metadata marshal metadata failed")
}
err = json.Unmarshal(medaBytes, &r)
if err != nil {
if err := taskcommon.UnmarshalMetadata(req.Metadata, &r); err != nil {
return nil, errors.Wrap(err, "unmarshal metadata failed")
}
return &r, nil
......@@ -291,20 +287,6 @@ func (a *TaskAdaptor) getAspectRatio(size string) string {
}
}
func defaultString(s, def string) string {
if strings.TrimSpace(s) == "" {
return def
}
return s
}
func defaultInt(v int, def int) int {
if v == 0 {
return def
}
return v
}
// ============================
// JWT helpers
// ============================
......@@ -340,7 +322,7 @@ func (a *TaskAdaptor) createJWTTokenWithKey(apiKey string) (string, error) {
func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) {
taskInfo := &relaycommon.TaskInfo{}
resPayload := responsePayload{}
err := json.Unmarshal(respBody, &resPayload)
err := common.Unmarshal(respBody, &resPayload)
if err != nil {
return nil, errors.Wrap(err, "failed to unmarshal response body")
}
......@@ -374,7 +356,7 @@ func isNewAPIRelay(apiKey string) bool {
func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, error) {
var klingResp responsePayload
if err := json.Unmarshal(originTask.Data, &klingResp); err != nil {
if err := common.Unmarshal(originTask.Data, &klingResp); err != nil {
return nil, errors.Wrap(err, "unmarshal kling task data failed")
}
......@@ -401,6 +383,5 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, erro
Code: fmt.Sprintf("%d", klingResp.Code),
}
}
jsonData, _ := common.Marshal(openAIVideo)
return jsonData, nil
return common.Marshal(openAIVideo)
}
relay/channel/task/sora/adaptor.go
package sora
import (
"bytes"
"fmt"
"io"
"mime/multipart"
"net/http"
"net/textproto"
"strconv"
"strings"
"github.com/QuantumNous/new-api/common"
......@@ -11,12 +15,13 @@ import (
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/relay/channel"
taskcommon "github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting/system_setting"
"github.com/gin-gonic/gin"
"github.com/pkg/errors"
"github.com/tidwall/sjson"
)
// ============================
......@@ -57,6 +62,7 @@ type responseTask struct {
// ============================
type TaskAdaptor struct {
taskcommon.BaseBilling
ChannelType int
apiKey string
baseURL string
......@@ -69,15 +75,15 @@ func (a *TaskAdaptor) Init(info *relaycommon.RelayInfo) {
}
func validateRemixRequest(c *gin.Context) *dto.TaskError {
var req struct {
Prompt string `json:"prompt"`
}
var req relaycommon.TaskSubmitReq
if err := common.UnmarshalBodyReusable(c, &req); err != nil {
return service.TaskErrorWrapperLocal(err, "invalid_request", http.StatusBadRequest)
}
if strings.TrimSpace(req.Prompt) == "" {
return service.TaskErrorWrapperLocal(fmt.Errorf("field prompt is required"), "invalid_request", http.StatusBadRequest)
}
// 存储原始请求到 context，与 ValidateMultipartDirect 路径保持一致
c.Set("task_request", req)
return nil
}
......@@ -88,6 +94,41 @@ func (a *TaskAdaptor) ValidateRequestAndSetAction(c *gin.Context, info *relaycom
return relaycommon.ValidateMultipartDirect(c, info)
}
// EstimateBilling 根据用户请求的 seconds 和 size 计算 OtherRatios。
func (a *TaskAdaptor) EstimateBilling(c *gin.Context, info *relaycommon.RelayInfo) map[string]float64 {
// remix 路径的 OtherRatios 已在 ResolveOriginTask 中设置
if info.Action == constant.TaskActionRemix {
return nil
}
req, err := relaycommon.GetTaskRequest(c)
if err != nil {
return nil
}
seconds, _ := strconv.Atoi(req.Seconds)
if seconds == 0 {
seconds = req.Duration
}
if seconds <= 0 {
seconds = 4
}
size := req.Size
if size == "" {
size = "720x1280"
}
ratios := map[string]float64{
"seconds": float64(seconds),
"size": 1,
}
if size == "1792x1024" || size == "1024x1792" {
ratios["size"] = 1.666667
}
return ratios
}
func (a *TaskAdaptor) BuildRequestURL(info *relaycommon.RelayInfo) (string, error) {
if info.Action == constant.TaskActionRemix {
return fmt.Sprintf("%s/v1/videos/%s/remix", a.baseURL, info.OriginTaskID), nil
......@@ -107,6 +148,74 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
if err != nil {
return nil, errors.Wrap(err, "get_request_body_failed")
}
cachedBody, err := storage.Bytes()
if err != nil {
return nil, errors.Wrap(err, "read_body_bytes_failed")
}
contentType := c.GetHeader("Content-Type")
if strings.HasPrefix(contentType, "application/json") {
var bodyMap map[string]interface{}
if err := common.Unmarshal(cachedBody, &bodyMap); err == nil {
bodyMap["model"] = info.UpstreamModelName
if newBody, err := common.Marshal(bodyMap); err == nil {
return bytes.NewReader(newBody), nil
}
}
return bytes.NewReader(cachedBody), nil
}
if strings.Contains(contentType, "multipart/form-data") {
formData, err := common.ParseMultipartFormReusable(c)
if err != nil {
return bytes.NewReader(cachedBody), nil
}
var buf bytes.Buffer
writer := multipart.NewWriter(&buf)
writer.WriteField("model", info.UpstreamModelName)
for key, values := range formData.Value {
if key == "model" {
continue
}
for _, v := range values {
writer.WriteField(key, v)
}
}
for fieldName, fileHeaders := range formData.File {
for _, fh := range fileHeaders {
f, err := fh.Open()
if err != nil {
continue
}
ct := fh.Header.Get("Content-Type")
if ct == "" || ct == "application/octet-stream" {
buf512 := make([]byte, 512)
n, _ := io.ReadFull(f, buf512)
ct = http.DetectContentType(buf512[:n])
// Re-open after sniffing so the full content is copied below
f.Close()
f, err = fh.Open()
if err != nil {
continue
}
}
h := make(textproto.MIMEHeader)
h.Set("Content-Disposition", fmt.Sprintf(`form-data; name="%s"; filename="%s"`, fieldName, fh.Filename))
h.Set("Content-Type", ct)
part, err := writer.CreatePart(h)
if err != nil {
f.Close()
continue
}
io.Copy(part, f)
f.Close()
}
}
writer.Close()
c.Request.Header.Set("Content-Type", writer.FormDataContentType())
return &buf, nil
}
return common.ReaderOnly(storage), nil
}
......@@ -116,7 +225,7 @@ func (a *TaskAdaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, req
}
// DoResponse handles upstream response, returns taskID etc.
func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, _ *relaycommon.RelayInfo) (taskID string, taskData []byte, taskErr *dto.TaskError) {
func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (taskID string, taskData []byte, taskErr *dto.TaskError) {
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
taskErr = service.TaskErrorWrapper(err, "read_response_body_failed", http.StatusInternalServerError)
......@@ -131,17 +240,20 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, _ *relayco
return
}
if dResp.ID == "" {
if dResp.TaskID == "" {
taskErr = service.TaskErrorWrapper(fmt.Errorf("task_id is empty"), "invalid_response", http.StatusInternalServerError)
return
}
dResp.ID = dResp.TaskID
dResp.TaskID = ""
upstreamID := dResp.ID
if upstreamID == "" {
upstreamID = dResp.TaskID
}
if upstreamID == "" {
taskErr = service.TaskErrorWrapper(fmt.Errorf("task_id is empty"), "invalid_response", http.StatusInternalServerError)
return
}
// 使用公开 task_xxxx ID 返回给客户端
dResp.ID = info.PublicTaskID
dResp.TaskID = info.PublicTaskID
c.JSON(http.StatusOK, dResp)
return dResp.ID, responseBody, nil
return upstreamID, responseBody, nil
}
// FetchTask fetch task status
......@@ -192,7 +304,7 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
taskResult.Status = model.TaskStatusInProgress
case "completed":
taskResult.Status = model.TaskStatusSuccess
taskResult.Url = fmt.Sprintf("%s/v1/videos/%s/content", system_setting.ServerAddress, resTask.ID)
// Url intentionally left empty — the caller constructs the proxy URL using the public task ID
case "failed", "cancelled":
taskResult.Status = model.TaskStatusFailure
if resTask.Error != nil {
......@@ -210,5 +322,10 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
}
func (a *TaskAdaptor) ConvertToOpenAIVideo(task *model.Task) ([]byte, error) {
return task.Data, nil
data := task.Data
var err error
if data, err = sjson.SetBytes(data, "id", task.TaskID); err != nil {
return nil, errors.Wrap(err, "set id failed")
}
return data, nil
}
relay/channel/task/suno/adaptor.go
......@@ -2,18 +2,16 @@ package suno
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"strings"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/relay/channel"
taskcommon "github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/service"
......@@ -21,11 +19,16 @@ import (
)
type TaskAdaptor struct {
taskcommon.BaseBilling
ChannelType int
}
// ParseTaskResult is not used for Suno tasks.
// Suno polling uses a dedicated batch-fetch path (service.UpdateSunoTasks) that
// receives dto.TaskResponse[[]dto.SunoDataResponse] from the upstream /fetch API.
// This differs from the per-task polling used by video adaptors.
func (a *TaskAdaptor) ParseTaskResult([]byte) (*relaycommon.TaskInfo, error) {
return nil, fmt.Errorf("not implement") // todo implement this method if needed
return nil, fmt.Errorf("suno uses batch polling via UpdateSunoTasks, ParseTaskResult is not applicable")
}
func (a *TaskAdaptor) Init(info *relaycommon.RelayInfo) {
......@@ -47,13 +50,13 @@ func (a *TaskAdaptor) ValidateRequestAndSetAction(c *gin.Context, info *relaycom
return
}
if sunoRequest.ContinueClipId != "" {
if sunoRequest.TaskID == "" {
taskErr = service.TaskErrorWrapperLocal(fmt.Errorf("task id is empty"), "invalid_request", http.StatusBadRequest)
return
}
info.OriginTaskID = sunoRequest.TaskID
}
//if sunoRequest.ContinueClipId != "" {
// if sunoRequest.TaskID == "" {
// taskErr = service.TaskErrorWrapperLocal(fmt.Errorf("task id is empty"), "invalid_request", http.StatusBadRequest)
// return
// }
// info.OriginTaskID = sunoRequest.TaskID
//}
info.Action = action
c.Set("task_request", sunoRequest)
......@@ -76,12 +79,9 @@ func (a *TaskAdaptor) BuildRequestHeader(c *gin.Context, req *http.Request, info
func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayInfo) (io.Reader, error) {
sunoRequest, ok := c.Get("task_request")
if !ok {
err := common.UnmarshalBodyReusable(c, &sunoRequest)
if err != nil {
return nil, err
}
return nil, fmt.Errorf("task_request not found in context")
}
data, err := json.Marshal(sunoRequest)
data, err := common.Marshal(sunoRequest)
if err != nil {
return nil, err
}
......@@ -99,7 +99,7 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
return
}
var sunoResponse dto.TaskResponse[string]
err = json.Unmarshal(responseBody, &sunoResponse)
err = common.Unmarshal(responseBody, &sunoResponse)
if err != nil {
taskErr = service.TaskErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError)
return
......@@ -109,17 +109,13 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
return
}
for k, v := range resp.Header {
c.Writer.Header().Set(k, v[0])
}
c.Writer.Header().Set("Content-Type", "application/json")
c.Writer.WriteHeader(resp.StatusCode)
_, err = io.Copy(c.Writer, bytes.NewBuffer(responseBody))
if err != nil {
taskErr = service.TaskErrorWrapper(err, "copy_response_body_failed", http.StatusInternalServerError)
return
// 使用公开 task_xxxx ID 替换上游 ID 返回给客户端
publicResponse := dto.TaskResponse[string]{
Code: sunoResponse.Code,
Message: sunoResponse.Message,
Data: info.PublicTaskID,
}
c.JSON(http.StatusOK, publicResponse)
return sunoResponse.Data, nil, nil
}
......@@ -134,7 +130,7 @@ func (a *TaskAdaptor) GetChannelName() string {
func (a *TaskAdaptor) FetchTask(baseUrl, key string, body map[string]any, proxy string) (*http.Response, error) {
requestUrl := fmt.Sprintf("%s/suno/fetch", baseUrl)
byteBody, err := json.Marshal(body)
byteBody, err := common.Marshal(body)
if err != nil {
return nil, err
}
......@@ -144,13 +140,6 @@ func (a *TaskAdaptor) FetchTask(baseUrl, key string, body map[string]any, proxy
common.SysLog(fmt.Sprintf("Get Task error: %v", err))
return nil, err
}
defer req.Body.Close()
// 设置超时时间
timeout := time.Second * 15
ctx, cancel := context.WithTimeout(context.Background(), timeout)
defer cancel()
// 使用带有超时的 context 创建新的请求
req = req.WithContext(ctx)
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer "+key)
client, err := service.GetHttpClientWithProxy(proxy)
......
relay/channel/task/taskcommon/helpers.go 0 → 100644
package taskcommon
import (
"encoding/base64"
"fmt"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/setting/system_setting"
"github.com/gin-gonic/gin"
)
// UnmarshalMetadata converts a map[string]any metadata to a typed struct via JSON round-trip.
// This replaces the repeated pattern: json.Marshal(metadata) → json.Unmarshal(bytes, &target).
func UnmarshalMetadata(metadata map[string]any, target any) error {
if metadata == nil {
return nil
}
metaBytes, err := common.Marshal(metadata)
if err != nil {
return fmt.Errorf("marshal metadata failed: %w", err)
}
if err := common.Unmarshal(metaBytes, target); err != nil {
return fmt.Errorf("unmarshal metadata failed: %w", err)
}
return nil
}
// DefaultString returns val if non-empty, otherwise fallback.
func DefaultString(val, fallback string) string {
if val == "" {
return fallback
}
return val
}
// DefaultInt returns val if non-zero, otherwise fallback.
func DefaultInt(val, fallback int) int {
if val == 0 {
return fallback
}
return val
}
// EncodeLocalTaskID encodes an upstream operation name to a URL-safe base64 string.
// Used by Gemini/Vertex to store upstream names as task IDs.
func EncodeLocalTaskID(name string) string {
return base64.RawURLEncoding.EncodeToString([]byte(name))
}
// DecodeLocalTaskID decodes a base64-encoded upstream operation name.
func DecodeLocalTaskID(id string) (string, error) {
b, err := base64.RawURLEncoding.DecodeString(id)
if err != nil {
return "", err
}
return string(b), nil
}
// BuildProxyURL constructs the video proxy URL using the public task ID.
// e.g., "https://your-server.com/v1/videos/task_xxxx/content"
func BuildProxyURL(taskID string) string {
return fmt.Sprintf("%s/v1/videos/%s/content", system_setting.ServerAddress, taskID)
}
// Status-to-progress mapping constants for polling updates.
const (
ProgressSubmitted = "10%"
ProgressQueued = "20%"
ProgressInProgress = "30%"
ProgressComplete = "100%"
)
// ---------------------------------------------------------------------------
// BaseBilling — embeddable no-op implementations for TaskAdaptor billing methods.
// Adaptors that do not need custom billing can embed this struct directly.
// ---------------------------------------------------------------------------
type BaseBilling struct{}
// EstimateBilling returns nil (no extra ratios; use base model price).
func (BaseBilling) EstimateBilling(_ *gin.Context, _ *relaycommon.RelayInfo) map[string]float64 {
return nil
}
// AdjustBillingOnSubmit returns nil (no submit-time adjustment).
func (BaseBilling) AdjustBillingOnSubmit(_ *relaycommon.RelayInfo, _ []byte) map[string]float64 {
return nil
}
// AdjustBillingOnComplete returns 0 (keep pre-charged amount).
func (BaseBilling) AdjustBillingOnComplete(_ *model.Task, _ *relaycommon.TaskInfo) int {
return 0
}
relay/channel/task/vertex/adaptor.go
......@@ -2,13 +2,12 @@ package vertex
import (
"bytes"
"encoding/base64"
"encoding/json"
"fmt"
"io"
"net/http"
"regexp"
"strings"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
......@@ -17,6 +16,7 @@ import (
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/relay/channel"
taskcommon "github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
vertexcore "github.com/QuantumNous/new-api/relay/channel/vertex"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/service"
......@@ -62,6 +62,7 @@ type operationResponse struct {
// ============================
type TaskAdaptor struct {
taskcommon.BaseBilling
ChannelType int
apiKey string
baseURL string
......@@ -82,10 +83,10 @@ func (a *TaskAdaptor) ValidateRequestAndSetAction(c *gin.Context, info *relaycom
// BuildRequestURL constructs the upstream URL.
func (a *TaskAdaptor) BuildRequestURL(info *relaycommon.RelayInfo) (string, error) {
adc := &vertexcore.Credentials{}
if err := json.Unmarshal([]byte(a.apiKey), adc); err != nil {
if err := common.Unmarshal([]byte(a.apiKey), adc); err != nil {
return "", fmt.Errorf("failed to decode credentials: %w", err)
}
modelName := info.OriginModelName
modelName := info.UpstreamModelName
if modelName == "" {
modelName = "veo-3.0-generate-001"
}
......@@ -116,7 +117,7 @@ func (a *TaskAdaptor) BuildRequestHeader(c *gin.Context, req *http.Request, info
req.Header.Set("Accept", "application/json")
adc := &vertexcore.Credentials{}
if err := json.Unmarshal([]byte(a.apiKey), adc); err != nil {
if err := common.Unmarshal([]byte(a.apiKey), adc); err != nil {
return fmt.Errorf("failed to decode credentials: %w", err)
}
......@@ -133,6 +134,28 @@ func (a *TaskAdaptor) BuildRequestHeader(c *gin.Context, req *http.Request, info
return nil
}
// EstimateBilling 根据用户请求中的 sampleCount 计算 OtherRatios。
func (a *TaskAdaptor) EstimateBilling(c *gin.Context, _ *relaycommon.RelayInfo) map[string]float64 {
sampleCount := 1
v, ok := c.Get("task_request")
if ok {
req := v.(relaycommon.TaskSubmitReq)
if req.Metadata != nil {
if sc, exists := req.Metadata["sampleCount"]; exists {
if i, ok := sc.(int); ok && i > 0 {
sampleCount = i
}
if f, ok := sc.(float64); ok && int(f) > 0 {
sampleCount = int(f)
}
}
}
}
return map[string]float64{
"sampleCount": float64(sampleCount),
}
}
// BuildRequestBody converts request into Vertex specific format.
func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayInfo) (io.Reader, error) {
v, ok := c.Get("task_request")
......@@ -166,25 +189,7 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
return nil, fmt.Errorf("sampleCount must be greater than 0")
}
// if req.Duration > 0 {
// body.Parameters["durationSeconds"] = req.Duration
// } else if req.Seconds != "" {
// seconds, err := strconv.Atoi(req.Seconds)
// if err != nil {
// return nil, errors.Wrap(err, "convert seconds to int failed")
// }
// body.Parameters["durationSeconds"] = seconds
// }
info.PriceData.OtherRatios = map[string]float64{
"sampleCount": float64(body.Parameters["sampleCount"].(int)),
}
// if v, ok := body.Parameters["durationSeconds"]; ok {
// info.PriceData.OtherRatios["durationSeconds"] = float64(v.(int))
// }
data, err := json.Marshal(body)
data, err := common.Marshal(body)
if err != nil {
return nil, err
}
......@@ -205,14 +210,19 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
_ = resp.Body.Close()
var s submitResponse
if err := json.Unmarshal(responseBody, &s); err != nil {
if err := common.Unmarshal(responseBody, &s); err != nil {
return "", nil, service.TaskErrorWrapper(err, "unmarshal_response_failed", http.StatusInternalServerError)
}
if strings.TrimSpace(s.Name) == "" {
return "", nil, service.TaskErrorWrapper(fmt.Errorf("missing operation name"), "invalid_response", http.StatusInternalServerError)
}
localID := encodeLocalTaskID(s.Name)
c.JSON(http.StatusOK, gin.H{"task_id": localID})
localID := taskcommon.EncodeLocalTaskID(s.Name)
ov := dto.NewOpenAIVideo()
ov.ID = info.PublicTaskID
ov.TaskID = info.PublicTaskID
ov.CreatedAt = time.Now().Unix()
ov.Model = info.OriginModelName
c.JSON(http.StatusOK, ov)
return localID, responseBody, nil
}
......@@ -225,7 +235,7 @@ func (a *TaskAdaptor) FetchTask(baseUrl, key string, body map[string]any, proxy
if !ok {
return nil, fmt.Errorf("invalid task_id")
}
upstreamName, err := decodeLocalTaskID(taskID)
upstreamName, err := taskcommon.DecodeLocalTaskID(taskID)
if err != nil {
return nil, fmt.Errorf("decode task_id failed: %w", err)
}
......@@ -245,12 +255,12 @@ func (a *TaskAdaptor) FetchTask(baseUrl, key string, body map[string]any, proxy
url = fmt.Sprintf("https://%s-aiplatform.googleapis.com/v1/projects/%s/locations/%s/publishers/google/models/%s:fetchPredictOperation", region, project, region, modelName)
}
payload := map[string]string{"operationName": upstreamName}
data, err := json.Marshal(payload)
data, err := common.Marshal(payload)
if err != nil {
return nil, err
}
adc := &vertexcore.Credentials{}
if err := json.Unmarshal([]byte(key), adc); err != nil {
if err := common.Unmarshal([]byte(key), adc); err != nil {
return nil, fmt.Errorf("failed to decode credentials: %w", err)
}
token, err := vertexcore.AcquireAccessToken(*adc, proxy)
......@@ -274,7 +284,7 @@ func (a *TaskAdaptor) FetchTask(baseUrl, key string, body map[string]any, proxy
func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) {
var op operationResponse
if err := json.Unmarshal(respBody, &op); err != nil {
if err := common.Unmarshal(respBody, &op); err != nil {
return nil, fmt.Errorf("unmarshal operation response failed: %w", err)
}
ti := &relaycommon.TaskInfo{}
......@@ -338,7 +348,10 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
}
func (a *TaskAdaptor) ConvertToOpenAIVideo(task *model.Task) ([]byte, error) {
upstreamName, err := decodeLocalTaskID(task.TaskID)
// Use GetUpstreamTaskID() to get the real upstream operation name for model extraction.
// task.TaskID is now a public task_xxxx ID, no longer a base64-encoded upstream name.
upstreamTaskID := task.GetUpstreamTaskID()
upstreamName, err := taskcommon.DecodeLocalTaskID(upstreamTaskID)
if err != nil {
upstreamName = ""
}
......@@ -353,8 +366,8 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(task *model.Task) ([]byte, error) {
v.SetProgressStr(task.Progress)
v.CreatedAt = task.CreatedAt
v.CompletedAt = task.UpdatedAt
if strings.HasPrefix(task.FailReason, "data:") && len(task.FailReason) > 0 {
v.SetMetadata("url", task.FailReason)
if resultURL := task.GetResultURL(); strings.HasPrefix(resultURL, "data:") && len(resultURL) > 0 {
v.SetMetadata("url", resultURL)
}
return common.Marshal(v)
......@@ -364,18 +377,6 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(task *model.Task) ([]byte, error) {
// helpers
// ============================
func encodeLocalTaskID(name string) string {
return base64.RawURLEncoding.EncodeToString([]byte(name))
}
func decodeLocalTaskID(local string) (string, error) {
b, err := base64.RawURLEncoding.DecodeString(local)
if err != nil {
return "", err
}
return string(b), nil
}
var regionRe = regexp.MustCompile(`locations/([a-z0-9-]+)/`)
func extractRegionFromOperationName(name string) string {
......
relay/channel/task/vidu/adaptor.go
......@@ -2,7 +2,6 @@ package vidu
import (
"bytes"
"encoding/json"
"fmt"
"io"
"net/http"
......@@ -16,6 +15,7 @@ import (
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/relay/channel"
taskcommon "github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/service"
......@@ -73,6 +73,7 @@ type creation struct {
// ============================
type TaskAdaptor struct {
taskcommon.BaseBilling
ChannelType int
baseURL string
}
......@@ -115,7 +116,7 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
}
req := v.(relaycommon.TaskSubmitReq)
body, err := a.convertToRequestPayload(&req)
body, err := a.convertToRequestPayload(&req, info)
if err != nil {
return nil, err
}
......@@ -127,7 +128,7 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
}
}
data, err := json.Marshal(body)
data, err := common.Marshal(body)
if err != nil {
return nil, err
}
......@@ -168,7 +169,7 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
}
var vResp responsePayload
err = json.Unmarshal(responseBody, &vResp)
err = common.Unmarshal(responseBody, &vResp)
if err != nil {
taskErr = service.TaskErrorWrapper(errors.Wrap(err, fmt.Sprintf("%s", responseBody)), "unmarshal_response_failed", http.StatusInternalServerError)
return
......@@ -180,8 +181,8 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
}
ov := dto.NewOpenAIVideo()
ov.ID = vResp.TaskId
ov.TaskID = vResp.TaskId
ov.ID = info.PublicTaskID
ov.TaskID = info.PublicTaskID
ov.CreatedAt = time.Now().Unix()
ov.Model = info.OriginModelName
c.JSON(http.StatusOK, ov)
......@@ -223,47 +224,27 @@ func (a *TaskAdaptor) GetChannelName() string {
// helpers
// ============================
func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*requestPayload, error) {
func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq, info *relaycommon.RelayInfo) (*requestPayload, error) {
r := requestPayload{
Model: defaultString(req.Model, "viduq1"),
Model: taskcommon.DefaultString(info.UpstreamModelName, "viduq1"),
Images: req.Images,
Prompt: req.Prompt,
Duration: defaultInt(req.Duration, 5),
Resolution: defaultString(req.Size, "1080p"),
Duration: taskcommon.DefaultInt(req.Duration, 5),
Resolution: taskcommon.DefaultString(req.Size, "1080p"),
MovementAmplitude: "auto",
Bgm: false,
}
metadata := req.Metadata
medaBytes, err := json.Marshal(metadata)
if err != nil {
return nil, errors.Wrap(err, "metadata marshal metadata failed")
}
err = json.Unmarshal(medaBytes, &r)
if err != nil {
if err := taskcommon.UnmarshalMetadata(req.Metadata, &r); err != nil {
return nil, errors.Wrap(err, "unmarshal metadata failed")
}
return &r, nil
}
func defaultString(value, defaultValue string) string {
if value == "" {
return defaultValue
}
return value
}
func defaultInt(value, defaultValue int) int {
if value == 0 {
return defaultValue
}
return value
}
func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) {
taskInfo := &relaycommon.TaskInfo{}
var taskResp taskResultResponse
err := json.Unmarshal(respBody, &taskResp)
err := common.Unmarshal(respBody, &taskResp)
if err != nil {
return nil, errors.Wrap(err, "failed to unmarshal response body")
}
......@@ -293,7 +274,7 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, error) {
var viduResp taskResultResponse
if err := json.Unmarshal(originTask.Data, &viduResp); err != nil {
if err := common.Unmarshal(originTask.Data, &viduResp); err != nil {
return nil, errors.Wrap(err, "unmarshal vidu task data failed")
}
......@@ -315,6 +296,5 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, erro
}
}
jsonData, _ := common.Marshal(openAIVideo)
return jsonData, nil
return common.Marshal(openAIVideo)
}
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or sign in to comment