volcengine-python-sdk-aicc 0.0.41__py3-none-any.whl

bytedance/volcengine_aicc_sdk/auc_inference_client.py

@@ -0,0 +1,305 @@
+import uuid
+"""
+seedance机密推理客户端
+"""
+__all__ = ["AUCClient"]
+
+import os
+import base64
+import requests
+import traceback
+from urllib.parse import unquote
+from cryptography.hazmat.primitives import serialization as crypto_serialization
+from cryptography.hazmat.primitives.asymmetric import rsa
+from cryptography.hazmat.backends import default_backend as crypto_default_backend
+from bytedance.jeddak_secure_channel.crypto import PrivateKey
+from .secure_http_client import SecureHttpClient
+from .common.tos_utils import TOSUtils
+from .common.error import AICCException
+from .common.log import logger
+from .secure_http_client import AICCConfig
+from .common.aes import decrypt_file
+
+
+class AUCClient:
+    """
+    AUC机密推理客户端.
+    """
+
+    def __init__(self,
+                 base_url: str = "",
+                 api_key: str = "",
+                 endpoint: str = "",
+                 aicc_config: AICCConfig = None,
+                 **kwargs):
+        self.base_url = base_url
+        self.api_key = api_key
+        self.endpoint = endpoint
+        self.aicc_config = aicc_config
+        self.tos_utils = TOSUtils()
+        self.kwargs = kwargs
+        self.is_secure = self.kwargs.get("is_secure", True)
+
+        if aicc_config:
+            self.secure_http_client = SecureHttpClient(aicc_config=aicc_config)
+        else:
+            if self.is_secure:
+                aicc_config = AICCConfig(ra_url=base_url, auth_token=api_key)
+                self.secure_http_client = SecureHttpClient(aicc_config=aicc_config)
+            else:
+                self.secure_http_client = SecureHttpClient()
+
+        # 是否开启AUC文件加密, 默认加密
+        self.is_enable_auc_encrypt = str(self.kwargs.get("is_enable_auc_encrypt", True)).lower() == "true"
+
+        # AUC文件的存储位置, 目前仅支持火山TOS
+        self.auc_file_storage_location = self.kwargs.get("auc_file_storage_location", "TOS").upper()
+
+        self.auc_file_encrypt_pub_key_content = None
+        self.auc_file_encrypt_priv_key_content = None
+
+        self.auc_file_encrypt_priv_key: PrivateKey = None
+
+        self.timeout = self.kwargs.get("timeout", 120.0)
+
+        # API Endpoints
+        self.CREATE_TASK_URL = "/auc/lark/submit"
+        self.QUERY_TASK_URL = "/auc/lark/query"
+
+    def _build_common_headers(self, request_id: str) -> dict[str, str]:
+        if not request_id:
+            request_id = str(uuid.uuid4())
+
+        return {
+            "Content-Type": "application/json",
+            "X-Api-App-Key": self.endpoint,
+            "Authorization": f"Bearer {self.api_key}",
+            "X-Request-Id": request_id,
+        }
+
+    def _get_base_url(self):
+        """
+        Get base URL with trailing slash removed
+        """
+        return self.base_url.rstrip("/")
+
+    def generate_key_pair(self, pub_key_path: str = "", priv_key_path: str = ""):
+        """
+        Generate RSA 4096 key pair and save to files
+
+        Args:
+            pub_key_path: Path to save public key
+            priv_key_path: Path to save private key
+
+        Returns:
+            Tuple of (public_key_path, private_key_path)
+        """
+        try:
+            # Generate RSA key pair
+            key = rsa.generate_private_key(
+                backend=crypto_default_backend(),
+                public_exponent=65537,
+                key_size=4096
+            )
+
+            # Save private key
+            private_key = key.private_bytes(
+                crypto_serialization.Encoding.PEM,
+                crypto_serialization.PrivateFormat.PKCS8,
+                crypto_serialization.NoEncryption()
+            )
+            with open(priv_key_path, 'wb') as f:
+                f.write(private_key)
+
+            # Save public key
+            public_key = key.public_key().public_bytes(
+                crypto_serialization.Encoding.PEM,
+                crypto_serialization.PublicFormat.SubjectPublicKeyInfo
+            )
+            with open(pub_key_path, 'wb') as f:
+                f.write(public_key)
+        except Exception as e:
+            logger.error(f"[Error] Failed to generate key pair: {e}")
+            raise AICCException("Failed to generate RSA key pair") from e
+
+    def set_file_encrypt_key(self, pub_key_path: str = "", priv_key_path: str = ""):
+        """
+        Set auc file encrypt key
+        Args:
+            pub_key_path: Public key for encrypt
+            priv_key_path: Private key for decrypt
+        """
+        if os.path.exists(pub_key_path):
+            with open(pub_key_path, 'r') as f:
+                content = f.read()
+                if '-----BEGIN' in content and '-----END' in content:
+                    self.auc_file_encrypt_pub_key_content = content
+
+        if os.path.exists(priv_key_path):
+            with open(priv_key_path, 'r') as f:
+                content = f.read()
+                if '-----BEGIN' in content and '-----END' in content:
+                    self.auc_file_encrypt_priv_key_content = content
+                    self.auc_file_encrypt_priv_key = \
+                        PrivateKey.from_private_pem(self.auc_file_encrypt_priv_key_content)
+
+    def create_task(self, data: dict, custom_header: dict = None, request_id: str = "") -> str:
+        """
+        :param data:
+        :param custom_header:
+        :return:
+        """
+        if not data:
+            raise AICCException("data must be provided")
+
+        if "model" not in data:
+            data["model"] = self.endpoint
+
+        base_url = self._get_base_url()
+        url = f"{base_url}{self.CREATE_TASK_URL}"
+
+        headers = self._build_common_headers(request_id)
+        if custom_header:
+            headers.update(custom_header)
+
+        # Add encryption headers if provided
+        if self.is_enable_auc_encrypt:
+            if not self.auc_file_encrypt_pub_key_content or not self.auc_file_encrypt_priv_key_content:
+                raise AICCException("pub_key and priv_key must be provided when Is-Enable-Encrypt is 'true'")
+
+            headers["Enable-TOS-Content-Result-Encryption"] = str(self.is_enable_auc_encrypt).lower()
+            headers["X-Encryption-Algorithm"] = "RSA_OAEP_4096_AES_256"
+            if self.auc_file_encrypt_pub_key_content:
+                # Base64 encode the PEM public key for safe HTTP header transmission
+                # This avoids issues with newlines and special characters in headers
+                pk_base64 = base64.b64encode(self.auc_file_encrypt_pub_key_content.encode()).decode()
+                headers["PK"] = pk_base64
+
+        # logger.info(f"[Request] url: {url}, headers: {headers}, data: {data}")
+
+        response = self.secure_http_client.post(url, headers=headers, json=data, timeout=self.timeout)
+        if response.status_code == 200:
+            result = response.json()
+            task_id = result.get("Data", {}).get("TaskID", "")
+            return task_id
+        else:
+            logger.error(f"[Error] url: {url}, response: {response.text}")
+            return ""
+
+    def query_task(self, task_id: str, request_id: str = "") -> dict:
+        """
+        Query auc task
+        API: POST /auc/lark/query
+        """
+
+        if not task_id:
+            raise AICCException("task_id must be provided")
+
+        base_url = self._get_base_url()
+        url = f"{base_url}{self.QUERY_TASK_URL}"
+
+        headers = self._build_common_headers(request_id)
+        payload = {"TaskID": task_id}
+
+        response = self.secure_http_client.post(url, headers=headers, json=payload, timeout=self.timeout)
+        if response.status_code == 200:
+            result = response.json()
+            return result
+        else:
+            try:
+                result = response.json()
+                return result
+            except Exception as e:
+                logger.error(f"[Error] url: {url}, Exception={e}")
+                raise AICCException("Failed to query auc task")
+
+    def download_file(self, task_id: str, auc_file_path: str) -> bool:
+        """
+        Download and decrypt encrypted auc
+        This function polls the task status until it succeeds, then downloads
+        and decrypts the auc using the provided private key.
+        Args:
+            task_id: AUC generation task ID
+            auc_file_path: Directory to save decrypted auc
+
+        Returns:
+            True if download succeeded, False otherwise
+        """
+        if not task_id:
+            raise AICCException("task_id must be provided")
+
+        if not auc_file_path:
+            raise AICCException("auc_file_path must be provided")
+
+        task_info = self.query_task(task_id)
+        status = task_info.get("Data", {}).get("Status", "")
+        if not status:
+            logger.error(f"task status is null, task_id={task_id}, query_task return {task_info}")
+            raise AICCException("task status is null")
+        if status != "success":
+            logger.info(f"task status is not success, task_id={task_id}, status={status}")
+            return False
+
+        result = task_info.get("Data", {}).get("Result") or {}
+        for label, key in (
+            ("AudioTranscriptionFile", "AudioTranscriptionFile"),
+            ("SummarizationFile", "SummarizationFile"),
+            ("ChapterFile", "ChapterFile"),
+            ("InformationExtractionFile", "InformationExtractionFile"),
+            ("TranslationFile", "TranslationFile"),
+        ):
+            file_url = result.get(key, "")
+            if file_url:
+                local_file_path = f"./{auc_file_path}_{label}.json"
+                self.download_file_by_url(file_url, local_file_path)
+
+    def download_file_by_url(self, file_url: str, local_file_path: str) -> bool:
+        """
+        Download auc file
+        Args:
+            file_url: File file_url
+            auc_file_path: Directory to save auc file
+
+        Returns:
+            Path to auc file, or None if failed
+        """
+        if not file_url:
+            raise AICCException("file_url must be provided")
+
+        if not local_file_path:
+            raise AICCException("local_file_path must be provided")
+
+        if os.path.exists(local_file_path) and os.path.getsize(local_file_path) > 0:
+            logger.error(f"[Error] File {local_file_path} already exists and is not empty")
+            return False
+
+        try:
+            response = requests.get(file_url, stream=True)
+            with open(local_file_path, "wb") as f:
+                for chunk in response.iter_content(chunk_size=8192):
+                    f.write(chunk)
+
+            if self.is_enable_auc_encrypt:
+                if not self.auc_file_encrypt_priv_key:
+                    logger.error("[Error] No private key set, cannot decrypt encrypted auc")
+                    return False
+
+                headers = response.headers
+                encrypted_key = headers.get("x-tos-meta-enc-dek")
+                if encrypted_key:
+                    try:
+                        encrypted_key_bytes = base64.b64decode(unquote(encrypted_key))
+                        encrypted_key = self.auc_file_encrypt_priv_key.decrypt(encrypted_key_bytes)
+                        tmp_file_path = local_file_path + ".tmp"
+                        decrypt_file(encrypted_key, local_file_path, tmp_file_path, 'b')
+                        os.remove(local_file_path)
+                        os.rename(tmp_file_path, local_file_path)
+                    except Exception as e:
+                        logger.error(f"[Error] Failed to decode encrypted_key: {e}")
+                        traceback.print_exc()
+                        return False
+        except Exception as e:
+            logger.error(f"[Error] Download failed: {e}")
+            traceback.print_exc()
+            return False
+        return True

