video-shuoshan/backend/utils/airdrama_client.py

"""Volcano Engine ARK video generation API client."""

import requests
from django.conf import settings


# API error code → user-friendly Chinese message
ERROR_MESSAGES = {
    # Input content moderation — 人脸/敏感内容
    'InputImageSensitiveContentDetected.PrivacyInformation': '参考图片中检测到真实人脸，请使用虚拟人像素材替代真人照片',
    'InputImageSensitiveContentDetected': '参考图片包含敏感内容，请更换图片后重试',
    'InputVideoSensitiveContentDetected': '参考视频包含敏感内容，请更换视频后重试',
    'InputTextSensitiveContentDetected': '提示词包含敏感内容，请修改后重试',
    'InputAudioSensitiveContentDetected': '参考音频包含敏感内容，请更换音频后重试',
    # Output content moderation
    'OutputVideoSensitiveContentDetected': '生成的视频包含敏感内容，已被系统拦截，请修改提示词后重试',
    'OutputImageSensitiveContentDetected': '生成的图片包含敏感内容，已被系统拦截',
    # Parameter errors
    'InvalidParameter': '请求参数无效，请检查输入内容',
    'InvalidImage': '图片格式或尺寸不符合要求，请检查后重试',
    'InvalidVideo': '视频格式或尺寸不符合要求，请检查后重试',
    'InvalidAudio': '音频格式不符合要求，请检查后重试',
    # Rate limit
    'RateLimitExceeded': '请求过于频繁，请稍后重试',
    'ConcurrencyLimitExceeded': '当前生成任务过多，请稍后重试',
    # Account & billing
    'InsufficientBalance': '平台账户余额不足，请联系管理员',
    # Asset errors
    'AssetNotFound': '引用的素材不存在或已被删除，请检查素材库',
    # Server errors
    'ServerOverloaded': '服务器繁忙，请稍后重试',
    'InternalError': '视频生成服务异常，请稍后重试',
    'Timeout': '生成超时，请重试',
}

# 关键词匹配：API 返回的 message 中包含这些关键词时，映射为对应中文提示
_MESSAGE_KEYWORDS = {
    'face': '检测到真实人脸，请使用虚拟人像素材替代真人照片',
    'privacy': '检测到真实人脸，请使用虚拟人像素材替代真人照片',
    'sensitive': '内容包含敏感信息，请修改后重试',
    'not found': '引用的素材不存在或已被删除，请检查素材库',
    'not valid': '请求参数无效，请检查输入内容',
}


class AirDramaAPIError(Exception):
    """Raised when video generation API returns an error response."""
    def __init__(self, code, message, status_code=400):
        self.code = code
        self.api_message = message
        self.status_code = status_code
        # 1. 精确匹配 error code
        friendly = ERROR_MESSAGES.get(code)
        if not friendly:
            # 2. 关键词匹配 message 内容
            msg_lower = (message or '').lower()
            for keyword, hint in _MESSAGE_KEYWORDS.items():
                if keyword in msg_lower:
                    friendly = hint
                    break
        self.user_message = friendly or '生成失败，请重试'
        super().__init__(self.user_message)


MODEL_MAP = {
    'seedance_2.0': 'doubao-seedance-2-0-260128',
    'seedance_2.0_fast': 'doubao-seedance-2-0-fast-260128',
}


def _headers():
    return {
        'Content-Type': 'application/json',
        'Authorization': f'Bearer {settings.ARK_API_KEY}',
    }


def create_task(prompt, model, content_items, aspect_ratio, duration,
                generate_audio=True, search_mode='off'):
    """Create a video generation task.

    Args:
        prompt: Text prompt for video generation.
        model: Model key ('airdrama' or 'airdrama_fast').
        content_items: List of media content dicts (image_url, video_url, audio_url).
        aspect_ratio: Video aspect ratio ('16:9', '9:16', etc.).
        duration: Video duration in seconds.
        generate_audio: Whether to generate audio with the video.
        search_mode: 'smart' to enable internet search, 'off' to disable.

    Returns:
        dict: API response with task id and status.
    """
    url = f'{settings.ARK_BASE_URL}/contents/generations/tasks'

    content = []
    if prompt:
        content.append({'type': 'text', 'text': prompt})
    content.extend(content_items)

    payload = {
        'model': MODEL_MAP.get(model, model),
        'content': content,
        'generate_audio': generate_audio,
        'ratio': aspect_ratio,
        'duration': duration,
        'watermark': False,
    }

    if search_mode and search_mode != 'off':
        payload['tools'] = [{'type': 'web_search'}]

    import logging
    logger = logging.getLogger(__name__)
    logger.info('AirDrama API payload: %s', {k: v for k, v in payload.items() if k != 'content'})

    resp = requests.post(url, json=payload, headers=_headers(), timeout=60)
    if resp.status_code != 200:
        # Extract human-readable error from API response
        try:
            err = resp.json().get('error', {})
            code = err.get('code', '')
            message = err.get('message', resp.text)
            logger.error('AirDrama API error: status=%s code=%s message=%s', resp.status_code, code, message)
        except Exception:
            code, message = '', resp.text
            logger.error('AirDrama API error: status=%s body=%s', resp.status_code, resp.text)
        raise AirDramaAPIError(code, message, resp.status_code)
    return resp.json()


def query_task(task_id):
    """Query a video generation task by its ARK task ID.

    Returns:
        dict: Task status, content (video URL on success), usage info.
    """
    url = f'{settings.ARK_BASE_URL}/contents/generations/tasks/{task_id}'
    resp = requests.get(url, headers=_headers(), timeout=30)
    resp.raise_for_status()
    return resp.json()


def extract_video_url(task_response):
    """Extract the video URL from a completed task response."""
    content = task_response.get('content')
    if not content:
        return None

    # content could be a list of items or a dict
    if isinstance(content, list):
        for item in content:
            if item.get('type') == 'video_url':
                return item.get('video_url', {}).get('url')
    elif isinstance(content, dict):
        if 'video_url' in content:
            url = content['video_url']
            return url.get('url') if isinstance(url, dict) else url

    return None


def map_status(ark_status):
    """Map ARK task status to our DB status."""
    mapping = {
        'running': 'processing',
        'submitted': 'queued',
        'queued': 'queued',
        'succeeded': 'completed',
        'failed': 'failed',
    }
    return mapping.get(ark_status, 'processing')