plugins/vertex_ai_responses_test.go
Rama main · réplica de lectura del repositorio original.
package plugins_test
import (
"encoding/base64"
"maps"
"slices"
"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"]))
})
}
}
// An inline video does not fit the stored operation, which keeps only its
// media type: the artifact is downloaded by fetching the operation again and
// reading the bytes at contentPath. A video written to Cloud Storage stays at
// its gcsUri.
func TestVertexAIDeliversInlineVideoByFetchingTheOperationAgain(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)
call := func(hook string, args ...any) any {
t.Helper()
value, err := plugin.Engine.Call(t.Context(), hook, args...)
require.NoError(t, err)
return value
}
const veo = "veo-3.0-generate-001"
operation := "projects/p1/locations/us-central1/publishers/google/models/" + veo + "/operations/op-1"
inline := map[string]any{"name": operation, "done": true, "response": map[string]any{"videos": []any{map[string]any{"bytesBase64Encoded": "AAAA", "mimeType": "video/webm"}}}}
stored := map[string]any{"name": operation, "done": true, "response": map[string]any{"videos": []any{map[string]any{"mimeType": "video/webm"}}}}
stored2 := map[string]any{"name": operation, "done": true, "response": map[string]any{"videos": []any{map[string]any{"gcsUri": "gs://bucket/video.mp4", "mimeType": "video/mp4"}}}}
assert.Equal(t, map[string]any{"status": "SUCCESS", "progress": "100%"}, decodePluginValue(t, call("parseTaskResult", map[string]any{}, inline)))
assert.Equal(t, map[string]any{"status": "SUCCESS", "progress": "100%", "remoteUrl": "gs://bucket/video.mp4"}, decodePluginValue(t, call("parseTaskResult", map[string]any{}, stored2)))
assert.Equal(t, []any{map[string]any{"key": "video", "type": "video", "mimeType": "video/webm"}}, call("listArtifacts", map[string]any{"status": "SUCCESS", "data": stored}))
assert.Empty(t, call("listArtifacts", map[string]any{"status": "SUCCESS", "data": stored2}))
assert.Empty(t, call("listArtifacts", map[string]any{"status": "IN_PROGRESS", "data": map[string]any{"name": operation}}))
descriptor := decodePluginValue(t, call("buildContentRequest", map[string]any{
"artifactKey": "video", "upstreamTaskId": base64.RawURLEncoding.EncodeToString([]byte(operation)), "data": stored, "baseUrl": "",
"authHeader": "Bearer ya29.vendor", "auth": map[string]any{"authHeader": "Bearer ya29.vendor", "projectId": "p1"},
"upstream": map[string]any{"kind": "vendor"}, "clientRequest": map[string]any{"method": "HEAD", "headers": map[string]any{}},
}))
assert.Equal(t, map[string]any{
"url": "https://us-central1-aiplatform.googleapis.com/v1/projects/p1/locations/us-central1/publishers/google/models/" + veo + ":fetchPredictOperation",
"method": "POST",
"headers": map[string]any{"Accept": "application/json", "Content-Type": "application/json", "Authorization": "Bearer ya29.vendor", "x-goog-user-project": "p1"},
"body": map[string]any{"operationName": operation},
"responseType": "json",
"contentPath": []any{"response", "videos", float64(0), "bytesBase64Encoded"},
"mimeType": "video/webm",
}, descriptor)
// v1beta1, google-genai's default version, is the same route.
v1 := registry.Generation().LookupDeclaredRoutes("POST", "/v1/projects/:p/locations/:l/publishers/google/models/:m")
v1beta1 := registry.Generation().LookupDeclaredRoutes("POST", "/v1beta1/projects/:p/locations/:l/publishers/google/models/:m")
require.Len(t, v1, 1)
require.Len(t, v1beta1, 1)
assert.Equal(t, v1[0].Route.Decode, v1beta1[0].Route.Decode)
assert.Equal(t, v1[0].Route.Render, v1beta1[0].Route.Render)
}
// Gemini Omni submits a background interaction on the global endpoint.
func TestVertexAIOmniBuildSubmitRequest(t *testing.T) {
plugin := loadVertexAIPlugin(t)
const omni = "gemini-omni-1.1-flash-preview"
call := func(args map[string]any) (map[string]any, error) {
t.Helper()
value, err := plugin.Engine.Call(t.Context(), "buildSubmitRequest", args)
if err != nil {
return nil, err
}
return decodePluginValue(t, value), nil
}
descriptor, err := call(map[string]any{
"model": "gemini-omni-1.1-flash", "upstreamModel": omni,
"authHeader": "Bearer token", "auth": map[string]any{"authHeader": "Bearer token", "projectId": "my-project"},
"requestBody": map[string]any{"prompt": "a cat on the beach", "duration": 6, "size": "720x1280", "metadata": map[string]any{"resolution": "360p"}},
})
require.NoError(t, err)
assert.Equal(t, "https://aiplatform.googleapis.com/v1beta1/projects/my-project/locations/global/interactions", descriptor["url"])
assert.Equal(t, "POST", descriptor["method"])
assert.Equal(t, "text_to_video", descriptor["action"])
assert.Equal(t, map[string]any{"Accept": "application/json", "Content-Type": "application/json", "Authorization": "Bearer token", "x-goog-user-project": "my-project"}, descriptor["headers"])
assert.Equal(t, map[string]any{
"model": omni,
"input": []any{map[string]any{"type": "text", "text": "a cat on the beach"}},
"response_format": []any{map[string]any{"type": "video", "aspect_ratio": "9:16", "resolution": "360p", "duration": "6s"}},
"generation_config": map[string]any{"video_config": map[string]any{"task": "text_to_video"}},
"background": true,
"store": true,
}, descriptor["body"])
descriptor, err = call(map[string]any{
"upstreamModel": omni, "baseUrl": "https://vertex.example/",
"authHeader": "Bearer token",
"auth": map[string]any{"authHeader": "Bearer token", "projectId": "my-project"},
"files": []any{map[string]any{"ref": "request_file:input_reference", "field": "input_reference", "mimeType": "image/png"}},
"requestBody": map[string]any{"prompt": "animate", "images": []any{map[string]any{"__fileRef": "request_file:input_reference", "encoding": "base64"}}},
})
require.NoError(t, err)
assert.Equal(t, "https://vertex.example/v1beta1/projects/my-project/locations/global/interactions", descriptor["url"])
assert.Equal(t, "image_to_video", descriptor["action"])
assert.Equal(t, []any{
map[string]any{"type": "image", "mime_type": "image/png", "data": map[string]any{"__fileRef": "request_file:input_reference", "encoding": "base64"}},
map[string]any{"type": "text", "text": "animate"},
}, descriptor["body"].(map[string]any)["input"])
// Several images are sent in order without a task hint, which would limit them to one.
descriptor, err = call(map[string]any{
"upstreamModel": "gemini-omni-1.2-flash", "authHeader": "Bearer t", "auth": map[string]any{"authHeader": "Bearer t", "projectId": "p"},
"requestBody": map[string]any{"prompt": "morph", "duration": 15, "images": []any{"data:image/png;base64,aGVsbG8=", "/9j/4AAQ"}},
})
require.NoError(t, err)
assert.Equal(t, map[string]any{
"model": "gemini-omni-1.2-flash",
"input": []any{
map[string]any{"type": "image", "mime_type": "image/png", "data": "aGVsbG8="},
map[string]any{"type": "image", "mime_type": "image/jpeg", "data": "/9j/4AAQ"},
map[string]any{"type": "text", "text": "morph"},
},
"response_format": []any{map[string]any{"type": "video", "aspect_ratio": "16:9", "resolution": "720p", "duration": "15s"}},
"background": true,
"store": true,
}, descriptor["body"])
_, err = call(map[string]any{"upstreamModel": omni, "authHeader": "Bearer t", "auth": map[string]any{"authHeader": "Bearer t", "projectId": "p"}, "requestBody": map[string]any{"prompt": "x", "duration": 16}})
assert.ErrorContains(t, err, "3 to 15")
_, err = call(map[string]any{"upstreamModel": omni, "authHeader": "Bearer t", "auth": map[string]any{"authHeader": "Bearer t", "projectId": "p"}, "requestBody": map[string]any{"prompt": "x", "resolution": "2k"}})
assert.ErrorContains(t, err, "resolution must be one of 360p")
_, err = call(map[string]any{"upstreamModel": omni, "upstream": map[string]any{"kind": "new_api"}, "authHeader": "Bearer t", "auth": map[string]any{}, "requestBody": map[string]any{"prompt": "x"}})
assert.ErrorContains(t, err, "New API gateway")
}
func TestVertexAIOmniQueryAndParse(t *testing.T) {
plugin := loadVertexAIPlugin(t)
const omni = "gemini-omni-1.1-flash-preview"
call := func(hook string, args ...any) map[string]any {
t.Helper()
value, err := plugin.Engine.Call(t.Context(), hook, args...)
require.NoError(t, err)
return decodePluginValue(t, value)
}
queryCtx := map[string]any{"taskId": "int_1", "model": "gemini-omni-1.1-flash", "upstreamModel": omni, "authHeader": "Bearer token", "auth": map[string]any{"authHeader": "Bearer token", "projectId": "p1"}}
descriptor := call("buildQueryRequest", queryCtx)
assert.Equal(t, "https://aiplatform.googleapis.com/v1beta1/projects/p1/locations/global/interactions/int_1", descriptor["url"])
assert.Equal(t, "GET", descriptor["method"])
submitted := call("parseSubmitResponse", queryCtx, map[string]any{"body": map[string]any{"id": "int_1", "status": "in_progress"}})
assert.Equal(t, "int_1", submitted["taskId"])
// The stored interaction keeps the inline video's media type only, so the
// artifact is served by fetching the interaction again and reading the
// model_output step.
redacted := map[string]any{"id": "int_1", "object": "interaction", "status": "completed", "steps": []any{
map[string]any{"type": "thought", "summary": []any{}},
map[string]any{"type": "model_output", "content": []any{map[string]any{"type": "video", "mime_type": "video/mp4"}}},
}}
inline := map[string]any{"id": "int_1", "object": "interaction", "status": "completed", "steps": []any{
map[string]any{"type": "model_output", "content": []any{map[string]any{"type": "video", "mime_type": "video/mp4", "data": "AAAA"}}},
}}
stored := map[string]any{"id": "int_1", "object": "interaction", "status": "completed", "steps": []any{
map[string]any{"type": "model_output", "content": []any{map[string]any{"type": "video", "mime_type": "video/mp4", "uri": "gs://b/v.mp4"}}},
}}
for _, tc := range []struct {
name string
body map[string]any
want map[string]any
}{
{"running", map[string]any{"id": "int_1", "status": "in_progress"}, map[string]any{"status": "IN_PROGRESS", "progress": "50%"}},
{"completed inline", inline, map[string]any{"status": "SUCCESS", "progress": "100%"}},
{"completed uri", stored, map[string]any{"status": "SUCCESS", "progress": "100%", "remoteUrl": "gs://b/v.mp4"}},
{"completed without video", map[string]any{"id": "int_1", "status": "completed", "steps": []any{}}, map[string]any{"status": "FAILURE", "progress": "100%", "reason": "interaction completed without a video"}},
{"failed", map[string]any{"id": "int_1", "status": "failed"}, map[string]any{"status": "FAILURE", "progress": "100%", "reason": "failed"}},
{"blocked", map[string]any{"id": "int_1", "status": "failed", "errors": []any{map[string]any{"code": "content_blocked", "message": "Request blocked for an unspecified policy reason."}}}, map[string]any{"status": "FAILURE", "progress": "100%", "reason": "content_blocked: Request blocked for an unspecified policy reason."}},
{"error", map[string]any{"error": map[string]any{"message": "quota exceeded"}}, map[string]any{"status": "FAILURE", "progress": "100%", "reason": "quota exceeded"}},
{"unrecognized", map[string]any{"foo": "bar"}, map[string]any{"status": "UNKNOWN", "reason": "unrecognized interaction state"}},
} {
t.Run(tc.name, func(t *testing.T) {
assert.Equal(t, tc.want, call("parseTaskResult", queryCtx, tc.body))
})
}
artifacts, err := plugin.Engine.Call(t.Context(), "listArtifacts", map[string]any{"status": "SUCCESS", "data": redacted})
require.NoError(t, err)
assert.Equal(t, []any{map[string]any{"key": "video", "type": "video", "mimeType": "video/mp4"}}, artifacts)
artifacts, err = plugin.Engine.Call(t.Context(), "listArtifacts", map[string]any{"status": "SUCCESS", "data": stored})
require.NoError(t, err)
assert.Empty(t, artifacts)
assert.Equal(t, map[string]any{
"url": "https://aiplatform.googleapis.com/v1beta1/projects/p1/locations/global/interactions/int_1",
"method": "GET",
"headers": map[string]any{"Accept": "application/json", "Content-Type": "application/json", "Authorization": "Bearer token", "x-goog-user-project": "p1"},
"responseType": "json",
"contentPath": []any{"steps", float64(1), "content", float64(0), "data"},
"mimeType": "video/mp4",
}, call("buildContentRequest", map[string]any{
"artifactKey": "video", "upstreamTaskId": "int_1", "data": redacted, "authHeader": "Bearer token",
"auth": map[string]any{"authHeader": "Bearer token", "projectId": "p1"}, "clientRequest": map[string]any{"method": "GET"},
}))
}
// The OpenAI video decoder leaves Gemini Omni's duration and resolution to the
// request builder, which allows what Veo does not.
func TestVertexAIOmniUsageAndVideoDecoder(t *testing.T) {
plugin := loadVertexAIPlugin(t)
const omni = "gemini-omni-1.1-flash-preview"
usage, err := plugin.Engine.Call(t.Context(), "extractUsage", map[string]any{"upstreamModel": omni, "requestBody": map[string]any{"prompt": "x", "size": "1920x1080"}})
require.NoError(t, err)
assert.Equal(t, map[string]any{"input_tokens": float64(0), "output_tokens": float64(0), "video_tokens": float64(8 * 8688)}, decodePluginValue(t, usage))
schema, _ := plugin.Meta.UsageForModel(omni)
assert.ElementsMatch(t, []string{"input_tokens", "output_tokens", "video_tokens"}, slices.Collect(maps.Keys(schema)))
schema12, _ := plugin.Meta.UsageForModel("gemini-omni-1.2-flash")
assert.Equal(t, schema, schema12)
// Thinking is billed as text output.
completed, err := plugin.Engine.Call(t.Context(), "extractUsageOnComplete", map[string]any{"upstreamModel": omni}, map[string]any{}, map[string]any{
"usage": map[string]any{
"total_input_tokens": 12, "total_thought_tokens": 40,
"output_tokens_by_modality": []any{map[string]any{"modality": "VIDEO", "tokens": 5793}, map[string]any{"modality": "text", "tokens": 7}},
},
})
require.NoError(t, err)
assert.Equal(t, map[string]any{"input_tokens": float64(12), "output_tokens": float64(47), "video_tokens": float64(5793)}, decodePluginValue(t, completed))
// Unreported facts are omitted so the host keeps the submit estimate.
partial, err := plugin.Engine.Call(t.Context(), "extractUsageOnComplete", map[string]any{"upstreamModel": omni}, map[string]any{}, map[string]any{"usage": map[string]any{"total_input_tokens": 12}})
require.NoError(t, err)
assert.Equal(t, map[string]any{"input_tokens": float64(12)}, decodePluginValue(t, partial))
veo, err := plugin.Engine.Call(t.Context(), "extractUsageOnComplete", map[string]any{"upstreamModel": "veo-3.0-generate-001"}, map[string]any{}, map[string]any{"done": true})
require.NoError(t, err)
assert.Nil(t, veo)
decoded, err := plugin.Engine.CallPath(t.Context(), "protocols", []string{"openai_video", "decodeRequest"}, map[string]any{
"model": "omni-alias", "upstreamModel": omni,
"body": map[string]any{"kind": "json", "value": map[string]any{"model": "omni-alias", "prompt": "a cat", "seconds": "3", "resolution": "360p"}},
})
require.NoError(t, err)
requestBody := decodePluginValue(t, decoded)["requestBody"].(map[string]any)
assert.Equal(t, "omni-alias", requestBody["model"])
assert.Equal(t, float64(3), requestBody["duration"])
assert.Equal(t, "360p", requestBody["resolution"])
}
func loadVertexAIPlugin(t *testing.T) *jsplugin.LoadedPlugin {
t.Helper()
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)
return plugin
}