AutoAndroidController/py/venv/Lib/site-packages/aistudio_sdk/requests/dataset.py

# !/usr/bin/env python3
# -*- coding: UTF-8 -*-
################################################################################
#
# Copyright (c) 2024 Baidu.com, Inc. All Rights Reserved
#
################################################################################
"""
本文件实现了请求产线任务

Authors: suoyi@baidu.com
Date:    2024/7/20
"""
import json
import requests
from aistudio_sdk import config, log
from baidubce.bce_client_configuration import BceClientConfiguration
from baidubce.auth.bce_credentials import BceCredentials
from baidubce.services.bos.bos_client import BosClient
import os
import threading
from concurrent.futures import ThreadPoolExecutor
from tqdm import tqdm
from urllib.parse import urljoin
from pathlib import Path
import re


class RequestDatasetException(Exception):
    """
    exception for requesting dataset server
    """
    pass
MAX_WORKERS_FILE = os.path.expanduser("~/.download_max_workers")

# 默认线程数
DEFAULT_MAX_WORKERS = 6


def get_max_workers():
    """max download worker"""
    try:
        with open(MAX_WORKERS_FILE, 'r') as f:
            return int(f.read().strip())
    except (Exception) as e:
        return DEFAULT_MAX_WORKERS


def post_request_get_file_ids(url, datasetId):
    """file info"""
    data = {"datasetId": datasetId}
    response = requests.post(url, data=data)
    response.raise_for_status()
    result = response.json().get("result", {})
    file_ids = result.get("fileIds", [])
    return file_ids


def load_token():
    """
    load
    """
    if not os.path.exists(config.TOKEN_FILE):
        return None
    with open(config.TOKEN_FILE, 'r') as f:
        return f.read().strip()


def _header_fill(params=None, token=''):
    """
    填充header
    """
    if token:
        auth = f'{token}'
    else:
        auth = f'{os.getenv("AISTUDIO_ACCESS_TOKEN", default="")}'
    headers = {
        'Content-Type': 'application/json',
        'Authorization': auth
    }
    if params:
        headers.update(params)
    return headers


def get_file_url(host, datasetId, fileId):
    """get url"""
    path = f"/llm/files/datasets/{datasetId}/file/{fileId}/download"
    url = urljoin(host, path)
    token = load_token()
    print(token)
    if token is not None:
        headers = _header_fill(token=token)
    else:
        headers = _header_fill()
    response = requests.get(url, headers=headers)
    response.raise_for_status()
    print(response.json())
    return response.json()["result"]["fileUrl"]

CHUNK_SIZE = 160 * 1024 * 1024  # 160MB


def parse_filename_from_cd(cd_header):
    """filename"""
    if not cd_header:
        return None
    fname = re.findall('filename="?([^";]+)"?', cd_header)
    return fname[0] if fname else None


def get_file_info(file_url):
    """获取文件大小和文件名"""
    r = requests.head(file_url, allow_redirects=True)
    r.raise_for_status()
    file_size = int(r.headers.get('Content-Length', 0))
    cd = r.headers.get("Content-Disposition", "")
    filename = parse_filename_from_cd(cd)
    if not filename:
        filename = os.path.basename(file_url.split("?")[0])
    return file_size, filename


def download_chunk(file_url, start, end, local_path, pbar, lock):
    """download"""
    headers = {'Range': f"bytes={start}-{end}"}
    response = requests.get(file_url, headers=headers, stream=True)
    response.raise_for_status()

    with open(local_path, 'rb+') as f:
        f.seek(start)
        for chunk in response.iter_content(chunk_size=8192):
            if chunk:
                f.write(chunk)
                with lock:
                    pbar.update(len(chunk))


