169 lines
8.3 KiB
Go
169 lines
8.3 KiB
Go
package plugins_test
|
|
|
|
import (
|
|
"net/http"
|
|
"testing"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/constant"
|
|
"github.com/QuantumNous/new-api/model"
|
|
"github.com/QuantumNous/new-api/pkg/jsplugin"
|
|
builtinplugins "github.com/QuantumNous/new-api/plugins"
|
|
"github.com/QuantumNous/new-api/relay/channel"
|
|
taskplugin "github.com/QuantumNous/new-api/relay/channel/task/jsplugin"
|
|
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestGoogleResponsesProtocol(t *testing.T) {
|
|
testVideoResponsesProtocol(t, videoResponsesTestCase{
|
|
pluginKey: "google",
|
|
model: "veo-3.1-fast-generate-preview",
|
|
requestBody: map[string]any{
|
|
"model": "veo-3.1-fast-generate-preview",
|
|
"input": []any{map[string]any{"role": "user", "content": []any{
|
|
map[string]any{"type": "input_text", "text": "animate this frame"},
|
|
map[string]any{"type": "input_image", "image_url": "data:image/png;base64,aGVsbG8="},
|
|
}}},
|
|
"seconds": 8,
|
|
"size": "1280x720",
|
|
},
|
|
wantAction: "image_to_video",
|
|
wantRequest: map[string]any{
|
|
"model": "veo-3.1-fast-generate-preview",
|
|
"prompt": "animate this frame",
|
|
"images": []any{"data:image/png;base64,aGVsbG8="},
|
|
"duration": float64(8),
|
|
"size": "1280x720",
|
|
"metadata": map[string]any{},
|
|
},
|
|
wantUsageKeys: []string{"resolution", "seconds"},
|
|
wantVendorName: "gemini",
|
|
})
|
|
}
|
|
|
|
// A New API channel reaches Veo through another gateway's /google routes. The
|
|
// gateway rebuilds the request this driver sent, names the operation after its
|
|
// own task so it is polled there, and serves the video file only its Gemini
|
|
// key can read.
|
|
func TestGoogleThroughNewAPIGateway(t *testing.T) {
|
|
source, err := builtinplugins.Source("google")
|
|
require.NoError(t, err)
|
|
plugin, err := jsplugin.NewRegistry().RegisterFactory(source, jsplugin.Options{Key: "google"})
|
|
require.NoError(t, err)
|
|
call := func(export string, members []string, args ...any) map[string]any {
|
|
t.Helper()
|
|
value, err := plugin.Engine.CallPath(t.Context(), export, members, args...)
|
|
require.NoError(t, err)
|
|
return decodePluginValue(t, value)
|
|
}
|
|
const veo = "veo-3.1-fast-generate-preview"
|
|
gateway := map[string]any{"kind": "new_api"}
|
|
submitParams := map[string]any{"model": veo + ":predictLongRunning"}
|
|
|
|
for _, tc := range []struct {
|
|
name string
|
|
requestBody map[string]any
|
|
}{
|
|
{"text with parameters", map[string]any{"model": veo, "prompt": "p", "duration": 6, "metadata": map[string]any{"resolution": "1080p", "aspectRatio": "9:16", "negativePrompt": "rain"}}},
|
|
{"first frame", map[string]any{"model": veo, "prompt": "p", "images": []any{"data:image/png;base64,aGVsbG8="}, "size": "1280x720"}},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
downstream := map[string]any{"requestBody": tc.requestBody, "model": veo, "upstreamModel": veo,
|
|
"baseUrl": "https://gateway.example", "apiKey": "sk-gateway", "upstream": gateway}
|
|
sent := call("buildSubmitRequest", nil, downstream)
|
|
assert.Equal(t, "https://gateway.example/google/v1beta/models/"+veo+":predictLongRunning", sent["url"])
|
|
assert.Equal(t, map[string]any{"Accept": "application/json", "Content-Type": "application/json", "Authorization": "Bearer sk-gateway"}, sent["headers"])
|
|
|
|
intent := call("native", []string{"predictLongRunning"}, map[string]any{"params": submitParams, "body": map[string]any{"kind": "json", "value": sent["body"]}})
|
|
assert.Equal(t, veo, intent["model"])
|
|
assert.Equal(t, sent["action"], intent["action"])
|
|
upstream := map[string]any{"requestBody": intent["requestBody"], "model": veo, "upstreamModel": veo,
|
|
"baseUrl": "https://generativelanguage.example", "apiKey": "vendor-key", "upstream": map[string]any{"kind": "vendor"}}
|
|
forwarded := call("buildSubmitRequest", nil, upstream)
|
|
assert.Equal(t, "https://generativelanguage.example/v1beta/models/"+veo+":predictLongRunning", forwarded["url"])
|
|
assert.Equal(t, "vendor-key", forwarded["headers"].(map[string]any)["x-goog-api-key"])
|
|
assert.Equal(t, sent["body"], forwarded["body"])
|
|
assert.Equal(t, call("extractUsage", nil, downstream), call("extractUsage", nil, upstream))
|
|
})
|
|
}
|
|
|
|
created := call("native", []string{"operation"}, map[string]any{"params": submitParams},
|
|
map[string]any{"task_id": "task_pub", "status": "SUBMITTED", "data": map[string]any{"name": "models/" + veo + "/operations/op-private"}})
|
|
assert.Equal(t, map[string]any{"name": "models/" + veo + "/operations/task_pub"}, created)
|
|
upstreamTaskID := call("parseSubmitResponse", nil, map[string]any{}, map[string]any{"body": created})["taskId"]
|
|
query := call("buildQueryRequest", nil, map[string]any{"taskId": upstreamTaskID, "baseUrl": "https://gateway.example", "apiKey": "sk-gateway", "upstream": gateway})
|
|
assert.Equal(t, "https://gateway.example/google/v1beta/models/"+veo+"/operations/task_pub", query["url"])
|
|
|
|
adaptor := taskplugin.New(plugin)
|
|
adaptor.Init(&relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{
|
|
ChannelType: constant.ChannelTypeNewAPI, ApiKey: "sk-gateway", ChannelBaseUrl: "https://gateway.example",
|
|
}})
|
|
queryParams := map[string]any{"model": veo, "task_id": "task_pub"}
|
|
for _, tc := range []struct {
|
|
name string
|
|
task map[string]any
|
|
wantStatus string
|
|
wantReason string
|
|
}{
|
|
{
|
|
name: "running", wantStatus: "IN_PROGRESS",
|
|
task: map[string]any{"status": "IN_PROGRESS", "data": map[string]any{"name": "models/" + veo + "/operations/op-private"}},
|
|
},
|
|
{
|
|
name: "succeeded", wantStatus: "SUCCESS",
|
|
task: map[string]any{"status": "SUCCESS", "data": map[string]any{"name": "models/" + veo + "/operations/op-private", "done": true,
|
|
"response": map[string]any{"generateVideoResponse": map[string]any{"generatedVideos": []any{
|
|
map[string]any{"video": map[string]any{"uri": "https://generativelanguage.example/v1beta/files/f1:download?alt=media"}},
|
|
}}}}},
|
|
},
|
|
{
|
|
name: "failed by the gateway poller", wantStatus: "FAILURE", wantReason: "task timed out",
|
|
task: map[string]any{"status": "FAILURE", "fail_reason": "task timed out", "data": map[string]any{"name": "models/" + veo + "/operations/op-private"}},
|
|
},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
tc.task["task_id"] = "task_pub"
|
|
rendered := call("native", []string{"operation"}, map[string]any{"params": queryParams}, tc.task)
|
|
assert.Equal(t, "models/"+veo+"/operations/task_pub", rendered["name"])
|
|
result := call("parseTaskResult", nil, map[string]any{}, rendered)
|
|
assert.Equal(t, tc.wantStatus, result["status"])
|
|
assert.Equal(t, tc.wantReason, common.Interface2String(result["reason"]))
|
|
if tc.wantStatus != "SUCCESS" {
|
|
return
|
|
}
|
|
data, err := common.Marshal(rendered)
|
|
require.NoError(t, err)
|
|
task := &model.Task{TaskID: "task_down", Status: model.TaskStatusSuccess, Data: data, PrivateData: model.TaskPrivateData{UpstreamTaskID: upstreamTaskID.(string)}}
|
|
artifacts, err := adaptor.ListArtifacts(task)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, []channel.TaskArtifact{{Key: "video", Type: "video"}}, artifacts)
|
|
descriptor, err := adaptor.BuildContentRequest(t.Context(), task, "video", channel.TaskArtifactClientRequest{Method: http.MethodGet})
|
|
require.NoError(t, err)
|
|
assert.Equal(t, &channel.TaskContentRequest{
|
|
URL: "https://gateway.example/v1/tasks/task_pub/artifacts/video/content",
|
|
Method: http.MethodGet,
|
|
Headers: map[string]string{"Authorization": "Bearer sk-gateway"},
|
|
}, descriptor)
|
|
})
|
|
}
|
|
|
|
t.Run("native bodies are checked before billing", func(t *testing.T) {
|
|
for _, tc := range []struct {
|
|
model string
|
|
body map[string]any
|
|
wantErr string
|
|
}{
|
|
{veo + ":generateVideos", map[string]any{"instances": []any{map[string]any{"prompt": "p"}}}, "models/{model}:predictLongRunning"},
|
|
{veo + ":predictLongRunning", map[string]any{"instances": []any{map[string]any{"prompt": "p", "lastFrame": map[string]any{}}}}, "instances[0].lastFrame is not supported"},
|
|
{veo + ":predictLongRunning", map[string]any{"instances": []any{map[string]any{"prompt": "p"}}, "parameters": map[string]any{"durationSeconds": 5}}, "durationSeconds must be one of 4, 6, or 8"},
|
|
} {
|
|
_, err := plugin.Engine.CallPath(t.Context(), "native", []string{"predictLongRunning"}, map[string]any{
|
|
"params": map[string]any{"model": tc.model}, "body": map[string]any{"kind": "json", "value": tc.body},
|
|
})
|
|
require.ErrorContains(t, err, tc.wantErr)
|
|
}
|
|
})
|
|
}
|