bytedance/volcengine_aicc_sdk/common/aes.py

@@ -0,0 +1,136 @@
+__all__ = [
+    "gcm_encrypt",
+    "gcm_decrypt",
+    "aes_encrypt",
+    "aes_decrypt",
+    "encrypt_file",
+    "decrypt_file"
+]
+
+import base64
+from typing_extensions import IO, Union
+
+import secrets
+
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
+
+
+SALT = "ai-handler-salt"
+
+
+def gcm_decrypt(key: bytes, ciphertext: bytes) -> bytes:
+    iv, ciphertext = ciphertext[:12], ciphertext[12:]
+    ciphertext, tag = ciphertext[:-16], ciphertext[-16:]
+    decryptor = Cipher(
+        algorithms.AES(key),
+        modes.GCM(iv, tag),
+    ).decryptor()
+
+    return decryptor.update(ciphertext) + decryptor.finalize()
+
+
+def gcm_encrypt(key: bytes, plaintext: bytes) -> bytes:
+    iv = secrets.token_bytes(12)
+    encryptor = Cipher(
+        algorithms.AES(key),
+        modes.GCM(iv),
+    ).encryptor()
+
+    ciphertext = encryptor.update(plaintext) + encryptor.finalize()
+    return iv + ciphertext + encryptor.tag
+
+
+def generate_key_from_string(seed_string: str) -> bytes:
+    backend = default_backend()
+    # 将字符串种子转换为字节
+    password = seed_string.encode()
+    salt = SALT.encode()
+    kdf = PBKDF2HMAC(
+        algorithm=hashes.SHA256(), length=32, salt=salt, iterations=100000, backend=backend
+    )
+    key = kdf.derive(password)
+    return key
+
+
+def aes_encrypt(aes_key: Union[bytes, str], data: Union[bytes, str]) -> bytes:
+    if isinstance(aes_key, str):
+        aes_key = aes_key.encode()
+    if isinstance(data, str):
+        data = data.encode()
+    return gcm_encrypt(aes_key, data)
+
+
+def aes_decrypt(aes_key: Union[bytes, str], data: Union[bytes, str]) -> bytes:
+    if isinstance(aes_key, str):
+        aes_key = aes_key.encode()
+    if isinstance(data, str):
+        data = data.encode()
+    return gcm_decrypt(aes_key, data)
+
+
+def encrypt_file(
+        aes_key: Union[bytes, str],
+        source_path: str,
+        dest_path: str,
+        mode: str = "b",
+) -> bool:
+    if mode not in ["b", "t"]:
+        raise Exception("mode argument error")
+
+    if isinstance(aes_key, str):
+        aes_key = aes_key.encode()
+
+    with open(source_path, f"r{mode}") as source, open(
+            dest_path, f"w{mode}"
+    ) as dest:
+        if mode == "b":
+            source_b: IO[bytes] = source
+            dest_b: IO[bytes] = dest
+            plaintext = source_b.read()
+            ciphertext = aes_encrypt(aes_key, plaintext)
+            dest_b.write(ciphertext)
+        else:
+            source_t: IO[str] = source
+            dest_t: IO[str] = dest
+            for plaintext in source_t:
+                plaintext = plaintext.rstrip("\n")
+                ciphertext = aes_encrypt(aes_key, plaintext)
+                dest_t.write(base64.b64encode(ciphertext).decode() + "\n")
+
+    return True
+
+
+def decrypt_file(
+        aes_key: Union[bytes, str],
+        source_path: str,
+        dest_path: str,
+        mode: str = "b",
+) -> bool:
+    if mode not in ["b", "t"]:
+        raise Exception("mode argument error")
+
+    if isinstance(aes_key, str):
+        aes_key = aes_key.encode()
+
+    with open(source_path, f"r{mode}") as source, open(
+            dest_path, f"w{mode}"
+    ) as dest:
+        if mode == "b":
+            source_b: IO[bytes] = source
+            dest_b: IO[bytes] = dest
+            ciphertext = source_b.read()
+            plaintext = aes_decrypt(aes_key, ciphertext)
+            dest_b.write(plaintext)
+        else:
+            source_t: IO[str] = source
+            dest_t: IO[str] = dest
+            for ciphertext in source_t:
+                ciphertext = ciphertext.rstrip("\n")
+                ciphertext = base64.b64decode(ciphertext.encode())
+                plaintext = aes_decrypt(aes_key, ciphertext)
+                dest_t.write(plaintext.decode() + "\n")
+                # dest_t.write(base64.b64encode(plaintext).decode() + "\n")
+    return True