def download_file_multithreaded(file_url, local_dir, max_workers=None):
    """multi thread"""
    if max_workers is None:
        max_workers = get_max_workers()

    # Step 1: Get file size and filename
    file_size, filename = get_file_info(file_url)
    local_path = os.path.join(local_dir, filename)
    os.makedirs(local_dir, exist_ok=True)

    # Step 2: Create empty file if not exists
    if not os.path.exists(local_path):
        with open(local_path, 'wb') as f:
            f.truncate(file_size)

    # Step 3: Calculate chunks
    chunks = []
    for i in range(0, file_size, CHUNK_SIZE):
        start = i
        end = min(i + CHUNK_SIZE - 1, file_size - 1)
        chunks.append((start, end))

    # Step 4: Prepare download tasks
    from threading import Lock
    pbar = tqdm(total=file_size, unit='B', unit_scale=True, desc=filename)
    lock = Lock()

    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = []
        for start, end in chunks:
            # 检查是否已下载该块
            if os.path.exists(local_path):
                current_size = os.path.getsize(local_path)
                if current_size >= end + 1:
                    pbar.update(end - start + 1)
                    continue
            futures.append(
                executor.submit(download_chunk, file_url, start, end, local_path, pbar, lock)
            )
        for f in futures:
            f.result()
    pbar.close()


def download_datasets(datasetId, local_dir=None):
    """old dataset"""
    if local_dir is None:
        local_dir = os.getenv("HOME")
    host = os.getenv("STUDIO_GIT_HOST", default=config.STUDIO_MODEL_API_URL_PREFIX_DEFAULT)
    url = f"{host}/studio/dataset/detail"
    download_all_files(url, host, datasetId, local_dir)


def download_all_files(url, host, datasetId, localDir):
    """
    all
    """
    file_ids = post_request_get_file_ids(url, datasetId)
    os.makedirs(localDir, exist_ok=True)

    tasks = []
    pbar_lock = threading.Lock()
    with ThreadPoolExecutor(max_workers=4) as executor:
        for fileId in file_ids:
            file_url = get_file_url(host, datasetId, fileId)
            tasks.append(executor.submit(download_file_multithreaded, file_url, localDir, pbar_lock))

    for task in tasks:
        task.result()

def bos_acl_dataset_file(
        token: str,
        bucket_name=None
    ):
    """
    申请ak/sk
    response:
    {
        "logId": "",
        "errorCode": 0,
        "errorMsg": "",
        "timestamp": 0,
        "result": {
            "accessKeyId": "",
            "secretAccessKey": "",
            "sessionToken": "",
            "fileKey": "",
            "serverTime": 0,
            "expiresIn": 0,
            "endpoint": "",
            "bucketName": ""
        }
    }
    """
    url = f"{config.STUDIO_MODEL_API_URL_PREFIX_DEFAULT}/llm/files/acl"
    headers = {
        "Authorization": f"{token}",
        "Content-Type": "application/json"
    }
    params = {}
    if bucket_name:
        params["bucketName"] = bucket_name

    response = requests.get(url, headers=headers, params=params)
    if response.status_code == 200:
        return response.json()
    else:
        raise RequestDatasetException(f"Failed to get bos acl: {response.text}")

def add_file_with_retry(token: str, file_origin_name: str, file_key: str, bucket_name=None, file_abs=None):
    """
    上传文件到指定的bucket，并返回文件ID。
    """
    for i in range(3):
        try:
            file_id = add_file(token, file_origin_name, file_key, bucket_name, file_abs)
            return file_id
        except RequestDatasetException as e:
            log.error(f"add file 失败，重试第{i+1}次")
            log.error(e)


def add_file(token: str, file_origin_name: str, file_key: str, bucket_name=None, file_abs=None):
    """
    上传文件到指定的bucket，并返回文件ID。

    Args:
        token (str): 认证token。
        file_origin_name (str): 文件的原始名称。
        file_key (str): 文件在存储中的键值。
        bucket_name (str, optional): 如果提供，则上传到此bucket，否则使用默认bucket。
        file_abs (str, optional): 文件的绝对路径，可选。

    Returns:
        dict: 包含操作结果的字典，其中包括logId, errorCode, errorMsg, timestamp和result（包含fileId）。

    Raises:
        HTTPError: 如果请求失败，抛出异常。
    """
    log.debug("add file..")
    url = f"{config.STUDIO_MODEL_API_URL_PREFIX_DEFAULT}/llm/files/addfile"
    headers = {
        "Authorization": f"{token}",
        "Content-Type": "application/json"
    }
    data = {
        "fileOriginName": file_origin_name,
        "fileKey": file_key,
        "bucketName": bucket_name,
        "fileAbs": file_abs
    }
    response = requests.post(url, headers=headers, json=data)

    if response.status_code == 200:
        if response.json().get("errorCode") == 0:
            log.debug(f"add file success")
            result = response.json()
            file_id = result.get("result", {}).get("fileId")
            return file_id
        else:
            log.error("落库失败")
            log.error(f"add file failed, response: {data} {response.text}")
            return None
    else:
        raise RequestDatasetException(f"Failed to add file: {response.text}")

