video-gen/video-gen-api/app/services/video_gen.py

import json
import logging
import os
import time

import httpx
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from volcenginesdkarkruntime import AsyncArk

from app.config import settings
from app.enums.private_portrait import PRIVATE_PORTRAIT_ASSET_URI_PREFIX
from app.models.video_engine import VideoEngine
from app.services.operation_log_service import build_exception_detail, log_ai_model_event
from app.utils.id_gen import generate_id
from app.types.generation.provider import (
    ProviderGenerationRecordLike,
    ProviderVideoEngineLike,
)

logger = logging.getLogger("videogen")



def _provider_log_context(engine, record, *, call_id: str, step_code: str) -> dict:
    generation_mode = str(getattr(record, "generation_mode", "") or "generation_record")
    owner_type = "chat_generation_task" if generation_mode != "generation_record" else "generation_record"
    return {
        "module": generation_mode,
        "step_code": step_code,
        "call_id": call_id,
        "source": "app.services.video_gen",
        "user_id": str(getattr(record, "user_id", "") or "") or None,
        "project_id": str(getattr(record, "project_id", "") or "") or None,
        "task_id": str(getattr(record, "id", "") or "") or None,
        "owner_type": owner_type,
        "owner_id": str(getattr(record, "id", "") or "") or None,
        "generation_attempt_no": int(getattr(record, "generation_attempt_no", 1) or 1),
        "model_config_id": str(getattr(engine, "id", "") or "") or None,
        "model_config_name": str(getattr(engine, "name", "") or "") or None,
        "model_name": str(getattr(engine, "model_name", "") or "") or None,
        "provider": str(getattr(engine, "provider", "") or "") or None,
        "api_base": str(getattr(engine, "api_base", "") or "") or None,
    }


async def get_active_engine(db: AsyncSession) -> VideoEngine:
    """Get the active video engine with highest priority."""
    result = await db.execute(
        select(VideoEngine)
        .where(VideoEngine.is_active == True, VideoEngine.deleted_at.is_(None))
        .order_by(VideoEngine.priority.desc())
        .limit(1)
    )
    engine = result.scalar_one_or_none()
    if not engine:
        raise ValueError("没有可用的视频引擎，请联系管理员配置")
    return engine


def _resolve_url(url: str) -> str:
    """Convert local path to base64 data URI, pass through remote URLs."""
    # if url.startswith("http"):
    #     return url
    # # Local file: read and encode as base64 data URI
    # file_path = os.path.join(settings.UPLOAD_LOCAL_PATH, url.replace("/uploads/", ""))
    # if not os.path.exists(file_path):
    #     return url
    # mime = mimetypes.guess_type(file_path)[0] or "application/octet-stream"
    # with open(file_path, "rb") as f:
    #     b64 = base64.b64encode(f.read()).decode()
    # return f"data:{mime};base64,{b64}"
    if url.startswith(("http://", "https://", "data:", PRIVATE_PORTRAIT_ASSET_URI_PREFIX)):
        return url
    return f"{settings.BASE_URL.rstrip('/')}/{url.lstrip('/')}"


