mirror of
https://github.com/QuantumNous/new-api.git
synced 2026-09-11 06:30:21 +00:00
204 lines
9.0 KiB
Go
204 lines
9.0 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
|
|
}
|