feat(gemini): 接入 Veo 官方异步生成协议
将 Gemini Veo 视频请求改为 predictLongRunning 提交与 Operation 轮询,支持任务恢复、首尾帧和参考图的官方 Base64 字段映射。 完成后由网关携带 API Key 从受信任的 Gemini 文件域下载视频,再进入现有对象存储转存链路,避免暴露鉴权地址。 验证:env -u AI_GATEWAY_TEST_DATABASE_URL go test ./... -count=1;go vet ./...;真实 Veo 3.1 任务生成并下载 1633544 字节视频。
This commit is contained in:
@@ -0,0 +1,267 @@
|
||||
package clients
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"reflect"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
func TestGeminiVeoRunUsesOfficialLongRunningProtocol(t *testing.T) {
|
||||
const operationName = "models/veo-3.1-generate-preview/operations/test-operation"
|
||||
videoPayload := []byte("test mp4 payload")
|
||||
var server *httptest.Server
|
||||
var pollCount atomic.Int32
|
||||
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Header.Get("x-goog-api-key") != "test-gemini-key" {
|
||||
t.Errorf("missing Gemini API key header for %s", r.URL.Path)
|
||||
}
|
||||
switch {
|
||||
case r.Method == http.MethodPost && r.URL.Path == "/v1beta/models/veo-3.1-generate-preview:predictLongRunning":
|
||||
var body map[string]any
|
||||
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
||||
t.Fatalf("decode Veo body: %v", err)
|
||||
}
|
||||
instances := body["instances"].([]any)
|
||||
if prompt := instances[0].(map[string]any)["prompt"]; prompt != "a paper boat on a lake" {
|
||||
t.Errorf("unexpected prompt: %v", prompt)
|
||||
}
|
||||
parameters := body["parameters"].(map[string]any)
|
||||
if parameters["durationSeconds"] != float64(4) || parameters["resolution"] != "720p" || parameters["aspectRatio"] != "16:9" {
|
||||
t.Errorf("unexpected parameters: %+v", parameters)
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"name": operationName})
|
||||
case r.Method == http.MethodGet && r.URL.Path == "/v1beta/"+operationName:
|
||||
if pollCount.Add(1) == 1 {
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"name": operationName, "done": false})
|
||||
return
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"name": operationName,
|
||||
"done": true,
|
||||
"response": map[string]any{"generateVideoResponse": map[string]any{"generatedSamples": []any{
|
||||
map[string]any{"video": map[string]any{"uri": server.URL + "/v1beta/files/generated-video:download", "mimeType": "video/mp4"}},
|
||||
}}},
|
||||
})
|
||||
case r.Method == http.MethodGet && r.URL.Path == "/v1beta/files/generated-video:download":
|
||||
w.Header().Set("Content-Type", "video/mp4")
|
||||
_, _ = w.Write(videoPayload)
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
var submitted string
|
||||
polled := 0
|
||||
response, err := (GeminiClient{}).Run(context.Background(), Request{
|
||||
Kind: "videos.generations",
|
||||
ModelType: "video_generate",
|
||||
Model: "veo-3.1",
|
||||
Body: map[string]any{
|
||||
"prompt": "a paper boat on a lake",
|
||||
"duration": 4,
|
||||
"resolution": "720p",
|
||||
"aspect_ratio": "16:9",
|
||||
},
|
||||
Candidate: testGeminiVeoCandidate(server.URL),
|
||||
OnRemoteTaskSubmitted: func(remoteTaskID string, _ map[string]any) error {
|
||||
submitted = remoteTaskID
|
||||
return nil
|
||||
},
|
||||
OnRemoteTaskPolled: func(remoteTaskID string, _ map[string]any) error {
|
||||
if remoteTaskID != operationName {
|
||||
t.Errorf("unexpected polled operation: %s", remoteTaskID)
|
||||
}
|
||||
polled++
|
||||
return nil
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("run Gemini Veo: %v", err)
|
||||
}
|
||||
if submitted != operationName || polled != 2 || pollCount.Load() != 2 {
|
||||
t.Fatalf("unexpected async callbacks: submitted=%q polled=%d requests=%d", submitted, polled, pollCount.Load())
|
||||
}
|
||||
if response.UpstreamProtocol != ProtocolGeminiVeo || response.RequestID != operationName {
|
||||
t.Fatalf("unexpected response metadata: %+v", response)
|
||||
}
|
||||
data := response.Result["data"].([]any)
|
||||
item := data[0].(map[string]any)
|
||||
if got := item["video_bytes"].([]byte); !reflect.DeepEqual(got, videoPayload) {
|
||||
t.Fatalf("unexpected video payload: %q", got)
|
||||
}
|
||||
if _, exposed := item["uri"]; exposed {
|
||||
t.Fatalf("upstream authenticated URI must not be exposed: %+v", item)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeminiVeoRunResumesOperationWithoutSubmittingAgain(t *testing.T) {
|
||||
const operationName = "models/veo-3.1-generate-preview/operations/resumed"
|
||||
var postCount atomic.Int32
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method == http.MethodPost {
|
||||
postCount.Add(1)
|
||||
http.Error(w, "unexpected submit", http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"name": operationName,
|
||||
"done": true,
|
||||
"response": map[string]any{"generateVideoResponse": map[string]any{"generatedSamples": []any{
|
||||
map[string]any{"video": map[string]any{"videoBytes": base64.StdEncoding.EncodeToString([]byte("resumed video")), "mimeType": "video/mp4"}},
|
||||
}}},
|
||||
})
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
response, err := (GeminiClient{}).Run(context.Background(), Request{
|
||||
Kind: "videos.generations",
|
||||
ModelType: "video_generate",
|
||||
Model: "veo-3.1",
|
||||
Body: map[string]any{"prompt": "ignored on resume"},
|
||||
Candidate: testGeminiVeoCandidate(server.URL),
|
||||
RemoteTaskID: operationName,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("resume Gemini Veo: %v", err)
|
||||
}
|
||||
if postCount.Load() != 0 || response.Result["id"] != operationName {
|
||||
t.Fatalf("resume submitted a new operation or returned wrong ID: posts=%d result=%+v", postCount.Load(), response.Result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeminiVeoBodyMapsImageAndParameters(t *testing.T) {
|
||||
first := "data:image/png;base64," + base64.StdEncoding.EncodeToString([]byte("first"))
|
||||
last := "data:image/jpeg;base64," + base64.StdEncoding.EncodeToString([]byte("last"))
|
||||
reference := "data:image/webp;base64," + base64.StdEncoding.EncodeToString([]byte("reference"))
|
||||
body, err := geminiVeoBody(Request{Body: map[string]any{
|
||||
"prompt": "animate",
|
||||
"first_frame": first,
|
||||
"last_frame": last,
|
||||
"reference_images": []any{reference},
|
||||
"duration": 8,
|
||||
"resolution": "2160p",
|
||||
"aspect_ratio": "9:16",
|
||||
"negative_prompt": "text",
|
||||
"person_generation": "allow_adult",
|
||||
"enhance_prompt": false,
|
||||
}})
|
||||
if err != nil {
|
||||
t.Fatalf("build Gemini Veo body: %v", err)
|
||||
}
|
||||
instance := body["instances"].([]any)[0].(map[string]any)
|
||||
image := instance["image"].(map[string]any)
|
||||
lastFrame := instance["lastFrame"].(map[string]any)
|
||||
references := instance["referenceImages"].([]any)
|
||||
if image["bytesBase64Encoded"] != base64.StdEncoding.EncodeToString([]byte("first")) || image["mimeType"] != "image/png" {
|
||||
t.Fatalf("unexpected first frame: %+v", image)
|
||||
}
|
||||
if lastFrame["mimeType"] != "image/jpeg" || len(references) != 1 {
|
||||
t.Fatalf("unexpected last/reference images: last=%+v references=%+v", lastFrame, references)
|
||||
}
|
||||
parameters := body["parameters"].(map[string]any)
|
||||
if parameters["resolution"] != "4k" || parameters["durationSeconds"] != 8 || parameters["enhancePrompt"] != false {
|
||||
t.Fatalf("unexpected Gemini Veo parameters: %+v", parameters)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeminiVeoRejectsInvalidDurationBeforeSubmission(t *testing.T) {
|
||||
_, err := geminiVeoBody(Request{Body: map[string]any{"prompt": "test", "duration": 5}})
|
||||
if err == nil || ErrorCode(err) != "invalid_parameter" || ErrorParam(err) != "duration" || !strings.Contains(err.Error(), "4, 6, or 8") {
|
||||
t.Fatalf("unexpected invalid duration error: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeminiVeoDownloadURLOnlyTrustsProviderAndGoogleAPIs(t *testing.T) {
|
||||
for _, test := range []struct {
|
||||
uri string
|
||||
trusted bool
|
||||
}{
|
||||
{uri: "https://gemini-proxy.example.com/v1beta/files/video", trusted: true},
|
||||
{uri: "https://generativelanguage.googleapis.com/v1beta/files/video", trusted: true},
|
||||
{uri: "https://storage.googleapis.com/signed/video", trusted: false},
|
||||
{uri: "https://attacker.example.com/video", trusted: false},
|
||||
} {
|
||||
_, trusted, err := geminiVeoDownloadURL("https://gemini-proxy.example.com/v1beta", test.uri)
|
||||
if err != nil {
|
||||
t.Fatalf("parse %s: %v", test.uri, err)
|
||||
}
|
||||
if trusted != test.trusted {
|
||||
t.Fatalf("unexpected trust for %s: got %t want %t", test.uri, trusted, test.trusted)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeminiVeoLive(t *testing.T) {
|
||||
apiKey := strings.TrimSpace(os.Getenv("GEMINI_VEO_LIVE_API_KEY"))
|
||||
if apiKey == "" {
|
||||
t.Skip("GEMINI_VEO_LIVE_API_KEY is not configured")
|
||||
}
|
||||
baseURL := strings.TrimSpace(os.Getenv("GEMINI_VEO_LIVE_BASE_URL"))
|
||||
model := strings.TrimSpace(os.Getenv("GEMINI_VEO_LIVE_MODEL"))
|
||||
if model == "" {
|
||||
model = "veo-3.1-generate-preview"
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 12*time.Minute)
|
||||
defer cancel()
|
||||
operationName := strings.TrimSpace(os.Getenv("GEMINI_VEO_LIVE_OPERATION"))
|
||||
response, err := (GeminiClient{}).Run(ctx, Request{
|
||||
Kind: "videos.generations",
|
||||
ModelType: "video_generate",
|
||||
Model: model,
|
||||
Body: map[string]any{
|
||||
"prompt": "A small red paper boat gently floating on calm water, static camera, natural daylight.",
|
||||
"duration": 4,
|
||||
"resolution": "720p",
|
||||
"aspect_ratio": "16:9",
|
||||
},
|
||||
Candidate: store.RuntimeModelCandidate{
|
||||
Provider: "gemini",
|
||||
BaseURL: baseURL,
|
||||
ProviderModelName: model,
|
||||
Credentials: map[string]any{"apiKey": apiKey},
|
||||
PlatformConfig: map[string]any{"pollIntervalMs": 10000, "pollTimeoutMs": 660000},
|
||||
},
|
||||
RemoteTaskID: operationName,
|
||||
OnRemoteTaskSubmitted: func(remoteTaskID string, _ map[string]any) error {
|
||||
t.Logf("live Gemini Veo submitted operation %s", remoteTaskID)
|
||||
return nil
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("live Gemini Veo request failed: %v", err)
|
||||
}
|
||||
data, _ := response.Result["data"].([]any)
|
||||
if len(data) != 1 {
|
||||
t.Fatalf("live Gemini Veo returned %d videos", len(data))
|
||||
}
|
||||
item, _ := data[0].(map[string]any)
|
||||
payload, _ := item["video_bytes"].([]byte)
|
||||
if len(payload) < 1024 {
|
||||
t.Fatalf("live Gemini Veo video is unexpectedly small: %d bytes", len(payload))
|
||||
}
|
||||
if mimeType := strings.TrimSpace(stringFromAny(item["mime_type"])); !strings.HasPrefix(mimeType, "video/") {
|
||||
t.Fatalf("live Gemini Veo returned unexpected MIME type: %q", mimeType)
|
||||
}
|
||||
t.Logf("live Gemini Veo accepted operation %s and downloaded %d video bytes", response.RequestID, len(payload))
|
||||
}
|
||||
|
||||
func testGeminiVeoCandidate(baseURL string) store.RuntimeModelCandidate {
|
||||
return store.RuntimeModelCandidate{
|
||||
Provider: "gemini",
|
||||
BaseURL: baseURL,
|
||||
ProviderModelName: "veo-3.1-generate-preview",
|
||||
Credentials: map[string]any{"apiKey": "test-gemini-key"},
|
||||
PlatformConfig: map[string]any{"pollIntervalMs": 1, "pollTimeoutMs": 1000},
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user