237 lines
8.6 KiB
Python
237 lines
8.6 KiB
Python
import base64
|
|
import json
|
|
import logging
|
|
import mimetypes
|
|
import os
|
|
from datetime import datetime, timezone
|
|
|
|
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.log_config import is_enabled, LOG_DIR, LOG_DATE_FORMAT, encrypt_data
|
|
from app.types.generation.provider import (
|
|
ProviderGenerationRecordLike,
|
|
ProviderVideoEngineLike,
|
|
)
|
|
|
|
logger = logging.getLogger("videogen")
|
|
|
|
|
|
def _log_video_request(engine: ProviderVideoEngineLike, record_id: str, request_data: dict):
|
|
"""Log video generation request to log/AiModel/YYYY-MM-DD.log"""
|
|
if not is_enabled():
|
|
return
|
|
try:
|
|
os.makedirs(LOG_DIR, exist_ok=True)
|
|
today = datetime.now().strftime(LOG_DATE_FORMAT)
|
|
log_file = os.path.join(LOG_DIR, f"{today}.log")
|
|
request_str = json.dumps(request_data, ensure_ascii=False)
|
|
request_encrypted = encrypt_data(request_data, True)
|
|
entry = {
|
|
"timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
|
|
"type": "video_gen_request",
|
|
"engine": engine.name,
|
|
"model": engine.model_name,
|
|
"record_id": record_id,
|
|
"request": request_encrypted,
|
|
"request_length": len(request_str),
|
|
}
|
|
with open(log_file, "a", encoding="utf-8") as f:
|
|
f.write(json.dumps(entry, ensure_ascii=False) + "\n")
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def _log_video_response(record_id: str, response_data: dict, error: str | None = None):
|
|
"""Log video generation response to log/AiModel/YYYY-MM-DD.log"""
|
|
if not is_enabled():
|
|
return
|
|
try:
|
|
os.makedirs(LOG_DIR, exist_ok=True)
|
|
today = datetime.now().strftime(LOG_DATE_FORMAT)
|
|
log_file = os.path.join(LOG_DIR, f"{today}.log")
|
|
response_encrypted = encrypt_data(response_data, True) if response_data else ""
|
|
|
|
entry = {
|
|
"timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
|
|
"type": "video_gen_response",
|
|
"record_id": record_id,
|
|
"response": response_encrypted,
|
|
"error": error,
|
|
}
|
|
with open(log_file, "a", encoding="utf-8") as f:
|
|
f.write(json.dumps(entry, ensure_ascii=False) + "\n")
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
|
|
|
|
|
|
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)
|
|
.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,
|
|
}
|
|
|
|
# Log request to AiModel log. include_media_references 只用于排查日志，不传给供应商 API。
|
|
_log_video_request(
|
|
engine,
|
|
record.id,
|
|
{**request_payload, "include_media_references": include_media_references},
|
|
)
|
|
|
|
try:
|
|
result = await client.content_generation.tasks.create(**request_payload)
|
|
task_id = result.id
|
|
_log_video_response(record.id, {"task_id": task_id})
|
|
except Exception as e:
|
|
_log_video_response(record.id, {}, str(e))
|
|
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) -> 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:
|
|
async for chunk in response.aiter_bytes(chunk_size=8192):
|
|
f.write(chunk)
|
|
return dest_path
|