def create_dataset_with_retry(token: str, dataset_name: str, file_ids: list,
                              dataset_type=1, dataset_abs="", dataset_license=1):
    """
    创建一个新的数据集，并返回数据集ID。
    """
    for i in range(3):
        try:
            dataset_id = create_dataset(token, dataset_name, file_ids, dataset_type, dataset_abs, dataset_license)
            return dataset_id
        except RequestDatasetException as e:
            log.error(f"create dataset 失败，重试第{i+1}次")
            log.error(e)


def create_dataset(token: str, dataset_name: str, file_ids: list, dataset_type=1, dataset_abs="", dataset_license=1):
    """
    创建一个新的数据集，并返回数据集ID。

    Args:
        token (str): 认证token。
        dataset_name (str): 数据集的名称。
        file_ids (list of int): 包含在数据集中的文件ID列表。
        dataset_type (int, optional): 数据集的类型，1 表示私有，2 表示公开。默认为0（私有）。
        dataset_abs (str, optional): 数据集的简介，可选。

    Returns:
        dict: 包含操作结果的字典，其中包括logId, errorCode, errorMsg, timestamp和result（包含datasetId）。
        None: 如果请求失败，返回None。

    """
    url = f"{config.STUDIO_MODEL_API_URL_PREFIX_DEFAULT}/llm/files/datasets"
    headers = {
        "Authorization": f"{token}",
        "Content-Type": "application/json"
    }

    data = {
        "datasetName": dataset_name,
        "datasetAbs": dataset_abs,
        "fileIds": file_ids,
        "datasetType": dataset_type,
        "protocolId": dataset_license
    }
    response = requests.post(url, headers=headers, json=data)


    if response.status_code == 200:
        log.debug(f"add file success")
        if response.json().get("errorCode") == 0:

            result = response.json()

            dataset_id = result.get("result", {}).get("datasetId")
            return dataset_id
        else:
            log.error(f"数据集创建失败:{response.json().get('errorMsg')}")
            log.debug(f"add file failed, response: {data} {response.text}")
            return None
    else:
        raise RequestDatasetException(f"Failed to create dataset: {response.text}")

def add_files_to_dataset_with_retry(token: str, dataset_id: int, file_ids: list):
    """
    向指定的数据集中添加文件。
    """
    for i in range(3):
        try:
            result = add_files_to_dataset(token, dataset_id, file_ids)
            return result
        except RequestDatasetException as e:
            log.error(f"add file to dataset 失败，重试第{i+1}次")
            log.error(e)

def add_files_to_dataset(token: str, dataset_id: int, file_ids: list):
    """
    向指定的数据集中添加文件。

    Args:
        token (str): 认证token。
        dataset_id (int): 数据集的ID。
        file_ids (list of int): 需要添加到数据集的文件ID列表。

    Returns:
        dict: 包含操作结果的字典，其中包括logId, errorCode, errorMsg, timestamp和result。
        None: 如果请求失败，返回None。

    """
    url = f"{config.STUDIO_MODEL_API_URL_PREFIX_DEFAULT}/llm/files/datasets/{dataset_id}/addfile"
    headers = {
        "Authorization": f"{token}",
        "Content-Type": "application/json"
    }

    data = {
        "fileIds": file_ids
    }
    response = requests.post(url, headers=headers, json=data)

    if response.status_code == 200:
        if response.json().get("errorCode") == 0:
            log.info(f"向数据集[{dataset_id}]中添加文件成功!")
            log.debug(f"向数据集[{dataset_id}]中添加文件成功[{file_ids}]")
            return response.json()
        else:
            log.error(f"添加文件失败: {response.json().get('errorMsg')}")
            log.debug(f"add file failed, response: {data} {response.text}")
            return None
    else:
        # log.error(f"Failed to add files to dataset: {response.text}")
        raise RequestDatasetException(f"Failed to add files to dataset: {response.text}")