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
This commit is contained in:
@@ -36,6 +36,32 @@ type TaskAdaptor interface {
|
||||
|
||||
ValidateRequestAndSetAction(c *gin.Context, info *relaycommon.RelayInfo) *dto.TaskError
|
||||
|
||||
// ── Billing ──────────────────────────────────────────────────────
|
||||
|
||||
// EstimateBilling returns OtherRatios for pre-charge based on user request.
|
||||
// Called after ValidateRequestAndSetAction, before price calculation.
|
||||
// Adaptors should extract duration, resolution, etc. from the parsed request
|
||||
// and return them as ratio multipliers (e.g. {"seconds": 5, "size": 1.666}).
|
||||
// Return nil to use the base model price without extra ratios.
|
||||
EstimateBilling(c *gin.Context, info *relaycommon.RelayInfo) map[string]float64
|
||||
|
||||
// AdjustBillingOnSubmit returns adjusted OtherRatios from the upstream
|
||||
// submit response. Called after a successful DoResponse.
|
||||
// If the upstream returned actual parameters that differ from the estimate
|
||||
// (e.g. actual seconds), return updated ratios so the caller can recalculate
|
||||
// the quota and settle the delta with the pre-charge.
|
||||
// Return nil if no adjustment is needed.
|
||||
AdjustBillingOnSubmit(info *relaycommon.RelayInfo, taskData []byte) map[string]float64
|
||||
|
||||
// AdjustBillingOnComplete returns the actual quota when a task reaches a
|
||||
// terminal state (success/failure) during polling.
|
||||
// Called by the polling loop after ParseTaskResult.
|
||||
// Return a positive value to trigger delta settlement (supplement / refund).
|
||||
// Return 0 to keep the pre-charged amount unchanged.
|
||||
AdjustBillingOnComplete(task *model.Task, taskResult *relaycommon.TaskInfo) int
|
||||
|
||||
// ── Request / Response ───────────────────────────────────────────
|
||||
|
||||
BuildRequestURL(info *relaycommon.RelayInfo) (string, error)
|
||||
BuildRequestHeader(c *gin.Context, req *http.Request, info *relaycommon.RelayInfo) error
|
||||
BuildRequestBody(c *gin.Context, info *relaycommon.RelayInfo) (io.Reader, error)
|
||||
@@ -46,9 +72,9 @@ type TaskAdaptor interface {
|
||||
GetModelList() []string
|
||||
GetChannelName() string
|
||||
|
||||
// FetchTask
|
||||
FetchTask(baseUrl, key string, body map[string]any, proxy string) (*http.Response, error)
|
||||
// ── Polling ──────────────────────────────────────────────────────
|
||||
|
||||
FetchTask(baseUrl, key string, body map[string]any, proxy string) (*http.Response, error)
|
||||
ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error)
|
||||
}
|
||||
|
||||
|
||||
@@ -61,8 +61,9 @@ var passthroughSkipHeaderNamesLower = map[string]struct{}{
|
||||
"cookie": {},
|
||||
|
||||
// Additional headers that should not be forwarded by name-matching passthrough rules.
|
||||
"host": {},
|
||||
"content-length": {},
|
||||
"host": {},
|
||||
"content-length": {},
|
||||
"accept-encoding": {},
|
||||
|
||||
// Do not passthrough credentials by wildcard/regex.
|
||||
"authorization": {},
|
||||
|
||||
@@ -110,3 +110,30 @@ func TestProcessHeaderOverride_RuntimeOverrideHasPriority(t *testing.T) {
|
||||
_, ok := headers["X-Legacy"]
|
||||
require.False(t, ok)
|
||||
}
|
||||
|
||||
func TestProcessHeaderOverride_PassthroughSkipsAcceptEncoding(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(recorder)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
||||
ctx.Request.Header.Set("X-Trace-Id", "trace-123")
|
||||
ctx.Request.Header.Set("Accept-Encoding", "gzip")
|
||||
|
||||
info := &relaycommon.RelayInfo{
|
||||
IsChannelTest: false,
|
||||
ChannelMeta: &relaycommon.ChannelMeta{
|
||||
HeadersOverride: map[string]any{
|
||||
"*": "",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
headers, err := processHeaderOverride(info, ctx)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "trace-123", headers["X-Trace-Id"])
|
||||
|
||||
_, hasAcceptEncoding := headers["Accept-Encoding"]
|
||||
require.False(t, hasAcceptEncoding)
|
||||
}
|
||||
|
||||
@@ -42,22 +42,7 @@ func GeminiTextGenerationHandler(c *gin.Context, info *relaycommon.RelayInfo, re
|
||||
}
|
||||
|
||||
// 计算使用量(基于 UsageMetadata)
|
||||
usage := dto.Usage{
|
||||
PromptTokens: geminiResponse.UsageMetadata.PromptTokenCount,
|
||||
CompletionTokens: geminiResponse.UsageMetadata.CandidatesTokenCount + geminiResponse.UsageMetadata.ThoughtsTokenCount,
|
||||
TotalTokens: geminiResponse.UsageMetadata.TotalTokenCount,
|
||||
}
|
||||
|
||||
usage.CompletionTokenDetails.ReasoningTokens = geminiResponse.UsageMetadata.ThoughtsTokenCount
|
||||
usage.PromptTokensDetails.CachedTokens = geminiResponse.UsageMetadata.CachedContentTokenCount
|
||||
|
||||
for _, detail := range geminiResponse.UsageMetadata.PromptTokensDetails {
|
||||
if detail.Modality == "AUDIO" {
|
||||
usage.PromptTokensDetails.AudioTokens = detail.TokenCount
|
||||
} else if detail.Modality == "TEXT" {
|
||||
usage.PromptTokensDetails.TextTokens = detail.TokenCount
|
||||
}
|
||||
}
|
||||
usage := buildUsageFromGeminiMetadata(geminiResponse.UsageMetadata, info.GetEstimatePromptTokens())
|
||||
|
||||
service.IOCopyBytesGracefully(c, resp, responseBody)
|
||||
|
||||
|
||||
@@ -1032,6 +1032,46 @@ func getResponseToolCall(item *dto.GeminiPart) *dto.ToolCallResponse {
|
||||
}
|
||||
}
|
||||
|
||||
func buildUsageFromGeminiMetadata(metadata dto.GeminiUsageMetadata, fallbackPromptTokens int) dto.Usage {
|
||||
promptTokens := metadata.PromptTokenCount + metadata.ToolUsePromptTokenCount
|
||||
if promptTokens <= 0 && fallbackPromptTokens > 0 {
|
||||
promptTokens = fallbackPromptTokens
|
||||
}
|
||||
|
||||
usage := dto.Usage{
|
||||
PromptTokens: promptTokens,
|
||||
CompletionTokens: metadata.CandidatesTokenCount + metadata.ThoughtsTokenCount,
|
||||
TotalTokens: metadata.TotalTokenCount,
|
||||
}
|
||||
usage.CompletionTokenDetails.ReasoningTokens = metadata.ThoughtsTokenCount
|
||||
usage.PromptTokensDetails.CachedTokens = metadata.CachedContentTokenCount
|
||||
|
||||
for _, detail := range metadata.PromptTokensDetails {
|
||||
if detail.Modality == "AUDIO" {
|
||||
usage.PromptTokensDetails.AudioTokens += detail.TokenCount
|
||||
} else if detail.Modality == "TEXT" {
|
||||
usage.PromptTokensDetails.TextTokens += detail.TokenCount
|
||||
}
|
||||
}
|
||||
for _, detail := range metadata.ToolUsePromptTokensDetails {
|
||||
if detail.Modality == "AUDIO" {
|
||||
usage.PromptTokensDetails.AudioTokens += detail.TokenCount
|
||||
} else if detail.Modality == "TEXT" {
|
||||
usage.PromptTokensDetails.TextTokens += detail.TokenCount
|
||||
}
|
||||
}
|
||||
|
||||
if usage.TotalTokens > 0 && usage.CompletionTokens <= 0 {
|
||||
usage.CompletionTokens = usage.TotalTokens - usage.PromptTokens
|
||||
}
|
||||
|
||||
if usage.PromptTokens > 0 && usage.PromptTokensDetails.TextTokens == 0 && usage.PromptTokensDetails.AudioTokens == 0 {
|
||||
usage.PromptTokensDetails.TextTokens = usage.PromptTokens
|
||||
}
|
||||
|
||||
return usage
|
||||
}
|
||||
|
||||
func responseGeminiChat2OpenAI(c *gin.Context, response *dto.GeminiChatResponse) *dto.OpenAITextResponse {
|
||||
fullTextResponse := dto.OpenAITextResponse{
|
||||
Id: helper.GetResponseID(c),
|
||||
@@ -1272,18 +1312,8 @@ func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http
|
||||
|
||||
// 更新使用量统计
|
||||
if geminiResponse.UsageMetadata.TotalTokenCount != 0 {
|
||||
usage.PromptTokens = geminiResponse.UsageMetadata.PromptTokenCount
|
||||
usage.CompletionTokens = geminiResponse.UsageMetadata.CandidatesTokenCount + geminiResponse.UsageMetadata.ThoughtsTokenCount
|
||||
usage.CompletionTokenDetails.ReasoningTokens = geminiResponse.UsageMetadata.ThoughtsTokenCount
|
||||
usage.TotalTokens = geminiResponse.UsageMetadata.TotalTokenCount
|
||||
usage.PromptTokensDetails.CachedTokens = geminiResponse.UsageMetadata.CachedContentTokenCount
|
||||
for _, detail := range geminiResponse.UsageMetadata.PromptTokensDetails {
|
||||
if detail.Modality == "AUDIO" {
|
||||
usage.PromptTokensDetails.AudioTokens = detail.TokenCount
|
||||
} else if detail.Modality == "TEXT" {
|
||||
usage.PromptTokensDetails.TextTokens = detail.TokenCount
|
||||
}
|
||||
}
|
||||
mappedUsage := buildUsageFromGeminiMetadata(geminiResponse.UsageMetadata, info.GetEstimatePromptTokens())
|
||||
*usage = mappedUsage
|
||||
}
|
||||
|
||||
return callback(data, &geminiResponse)
|
||||
@@ -1295,11 +1325,6 @@ func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http
|
||||
}
|
||||
}
|
||||
|
||||
usage.PromptTokensDetails.TextTokens = usage.PromptTokens
|
||||
if usage.TotalTokens > 0 {
|
||||
usage.CompletionTokens = usage.TotalTokens - usage.PromptTokens
|
||||
}
|
||||
|
||||
if usage.CompletionTokens <= 0 {
|
||||
if info.ReceivedResponseCount > 0 {
|
||||
usage = service.ResponseText2Usage(c, responseText.String(), info.UpstreamModelName, info.GetEstimatePromptTokens())
|
||||
@@ -1416,21 +1441,7 @@ func GeminiChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.R
|
||||
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
|
||||
}
|
||||
if len(geminiResponse.Candidates) == 0 {
|
||||
usage := dto.Usage{
|
||||
PromptTokens: geminiResponse.UsageMetadata.PromptTokenCount,
|
||||
}
|
||||
usage.CompletionTokenDetails.ReasoningTokens = geminiResponse.UsageMetadata.ThoughtsTokenCount
|
||||
usage.PromptTokensDetails.CachedTokens = geminiResponse.UsageMetadata.CachedContentTokenCount
|
||||
for _, detail := range geminiResponse.UsageMetadata.PromptTokensDetails {
|
||||
if detail.Modality == "AUDIO" {
|
||||
usage.PromptTokensDetails.AudioTokens = detail.TokenCount
|
||||
} else if detail.Modality == "TEXT" {
|
||||
usage.PromptTokensDetails.TextTokens = detail.TokenCount
|
||||
}
|
||||
}
|
||||
if usage.PromptTokens <= 0 {
|
||||
usage.PromptTokens = info.GetEstimatePromptTokens()
|
||||
}
|
||||
usage := buildUsageFromGeminiMetadata(geminiResponse.UsageMetadata, info.GetEstimatePromptTokens())
|
||||
|
||||
var newAPIError *types.NewAPIError
|
||||
if geminiResponse.PromptFeedback != nil && geminiResponse.PromptFeedback.BlockReason != nil {
|
||||
@@ -1466,23 +1477,7 @@ func GeminiChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.R
|
||||
}
|
||||
fullTextResponse := responseGeminiChat2OpenAI(c, &geminiResponse)
|
||||
fullTextResponse.Model = info.UpstreamModelName
|
||||
usage := dto.Usage{
|
||||
PromptTokens: geminiResponse.UsageMetadata.PromptTokenCount,
|
||||
CompletionTokens: geminiResponse.UsageMetadata.CandidatesTokenCount,
|
||||
TotalTokens: geminiResponse.UsageMetadata.TotalTokenCount,
|
||||
}
|
||||
|
||||
usage.CompletionTokenDetails.ReasoningTokens = geminiResponse.UsageMetadata.ThoughtsTokenCount
|
||||
usage.PromptTokensDetails.CachedTokens = geminiResponse.UsageMetadata.CachedContentTokenCount
|
||||
usage.CompletionTokens = usage.TotalTokens - usage.PromptTokens
|
||||
|
||||
for _, detail := range geminiResponse.UsageMetadata.PromptTokensDetails {
|
||||
if detail.Modality == "AUDIO" {
|
||||
usage.PromptTokensDetails.AudioTokens = detail.TokenCount
|
||||
} else if detail.Modality == "TEXT" {
|
||||
usage.PromptTokensDetails.TextTokens = detail.TokenCount
|
||||
}
|
||||
}
|
||||
usage := buildUsageFromGeminiMetadata(geminiResponse.UsageMetadata, info.GetEstimatePromptTokens())
|
||||
|
||||
fullTextResponse.Usage = usage
|
||||
|
||||
|
||||
@@ -0,0 +1,333 @@
|
||||
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)
|
||||
}
|
||||
@@ -10,6 +10,7 @@ import (
|
||||
|
||||
"github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/relay/channel"
|
||||
"github.com/QuantumNous/new-api/relay/channel/claude"
|
||||
"github.com/QuantumNous/new-api/relay/channel/openai"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relay/constant"
|
||||
@@ -26,7 +27,8 @@ func (a *Adaptor) ConvertGeminiRequest(*gin.Context, *relaycommon.RelayInfo, *dt
|
||||
}
|
||||
|
||||
func (a *Adaptor) ConvertClaudeRequest(c *gin.Context, info *relaycommon.RelayInfo, req *dto.ClaudeRequest) (any, error) {
|
||||
return nil, errors.New("not implemented")
|
||||
adaptor := claude.Adaptor{}
|
||||
return adaptor.ConvertClaudeRequest(c, info, req)
|
||||
}
|
||||
|
||||
func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.AudioRequest) (io.Reader, error) {
|
||||
@@ -119,8 +121,14 @@ func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycom
|
||||
return handleTTSResponse(c, resp, info)
|
||||
}
|
||||
|
||||
adaptor := openai.Adaptor{}
|
||||
return adaptor.DoResponse(c, resp, info)
|
||||
switch info.RelayFormat {
|
||||
case types.RelayFormatClaude:
|
||||
adaptor := claude.Adaptor{}
|
||||
return adaptor.DoResponse(c, resp, info)
|
||||
default:
|
||||
adaptor := openai.Adaptor{}
|
||||
return adaptor.DoResponse(c, resp, info)
|
||||
}
|
||||
}
|
||||
|
||||
func (a *Adaptor) GetModelList() []string {
|
||||
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
channelconstant "github.com/QuantumNous/new-api/constant"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relay/constant"
|
||||
"github.com/QuantumNous/new-api/types"
|
||||
)
|
||||
|
||||
func GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
|
||||
@@ -13,13 +14,17 @@ func GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
|
||||
if baseUrl == "" {
|
||||
baseUrl = channelconstant.ChannelBaseURLs[channelconstant.ChannelTypeMiniMax]
|
||||
}
|
||||
|
||||
switch info.RelayMode {
|
||||
case constant.RelayModeChatCompletions:
|
||||
return fmt.Sprintf("%s/v1/text/chatcompletion_v2", baseUrl), nil
|
||||
case constant.RelayModeAudioSpeech:
|
||||
return fmt.Sprintf("%s/v1/t2a_v2", baseUrl), nil
|
||||
switch info.RelayFormat {
|
||||
case types.RelayFormatClaude:
|
||||
return fmt.Sprintf("%s/anthropic/v1/messages", info.ChannelBaseUrl), nil
|
||||
default:
|
||||
return "", fmt.Errorf("unsupported relay mode: %d", info.RelayMode)
|
||||
switch info.RelayMode {
|
||||
case constant.RelayModeChatCompletions:
|
||||
return fmt.Sprintf("%s/v1/text/chatcompletion_v2", baseUrl), nil
|
||||
case constant.RelayModeAudioSpeech:
|
||||
return fmt.Sprintf("%s/v1/t2a_v2", baseUrl), nil
|
||||
default:
|
||||
return "", fmt.Errorf("unsupported relay mode: %d", info.RelayMode)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -13,6 +13,7 @@ import (
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/relay/channel"
|
||||
"github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/samber/lo"
|
||||
@@ -108,10 +109,10 @@ type AliMetadata struct {
|
||||
// ============================
|
||||
|
||||
type TaskAdaptor struct {
|
||||
taskcommon.BaseBilling
|
||||
ChannelType int
|
||||
apiKey string
|
||||
baseURL string
|
||||
aliReq *AliVideoRequest
|
||||
}
|
||||
|
||||
func (a *TaskAdaptor) Init(info *relaycommon.RelayInfo) {
|
||||
@@ -121,17 +122,7 @@ func (a *TaskAdaptor) Init(info *relaycommon.RelayInfo) {
|
||||
}
|
||||
|
||||
func (a *TaskAdaptor) ValidateRequestAndSetAction(c *gin.Context, info *relaycommon.RelayInfo) (taskErr *dto.TaskError) {
|
||||
// 阿里通义万相支持 JSON 格式,不使用 multipart
|
||||
var taskReq relaycommon.TaskSubmitReq
|
||||
if err := common.UnmarshalBodyReusable(c, &taskReq); err != nil {
|
||||
return service.TaskErrorWrapper(err, "unmarshal_task_request_failed", http.StatusBadRequest)
|
||||
}
|
||||
aliReq, err := a.convertToAliRequest(info, taskReq)
|
||||
if err != nil {
|
||||
return service.TaskErrorWrapper(err, "convert_to_ali_request_failed", http.StatusInternalServerError)
|
||||
}
|
||||
a.aliReq = aliReq
|
||||
logger.LogJson(c, "ali video request body", aliReq)
|
||||
// ValidateMultipartDirect 负责解析并将原始 TaskSubmitReq 存入 context
|
||||
return relaycommon.ValidateMultipartDirect(c, info)
|
||||
}
|
||||
|
||||
@@ -148,11 +139,21 @@ func (a *TaskAdaptor) BuildRequestHeader(c *gin.Context, req *http.Request, info
|
||||
}
|
||||
|
||||
func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayInfo) (io.Reader, error) {
|
||||
bodyBytes, err := common.Marshal(a.aliReq)
|
||||
taskReq, err := relaycommon.GetTaskRequest(c)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "get_task_request_failed")
|
||||
}
|
||||
|
||||
aliReq, err := a.convertToAliRequest(info, taskReq)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "convert_to_ali_request_failed")
|
||||
}
|
||||
logger.LogJson(c, "ali video request body", aliReq)
|
||||
|
||||
bodyBytes, err := common.Marshal(aliReq)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "marshal_ali_request_failed")
|
||||
}
|
||||
|
||||
return bytes.NewReader(bodyBytes), nil
|
||||
}
|
||||
|
||||
@@ -252,8 +253,12 @@ func ProcessAliOtherRatios(aliReq *AliVideoRequest) (map[string]float64, error)
|
||||
}
|
||||
|
||||
func (a *TaskAdaptor) convertToAliRequest(info *relaycommon.RelayInfo, req relaycommon.TaskSubmitReq) (*AliVideoRequest, error) {
|
||||
upstreamModel := req.Model
|
||||
if info.IsModelMapped {
|
||||
upstreamModel = info.UpstreamModelName
|
||||
}
|
||||
aliReq := &AliVideoRequest{
|
||||
Model: req.Model,
|
||||
Model: upstreamModel,
|
||||
Input: AliVideoInput{
|
||||
Prompt: req.Prompt,
|
||||
ImgURL: req.InputReference,
|
||||
@@ -331,23 +336,37 @@ func (a *TaskAdaptor) convertToAliRequest(info *relaycommon.RelayInfo, req relay
|
||||
}
|
||||
}
|
||||
|
||||
if aliReq.Model != req.Model {
|
||||
if aliReq.Model != upstreamModel {
|
||||
return nil, errors.New("can't change model with metadata")
|
||||
}
|
||||
|
||||
info.PriceData.OtherRatios = map[string]float64{
|
||||
return aliReq, nil
|
||||
}
|
||||
|
||||
// EstimateBilling 根据用户请求参数计算 OtherRatios(时长、分辨率等)。
|
||||
// 在 ValidateRequestAndSetAction 之后、价格计算之前调用。
|
||||
func (a *TaskAdaptor) EstimateBilling(c *gin.Context, info *relaycommon.RelayInfo) map[string]float64 {
|
||||
taskReq, err := relaycommon.GetTaskRequest(c)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
aliReq, err := a.convertToAliRequest(info, taskReq)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
otherRatios := map[string]float64{
|
||||
"seconds": float64(aliReq.Parameters.Duration),
|
||||
}
|
||||
|
||||
ratios, err := ProcessAliOtherRatios(aliReq)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
return otherRatios
|
||||
}
|
||||
for s, f := range ratios {
|
||||
info.PriceData.OtherRatios[s] = f
|
||||
for k, v := range ratios {
|
||||
otherRatios[k] = v
|
||||
}
|
||||
|
||||
return aliReq, nil
|
||||
return otherRatios
|
||||
}
|
||||
|
||||
// DoRequest delegates to common helper
|
||||
@@ -384,7 +403,8 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
|
||||
|
||||
// 转换为 OpenAI 格式响应
|
||||
openAIResp := dto.NewOpenAIVideo()
|
||||
openAIResp.ID = aliResp.Output.TaskID
|
||||
openAIResp.ID = info.PublicTaskID
|
||||
openAIResp.TaskID = info.PublicTaskID
|
||||
openAIResp.Model = c.GetString("model")
|
||||
if openAIResp.Model == "" && info != nil {
|
||||
openAIResp.Model = info.OriginModelName
|
||||
|
||||
@@ -2,7 +2,6 @@ package doubao
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
@@ -14,6 +13,7 @@ import (
|
||||
"github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/relay/channel"
|
||||
taskcommon "github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
|
||||
@@ -89,6 +89,7 @@ type responseTask struct {
|
||||
// ============================
|
||||
|
||||
type TaskAdaptor struct {
|
||||
taskcommon.BaseBilling
|
||||
ChannelType int
|
||||
apiKey string
|
||||
baseURL string
|
||||
@@ -130,8 +131,12 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "convert request payload failed")
|
||||
}
|
||||
info.UpstreamModelName = body.Model
|
||||
data, err := json.Marshal(body)
|
||||
if info.IsModelMapped {
|
||||
body.Model = info.UpstreamModelName
|
||||
} else {
|
||||
info.UpstreamModelName = body.Model
|
||||
}
|
||||
data, err := common.Marshal(body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -154,7 +159,7 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
|
||||
|
||||
// Parse Doubao response
|
||||
var dResp responsePayload
|
||||
if err := json.Unmarshal(responseBody, &dResp); err != nil {
|
||||
if err := common.Unmarshal(responseBody, &dResp); err != nil {
|
||||
taskErr = service.TaskErrorWrapper(errors.Wrapf(err, "body: %s", responseBody), "unmarshal_response_body_failed", http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
@@ -165,8 +170,8 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
|
||||
}
|
||||
|
||||
ov := dto.NewOpenAIVideo()
|
||||
ov.ID = dResp.ID
|
||||
ov.TaskID = dResp.ID
|
||||
ov.ID = info.PublicTaskID
|
||||
ov.TaskID = info.PublicTaskID
|
||||
ov.CreatedAt = time.Now().Unix()
|
||||
ov.Model = info.OriginModelName
|
||||
|
||||
@@ -234,12 +239,7 @@ func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*
|
||||
}
|
||||
|
||||
metadata := req.Metadata
|
||||
medaBytes, err := json.Marshal(metadata)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "metadata marshal metadata failed")
|
||||
}
|
||||
err = json.Unmarshal(medaBytes, &r)
|
||||
if err != nil {
|
||||
if err := taskcommon.UnmarshalMetadata(metadata, &r); err != nil {
|
||||
return nil, errors.Wrap(err, "unmarshal metadata failed")
|
||||
}
|
||||
|
||||
@@ -248,7 +248,7 @@ func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*
|
||||
|
||||
func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) {
|
||||
resTask := responseTask{}
|
||||
if err := json.Unmarshal(respBody, &resTask); err != nil {
|
||||
if err := common.Unmarshal(respBody, &resTask); err != nil {
|
||||
return nil, errors.Wrap(err, "unmarshal task result failed")
|
||||
}
|
||||
|
||||
@@ -286,7 +286,7 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
|
||||
|
||||
func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, error) {
|
||||
var dResp responseTask
|
||||
if err := json.Unmarshal(originTask.Data, &dResp); err != nil {
|
||||
if err := common.Unmarshal(originTask.Data, &dResp); err != nil {
|
||||
return nil, errors.Wrap(err, "unmarshal doubao task data failed")
|
||||
}
|
||||
|
||||
@@ -307,6 +307,5 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, erro
|
||||
}
|
||||
}
|
||||
|
||||
jsonData, _ := common.Marshal(openAIVideo)
|
||||
return jsonData, nil
|
||||
return common.Marshal(openAIVideo)
|
||||
}
|
||||
|
||||
@@ -2,8 +2,6 @@ package gemini
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
@@ -16,10 +14,10 @@ import (
|
||||
"github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/relay/channel"
|
||||
taskcommon "github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/setting/model_setting"
|
||||
"github.com/QuantumNous/new-api/setting/system_setting"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/pkg/errors"
|
||||
)
|
||||
@@ -87,6 +85,7 @@ type operationResponse struct {
|
||||
// ============================
|
||||
|
||||
type TaskAdaptor struct {
|
||||
taskcommon.BaseBilling
|
||||
ChannelType int
|
||||
apiKey string
|
||||
baseURL string
|
||||
@@ -106,7 +105,7 @@ func (a *TaskAdaptor) ValidateRequestAndSetAction(c *gin.Context, info *relaycom
|
||||
|
||||
// BuildRequestURL constructs the upstream URL.
|
||||
func (a *TaskAdaptor) BuildRequestURL(info *relaycommon.RelayInfo) (string, error) {
|
||||
modelName := info.OriginModelName
|
||||
modelName := info.UpstreamModelName
|
||||
version := model_setting.GetGeminiVersionSetting(modelName)
|
||||
|
||||
return fmt.Sprintf(
|
||||
@@ -145,16 +144,11 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
|
||||
}
|
||||
|
||||
metadata := req.Metadata
|
||||
medaBytes, err := json.Marshal(metadata)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "metadata marshal metadata failed")
|
||||
}
|
||||
err = json.Unmarshal(medaBytes, &body.Parameters)
|
||||
if err != nil {
|
||||
if err := taskcommon.UnmarshalMetadata(metadata, &body.Parameters); err != nil {
|
||||
return nil, errors.Wrap(err, "unmarshal metadata failed")
|
||||
}
|
||||
|
||||
data, err := json.Marshal(body)
|
||||
data, err := common.Marshal(body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -175,16 +169,16 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
|
||||
_ = resp.Body.Close()
|
||||
|
||||
var s submitResponse
|
||||
if err := json.Unmarshal(responseBody, &s); err != nil {
|
||||
if err := common.Unmarshal(responseBody, &s); err != nil {
|
||||
return "", nil, service.TaskErrorWrapper(err, "unmarshal_response_failed", http.StatusInternalServerError)
|
||||
}
|
||||
if strings.TrimSpace(s.Name) == "" {
|
||||
return "", nil, service.TaskErrorWrapper(fmt.Errorf("missing operation name"), "invalid_response", http.StatusInternalServerError)
|
||||
}
|
||||
taskID = encodeLocalTaskID(s.Name)
|
||||
taskID = taskcommon.EncodeLocalTaskID(s.Name)
|
||||
ov := dto.NewOpenAIVideo()
|
||||
ov.ID = taskID
|
||||
ov.TaskID = taskID
|
||||
ov.ID = info.PublicTaskID
|
||||
ov.TaskID = info.PublicTaskID
|
||||
ov.CreatedAt = time.Now().Unix()
|
||||
ov.Model = info.OriginModelName
|
||||
c.JSON(http.StatusOK, ov)
|
||||
@@ -206,7 +200,7 @@ func (a *TaskAdaptor) FetchTask(baseUrl, key string, body map[string]any, proxy
|
||||
return nil, fmt.Errorf("invalid task_id")
|
||||
}
|
||||
|
||||
upstreamName, err := decodeLocalTaskID(taskID)
|
||||
upstreamName, err := taskcommon.DecodeLocalTaskID(taskID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("decode task_id failed: %w", err)
|
||||
}
|
||||
@@ -232,7 +226,7 @@ func (a *TaskAdaptor) FetchTask(baseUrl, key string, body map[string]any, proxy
|
||||
|
||||
func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) {
|
||||
var op operationResponse
|
||||
if err := json.Unmarshal(respBody, &op); err != nil {
|
||||
if err := common.Unmarshal(respBody, &op); err != nil {
|
||||
return nil, fmt.Errorf("unmarshal operation response failed: %w", err)
|
||||
}
|
||||
|
||||
@@ -254,9 +248,8 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
|
||||
ti.Status = model.TaskStatusSuccess
|
||||
ti.Progress = "100%"
|
||||
|
||||
taskID := encodeLocalTaskID(op.Name)
|
||||
ti.TaskID = taskID
|
||||
ti.Url = fmt.Sprintf("%s/v1/videos/%s/content", system_setting.ServerAddress, taskID)
|
||||
ti.TaskID = taskcommon.EncodeLocalTaskID(op.Name)
|
||||
// Url intentionally left empty — the caller constructs the proxy URL using the public task ID
|
||||
|
||||
// Extract URL from generateVideoResponse if available
|
||||
if len(op.Response.GenerateVideoResponse.GeneratedSamples) > 0 {
|
||||
@@ -269,7 +262,10 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
|
||||
}
|
||||
|
||||
func (a *TaskAdaptor) ConvertToOpenAIVideo(task *model.Task) ([]byte, error) {
|
||||
upstreamName, err := decodeLocalTaskID(task.TaskID)
|
||||
// Use GetUpstreamTaskID() to get the real upstream operation name for model extraction.
|
||||
// task.TaskID is now a public task_xxxx ID, no longer a base64-encoded upstream name.
|
||||
upstreamTaskID := task.GetUpstreamTaskID()
|
||||
upstreamName, err := taskcommon.DecodeLocalTaskID(upstreamTaskID)
|
||||
if err != nil {
|
||||
upstreamName = ""
|
||||
}
|
||||
@@ -297,18 +293,6 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(task *model.Task) ([]byte, error) {
|
||||
// helpers
|
||||
// ============================
|
||||
|
||||
func encodeLocalTaskID(name string) string {
|
||||
return base64.RawURLEncoding.EncodeToString([]byte(name))
|
||||
}
|
||||
|
||||
func decodeLocalTaskID(local string) (string, error) {
|
||||
b, err := base64.RawURLEncoding.DecodeString(local)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return string(b), nil
|
||||
}
|
||||
|
||||
var modelRe = regexp.MustCompile(`models/([^/]+)/operations/`)
|
||||
|
||||
func extractModelFromOperationName(name string) string {
|
||||
|
||||
@@ -2,7 +2,6 @@ package hailuo
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
@@ -18,12 +17,14 @@ import (
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/relay/channel"
|
||||
taskcommon "github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
)
|
||||
|
||||
// https://platform.minimaxi.com/docs/api-reference/video-generation-intro
|
||||
type TaskAdaptor struct {
|
||||
taskcommon.BaseBilling
|
||||
ChannelType int
|
||||
apiKey string
|
||||
baseURL string
|
||||
@@ -60,12 +61,12 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
|
||||
return nil, fmt.Errorf("invalid request type in context")
|
||||
}
|
||||
|
||||
body, err := a.convertToRequestPayload(&req)
|
||||
body, err := a.convertToRequestPayload(&req, info)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "convert request payload failed")
|
||||
}
|
||||
|
||||
data, err := json.Marshal(body)
|
||||
data, err := common.Marshal(body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -86,7 +87,7 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
|
||||
_ = resp.Body.Close()
|
||||
|
||||
var hResp VideoResponse
|
||||
if err := json.Unmarshal(responseBody, &hResp); err != nil {
|
||||
if err := common.Unmarshal(responseBody, &hResp); err != nil {
|
||||
taskErr = service.TaskErrorWrapper(errors.Wrapf(err, "body: %s", responseBody), "unmarshal_response_body_failed", http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
@@ -101,8 +102,8 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
|
||||
}
|
||||
|
||||
ov := dto.NewOpenAIVideo()
|
||||
ov.ID = hResp.TaskID
|
||||
ov.TaskID = hResp.TaskID
|
||||
ov.ID = info.PublicTaskID
|
||||
ov.TaskID = info.PublicTaskID
|
||||
ov.CreatedAt = time.Now().Unix()
|
||||
ov.Model = info.OriginModelName
|
||||
|
||||
@@ -141,8 +142,8 @@ func (a *TaskAdaptor) GetChannelName() string {
|
||||
return ChannelName
|
||||
}
|
||||
|
||||
func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*VideoRequest, error) {
|
||||
modelConfig := GetModelConfig(req.Model)
|
||||
func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq, info *relaycommon.RelayInfo) (*VideoRequest, error) {
|
||||
modelConfig := GetModelConfig(info.UpstreamModelName)
|
||||
duration := DefaultDuration
|
||||
if req.Duration > 0 {
|
||||
duration = req.Duration
|
||||
@@ -153,7 +154,7 @@ func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*
|
||||
}
|
||||
|
||||
videoRequest := &VideoRequest{
|
||||
Model: req.Model,
|
||||
Model: info.UpstreamModelName,
|
||||
Prompt: req.Prompt,
|
||||
Duration: &duration,
|
||||
Resolution: resolution,
|
||||
@@ -182,7 +183,7 @@ func (a *TaskAdaptor) parseResolutionFromSize(size string, modelConfig ModelConf
|
||||
|
||||
func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) {
|
||||
resTask := QueryTaskResponse{}
|
||||
if err := json.Unmarshal(respBody, &resTask); err != nil {
|
||||
if err := common.Unmarshal(respBody, &resTask); err != nil {
|
||||
return nil, errors.Wrap(err, "unmarshal task result failed")
|
||||
}
|
||||
|
||||
@@ -224,7 +225,7 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
|
||||
|
||||
func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, error) {
|
||||
var hailuoResp QueryTaskResponse
|
||||
if err := json.Unmarshal(originTask.Data, &hailuoResp); err != nil {
|
||||
if err := common.Unmarshal(originTask.Data, &hailuoResp); err != nil {
|
||||
return nil, errors.Wrap(err, "unmarshal hailuo task data failed")
|
||||
}
|
||||
|
||||
@@ -271,7 +272,7 @@ func (a *TaskAdaptor) buildVideoURL(_, fileID string) string {
|
||||
}
|
||||
|
||||
var retrieveResp RetrieveFileResponse
|
||||
if err := json.Unmarshal(responseBody, &retrieveResp); err != nil {
|
||||
if err := common.Unmarshal(responseBody, &retrieveResp); err != nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
|
||||
@@ -6,7 +6,6 @@ import (
|
||||
"crypto/sha256"
|
||||
"encoding/base64"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
@@ -25,6 +24,7 @@ import (
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/relay/channel"
|
||||
taskcommon "github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
)
|
||||
@@ -77,6 +77,7 @@ const (
|
||||
// ============================
|
||||
|
||||
type TaskAdaptor struct {
|
||||
taskcommon.BaseBilling
|
||||
ChannelType int
|
||||
accessKey string
|
||||
secretKey string
|
||||
@@ -164,11 +165,11 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
|
||||
}
|
||||
}
|
||||
|
||||
body, err := a.convertToRequestPayload(&req)
|
||||
body, err := a.convertToRequestPayload(&req, info)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "convert request payload failed")
|
||||
}
|
||||
data, err := json.Marshal(body)
|
||||
data, err := common.Marshal(body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -191,7 +192,7 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
|
||||
|
||||
// Parse Jimeng response
|
||||
var jResp responsePayload
|
||||
if err := json.Unmarshal(responseBody, &jResp); err != nil {
|
||||
if err := common.Unmarshal(responseBody, &jResp); err != nil {
|
||||
taskErr = service.TaskErrorWrapper(errors.Wrapf(err, "body: %s", responseBody), "unmarshal_response_body_failed", http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
@@ -202,8 +203,8 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
|
||||
}
|
||||
|
||||
ov := dto.NewOpenAIVideo()
|
||||
ov.ID = jResp.Data.TaskID
|
||||
ov.TaskID = jResp.Data.TaskID
|
||||
ov.ID = info.PublicTaskID
|
||||
ov.TaskID = info.PublicTaskID
|
||||
ov.CreatedAt = time.Now().Unix()
|
||||
ov.Model = info.OriginModelName
|
||||
c.JSON(http.StatusOK, ov)
|
||||
@@ -225,7 +226,7 @@ func (a *TaskAdaptor) FetchTask(baseUrl, key string, body map[string]any, proxy
|
||||
"req_key": "jimeng_vgfm_t2v_l20", // This is fixed value from doc: https://www.volcengine.com/docs/85621/1544774
|
||||
"task_id": taskID,
|
||||
}
|
||||
payloadBytes, err := json.Marshal(payload)
|
||||
payloadBytes, err := common.Marshal(payload)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "marshal fetch task payload failed")
|
||||
}
|
||||
@@ -377,9 +378,9 @@ func hmacSHA256(key []byte, data []byte) []byte {
|
||||
return h.Sum(nil)
|
||||
}
|
||||
|
||||
func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*requestPayload, error) {
|
||||
func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq, info *relaycommon.RelayInfo) (*requestPayload, error) {
|
||||
r := requestPayload{
|
||||
ReqKey: req.Model,
|
||||
ReqKey: info.UpstreamModelName,
|
||||
Prompt: req.Prompt,
|
||||
}
|
||||
|
||||
@@ -398,13 +399,7 @@ func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*
|
||||
r.BinaryDataBase64 = req.Images
|
||||
}
|
||||
}
|
||||
metadata := req.Metadata
|
||||
medaBytes, err := json.Marshal(metadata)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "metadata marshal metadata failed")
|
||||
}
|
||||
err = json.Unmarshal(medaBytes, &r)
|
||||
if err != nil {
|
||||
if err := taskcommon.UnmarshalMetadata(req.Metadata, &r); err != nil {
|
||||
return nil, errors.Wrap(err, "unmarshal metadata failed")
|
||||
}
|
||||
|
||||
@@ -432,7 +427,7 @@ func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*
|
||||
|
||||
func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) {
|
||||
resTask := responseTask{}
|
||||
if err := json.Unmarshal(respBody, &resTask); err != nil {
|
||||
if err := common.Unmarshal(respBody, &resTask); err != nil {
|
||||
return nil, errors.Wrap(err, "unmarshal task result failed")
|
||||
}
|
||||
taskResult := relaycommon.TaskInfo{}
|
||||
@@ -458,7 +453,7 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
|
||||
|
||||
func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, error) {
|
||||
var jimengResp responseTask
|
||||
if err := json.Unmarshal(originTask.Data, &jimengResp); err != nil {
|
||||
if err := common.Unmarshal(originTask.Data, &jimengResp); err != nil {
|
||||
return nil, errors.Wrap(err, "unmarshal jimeng task data failed")
|
||||
}
|
||||
|
||||
@@ -477,8 +472,7 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, erro
|
||||
}
|
||||
}
|
||||
|
||||
jsonData, _ := common.Marshal(openAIVideo)
|
||||
return jsonData, nil
|
||||
return common.Marshal(openAIVideo)
|
||||
}
|
||||
|
||||
func isNewAPIRelay(apiKey string) bool {
|
||||
|
||||
@@ -2,7 +2,6 @@ package kling
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
@@ -21,6 +20,7 @@ import (
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/relay/channel"
|
||||
taskcommon "github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
)
|
||||
@@ -97,6 +97,7 @@ type responsePayload struct {
|
||||
// ============================
|
||||
|
||||
type TaskAdaptor struct {
|
||||
taskcommon.BaseBilling
|
||||
ChannelType int
|
||||
apiKey string
|
||||
baseURL string
|
||||
@@ -149,14 +150,14 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
|
||||
}
|
||||
req := v.(relaycommon.TaskSubmitReq)
|
||||
|
||||
body, err := a.convertToRequestPayload(&req)
|
||||
body, err := a.convertToRequestPayload(&req, info)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if body.Image == "" && body.ImageTail == "" {
|
||||
c.Set("action", constant.TaskActionTextGenerate)
|
||||
}
|
||||
data, err := json.Marshal(body)
|
||||
data, err := common.Marshal(body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -180,7 +181,7 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
|
||||
}
|
||||
|
||||
var kResp responsePayload
|
||||
err = json.Unmarshal(responseBody, &kResp)
|
||||
err = common.Unmarshal(responseBody, &kResp)
|
||||
if err != nil {
|
||||
taskErr = service.TaskErrorWrapper(err, "unmarshal_response_failed", http.StatusInternalServerError)
|
||||
return
|
||||
@@ -190,8 +191,8 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
|
||||
return
|
||||
}
|
||||
ov := dto.NewOpenAIVideo()
|
||||
ov.ID = kResp.Data.TaskId
|
||||
ov.TaskID = kResp.Data.TaskId
|
||||
ov.ID = info.PublicTaskID
|
||||
ov.TaskID = info.PublicTaskID
|
||||
ov.CreatedAt = time.Now().Unix()
|
||||
ov.Model = info.OriginModelName
|
||||
c.JSON(http.StatusOK, ov)
|
||||
@@ -247,15 +248,15 @@ func (a *TaskAdaptor) GetChannelName() string {
|
||||
// helpers
|
||||
// ============================
|
||||
|
||||
func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*requestPayload, error) {
|
||||
func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq, info *relaycommon.RelayInfo) (*requestPayload, error) {
|
||||
r := requestPayload{
|
||||
Prompt: req.Prompt,
|
||||
Image: req.Image,
|
||||
Mode: defaultString(req.Mode, "std"),
|
||||
Duration: fmt.Sprintf("%d", defaultInt(req.Duration, 5)),
|
||||
Mode: taskcommon.DefaultString(req.Mode, "std"),
|
||||
Duration: fmt.Sprintf("%d", taskcommon.DefaultInt(req.Duration, 5)),
|
||||
AspectRatio: a.getAspectRatio(req.Size),
|
||||
ModelName: req.Model,
|
||||
Model: req.Model, // Keep consistent with model_name, double writing improves compatibility
|
||||
ModelName: info.UpstreamModelName,
|
||||
Model: info.UpstreamModelName,
|
||||
CfgScale: 0.5,
|
||||
StaticMask: "",
|
||||
DynamicMasks: []DynamicMask{},
|
||||
@@ -265,14 +266,9 @@ func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*
|
||||
}
|
||||
if r.ModelName == "" {
|
||||
r.ModelName = "kling-v1"
|
||||
r.Model = "kling-v1"
|
||||
}
|
||||
metadata := req.Metadata
|
||||
medaBytes, err := json.Marshal(metadata)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "metadata marshal metadata failed")
|
||||
}
|
||||
err = json.Unmarshal(medaBytes, &r)
|
||||
if err != nil {
|
||||
if err := taskcommon.UnmarshalMetadata(req.Metadata, &r); err != nil {
|
||||
return nil, errors.Wrap(err, "unmarshal metadata failed")
|
||||
}
|
||||
return &r, nil
|
||||
@@ -291,20 +287,6 @@ func (a *TaskAdaptor) getAspectRatio(size string) string {
|
||||
}
|
||||
}
|
||||
|
||||
func defaultString(s, def string) string {
|
||||
if strings.TrimSpace(s) == "" {
|
||||
return def
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func defaultInt(v int, def int) int {
|
||||
if v == 0 {
|
||||
return def
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
// ============================
|
||||
// JWT helpers
|
||||
// ============================
|
||||
@@ -340,7 +322,7 @@ func (a *TaskAdaptor) createJWTTokenWithKey(apiKey string) (string, error) {
|
||||
func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) {
|
||||
taskInfo := &relaycommon.TaskInfo{}
|
||||
resPayload := responsePayload{}
|
||||
err := json.Unmarshal(respBody, &resPayload)
|
||||
err := common.Unmarshal(respBody, &resPayload)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "failed to unmarshal response body")
|
||||
}
|
||||
@@ -374,7 +356,7 @@ func isNewAPIRelay(apiKey string) bool {
|
||||
|
||||
func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, error) {
|
||||
var klingResp responsePayload
|
||||
if err := json.Unmarshal(originTask.Data, &klingResp); err != nil {
|
||||
if err := common.Unmarshal(originTask.Data, &klingResp); err != nil {
|
||||
return nil, errors.Wrap(err, "unmarshal kling task data failed")
|
||||
}
|
||||
|
||||
@@ -401,6 +383,5 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, erro
|
||||
Code: fmt.Sprintf("%d", klingResp.Code),
|
||||
}
|
||||
}
|
||||
jsonData, _ := common.Marshal(openAIVideo)
|
||||
return jsonData, nil
|
||||
return common.Marshal(openAIVideo)
|
||||
}
|
||||
|
||||
@@ -1,9 +1,13 @@
|
||||
package sora
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"io"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"net/textproto"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
@@ -11,12 +15,13 @@ import (
|
||||
"github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/relay/channel"
|
||||
taskcommon "github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/setting/system_setting"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/pkg/errors"
|
||||
"github.com/tidwall/sjson"
|
||||
)
|
||||
|
||||
// ============================
|
||||
@@ -57,6 +62,7 @@ type responseTask struct {
|
||||
// ============================
|
||||
|
||||
type TaskAdaptor struct {
|
||||
taskcommon.BaseBilling
|
||||
ChannelType int
|
||||
apiKey string
|
||||
baseURL string
|
||||
@@ -69,15 +75,15 @@ func (a *TaskAdaptor) Init(info *relaycommon.RelayInfo) {
|
||||
}
|
||||
|
||||
func validateRemixRequest(c *gin.Context) *dto.TaskError {
|
||||
var req struct {
|
||||
Prompt string `json:"prompt"`
|
||||
}
|
||||
var req relaycommon.TaskSubmitReq
|
||||
if err := common.UnmarshalBodyReusable(c, &req); err != nil {
|
||||
return service.TaskErrorWrapperLocal(err, "invalid_request", http.StatusBadRequest)
|
||||
}
|
||||
if strings.TrimSpace(req.Prompt) == "" {
|
||||
return service.TaskErrorWrapperLocal(fmt.Errorf("field prompt is required"), "invalid_request", http.StatusBadRequest)
|
||||
}
|
||||
// 存储原始请求到 context,与 ValidateMultipartDirect 路径保持一致
|
||||
c.Set("task_request", req)
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -88,6 +94,41 @@ func (a *TaskAdaptor) ValidateRequestAndSetAction(c *gin.Context, info *relaycom
|
||||
return relaycommon.ValidateMultipartDirect(c, info)
|
||||
}
|
||||
|
||||
// EstimateBilling 根据用户请求的 seconds 和 size 计算 OtherRatios。
|
||||
func (a *TaskAdaptor) EstimateBilling(c *gin.Context, info *relaycommon.RelayInfo) map[string]float64 {
|
||||
// remix 路径的 OtherRatios 已在 ResolveOriginTask 中设置
|
||||
if info.Action == constant.TaskActionRemix {
|
||||
return nil
|
||||
}
|
||||
|
||||
req, err := relaycommon.GetTaskRequest(c)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
seconds, _ := strconv.Atoi(req.Seconds)
|
||||
if seconds == 0 {
|
||||
seconds = req.Duration
|
||||
}
|
||||
if seconds <= 0 {
|
||||
seconds = 4
|
||||
}
|
||||
|
||||
size := req.Size
|
||||
if size == "" {
|
||||
size = "720x1280"
|
||||
}
|
||||
|
||||
ratios := map[string]float64{
|
||||
"seconds": float64(seconds),
|
||||
"size": 1,
|
||||
}
|
||||
if size == "1792x1024" || size == "1024x1792" {
|
||||
ratios["size"] = 1.666667
|
||||
}
|
||||
return ratios
|
||||
}
|
||||
|
||||
func (a *TaskAdaptor) BuildRequestURL(info *relaycommon.RelayInfo) (string, error) {
|
||||
if info.Action == constant.TaskActionRemix {
|
||||
return fmt.Sprintf("%s/v1/videos/%s/remix", a.baseURL, info.OriginTaskID), nil
|
||||
@@ -107,6 +148,74 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "get_request_body_failed")
|
||||
}
|
||||
cachedBody, err := storage.Bytes()
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "read_body_bytes_failed")
|
||||
}
|
||||
contentType := c.GetHeader("Content-Type")
|
||||
|
||||
if strings.HasPrefix(contentType, "application/json") {
|
||||
var bodyMap map[string]interface{}
|
||||
if err := common.Unmarshal(cachedBody, &bodyMap); err == nil {
|
||||
bodyMap["model"] = info.UpstreamModelName
|
||||
if newBody, err := common.Marshal(bodyMap); err == nil {
|
||||
return bytes.NewReader(newBody), nil
|
||||
}
|
||||
}
|
||||
return bytes.NewReader(cachedBody), nil
|
||||
}
|
||||
|
||||
if strings.Contains(contentType, "multipart/form-data") {
|
||||
formData, err := common.ParseMultipartFormReusable(c)
|
||||
if err != nil {
|
||||
return bytes.NewReader(cachedBody), nil
|
||||
}
|
||||
var buf bytes.Buffer
|
||||
writer := multipart.NewWriter(&buf)
|
||||
writer.WriteField("model", info.UpstreamModelName)
|
||||
for key, values := range formData.Value {
|
||||
if key == "model" {
|
||||
continue
|
||||
}
|
||||
for _, v := range values {
|
||||
writer.WriteField(key, v)
|
||||
}
|
||||
}
|
||||
for fieldName, fileHeaders := range formData.File {
|
||||
for _, fh := range fileHeaders {
|
||||
f, err := fh.Open()
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
ct := fh.Header.Get("Content-Type")
|
||||
if ct == "" || ct == "application/octet-stream" {
|
||||
buf512 := make([]byte, 512)
|
||||
n, _ := io.ReadFull(f, buf512)
|
||||
ct = http.DetectContentType(buf512[:n])
|
||||
// Re-open after sniffing so the full content is copied below
|
||||
f.Close()
|
||||
f, err = fh.Open()
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
}
|
||||
h := make(textproto.MIMEHeader)
|
||||
h.Set("Content-Disposition", fmt.Sprintf(`form-data; name="%s"; filename="%s"`, fieldName, fh.Filename))
|
||||
h.Set("Content-Type", ct)
|
||||
part, err := writer.CreatePart(h)
|
||||
if err != nil {
|
||||
f.Close()
|
||||
continue
|
||||
}
|
||||
io.Copy(part, f)
|
||||
f.Close()
|
||||
}
|
||||
}
|
||||
writer.Close()
|
||||
c.Request.Header.Set("Content-Type", writer.FormDataContentType())
|
||||
return &buf, nil
|
||||
}
|
||||
|
||||
return common.ReaderOnly(storage), nil
|
||||
}
|
||||
|
||||
@@ -116,7 +225,7 @@ func (a *TaskAdaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, req
|
||||
}
|
||||
|
||||
// DoResponse handles upstream response, returns taskID etc.
|
||||
func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, _ *relaycommon.RelayInfo) (taskID string, taskData []byte, taskErr *dto.TaskError) {
|
||||
func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (taskID string, taskData []byte, taskErr *dto.TaskError) {
|
||||
responseBody, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
taskErr = service.TaskErrorWrapper(err, "read_response_body_failed", http.StatusInternalServerError)
|
||||
@@ -131,17 +240,20 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, _ *relayco
|
||||
return
|
||||
}
|
||||
|
||||
if dResp.ID == "" {
|
||||
if dResp.TaskID == "" {
|
||||
taskErr = service.TaskErrorWrapper(fmt.Errorf("task_id is empty"), "invalid_response", http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
dResp.ID = dResp.TaskID
|
||||
dResp.TaskID = ""
|
||||
upstreamID := dResp.ID
|
||||
if upstreamID == "" {
|
||||
upstreamID = dResp.TaskID
|
||||
}
|
||||
if upstreamID == "" {
|
||||
taskErr = service.TaskErrorWrapper(fmt.Errorf("task_id is empty"), "invalid_response", http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
// 使用公开 task_xxxx ID 返回给客户端
|
||||
dResp.ID = info.PublicTaskID
|
||||
dResp.TaskID = info.PublicTaskID
|
||||
c.JSON(http.StatusOK, dResp)
|
||||
return dResp.ID, responseBody, nil
|
||||
return upstreamID, responseBody, nil
|
||||
}
|
||||
|
||||
// FetchTask fetch task status
|
||||
@@ -192,7 +304,7 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
|
||||
taskResult.Status = model.TaskStatusInProgress
|
||||
case "completed":
|
||||
taskResult.Status = model.TaskStatusSuccess
|
||||
taskResult.Url = fmt.Sprintf("%s/v1/videos/%s/content", system_setting.ServerAddress, resTask.ID)
|
||||
// Url intentionally left empty — the caller constructs the proxy URL using the public task ID
|
||||
case "failed", "cancelled":
|
||||
taskResult.Status = model.TaskStatusFailure
|
||||
if resTask.Error != nil {
|
||||
@@ -210,5 +322,10 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
|
||||
}
|
||||
|
||||
func (a *TaskAdaptor) ConvertToOpenAIVideo(task *model.Task) ([]byte, error) {
|
||||
return task.Data, nil
|
||||
data := task.Data
|
||||
var err error
|
||||
if data, err = sjson.SetBytes(data, "id", task.TaskID); err != nil {
|
||||
return nil, errors.Wrap(err, "set id failed")
|
||||
}
|
||||
return data, nil
|
||||
}
|
||||
|
||||
@@ -2,18 +2,16 @@ package suno
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/relay/channel"
|
||||
taskcommon "github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
|
||||
@@ -21,11 +19,16 @@ import (
|
||||
)
|
||||
|
||||
type TaskAdaptor struct {
|
||||
taskcommon.BaseBilling
|
||||
ChannelType int
|
||||
}
|
||||
|
||||
// ParseTaskResult is not used for Suno tasks.
|
||||
// Suno polling uses a dedicated batch-fetch path (service.UpdateSunoTasks) that
|
||||
// receives dto.TaskResponse[[]dto.SunoDataResponse] from the upstream /fetch API.
|
||||
// This differs from the per-task polling used by video adaptors.
|
||||
func (a *TaskAdaptor) ParseTaskResult([]byte) (*relaycommon.TaskInfo, error) {
|
||||
return nil, fmt.Errorf("not implement") // todo implement this method if needed
|
||||
return nil, fmt.Errorf("suno uses batch polling via UpdateSunoTasks, ParseTaskResult is not applicable")
|
||||
}
|
||||
|
||||
func (a *TaskAdaptor) Init(info *relaycommon.RelayInfo) {
|
||||
@@ -47,13 +50,13 @@ func (a *TaskAdaptor) ValidateRequestAndSetAction(c *gin.Context, info *relaycom
|
||||
return
|
||||
}
|
||||
|
||||
if sunoRequest.ContinueClipId != "" {
|
||||
if sunoRequest.TaskID == "" {
|
||||
taskErr = service.TaskErrorWrapperLocal(fmt.Errorf("task id is empty"), "invalid_request", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
info.OriginTaskID = sunoRequest.TaskID
|
||||
}
|
||||
//if sunoRequest.ContinueClipId != "" {
|
||||
// if sunoRequest.TaskID == "" {
|
||||
// taskErr = service.TaskErrorWrapperLocal(fmt.Errorf("task id is empty"), "invalid_request", http.StatusBadRequest)
|
||||
// return
|
||||
// }
|
||||
// info.OriginTaskID = sunoRequest.TaskID
|
||||
//}
|
||||
|
||||
info.Action = action
|
||||
c.Set("task_request", sunoRequest)
|
||||
@@ -76,12 +79,9 @@ func (a *TaskAdaptor) BuildRequestHeader(c *gin.Context, req *http.Request, info
|
||||
func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayInfo) (io.Reader, error) {
|
||||
sunoRequest, ok := c.Get("task_request")
|
||||
if !ok {
|
||||
err := common.UnmarshalBodyReusable(c, &sunoRequest)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return nil, fmt.Errorf("task_request not found in context")
|
||||
}
|
||||
data, err := json.Marshal(sunoRequest)
|
||||
data, err := common.Marshal(sunoRequest)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -99,7 +99,7 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
|
||||
return
|
||||
}
|
||||
var sunoResponse dto.TaskResponse[string]
|
||||
err = json.Unmarshal(responseBody, &sunoResponse)
|
||||
err = common.Unmarshal(responseBody, &sunoResponse)
|
||||
if err != nil {
|
||||
taskErr = service.TaskErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError)
|
||||
return
|
||||
@@ -109,17 +109,13 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
|
||||
return
|
||||
}
|
||||
|
||||
for k, v := range resp.Header {
|
||||
c.Writer.Header().Set(k, v[0])
|
||||
}
|
||||
c.Writer.Header().Set("Content-Type", "application/json")
|
||||
c.Writer.WriteHeader(resp.StatusCode)
|
||||
|
||||
_, err = io.Copy(c.Writer, bytes.NewBuffer(responseBody))
|
||||
if err != nil {
|
||||
taskErr = service.TaskErrorWrapper(err, "copy_response_body_failed", http.StatusInternalServerError)
|
||||
return
|
||||
// 使用公开 task_xxxx ID 替换上游 ID 返回给客户端
|
||||
publicResponse := dto.TaskResponse[string]{
|
||||
Code: sunoResponse.Code,
|
||||
Message: sunoResponse.Message,
|
||||
Data: info.PublicTaskID,
|
||||
}
|
||||
c.JSON(http.StatusOK, publicResponse)
|
||||
|
||||
return sunoResponse.Data, nil, nil
|
||||
}
|
||||
@@ -134,7 +130,7 @@ func (a *TaskAdaptor) GetChannelName() string {
|
||||
|
||||
func (a *TaskAdaptor) FetchTask(baseUrl, key string, body map[string]any, proxy string) (*http.Response, error) {
|
||||
requestUrl := fmt.Sprintf("%s/suno/fetch", baseUrl)
|
||||
byteBody, err := json.Marshal(body)
|
||||
byteBody, err := common.Marshal(body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -144,13 +140,6 @@ func (a *TaskAdaptor) FetchTask(baseUrl, key string, body map[string]any, proxy
|
||||
common.SysLog(fmt.Sprintf("Get Task error: %v", err))
|
||||
return nil, err
|
||||
}
|
||||
defer req.Body.Close()
|
||||
// 设置超时时间
|
||||
timeout := time.Second * 15
|
||||
ctx, cancel := context.WithTimeout(context.Background(), timeout)
|
||||
defer cancel()
|
||||
// 使用带有超时的 context 创建新的请求
|
||||
req = req.WithContext(ctx)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Authorization", "Bearer "+key)
|
||||
client, err := service.GetHttpClientWithProxy(proxy)
|
||||
|
||||
@@ -0,0 +1,95 @@
|
||||
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
|
||||
}
|
||||
@@ -2,13 +2,12 @@ package vertex
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"regexp"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
@@ -17,6 +16,7 @@ import (
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/relay/channel"
|
||||
taskcommon "github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
|
||||
vertexcore "github.com/QuantumNous/new-api/relay/channel/vertex"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
@@ -62,6 +62,7 @@ type operationResponse struct {
|
||||
// ============================
|
||||
|
||||
type TaskAdaptor struct {
|
||||
taskcommon.BaseBilling
|
||||
ChannelType int
|
||||
apiKey string
|
||||
baseURL string
|
||||
@@ -82,10 +83,10 @@ func (a *TaskAdaptor) ValidateRequestAndSetAction(c *gin.Context, info *relaycom
|
||||
// BuildRequestURL constructs the upstream URL.
|
||||
func (a *TaskAdaptor) BuildRequestURL(info *relaycommon.RelayInfo) (string, error) {
|
||||
adc := &vertexcore.Credentials{}
|
||||
if err := json.Unmarshal([]byte(a.apiKey), adc); err != nil {
|
||||
if err := common.Unmarshal([]byte(a.apiKey), adc); err != nil {
|
||||
return "", fmt.Errorf("failed to decode credentials: %w", err)
|
||||
}
|
||||
modelName := info.OriginModelName
|
||||
modelName := info.UpstreamModelName
|
||||
if modelName == "" {
|
||||
modelName = "veo-3.0-generate-001"
|
||||
}
|
||||
@@ -116,7 +117,7 @@ func (a *TaskAdaptor) BuildRequestHeader(c *gin.Context, req *http.Request, info
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
adc := &vertexcore.Credentials{}
|
||||
if err := json.Unmarshal([]byte(a.apiKey), adc); err != nil {
|
||||
if err := common.Unmarshal([]byte(a.apiKey), adc); err != nil {
|
||||
return fmt.Errorf("failed to decode credentials: %w", err)
|
||||
}
|
||||
|
||||
@@ -133,6 +134,28 @@ func (a *TaskAdaptor) BuildRequestHeader(c *gin.Context, req *http.Request, info
|
||||
return nil
|
||||
}
|
||||
|
||||
// EstimateBilling 根据用户请求中的 sampleCount 计算 OtherRatios。
|
||||
func (a *TaskAdaptor) EstimateBilling(c *gin.Context, _ *relaycommon.RelayInfo) map[string]float64 {
|
||||
sampleCount := 1
|
||||
v, ok := c.Get("task_request")
|
||||
if ok {
|
||||
req := v.(relaycommon.TaskSubmitReq)
|
||||
if req.Metadata != nil {
|
||||
if sc, exists := req.Metadata["sampleCount"]; exists {
|
||||
if i, ok := sc.(int); ok && i > 0 {
|
||||
sampleCount = i
|
||||
}
|
||||
if f, ok := sc.(float64); ok && int(f) > 0 {
|
||||
sampleCount = int(f)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return map[string]float64{
|
||||
"sampleCount": float64(sampleCount),
|
||||
}
|
||||
}
|
||||
|
||||
// BuildRequestBody converts request into Vertex specific format.
|
||||
func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayInfo) (io.Reader, error) {
|
||||
v, ok := c.Get("task_request")
|
||||
@@ -166,25 +189,7 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
|
||||
return nil, fmt.Errorf("sampleCount must be greater than 0")
|
||||
}
|
||||
|
||||
// if req.Duration > 0 {
|
||||
// body.Parameters["durationSeconds"] = req.Duration
|
||||
// } else if req.Seconds != "" {
|
||||
// seconds, err := strconv.Atoi(req.Seconds)
|
||||
// if err != nil {
|
||||
// return nil, errors.Wrap(err, "convert seconds to int failed")
|
||||
// }
|
||||
// body.Parameters["durationSeconds"] = seconds
|
||||
// }
|
||||
|
||||
info.PriceData.OtherRatios = map[string]float64{
|
||||
"sampleCount": float64(body.Parameters["sampleCount"].(int)),
|
||||
}
|
||||
|
||||
// if v, ok := body.Parameters["durationSeconds"]; ok {
|
||||
// info.PriceData.OtherRatios["durationSeconds"] = float64(v.(int))
|
||||
// }
|
||||
|
||||
data, err := json.Marshal(body)
|
||||
data, err := common.Marshal(body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -205,14 +210,19 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
|
||||
_ = resp.Body.Close()
|
||||
|
||||
var s submitResponse
|
||||
if err := json.Unmarshal(responseBody, &s); err != nil {
|
||||
if err := common.Unmarshal(responseBody, &s); err != nil {
|
||||
return "", nil, service.TaskErrorWrapper(err, "unmarshal_response_failed", http.StatusInternalServerError)
|
||||
}
|
||||
if strings.TrimSpace(s.Name) == "" {
|
||||
return "", nil, service.TaskErrorWrapper(fmt.Errorf("missing operation name"), "invalid_response", http.StatusInternalServerError)
|
||||
}
|
||||
localID := encodeLocalTaskID(s.Name)
|
||||
c.JSON(http.StatusOK, gin.H{"task_id": localID})
|
||||
localID := taskcommon.EncodeLocalTaskID(s.Name)
|
||||
ov := dto.NewOpenAIVideo()
|
||||
ov.ID = info.PublicTaskID
|
||||
ov.TaskID = info.PublicTaskID
|
||||
ov.CreatedAt = time.Now().Unix()
|
||||
ov.Model = info.OriginModelName
|
||||
c.JSON(http.StatusOK, ov)
|
||||
return localID, responseBody, nil
|
||||
}
|
||||
|
||||
@@ -225,7 +235,7 @@ func (a *TaskAdaptor) FetchTask(baseUrl, key string, body map[string]any, proxy
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("invalid task_id")
|
||||
}
|
||||
upstreamName, err := decodeLocalTaskID(taskID)
|
||||
upstreamName, err := taskcommon.DecodeLocalTaskID(taskID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("decode task_id failed: %w", err)
|
||||
}
|
||||
@@ -245,12 +255,12 @@ func (a *TaskAdaptor) FetchTask(baseUrl, key string, body map[string]any, proxy
|
||||
url = fmt.Sprintf("https://%s-aiplatform.googleapis.com/v1/projects/%s/locations/%s/publishers/google/models/%s:fetchPredictOperation", region, project, region, modelName)
|
||||
}
|
||||
payload := map[string]string{"operationName": upstreamName}
|
||||
data, err := json.Marshal(payload)
|
||||
data, err := common.Marshal(payload)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
adc := &vertexcore.Credentials{}
|
||||
if err := json.Unmarshal([]byte(key), adc); err != nil {
|
||||
if err := common.Unmarshal([]byte(key), adc); err != nil {
|
||||
return nil, fmt.Errorf("failed to decode credentials: %w", err)
|
||||
}
|
||||
token, err := vertexcore.AcquireAccessToken(*adc, proxy)
|
||||
@@ -274,7 +284,7 @@ func (a *TaskAdaptor) FetchTask(baseUrl, key string, body map[string]any, proxy
|
||||
|
||||
func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) {
|
||||
var op operationResponse
|
||||
if err := json.Unmarshal(respBody, &op); err != nil {
|
||||
if err := common.Unmarshal(respBody, &op); err != nil {
|
||||
return nil, fmt.Errorf("unmarshal operation response failed: %w", err)
|
||||
}
|
||||
ti := &relaycommon.TaskInfo{}
|
||||
@@ -338,7 +348,10 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
|
||||
}
|
||||
|
||||
func (a *TaskAdaptor) ConvertToOpenAIVideo(task *model.Task) ([]byte, error) {
|
||||
upstreamName, err := decodeLocalTaskID(task.TaskID)
|
||||
// Use GetUpstreamTaskID() to get the real upstream operation name for model extraction.
|
||||
// task.TaskID is now a public task_xxxx ID, no longer a base64-encoded upstream name.
|
||||
upstreamTaskID := task.GetUpstreamTaskID()
|
||||
upstreamName, err := taskcommon.DecodeLocalTaskID(upstreamTaskID)
|
||||
if err != nil {
|
||||
upstreamName = ""
|
||||
}
|
||||
@@ -353,8 +366,8 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(task *model.Task) ([]byte, error) {
|
||||
v.SetProgressStr(task.Progress)
|
||||
v.CreatedAt = task.CreatedAt
|
||||
v.CompletedAt = task.UpdatedAt
|
||||
if strings.HasPrefix(task.FailReason, "data:") && len(task.FailReason) > 0 {
|
||||
v.SetMetadata("url", task.FailReason)
|
||||
if resultURL := task.GetResultURL(); strings.HasPrefix(resultURL, "data:") && len(resultURL) > 0 {
|
||||
v.SetMetadata("url", resultURL)
|
||||
}
|
||||
|
||||
return common.Marshal(v)
|
||||
@@ -364,18 +377,6 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(task *model.Task) ([]byte, error) {
|
||||
// helpers
|
||||
// ============================
|
||||
|
||||
func encodeLocalTaskID(name string) string {
|
||||
return base64.RawURLEncoding.EncodeToString([]byte(name))
|
||||
}
|
||||
|
||||
func decodeLocalTaskID(local string) (string, error) {
|
||||
b, err := base64.RawURLEncoding.DecodeString(local)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return string(b), nil
|
||||
}
|
||||
|
||||
var regionRe = regexp.MustCompile(`locations/([a-z0-9-]+)/`)
|
||||
|
||||
func extractRegionFromOperationName(name string) string {
|
||||
|
||||
@@ -2,7 +2,6 @@ package vidu
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
@@ -16,6 +15,7 @@ import (
|
||||
"github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/relay/channel"
|
||||
taskcommon "github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
|
||||
@@ -73,6 +73,7 @@ type creation struct {
|
||||
// ============================
|
||||
|
||||
type TaskAdaptor struct {
|
||||
taskcommon.BaseBilling
|
||||
ChannelType int
|
||||
baseURL string
|
||||
}
|
||||
@@ -115,7 +116,7 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
|
||||
}
|
||||
req := v.(relaycommon.TaskSubmitReq)
|
||||
|
||||
body, err := a.convertToRequestPayload(&req)
|
||||
body, err := a.convertToRequestPayload(&req, info)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -127,7 +128,7 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
|
||||
}
|
||||
}
|
||||
|
||||
data, err := json.Marshal(body)
|
||||
data, err := common.Marshal(body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -168,7 +169,7 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
|
||||
}
|
||||
|
||||
var vResp responsePayload
|
||||
err = json.Unmarshal(responseBody, &vResp)
|
||||
err = common.Unmarshal(responseBody, &vResp)
|
||||
if err != nil {
|
||||
taskErr = service.TaskErrorWrapper(errors.Wrap(err, fmt.Sprintf("%s", responseBody)), "unmarshal_response_failed", http.StatusInternalServerError)
|
||||
return
|
||||
@@ -180,8 +181,8 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
|
||||
}
|
||||
|
||||
ov := dto.NewOpenAIVideo()
|
||||
ov.ID = vResp.TaskId
|
||||
ov.TaskID = vResp.TaskId
|
||||
ov.ID = info.PublicTaskID
|
||||
ov.TaskID = info.PublicTaskID
|
||||
ov.CreatedAt = time.Now().Unix()
|
||||
ov.Model = info.OriginModelName
|
||||
c.JSON(http.StatusOK, ov)
|
||||
@@ -223,47 +224,27 @@ func (a *TaskAdaptor) GetChannelName() string {
|
||||
// helpers
|
||||
// ============================
|
||||
|
||||
func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*requestPayload, error) {
|
||||
func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq, info *relaycommon.RelayInfo) (*requestPayload, error) {
|
||||
r := requestPayload{
|
||||
Model: defaultString(req.Model, "viduq1"),
|
||||
Model: taskcommon.DefaultString(info.UpstreamModelName, "viduq1"),
|
||||
Images: req.Images,
|
||||
Prompt: req.Prompt,
|
||||
Duration: defaultInt(req.Duration, 5),
|
||||
Resolution: defaultString(req.Size, "1080p"),
|
||||
Duration: taskcommon.DefaultInt(req.Duration, 5),
|
||||
Resolution: taskcommon.DefaultString(req.Size, "1080p"),
|
||||
MovementAmplitude: "auto",
|
||||
Bgm: false,
|
||||
}
|
||||
metadata := req.Metadata
|
||||
medaBytes, err := json.Marshal(metadata)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "metadata marshal metadata failed")
|
||||
}
|
||||
err = json.Unmarshal(medaBytes, &r)
|
||||
if err != nil {
|
||||
if err := taskcommon.UnmarshalMetadata(req.Metadata, &r); err != nil {
|
||||
return nil, errors.Wrap(err, "unmarshal metadata failed")
|
||||
}
|
||||
return &r, nil
|
||||
}
|
||||
|
||||
func defaultString(value, defaultValue string) string {
|
||||
if value == "" {
|
||||
return defaultValue
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
func defaultInt(value, defaultValue int) int {
|
||||
if value == 0 {
|
||||
return defaultValue
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) {
|
||||
taskInfo := &relaycommon.TaskInfo{}
|
||||
|
||||
var taskResp taskResultResponse
|
||||
err := json.Unmarshal(respBody, &taskResp)
|
||||
err := common.Unmarshal(respBody, &taskResp)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "failed to unmarshal response body")
|
||||
}
|
||||
@@ -293,7 +274,7 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
|
||||
|
||||
func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, error) {
|
||||
var viduResp taskResultResponse
|
||||
if err := json.Unmarshal(originTask.Data, &viduResp); err != nil {
|
||||
if err := common.Unmarshal(originTask.Data, &viduResp); err != nil {
|
||||
return nil, errors.Wrap(err, "unmarshal vidu task data failed")
|
||||
}
|
||||
|
||||
@@ -315,6 +296,5 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, erro
|
||||
}
|
||||
}
|
||||
|
||||
jsonData, _ := common.Marshal(openAIVideo)
|
||||
return jsonData, nil
|
||||
return common.Marshal(openAIVideo)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user