Add Grok and VEO video workflow support

2026-06-04 16:06:23 +08:00
parent 5d9aff9ca7
commit 30e9603f77
24 changed files with 3752 additions and 528 deletions
@@ -0,0 +1,483 @@
+"""
+Grok Video API client.
+
+Flow:
+  1. POST /v1/videos
+  2. GET  /v1/videos/{task_id}
+  3. GET  /v1/videos/{task_id}/content, or download a URL from the status body
+"""
+
+import asyncio
+import base64
+import json
+import os
+import re
+import time
+from typing import Any, Callable, Dict, List, Optional
+
+import aiohttp
+
+from .base_client import BaseAPIClient
+from ..utils.config import get_api_base_url, get_api_key_or_raise
+from ..utils.http_error import RETRYABLE_STATUS_CODES, get_friendly_message
+from ..utils.video_task import (
+    check_interrupt,
+    extract_error_message,
+    extract_progress,
+    extract_status,
+    extract_video_url,
+    interruptible_sleep,
+    is_failure_status,
+    is_success_status,
+    run_with_interrupt,
+)
+
+
+class GrokVideoClient(BaseAPIClient):
+    CREATE_ENDPOINT = "/v1/videos"
+    STATUS_ENDPOINT = "/v1/videos/{task_id}"
+    CONTENT_ENDPOINT = "/v1/videos/{task_id}/content"
+
+    MODEL_OPTIONS = ["grok-imagine-video-1.5-preview", "grok-imagine-1.0-video"]
+    ASPECT_RATIO_OPTIONS = ["1:1", "16:9", "9:16", "4:3", "3:4", "3:2", "2:3"]
+    QUALITY_OPTIONS = ["720p"]
+    MODEL_SECONDS_OPTIONS = {
+        "grok-imagine-1.0-video": [6, 10, 12, 16, 20],
+    }
+    QUALITY_API_MAP = {
+        "720p": "high",
+        "high": "high",
+    }
+
+    SUCCESS_STATUSES = {"complete", "completed", "succeed", "succeeded", "success", "done", "finished"}
+    FAILURE_STATUSES = {"fail", "failed", "failure", "error", "expired", "timeout", "cancelled", "canceled"}
+
+    def __init__(self, base_url: Optional[str] = None):
+        api_key = get_api_key_or_raise("O1KEY_API_KEY")
+        resolved_base_url = (base_url or "").strip() or get_api_base_url()
+        super().__init__(base_url=resolved_base_url.rstrip("/"), api_key=api_key)
+
+    def get_endpoint(self, **kwargs) -> str:
+        return self.CREATE_ENDPOINT
+
+    def build_request_body(self, **kwargs) -> Dict[str, Any]:
+        return self.build_video_body(**kwargs)
+
+    def parse_response(self, response: Dict[str, Any]) -> Any:
+        return response
+
+    @classmethod
+    def build_video_body(
+        cls,
+        prompt: str,
+        model: str,
+        aspect_ratio: str,
+        seconds: int,
+        quality: str = "720p",
+        images: Optional[List[str]] = None,
+    ) -> Dict[str, Any]:
+        prompt = (prompt or "").strip()
+        if not prompt:
+            raise ValueError("提示词不能为空。")
+
+        if model not in cls.MODEL_OPTIONS:
+            raise ValueError(f"模型仅支持: {', '.join(cls.MODEL_OPTIONS)}。")
+
+        if aspect_ratio not in cls.ASPECT_RATIO_OPTIONS:
+            raise ValueError(f"宽高比仅支持: {', '.join(cls.ASPECT_RATIO_OPTIONS)}。")
+
+        try:
+            seconds_value = int(seconds)
+        except (TypeError, ValueError):
+            raise ValueError("秒数必须是整数。") from None
+
+        allowed_seconds = cls.MODEL_SECONDS_OPTIONS.get(model)
+        if allowed_seconds is not None:
+            if seconds_value not in allowed_seconds:
+                raise ValueError(
+                    f"模型 {model} 仅支持秒数: "
+                    f"{', '.join(str(s) for s in allowed_seconds)}。"
+                    "请修改为正确的秒数后再发起请求。"
+                )
+        elif seconds_value < 5 or seconds_value > 15:
+            raise ValueError("秒数仅支持 5 到 15。")
+
+        api_quality = cls.QUALITY_API_MAP.get(str(quality), str(quality))
+        if api_quality != "high":
+            raise ValueError("画质仅支持 720p。")
+
+        body: Dict[str, Any] = {
+            "model": model,
+            "prompt": prompt,
+            "aspect_ratio": aspect_ratio,
+            "seconds": str(seconds_value),
+            "quality": api_quality,
+        }
+
+        image_list = [img for img in (images or []) if img]
+        if image_list:
+            body["images"] = image_list[:3]
+
+        return body
+
+    @staticmethod
+    def _safe_task_filename(task_id: str) -> str:
+        safe = re.sub(r"[^A-Za-z0-9_.-]+", "_", task_id).strip("._")
+        return safe or "grok_video"
+
+    @staticmethod
+    def _mask_body_for_log(body: Dict[str, Any]) -> Dict[str, Any]:
+        log_body = dict(body)
+        images = log_body.get("images")
+        if isinstance(images, list):
+            log_body["images"] = [f"<data-url chars={len(item)}>" for item in images]
+        return log_body
+
+    @staticmethod
+    def _extract_task_id(payload: Dict[str, Any]) -> Optional[str]:
+        sources = [payload]
+        data = payload.get("data")
+        if isinstance(data, dict):
+            sources.append(data)
+
+        for source in sources:
+            for key in ("id", "task_id", "video_id"):
+                value = source.get(key)
+                if value:
+                    return str(value)
+        return None
+
+    @staticmethod
+    def _format_http_error(endpoint: str, status: int, error_text: str, task_id: Optional[str] = None) -> str:
+        message = get_friendly_message(status, error_text)
+        parts = [
+            "Grok Video 请求失败。",
+            f"endpoint: {endpoint}",
+            f"http_status: {status}",
+        ]
+        if task_id:
+            parts.append(f"task_id: {task_id}")
+        if message:
+            parts.append(f"message: {message}")
+        return "\n".join(parts)
+
+    @classmethod
+    def _format_task_failure(cls, task_id: str, payload: Dict[str, Any]) -> str:
+        return "\n".join(
+            [
+                "Grok Video 任务失败。",
+                f"endpoint: {cls.STATUS_ENDPOINT.format(task_id=task_id)}",
+                f"task_id: {task_id}",
+                f"message: {extract_error_message(payload)}",
+            ]
+        )
+
+    async def _request_json_with_retry(
+        self,
+        method: str,
+        endpoint: str,
+        session: aiohttp.ClientSession,
+        task_id: Optional[str] = None,
+        json_body: Optional[Dict[str, Any]] = None,
+        max_retries: int = 3,
+        timeout_seconds: int = 120,
+    ) -> Dict[str, Any]:
+        url = f"{self.base_url}{endpoint}"
+        headers = self.get_headers(use_bearer_token=True)
+        timeout = aiohttp.ClientTimeout(total=timeout_seconds, connect=30, sock_read=timeout_seconds)
+
+        last_status = 0
+        last_text = ""
+
+        for attempt in range(max_retries + 1):
+            check_interrupt()
+            response = None
+            try:
+                response = await run_with_interrupt(
+                    session.request(method, url, json=json_body, headers=headers, timeout=timeout)
+                )
+                text = await run_with_interrupt(response.text())
+                last_status = response.status
+                last_text = text
+
+                if 200 <= response.status < 300:
+                    if not text.strip():
+                        return {}
+                    try:
+                        return json.loads(text)
+                    except Exception:
+                        raise RuntimeError(f"Grok Video 响应 JSON 解析失败，原始内容：{text[:500]}") from None
+
+                if response.status in RETRYABLE_STATUS_CODES and attempt < max_retries:
+                    delay = min(2 ** attempt, 8)
+                    print(
+                        f"Grok Video：{get_friendly_message(response.status)} "
+                        f"{delay}s 后重试 ({attempt + 1}/{max_retries})..."
+                    )
+                    await interruptible_sleep(delay)
+                    continue
+
+                break
+
+            except (aiohttp.ClientError, asyncio.TimeoutError) as e:
+                if attempt < max_retries:
+                    delay = min(2 ** attempt, 8)
+                    print(f"Grok Video：网络错误，{delay}s 后重试 ({attempt + 1}/{max_retries})...")
+                    await interruptible_sleep(delay)
+                    continue
+                raise RuntimeError(f"Grok Video 网络错误: {e}") from None
+
+            finally:
+                if response is not None:
+                    response.release()
+
+        raise RuntimeError(self._format_http_error(endpoint, last_status, last_text, task_id=task_id))
+
+    async def create_video_async(
+        self,
+        body: Dict[str, Any],
+        session: aiohttp.ClientSession,
+    ) -> Dict[str, Any]:
+        print("Grok Video：正在提交任务...")
+        return await self._request_json_with_retry(
+            "POST",
+            self.CREATE_ENDPOINT,
+            session=session,
+            json_body=body,
+            timeout_seconds=180,
+        )
+
+    async def poll_video_status_async(
+        self,
+        task_id: str,
+        session: aiohttp.ClientSession,
+        poll_interval: int = 5,
+        timeout: int = 900,
+        progress_callback: Optional[Callable[[int, str, float], None]] = None,
+    ) -> Dict[str, Any]:
+        endpoint = self.STATUS_ENDPOINT.format(task_id=task_id)
+        start = time.time()
+        interval = max(1, int(poll_interval))
+
+        await interruptible_sleep(interval)
+
+        while True:
+            data = await self._request_json_with_retry(
+                "GET",
+                endpoint,
+                session=session,
+                task_id=task_id,
+                timeout_seconds=60,
+            )
+
+            status = extract_status(data)
+            progress = extract_progress(data)
+            elapsed = time.time() - start
+
+            if progress_callback:
+                progress_callback(progress, status, elapsed)
+
+            if status in self.SUCCESS_STATUSES or is_success_status(status):
+                return data
+
+            if status in self.FAILURE_STATUSES or is_failure_status(status, data):
+                raise RuntimeError(self._format_task_failure(task_id, data))
+
+            if elapsed >= timeout:
+                raise TimeoutError(
+                    "Grok Video 任务轮询超时；任务未被标记为失败，可用 task_id 继续查询。\n"
+                    f"endpoint: {endpoint}\n"
+                    f"task_id: {task_id}\n"
+                    f"status: {status or 'unknown'}\n"
+                    f"timeout: {timeout}s"
+                )
+
+            await interruptible_sleep(min(interval, max(0.0, timeout - elapsed)))
+
+    async def _download_url_to_file(
+        self,
+        url: str,
+        save_path: str,
+        session: aiohttp.ClientSession,
+        max_retries: int = 3,
+    ) -> str:
+        timeout = aiohttp.ClientTimeout(total=900, connect=30, sock_read=900)
+        last_status = 0
+        last_text = ""
+        headers = None
+        resolved_url = url
+
+        if url.startswith("data:"):
+            if "," not in url:
+                raise RuntimeError("Grok Video 下载失败：data URL 格式无效。")
+            _, b64_data = url.split(",", 1)
+            os.makedirs(os.path.dirname(save_path), exist_ok=True)
+            with open(save_path, "wb") as f:
+                f.write(base64.b64decode(b64_data))
+            if not os.path.isfile(save_path) or os.path.getsize(save_path) <= 0:
+                raise RuntimeError("Grok Video 下载失败：保存后的文件为空。")
+            return save_path
+
+        if url.startswith("/"):
+            resolved_url = f"{self.base_url}{url}"
+            headers = self.get_headers(use_bearer_token=True)
+
+        for attempt in range(max_retries + 1):
+            check_interrupt()
+            async with session.get(
+                resolved_url,
+                headers=headers,
+                timeout=timeout,
+                allow_redirects=True,
+            ) as response:
+                if 200 <= response.status < 300:
+                    os.makedirs(os.path.dirname(save_path), exist_ok=True)
+                    with open(save_path, "wb") as f:
+                        async for chunk in response.content.iter_chunked(1024 * 1024):
+                            check_interrupt()
+                            if chunk:
+                                f.write(chunk)
+
+                    if not os.path.isfile(save_path) or os.path.getsize(save_path) <= 0:
+                        raise RuntimeError("Grok Video 下载失败：保存后的文件为空。")
+                    return save_path
+
+                last_status = response.status
+                last_text = await response.text()
+                if response.status not in RETRYABLE_STATUS_CODES or attempt >= max_retries:
+                    break
+
+            delay = min(2 ** attempt, 8)
+            print(f"Grok Video：下载重试 {attempt + 1}/{max_retries}，{delay}s 后继续...")
+            await interruptible_sleep(delay)
+
+        raise RuntimeError(self._format_http_error("download_url", last_status, last_text))
+
+    async def download_video_async(
+        self,
+        task_id: str,
+        save_path: str,
+        session: aiohttp.ClientSession,
+    ) -> str:
+        endpoint = self.CONTENT_ENDPOINT.format(task_id=task_id)
+        url = f"{self.base_url}{endpoint}"
+        headers = self.get_headers(use_bearer_token=True)
+        timeout = aiohttp.ClientTimeout(total=900, connect=30, sock_read=900)
+
+        last_status = 0
+        last_text = ""
+
+        for attempt in range(4):
+            check_interrupt()
+            async with session.get(url, headers=headers, timeout=timeout, allow_redirects=True) as response:
+                if 200 <= response.status < 300:
+                    content_type = response.headers.get("Content-Type", "").lower()
+                    if "application/json" in content_type:
+                        data = await response.json(content_type=None)
+                        download_url = extract_video_url(data)
+                        if not download_url:
+                            raise RuntimeError(
+                                "Grok Video 下载失败：content 响应为 JSON，但未包含视频 URL。\n"
+                                f"endpoint: {endpoint}\n"
+                                f"task_id: {task_id}"
+                            )
+                        return await self._download_url_to_file(download_url, save_path, session)
+
+                    os.makedirs(os.path.dirname(save_path), exist_ok=True)
+                    with open(save_path, "wb") as f:
+                        async for chunk in response.content.iter_chunked(1024 * 1024):
+                            check_interrupt()
+                            if chunk:
+                                f.write(chunk)
+
+                    if not os.path.isfile(save_path) or os.path.getsize(save_path) <= 0:
+                        raise RuntimeError(
+                            "Grok Video 下载失败：保存后的文件为空。\n"
+                            f"endpoint: {endpoint}\n"
+                            f"task_id: {task_id}"
+                        )
+                    return save_path
+
+                last_status = response.status
+                last_text = await response.text()
+                if response.status not in RETRYABLE_STATUS_CODES or attempt >= 3:
+                    break
+
+            delay = min(2 ** attempt, 8)
+            print(f"Grok Video：content 下载重试 {attempt + 1}/3，{delay}s 后继续...")
+            await interruptible_sleep(delay)
+
+        raise RuntimeError(self._format_http_error(endpoint, last_status, last_text, task_id=task_id))
+
+    def generate_video_sync(
+        self,
+        prompt: str,
+        model: str,
+        aspect_ratio: str,
+        seconds: int,
+        quality: str,
+        images: Optional[List[str]],
+        output_dir: Optional[str] = None,
+        save_path: Optional[str] = None,
+        poll_interval: int = 5,
+        timeout: int = 900,
+        progress_callback: Optional[Callable[[int, str, float], None]] = None,
+    ) -> Dict[str, Any]:
+        async def _run():
+            async with self._make_session() as session:
+                body = self.build_video_body(
+                    prompt=prompt,
+                    model=model,
+                    aspect_ratio=aspect_ratio,
+                    seconds=seconds,
+                    quality=quality,
+                    images=images,
+                )
+
+                create_response = await self.create_video_async(body, session)
+                task_id = self._extract_task_id(create_response) or ""
+                if not task_id:
+                    raise RuntimeError(
+                        "Grok Video 未返回任务 ID。\n"
+                        f"endpoint: {self.CREATE_ENDPOINT}\n"
+                        f"response: {json.dumps(create_response, ensure_ascii=False)[:1200]}"
+                    )
+
+                print(f"Grok Video：任务已提交，任务ID：{task_id}")
+                print("Grok Video：视频生成中...")
+                status_response = await self.poll_video_status_async(
+                    task_id=task_id,
+                    session=session,
+                    poll_interval=poll_interval,
+                    timeout=timeout,
+                    progress_callback=progress_callback,
+                )
+
+                video_url = extract_video_url(status_response)
+                print("Grok Video：视频生成完成，正在下载...")
+                if save_path is None:
+                    resolved_output_dir = output_dir or os.getcwd()
+                    os.makedirs(resolved_output_dir, exist_ok=True)
+                    target_path = os.path.join(
+                        resolved_output_dir,
+                        f"{self._safe_task_filename(task_id)}.mp4",
+                    )
+                else:
+                    target_path = save_path
+
+                if video_url:
+                    video_path = await self._download_url_to_file(video_url, target_path, session)
+                else:
+                    video_path = await self.download_video_async(task_id, target_path, session)
+
+                return {
+                    "task_id": task_id,
+                    "status": extract_status(status_response),
+                    "video_path": video_path,
+                    "raw_json": {
+                        "create": create_response,
+                        "status": status_response,
+                    },
+                }
+
+        return self.run_async_in_thread(_run())