diff --git a/skills/public/video-generation/scripts/generate.py b/skills/public/video-generation/scripts/generate.py index 182bda70a3..ff86873d38 100644 --- a/skills/public/video-generation/scripts/generate.py +++ b/skills/public/video-generation/scripts/generate.py @@ -82,6 +82,41 @@ def _poll_video_task(host: str, auth: str, task_id: str, max_attempts: int = 120 raise Exception(f"MiniMax video task {task_id} timed out after {max_attempts} polls") +def _poll_gemini_operation(api_key: str, operation_name: str, max_attempts: int = 240, interval: int = 3) -> str: + """Wait for a Veo long-running operation and return the generated video URI. + + Bounded like ``_poll_video_task`` so a stuck or rejected operation fails with a + diagnosis instead of polling forever. The default budget (12 minutes) leaves + room for queueing, which routinely outlasts a short generation. + """ + for _ in range(max_attempts): + response = requests.get( + f"https://generativelanguage.googleapis.com/v1beta/{operation_name}", + headers={"x-goog-api-key": api_key}, + timeout=30, + ) + response.raise_for_status() + payload = response.json() + if not payload.get("done", False): + time.sleep(interval) + continue + error = payload.get("error") or {} + if error: + raise Exception(f"Gemini video operation {operation_name} failed: {error.get('code')} {error.get('message')}") + video_response = (payload.get("response") or {}).get("generateVideoResponse") or {} + samples = video_response.get("generatedSamples") or [] + if not samples: + # A filtered prompt finishes successfully without any sample. + reasons = "; ".join(video_response.get("raiMediaFilteredReasons") or []) + detail = f": {reasons}" if reasons else "" + raise Exception(f"Gemini video operation {operation_name} returned no generated sample{detail}") + uri = (samples[0].get("video") or {}).get("uri") + if not uri: + raise Exception(f"Gemini video operation {operation_name} returned a sample without a video URI") + return uri + raise Exception(f"Gemini video operation {operation_name} timed out after {max_attempts} polls") + + def _retrieve_file_url(host: str, auth: str, file_id: str) -> str: response = requests.get( f"{host}/v1/files/retrieve", @@ -163,19 +198,8 @@ def _generate_video_gemini(prompt: str, reference_images: list[str], output_file response.raise_for_status() data = response.json() operation_name = data["name"] - while True: - response = requests.get( - f"https://generativelanguage.googleapis.com/v1beta/{operation_name}", - headers={"x-goog-api-key": gemini_api_key}, - timeout=30, - ) - response.raise_for_status() - data = response.json() - if data.get("done", False): - sample = data["response"]["generateVideoResponse"]["generatedSamples"][0] - download(sample["video"]["uri"], output_file) - break - time.sleep(3) + download_url = _poll_gemini_operation(gemini_api_key, operation_name) + download(download_url, output_file) return f"The video has been generated successfully to {output_file}" diff --git a/tests/skills/test_video_generation.py b/tests/skills/test_video_generation.py index e016e11552..f5fc2665d0 100644 --- a/tests/skills/test_video_generation.py +++ b/tests/skills/test_video_generation.py @@ -191,6 +191,91 @@ def test_gemini_post_raises_on_http_error(monkeypatch, tmp_path): vid.generate_video(str(pf), [], str(tmp_path / "v.mp4"), "16:9") +def test_gemini_poll_is_bounded(monkeypatch): + # The MiniMax path bounds its polling; the Gemini operation path must too. + polls = {"count": 0} + + def fake_get(url, headers=None, **kw): + polls["count"] += 1 + return FakeResp(json_data={"done": False}) + + monkeypatch.setattr(vid.requests, "get", fake_get) + with pytest.raises(Exception, match="timed out"): + vid._poll_gemini_operation("g", "operations/op", max_attempts=3, interval=0) + assert polls["count"] == 3 + + +def test_gemini_operation_error_is_reported(monkeypatch, tmp_path): + # A finished operation carrying an error must not surface as a KeyError on + # the missing "response" field. + monkeypatch.setenv("GEMINI_API_KEY", "g") + + def fake_post(url, headers=None, json=None, **kw): + return FakeResp(json_data={"name": "operations/op"}) + + def fake_get(url, headers=None, **kw): + return FakeResp(json_data={"done": True, "error": {"code": 3, "message": "prompt rejected"}}) + + monkeypatch.setattr(vid.requests, "post", fake_post) + monkeypatch.setattr(vid.requests, "get", fake_get) + pf = tmp_path / "p.json" + pf.write_text("x", encoding="utf-8") + out = tmp_path / "v.mp4" + with pytest.raises(Exception, match="prompt rejected"): + vid.generate_video(str(pf), [], str(out), "16:9") + assert not out.exists() + + +def test_gemini_filtered_sample_reports_reason(monkeypatch, tmp_path): + monkeypatch.setenv("GEMINI_API_KEY", "g") + reason = "The generated video was filtered by the safety system." + + def fake_post(url, headers=None, json=None, **kw): + return FakeResp(json_data={"name": "operations/op"}) + + def fake_get(url, headers=None, **kw): + return FakeResp( + json_data={ + "done": True, + "response": { + "generateVideoResponse": { + "raiMediaFilteredCount": 1, + "raiMediaFilteredReasons": [reason], + } + }, + } + ) + + monkeypatch.setattr(vid.requests, "post", fake_post) + monkeypatch.setattr(vid.requests, "get", fake_get) + pf = tmp_path / "p.json" + pf.write_text("x", encoding="utf-8") + with pytest.raises(Exception, match="filtered by the safety system"): + vid.generate_video(str(pf), [], str(tmp_path / "v.mp4"), "16:9") + + +def test_gemini_sample_without_uri_is_reported(monkeypatch, tmp_path): + monkeypatch.setenv("GEMINI_API_KEY", "g") + + def fake_post(url, headers=None, json=None, **kw): + return FakeResp(json_data={"name": "operations/op"}) + + def fake_get(url, headers=None, **kw): + return FakeResp( + json_data={ + "done": True, + "response": {"generateVideoResponse": {"generatedSamples": [{"video": {}}]}}, + } + ) + + monkeypatch.setattr(vid.requests, "post", fake_post) + monkeypatch.setattr(vid.requests, "get", fake_get) + pf = tmp_path / "p.json" + pf.write_text("x", encoding="utf-8") + with pytest.raises(Exception, match="without a video URI"): + vid.generate_video(str(pf), [], str(tmp_path / "v.mp4"), "16:9") + + def test_gemini_forwards_aspect_ratio_to_predict_request(monkeypatch, tmp_path): monkeypatch.setenv("GEMINI_API_KEY", "g") captured = {}