Tool_Agent/tools/local/veo_video_gen/veo_client.py

import requests
import time
from typing import Dict, Any


class VEOClient:
    """VEO 3.1 异步调用客户端"""

    def __init__(self, api_key: str, base_url: str = "https://api.apiyi.com"):
        self.base_url = base_url
        self.headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {api_key}"
        }

    def create_video(self, prompt: str, model: str = "veo-3.1-fast") -> str:
        """创建视频任务，返回 video_id"""
        response = requests.post(
            f"{self.base_url}/v1/videos",
            headers=self.headers,
            json={"prompt": prompt, "model": model}
        )
        response.raise_for_status()
        return response.json()["id"]

    def get_status(self, video_id: str) -> Dict[str, Any]:
        """查询任务状态，completed 时响应中直接含 video_url"""
        response = requests.get(
            f"{self.base_url}/v1/videos/{video_id}",
            headers=self.headers
        )
        response.raise_for_status()
        return response.json()

    def wait_for_completion(
        self,
        video_id: str,
        timeout: int = 600,
        interval: int = 5
    ) -> Dict[str, Any]:
        """轮询直到完成，返回包含 video_url 的状态数据"""
        start_time = time.time()
        while time.time() - start_time < timeout:
            data = self.get_status(video_id)
            status = data.get("status")

            if status == "completed":
                return data
            elif status == "failed":
                raise Exception(f"视频生成失败: {data}")

            time.sleep(interval)

        raise TimeoutError("等待超时")