dsProject/dsLightRag/Util/ObsUtil.py

# pip install esdk-obs-python --trusted-host pypi.org
import traceback
import logging
import base64  # 添加base64模块导入
from obs import ObsClient
from obs import PutObjectHeader

from Config.Config import OBS_AK, OBS_SK, OBS_SERVER, OBS_BUCKET

# 配置日志
logger = logging.getLogger(__name__)


class ObsUploader:
    def __init__(self, access_key_id=OBS_AK, secret_access_key=OBS_SK, server="https://" + OBS_SERVER):
        """
        初始化OBS客户端
        :param access_key_id: 访问密钥ID
        :param secret_access_key: 秘密访问密钥
        :param server: OBS服务器地址
        """
        self.obs_client = ObsClient(
            access_key_id=access_key_id,
            secret_access_key=secret_access_key,
            server=server
        )

    def upload_file(self, object_key, file_path, bucket_name=OBS_BUCKET, metadata=None, headers=None):
        """
        上传文件到OBS
        :param bucket_name: 桶名称
        :param object_key: 对象键（上传后的文件名）
        :param file_path: 本地文件路径
        :param metadata: 元数据
        :param headers: 请求头
        :return: 成功返回True和响应信息，失败返回False和错误信息
        """
        try:
            # 初始化默认参数
            if metadata is None:
                metadata = {}
            if headers is None:
                headers = PutObjectHeader()

            # 执行文件上传
            resp = self.obs_client.putFile(
                bucket_name,
                object_key,
                file_path,
                metadata,
                headers
            )

            # 检查响应状态
            if resp.status < 300:
                logger.info(f'文件上传成功: {file_path} -> {bucket_name}/{object_key}')
                logger.info(f'requestId: {resp.requestId}')
                logger.info(f'etag: {resp.body.etag}')
                return True, {
                    'requestId': resp.requestId,
                    'etag': resp.body.etag,
                    'versionId': resp.body.versionId,
                    'storageClass': resp.body.storageClass
                }
            else:
                error_msg = f'文件上传失败: {resp.errorCode} - {resp.errorMessage}'
                logger.error(error_msg)
                logger.error(f'requestId: {resp.requestId}')
                return False, {
                    'requestId': resp.requestId,
                    'errorCode': resp.errorCode,
                    'errorMessage': resp.errorMessage
                }
        except Exception as e:
            error_msg = f'文件上传异常: {str(e)}'
            logger.error(error_msg)
            logger.error(traceback.format_exc())
            return False, {'error': str(e), 'traceback': traceback.format_exc()}

    def upload_base64_image(self, object_key, base64_data, bucket_name=OBS_BUCKET, metadata=None, headers=None):
        """
        上传Base64编码的图片或原始二进制图片数据到OBS
        :param bucket_name: 桶名称
        :param object_key: 对象键（上传后的文件名）
        :param base64_data: Base64编码的图片字符串、字节形式或原始二进制图片数据
        :param metadata: 元数据
        :param headers: 请求头
        :return: 成功返回True和响应信息，失败返回False和错误信息
        """
        try:
            # 初始化默认参数
            if metadata is None:
                metadata = {}
            if headers is None:
                headers = PutObjectHeader()

            # 处理数据
            if isinstance(base64_data, bytes):
                try:
                    # 尝试解码为UTF-8字符串（假设是Base64编码的字符串的字节形式）
                    base64_string = base64_data.decode('utf-8')
                    # 移除Base64前缀（如果有）
                    if 'base64,' in base64_string:
                        base64_string = base64_string.split('base64,')[1]
                    # 解码Base64数据
                    image_data = base64.b64decode(base64_string)
                except UnicodeDecodeError:
                    # 解码失败，假设是原始二进制图片数据
                    image_data = base64_data
            else:
                # 假设是字符串
                base64_string = str(base64_data)
                # 移除Base64前缀（如果有）
                if 'base64,' in base64_string:
                    base64_string = base64_string.split('base64,')[1]
                # 解码Base64数据
                image_data = base64.b64decode(base64_string)

            # 执行对象上传
            resp = self.obs_client.putObject(
                bucket_name,
                object_key,
                content=image_data,
                metadata=metadata,
                headers=headers
            )

            # 检查响应状态
            if resp.status < 300:
                logger.info(f'图片上传成功: -> {bucket_name}/{object_key}')
                logger.info(f'requestId: {resp.requestId}')
                logger.info(f'etag: {resp.body.etag}')
                return True, {
                    'requestId': resp.requestId,
                    'etag': resp.body.etag,
                    'versionId': resp.body.versionId,
                    'storageClass': resp.body.storageClass
                }
            else:
                error_msg = f'图片上传失败: {resp.errorCode} - {resp.errorMessage}'
                logger.error(error_msg)
                logger.error(f'requestId: {resp.requestId}')
                return False, {
                    'requestId': resp.requestId,
                    'errorCode': resp.errorCode,
                    'errorMessage': resp.errorMessage
                }
        except Exception as e:
            error_msg = f'图片上传异常: {str(e)}'
            logger.error(error_msg)
            logger.error(traceback.format_exc())
            return False, {'error': str(e), 'traceback': traceback.format_exc()}