281 lines
14 KiB
Go
281 lines
14 KiB
Go
package plugins_test
|
|
|
|
import (
|
|
"testing"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/pkg/jsplugin"
|
|
builtinplugins "github.com/QuantumNous/new-api/plugins"
|
|
"github.com/QuantumNous/new-api/relay"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestVertexAIResponsesProtocol(t *testing.T) {
|
|
source, err := builtinplugins.Source("vertex-ai")
|
|
require.NoError(t, err)
|
|
registry := jsplugin.NewRegistry()
|
|
plugin, err := registry.RegisterFactory(source, jsplugin.Options{Key: "vertex-ai"})
|
|
require.NoError(t, err)
|
|
|
|
t.Run("claims every model", func(t *testing.T) {
|
|
for _, model := range plugin.Meta.Models {
|
|
binding, found := registry.Generation().LookupEndpoint("POST", "/v1/responses", model)
|
|
require.True(t, found, model)
|
|
assert.Same(t, plugin, binding.Plugin)
|
|
assert.Equal(t, "openai_responses", binding.Protocol)
|
|
}
|
|
})
|
|
|
|
t.Run("shares models with Gemini without losing either provider", func(t *testing.T) {
|
|
candidates := jsplugin.DefaultRegistry.Generation().LookupEndpointCandidates("POST", "/v1/responses", "veo-3.0-generate-001")
|
|
require.Len(t, candidates, 2)
|
|
assert.Equal(t, "google", candidates[0].Plugin.Meta.Key)
|
|
assert.Equal(t, "vertex-ai", candidates[1].Plugin.Meta.Key)
|
|
|
|
request := map[string]any{"model": "veo-3.0-generate-001", "body": map[string]any{"kind": "json", "value": map[string]any{
|
|
"model": "veo-3.0-generate-001", "input": "waves", "seconds": 8, "size": "1280x720",
|
|
}}}
|
|
first, callErr := candidates[0].Plugin.Engine.CallPath(t.Context(), "protocols", []string{"openai_responses", "decodeRequest"}, request)
|
|
require.NoError(t, callErr)
|
|
second, callErr := candidates[1].Plugin.Engine.CallPath(t.Context(), "protocols", []string{"openai_responses", "decodeRequest"}, request)
|
|
require.NoError(t, callErr)
|
|
assert.Equal(t, decodePluginValue(t, first), decodePluginValue(t, second))
|
|
})
|
|
|
|
t.Run("declares documented usage facts", func(t *testing.T) {
|
|
require.Len(t, plugin.Meta.UsageSchema, 3)
|
|
for _, key := range []string{"seconds", "resolution", "generate_audio"} {
|
|
schema, exists := plugin.Meta.UsageSchema[key]
|
|
require.True(t, exists, key)
|
|
assert.NotEmpty(t, schema.Description, key)
|
|
}
|
|
assert.Equal(t, []string{"720p", "1080p", "4k"}, plugin.Meta.UsageSchema["resolution"].Enum)
|
|
assert.Equal(t, "boolean", plugin.Meta.UsageSchema["generate_audio"].Type)
|
|
})
|
|
|
|
callProtocol := func(t *testing.T, hook string, args ...any) any {
|
|
t.Helper()
|
|
value, callErr := plugin.Engine.CallPath(t.Context(), "protocols", []string{"openai_responses", hook}, args...)
|
|
require.NoError(t, callErr)
|
|
return value
|
|
}
|
|
decodeMap := func(t *testing.T, value any) map[string]any {
|
|
t.Helper()
|
|
encoded, marshalErr := common.Marshal(value)
|
|
require.NoError(t, marshalErr)
|
|
var decoded map[string]any
|
|
require.NoError(t, common.Unmarshal(encoded, &decoded))
|
|
return decoded
|
|
}
|
|
|
|
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": "1920x1080",
|
|
"resolution": "1080P",
|
|
}
|
|
|
|
t.Run("parses Responses input", func(t *testing.T) {
|
|
resolved := decodeMap(t, callProtocol(t, "decodeRequest", map[string]any{"model": requestBody["model"], "body": map[string]any{"kind": "json", "value": requestBody}}))
|
|
assert.Equal(t, "veo-3.1-fast-generate-preview", resolved["model"])
|
|
assert.Equal(t, "image_to_video", resolved["action"])
|
|
assert.Equal(t, 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": "1920x1080",
|
|
"metadata": map[string]any{"resolution": "1080P"},
|
|
}, resolved["requestBody"])
|
|
})
|
|
|
|
t.Run("rejects malformed input", func(t *testing.T) {
|
|
_, callErr := plugin.Engine.CallPath(t.Context(), "protocols", []string{"openai_responses", "decodeRequest"}, map[string]any{
|
|
"model": "veo-3.0-generate-001", "body": map[string]any{"kind": "json", "value": map[string]any{"model": "veo-3.0-generate-001", "input": map[string]any{"text": "bad"}}},
|
|
})
|
|
require.ErrorContains(t, callErr, "input must be a string or array")
|
|
|
|
_, callErr = plugin.Engine.CallPath(t.Context(), "protocols", []string{"openai_responses", "decodeRequest"}, map[string]any{
|
|
"model": "veo-3.0-generate-001", "body": map[string]any{"kind": "json", "value": map[string]any{"model": "veo-3.0-generate-001", "input": []any{map[string]any{"type": "input_image", "image_url": "https://example.com/frame.png"}}}},
|
|
})
|
|
require.ErrorContains(t, callErr, "input image must be a data URL or base64 value")
|
|
})
|
|
|
|
t.Run("extracts schema-declared usage", func(t *testing.T) {
|
|
value, callErr := plugin.Engine.Call(t.Context(), "extractUsage", map[string]any{
|
|
"requestBody": map[string]any{"duration": 8, "size": "1920x1080", "metadata": map[string]any{}},
|
|
"usagePurpose": "facts",
|
|
})
|
|
require.NoError(t, callErr)
|
|
assert.Equal(t, map[string]any{"seconds": int64(8), "resolution": "1080p", "generate_audio": true}, value)
|
|
|
|
value, callErr = plugin.Engine.Call(t.Context(), "extractUsageOnComplete", nil, map[string]any{}, map[string]any{
|
|
"response": map[string]any{"videos": []any{map[string]any{"durationSeconds": 7, "resolution": "4K"}}},
|
|
})
|
|
require.NoError(t, callErr)
|
|
assert.Nil(t, value)
|
|
})
|
|
|
|
protocolContext := map[string]any{
|
|
"requestBody": map[string]any{"model": "veo-3.1-fast-generate-preview", "duration": 8},
|
|
"stream": true,
|
|
}
|
|
successTask := map[string]any{
|
|
"task_id": "task-public", "status": "SUCCESS", "progress": "100%", "created_at": 10, "updated_at": 20,
|
|
"data": map[string]any{"url": "data:video/mp4;base64,MUST_NOT_LEAK"},
|
|
}
|
|
|
|
t.Run("renders stream state transitions", func(t *testing.T) {
|
|
progressValue := callProtocol(t, "renderEvents", protocolContext, map[string]any{"status": "IN_PROGRESS", "progress": "42%"})
|
|
progress, decodeErr := relay.DecodePluginProtocolEventResult(progressValue, relay.DefaultPluginProtocolLimits())
|
|
require.NoError(t, decodeErr)
|
|
require.Len(t, progress.Events, 1)
|
|
require.NotNil(t, progress.Events[0].Progress)
|
|
assert.Equal(t, float64(42), *progress.Events[0].Progress)
|
|
|
|
duplicateValue := callProtocol(t, "renderEvents", protocolContext, map[string]any{"status": "IN_PROGRESS", "progress": "42%"}, map[string]any{"status": "IN_PROGRESS", "progress": float64(42)})
|
|
duplicate, decodeErr := relay.DecodePluginProtocolEventResult(duplicateValue, relay.DefaultPluginProtocolLimits())
|
|
require.NoError(t, decodeErr)
|
|
assert.Empty(t, duplicate.Events)
|
|
assert.False(t, duplicate.Done)
|
|
|
|
failureValue := callProtocol(t, "renderEvents", protocolContext, map[string]any{"status": "FAILURE", "fail_reason": "blocked"})
|
|
failure, decodeErr := relay.DecodePluginProtocolEventResult(failureValue, relay.DefaultPluginProtocolLimits())
|
|
require.NoError(t, decodeErr)
|
|
require.Len(t, failure.Events, 1)
|
|
assert.Equal(t, "error", failure.Events[0].Type)
|
|
assert.True(t, failure.Done)
|
|
|
|
successValue := callProtocol(t, "renderEvents", protocolContext, successTask)
|
|
success, decodeErr := relay.DecodePluginProtocolEventResult(successValue, relay.DefaultPluginProtocolLimits())
|
|
require.NoError(t, decodeErr)
|
|
require.Len(t, success.Events, 1)
|
|
assert.True(t, success.Done)
|
|
var text string
|
|
require.NoError(t, common.Unmarshal(success.Events[0].Data, &text))
|
|
assert.Contains(t, text, "veo-3.1-fast-generate-preview")
|
|
assert.Contains(t, text, "8 seconds")
|
|
assert.Contains(t, text, "/v1/videos")
|
|
assert.NotContains(t, text, "MUST_NOT_LEAK")
|
|
assert.NotContains(t, text, "<video")
|
|
})
|
|
|
|
t.Run("renders a degraded final response without an artifact", func(t *testing.T) {
|
|
value := callProtocol(t, "renderFinal", protocolContext, successTask)
|
|
machine := relay.NewPluginResponsesMachine("task-public", "veo-3.1-fast-generate-preview", 10, relay.DefaultPluginProtocolLimits())
|
|
response, finalErr := machine.FinalResponse(value, "SUCCESS")
|
|
require.NoError(t, finalErr)
|
|
assert.Equal(t, "completed", response["status"])
|
|
output, ok := response["output"].([]any)
|
|
require.True(t, ok)
|
|
require.Len(t, output, 1)
|
|
message, ok := output[0].(map[string]any)
|
|
require.True(t, ok)
|
|
content, ok := message["content"].([]any)
|
|
require.True(t, ok)
|
|
require.Len(t, content, 1)
|
|
part, ok := content[0].(map[string]any)
|
|
require.True(t, ok)
|
|
text, ok := part["text"].(string)
|
|
require.True(t, ok)
|
|
assert.Contains(t, text, "/v1/videos")
|
|
assert.NotContains(t, text, "MUST_NOT_LEAK")
|
|
assert.NotContains(t, text, "<video")
|
|
metadata, ok := response["metadata"].(map[string]string)
|
|
require.True(t, ok)
|
|
assert.Equal(t, "vertex", metadata["vendor"])
|
|
assert.Equal(t, "native_videos", metadata["artifact_mode"])
|
|
})
|
|
}
|
|
|
|
func decodePluginValue(t *testing.T, value any) map[string]any {
|
|
t.Helper()
|
|
encoded, err := common.Marshal(value)
|
|
require.NoError(t, err)
|
|
var decoded map[string]any
|
|
require.NoError(t, common.Unmarshal(encoded, &decoded))
|
|
return decoded
|
|
}
|
|
|
|
// A New API channel reaches Vertex AI through another gateway's /vertex-ai
|
|
// route, which runs the request in the gateway's own project. Operation names
|
|
// carry the caller's path and the gateway task id, never the gateway's project,
|
|
// and fetchPredictOperation on such a name reads that gateway task.
|
|
func TestVertexAIThroughNewAPIGateway(t *testing.T) {
|
|
source, err := builtinplugins.Source("vertex-ai")
|
|
require.NoError(t, err)
|
|
plugin, err := jsplugin.NewRegistry().RegisterFactory(source, jsplugin.Options{Key: "vertex-ai"})
|
|
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.0-generate-001"
|
|
const route = "https://gateway.example/vertex-ai/v1/projects/-/locations/global/publishers/google/models/" + veo
|
|
gateway := map[string]any{"kind": "new_api"}
|
|
gatewayAuth := map[string]any{"authHeader": "Bearer sk-gateway"}
|
|
|
|
downstream := map[string]any{
|
|
"requestBody": map[string]any{"model": veo, "prompt": "p", "duration": 6, "images": []any{"data:image/png;base64,aGVsbG8="},
|
|
"metadata": map[string]any{"resolution": "1080p", "generateAudio": false}},
|
|
"model": veo, "upstreamModel": veo, "baseUrl": "https://gateway.example", "apiKey": "sk-gateway",
|
|
"authHeader": "Bearer sk-gateway", "auth": gatewayAuth, "upstream": gateway,
|
|
}
|
|
sent := call("buildSubmitRequest", nil, downstream)
|
|
assert.Equal(t, route+":predictLongRunning", sent["url"])
|
|
assert.Equal(t, map[string]any{"Accept": "application/json", "Content-Type": "application/json", "Authorization": "Bearer sk-gateway"}, sent["headers"])
|
|
|
|
submitParams := map[string]any{"project": "-", "location": "global", "model": veo + ":predictLongRunning"}
|
|
intent := call("native", []string{"modelMethod"}, map[string]any{"params": submitParams, "body": map[string]any{"kind": "json", "value": sent["body"]}})
|
|
assert.Equal(t, "submit", intent["kind"])
|
|
assert.Equal(t, sent["action"], intent["action"])
|
|
upstream := map[string]any{"requestBody": intent["requestBody"], "model": veo, "upstreamModel": veo, "baseUrl": "https://aiplatform.example",
|
|
"authHeader": "Bearer ya29.vendor", "auth": map[string]any{"authHeader": "Bearer ya29.vendor", "projectId": "gateway-project"},
|
|
"upstream": map[string]any{"kind": "vendor"}}
|
|
forwarded := call("buildSubmitRequest", nil, upstream)
|
|
assert.Equal(t, "https://aiplatform.example/v1/projects/gateway-project/locations/global/publishers/google/models/"+veo+":predictLongRunning", forwarded["url"])
|
|
assert.Equal(t, "gateway-project", forwarded["headers"].(map[string]any)["x-goog-user-project"])
|
|
assert.Equal(t, sent["body"], forwarded["body"])
|
|
assert.Equal(t, call("extractUsage", nil, downstream), call("extractUsage", nil, upstream))
|
|
|
|
vendorOperation := "projects/gateway-project/locations/us-central1/publishers/google/models/" + veo + "/operations/op-private"
|
|
created := call("native", []string{"operation"}, map[string]any{"params": submitParams},
|
|
map[string]any{"task_id": "task_pub", "status": "SUBMITTED", "data": map[string]any{"name": vendorOperation}})
|
|
assert.Equal(t, map[string]any{"name": "projects/-/locations/global/publishers/google/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",
|
|
"authHeader": "Bearer sk-gateway", "auth": gatewayAuth, "upstream": gateway})
|
|
assert.Equal(t, route+":fetchPredictOperation", query["url"])
|
|
fetchParams := map[string]any{"project": "-", "location": "global", "model": veo + ":fetchPredictOperation"}
|
|
assert.Equal(t, map[string]any{"kind": "query", "taskIds": []any{"task_pub"}},
|
|
call("native", []string{"modelMethod"}, map[string]any{"params": fetchParams, "body": map[string]any{"kind": "json", "value": query["body"]}}))
|
|
|
|
for _, tc := range []struct {
|
|
name string
|
|
task map[string]any
|
|
wantStatus string
|
|
wantReason string
|
|
}{
|
|
{"running", map[string]any{"status": "IN_PROGRESS", "data": map[string]any{"name": vendorOperation}}, "IN_PROGRESS", ""},
|
|
{"succeeded", map[string]any{"status": "SUCCESS", "data": map[string]any{"name": vendorOperation, "done": true,
|
|
"response": map[string]any{"videos": []any{map[string]any{"mimeType": "video/mp4"}}}}}, "SUCCESS", ""},
|
|
{"failed by the gateway poller", map[string]any{"status": "FAILURE", "fail_reason": "task timed out", "data": map[string]any{"name": vendorOperation}}, "FAILURE", "task timed out"},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
tc.task["task_id"] = "task_pub"
|
|
rendered := call("native", []string{"operation"}, map[string]any{"params": fetchParams}, []any{tc.task})
|
|
assert.Equal(t, "projects/-/locations/global/publishers/google/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"]))
|
|
})
|
|
}
|
|
}
|