53 lines
1.9 KiB
Python
53 lines
1.9 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
from dataclasses import dataclass
|
|
from datetime import datetime
|
|
|
|
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_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"
|
|
|
|
|
|
async def download_generation_result(record: ChatGenerationTask) -> DownloadedGenerationResult:
|
|
if not record.remote_result_url:
|
|
raise ValueError("缺少远程结果URL")
|
|
|
|
date_dir = datetime.now().strftime("%Y/%m/%d")
|
|
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(record.remote_result_url, 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(record.remote_result_url, dest)
|
|
return DownloadedGenerationResult(
|
|
url=f"/generate/videos/{date_dir}/{record.id}.mp4",
|
|
storage_path=dest,
|
|
file_size_bytes=safe_file_size(dest),
|
|
resource_type="video",
|
|
)
|