bytedance/volcengine_aicc_sdk/common/constant.py

@@ -0,0 +1,4 @@
+# -*- coding: utf-8 -*-
+
+class Constants:
+    LOG_CONFIG = "{\"dir\": \"aicc_log\"}"

bytedance/volcengine_aicc_sdk/common/error.py

@@ -0,0 +1,21 @@
+"""
+安全通信错误类型.
+"""
+
+__all__ = [
+    "AICCException"
+]
+
+
+class AICCException(Exception):
+    """AICC 异常"""
+
+    def __init__(self, msg):
+        Exception.__init__(self)
+        self._msg = msg
+
+    def get_msg(self):
+        return self._msg
+
+    def __str__(self):
+        return self.get_msg()

bytedance/volcengine_aicc_sdk/common/log/__init__.py

@@ -0,0 +1,24 @@
+__all__ = [
+    "DEBUG",
+    "INFO",
+    "WARNING",
+    "ERROR",
+    "CRITICAL",
+    "FORMAT",
+    "LogConfig",
+    "debug",
+    "info",
+    "warning",
+    "error",
+    "critical",
+    "init_log_config"
+]
+
+from .config import DEBUG, INFO, WARNING, ERROR, CRITICAL, FORMAT, LogConfig
+from .logger import init_log_config, info, debug, warning, error, critical
+
+_initialized = False
+if not _initialized:
+    config = LogConfig()
+    init_log_config(config)
+    _initialized = True

