mirror of
https://github.com/QuantumNous/new-api.git
synced 2026-10-28 08:35:33 +08:00
CaIon
db80ac8e0c
Before billing, the gateway walked the whole decoded request of a task plugin and rejected duration, seconds and count members that were out of range at any depth, including vendor parameters passed through metadata. A plugin had to declare duration-auto@1 for the gateway to let a duration of -1 through. The gateway no longer reads billing quantities from the request. It still bounds every usage fact a plugin returns before quota is calculated: second facts at most 3600, count facts at most 128, token and credit facts saturated at the quota limit. Each plugin validates its own request and derives its facts from the request it sends upstream. duration-auto@1 is still accepted and has no effect. Kling now bills the duration it sends. metadata.duration, which is the whole body on native routes, overrides the top-level duration and must be a whole number of seconds from 1 to 3600. Before, a metadata duration of 10 was sent upstream while 5 seconds were billed.
105 lines
4.3 KiB
Go
105 lines
4.3 KiB
Go
package plugins_test
|
|
|
|
import (
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/dto"
|
|
"github.com/QuantumNous/new-api/pkg/jsplugin"
|
|
builtinplugins "github.com/QuantumNous/new-api/plugins"
|
|
taskplugin "github.com/QuantumNous/new-api/relay/channel/task/jsplugin"
|
|
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestKlingResponsesProtocol(t *testing.T) {
|
|
testVideoResponsesProtocol(t, videoResponsesTestCase{
|
|
pluginKey: "kling",
|
|
model: "kling-v2-master",
|
|
requestBody: map[string]any{
|
|
"model": "kling-v2-master",
|
|
"input": []any{map[string]any{"role": "user", "content": []any{
|
|
map[string]any{"type": "input_text", "text": "camera orbit"},
|
|
map[string]any{"type": "input_image", "image_url": "https://cdn.example/frame.png"},
|
|
}}},
|
|
"seconds": 10,
|
|
"metadata": map[string]any{
|
|
"mode": "pro",
|
|
},
|
|
},
|
|
wantAction: "image_to_video",
|
|
wantRequest: map[string]any{
|
|
"model": "kling-v2-master",
|
|
"prompt": "camera orbit",
|
|
"image": "https://cdn.example/frame.png",
|
|
"duration": float64(10),
|
|
"metadata": map[string]any{"mode": "pro"},
|
|
},
|
|
wantUsageKeys: []string{"units"},
|
|
wantVendorName: "kling",
|
|
})
|
|
|
|
// Kling bills the duration it sends, and metadata.duration (the whole body
|
|
// on native routes) overrides the top-level one, so the plugin bounds that
|
|
// value before the host reserves quota.
|
|
t.Run("bills the sent duration and rejects out-of-range durations before billing", func(t *testing.T) {
|
|
source, err := builtinplugins.Source("kling")
|
|
require.NoError(t, err)
|
|
plugin, err := jsplugin.NewRegistry().RegisterFactory(source, jsplugin.Options{Key: "kling"})
|
|
require.NoError(t, err)
|
|
submit := func(t *testing.T, decoder []string, body map[string]any) (*taskplugin.TaskAdaptor, *gin.Context, *relaycommon.RelayInfo, *dto.TaskError) {
|
|
t.Helper()
|
|
value, err := plugin.Engine.CallPath(t.Context(), decoder[0], decoder[1:], map[string]any{"model": "kling-v1", "body": map[string]any{"kind": "json", "value": body}})
|
|
require.NoError(t, err)
|
|
info := &relaycommon.RelayInfo{
|
|
ChannelMeta: &relaycommon.ChannelMeta{ChannelBaseUrl: "https://api-singapore.klingai.com", ApiKey: "ak|sk", UpstreamModelName: "kling-v1"},
|
|
OriginModelName: "kling-v1",
|
|
TaskRelayInfo: &relaycommon.TaskRelayInfo{PublicTaskID: "task_public"},
|
|
}
|
|
adaptor := taskplugin.New(plugin)
|
|
adaptor.Init(info)
|
|
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos", nil)
|
|
c.Set("task_request", alibabaObject(t, value)["requestBody"])
|
|
return adaptor, c, info, adaptor.ValidateRequestAndSetAction(c, info)
|
|
}
|
|
native := []string{"native", "decodeSubmit"}
|
|
openaiVideo := []string{"protocols", "openai_video", "decodeRequest"}
|
|
|
|
for _, tc := range []struct {
|
|
name string
|
|
decoder []string
|
|
body map[string]any
|
|
}{
|
|
{"native duration above 3600", native, map[string]any{"model_name": "kling-v1", "prompt": "a cat", "duration": 1e12}},
|
|
{"native fractional duration", native, map[string]any{"model_name": "kling-v1", "prompt": "a cat", "duration": "5.5"}},
|
|
{"metadata duration above 3600 behind valid seconds", openaiVideo, map[string]any{"model": "kling-v1", "prompt": "a cat", "seconds": "5", "metadata": map[string]any{"duration": 3601}}},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
_, _, _, taskErr := submit(t, tc.decoder, tc.body)
|
|
require.NotNil(t, taskErr)
|
|
assert.Equal(t, http.StatusBadRequest, taskErr.StatusCode)
|
|
assert.ErrorContains(t, taskErr.Error, "at most 3600")
|
|
})
|
|
}
|
|
|
|
adaptor, c, info, taskErr := submit(t, openaiVideo, map[string]any{"model": "kling-v1", "prompt": "a cat", "seconds": "5", "metadata": map[string]any{"duration": "10"}})
|
|
require.Nil(t, taskErr)
|
|
reader, err := adaptor.BuildRequestBody(c, info)
|
|
require.NoError(t, err)
|
|
encoded, err := io.ReadAll(reader)
|
|
require.NoError(t, err)
|
|
var sent map[string]any
|
|
require.NoError(t, common.Unmarshal(encoded, &sent))
|
|
assert.Equal(t, "10", sent["duration"])
|
|
facts, err := adaptor.ExtractUsageFactsValidated(c, info)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, 2.0, facts["units"], "kling-v1 std costs 0.2 credits per second")
|
|
})
|
|
}
|