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 @@ ...@@ -30,8 +30,8 @@
</p> </p>
<p align="center"> <p align="center">
<a href="https://trendshift.io/repositories/8227" target="_blank"> <a href="https://trendshift.io/repositories/20180" 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"/> <img src="https://trendshift.io/api/badge/repositories/20180" alt="QuantumNous%2Fnew-api | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/>
</a> </a>
<br> <br>
<a href="https://hellogithub.com/repository/QuantumNous/new-api" target="_blank"> <a href="https://hellogithub.com/repository/QuantumNous/new-api" target="_blank">
......
README.ja.md
...@@ -30,8 +30,8 @@ ...@@ -30,8 +30,8 @@
</p> </p>
<p align="center"> <p align="center">
<a href="https://trendshift.io/repositories/8227" target="_blank"> <a href="https://trendshift.io/repositories/20180" 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"/> <img src="https://trendshift.io/api/badge/repositories/20180" alt="QuantumNous%2Fnew-api | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/>
</a> </a>
<br> <br>
<a href="https://hellogithub.com/repository/QuantumNous/new-api" target="_blank"> <a href="https://hellogithub.com/repository/QuantumNous/new-api" target="_blank">
......
README.md
...@@ -30,8 +30,8 @@ ...@@ -30,8 +30,8 @@
</p> </p>
<p align="center"> <p align="center">
<a href="https://trendshift.io/repositories/8227" target="_blank"> <a href="https://trendshift.io/repositories/20180" 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"/> <img src="https://trendshift.io/api/badge/repositories/20180" alt="QuantumNous%2Fnew-api | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/>
</a> </a>
<br> <br>
<a href="https://hellogithub.com/repository/QuantumNous/new-api" target="_blank"> <a href="https://hellogithub.com/repository/QuantumNous/new-api" target="_blank">
......
README.zh_CN.md
...@@ -30,8 +30,8 @@ ...@@ -30,8 +30,8 @@
</p> </p>
<p align="center"> <p align="center">
<a href="https://trendshift.io/repositories/8227" target="_blank"> <a href="https://trendshift.io/repositories/20180" 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"/> <img src="https://trendshift.io/api/badge/repositories/20180" alt="QuantumNous%2Fnew-api | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/>
</a> </a>
<br> <br>
<a href="https://hellogithub.com/repository/QuantumNous/new-api" target="_blank"> <a href="https://hellogithub.com/repository/QuantumNous/new-api" target="_blank">
......
README.zh_TW.md
...@@ -30,8 +30,8 @@ ...@@ -30,8 +30,8 @@
</p> </p>
<p align="center"> <p align="center">
<a href="https://trendshift.io/repositories/8227" target="_blank"> <a href="https://trendshift.io/repositories/20180" 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"/> <img src="https://trendshift.io/api/badge/repositories/20180" alt="QuantumNous%2Fnew-api | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/>
</a> </a>
<br> <br>
<a href="https://hellogithub.com/repository/QuantumNous/new-api" target="_blank"> <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) { ...@@ -243,7 +243,15 @@ func ParseMultipartFormReusable(c *gin.Context) (*multipart.Form, error) {
return nil, err 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) boundary, err := parseBoundary(contentType)
if err != nil { if err != nil {
return nil, err return nil, err
...@@ -295,7 +303,13 @@ func parseFormData(data []byte, v any) error { ...@@ -295,7 +303,13 @@ func parseFormData(data []byte, v any) error {
} }
func parseMultipartFormData(c *gin.Context, 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) boundary, err := parseBoundary(contentType)
if err != nil { if err != nil {
if errors.Is(err, errBoundaryNotFound) { if errors.Is(err, errBoundaryNotFound) {
......
common/init.go
...@@ -145,6 +145,8 @@ func initConstantEnv() { ...@@ -145,6 +145,8 @@ func initConstantEnv() {
constant.ErrorLogEnabled = GetEnvOrDefaultBool("ERROR_LOG_ENABLED", false) constant.ErrorLogEnabled = GetEnvOrDefaultBool("ERROR_LOG_ENABLED", false)
// 任务轮询时查询的最大数量 // 任务轮询时查询的最大数量
constant.TaskQueryLimit = GetEnvOrDefault("TASK_QUERY_LIMIT", 1000) constant.TaskQueryLimit = GetEnvOrDefault("TASK_QUERY_LIMIT", 1000)
// 异步任务超时时间（分钟），超过此时间未完成的任务将被标记为失败并退款。0 表示禁用。
constant.TaskTimeoutMinutes = GetEnvOrDefault("TASK_TIMEOUT_MINUTES", 1440)
soraPatchStr := GetEnvOrDefaultString("TASK_PRICE_PATCH", "") soraPatchStr := GetEnvOrDefaultString("TASK_PRICE_PATCH", "")
if soraPatchStr != "" { if soraPatchStr != "" {
......
constant/env.go
...@@ -16,6 +16,7 @@ var NotificationLimitDurationMinute int ...@@ -16,6 +16,7 @@ var NotificationLimitDurationMinute int
var GenerateDefaultToken bool var GenerateDefaultToken bool
var ErrorLogEnabled bool var ErrorLogEnabled bool
var TaskQueryLimit int var TaskQueryLimit int
var TaskTimeoutMinutes int
// temporary variable for sora patch, will be removed in future // temporary variable for sora patch, will be removed in future
var TaskPricePatches []string var TaskPricePatches []string
......
controller/codex_oauth.go
...@@ -145,6 +145,7 @@ func completeCodexOAuthWithChannelID(c *gin.Context, channelID int) { ...@@ -145,6 +145,7 @@ func completeCodexOAuthWithChannelID(c *gin.Context, channelID int) {
return return
} }
channelProxy := ""
if channelID > 0 { if channelID > 0 {
ch, err := model.GetChannelById(channelID, false) ch, err := model.GetChannelById(channelID, false)
if err != nil { if err != nil {
...@@ -159,6 +160,7 @@ func completeCodexOAuthWithChannelID(c *gin.Context, channelID int) { ...@@ -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"}) c.JSON(http.StatusOK, gin.H{"success": false, "message": "channel type is not Codex"})
return return
} }
channelProxy = ch.GetSetting().Proxy
} }
session := sessions.Default(c) session := sessions.Default(c)
...@@ -176,7 +178,7 @@ func completeCodexOAuthWithChannelID(c *gin.Context, channelID int) { ...@@ -176,7 +178,7 @@ func completeCodexOAuthWithChannelID(c *gin.Context, channelID int) {
ctx, cancel := context.WithTimeout(c.Request.Context(), 15*time.Second) ctx, cancel := context.WithTimeout(c.Request.Context(), 15*time.Second)
defer cancel() defer cancel()
tokenRes, err := service.ExchangeCodexAuthorizationCode(ctx, code, verifier) tokenRes, err := service.ExchangeCodexAuthorizationCodeWithProxy(ctx, code, verifier, channelProxy)
if err != nil { if err != nil {
common.SysError("failed to exchange codex authorization code: " + err.Error()) common.SysError("failed to exchange codex authorization code: " + err.Error())
c.JSON(http.StatusOK, gin.H{"success": false, "message": "授权码交换失败，请重试"}) c.JSON(http.StatusOK, gin.H{"success": false, "message": "授权码交换失败，请重试"})
......
controller/codex_usage.go
...@@ -2,7 +2,6 @@ package controller ...@@ -2,7 +2,6 @@ package controller
import ( import (
"context" "context"
"encoding/json"
"fmt" "fmt"
"net/http" "net/http"
"strconv" "strconv"
...@@ -80,7 +79,7 @@ func GetCodexChannelUsage(c *gin.Context) { ...@@ -80,7 +79,7 @@ func GetCodexChannelUsage(c *gin.Context) {
refreshCtx, refreshCancel := context.WithTimeout(c.Request.Context(), 10*time.Second) refreshCtx, refreshCancel := context.WithTimeout(c.Request.Context(), 10*time.Second)
defer refreshCancel() defer refreshCancel()
res, refreshErr := service.RefreshCodexOAuthToken(refreshCtx, oauthKey.RefreshToken) res, refreshErr := service.RefreshCodexOAuthTokenWithProxy(refreshCtx, oauthKey.RefreshToken, ch.GetSetting().Proxy)
if refreshErr == nil { if refreshErr == nil {
oauthKey.AccessToken = res.AccessToken oauthKey.AccessToken = res.AccessToken
oauthKey.RefreshToken = res.RefreshToken oauthKey.RefreshToken = res.RefreshToken
...@@ -109,7 +108,7 @@ func GetCodexChannelUsage(c *gin.Context) { ...@@ -109,7 +108,7 @@ func GetCodexChannelUsage(c *gin.Context) {
} }
var payload any var payload any
if json.Unmarshal(body, &payload) != nil { if common.Unmarshal(body, &payload) != nil {
payload = string(body) payload = string(body)
} }
......
controller/midjourney.go
...@@ -105,13 +105,13 @@ func UpdateMidjourneyTaskBulk() { ...@@ -105,13 +105,13 @@ func UpdateMidjourneyTaskBulk() {
} }
responseBody, err := io.ReadAll(resp.Body) responseBody, err := io.ReadAll(resp.Body)
if err != nil { 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 continue
} }
var responseItems []dto.MidjourneyDto var responseItems []dto.MidjourneyDto
err = json.Unmarshal(responseBody, &responseItems) err = json.Unmarshal(responseBody, &responseItems)
if err != nil { 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 continue
} }
resp.Body.Close() resp.Body.Close()
...@@ -130,6 +130,7 @@ func UpdateMidjourneyTaskBulk() { ...@@ -130,6 +130,7 @@ func UpdateMidjourneyTaskBulk() {
if !checkMjTaskNeedUpdate(task, responseItem) { if !checkMjTaskNeedUpdate(task, responseItem) {
continue continue
} }
preStatus := task.Status
task.Code = 1 task.Code = 1
task.Progress = responseItem.Progress task.Progress = responseItem.Progress
task.PromptEn = responseItem.PromptEn task.PromptEn = responseItem.PromptEn
...@@ -172,18 +173,26 @@ func UpdateMidjourneyTaskBulk() { ...@@ -172,18 +173,26 @@ func UpdateMidjourneyTaskBulk() {
shouldReturnQuota = true shouldReturnQuota = true
} }
} }
err = task.Update() won, err := task.UpdateWithStatus(preStatus)
if err != nil { if err != nil {
logger.LogError(ctx, "UpdateMidjourneyTask task error: "+err.Error()) logger.LogError(ctx, "UpdateMidjourneyTask task error: "+err.Error())
} else { } else if won && shouldReturnQuota {
if shouldReturnQuota { err = model.IncreaseUserQuota(task.UserId, task.Quota, false)
err = model.IncreaseUserQuota(task.UserId, task.Quota, false) if err != nil {
if err != nil { logger.LogError(ctx, "fail to increase user quota: "+err.Error())
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)
} }
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) { ...@@ -134,8 +134,10 @@ func GetStatus(c *gin.Context) {
customProviders := oauth.GetEnabledCustomProviders() customProviders := oauth.GetEnabledCustomProviders()
if len(customProviders) > 0 { if len(customProviders) > 0 {
type CustomOAuthInfo struct { type CustomOAuthInfo struct {
Id int `json:"id"`
Name string `json:"name"` Name string `json:"name"`
Slug string `json:"slug"` Slug string `json:"slug"`
Icon string `json:"icon"`
ClientId string `json:"client_id"` ClientId string `json:"client_id"`
AuthorizationEndpoint string `json:"authorization_endpoint"` AuthorizationEndpoint string `json:"authorization_endpoint"`
Scopes string `json:"scopes"` Scopes string `json:"scopes"`
...@@ -144,8 +146,10 @@ func GetStatus(c *gin.Context) { ...@@ -144,8 +146,10 @@ func GetStatus(c *gin.Context) {
for _, p := range customProviders { for _, p := range customProviders {
config := p.GetConfig() config := p.GetConfig()
providersInfo = append(providersInfo, CustomOAuthInfo{ providersInfo = append(providersInfo, CustomOAuthInfo{
Id: config.Id,
Name: config.Name, Name: config.Name,
Slug: config.Slug, Slug: config.Slug,
Icon: config.Icon,
ClientId: config.ClientId, ClientId: config.ClientId,
AuthorizationEndpoint: config.AuthorizationEndpoint, AuthorizationEndpoint: config.AuthorizationEndpoint,
Scopes: config.Scopes, Scopes: config.Scopes,
......
controller/oauth.go
...@@ -237,6 +237,16 @@ func findOrCreateOAuthUser(c *gin.Context, provider oauth.Provider, oauthUser *o ...@@ -237,6 +237,16 @@ func findOrCreateOAuthUser(c *gin.Context, provider oauth.Provider, oauthUser *o
// Set up new user // Set up new user
user.Username = provider.GetProviderPrefix() + strconv.Itoa(model.GetMaxUserId()+1) 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 != "" { if oauthUser.DisplayName != "" {
user.DisplayName = oauthUser.DisplayName user.DisplayName = oauthUser.DisplayName
} else if oauthUser.Username != "" { } else if oauthUser.Username != "" {
...@@ -295,12 +305,12 @@ func findOrCreateOAuthUser(c *gin.Context, provider oauth.Provider, oauthUser *o ...@@ -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 // Set the provider user ID on the user model and update
provider.SetProviderUserID(user, oauthUser.ProviderUserID) provider.SetProviderUserID(user, oauthUser.ProviderUserID)
if err := tx.Model(user).Updates(map[string]interface{}{ if err := tx.Model(user).Updates(map[string]interface{}{
"github_id": user.GitHubId, "github_id": user.GitHubId,
"discord_id": user.DiscordId, "discord_id": user.DiscordId,
"oidc_id": user.OidcId, "oidc_id": user.OidcId,
"linux_do_id": user.LinuxDOId, "linux_do_id": user.LinuxDOId,
"wechat_id": user.WeChatId, "wechat_id": user.WeChatId,
"telegram_id": user.TelegramId, "telegram_id": user.TelegramId,
}).Error; err != nil { }).Error; err != nil {
return err return err
} }
...@@ -340,6 +350,8 @@ func handleOAuthError(c *gin.Context, err error) { ...@@ -340,6 +350,8 @@ func handleOAuthError(c *gin.Context, err error) {
} else { } else {
common.ApiErrorI18n(c, e.MsgKey) common.ApiErrorI18n(c, e.MsgKey)
} }
case *oauth.AccessDeniedError:
common.ApiErrorMsg(c, e.Message)
case *oauth.TrustLevelError: case *oauth.TrustLevelError:
common.ApiErrorI18n(c, i18n.MsgOAuthTrustLevelLow) common.ApiErrorI18n(c, i18n.MsgOAuthTrustLevelLow)
default: default:
......
controller/relay.go
...@@ -455,72 +455,147 @@ func RelayNotFound(c *gin.Context) { ...@@ -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) { 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) relayInfo, err := relaycommon.GenRelayInfo(c, types.RelayFormatTask, nil, nil)
if err != nil { if err != nil {
c.JSON(http.StatusInternalServerError, &dto.TaskError{
Code: "gen_relay_info_failed",
Message: err.Error(),
StatusCode: http.StatusInternalServerError,
})
return return
} }
taskErr := taskRelayHandler(c, relayInfo)
if taskErr == nil { if taskErr := relay.ResolveOriginTask(c, relayInfo); taskErr != nil {
retryTimes = 0 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{ retryParam := &service.RetryParam{
Ctx: c, Ctx: c,
TokenGroup: relayInfo.TokenGroup, TokenGroup: relayInfo.TokenGroup,
ModelName: relayInfo.OriginModelName, ModelName: relayInfo.OriginModelName,
Retry: common.GetPointer(0), Retry: common.GetPointer(0),
} }
for ; shouldRetryTaskRelay(c, channelId, taskErr, retryTimes) && retryParam.GetRetry() < retryTimes; retryParam.IncreaseRetry() {
channel, newAPIError := getChannel(c, relayInfo, retryParam) for ; retryParam.GetRetry() <= common.RetryTimes; retryParam.IncreaseRetry() {
if newAPIError != nil { var channel *model.Channel
logger.LogError(c, fmt.Sprintf("CacheGetRandomSatisfiedChannel failed: %s", newAPIError.Error()))
taskErr = service.TaskErrorWrapperLocal(newAPIError.Err, "get_channel_failed", http.StatusInternalServerError) if lockedCh, ok := relayInfo.LockedChannel.(*model.Channel); ok && lockedCh != nil {
break 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") addUsedChannel(c, channel.Id)
useChannel = append(useChannel, fmt.Sprintf("%d", channelId)) bodyStorage, bodyErr := common.GetBodyStorage(c)
c.Set("use_channel", useChannel) if bodyErr != nil {
logger.LogInfo(c, fmt.Sprintf("using channel #%d to retry (remain times %d)", channel.Id, retryParam.GetRetry())) if common.IsRequestBodyTooLargeError(bodyErr) || errors.Is(bodyErr, common.ErrRequestBodyTooLarge) {
//middleware.SetupContextForSelectedChannel(c, channel, originalModel) taskErr = service.TaskErrorWrapperLocal(bodyErr, "read_request_body_failed", http.StatusRequestEntityTooLarge)
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)
} else { } else {
taskErr = service.TaskErrorWrapperLocal(err, "read_request_body_failed", http.StatusBadRequest) taskErr = service.TaskErrorWrapperLocal(bodyErr, "read_request_body_failed", http.StatusBadRequest)
} }
break break
} }
c.Request.Body = io.NopCloser(bodyStorage) 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") useChannel := c.GetStringSlice("use_channel")
if len(useChannel) > 1 { if len(useChannel) > 1 {
retryLogStr := fmt.Sprintf("重试：%s", strings.Trim(strings.Join(strings.Fields(fmt.Sprint(useChannel)), "->"), "[]")) retryLogStr := fmt.Sprintf("重试：%s", strings.Trim(strings.Join(strings.Fields(fmt.Sprint(useChannel)), "->"), "[]"))
logger.LogInfo(c, retryLogStr) 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 { // respondTaskError 统一输出 Task 错误响应（含 429 限流提示改写）
var err *dto.TaskError func respondTaskError(c *gin.Context, taskErr *dto.TaskError) {
switch relayInfo.RelayMode { if taskErr.StatusCode == http.StatusTooManyRequests {
case relayconstant.RelayModeSunoFetch, relayconstant.RelayModeSunoFetchByID, relayconstant.RelayModeVideoFetchByID: taskErr.Message = "当前分组上游负载已饱和，请稍后再试"
err = relay.RelayTaskFetch(c, relayInfo.RelayMode)
default:
err = relay.RelayTaskSubmit(c, relayInfo)
} }
return err c.JSON(taskErr.StatusCode, taskErr)
} }
func shouldRetryTaskRelay(c *gin.Context, channelId int, taskErr *dto.TaskError, retryTimes int) bool { 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, ...@@ -544,7 +619,7 @@ func shouldRetryTaskRelay(c *gin.Context, channelId int, taskErr *dto.TaskError,
} }
if taskErr.StatusCode/100 == 5 { if taskErr.StatusCode/100 == 5 {
// 超时不重试 // 超时不重试
if taskErr.StatusCode == 504 || taskErr.StatusCode == 524 { if operation_setting.IsAlwaysSkipRetryStatusCode(taskErr.StatusCode) {
return false return false
} }
return true return true
......
controller/task.go
package controller package controller
import ( import (
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"sort"
"strconv" "strconv"
"time"
"github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto" "github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/relay" "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/gin-gonic/gin"
"github.com/samber/lo"
) )
// UpdateTaskBulk 薄入口，实际轮询逻辑在 service 层
func UpdateTaskBulk() { func UpdateTaskBulk() {
//revocer service.TaskPollingLoop()
//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
} }
func GetAllTask(c *gin.Context) { func GetAllTask(c *gin.Context) {
...@@ -247,7 +38,7 @@ func GetAllTask(c *gin.Context) { ...@@ -247,7 +38,7 @@ func GetAllTask(c *gin.Context) {
items := model.TaskGetAllTasks(pageInfo.GetStartIdx(), pageInfo.GetPageSize(), queryParams) items := model.TaskGetAllTasks(pageInfo.GetStartIdx(), pageInfo.GetPageSize(), queryParams)
total := model.TaskCountAllTasks(queryParams) total := model.TaskCountAllTasks(queryParams)
pageInfo.SetTotal(int(total)) pageInfo.SetTotal(int(total))
pageInfo.SetItems(items) pageInfo.SetItems(tasksToDto(items, true))
common.ApiSuccess(c, pageInfo) common.ApiSuccess(c, pageInfo)
} }
...@@ -271,6 +62,33 @@ func GetUserTask(c *gin.Context) { ...@@ -271,6 +62,33 @@ func GetUserTask(c *gin.Context) {
items := model.TaskGetAllUserTask(userId, pageInfo.GetStartIdx(), pageInfo.GetPageSize(), queryParams) items := model.TaskGetAllUserTask(userId, pageInfo.GetStartIdx(), pageInfo.GetPageSize(), queryParams)
total := model.TaskCountAllUserTask(userId, queryParams) total := model.TaskCountAllUserTask(userId, queryParams)
pageInfo.SetTotal(int(total)) pageInfo.SetTotal(int(total))
pageInfo.SetItems(items) pageInfo.SetItems(tasksToDto(items, false))
common.ApiSuccess(c, pageInfo) 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) { ...@@ -582,6 +582,44 @@ func UpdateUser(c *gin.Context) {
return 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) { func UpdateSelf(c *gin.Context) {
var requestData map[string]interface{} var requestData map[string]interface{}
err := json.NewDecoder(c.Request.Body).Decode(&requestData) err := json.NewDecoder(c.Request.Body).Decode(&requestData)
......
controller/video_proxy.go
...@@ -16,59 +16,44 @@ import ( ...@@ -16,59 +16,44 @@ import (
"github.com/gin-gonic/gin" "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) { func VideoProxy(c *gin.Context) {
taskID := c.Param("task_id") taskID := c.Param("task_id")
if taskID == "" { if taskID == "" {
c.JSON(http.StatusBadRequest, gin.H{ videoProxyError(c, http.StatusBadRequest, "invalid_request_error", "task_id is required")
"error": gin.H{
"message": "task_id is required",
"type": "invalid_request_error",
},
})
return return
} }
task, exists, err := model.GetByOnlyTaskId(taskID) task, exists, err := model.GetByOnlyTaskId(taskID)
if err != nil { if err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to query task %s: %s", taskID, err.Error())) logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to query task %s: %s", taskID, err.Error()))
c.JSON(http.StatusInternalServerError, gin.H{ videoProxyError(c, http.StatusInternalServerError, "server_error", "Failed to query task")
"error": gin.H{
"message": "Failed to query task",
"type": "server_error",
},
})
return return
} }
if !exists || task == nil { if !exists || task == nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to get task %s: %v", taskID, err)) videoProxyError(c, http.StatusNotFound, "invalid_request_error", "Task not found")
c.JSON(http.StatusNotFound, gin.H{
"error": gin.H{
"message": "Task not found",
"type": "invalid_request_error",
},
})
return return
} }
if task.Status != model.TaskStatusSuccess { if task.Status != model.TaskStatusSuccess {
c.JSON(http.StatusBadRequest, gin.H{ videoProxyError(c, http.StatusBadRequest, "invalid_request_error",
"error": gin.H{ fmt.Sprintf("Task is not completed yet, current status: %s", task.Status))
"message": fmt.Sprintf("Task is not completed yet, current status: %s", task.Status),
"type": "invalid_request_error",
},
})
return return
} }
channel, err := model.CacheGetChannel(task.ChannelId) channel, err := model.CacheGetChannel(task.ChannelId)
if err != nil { if err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to get task %s: not found", taskID)) logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to get channel for task %s: %s", taskID, err.Error()))
c.JSON(http.StatusInternalServerError, gin.H{ videoProxyError(c, http.StatusInternalServerError, "server_error", "Failed to retrieve channel information")
"error": gin.H{
"message": "Failed to retrieve channel information",
"type": "server_error",
},
})
return return
} }
baseURL := channel.GetBaseURL() baseURL := channel.GetBaseURL()
...@@ -81,12 +66,7 @@ func VideoProxy(c *gin.Context) { ...@@ -81,12 +66,7 @@ func VideoProxy(c *gin.Context) {
client, err := service.GetHttpClientWithProxy(proxy) client, err := service.GetHttpClientWithProxy(proxy)
if err != nil { if err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to create proxy client for task %s: %s", taskID, err.Error())) 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{ videoProxyError(c, http.StatusInternalServerError, "server_error", "Failed to create proxy client")
"error": gin.H{
"message": "Failed to create proxy client",
"type": "server_error",
},
})
return return
} }
...@@ -95,12 +75,7 @@ func VideoProxy(c *gin.Context) { ...@@ -95,12 +75,7 @@ func VideoProxy(c *gin.Context) {
req, err := http.NewRequestWithContext(ctx, http.MethodGet, "", nil) req, err := http.NewRequestWithContext(ctx, http.MethodGet, "", nil)
if err != nil { if err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to create request: %s", err.Error())) logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to create request: %s", err.Error()))
c.JSON(http.StatusInternalServerError, gin.H{ videoProxyError(c, http.StatusInternalServerError, "server_error", "Failed to create proxy request")
"error": gin.H{
"message": "Failed to create proxy request",
"type": "server_error",
},
})
return return
} }
...@@ -109,68 +84,43 @@ func VideoProxy(c *gin.Context) { ...@@ -109,68 +84,43 @@ func VideoProxy(c *gin.Context) {
apiKey := task.PrivateData.Key apiKey := task.PrivateData.Key
if apiKey == "" { if apiKey == "" {
logger.LogError(c.Request.Context(), fmt.Sprintf("Missing stored API key for Gemini task %s", taskID)) logger.LogError(c.Request.Context(), fmt.Sprintf("Missing stored API key for Gemini task %s", taskID))
c.JSON(http.StatusInternalServerError, gin.H{ videoProxyError(c, http.StatusInternalServerError, "server_error", "API key not stored for task")
"error": gin.H{
"message": "API key not stored for task",
"type": "server_error",
},
})
return return
} }
videoURL, err = getGeminiVideoURL(channel, task, apiKey) videoURL, err = getGeminiVideoURL(channel, task, apiKey)
if err != nil { if err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to resolve Gemini video URL for task %s: %s", taskID, err.Error())) 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{ videoProxyError(c, http.StatusBadGateway, "server_error", "Failed to resolve Gemini video URL")
"error": gin.H{
"message": "Failed to resolve Gemini video URL",
"type": "server_error",
},
})
return return
} }
req.Header.Set("x-goog-api-key", apiKey) req.Header.Set("x-goog-api-key", apiKey)
case constant.ChannelTypeOpenAI, constant.ChannelTypeSora: 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) req.Header.Set("Authorization", "Bearer "+channel.Key)
default: default:
// Video URL is directly in task.FailReason // Video URL is stored in PrivateData.ResultURL (fallback to FailReason for old data)
videoURL = task.FailReason videoURL = task.GetResultURL()
} }
req.URL, err = url.Parse(videoURL) req.URL, err = url.Parse(videoURL)
if err != nil { if err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to parse URL %s: %s", videoURL, err.Error())) logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to parse URL %s: %s", videoURL, err.Error()))
c.JSON(http.StatusInternalServerError, gin.H{ videoProxyError(c, http.StatusInternalServerError, "server_error", "Failed to create proxy request")
"error": gin.H{
"message": "Failed to create proxy request",
"type": "server_error",
},
})
return return
} }
resp, err := client.Do(req) resp, err := client.Do(req)
if err != nil { if err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to fetch video from %s: %s", videoURL, err.Error())) logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to fetch video from %s: %s", videoURL, err.Error()))
c.JSON(http.StatusBadGateway, gin.H{ videoProxyError(c, http.StatusBadGateway, "server_error", "Failed to fetch video content")
"error": gin.H{
"message": "Failed to fetch video content",
"type": "server_error",
},
})
return return
} }
defer resp.Body.Close() defer resp.Body.Close()
if resp.StatusCode != http.StatusOK { if resp.StatusCode != http.StatusOK {
logger.LogError(c.Request.Context(), fmt.Sprintf("Upstream returned status %d for %s", resp.StatusCode, videoURL)) logger.LogError(c.Request.Context(), fmt.Sprintf("Upstream returned status %d for %s", resp.StatusCode, videoURL))
c.JSON(http.StatusBadGateway, gin.H{ videoProxyError(c, http.StatusBadGateway, "server_error",
"error": gin.H{ fmt.Sprintf("Upstream service returned status %d", resp.StatusCode))
"message": fmt.Sprintf("Upstream service returned status %d", resp.StatusCode),
"type": "server_error",
},
})
return return
} }
...@@ -180,10 +130,9 @@ func VideoProxy(c *gin.Context) { ...@@ -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) c.Writer.WriteHeader(resp.StatusCode)
_, err = io.Copy(c.Writer, resp.Body) if _, err = io.Copy(c.Writer, resp.Body); err != nil {
if err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to stream video content: %s", err.Error())) logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to stream video content: %s", err.Error()))
} }
} }
controller/video_proxy_gemini.go
package controller package controller
import ( import (
"encoding/json"
"fmt" "fmt"
"io" "io"
"strconv" "strconv"
"strings" "strings"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/relay" "github.com/QuantumNous/new-api/relay"
...@@ -37,7 +37,7 @@ func getGeminiVideoURL(channel *model.Channel, task *model.Task, apiKey string) ...@@ -37,7 +37,7 @@ func getGeminiVideoURL(channel *model.Channel, task *model.Task, apiKey string)
proxy := channel.GetSetting().Proxy proxy := channel.GetSetting().Proxy
resp, err := adaptor.FetchTask(baseURL, apiKey, map[string]any{ resp, err := adaptor.FetchTask(baseURL, apiKey, map[string]any{
"task_id": task.TaskID, "task_id": task.GetUpstreamTaskID(),
"action": task.Action, "action": task.Action,
}, proxy) }, proxy)
if err != nil { if err != nil {
...@@ -71,7 +71,7 @@ func extractGeminiVideoURLFromTaskData(task *model.Task) string { ...@@ -71,7 +71,7 @@ func extractGeminiVideoURLFromTaskData(task *model.Task) string {
return "" return ""
} }
var payload map[string]any 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 ""
} }
return extractGeminiVideoURLFromMap(payload) return extractGeminiVideoURLFromMap(payload)
...@@ -79,7 +79,7 @@ func extractGeminiVideoURLFromTaskData(task *model.Task) string { ...@@ -79,7 +79,7 @@ func extractGeminiVideoURLFromTaskData(task *model.Task) string {
func extractGeminiVideoURLFromPayload(body []byte) string { func extractGeminiVideoURLFromPayload(body []byte) string {
var payload map[string]any var payload map[string]any
if err := json.Unmarshal(body, &payload); err != nil { if err := common.Unmarshal(body, &payload); err != nil {
return "" return ""
} }
return extractGeminiVideoURLFromMap(payload) return extractGeminiVideoURLFromMap(payload)
......
dto/channel_settings.go
...@@ -24,14 +24,16 @@ const ( ...@@ -24,14 +24,16 @@ const (
) )
type ChannelOtherSettings struct { type ChannelOtherSettings struct {
AzureResponsesVersion string `json:"azure_responses_version,omitempty"` AzureResponsesVersion string `json:"azure_responses_version,omitempty"`
VertexKeyType VertexKeyType `json:"vertex_key_type,omitempty"` // "json" or "api_key" VertexKeyType VertexKeyType `json:"vertex_key_type,omitempty"` // "json" or "api_key"
OpenRouterEnterprise *bool `json:"openrouter_enterprise,omitempty"` OpenRouterEnterprise *bool `json:"openrouter_enterprise,omitempty"`
ClaudeBetaQuery bool `json:"claude_beta_query,omitempty"` // Claude 渠道是否强制追加 ?beta=true ClaudeBetaQuery bool `json:"claude_beta_query,omitempty"` // Claude 渠道是否强制追加 ?beta=true
AllowServiceTier bool `json:"allow_service_tier,omitempty"` // 是否允许 service_tier 透传（默认过滤以避免额外计费） AllowServiceTier bool `json:"allow_service_tier,omitempty"` // 是否允许 service_tier 透传（默认过滤以避免额外计费）
DisableStore bool `json:"disable_store,omitempty"` // 是否禁用 store 透传（默认允许透传，禁用后可能导致 Codex 无法使用） AllowInferenceGeo bool `json:"allow_inference_geo,omitempty"` // 是否允许 inference_geo 透传（仅 Claude，默认过滤以满足数据驻留合规）
AllowSafetyIdentifier bool `json:"allow_safety_identifier,omitempty"` // 是否允许 safety_identifier 透传（默认过滤以保护用户隐私） DisableStore bool `json:"disable_store,omitempty"` // 是否禁用 store 透传（默认允许透传，禁用后可能导致 Codex 无法使用）
AwsKeyType AwsKeyType `json:"aws_key_type,omitempty"` 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 { func (s *ChannelOtherSettings) IsOpenRouterEnterprise() bool {
......
dto/claude.go
...@@ -190,10 +190,13 @@ type ClaudeToolChoice struct { ...@@ -190,10 +190,13 @@ type ClaudeToolChoice struct {
} }
type ClaudeRequest struct { type ClaudeRequest struct {
Model string `json:"model"` Model string `json:"model"`
Prompt string `json:"prompt,omitempty"` Prompt string `json:"prompt,omitempty"`
System any `json:"system,omitempty"` System any `json:"system,omitempty"`
Messages []ClaudeMessage `json:"messages,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"` MaxTokens uint `json:"max_tokens,omitempty"`
MaxTokensToSample uint `json:"max_tokens_to_sample,omitempty"` MaxTokensToSample uint `json:"max_tokens_to_sample,omitempty"`
StopSequences []string `json:"stop_sequences,omitempty"` StopSequences []string `json:"stop_sequences,omitempty"`
...@@ -210,7 +213,8 @@ type ClaudeRequest struct { ...@@ -210,7 +213,8 @@ type ClaudeRequest struct {
Thinking *Thinking `json:"thinking,omitempty"` Thinking *Thinking `json:"thinking,omitempty"`
McpServers json.RawMessage `json:"mcp_servers,omitempty"` McpServers json.RawMessage `json:"mcp_servers,omitempty"`
Metadata json.RawMessage `json:"metadata,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"` ServiceTier string `json:"service_tier,omitempty"`
} }
......
dto/gemini.go
...@@ -324,25 +324,26 @@ type GeminiChatTool struct { ...@@ -324,25 +324,26 @@ type GeminiChatTool struct {
} }
type GeminiChatGenerationConfig struct { type GeminiChatGenerationConfig struct {
Temperature *float64 `json:"temperature,omitempty"` Temperature *float64 `json:"temperature,omitempty"`
TopP float64 `json:"topP,omitempty"` TopP float64 `json:"topP,omitempty"`
TopK float64 `json:"topK,omitempty"` TopK float64 `json:"topK,omitempty"`
MaxOutputTokens uint `json:"maxOutputTokens,omitempty"` MaxOutputTokens uint `json:"maxOutputTokens,omitempty"`
CandidateCount int `json:"candidateCount,omitempty"` CandidateCount int `json:"candidateCount,omitempty"`
StopSequences []string `json:"stopSequences,omitempty"` StopSequences []string `json:"stopSequences,omitempty"`
ResponseMimeType string `json:"responseMimeType,omitempty"` ResponseMimeType string `json:"responseMimeType,omitempty"`
ResponseSchema any `json:"responseSchema,omitempty"` ResponseSchema any `json:"responseSchema,omitempty"`
ResponseJsonSchema json.RawMessage `json:"responseJsonSchema,omitempty"` ResponseJsonSchema json.RawMessage `json:"responseJsonSchema,omitempty"`
PresencePenalty *float32 `json:"presencePenalty,omitempty"` PresencePenalty *float32 `json:"presencePenalty,omitempty"`
FrequencyPenalty *float32 `json:"frequencyPenalty,omitempty"` FrequencyPenalty *float32 `json:"frequencyPenalty,omitempty"`
ResponseLogprobs bool `json:"responseLogprobs,omitempty"` ResponseLogprobs bool `json:"responseLogprobs,omitempty"`
Logprobs *int32 `json:"logprobs,omitempty"` Logprobs *int32 `json:"logprobs,omitempty"`
MediaResolution MediaResolution `json:"mediaResolution,omitempty"` EnableEnhancedCivicAnswers *bool `json:"enableEnhancedCivicAnswers,omitempty"`
Seed int64 `json:"seed,omitempty"` MediaResolution MediaResolution `json:"mediaResolution,omitempty"`
ResponseModalities []string `json:"responseModalities,omitempty"` Seed int64 `json:"seed,omitempty"`
ThinkingConfig *GeminiThinkingConfig `json:"thinkingConfig,omitempty"` ResponseModalities []string `json:"responseModalities,omitempty"`
SpeechConfig json.RawMessage `json:"speechConfig,omitempty"` // RawMessage to allow flexible speech config ThinkingConfig *GeminiThinkingConfig `json:"thinkingConfig,omitempty"`
ImageConfig json.RawMessage `json:"imageConfig,omitempty"` // RawMessage to allow flexible image config 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. // UnmarshalJSON allows GeminiChatGenerationConfig to accept both snake_case and camelCase fields.
...@@ -350,22 +351,23 @@ func (c *GeminiChatGenerationConfig) UnmarshalJSON(data []byte) error { ...@@ -350,22 +351,23 @@ func (c *GeminiChatGenerationConfig) UnmarshalJSON(data []byte) error {
type Alias GeminiChatGenerationConfig type Alias GeminiChatGenerationConfig
var aux struct { var aux struct {
Alias Alias
TopPSnake float64 `json:"top_p,omitempty"` TopPSnake float64 `json:"top_p,omitempty"`
TopKSnake float64 `json:"top_k,omitempty"` TopKSnake float64 `json:"top_k,omitempty"`
MaxOutputTokensSnake uint `json:"max_output_tokens,omitempty"` MaxOutputTokensSnake uint `json:"max_output_tokens,omitempty"`
CandidateCountSnake int `json:"candidate_count,omitempty"` CandidateCountSnake int `json:"candidate_count,omitempty"`
StopSequencesSnake []string `json:"stop_sequences,omitempty"` StopSequencesSnake []string `json:"stop_sequences,omitempty"`
ResponseMimeTypeSnake string `json:"response_mime_type,omitempty"` ResponseMimeTypeSnake string `json:"response_mime_type,omitempty"`
ResponseSchemaSnake any `json:"response_schema,omitempty"` ResponseSchemaSnake any `json:"response_schema,omitempty"`
ResponseJsonSchemaSnake json.RawMessage `json:"response_json_schema,omitempty"` ResponseJsonSchemaSnake json.RawMessage `json:"response_json_schema,omitempty"`
PresencePenaltySnake *float32 `json:"presence_penalty,omitempty"` PresencePenaltySnake *float32 `json:"presence_penalty,omitempty"`
FrequencyPenaltySnake *float32 `json:"frequency_penalty,omitempty"` FrequencyPenaltySnake *float32 `json:"frequency_penalty,omitempty"`
ResponseLogprobsSnake bool `json:"response_logprobs,omitempty"` ResponseLogprobsSnake bool `json:"response_logprobs,omitempty"`
MediaResolutionSnake MediaResolution `json:"media_resolution,omitempty"` EnableEnhancedCivicAnswersSnake *bool `json:"enable_enhanced_civic_answers,omitempty"`
ResponseModalitiesSnake []string `json:"response_modalities,omitempty"` MediaResolutionSnake MediaResolution `json:"media_resolution,omitempty"`
ThinkingConfigSnake *GeminiThinkingConfig `json:"thinking_config,omitempty"` ResponseModalitiesSnake []string `json:"response_modalities,omitempty"`
SpeechConfigSnake json.RawMessage `json:"speech_config,omitempty"` ThinkingConfigSnake *GeminiThinkingConfig `json:"thinking_config,omitempty"`
ImageConfigSnake json.RawMessage `json:"image_config,omitempty"` SpeechConfigSnake json.RawMessage `json:"speech_config,omitempty"`
ImageConfigSnake json.RawMessage `json:"image_config,omitempty"`
} }
if err := common.Unmarshal(data, &aux); err != nil { if err := common.Unmarshal(data, &aux); err != nil {
...@@ -408,6 +410,9 @@ func (c *GeminiChatGenerationConfig) UnmarshalJSON(data []byte) error { ...@@ -408,6 +410,9 @@ func (c *GeminiChatGenerationConfig) UnmarshalJSON(data []byte) error {
if aux.ResponseLogprobsSnake { if aux.ResponseLogprobsSnake {
c.ResponseLogprobs = aux.ResponseLogprobsSnake c.ResponseLogprobs = aux.ResponseLogprobsSnake
} }
if aux.EnableEnhancedCivicAnswersSnake != nil {
c.EnableEnhancedCivicAnswers = aux.EnableEnhancedCivicAnswersSnake
}
if aux.MediaResolutionSnake != "" { if aux.MediaResolutionSnake != "" {
c.MediaResolution = aux.MediaResolutionSnake c.MediaResolution = aux.MediaResolutionSnake
} }
...@@ -453,12 +458,14 @@ type GeminiChatResponse struct { ...@@ -453,12 +458,14 @@ type GeminiChatResponse struct {
} }
type GeminiUsageMetadata struct { type GeminiUsageMetadata struct {
PromptTokenCount int `json:"promptTokenCount"` PromptTokenCount int `json:"promptTokenCount"`
CandidatesTokenCount int `json:"candidatesTokenCount"` ToolUsePromptTokenCount int `json:"toolUsePromptTokenCount"`
TotalTokenCount int `json:"totalTokenCount"` CandidatesTokenCount int `json:"candidatesTokenCount"`
ThoughtsTokenCount int `json:"thoughtsTokenCount"` TotalTokenCount int `json:"totalTokenCount"`
CachedContentTokenCount int `json:"cachedContentTokenCount"` ThoughtsTokenCount int `json:"thoughtsTokenCount"`
PromptTokensDetails []GeminiPromptTokensDetails `json:"promptTokensDetails"` CachedContentTokenCount int `json:"cachedContentTokenCount"`
PromptTokensDetails []GeminiPromptTokensDetails `json:"promptTokensDetails"`
ToolUsePromptTokensDetails []GeminiPromptTokensDetails `json:"toolUsePromptTokensDetails"`
} }
type GeminiPromptTokensDetails struct { type GeminiPromptTokensDetails struct {
......
dto/openai_request.go
...@@ -54,18 +54,22 @@ type GeneralOpenAIRequest struct { ...@@ -54,18 +54,22 @@ type GeneralOpenAIRequest struct {
ParallelTooCalls *bool `json:"parallel_tool_calls,omitempty"` ParallelTooCalls *bool `json:"parallel_tool_calls,omitempty"`
Tools []ToolCallRequest `json:"tools,omitempty"` Tools []ToolCallRequest `json:"tools,omitempty"`
ToolChoice any `json:"tool_choice,omitempty"` ToolChoice any `json:"tool_choice,omitempty"`
FunctionCall json.RawMessage `json:"function_call,omitempty"`
User string `json:"user,omitempty"` User string `json:"user,omitempty"`
LogProbs bool `json:"logprobs,omitempty"` // ServiceTier specifies upstream service level and may affect billing.
TopLogProbs int `json:"top_logprobs,omitempty"` // This field is filtered by default and can be enabled via channel setting allow_service_tier.
Dimensions int `json:"dimensions,omitempty"` ServiceTier string `json:"service_tier,omitempty"`
Modalities json.RawMessage `json:"modalities,omitempty"` LogProbs bool `json:"logprobs,omitempty"`
Audio json.RawMessage `json:"audio,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 发送用户标识信息，默认过滤以保护用户隐私 // 注意：此字段会向 OpenAI 发送用户标识信息，默认过滤，可通过 allow_safety_identifier 开启
SafetyIdentifier string `json:"safety_identifier,omitempty"` 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. // Whether or not to store the output of this chat completion request for use in our model distillation or evals products.
// 是否存储此次请求数据供 OpenAI 用于评估和优化产品 // 是否存储此次请求数据供 OpenAI 用于评估和优化产品
// 注意：默认过滤此字段以保护用户隐私，但过滤后可能导致 Codex 无法正常使用 // 注意：默认允许透传，可通过 disable_store 禁用；禁用后可能导致 Codex 无法正常使用
Store json.RawMessage `json:"store,omitempty"` 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 // 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"` PromptCacheKey string `json:"prompt_cache_key,omitempty"`
...@@ -261,6 +265,9 @@ type FunctionRequest struct { ...@@ -261,6 +265,9 @@ type FunctionRequest struct {
type StreamOptions struct { type StreamOptions struct {
IncludeUsage bool `json:"include_usage,omitempty"` 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 { func (r *GeneralOpenAIRequest) GetMaxTokens() uint {
...@@ -799,30 +806,42 @@ type WebSearchOptions struct { ...@@ -799,30 +806,42 @@ type WebSearchOptions struct {
// https://platform.openai.com/docs/api-reference/responses/create // https://platform.openai.com/docs/api-reference/responses/create
type OpenAIResponsesRequest struct { type OpenAIResponsesRequest struct {
Model string `json:"model"` Model string `json:"model"`
Input json.RawMessage `json:"input,omitempty"` Input json.RawMessage `json:"input,omitempty"`
Include json.RawMessage `json:"include,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"` Instructions json.RawMessage `json:"instructions,omitempty"`
MaxOutputTokens uint `json:"max_output_tokens,omitempty"` MaxOutputTokens uint `json:"max_output_tokens,omitempty"`
TopLogProbs *int `json:"top_logprobs,omitempty"`
Metadata json.RawMessage `json:"metadata,omitempty"` Metadata json.RawMessage `json:"metadata,omitempty"`
ParallelToolCalls json.RawMessage `json:"parallel_tool_calls,omitempty"` ParallelToolCalls json.RawMessage `json:"parallel_tool_calls,omitempty"`
PreviousResponseID string `json:"previous_response_id,omitempty"` PreviousResponseID string `json:"previous_response_id,omitempty"`
Reasoning *Reasoning `json:"reasoning,omitempty"` Reasoning *Reasoning `json:"reasoning,omitempty"`
// 服务层级字段，用于指定 API 服务等级。允许透传可能导致实际计费高于预期，默认应过滤 // ServiceTier specifies upstream service level and may affect billing.
ServiceTier string `json:"service_tier,omitempty"` // 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"` Store json.RawMessage `json:"store,omitempty"`
PromptCacheKey json.RawMessage `json:"prompt_cache_key,omitempty"` PromptCacheKey json.RawMessage `json:"prompt_cache_key,omitempty"`
PromptCacheRetention json.RawMessage `json:"prompt_cache_retention,omitempty"` PromptCacheRetention json.RawMessage `json:"prompt_cache_retention,omitempty"`
Stream bool `json:"stream,omitempty"` // SafetyIdentifier carries client identity for policy abuse detection.
Temperature *float64 `json:"temperature,omitempty"` // This field is filtered by default and can be enabled via channel setting allow_safety_identifier.
Text json.RawMessage `json:"text,omitempty"` SafetyIdentifier string `json:"safety_identifier,omitempty"`
ToolChoice json.RawMessage `json:"tool_choice,omitempty"` Stream bool `json:"stream,omitempty"`
Tools json.RawMessage `json:"tools,omitempty"` // 需要处理的参数很少，MCP 参数太多不确定，所以用 map StreamOptions *StreamOptions `json:"stream_options,omitempty"`
TopP *float64 `json:"top_p,omitempty"` Temperature *float64 `json:"temperature,omitempty"`
Truncation string `json:"truncation,omitempty"` Text json.RawMessage `json:"text,omitempty"`
User string `json:"user,omitempty"` ToolChoice json.RawMessage `json:"tool_choice,omitempty"`
MaxToolCalls uint `json:"max_tool_calls,omitempty"` Tools json.RawMessage `json:"tools,omitempty"` // 需要处理的参数很少，MCP 参数太多不确定，所以用 map
Prompt json.RawMessage `json:"prompt,omitempty"` 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 // qwen
EnableThinking json.RawMessage `json:"enable_thinking,omitempty"` EnableThinking json.RawMessage `json:"enable_thinking,omitempty"`
// perplexity // perplexity
......
dto/suno.go
...@@ -4,10 +4,6 @@ import ( ...@@ -4,10 +4,6 @@ import (
"encoding/json" "encoding/json"
) )
type TaskData interface {
SunoDataResponse | []SunoDataResponse | string | any
}
type SunoSubmitReq struct { type SunoSubmitReq struct {
GptDescriptionPrompt string `json:"gpt_description_prompt,omitempty"` GptDescriptionPrompt string `json:"gpt_description_prompt,omitempty"`
Prompt string `json:"prompt,omitempty"` Prompt string `json:"prompt,omitempty"`
...@@ -20,10 +16,6 @@ type SunoSubmitReq struct { ...@@ -20,10 +16,6 @@ type SunoSubmitReq struct {
MakeInstrumental bool `json:"make_instrumental"` MakeInstrumental bool `json:"make_instrumental"`
} }
type FetchReq struct {
IDs []string `json:"ids"`
}
type SunoDataResponse struct { type SunoDataResponse struct {
TaskID string `json:"task_id" gorm:"type:varchar(50);index"` TaskID string `json:"task_id" gorm:"type:varchar(50);index"`
Action string `json:"action" gorm:"type:varchar(40);index"` // 任务类型, song, lyrics, description-mode Action string `json:"action" gorm:"type:varchar(40);index"` // 任务类型, song, lyrics, description-mode
...@@ -66,30 +58,6 @@ type SunoLyrics struct { ...@@ -66,30 +58,6 @@ type SunoLyrics struct {
Text string `json:"text"` 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 { type SunoGoAPISubmitReq struct {
CustomMode bool `json:"custom_mode"` CustomMode bool `json:"custom_mode"`
......
dto/task.go
package dto package dto
import (
"encoding/json"
)
type TaskError struct { type TaskError struct {
Code string `json:"code"` Code string `json:"code"`
Message string `json:"message"` Message string `json:"message"`
...@@ -8,3 +12,46 @@ type TaskError struct { ...@@ -8,3 +12,46 @@ type TaskError struct {
LocalError bool `json:"-"` LocalError bool `json:"-"`
Error error `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 ...@@ -2,7 +2,6 @@ package logger
import ( import (
"context" "context"
"encoding/json"
"fmt" "fmt"
"io" "io"
"log" "log"
...@@ -151,7 +150,7 @@ func FormatQuota(quota int) string { ...@@ -151,7 +150,7 @@ func FormatQuota(quota int) string {
// LogJson 仅供测试使用 only for test // LogJson 仅供测试使用 only for test
func LogJson(ctx context.Context, msg string, obj any) { func LogJson(ctx context.Context, msg string, obj any) {
jsonStr, err := json.Marshal(obj) jsonStr, err := common.Marshal(obj)
if err != nil { if err != nil {
LogError(ctx, fmt.Sprintf("json marshal failed: %s", err.Error())) LogError(ctx, fmt.Sprintf("json marshal failed: %s", err.Error()))
return return
......
main.go
...@@ -19,6 +19,7 @@ import ( ...@@ -19,6 +19,7 @@ import (
"github.com/QuantumNous/new-api/middleware" "github.com/QuantumNous/new-api/middleware"
"github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/oauth" "github.com/QuantumNous/new-api/oauth"
"github.com/QuantumNous/new-api/relay"
"github.com/QuantumNous/new-api/router" "github.com/QuantumNous/new-api/router"
"github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/service"
_ "github.com/QuantumNous/new-api/setting/performance_setting" _ "github.com/QuantumNous/new-api/setting/performance_setting"
...@@ -111,6 +112,15 @@ func main() { ...@@ -111,6 +112,15 @@ func main() {
// Subscription quota reset task (daily/weekly/monthly/custom) // Subscription quota reset task (daily/weekly/monthly/custom)
service.StartSubscriptionQuotaResetTask() 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 { if common.IsMasterNode && constant.UpdateTask {
gopool.Go(func() { gopool.Go(func() {
controller.UpdateMidjourneyTaskBulk() controller.UpdateMidjourneyTaskBulk()
......
middleware/auth.go
...@@ -170,6 +170,24 @@ func WssAuth(c *gin.Context) { ...@@ -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 宽松版本的令牌认证中间件，用于只读查询接口。 // TokenAuthReadOnly 宽松版本的令牌认证中间件，用于只读查询接口。
// 只验证令牌 key 是否存在，不检查令牌状态、过期时间和额度。 // 只验证令牌 key 是否存在，不检查令牌状态、过期时间和额度。
// 即使令牌已过期、已耗尽或已禁用，也允许访问。 // 即使令牌已过期、已耗尽或已禁用，也允许访问。
......
middleware/logger.go
...@@ -7,14 +7,28 @@ import ( ...@@ -7,14 +7,28 @@ import (
"github.com/gin-gonic/gin" "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) { func SetUpLogger(server *gin.Engine) {
server.Use(gin.LoggerWithFormatter(func(param gin.LogFormatterParams) string { server.Use(gin.LoggerWithFormatter(func(param gin.LogFormatterParams) string {
var requestID string var requestID string
if param.Keys != nil { 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"), param.TimeStamp.Format("2006/01/02 - 15:04:05"),
tag,
requestID, requestID,
param.StatusCode, param.StatusCode,
param.Latency, param.Latency,
......
model/custom_oauth_provider.go
...@@ -2,32 +2,65 @@ package model ...@@ -2,32 +2,65 @@ package model
import ( import (
"errors" "errors"
"fmt"
"strings" "strings"
"time" "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 // CustomOAuthProvider stores configuration for custom OAuth providers
type CustomOAuthProvider struct { type CustomOAuthProvider struct {
Id int `json:"id" gorm:"primaryKey"` Id int `json:"id" gorm:"primaryKey"`
Name string `json:"name" gorm:"type:varchar(64);not null"` // Display name, e.g., "GitHub Enterprise" 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" 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 Icon string `json:"icon" gorm:"type:varchar(128);default:''"` // Icon name from @lobehub/icons
ClientId string `json:"client_id" gorm:"type:varchar(256)"` // OAuth client ID Enabled bool `json:"enabled" gorm:"default:false"` // Whether this provider is enabled
ClientSecret string `json:"-" gorm:"type:varchar(512)"` // OAuth client secret (not returned to frontend) ClientId string `json:"client_id" gorm:"type:varchar(256)"` // OAuth client ID
AuthorizationEndpoint string `json:"authorization_endpoint" gorm:"type:varchar(512)"` // Authorization URL ClientSecret string `json:"-" gorm:"type:varchar(512)"` // OAuth client secret (not returned to frontend)
TokenEndpoint string `json:"token_endpoint" gorm:"type:varchar(512)"` // Token exchange URL AuthorizationEndpoint string `json:"authorization_endpoint" gorm:"type:varchar(512)"` // Authorization URL
UserInfoEndpoint string `json:"user_info_endpoint" gorm:"type:varchar(512)"` // User info URL TokenEndpoint string `json:"token_endpoint" gorm:"type:varchar(512)"` // Token exchange URL
Scopes string `json:"scopes" gorm:"type:varchar(256);default:'openid profile email'"` // OAuth scopes 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) // 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" 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 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 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 EmailField string `json:"email_field" gorm:"type:varchar(128);default:'email'"` // Email field path
// Advanced options // Advanced options
WellKnown string `json:"well_known" gorm:"type:varchar(512)"` // OIDC discovery endpoint (optional) 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) 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"` CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"` UpdatedAt time.Time `json:"updated_at"`
...@@ -158,6 +191,57 @@ func validateCustomOAuthProvider(provider *CustomOAuthProvider) error { ...@@ -158,6 +191,57 @@ func validateCustomOAuthProvider(provider *CustomOAuthProvider) error {
if provider.Scopes == "" { if provider.Scopes == "" {
provider.Scopes = "openid profile email" 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 return nil
} }
model/log.go
...@@ -199,6 +199,49 @@ func RecordConsumeLog(c *gin.Context, userId int, params RecordConsumeLogParams) ...@@ -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) { 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 var tx *gorm.DB
if logType == LogTypeUnknown { if logType == LogTypeUnknown {
...@@ -252,8 +295,24 @@ func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName ...@@ -252,8 +295,24 @@ func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName
Id int `gorm:"column:id"` Id int `gorm:"column:id"`
Name string `gorm:"column:name"` Name string `gorm:"column:name"`
} }
if err = DB.Table("channels").Select("id, name").Where("id IN ?", channelIds.Items()).Find(&channels).Error; err != nil { if common.MemoryCacheEnabled {
return logs, total, err // 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)) channelMap := make(map[int]string, len(channels))
for _, channel := range channels { for _, channel := range channels {
......
model/midjourney.go
...@@ -157,6 +157,19 @@ func (midjourney *Midjourney) Update() error { ...@@ -157,6 +157,19 @@ func (midjourney *Midjourney) Update() error {
return err 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 { func MjBulkUpdate(mjIds []string, params map[string]any) error {
return DB.Model(&Midjourney{}). return DB.Model(&Midjourney{}).
Where("mj_id in (?)", mjIds). Where("mj_id in (?)", mjIds).
......
model/task.go
package model package model
import ( import (
"bytes"
"database/sql/driver" "database/sql/driver"
"encoding/json" "encoding/json"
"time" "time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto" "github.com/QuantumNous/new-api/dto"
commonRelay "github.com/QuantumNous/new-api/relay/common" commonRelay "github.com/QuantumNous/new-api/relay/common"
...@@ -64,13 +66,12 @@ type Task struct { ...@@ -64,13 +66,12 @@ type Task struct {
} }
func (t *Task) SetData(data any) { func (t *Task) SetData(data any) {
b, _ := json.Marshal(data) b, _ := common.Marshal(data)
t.Data = json.RawMessage(b) t.Data = json.RawMessage(b)
} }
func (t *Task) GetData(v any) error { func (t *Task) GetData(v any) error {
err := json.Unmarshal(t.Data, &v) return common.Unmarshal(t.Data, &v)
return err
} }
type Properties struct { type Properties struct {
...@@ -85,18 +86,59 @@ func (m *Properties) Scan(val interface{}) error { ...@@ -85,18 +86,59 @@ func (m *Properties) Scan(val interface{}) error {
*m = Properties{} *m = Properties{}
return nil return nil
} }
return json.Unmarshal(bytesValue, m) return common.Unmarshal(bytesValue, m)
} }
func (m Properties) Value() (driver.Value, error) { func (m Properties) Value() (driver.Value, error) {
if m == (Properties{}) { if m == (Properties{}) {
return nil, nil return nil, nil
} }
return json.Marshal(m) return common.Marshal(m)
} }
type TaskPrivateData struct { 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 { func (p *TaskPrivateData) Scan(val interface{}) error {
...@@ -104,14 +146,14 @@ func (p *TaskPrivateData) Scan(val interface{}) error { ...@@ -104,14 +146,14 @@ func (p *TaskPrivateData) Scan(val interface{}) error {
if len(bytesValue) == 0 { if len(bytesValue) == 0 {
return nil return nil
} }
return json.Unmarshal(bytesValue, p) return common.Unmarshal(bytesValue, p)
} }
func (p TaskPrivateData) Value() (driver.Value, error) { func (p TaskPrivateData) Value() (driver.Value, error) {
if (p == TaskPrivateData{}) { if (p == TaskPrivateData{}) {
return nil, nil return nil, nil
} }
return json.Marshal(p) return common.Marshal(p)
} }
// SyncTaskQueryParams 用于包含所有搜索条件的结构体，可以根据需求添加更多字段 // SyncTaskQueryParams 用于包含所有搜索条件的结构体，可以根据需求添加更多字段
...@@ -142,7 +184,16 @@ func InitTask(platform constant.TaskPlatform, relayInfo *commonRelay.RelayInfo) ...@@ -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{ t := &Task{
TaskID: taskID,
UserId: relayInfo.UserId, UserId: relayInfo.UserId,
Group: relayInfo.UsingGroup, Group: relayInfo.UsingGroup,
SubmitTime: time.Now().Unix(), SubmitTime: time.Now().Unix(),
...@@ -237,6 +288,20 @@ func TaskGetAllTasks(startIdx int, num int, queryParams SyncTaskQueryParams) []* ...@@ -237,6 +288,20 @@ func TaskGetAllTasks(startIdx int, num int, queryParams SyncTaskQueryParams) []*
return tasks 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 { func GetAllUnFinishSyncTasks(limit int) []*Task {
var tasks []*Task var tasks []*Task
var err error var err error
...@@ -291,40 +356,70 @@ func GetByTaskIds(userId int, taskIds []any) ([]*Task, error) { ...@@ -291,40 +356,70 @@ func GetByTaskIds(userId int, taskIds []any) ([]*Task, error) {
return task, nil 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 { func (Task *Task) Insert() error {
var err error var err error
err = DB.Create(Task).Error err = DB.Create(Task).Error
return err 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 { func (Task *Task) Update() error {
var err error var err error
err = DB.Save(Task).Error err = DB.Save(Task).Error
return err return err
} }
func TaskBulkUpdate(TaskIds []string, params map[string]any) error { // UpdateWithStatus performs a conditional UPDATE guarded by fromStatus (CAS).
if len(TaskIds) == 0 { // Returns (true, nil) if this caller won the update, (false, nil) if
return nil // another process already moved the task out of fromStatus.
} //
return DB.Model(&Task{}). // Uses Model().Select("*").Updates() instead of Save() because GORM's Save
Where("task_id in (?)", TaskIds). // falls back to INSERT ON CONFLICT when the WHERE-guarded UPDATE matches
Updates(params).Error // 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)
func TaskBulkUpdateByTaskIds(taskIDs []int64, params map[string]any) error { if result.Error != nil {
if len(taskIDs) == 0 { return false, result.Error
return nil }
} return result.RowsAffected > 0, nil
return DB.Model(&Task{}).
Where("id in (?)", taskIDs).
Updates(params).Error
} }
// 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 { func TaskBulkUpdateByID(ids []int64, params map[string]any) error {
if len(ids) == 0 { if len(ids) == 0 {
return nil return nil
...@@ -339,37 +434,6 @@ type TaskQuotaUsage struct { ...@@ -339,37 +434,6 @@ type TaskQuotaUsage struct {
Count float64 `json:"count"` 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) // TaskCountAllTasks returns total tasks that match the given query params (admin usage)
func TaskCountAllTasks(queryParams SyncTaskQueryParams) int64 { func TaskCountAllTasks(queryParams SyncTaskQueryParams) int64 {
var total int64 var total int64
...@@ -438,6 +502,6 @@ func (t *Task) ToOpenAIVideo() *dto.OpenAIVideo { ...@@ -438,6 +502,6 @@ func (t *Task) ToOpenAIVideo() *dto.OpenAIVideo {
openAIVideo.SetProgressStr(t.Progress) openAIVideo.SetProgressStr(t.Progress)
openAIVideo.CreatedAt = t.CreatedAt openAIVideo.CreatedAt = t.CreatedAt
openAIVideo.CompletedAt = t.UpdatedAt openAIVideo.CompletedAt = t.UpdatedAt
openAIVideo.SetMetadata("url", t.FailReason) openAIVideo.SetMetadata("url", t.GetResultURL())
return openAIVideo 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) { ...@@ -360,7 +360,7 @@ func DeleteTokenById(id int, userId int) (err error) {
return token.Delete() 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 { if quota < 0 {
return errors.New("quota 不能为负数！") return errors.New("quota 不能为负数！")
} }
...@@ -373,10 +373,10 @@ func IncreaseTokenQuota(id int, key string, quota int) (err error) { ...@@ -373,10 +373,10 @@ func IncreaseTokenQuota(id int, key string, quota int) (err error) {
}) })
} }
if common.BatchUpdateEnabled { if common.BatchUpdateEnabled {
addNewRecord(BatchUpdateTypeTokenQuota, id, quota) addNewRecord(BatchUpdateTypeTokenQuota, tokenId, quota)
return nil return nil
} }
return increaseTokenQuota(id, quota) return increaseTokenQuota(tokenId, quota)
} }
func increaseTokenQuota(id int, quota int) (err error) { func increaseTokenQuota(id int, quota int) (err error) {
......
model/user.go
package model package model
import ( import (
"database/sql"
"encoding/json" "encoding/json"
"errors" "errors"
"fmt" "fmt"
...@@ -15,6 +16,8 @@ import ( ...@@ -15,6 +16,8 @@ import (
"gorm.io/gorm" "gorm.io/gorm"
) )
const UserNameMaxLength = 20
// User if you add sensitive fields, don't forget to clean them in setupLogin function. // 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! // Otherwise, the sensitive information will be saved on local storage in plain text!
type User struct { type User struct {
...@@ -536,6 +539,37 @@ func (user *User) Edit(updatePassword bool) error { ...@@ -536,6 +539,37 @@ func (user *User) Edit(updatePassword bool) error {
return updateUserCache(*user) 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 { func (user *User) Delete() error {
if user.Id == 0 { if user.Id == 0 {
return errors.New("id 为空！") return errors.New("id 为空！")
...@@ -820,10 +854,17 @@ func GetUserSetting(id int, fromDB bool) (settingMap dto.UserSetting, err error) ...@@ -820,10 +854,17 @@ func GetUserSetting(id int, fromDB bool) (settingMap dto.UserSetting, err error)
// Don't return error - fall through to DB // Don't return error - fall through to DB
} }
fromDB = true 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 { if err != nil {
return settingMap, err return settingMap, err
} }
if safeSetting.Valid {
setting = safeSetting.String
} else {
setting = ""
}
userBase := &UserBase{ userBase := &UserBase{
Setting: setting, Setting: setting,
} }
......
oauth/types.go
...@@ -57,3 +57,12 @@ func NewOAuthErrorWithRaw(msgKey string, params map[string]any, rawError string) ...@@ -57,3 +57,12 @@ func NewOAuthErrorWithRaw(msgKey string, params map[string]any, rawError string)
RawError: rawError, 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 { ...@@ -36,6 +36,32 @@ type TaskAdaptor interface {
ValidateRequestAndSetAction(c *gin.Context, info *relaycommon.RelayInfo) *dto.TaskError 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) BuildRequestURL(info *relaycommon.RelayInfo) (string, error)
BuildRequestHeader(c *gin.Context, req *http.Request, info *relaycommon.RelayInfo) error BuildRequestHeader(c *gin.Context, req *http.Request, info *relaycommon.RelayInfo) error
BuildRequestBody(c *gin.Context, info *relaycommon.RelayInfo) (io.Reader, error) BuildRequestBody(c *gin.Context, info *relaycommon.RelayInfo) (io.Reader, error)
...@@ -46,9 +72,9 @@ type TaskAdaptor interface { ...@@ -46,9 +72,9 @@ type TaskAdaptor interface {
GetModelList() []string GetModelList() []string
GetChannelName() string GetChannelName() string
// FetchTask // ── Polling ──────────────────────────────────────────────────────
FetchTask(baseUrl, key string, body map[string]any, proxy string) (*http.Response, error)
FetchTask(baseUrl, key string, body map[string]any, proxy string) (*http.Response, error)
ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error)
} }
......
relay/channel/api_request.go
...@@ -61,8 +61,9 @@ var passthroughSkipHeaderNamesLower = map[string]struct{}{ ...@@ -61,8 +61,9 @@ var passthroughSkipHeaderNamesLower = map[string]struct{}{
"cookie": {}, "cookie": {},
// Additional headers that should not be forwarded by name-matching passthrough rules. // Additional headers that should not be forwarded by name-matching passthrough rules.
"host": {}, "host": {},
"content-length": {}, "content-length": {},
"accept-encoding": {},
// Do not passthrough credentials by wildcard/regex. // Do not passthrough credentials by wildcard/regex.
"authorization": {}, "authorization": {},
......
relay/channel/api_request_test.go
...@@ -110,3 +110,30 @@ func TestProcessHeaderOverride_RuntimeOverrideHasPriority(t *testing.T) { ...@@ -110,3 +110,30 @@ func TestProcessHeaderOverride_RuntimeOverrideHasPriority(t *testing.T) {
_, ok := headers["X-Legacy"] _, ok := headers["X-Legacy"]
require.False(t, ok) 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 ...@@ -42,22 +42,7 @@ func GeminiTextGenerationHandler(c *gin.Context, info *relaycommon.RelayInfo, re
} }
// 计算使用量（基于 UsageMetadata） // 计算使用量（基于 UsageMetadata）
usage := dto.Usage{ usage := buildUsageFromGeminiMetadata(geminiResponse.UsageMetadata, info.GetEstimatePromptTokens())
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
}
}
service.IOCopyBytesGracefully(c, resp, responseBody) service.IOCopyBytesGracefully(c, resp, responseBody)
......
relay/channel/gemini/relay-gemini.go
...@@ -1032,6 +1032,46 @@ func getResponseToolCall(item *dto.GeminiPart) *dto.ToolCallResponse { ...@@ -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 { func responseGeminiChat2OpenAI(c *gin.Context, response *dto.GeminiChatResponse) *dto.OpenAITextResponse {
fullTextResponse := dto.OpenAITextResponse{ fullTextResponse := dto.OpenAITextResponse{
Id: helper.GetResponseID(c), Id: helper.GetResponseID(c),
...@@ -1272,18 +1312,8 @@ func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http ...@@ -1272,18 +1312,8 @@ func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http
// 更新使用量统计 // 更新使用量统计
if geminiResponse.UsageMetadata.TotalTokenCount != 0 { if geminiResponse.UsageMetadata.TotalTokenCount != 0 {
usage.PromptTokens = geminiResponse.UsageMetadata.PromptTokenCount mappedUsage := buildUsageFromGeminiMetadata(geminiResponse.UsageMetadata, info.GetEstimatePromptTokens())
usage.CompletionTokens = geminiResponse.UsageMetadata.CandidatesTokenCount + geminiResponse.UsageMetadata.ThoughtsTokenCount *usage = mappedUsage
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
}
}
} }
return callback(data, &geminiResponse) return callback(data, &geminiResponse)
...@@ -1295,11 +1325,6 @@ func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http ...@@ -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 usage.CompletionTokens <= 0 {
if info.ReceivedResponseCount > 0 { if info.ReceivedResponseCount > 0 {
usage = service.ResponseText2Usage(c, responseText.String(), info.UpstreamModelName, info.GetEstimatePromptTokens()) 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 ...@@ -1416,21 +1441,7 @@ func GeminiChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.R
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
} }
if len(geminiResponse.Candidates) == 0 { if len(geminiResponse.Candidates) == 0 {
usage := dto.Usage{ usage := buildUsageFromGeminiMetadata(geminiResponse.UsageMetadata, info.GetEstimatePromptTokens())
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()
}
var newAPIError *types.NewAPIError var newAPIError *types.NewAPIError
if geminiResponse.PromptFeedback != nil && geminiResponse.PromptFeedback.BlockReason != nil { if geminiResponse.PromptFeedback != nil && geminiResponse.PromptFeedback.BlockReason != nil {
...@@ -1466,23 +1477,7 @@ func GeminiChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.R ...@@ -1466,23 +1477,7 @@ func GeminiChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.R
} }
fullTextResponse := responseGeminiChat2OpenAI(c, &geminiResponse) fullTextResponse := responseGeminiChat2OpenAI(c, &geminiResponse)
fullTextResponse.Model = info.UpstreamModelName fullTextResponse.Model = info.UpstreamModelName
usage := dto.Usage{ usage := buildUsageFromGeminiMetadata(geminiResponse.UsageMetadata, info.GetEstimatePromptTokens())
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
}
}
fullTextResponse.Usage = usage 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 ( ...@@ -10,6 +10,7 @@ import (
"github.com/QuantumNous/new-api/dto" "github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/relay/channel" "github.com/QuantumNous/new-api/relay/channel"
"github.com/QuantumNous/new-api/relay/channel/claude"
"github.com/QuantumNous/new-api/relay/channel/openai" "github.com/QuantumNous/new-api/relay/channel/openai"
relaycommon "github.com/QuantumNous/new-api/relay/common" relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/constant" "github.com/QuantumNous/new-api/relay/constant"
...@@ -26,7 +27,8 @@ func (a *Adaptor) ConvertGeminiRequest(*gin.Context, *relaycommon.RelayInfo, *dt ...@@ -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) { 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) { 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 ...@@ -119,8 +121,14 @@ func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycom
return handleTTSResponse(c, resp, info) return handleTTSResponse(c, resp, info)
} }
adaptor := openai.Adaptor{} switch info.RelayFormat {
return adaptor.DoResponse(c, resp, info) 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 { func (a *Adaptor) GetModelList() []string {
......
relay/channel/minimax/relay-minimax.go
...@@ -6,6 +6,7 @@ import ( ...@@ -6,6 +6,7 @@ import (
channelconstant "github.com/QuantumNous/new-api/constant" channelconstant "github.com/QuantumNous/new-api/constant"
relaycommon "github.com/QuantumNous/new-api/relay/common" relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/constant" "github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/types"
) )
func GetRequestURL(info *relaycommon.RelayInfo) (string, error) { func GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
...@@ -13,13 +14,17 @@ func GetRequestURL(info *relaycommon.RelayInfo) (string, error) { ...@@ -13,13 +14,17 @@ func GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
if baseUrl == "" { if baseUrl == "" {
baseUrl = channelconstant.ChannelBaseURLs[channelconstant.ChannelTypeMiniMax] baseUrl = channelconstant.ChannelBaseURLs[channelconstant.ChannelTypeMiniMax]
} }
switch info.RelayFormat {
switch info.RelayMode { case types.RelayFormatClaude:
case constant.RelayModeChatCompletions: return fmt.Sprintf("%s/anthropic/v1/messages", info.ChannelBaseUrl), nil
return fmt.Sprintf("%s/v1/text/chatcompletion_v2", baseUrl), nil
case constant.RelayModeAudioSpeech:
return fmt.Sprintf("%s/v1/t2a_v2", baseUrl), nil
default: 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 ( ...@@ -13,6 +13,7 @@ import (
"github.com/QuantumNous/new-api/logger" "github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/relay/channel" "github.com/QuantumNous/new-api/relay/channel"
"github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
relaycommon "github.com/QuantumNous/new-api/relay/common" relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/service"
"github.com/samber/lo" "github.com/samber/lo"
...@@ -108,10 +109,10 @@ type AliMetadata struct { ...@@ -108,10 +109,10 @@ type AliMetadata struct {
// ============================ // ============================
type TaskAdaptor struct { type TaskAdaptor struct {
taskcommon.BaseBilling
ChannelType int ChannelType int
apiKey string apiKey string
baseURL string baseURL string
aliReq *AliVideoRequest
} }
func (a *TaskAdaptor) Init(info *relaycommon.RelayInfo) { func (a *TaskAdaptor) Init(info *relaycommon.RelayInfo) {
...@@ -121,17 +122,7 @@ 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) { func (a *TaskAdaptor) ValidateRequestAndSetAction(c *gin.Context, info *relaycommon.RelayInfo) (taskErr *dto.TaskError) {
// 阿里通义万相支持 JSON 格式，不使用 multipart // ValidateMultipartDirect 负责解析并将原始 TaskSubmitReq 存入 context
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)
return relaycommon.ValidateMultipartDirect(c, info) return relaycommon.ValidateMultipartDirect(c, info)
} }
...@@ -148,11 +139,21 @@ func (a *TaskAdaptor) BuildRequestHeader(c *gin.Context, req *http.Request, 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) { 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 { 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 return bytes.NewReader(bodyBytes), nil
} }
...@@ -252,8 +253,12 @@ func ProcessAliOtherRatios(aliReq *AliVideoRequest) (map[string]float64, error) ...@@ -252,8 +253,12 @@ func ProcessAliOtherRatios(aliReq *AliVideoRequest) (map[string]float64, error)
} }
func (a *TaskAdaptor) convertToAliRequest(info *relaycommon.RelayInfo, req relaycommon.TaskSubmitReq) (*AliVideoRequest, error) { func (a *TaskAdaptor) convertToAliRequest(info *relaycommon.RelayInfo, req relaycommon.TaskSubmitReq) (*AliVideoRequest, error) {
upstreamModel := req.Model
if info.IsModelMapped {
upstreamModel = info.UpstreamModelName
}
aliReq := &AliVideoRequest{ aliReq := &AliVideoRequest{
Model: req.Model, Model: upstreamModel,
Input: AliVideoInput{ Input: AliVideoInput{
Prompt: req.Prompt, Prompt: req.Prompt,
ImgURL: req.InputReference, ImgURL: req.InputReference,
...@@ -331,23 +336,37 @@ func (a *TaskAdaptor) convertToAliRequest(info *relaycommon.RelayInfo, req relay ...@@ -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") return nil, errors.New("can't change model with metadata")
} }
info.PriceData.OtherRatios = map[string]float64{ return aliReq, nil
"seconds": float64(aliReq.Parameters.Duration), }
// 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) ratios, err := ProcessAliOtherRatios(aliReq)
if err != nil { if err != nil {
return nil, err return otherRatios
} }
for s, f := range ratios { for k, v := range ratios {
info.PriceData.OtherRatios[s] = f otherRatios[k] = v
} }
return otherRatios
return aliReq, nil
} }
// DoRequest delegates to common helper // DoRequest delegates to common helper
...@@ -384,7 +403,8 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela ...@@ -384,7 +403,8 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
// 转换为 OpenAI 格式响应 // 转换为 OpenAI 格式响应
openAIResp := dto.NewOpenAIVideo() openAIResp := dto.NewOpenAIVideo()
openAIResp.ID = aliResp.Output.TaskID openAIResp.ID = info.PublicTaskID
openAIResp.TaskID = info.PublicTaskID
openAIResp.Model = c.GetString("model") openAIResp.Model = c.GetString("model")
if openAIResp.Model == "" && info != nil { if openAIResp.Model == "" && info != nil {
openAIResp.Model = info.OriginModelName openAIResp.Model = info.OriginModelName
......
relay/channel/task/doubao/adaptor.go
...@@ -2,7 +2,6 @@ package doubao ...@@ -2,7 +2,6 @@ package doubao
import ( import (
"bytes" "bytes"
"encoding/json"
"fmt" "fmt"
"io" "io"
"net/http" "net/http"
...@@ -14,6 +13,7 @@ import ( ...@@ -14,6 +13,7 @@ import (
"github.com/QuantumNous/new-api/dto" "github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/relay/channel" "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" relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/service"
...@@ -89,6 +89,7 @@ type responseTask struct { ...@@ -89,6 +89,7 @@ type responseTask struct {
// ============================ // ============================
type TaskAdaptor struct { type TaskAdaptor struct {
taskcommon.BaseBilling
ChannelType int ChannelType int
apiKey string apiKey string
baseURL string baseURL string
...@@ -130,8 +131,12 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn ...@@ -130,8 +131,12 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
if err != nil { if err != nil {
return nil, errors.Wrap(err, "convert request payload failed") return nil, errors.Wrap(err, "convert request payload failed")
} }
info.UpstreamModelName = body.Model if info.IsModelMapped {
data, err := json.Marshal(body) body.Model = info.UpstreamModelName
} else {
info.UpstreamModelName = body.Model
}
data, err := common.Marshal(body)
if err != nil { if err != nil {
return nil, err return nil, err
} }
...@@ -154,7 +159,7 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela ...@@ -154,7 +159,7 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
// Parse Doubao response // Parse Doubao response
var dResp responsePayload 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) taskErr = service.TaskErrorWrapper(errors.Wrapf(err, "body: %s", responseBody), "unmarshal_response_body_failed", http.StatusInternalServerError)
return return
} }
...@@ -165,8 +170,8 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela ...@@ -165,8 +170,8 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
} }
ov := dto.NewOpenAIVideo() ov := dto.NewOpenAIVideo()
ov.ID = dResp.ID ov.ID = info.PublicTaskID
ov.TaskID = dResp.ID ov.TaskID = info.PublicTaskID
ov.CreatedAt = time.Now().Unix() ov.CreatedAt = time.Now().Unix()
ov.Model = info.OriginModelName ov.Model = info.OriginModelName
...@@ -234,12 +239,7 @@ func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (* ...@@ -234,12 +239,7 @@ func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*
} }
metadata := req.Metadata metadata := req.Metadata
medaBytes, err := json.Marshal(metadata) if err := taskcommon.UnmarshalMetadata(metadata, &r); err != nil {
if err != nil {
return nil, errors.Wrap(err, "metadata marshal metadata failed")
}
err = json.Unmarshal(medaBytes, &r)
if err != nil {
return nil, errors.Wrap(err, "unmarshal metadata failed") return nil, errors.Wrap(err, "unmarshal metadata failed")
} }
...@@ -248,7 +248,7 @@ func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (* ...@@ -248,7 +248,7 @@ func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*
func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) { func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) {
resTask := responseTask{} 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") return nil, errors.Wrap(err, "unmarshal task result failed")
} }
...@@ -286,7 +286,7 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e ...@@ -286,7 +286,7 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, error) { func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, error) {
var dResp responseTask 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") return nil, errors.Wrap(err, "unmarshal doubao task data failed")
} }
...@@ -307,6 +307,5 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, erro ...@@ -307,6 +307,5 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, erro
} }
} }
jsonData, _ := common.Marshal(openAIVideo) return common.Marshal(openAIVideo)
return jsonData, nil
} }
relay/channel/task/gemini/adaptor.go
...@@ -2,8 +2,6 @@ package gemini ...@@ -2,8 +2,6 @@ package gemini
import ( import (
"bytes" "bytes"
"encoding/base64"
"encoding/json"
"fmt" "fmt"
"io" "io"
"net/http" "net/http"
...@@ -16,10 +14,10 @@ import ( ...@@ -16,10 +14,10 @@ import (
"github.com/QuantumNous/new-api/dto" "github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/relay/channel" "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" relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting/model_setting" "github.com/QuantumNous/new-api/setting/model_setting"
"github.com/QuantumNous/new-api/setting/system_setting"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/pkg/errors" "github.com/pkg/errors"
) )
...@@ -87,6 +85,7 @@ type operationResponse struct { ...@@ -87,6 +85,7 @@ type operationResponse struct {
// ============================ // ============================
type TaskAdaptor struct { type TaskAdaptor struct {
taskcommon.BaseBilling
ChannelType int ChannelType int
apiKey string apiKey string
baseURL string baseURL string
...@@ -106,7 +105,7 @@ func (a *TaskAdaptor) ValidateRequestAndSetAction(c *gin.Context, info *relaycom ...@@ -106,7 +105,7 @@ func (a *TaskAdaptor) ValidateRequestAndSetAction(c *gin.Context, info *relaycom
// BuildRequestURL constructs the upstream URL. // BuildRequestURL constructs the upstream URL.
func (a *TaskAdaptor) BuildRequestURL(info *relaycommon.RelayInfo) (string, error) { func (a *TaskAdaptor) BuildRequestURL(info *relaycommon.RelayInfo) (string, error) {
modelName := info.OriginModelName modelName := info.UpstreamModelName
version := model_setting.GetGeminiVersionSetting(modelName) version := model_setting.GetGeminiVersionSetting(modelName)
return fmt.Sprintf( return fmt.Sprintf(
...@@ -145,16 +144,11 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn ...@@ -145,16 +144,11 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
} }
metadata := req.Metadata metadata := req.Metadata
medaBytes, err := json.Marshal(metadata) if err := taskcommon.UnmarshalMetadata(metadata, &body.Parameters); err != nil {
if err != nil {
return nil, errors.Wrap(err, "metadata marshal metadata failed")
}
err = json.Unmarshal(medaBytes, &body.Parameters)
if err != nil {
return nil, errors.Wrap(err, "unmarshal metadata failed") return nil, errors.Wrap(err, "unmarshal metadata failed")
} }
data, err := json.Marshal(body) data, err := common.Marshal(body)
if err != nil { if err != nil {
return nil, err return nil, err
} }
...@@ -175,16 +169,16 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela ...@@ -175,16 +169,16 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
_ = resp.Body.Close() _ = resp.Body.Close()
var s submitResponse 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) return "", nil, service.TaskErrorWrapper(err, "unmarshal_response_failed", http.StatusInternalServerError)
} }
if strings.TrimSpace(s.Name) == "" { if strings.TrimSpace(s.Name) == "" {
return "", nil, service.TaskErrorWrapper(fmt.Errorf("missing operation name"), "invalid_response", http.StatusInternalServerError) 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 := dto.NewOpenAIVideo()
ov.ID = taskID ov.ID = info.PublicTaskID
ov.TaskID = taskID ov.TaskID = info.PublicTaskID
ov.CreatedAt = time.Now().Unix() ov.CreatedAt = time.Now().Unix()
ov.Model = info.OriginModelName ov.Model = info.OriginModelName
c.JSON(http.StatusOK, ov) c.JSON(http.StatusOK, ov)
...@@ -206,7 +200,7 @@ func (a *TaskAdaptor) FetchTask(baseUrl, key string, body map[string]any, proxy ...@@ -206,7 +200,7 @@ func (a *TaskAdaptor) FetchTask(baseUrl, key string, body map[string]any, proxy
return nil, fmt.Errorf("invalid task_id") return nil, fmt.Errorf("invalid task_id")
} }
upstreamName, err := decodeLocalTaskID(taskID) upstreamName, err := taskcommon.DecodeLocalTaskID(taskID)
if err != nil { if err != nil {
return nil, fmt.Errorf("decode task_id failed: %w", err) 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 ...@@ -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) { func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) {
var op operationResponse 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) return nil, fmt.Errorf("unmarshal operation response failed: %w", err)
} }
...@@ -254,9 +248,8 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e ...@@ -254,9 +248,8 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
ti.Status = model.TaskStatusSuccess ti.Status = model.TaskStatusSuccess
ti.Progress = "100%" ti.Progress = "100%"
taskID := encodeLocalTaskID(op.Name) ti.TaskID = taskcommon.EncodeLocalTaskID(op.Name)
ti.TaskID = taskID // Url intentionally left empty — the caller constructs the proxy URL using the public task ID
ti.Url = fmt.Sprintf("%s/v1/videos/%s/content", system_setting.ServerAddress, taskID)
// Extract URL from generateVideoResponse if available // Extract URL from generateVideoResponse if available
if len(op.Response.GenerateVideoResponse.GeneratedSamples) > 0 { if len(op.Response.GenerateVideoResponse.GeneratedSamples) > 0 {
...@@ -269,7 +262,10 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e ...@@ -269,7 +262,10 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
} }
func (a *TaskAdaptor) ConvertToOpenAIVideo(task *model.Task) ([]byte, error) { 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 { if err != nil {
upstreamName = "" upstreamName = ""
} }
...@@ -297,18 +293,6 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(task *model.Task) ([]byte, error) { ...@@ -297,18 +293,6 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(task *model.Task) ([]byte, error) {
// helpers // 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/`) var modelRe = regexp.MustCompile(`models/([^/]+)/operations/`)
func extractModelFromOperationName(name string) string { func extractModelFromOperationName(name string) string {
......
relay/channel/task/hailuo/adaptor.go
...@@ -2,7 +2,6 @@ package hailuo ...@@ -2,7 +2,6 @@ package hailuo
import ( import (
"bytes" "bytes"
"encoding/json"
"fmt" "fmt"
"io" "io"
"net/http" "net/http"
...@@ -18,12 +17,14 @@ import ( ...@@ -18,12 +17,14 @@ import (
"github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto" "github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/relay/channel" "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" relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/service"
) )
// https://platform.minimaxi.com/docs/api-reference/video-generation-intro // https://platform.minimaxi.com/docs/api-reference/video-generation-intro
type TaskAdaptor struct { type TaskAdaptor struct {
taskcommon.BaseBilling
ChannelType int ChannelType int
apiKey string apiKey string
baseURL string baseURL string
...@@ -60,12 +61,12 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn ...@@ -60,12 +61,12 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
return nil, fmt.Errorf("invalid request type in context") return nil, fmt.Errorf("invalid request type in context")
} }
body, err := a.convertToRequestPayload(&req) body, err := a.convertToRequestPayload(&req, info)
if err != nil { if err != nil {
return nil, errors.Wrap(err, "convert request payload failed") return nil, errors.Wrap(err, "convert request payload failed")
} }
data, err := json.Marshal(body) data, err := common.Marshal(body)
if err != nil { if err != nil {
return nil, err return nil, err
} }
...@@ -86,7 +87,7 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela ...@@ -86,7 +87,7 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
_ = resp.Body.Close() _ = resp.Body.Close()
var hResp VideoResponse 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) taskErr = service.TaskErrorWrapper(errors.Wrapf(err, "body: %s", responseBody), "unmarshal_response_body_failed", http.StatusInternalServerError)
return return
} }
...@@ -101,8 +102,8 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela ...@@ -101,8 +102,8 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
} }
ov := dto.NewOpenAIVideo() ov := dto.NewOpenAIVideo()
ov.ID = hResp.TaskID ov.ID = info.PublicTaskID
ov.TaskID = hResp.TaskID ov.TaskID = info.PublicTaskID
ov.CreatedAt = time.Now().Unix() ov.CreatedAt = time.Now().Unix()
ov.Model = info.OriginModelName ov.Model = info.OriginModelName
...@@ -141,8 +142,8 @@ func (a *TaskAdaptor) GetChannelName() string { ...@@ -141,8 +142,8 @@ func (a *TaskAdaptor) GetChannelName() string {
return ChannelName return ChannelName
} }
func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*VideoRequest, error) { func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq, info *relaycommon.RelayInfo) (*VideoRequest, error) {
modelConfig := GetModelConfig(req.Model) modelConfig := GetModelConfig(info.UpstreamModelName)
duration := DefaultDuration duration := DefaultDuration
if req.Duration > 0 { if req.Duration > 0 {
duration = req.Duration duration = req.Duration
...@@ -153,7 +154,7 @@ func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (* ...@@ -153,7 +154,7 @@ func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*
} }
videoRequest := &VideoRequest{ videoRequest := &VideoRequest{
Model: req.Model, Model: info.UpstreamModelName,
Prompt: req.Prompt, Prompt: req.Prompt,
Duration: &duration, Duration: &duration,
Resolution: resolution, Resolution: resolution,
...@@ -182,7 +183,7 @@ func (a *TaskAdaptor) parseResolutionFromSize(size string, modelConfig ModelConf ...@@ -182,7 +183,7 @@ func (a *TaskAdaptor) parseResolutionFromSize(size string, modelConfig ModelConf
func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) { func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) {
resTask := QueryTaskResponse{} 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") return nil, errors.Wrap(err, "unmarshal task result failed")
} }
...@@ -224,7 +225,7 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e ...@@ -224,7 +225,7 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, error) { func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, error) {
var hailuoResp QueryTaskResponse 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") return nil, errors.Wrap(err, "unmarshal hailuo task data failed")
} }
...@@ -271,7 +272,7 @@ func (a *TaskAdaptor) buildVideoURL(_, fileID string) string { ...@@ -271,7 +272,7 @@ func (a *TaskAdaptor) buildVideoURL(_, fileID string) string {
} }
var retrieveResp RetrieveFileResponse var retrieveResp RetrieveFileResponse
if err := json.Unmarshal(responseBody, &retrieveResp); err != nil { if err := common.Unmarshal(responseBody, &retrieveResp); err != nil {
return "" return ""
} }
......
relay/channel/task/jimeng/adaptor.go
...@@ -6,7 +6,6 @@ import ( ...@@ -6,7 +6,6 @@ import (
"crypto/sha256" "crypto/sha256"
"encoding/base64" "encoding/base64"
"encoding/hex" "encoding/hex"
"encoding/json"
"fmt" "fmt"
"io" "io"
"net/http" "net/http"
...@@ -25,6 +24,7 @@ import ( ...@@ -25,6 +24,7 @@ import (
"github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto" "github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/relay/channel" "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" relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/service"
) )
...@@ -77,6 +77,7 @@ const ( ...@@ -77,6 +77,7 @@ const (
// ============================ // ============================
type TaskAdaptor struct { type TaskAdaptor struct {
taskcommon.BaseBilling
ChannelType int ChannelType int
accessKey string accessKey string
secretKey string secretKey string
...@@ -164,11 +165,11 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn ...@@ -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 { if err != nil {
return nil, errors.Wrap(err, "convert request payload failed") return nil, errors.Wrap(err, "convert request payload failed")
} }
data, err := json.Marshal(body) data, err := common.Marshal(body)
if err != nil { if err != nil {
return nil, err return nil, err
} }
...@@ -191,7 +192,7 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela ...@@ -191,7 +192,7 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
// Parse Jimeng response // Parse Jimeng response
var jResp responsePayload 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) taskErr = service.TaskErrorWrapper(errors.Wrapf(err, "body: %s", responseBody), "unmarshal_response_body_failed", http.StatusInternalServerError)
return return
} }
...@@ -202,8 +203,8 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela ...@@ -202,8 +203,8 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
} }
ov := dto.NewOpenAIVideo() ov := dto.NewOpenAIVideo()
ov.ID = jResp.Data.TaskID ov.ID = info.PublicTaskID
ov.TaskID = jResp.Data.TaskID ov.TaskID = info.PublicTaskID
ov.CreatedAt = time.Now().Unix() ov.CreatedAt = time.Now().Unix()
ov.Model = info.OriginModelName ov.Model = info.OriginModelName
c.JSON(http.StatusOK, ov) c.JSON(http.StatusOK, ov)
...@@ -225,7 +226,7 @@ func (a *TaskAdaptor) FetchTask(baseUrl, key string, body map[string]any, proxy ...@@ -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 "req_key": "jimeng_vgfm_t2v_l20", // This is fixed value from doc: https://www.volcengine.com/docs/85621/1544774
"task_id": taskID, "task_id": taskID,
} }
payloadBytes, err := json.Marshal(payload) payloadBytes, err := common.Marshal(payload)
if err != nil { if err != nil {
return nil, errors.Wrap(err, "marshal fetch task payload failed") return nil, errors.Wrap(err, "marshal fetch task payload failed")
} }
...@@ -377,9 +378,9 @@ func hmacSHA256(key []byte, data []byte) []byte { ...@@ -377,9 +378,9 @@ func hmacSHA256(key []byte, data []byte) []byte {
return h.Sum(nil) 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{ r := requestPayload{
ReqKey: req.Model, ReqKey: info.UpstreamModelName,
Prompt: req.Prompt, Prompt: req.Prompt,
} }
...@@ -398,13 +399,7 @@ func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (* ...@@ -398,13 +399,7 @@ func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*
r.BinaryDataBase64 = req.Images r.BinaryDataBase64 = req.Images
} }
} }
metadata := req.Metadata if err := taskcommon.UnmarshalMetadata(req.Metadata, &r); err != nil {
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 {
return nil, errors.Wrap(err, "unmarshal metadata failed") return nil, errors.Wrap(err, "unmarshal metadata failed")
} }
...@@ -432,7 +427,7 @@ func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (* ...@@ -432,7 +427,7 @@ func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*
func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) { func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) {
resTask := responseTask{} 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") return nil, errors.Wrap(err, "unmarshal task result failed")
} }
taskResult := relaycommon.TaskInfo{} taskResult := relaycommon.TaskInfo{}
...@@ -458,7 +453,7 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e ...@@ -458,7 +453,7 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, error) { func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, error) {
var jimengResp responseTask 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") return nil, errors.Wrap(err, "unmarshal jimeng task data failed")
} }
...@@ -477,8 +472,7 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, erro ...@@ -477,8 +472,7 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, erro
} }
} }
jsonData, _ := common.Marshal(openAIVideo) return common.Marshal(openAIVideo)
return jsonData, nil
} }
func isNewAPIRelay(apiKey string) bool { func isNewAPIRelay(apiKey string) bool {
......
relay/channel/task/kling/adaptor.go
...@@ -2,7 +2,6 @@ package kling ...@@ -2,7 +2,6 @@ package kling
import ( import (
"bytes" "bytes"
"encoding/json"
"fmt" "fmt"
"io" "io"
"net/http" "net/http"
...@@ -21,6 +20,7 @@ import ( ...@@ -21,6 +20,7 @@ import (
"github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto" "github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/relay/channel" "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" relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/service"
) )
...@@ -97,6 +97,7 @@ type responsePayload struct { ...@@ -97,6 +97,7 @@ type responsePayload struct {
// ============================ // ============================
type TaskAdaptor struct { type TaskAdaptor struct {
taskcommon.BaseBilling
ChannelType int ChannelType int
apiKey string apiKey string
baseURL string baseURL string
...@@ -149,14 +150,14 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn ...@@ -149,14 +150,14 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
} }
req := v.(relaycommon.TaskSubmitReq) req := v.(relaycommon.TaskSubmitReq)
body, err := a.convertToRequestPayload(&req) body, err := a.convertToRequestPayload(&req, info)
if err != nil { if err != nil {
return nil, err return nil, err
} }
if body.Image == "" && body.ImageTail == "" { if body.Image == "" && body.ImageTail == "" {
c.Set("action", constant.TaskActionTextGenerate) c.Set("action", constant.TaskActionTextGenerate)
} }
data, err := json.Marshal(body) data, err := common.Marshal(body)
if err != nil { if err != nil {
return nil, err return nil, err
} }
...@@ -180,7 +181,7 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela ...@@ -180,7 +181,7 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
} }
var kResp responsePayload var kResp responsePayload
err = json.Unmarshal(responseBody, &kResp) err = common.Unmarshal(responseBody, &kResp)
if err != nil { if err != nil {
taskErr = service.TaskErrorWrapper(err, "unmarshal_response_failed", http.StatusInternalServerError) taskErr = service.TaskErrorWrapper(err, "unmarshal_response_failed", http.StatusInternalServerError)
return return
...@@ -190,8 +191,8 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela ...@@ -190,8 +191,8 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
return return
} }
ov := dto.NewOpenAIVideo() ov := dto.NewOpenAIVideo()
ov.ID = kResp.Data.TaskId ov.ID = info.PublicTaskID
ov.TaskID = kResp.Data.TaskId ov.TaskID = info.PublicTaskID
ov.CreatedAt = time.Now().Unix() ov.CreatedAt = time.Now().Unix()
ov.Model = info.OriginModelName ov.Model = info.OriginModelName
c.JSON(http.StatusOK, ov) c.JSON(http.StatusOK, ov)
...@@ -247,15 +248,15 @@ func (a *TaskAdaptor) GetChannelName() string { ...@@ -247,15 +248,15 @@ func (a *TaskAdaptor) GetChannelName() string {
// helpers // helpers
// ============================ // ============================
func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*requestPayload, error) { func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq, info *relaycommon.RelayInfo) (*requestPayload, error) {
r := requestPayload{ r := requestPayload{
Prompt: req.Prompt, Prompt: req.Prompt,
Image: req.Image, Image: req.Image,
Mode: defaultString(req.Mode, "std"), Mode: taskcommon.DefaultString(req.Mode, "std"),
Duration: fmt.Sprintf("%d", defaultInt(req.Duration, 5)), Duration: fmt.Sprintf("%d", taskcommon.DefaultInt(req.Duration, 5)),
AspectRatio: a.getAspectRatio(req.Size), AspectRatio: a.getAspectRatio(req.Size),
ModelName: req.Model, ModelName: info.UpstreamModelName,
Model: req.Model, // Keep consistent with model_name, double writing improves compatibility Model: info.UpstreamModelName,
CfgScale: 0.5, CfgScale: 0.5,
StaticMask: "", StaticMask: "",
DynamicMasks: []DynamicMask{}, DynamicMasks: []DynamicMask{},
...@@ -265,14 +266,9 @@ func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (* ...@@ -265,14 +266,9 @@ func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*
} }
if r.ModelName == "" { if r.ModelName == "" {
r.ModelName = "kling-v1" r.ModelName = "kling-v1"
r.Model = "kling-v1"
} }
metadata := req.Metadata if err := taskcommon.UnmarshalMetadata(req.Metadata, &r); err != nil {
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 {
return nil, errors.Wrap(err, "unmarshal metadata failed") return nil, errors.Wrap(err, "unmarshal metadata failed")
} }
return &r, nil return &r, nil
...@@ -291,20 +287,6 @@ func (a *TaskAdaptor) getAspectRatio(size string) string { ...@@ -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 // JWT helpers
// ============================ // ============================
...@@ -340,7 +322,7 @@ func (a *TaskAdaptor) createJWTTokenWithKey(apiKey string) (string, error) { ...@@ -340,7 +322,7 @@ func (a *TaskAdaptor) createJWTTokenWithKey(apiKey string) (string, error) {
func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) { func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) {
taskInfo := &relaycommon.TaskInfo{} taskInfo := &relaycommon.TaskInfo{}
resPayload := responsePayload{} resPayload := responsePayload{}
err := json.Unmarshal(respBody, &resPayload) err := common.Unmarshal(respBody, &resPayload)
if err != nil { if err != nil {
return nil, errors.Wrap(err, "failed to unmarshal response body") return nil, errors.Wrap(err, "failed to unmarshal response body")
} }
...@@ -374,7 +356,7 @@ func isNewAPIRelay(apiKey string) bool { ...@@ -374,7 +356,7 @@ func isNewAPIRelay(apiKey string) bool {
func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, error) { func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, error) {
var klingResp responsePayload 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") return nil, errors.Wrap(err, "unmarshal kling task data failed")
} }
...@@ -401,6 +383,5 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, erro ...@@ -401,6 +383,5 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, erro
Code: fmt.Sprintf("%d", klingResp.Code), Code: fmt.Sprintf("%d", klingResp.Code),
} }
} }
jsonData, _ := common.Marshal(openAIVideo) return common.Marshal(openAIVideo)
return jsonData, nil
} }
relay/channel/task/sora/adaptor.go
package sora package sora
import ( import (
"bytes"
"fmt" "fmt"
"io" "io"
"mime/multipart"
"net/http" "net/http"
"net/textproto"
"strconv"
"strings" "strings"
"github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/common"
...@@ -11,12 +15,13 @@ import ( ...@@ -11,12 +15,13 @@ import (
"github.com/QuantumNous/new-api/dto" "github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/relay/channel" "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" relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting/system_setting"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/pkg/errors" "github.com/pkg/errors"
"github.com/tidwall/sjson"
) )
// ============================ // ============================
...@@ -57,6 +62,7 @@ type responseTask struct { ...@@ -57,6 +62,7 @@ type responseTask struct {
// ============================ // ============================
type TaskAdaptor struct { type TaskAdaptor struct {
taskcommon.BaseBilling
ChannelType int ChannelType int
apiKey string apiKey string
baseURL string baseURL string
...@@ -69,15 +75,15 @@ func (a *TaskAdaptor) Init(info *relaycommon.RelayInfo) { ...@@ -69,15 +75,15 @@ func (a *TaskAdaptor) Init(info *relaycommon.RelayInfo) {
} }
func validateRemixRequest(c *gin.Context) *dto.TaskError { func validateRemixRequest(c *gin.Context) *dto.TaskError {
var req struct { var req relaycommon.TaskSubmitReq
Prompt string `json:"prompt"`
}
if err := common.UnmarshalBodyReusable(c, &req); err != nil { if err := common.UnmarshalBodyReusable(c, &req); err != nil {
return service.TaskErrorWrapperLocal(err, "invalid_request", http.StatusBadRequest) return service.TaskErrorWrapperLocal(err, "invalid_request", http.StatusBadRequest)
} }
if strings.TrimSpace(req.Prompt) == "" { if strings.TrimSpace(req.Prompt) == "" {
return service.TaskErrorWrapperLocal(fmt.Errorf("field prompt is required"), "invalid_request", http.StatusBadRequest) return service.TaskErrorWrapperLocal(fmt.Errorf("field prompt is required"), "invalid_request", http.StatusBadRequest)
} }
// 存储原始请求到 context，与 ValidateMultipartDirect 路径保持一致
c.Set("task_request", req)
return nil return nil
} }
...@@ -88,6 +94,41 @@ func (a *TaskAdaptor) ValidateRequestAndSetAction(c *gin.Context, info *relaycom ...@@ -88,6 +94,41 @@ func (a *TaskAdaptor) ValidateRequestAndSetAction(c *gin.Context, info *relaycom
return relaycommon.ValidateMultipartDirect(c, info) 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) { func (a *TaskAdaptor) BuildRequestURL(info *relaycommon.RelayInfo) (string, error) {
if info.Action == constant.TaskActionRemix { if info.Action == constant.TaskActionRemix {
return fmt.Sprintf("%s/v1/videos/%s/remix", a.baseURL, info.OriginTaskID), nil 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 ...@@ -107,6 +148,74 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
if err != nil { if err != nil {
return nil, errors.Wrap(err, "get_request_body_failed") 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 return common.ReaderOnly(storage), nil
} }
...@@ -116,7 +225,7 @@ func (a *TaskAdaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, req ...@@ -116,7 +225,7 @@ func (a *TaskAdaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, req
} }
// DoResponse handles upstream response, returns taskID etc. // 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) responseBody, err := io.ReadAll(resp.Body)
if err != nil { if err != nil {
taskErr = service.TaskErrorWrapper(err, "read_response_body_failed", http.StatusInternalServerError) 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 ...@@ -131,17 +240,20 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, _ *relayco
return return
} }
if dResp.ID == "" { upstreamID := dResp.ID
if dResp.TaskID == "" { if upstreamID == "" {
taskErr = service.TaskErrorWrapper(fmt.Errorf("task_id is empty"), "invalid_response", http.StatusInternalServerError) upstreamID = dResp.TaskID
return }
} if upstreamID == "" {
dResp.ID = dResp.TaskID taskErr = service.TaskErrorWrapper(fmt.Errorf("task_id is empty"), "invalid_response", http.StatusInternalServerError)
dResp.TaskID = "" return
} }
// 使用公开 task_xxxx ID 返回给客户端
dResp.ID = info.PublicTaskID
dResp.TaskID = info.PublicTaskID
c.JSON(http.StatusOK, dResp) c.JSON(http.StatusOK, dResp)
return dResp.ID, responseBody, nil return upstreamID, responseBody, nil
} }
// FetchTask fetch task status // FetchTask fetch task status
...@@ -192,7 +304,7 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e ...@@ -192,7 +304,7 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
taskResult.Status = model.TaskStatusInProgress taskResult.Status = model.TaskStatusInProgress
case "completed": case "completed":
taskResult.Status = model.TaskStatusSuccess 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": case "failed", "cancelled":
taskResult.Status = model.TaskStatusFailure taskResult.Status = model.TaskStatusFailure
if resTask.Error != nil { if resTask.Error != nil {
...@@ -210,5 +322,10 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e ...@@ -210,5 +322,10 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
} }
func (a *TaskAdaptor) ConvertToOpenAIVideo(task *model.Task) ([]byte, error) { 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 ...@@ -2,18 +2,16 @@ package suno
import ( import (
"bytes" "bytes"
"context"
"encoding/json"
"fmt" "fmt"
"io" "io"
"net/http" "net/http"
"strings" "strings"
"time"
"github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto" "github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/relay/channel" "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" relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/service"
...@@ -21,11 +19,16 @@ import ( ...@@ -21,11 +19,16 @@ import (
) )
type TaskAdaptor struct { type TaskAdaptor struct {
taskcommon.BaseBilling
ChannelType int 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) { 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) { func (a *TaskAdaptor) Init(info *relaycommon.RelayInfo) {
...@@ -47,13 +50,13 @@ func (a *TaskAdaptor) ValidateRequestAndSetAction(c *gin.Context, info *relaycom ...@@ -47,13 +50,13 @@ func (a *TaskAdaptor) ValidateRequestAndSetAction(c *gin.Context, info *relaycom
return return
} }
if sunoRequest.ContinueClipId != "" { //if sunoRequest.ContinueClipId != "" {
if sunoRequest.TaskID == "" { // if sunoRequest.TaskID == "" {
taskErr = service.TaskErrorWrapperLocal(fmt.Errorf("task id is empty"), "invalid_request", http.StatusBadRequest) // taskErr = service.TaskErrorWrapperLocal(fmt.Errorf("task id is empty"), "invalid_request", http.StatusBadRequest)
return // return
} // }
info.OriginTaskID = sunoRequest.TaskID // info.OriginTaskID = sunoRequest.TaskID
} //}
info.Action = action info.Action = action
c.Set("task_request", sunoRequest) c.Set("task_request", sunoRequest)
...@@ -76,12 +79,9 @@ func (a *TaskAdaptor) BuildRequestHeader(c *gin.Context, req *http.Request, info ...@@ -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) { func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayInfo) (io.Reader, error) {
sunoRequest, ok := c.Get("task_request") sunoRequest, ok := c.Get("task_request")
if !ok { if !ok {
err := common.UnmarshalBodyReusable(c, &sunoRequest) return nil, fmt.Errorf("task_request not found in context")
if err != nil {
return nil, err
}
} }
data, err := json.Marshal(sunoRequest) data, err := common.Marshal(sunoRequest)
if err != nil { if err != nil {
return nil, err return nil, err
} }
...@@ -99,7 +99,7 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela ...@@ -99,7 +99,7 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
return return
} }
var sunoResponse dto.TaskResponse[string] var sunoResponse dto.TaskResponse[string]
err = json.Unmarshal(responseBody, &sunoResponse) err = common.Unmarshal(responseBody, &sunoResponse)
if err != nil { if err != nil {
taskErr = service.TaskErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError) taskErr = service.TaskErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError)
return return
...@@ -109,17 +109,13 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela ...@@ -109,17 +109,13 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
return return
} }
for k, v := range resp.Header { // 使用公开 task_xxxx ID 替换上游 ID 返回给客户端
c.Writer.Header().Set(k, v[0]) publicResponse := dto.TaskResponse[string]{
} Code: sunoResponse.Code,
c.Writer.Header().Set("Content-Type", "application/json") Message: sunoResponse.Message,
c.Writer.WriteHeader(resp.StatusCode) Data: info.PublicTaskID,
_, err = io.Copy(c.Writer, bytes.NewBuffer(responseBody))
if err != nil {
taskErr = service.TaskErrorWrapper(err, "copy_response_body_failed", http.StatusInternalServerError)
return
} }
c.JSON(http.StatusOK, publicResponse)
return sunoResponse.Data, nil, nil return sunoResponse.Data, nil, nil
} }
...@@ -134,7 +130,7 @@ func (a *TaskAdaptor) GetChannelName() string { ...@@ -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) { func (a *TaskAdaptor) FetchTask(baseUrl, key string, body map[string]any, proxy string) (*http.Response, error) {
requestUrl := fmt.Sprintf("%s/suno/fetch", baseUrl) requestUrl := fmt.Sprintf("%s/suno/fetch", baseUrl)
byteBody, err := json.Marshal(body) byteBody, err := common.Marshal(body)
if err != nil { if err != nil {
return nil, err return nil, err
} }
...@@ -144,13 +140,6 @@ func (a *TaskAdaptor) FetchTask(baseUrl, key string, body map[string]any, proxy ...@@ -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)) common.SysLog(fmt.Sprintf("Get Task error: %v", err))
return nil, 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("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer "+key) req.Header.Set("Authorization", "Bearer "+key)
client, err := service.GetHttpClientWithProxy(proxy) 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 ...@@ -2,13 +2,12 @@ package vertex
import ( import (
"bytes" "bytes"
"encoding/base64"
"encoding/json"
"fmt" "fmt"
"io" "io"
"net/http" "net/http"
"regexp" "regexp"
"strings" "strings"
"time"
"github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/model"
...@@ -17,6 +16,7 @@ import ( ...@@ -17,6 +16,7 @@ import (
"github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto" "github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/relay/channel" "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" vertexcore "github.com/QuantumNous/new-api/relay/channel/vertex"
relaycommon "github.com/QuantumNous/new-api/relay/common" relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/service"
...@@ -62,6 +62,7 @@ type operationResponse struct { ...@@ -62,6 +62,7 @@ type operationResponse struct {
// ============================ // ============================
type TaskAdaptor struct { type TaskAdaptor struct {
taskcommon.BaseBilling
ChannelType int ChannelType int
apiKey string apiKey string
baseURL string baseURL string
...@@ -82,10 +83,10 @@ func (a *TaskAdaptor) ValidateRequestAndSetAction(c *gin.Context, info *relaycom ...@@ -82,10 +83,10 @@ func (a *TaskAdaptor) ValidateRequestAndSetAction(c *gin.Context, info *relaycom
// BuildRequestURL constructs the upstream URL. // BuildRequestURL constructs the upstream URL.
func (a *TaskAdaptor) BuildRequestURL(info *relaycommon.RelayInfo) (string, error) { func (a *TaskAdaptor) BuildRequestURL(info *relaycommon.RelayInfo) (string, error) {
adc := &vertexcore.Credentials{} 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) return "", fmt.Errorf("failed to decode credentials: %w", err)
} }
modelName := info.OriginModelName modelName := info.UpstreamModelName
if modelName == "" { if modelName == "" {
modelName = "veo-3.0-generate-001" modelName = "veo-3.0-generate-001"
} }
...@@ -116,7 +117,7 @@ func (a *TaskAdaptor) BuildRequestHeader(c *gin.Context, req *http.Request, info ...@@ -116,7 +117,7 @@ func (a *TaskAdaptor) BuildRequestHeader(c *gin.Context, req *http.Request, info
req.Header.Set("Accept", "application/json") req.Header.Set("Accept", "application/json")
adc := &vertexcore.Credentials{} 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) return fmt.Errorf("failed to decode credentials: %w", err)
} }
...@@ -133,6 +134,28 @@ func (a *TaskAdaptor) BuildRequestHeader(c *gin.Context, req *http.Request, info ...@@ -133,6 +134,28 @@ func (a *TaskAdaptor) BuildRequestHeader(c *gin.Context, req *http.Request, info
return nil 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. // BuildRequestBody converts request into Vertex specific format.
func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayInfo) (io.Reader, error) { func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayInfo) (io.Reader, error) {
v, ok := c.Get("task_request") v, ok := c.Get("task_request")
...@@ -166,25 +189,7 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn ...@@ -166,25 +189,7 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
return nil, fmt.Errorf("sampleCount must be greater than 0") return nil, fmt.Errorf("sampleCount must be greater than 0")
} }
// if req.Duration > 0 { data, err := common.Marshal(body)
// 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)
if err != nil { if err != nil {
return nil, err return nil, err
} }
...@@ -205,14 +210,19 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela ...@@ -205,14 +210,19 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
_ = resp.Body.Close() _ = resp.Body.Close()
var s submitResponse 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) return "", nil, service.TaskErrorWrapper(err, "unmarshal_response_failed", http.StatusInternalServerError)
} }
if strings.TrimSpace(s.Name) == "" { if strings.TrimSpace(s.Name) == "" {
return "", nil, service.TaskErrorWrapper(fmt.Errorf("missing operation name"), "invalid_response", http.StatusInternalServerError) return "", nil, service.TaskErrorWrapper(fmt.Errorf("missing operation name"), "invalid_response", http.StatusInternalServerError)
} }
localID := encodeLocalTaskID(s.Name) localID := taskcommon.EncodeLocalTaskID(s.Name)
c.JSON(http.StatusOK, gin.H{"task_id": localID}) 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 return localID, responseBody, nil
} }
...@@ -225,7 +235,7 @@ func (a *TaskAdaptor) FetchTask(baseUrl, key string, body map[string]any, proxy ...@@ -225,7 +235,7 @@ func (a *TaskAdaptor) FetchTask(baseUrl, key string, body map[string]any, proxy
if !ok { if !ok {
return nil, fmt.Errorf("invalid task_id") return nil, fmt.Errorf("invalid task_id")
} }
upstreamName, err := decodeLocalTaskID(taskID) upstreamName, err := taskcommon.DecodeLocalTaskID(taskID)
if err != nil { if err != nil {
return nil, fmt.Errorf("decode task_id failed: %w", err) 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 ...@@ -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) 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} payload := map[string]string{"operationName": upstreamName}
data, err := json.Marshal(payload) data, err := common.Marshal(payload)
if err != nil { if err != nil {
return nil, err return nil, err
} }
adc := &vertexcore.Credentials{} 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) return nil, fmt.Errorf("failed to decode credentials: %w", err)
} }
token, err := vertexcore.AcquireAccessToken(*adc, proxy) token, err := vertexcore.AcquireAccessToken(*adc, proxy)
...@@ -274,7 +284,7 @@ func (a *TaskAdaptor) FetchTask(baseUrl, key string, body map[string]any, 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) { func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) {
var op operationResponse 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) return nil, fmt.Errorf("unmarshal operation response failed: %w", err)
} }
ti := &relaycommon.TaskInfo{} ti := &relaycommon.TaskInfo{}
...@@ -338,7 +348,10 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e ...@@ -338,7 +348,10 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
} }
func (a *TaskAdaptor) ConvertToOpenAIVideo(task *model.Task) ([]byte, error) { 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 { if err != nil {
upstreamName = "" upstreamName = ""
} }
...@@ -353,8 +366,8 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(task *model.Task) ([]byte, error) { ...@@ -353,8 +366,8 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(task *model.Task) ([]byte, error) {
v.SetProgressStr(task.Progress) v.SetProgressStr(task.Progress)
v.CreatedAt = task.CreatedAt v.CreatedAt = task.CreatedAt
v.CompletedAt = task.UpdatedAt v.CompletedAt = task.UpdatedAt
if strings.HasPrefix(task.FailReason, "data:") && len(task.FailReason) > 0 { if resultURL := task.GetResultURL(); strings.HasPrefix(resultURL, "data:") && len(resultURL) > 0 {
v.SetMetadata("url", task.FailReason) v.SetMetadata("url", resultURL)
} }
return common.Marshal(v) return common.Marshal(v)
...@@ -364,18 +377,6 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(task *model.Task) ([]byte, error) { ...@@ -364,18 +377,6 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(task *model.Task) ([]byte, error) {
// helpers // 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-]+)/`) var regionRe = regexp.MustCompile(`locations/([a-z0-9-]+)/`)
func extractRegionFromOperationName(name string) string { func extractRegionFromOperationName(name string) string {
......
relay/channel/task/vidu/adaptor.go
...@@ -2,7 +2,6 @@ package vidu ...@@ -2,7 +2,6 @@ package vidu
import ( import (
"bytes" "bytes"
"encoding/json"
"fmt" "fmt"
"io" "io"
"net/http" "net/http"
...@@ -16,6 +15,7 @@ import ( ...@@ -16,6 +15,7 @@ import (
"github.com/QuantumNous/new-api/dto" "github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/relay/channel" "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" relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/service"
...@@ -73,6 +73,7 @@ type creation struct { ...@@ -73,6 +73,7 @@ type creation struct {
// ============================ // ============================
type TaskAdaptor struct { type TaskAdaptor struct {
taskcommon.BaseBilling
ChannelType int ChannelType int
baseURL string baseURL string
} }
...@@ -115,7 +116,7 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn ...@@ -115,7 +116,7 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
} }
req := v.(relaycommon.TaskSubmitReq) req := v.(relaycommon.TaskSubmitReq)
body, err := a.convertToRequestPayload(&req) body, err := a.convertToRequestPayload(&req, info)
if err != nil { if err != nil {
return nil, err return nil, err
} }
...@@ -127,7 +128,7 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn ...@@ -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 { if err != nil {
return nil, err return nil, err
} }
...@@ -168,7 +169,7 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela ...@@ -168,7 +169,7 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
} }
var vResp responsePayload var vResp responsePayload
err = json.Unmarshal(responseBody, &vResp) err = common.Unmarshal(responseBody, &vResp)
if err != nil { if err != nil {
taskErr = service.TaskErrorWrapper(errors.Wrap(err, fmt.Sprintf("%s", responseBody)), "unmarshal_response_failed", http.StatusInternalServerError) taskErr = service.TaskErrorWrapper(errors.Wrap(err, fmt.Sprintf("%s", responseBody)), "unmarshal_response_failed", http.StatusInternalServerError)
return return
...@@ -180,8 +181,8 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela ...@@ -180,8 +181,8 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
} }
ov := dto.NewOpenAIVideo() ov := dto.NewOpenAIVideo()
ov.ID = vResp.TaskId ov.ID = info.PublicTaskID
ov.TaskID = vResp.TaskId ov.TaskID = info.PublicTaskID
ov.CreatedAt = time.Now().Unix() ov.CreatedAt = time.Now().Unix()
ov.Model = info.OriginModelName ov.Model = info.OriginModelName
c.JSON(http.StatusOK, ov) c.JSON(http.StatusOK, ov)
...@@ -223,47 +224,27 @@ func (a *TaskAdaptor) GetChannelName() string { ...@@ -223,47 +224,27 @@ func (a *TaskAdaptor) GetChannelName() string {
// helpers // helpers
// ============================ // ============================
func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*requestPayload, error) { func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq, info *relaycommon.RelayInfo) (*requestPayload, error) {
r := requestPayload{ r := requestPayload{
Model: defaultString(req.Model, "viduq1"), Model: taskcommon.DefaultString(info.UpstreamModelName, "viduq1"),
Images: req.Images, Images: req.Images,
Prompt: req.Prompt, Prompt: req.Prompt,
Duration: defaultInt(req.Duration, 5), Duration: taskcommon.DefaultInt(req.Duration, 5),
Resolution: defaultString(req.Size, "1080p"), Resolution: taskcommon.DefaultString(req.Size, "1080p"),
MovementAmplitude: "auto", MovementAmplitude: "auto",
Bgm: false, Bgm: false,
} }
metadata := req.Metadata if err := taskcommon.UnmarshalMetadata(req.Metadata, &r); err != nil {
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 {
return nil, errors.Wrap(err, "unmarshal metadata failed") return nil, errors.Wrap(err, "unmarshal metadata failed")
} }
return &r, nil 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) { func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) {
taskInfo := &relaycommon.TaskInfo{} taskInfo := &relaycommon.TaskInfo{}
var taskResp taskResultResponse var taskResp taskResultResponse
err := json.Unmarshal(respBody, &taskResp) err := common.Unmarshal(respBody, &taskResp)
if err != nil { if err != nil {
return nil, errors.Wrap(err, "failed to unmarshal response body") return nil, errors.Wrap(err, "failed to unmarshal response body")
} }
...@@ -293,7 +274,7 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e ...@@ -293,7 +274,7 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, error) { func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, error) {
var viduResp taskResultResponse 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") return nil, errors.Wrap(err, "unmarshal vidu task data failed")
} }
...@@ -315,6 +296,5 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, erro ...@@ -315,6 +296,5 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, erro
} }
} }
jsonData, _ := common.Marshal(openAIVideo) return common.Marshal(openAIVideo)
return jsonData, nil
} }
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