chat-ai-api/app/core/oss_client.py

"""
阿里云 OSS 客户端
用于上传和管理租户的知识库文档
采用"单 Bucket + 租户前缀"的方式
"""
import oss2
import os
from typing import Optional, Tuple
from datetime import datetime


class OSSClient:
    def __init__(self):
        # 从环境变量读取阿里云 OSS 配置
        self.access_key_id = os.getenv("ALIYUN_OSS_ACCESS_KEY_ID")
        self.access_key_secret = os.getenv("ALIYUN_OSS_ACCESS_KEY_SECRET")
        self.endpoint = os.getenv("ALIYUN_OSS_ENDPOINT", "oss-cn-hangzhou.aliyuncs.com")
        self.bucket_name = os.getenv("ALIYUN_OSS_BUCKET", "chat-ai-knowledge")

        # 初始化 OSS 认证和 Bucket
        auth = oss2.Auth(self.access_key_id, self.access_key_secret)
        self.bucket = oss2.Bucket(auth, self.endpoint, self.bucket_name)

    def upload_file(self, tenant_id: str, file_content: bytes, filename: str) -> Tuple[bool, str]:
        """
        上传文件到 OSS
        :param tenant_id: 租户ID (如 com.xiaoyouhui)
        :param file_content: 文件内容（字节）
        :param filename: 原始文件名
        :return: (是否成功, OSS路径或错误信息)
        """
        try:
            # 构造 OSS 路径: {tenant_id}/documents/{timestamp}_{filename}
            timestamp = datetime.utcnow().strftime("%Y%m%d%H%M%S")
            oss_path = f"{tenant_id}/documents/{timestamp}_{filename}"

            # 上传文件
            result = self.bucket.put_object(oss_path, file_content)

            if result.status == 200:
                return True, oss_path
            else:
                return False, f"上传失败，状态码: {result.status}"
        except Exception as e:
            return False, f"上传异常: {str(e)}"

    def download_file(self, oss_path: str) -> Optional[bytes]:
        """
        从 OSS 下载文件
        :param oss_path: OSS 中的文件路径
        :return: 文件内容（字节）或 None
        """
        try:
            result = self.bucket.get_object(oss_path)
            return result.read()
        except Exception as e:
            print(f"下载文件失败: {str(e)}")
            return None

    def delete_file(self, oss_path: str) -> bool:
        """
        删除 OSS 中的文件
        :param oss_path: OSS 中的文件路径
        :return: 是否成功
        """
        try:
            self.bucket.delete_object(oss_path)
            return True
        except Exception as e:
            print(f"删除文件失败: {str(e)}")
            return False

    def get_file_url(self, oss_path: str, expires: int = 3600) -> Optional[str]:
        """
        生成文件的临时访问 URL
        :param oss_path: OSS 中的文件路径
        :param expires: 过期时间（秒），默认1小时
        :return: 临时访问URL
        """
        try:
            url = self.bucket.sign_url('GET', oss_path, expires)
            return url
        except Exception as e:
            print(f"生成URL失败: {str(e)}")
            return None


# 全局单例
_oss_client = None


def get_oss_client() -> OSSClient:
    """获取 OSS 客户端单例"""
    global _oss_client
    if _oss_client is None:
        _oss_client = OSSClient()
    return _oss_client