bytedance/volcengine_aicc_sdk/common/log/config.py

@@ -0,0 +1,31 @@
+__all__ = [
+    "DEBUG",
+    "INFO",
+    "WARNING",
+    "ERROR",
+    "CRITICAL",
+    "FORMAT",
+    "LogConfig"
+]
+
+# Log level constants
+DEBUG = "DEBUG"
+INFO = "INFO"
+WARNING = "WARNING"
+ERROR = "ERROR"
+CRITICAL = "CRITICAL"
+# Log output format
+FORMAT = "{time:YYYY-MM-DD HH:mm:ss} | {level} | {file.path}:{line} | {message}"
+
+
+class LogConfig:
+    """Log configuration class."""
+
+    def __init__(self):
+        self.rotation = 100  # Size limit for log files in MB
+        self.retention = 7  # Days to retain log files
+        self.compression = "zip"
+        self.dir = "aicc_log"
+        self.filename = "aicc.log"
+        self.format = FORMAT
+        self.level = INFO

bytedance/volcengine_aicc_sdk/common/log/logger.py

@@ -0,0 +1,88 @@
+__all__ = [
+    "debug",
+    "info",
+    "warning",
+    "error",
+    "critical",
+    "init_log_config"
+]
+
+import os.path
+import sys
+from typing import Any
+
+from loguru import logger
+
+from . import LogConfig, FORMAT
+
+
+def init_log_config(*configs: LogConfig):
+    """Initialize logging configuration with given configs."""
+    logger.remove()
+
+    log_level = "INFO"
+    if os.environ.get("log_level"):
+        log_level = os.environ.get("log_level")
+    if os.environ.get("LOG_LEVEL"):
+        log_level = os.environ.get("LOG_LEVEL")
+
+    logger.add(sink=sys.stdout, format=FORMAT, level=log_level)
+
+    for config in configs:
+        if not config.dir:
+            config.dir = os.environ.get("LOG_DIR", "aicc_log")
+
+        if not config.filename:
+            config.filename = "aicc.log"
+
+        os.makedirs(config.dir, exist_ok=True)
+        fullpath = os.path.join(config.dir, config.filename)
+
+        if config.rotation <= 0:
+            config.rotation = 100
+        rotation = f"{config.rotation} MB"
+
+        if config.retention <= 0:
+            config.retention = 7
+        retention = f"{config.retention} days"
+
+        logger.add(
+            sink=fullpath,
+            rotation=rotation,
+            retention=retention,
+            compression=config.compression,
+            format=config.format,
+            level=config.level
+        )
+
+
+def msg_proc(message: str) -> str:
+    if message:
+        return f"[aicc]:{message}"
+    else:
+        return message
+
+
+def debug(message: str, *args: Any, **kwargs: Any) -> None:
+    """Log debug message."""
+    logger.opt(depth=1).debug(msg_proc(message), *args, **kwargs)
+
+
+def info(message: str, *args: Any, **kwargs: Any) -> None:
+    """Log info message."""
+    logger.opt(depth=1).info(msg_proc(message), *args, **kwargs)
+
+
+def warning(message: str, *args: Any, **kwargs: Any) -> None:
+    """Log warning message."""
+    logger.opt(depth=1).warning(msg_proc(message), *args, **kwargs)
+
+
+def error(message: str, *args: Any, **kwargs: Any) -> None:
+    """Log error message."""
+    logger.opt(depth=1).error(msg_proc(message), *args, **kwargs)
+
+
+def critical(message: str, *args: Any, **kwargs: Any) -> None:
+    """Log critical message."""
+    logger.opt(depth=1).critical(msg_proc(message), *args, **kwargs)