ai-flow-python/app/util/download_util.py

import logging
import os
from pathlib import Path
from urllib.parse import urlparse

import httpx

logger = logging.getLogger(__name__)


def download_file(url: str, target_dir: str) -> str | None:
    if not url or not target_dir:
        logger.error("URL or target directory is blank. URL: %s, TargetDir: %s", url, target_dir)
        return None

    try:
        os.makedirs(target_dir, exist_ok=True)
        filename = _get_filename_from_url(url)
        if not filename:
            filename = f"downloaded_file_{os.getpid()}"

        filepath = os.path.join(target_dir, filename)
        if os.path.isfile(filepath):
            logger.info("File already exists, skipping download: %s", filepath)
            return filepath

        logger.info("Starting download from: %s to: %s", url, filepath)
        with httpx.Client(timeout=60, follow_redirects=True) as client:
            resp = client.get(url)
            resp.raise_for_status()
            with open(filepath, "wb") as f:
                f.write(resp.content)

        if os.path.isfile(filepath) and os.path.getsize(filepath) > 0:
            logger.info("Successfully downloaded file: %s (%d bytes)", filepath, os.path.getsize(filepath))
            return filepath

        logger.error("Download failed: file not found or empty after download")
        return None
    except Exception as e:
        logger.error("Failed to download file from %s: %s", url, e)
        return None


def download_files(urls: list[str], workspace: str, prefix_dir: str | None = None) -> list[str]:
    downloaded: list[str] = []
    if not urls:
        return downloaded

    target_dir = workspace + ("/" + prefix_dir if prefix_dir else "")
    logger.info("Starting batch download of %d files to directory: %s", len(urls), target_dir)

    for url in urls:
        try:
            local_path = download_file(url, target_dir)
            if local_path:
                rel = get_relative_path(local_path, workspace)
                downloaded.append(rel)
        except Exception as e:
            logger.error("Failed to download file from URL: %s, %s", url, e)

    logger.info("Batch download completed. Successfully downloaded %d/%d files", len(downloaded), len(urls))
    return downloaded


def get_relative_path(absolute_path: str, base_dir: str) -> str:
    try:
        return os.path.relpath(absolute_path, base_dir).replace("\\", "/")
    except Exception:
        return os.path.basename(absolute_path)


def _get_filename_from_url(url: str) -> str | None:
    try:
        path = urlparse(url).path
        if path:
            name = path.rsplit("/", 1)[-1]
            if name:
                return name
    except Exception:
        pass
    return None