ai-flow-python/app/service/task/impl/google_veo3_task.py

import logging
import os
import time

from app.config import settings
from app.service.task.base_task import BaseTask
from app.service.task.task_param import TaskParam
from app.service.task.task_result import TaskResult
from app.util.file_util import generate_unique_filename, write_bytes
from app.util.http_util import logged_client

logger = logging.getLogger(__name__)


class GoogleVeo3Task(BaseTask):
    def get_type(self) -> str:
        return "google_veo3"

    def _headers(self) -> dict[str, str]:
        return {
            "x-goog-api-key": settings.gemini_api_key,
            "Content-Type": "application/json",
        }

    def execute(self, param: TaskParam) -> TaskResult:
        t0 = time.time()
        try:
            op_name = self._create_request(param.inputs)
            logger.info("[google_veo3] Operation created: %s", op_name)

            video_uri = self._poll_until_done(op_name)
            if not video_uri:
                return TaskResult.failure("Failed to get video download link")

            logger.info("[google_veo3] Downloading video: %s", video_uri)
            video_bytes = self._download_video(video_uri)
            video_name = generate_unique_filename(".mp4")
            video_rel = os.path.join(param.step_id, video_name)
            write_bytes(os.path.join(param.workspace, video_rel), video_bytes)

            logger.info("[google_veo3] Completed: video=%s, size=%d bytes, total_elapsed=%.2fs",
                         video_rel, len(video_bytes), time.time() - t0)
            return TaskResult.success({"videos": [video_rel]})
        except Exception as e:
            logger.error("[google_veo3] Execution failed in %.2fs: %s", time.time() - t0, e, exc_info=True)
            return TaskResult.failure(f"GoogleVeo3Task execution failed: {e}")

    def _create_request(self, inputs: dict) -> str:
        endpoint = f"{settings.gemini_base_url}/models/{settings.gemini_veo3_model}:predictLongRunning"
        logger.info("[google_veo3] Submitting generation request: POST %s", endpoint)
        with logged_client(timeout=settings.http_timeout) as client:
            resp = client.post(endpoint, headers=self._headers(), json=inputs)
            logger.info("[google_veo3] Submit response: status=%d", resp.status_code)
            resp.raise_for_status()
            data = resp.json()
        op_name = data.get("name")
        if not op_name:
            raise RuntimeError(f"Operation name not found: {data}")
        return op_name

    def _poll_until_done(self, operation_name: str) -> str | None:
        url = f"{settings.gemini_base_url}/{operation_name}"
        max_retries = 60
        poll_t0 = time.time()

        for i in range(max_retries):
            try:
                with logged_client(timeout=settings.http_timeout) as client:
                    resp = client.get(url, headers=self._headers())
                    if not resp.is_success:
                        logger.warning("[google_veo3] Poll #%d: HTTP %d, elapsed=%.1fs",
                                        i + 1, resp.status_code, time.time() - poll_t0)
                        time.sleep(10)
                        continue
                    data = resp.json()
            except Exception as e:
                logger.warning("[google_veo3] Poll #%d: exception=%s, elapsed=%.1fs",
                                i + 1, e, time.time() - poll_t0)
                time.sleep(10)
                continue

            done = data.get("done")
            logger.info("[google_veo3] Poll #%d: done=%s, elapsed=%.1fs",
                         i + 1, done, time.time() - poll_t0)
            if done is None:
                time.sleep(10)
                continue
            if done:
                try:
                    samples = data["response"]["generateVideoResponse"]["generatedSamples"]
                    return samples[0]["video"]["uri"]
                except (KeyError, IndexError):
                    raise RuntimeError("Task completed but video URI not found")
            time.sleep(10)

        raise RuntimeError("Polling timeout after %d attempts" % max_retries)

    def _download_video(self, video_uri: str) -> bytes:
        t0 = time.time()
        with logged_client(timeout=settings.http_timeout, follow_redirects=True) as client:
            resp = client.get(video_uri, headers=self._headers())
            resp.raise_for_status()
            logger.info("[google_veo3] Video downloaded: %d bytes in %.2fs",
                         len(resp.content), time.time() - t0)
            return resp.content