mirror of
https://github.com/alibaba/higress.git
synced 2026-05-23 12:17:30 +08:00
Add Kling provider support (#3742)
Signed-off-by: wydream <yaodiwu618@gmail.com>
This commit is contained in:
woody
GitHub
418
plugins/wasm-go/extensions/ai-proxy/test/kling.go
Normal file
418
plugins/wasm-go/extensions/ai-proxy/test/kling.go
Normal file
@@ -0,0 +1,418 @@
|
||||
package test
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/higress-group/proxy-wasm-go-sdk/proxywasm/types"
|
||||
"github.com/higress-group/wasm-go/pkg/test"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
var klingOfficialConfig = func() json.RawMessage {
|
||||
data, _ := json.Marshal(map[string]interface{}{
|
||||
"provider": map[string]interface{}{
|
||||
"type": "kling",
|
||||
"klingAccessKey": "kling-ak-test",
|
||||
"klingSecretKey": "kling-sk-test",
|
||||
"klingTokenRefreshAhead": 60,
|
||||
"modelMapping": map[string]string{
|
||||
"client-video": "kling-v2-1",
|
||||
},
|
||||
},
|
||||
})
|
||||
return data
|
||||
}()
|
||||
|
||||
var klingGatewayConfig = func() json.RawMessage {
|
||||
data, _ := json.Marshal(map[string]interface{}{
|
||||
"provider": map[string]interface{}{
|
||||
"type": "kling",
|
||||
"apiTokens": []string{"gateway-token"},
|
||||
"providerDomain": "api.302.ai",
|
||||
"providerBasePath": "/klingai",
|
||||
"modelMapping": map[string]string{
|
||||
"client-video": "kling-v2-1",
|
||||
},
|
||||
},
|
||||
})
|
||||
return data
|
||||
}()
|
||||
|
||||
var klingGatewayCustomImageRetrieveConfig = func() json.RawMessage {
|
||||
data, _ := json.Marshal(map[string]interface{}{
|
||||
"provider": map[string]interface{}{
|
||||
"type": "kling",
|
||||
"apiTokens": []string{"gateway-token"},
|
||||
"providerDomain": "api.302.ai",
|
||||
"providerBasePath": "/klingai",
|
||||
"modelMapping": map[string]string{
|
||||
"client-video": "kling-v2-1",
|
||||
},
|
||||
"capabilities": map[string]string{
|
||||
"openai/v1/videos": "/gateway/text2video?mode=text",
|
||||
"kling/v1/image2video": "/gateway/image2video?mode=image",
|
||||
"kling/v1/retrieveimagevideo": "/gateway/image-tasks/{video_id}?version=1",
|
||||
},
|
||||
},
|
||||
})
|
||||
return data
|
||||
}()
|
||||
|
||||
var klingOriginalConfig = func() json.RawMessage {
|
||||
data, _ := json.Marshal(map[string]interface{}{
|
||||
"provider": map[string]interface{}{
|
||||
"type": "kling",
|
||||
"apiTokens": []string{"gateway-token"},
|
||||
"providerDomain": "api.302.ai",
|
||||
"providerBasePath": "/klingai",
|
||||
"protocol": "original",
|
||||
},
|
||||
})
|
||||
return data
|
||||
}()
|
||||
|
||||
func RunKlingParseConfigTests(t *testing.T) {
|
||||
test.RunGoTest(t, func(t *testing.T) {
|
||||
t.Run("kling official config", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(klingOfficialConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
config, err := host.GetMatchConfig()
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, config)
|
||||
})
|
||||
|
||||
t.Run("kling gateway config", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(klingGatewayConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
config, err := host.GetMatchConfig()
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, config)
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
func RunKlingOnHttpRequestHeadersTests(t *testing.T) {
|
||||
test.RunTest(t, func(t *testing.T) {
|
||||
t.Run("official mode sets jwt bearer and default host", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(klingOfficialConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
action := host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/v1/videos"},
|
||||
{":method", "POST"},
|
||||
{"Content-Type", "application/json"},
|
||||
})
|
||||
require.Equal(t, types.HeaderStopIteration, action)
|
||||
|
||||
requestHeaders := host.GetRequestHeaders()
|
||||
hostValue, hasHost := test.GetHeaderValue(requestHeaders, ":authority")
|
||||
require.True(t, hasHost)
|
||||
require.Equal(t, "api-singapore.klingai.com", hostValue)
|
||||
|
||||
authValue, hasAuth := test.GetHeaderValue(requestHeaders, "Authorization")
|
||||
require.True(t, hasAuth)
|
||||
require.True(t, strings.HasPrefix(authValue, "Bearer "))
|
||||
require.Len(t, strings.Split(strings.TrimPrefix(authValue, "Bearer "), "."), 3)
|
||||
|
||||
pathValue, hasPath := test.GetHeaderValue(requestHeaders, ":path")
|
||||
require.True(t, hasPath)
|
||||
require.Equal(t, "/v1/videos/text2video", pathValue)
|
||||
})
|
||||
|
||||
t.Run("providerDomain and providerBasePath apply to gateway mode", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(klingGatewayConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
action := host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/v1/videos"},
|
||||
{":method", "POST"},
|
||||
{"Content-Type", "application/json"},
|
||||
})
|
||||
require.Equal(t, types.HeaderStopIteration, action)
|
||||
|
||||
requestHeaders := host.GetRequestHeaders()
|
||||
hostValue, hasHost := test.GetHeaderValue(requestHeaders, ":authority")
|
||||
require.True(t, hasHost)
|
||||
require.Equal(t, "api.302.ai", hostValue)
|
||||
|
||||
authValue, hasAuth := test.GetHeaderValue(requestHeaders, "Authorization")
|
||||
require.True(t, hasAuth)
|
||||
require.Equal(t, "Bearer gateway-token", authValue)
|
||||
|
||||
pathValue, hasPath := test.GetHeaderValue(requestHeaders, ":path")
|
||||
require.True(t, hasPath)
|
||||
require.Equal(t, "/klingai/v1/videos/text2video", pathValue)
|
||||
})
|
||||
|
||||
t.Run("retrieve video query path is mapped under providerBasePath", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(klingGatewayConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
action := host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/v1/videos/task-123?with_status=true"},
|
||||
{":method", "GET"},
|
||||
})
|
||||
require.True(t, action == types.ActionContinue || action == types.HeaderStopIteration)
|
||||
|
||||
requestHeaders := host.GetRequestHeaders()
|
||||
pathValue, hasPath := test.GetHeaderValue(requestHeaders, ":path")
|
||||
require.True(t, hasPath)
|
||||
require.Equal(t, "/klingai/v1/videos/text2video/task-123?with_status=true", pathValue)
|
||||
})
|
||||
|
||||
t.Run("prefixed image task query path is mapped to image endpoint", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(klingGatewayConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
action := host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/v1/videos/kling-i2v-task-123?with_status=true"},
|
||||
{":method", "GET"},
|
||||
})
|
||||
require.True(t, action == types.ActionContinue || action == types.HeaderStopIteration)
|
||||
|
||||
requestHeaders := host.GetRequestHeaders()
|
||||
pathValue, hasPath := test.GetHeaderValue(requestHeaders, ":path")
|
||||
require.True(t, hasPath)
|
||||
require.Equal(t, "/klingai/v1/videos/image2video/task-123?with_status=true", pathValue)
|
||||
})
|
||||
|
||||
t.Run("prefixed image task query strips task type hint", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(klingGatewayConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
action := host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/v1/videos/kling-i2v-task-123?kling_task_type=image2video&with_status=true"},
|
||||
{":method", "GET"},
|
||||
})
|
||||
require.True(t, action == types.ActionContinue || action == types.HeaderStopIteration)
|
||||
|
||||
requestHeaders := host.GetRequestHeaders()
|
||||
pathValue, hasPath := test.GetHeaderValue(requestHeaders, ":path")
|
||||
require.True(t, hasPath)
|
||||
require.Equal(t, "/klingai/v1/videos/image2video/task-123?with_status=true", pathValue)
|
||||
})
|
||||
|
||||
t.Run("raw image task query path uses explicit task type hint", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(klingGatewayConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
action := host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/v1/videos/raw-task-123?kling_task_type=image2video&with_status=true"},
|
||||
{":method", "GET"},
|
||||
})
|
||||
require.True(t, action == types.ActionContinue || action == types.HeaderStopIteration)
|
||||
|
||||
requestHeaders := host.GetRequestHeaders()
|
||||
pathValue, hasPath := test.GetHeaderValue(requestHeaders, ":path")
|
||||
require.True(t, hasPath)
|
||||
require.Equal(t, "/klingai/v1/videos/image2video/raw-task-123?with_status=true", pathValue)
|
||||
})
|
||||
|
||||
t.Run("raw retrieve strips unknown task type hint before fallback mapping", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(klingGatewayConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
action := host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/v1/videos/raw-task-123?kling_task_type=bad&with_status=true"},
|
||||
{":method", "GET"},
|
||||
})
|
||||
require.True(t, action == types.ActionContinue || action == types.HeaderStopIteration)
|
||||
|
||||
requestHeaders := host.GetRequestHeaders()
|
||||
pathValue, hasPath := test.GetHeaderValue(requestHeaders, ":path")
|
||||
require.True(t, hasPath)
|
||||
require.Equal(t, "/klingai/v1/videos/text2video/raw-task-123?with_status=true", pathValue)
|
||||
})
|
||||
|
||||
t.Run("image retrieve path uses configured image capability", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(klingGatewayCustomImageRetrieveConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
action := host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/v1/videos/kling-i2v-task-123?with_status=true"},
|
||||
{":method", "GET"},
|
||||
})
|
||||
require.True(t, action == types.ActionContinue || action == types.HeaderStopIteration)
|
||||
|
||||
requestHeaders := host.GetRequestHeaders()
|
||||
pathValue, hasPath := test.GetHeaderValue(requestHeaders, ":path")
|
||||
require.True(t, hasPath)
|
||||
require.Equal(t, "/klingai/gateway/image-tasks/task-123?version=1&with_status=true", pathValue)
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
func RunKlingOnHttpRequestBodyTests(t *testing.T) {
|
||||
test.RunTest(t, func(t *testing.T) {
|
||||
t.Run("text to video keeps text endpoint and maps model_name", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(klingGatewayConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/v1/videos?gateway_param=1"},
|
||||
{":method", "POST"},
|
||||
{"Content-Type", "application/json"},
|
||||
})
|
||||
|
||||
action := host.CallOnHttpRequestBody([]byte(`{"model":"client-video","prompt":"sunrise","duration":"5"}`))
|
||||
require.Equal(t, types.ActionContinue, action)
|
||||
|
||||
requestHeaders := host.GetRequestHeaders()
|
||||
pathValue, hasPath := test.GetHeaderValue(requestHeaders, ":path")
|
||||
require.True(t, hasPath)
|
||||
require.Equal(t, "/klingai/v1/videos/text2video?gateway_param=1", pathValue)
|
||||
|
||||
processedBody := host.GetRequestBody()
|
||||
require.Equal(t, "kling-v2-1", gjson.GetBytes(processedBody, "model_name").String())
|
||||
require.False(t, gjson.GetBytes(processedBody, "model").Exists())
|
||||
require.Equal(t, "sunrise", gjson.GetBytes(processedBody, "prompt").String())
|
||||
})
|
||||
|
||||
t.Run("image to video switches endpoint after body inspection", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(klingGatewayConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/v1/videos?gateway_param=1"},
|
||||
{":method", "POST"},
|
||||
{"Content-Type", "application/json"},
|
||||
})
|
||||
|
||||
action := host.CallOnHttpRequestBody([]byte(`{"model":"client-video","prompt":"animate","image":"https://example.com/a.png"}`))
|
||||
require.Equal(t, types.ActionContinue, action)
|
||||
|
||||
requestHeaders := host.GetRequestHeaders()
|
||||
pathValue, hasPath := test.GetHeaderValue(requestHeaders, ":path")
|
||||
require.True(t, hasPath)
|
||||
require.Equal(t, "/klingai/v1/videos/image2video?gateway_param=1", pathValue)
|
||||
|
||||
processedBody := host.GetRequestBody()
|
||||
require.Equal(t, "kling-v2-1", gjson.GetBytes(processedBody, "model_name").String())
|
||||
require.Equal(t, "https://example.com/a.png", gjson.GetBytes(processedBody, "image").String())
|
||||
})
|
||||
|
||||
t.Run("image to video uses configured image capability and merges query", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(klingGatewayCustomImageRetrieveConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/v1/videos?gateway_param=1"},
|
||||
{":method", "POST"},
|
||||
{"Content-Type", "application/json"},
|
||||
})
|
||||
|
||||
action := host.CallOnHttpRequestBody([]byte(`{"model":"client-video","prompt":"animate","image":"https://example.com/a.png"}`))
|
||||
require.Equal(t, types.ActionContinue, action)
|
||||
|
||||
requestHeaders := host.GetRequestHeaders()
|
||||
pathValue, hasPath := test.GetHeaderValue(requestHeaders, ":path")
|
||||
require.True(t, hasPath)
|
||||
require.Equal(t, "/klingai/gateway/image2video?mode=image&gateway_param=1", pathValue)
|
||||
|
||||
processedBody := host.GetRequestBody()
|
||||
require.Equal(t, "kling-v2-1", gjson.GetBytes(processedBody, "model_name").String())
|
||||
})
|
||||
|
||||
t.Run("original protocol does not expose request body handler", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(klingOriginalConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
action := host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/v1/videos/image2video"},
|
||||
{":method", "POST"},
|
||||
{"Content-Type", "application/json"},
|
||||
{"Content-Length", "64"},
|
||||
})
|
||||
require.Equal(t, types.ActionContinue, action)
|
||||
|
||||
requestHeaders := host.GetRequestHeaders()
|
||||
pathValue, hasPath := test.GetHeaderValue(requestHeaders, ":path")
|
||||
require.True(t, hasPath)
|
||||
require.Equal(t, "/klingai/v1/videos/image2video", pathValue)
|
||||
contentLengthValue, hasContentLength := test.GetHeaderValue(requestHeaders, "Content-Length")
|
||||
require.True(t, hasContentLength)
|
||||
require.Equal(t, "64", contentLengthValue)
|
||||
})
|
||||
|
||||
t.Run("original protocol recognizes native retrieve video path", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(klingOriginalConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
action := host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/v1/videos/text2video/task-123"},
|
||||
{":method", "GET"},
|
||||
})
|
||||
require.True(t, action == types.ActionContinue || action == types.HeaderStopIteration)
|
||||
|
||||
requestHeaders := host.GetRequestHeaders()
|
||||
pathValue, hasPath := test.GetHeaderValue(requestHeaders, ":path")
|
||||
require.True(t, hasPath)
|
||||
require.Equal(t, "/klingai/v1/videos/text2video/task-123", pathValue)
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
func RunKlingOnHttpResponseBodyTests(t *testing.T) {
|
||||
test.RunTest(t, func(t *testing.T) {
|
||||
t.Run("image creation response prefixes task id", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(klingGatewayConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/v1/videos"},
|
||||
{":method", "POST"},
|
||||
{"Content-Type", "application/json"},
|
||||
})
|
||||
action := host.CallOnHttpRequestBody([]byte(`{"model":"client-video","prompt":"animate","image":"https://example.com/a.png"}`))
|
||||
require.Equal(t, types.ActionContinue, action)
|
||||
|
||||
require.NoError(t, host.SetProperty([]string{"response", "code_details"}, []byte("via_upstream")))
|
||||
action = host.CallOnHttpResponseHeaders([][2]string{
|
||||
{":status", "200"},
|
||||
{"Content-Type", "application/json"},
|
||||
})
|
||||
require.Equal(t, types.ActionContinue, action)
|
||||
action = host.CallOnHttpResponseBody([]byte(`{"id":"root-task","data":{"task_id":"task-123"}}`))
|
||||
require.Equal(t, types.ActionContinue, action)
|
||||
|
||||
processedBody := host.GetResponseBody()
|
||||
require.Equal(t, "root-task", gjson.GetBytes(processedBody, "id").String())
|
||||
require.Equal(t, "kling-i2v-task-123", gjson.GetBytes(processedBody, "data.task_id").String())
|
||||
})
|
||||
})
|
||||
}
|
||||
@@ -83,6 +83,15 @@ func RunMapRequestPathByCapabilityTests(t *testing.T) {
|
||||
},
|
||||
expected: "/v1/videos/video-xyz",
|
||||
},
|
||||
{
|
||||
name: "video placeholder is replaced in nested provider path",
|
||||
apiName: "openai/v1/retrievevideo",
|
||||
origin: "/openai/v1/videos/video-xyz",
|
||||
mapping: map[string]string{
|
||||
"openai/v1/retrievevideo": "/v1/videos/text2video/{video_id}",
|
||||
},
|
||||
expected: "/v1/videos/text2video/video-xyz",
|
||||
},
|
||||
{
|
||||
name: "video content placeholder with query",
|
||||
apiName: "openai/v1/retrievevideocontent",
|
||||
|
||||
Reference in New Issue
Block a user