from __future__ import annotations import os import uuid from dataclasses import dataclass from datetime import datetime, timezone from app.config import settings from app.models.chat_generation_task import ChatGenerationTask from app.services.image_gen import download_image from app.services.provider_limit import provider_limit from app.services.resource_accounting_service import safe_file_size from app.services.video_cover_service import create_video_cover_for_local_video from app.services.video_gen import download_video @dataclass(slots=True) class DownloadedGenerationResult: url: str storage_path: str | None file_size_bytes: int resource_type: str storage_type: str = "local" cover_url: str | None = None cover_storage_path: str | None = None def _to_aware_utc(value: datetime | None) -> datetime | None: if value is None: return None if value.tzinfo is None: return value.replace(tzinfo=timezone.utc) return value.astimezone(timezone.utc) def _build_storage_date_dir(record: ChatGenerationTask) -> str: fixed = (getattr(record, "download_storage_date_dir", None) or "").strip().strip("/") if fixed: return fixed created_at = _to_aware_utc(getattr(record, "created_at", None)) or datetime.now(timezone.utc) return created_at.strftime("%Y/%m/%d") def _make_part_path(final_path: str) -> str: return f"{final_path}.{uuid.uuid4().hex}.part" def _is_valid_file(path: str | None) -> bool: if not path: return False try: return os.path.isfile(path) and os.path.getsize(path) > 0 except OSError: return False def _safe_remove(path: str | None) -> None: if not path: return try: if os.path.exists(path): os.remove(path) except OSError: pass async def _download_image_atomically(remote_url: str, final_path: str) -> str: if _is_valid_file(final_path): return final_path os.makedirs(os.path.dirname(final_path), exist_ok=True) part_path = _make_part_path(final_path) try: await download_image(remote_url, part_path) if not _is_valid_file(part_path): raise RuntimeError("图片下载完成但临时文件为空") os.replace(part_path, final_path) return final_path except Exception: _safe_remove(part_path) raise async def _download_video_atomically(remote_url: str, final_path: str) -> str: if _is_valid_file(final_path): return final_path os.makedirs(os.path.dirname(final_path), exist_ok=True) part_path = _make_part_path(final_path) try: await download_video(remote_url, part_path) if not _is_valid_file(part_path): raise RuntimeError("视频下载完成但临时文件为空") os.replace(part_path, final_path) return final_path except Exception: _safe_remove(part_path) raise async def download_generation_result(record: ChatGenerationTask) -> DownloadedGenerationResult: if not record.remote_result_url: raise ValueError("缺少远程结果URL") date_dir = _build_storage_date_dir(record) if record.gen_type == "image": dest_dir = os.path.join(settings.STORAGE_IMAGE_LOCAL_PATH, date_dir) os.makedirs(dest_dir, exist_ok=True) dest = os.path.join(dest_dir, f"{record.id}.png") async with provider_limit("result_download", settings.RESULT_DOWNLOAD_MAX_CONCURRENCY): await _download_image_atomically(record.remote_result_url if record.remote_result_url else "", dest) return DownloadedGenerationResult( url=f"/generate/images/{date_dir}/{record.id}.png", storage_path=dest, file_size_bytes=safe_file_size(dest), resource_type="image", ) dest_dir = os.path.join(settings.STORAGE_LOCAL_PATH, date_dir) os.makedirs(dest_dir, exist_ok=True) dest = os.path.join(dest_dir, f"{record.id}.mp4") async with provider_limit("result_download", settings.RESULT_DOWNLOAD_MAX_CONCURRENCY): await _download_video_atomically(record.remote_result_url if record.remote_result_url else "", dest) cover_url, cover_storage_path = create_video_cover_for_local_video( record_id=record.id, video_path=dest, date_dir=date_dir, log_prefix=f"ChatGenerationTask视频封面生成 task_id={record.id}", ) return DownloadedGenerationResult( url=f"/generate/videos/{date_dir}/{record.id}.mp4", storage_path=dest, file_size_bytes=safe_file_size(dest), resource_type="video", cover_url=cover_url, cover_storage_path=cover_storage_path, )