async def submit_video_task(
    db: AsyncSession,
    engine: ProviderVideoEngineLike,
    record: ProviderGenerationRecordLike,
    *,
    include_media_references: bool,
) -> str:
    """Submit a video generation task via Ark SDK. Returns task_id."""
    client = AsyncArk(
        base_url=engine.api_base,
        api_key=engine.api_key,
    )

    content = [{"type": "text", "text": record.optimized_prompt or record.original_prompt}]

    # GenerationRecord 的附件已在 API 提词优化阶段参与过 optimized_prompt 生成；
    # ChatGenerationTask 不走 API 提词优化，所以创建供应商任务时仍需携带附件。
    if include_media_references and record.media_references:
        try:
            refs = json.loads(record.media_references)
            for ref in refs:
                ref_type = ref.get("type")
                ref_url = ref.get("url", "")
                ref_role = ref.get("role")
                if ref_type == "image" and ref_url:
                    resolved = _resolve_url(ref_url)
                    role = ref_role if ref_role in ("first_frame", "last_frame") else "reference_image"
                    content.append({"type": "image_url", "image_url": {"url": resolved}, "role": role})
                elif ref_type == "video" and ref_url:
                    resolved = _resolve_url(ref_url)
                    role = ref_role if ref_role else "reference_video"
                    content.append({"type": "video_url", "video_url": {"url": resolved}, "role": role})
                elif ref_type == "audio" and ref_url:
                    resolved = _resolve_url(ref_url)
                    role = ref_role if ref_role else "reference_audio"
                    content.append({"type": "audio_url", "audio_url": {"url": resolved}, "role": role})
        except (json.JSONDecodeError, TypeError):
            pass

    request_payload = {
        "model": engine.model_name,
        "content": content,
        "ratio": record.aspect_ratio,
        "duration": record.duration,
        "resolution": getattr(record, "provider_generation_resolution", None) or record.resolution,
        "generate_audio": True,
        "watermark": False,
    }

    call_id = generate_id()
    started = time.perf_counter()
    log_context = _provider_log_context(
        engine,
        record,
        call_id=call_id,
        step_code="video_create",
    )
    log_ai_model_event(
        event_type="REQUEST",
        event_phase="REQUEST",
        event_status="started",
        remote_action="video_create",
        request={**request_payload, "include_media_references": include_media_references},
        **log_context,
    )

    try:
        result = await client.content_generation.tasks.create(**request_payload)
        task_id = result.id
        log_ai_model_event(
            event_type="RESPONSE",
            event_phase="RESPONSE",
            event_status="success",
            remote_action="video_create",
            remote_request_id=task_id,
            latency_ms=int((time.perf_counter() - started) * 1000),
            response={"task_id": task_id},
            **log_context,
        )
    except Exception as exc:
        log_ai_model_event(
            event_type="ERROR",
            event_phase="ERROR",
            event_status="failed",
            remote_action="video_create",
            latency_ms=int((time.perf_counter() - started) * 1000),
            detail=build_exception_detail(exc),
            error=str(exc),
            **log_context,
        )
        raise
    finally:
        await client.close()

    return task_id


async def poll_task_status(engine: VideoEngine, task_id: str) -> dict:
    """Query task status via Ark SDK. Returns {status, video_url, response_data}."""
    client = AsyncArk(
        base_url=engine.api_base,
        api_key=engine.api_key,
    )

    result = await client.content_generation.tasks.get(task_id=task_id)
    await client.close()

    # Serialize response
    response_dict = {
        "id": result.id,
        "model": result.model,
        "status": result.status,
        "created_at": result.created_at,
        "updated_at": result.updated_at,
    }

    video_url = None
    video_tokens = 0
    if result.status == "succeeded" and result.content:
        video_url = getattr(result.content, "video_url", None)
        response_dict["video_url"] = video_url
        response_dict["duration"] = getattr(result, "duration", None)
        response_dict["ratio"] = getattr(result, "ratio", None)
        response_dict["resolution"] = getattr(result, "resolution", None)
        # Extract usage info if present
        usage = getattr(result, "usage", None)
        if usage:
            response_dict["usage"] = {
                "input_tokens": getattr(usage, "input_tokens", 0),
                "output_tokens": getattr(usage, "output_tokens", 0),
                "total_tokens": getattr(usage, "total_tokens", 0),
            }
            video_tokens = getattr(usage, "total_tokens", 0)
    elif result.status == "failed":
        response_dict["error"] = str(getattr(result, "error", "视频生成失败"))

    return {
        "status": result.status,
        "video_url": video_url,
        "video_tokens": video_tokens,
        "response_data": json.dumps(response_dict, ensure_ascii=False, default=str),
        "error": response_dict.get("error"),
    }


async def download_video(
    video_url: str,
    dest_path: str,
    *,
    execution_guard=None,
) -> str:
    """Download video to local storage."""
    import os

    os.makedirs(os.path.dirname(dest_path), exist_ok=True)

    async with httpx.AsyncClient(timeout=300) as client:
        async with client.stream("GET", video_url) as response:
            response.raise_for_status()
            with open(dest_path, "wb") as f:
                chunk_no = 0
                async for chunk in response.aiter_bytes(chunk_size=8192):
                    f.write(chunk)
                    chunk_no += 1
                    if execution_guard is not None and chunk_no % 32 == 0:
                        await execution_guard()
    if execution_guard is not None:
        await execution_guard()
    return dest_path