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

bytedance/volcengine_aicc_sdk/common/tos_utils.py

@@ -0,0 +1,341 @@
+__all__ = ["TOSUtils"]
+
+from urllib import request
+from .aes import decrypt_file
+import os
+import wget
+import time
+import httpx
+from .error import AICCException
+from .log import logger
+
+
+class TOSUtils:
+    def __init__(self, workdir: str = "./tmp"):
+        self.workdir = workdir
+        if not os.path.exists(self.workdir):
+            os.makedirs(self.workdir, exist_ok=True)
+
+        self.default_large_file_threshold = 100 * 1024 * 1024
+        self.default_download_timeout = 300.0
+        self.default_chunk_size = 1024 * 1024
+
+        self.max_retries = 3
+        self.retry_delay = 2  # seconds
+
+    def _download_small_file_wget(self, tos_url: str, local_file_path: str) -> bool:
+        """
+        Download small files using Python wget module
+        
+        Args:
+            tos_url: TOS文件路径
+            local_file_path: 本地文件路径
+            
+        Returns:
+            bool: 是否下载成功
+        """
+        if not tos_url or not local_file_path:
+            raise AICCException("TOS URL and local file path are required")
+
+        if os.path.exists(local_file_path):
+            raise AICCException("Local file already exists, please delete it first")
+        
+        for attempt in range(1, self.max_retries + 1):
+            try:
+                logger.info(f"Download attempt {attempt}/{self.max_retries}")
+                
+                # Download file using wget
+                wget.download(tos_url, local_file_path)
+                
+                # Verify file was downloaded
+                if not os.path.exists(local_file_path):
+                    raise AICCException("File download failed, try again...")
+                
+                file_size = os.path.getsize(local_file_path)
+                if file_size == 0:
+                    raise AICCException("Downloaded file is empty, try again...")
+                
+                logger.info(f"Download successful, file size: {file_size} bytes")
+                return True
+                
+            except Exception as e:
+                logger.warning(f"Download attempt {attempt} failed: {str(e)}, try again...")
+                if os.path.exists(local_file_path):
+                    try:
+                        os.remove(local_file_path)
+                    except:
+                        pass
+                
+                if attempt == self.max_retries:
+                    raise AICCException(
+                        f"Failed to download encrypted file after {self.max_retries} attempts\n"
+                        f"Last error: {str(e)}\n"
+                        "Please check:\n"
+                        "  1. TOS URL is accessible\n"
+                        "  2. Network connection is stable\n"
+                        "  3. File exists on TOS"
+                    )
+                
+                logger.info(f"Waiting {self.retry_delay} seconds before retry...")
+                time.sleep(self.retry_delay)
+                self.retry_delay *= 2
+        
+        raise AICCException("Download failed unexpectedly")
+
+    def _download_small_file_httpx(self, tos_url: str, local_file_path: str) -> bool:
+        """
+        Download small files using httpx (fallback when wget is not available)
+        
+        Args:
+            tos_url: TOS文件路径
+            local_file_path: Path to save the downloaded file            
+        Returns:
+            bool: 是否下载成功
+        """
+        if not tos_url or not local_file_path:
+            raise AICCException("TOS URL and local file path are required")
+
+        if os.path.exists(local_file_path):
+            raise AICCException("Local file already exists, please delete it first")
+        for attempt in range(1, self.max_retries + 1):
+            try:
+                logger.info(f"Download attempt {attempt}/{self.max_retries}")
+                
+                # Download file using httpx
+                with httpx.stream('GET', tos_url, timeout=self.default_download_timeout, follow_redirects=True) as response:
+                    response.raise_for_status()
+                    
+                    with open(local_file_path, 'wb') as f:
+                        for chunk in response.iter_bytes(chunk_size=self.default_chunk_size):
+                            f.write(chunk)
+                
+                # Verify file was downloaded
+                if not os.path.exists(local_file_path):
+                    raise AICCException("File download failed, try again...")
+                
+                file_size = os.path.getsize(local_file_path)
+                if file_size == 0:
+                    raise AICCException("Downloaded file is empty")
+                
+                logger.info(f"Download successful, file size: {file_size} bytes")
+                return True
+                
+            except Exception as e:
+                logger.warning(f"Download attempt {attempt} failed: {str(e)}")
+                
+                # Clean up failed download
+                if os.path.exists(local_file_path):
+                    try:
+                        os.remove(local_file_path)
+                    except:
+                        pass
+                
+                # If this was the last attempt, raise error
+                if attempt == self.max_retries:
+                    raise AICCException(
+                        f"Failed to download encrypted file after {self.max_retries} attempts\n"
+                        f"Last error: {str(e)}\n"
+                        "Please check:\n"
+                        "  1. TOS URL is accessible\n"
+                        "  2. Network connection is stable\n"
+                        "  3. File exists on TOS"
+                    )
+                
+                # Wait before retry
+                logger.info(f"Waiting {self.retry_delay} seconds before retry...")
+                time.sleep(self.retry_delay)
+                self.retry_delay *= 2
+
+        # Should never reach here, but satisfy type checker
+        raise AICCException("Download failed unexpectedly")
+
+    def _download_large_file_streaming(self, tos_url: str, local_file_path: str, expected_size: int) -> bool:
+        """
+        Download large files using httpx streaming
+        
+        Args:
+            tos_url: TOS文件路径
+            local_file_path: Path to save the downloaded file            
+            expected_size: Expected file size in bytes
+            
+        Returns:
+            Path to the downloaded file
+        """
+        if not tos_url or not local_file_path or expected_size <= 0:
+            raise AICCException("TOS URL, local file path, and expected file size are required and must be greater than 0")
+
+        if os.path.exists(local_file_path):
+            raise AICCException("Local file already exists, please delete it first")
+
+        for attempt in range(1, self.max_retries + 1):
+            try:
+                logger.info(f"Streaming download attempt {attempt}/{self.max_retries}")
+                
+                # Streaming download with timeout
+                with httpx.stream('GET', tos_url, timeout=self.default_download_timeout, follow_redirects=True) as response:
+                    response.raise_for_status()
+                    
+                    downloaded_bytes = 0
+                    
+                    with open(local_file_path, 'wb') as f:
+                        for chunk in response.iter_bytes(chunk_size=self.default_chunk_size):
+                            f.write(chunk)
+                            downloaded_bytes += len(chunk)
+                            
+                            # Log progress every 10MB
+                            if downloaded_bytes % (10 * 1024 * 1024) < self.default_chunk_size:
+                                progress = (downloaded_bytes / expected_size) * 100 if expected_size > 0 else 0
+                                logger.info(f"Download progress: {progress:.1f}% ({downloaded_bytes / (1024*1024):.1f} MB / {expected_size / (1024*1024):.1f} MB)")
+                
+                # Verify file was downloaded
+                if not os.path.exists(local_file_path):
+                    raise AICCException("File download failed, try again...")
+                
+                file_size = os.path.getsize(local_file_path)
+                if file_size == 0:
+                    raise AICCException("Downloaded file is empty")
+                
+                logger.info(f"Download successful, file size: {file_size} bytes")
+                return True
+                
+            except Exception as e:
+                logger.warning(f"Streaming download attempt {attempt} failed: {str(e)}")
+                
+                # Clean up failed download
+                if os.path.exists(local_file_path):
+                    try:
+                        os.remove(local_file_path)
+                    except:
+                        pass
+                
+                # If this was the last attempt, raise error
+                if attempt == self.max_retries:
+                    raise AICCException(
+                        f"Failed to download large file after {self.max_retries} attempts\n"
+                        f"Last error: {str(e)}\n"
+                        "Please check:\n"
+                        "  1. TOS URL is accessible\n"
+                        "  2. Network connection is stable\n"
+                        "  3. File exists on TOS\n"
+                        "  4. Timeout is sufficient for file size"
+                    )
+                
+                # Wait before retry
+                logger.info(f"Waiting {self.retry_delay} seconds before retry...")
+                time.sleep(self.retry_delay)
+                self.retry_delay *= 2
+        
+        # Should never reach here, but satisfy type checker
+        raise AICCException("Download failed unexpectedly")
+
+    def _extract_filename_from_url(self, url: str) -> str:
+        """
+        Extract filename from TOS URL
+        
+        Args:
+            url: TOS URL
+            
+        Returns:
+            Filename
+        """
+        # Remove query parameters
+        path = url.split('?')[0]
+        filename = os.path.basename(path)
+        return filename
+    
+    def download_file(self, tos_url: str, local_file_path: str, encrypted_key: bytes) -> bool:
+        """
+        下载文件到本地路径, encrypted_key有值时需要先解密文件, 再保存到本地路径
+        :return: 是否成功
+        :param tos_url: TOS文件路径
+        :param local_file_path: 本地文件路径
+        :param encrypted_key: 加密密钥, 为空时表示文件未加密
+        """
+        if not tos_url or not local_file_path:
+            raise AICCException("TOS URL and local file path are required")
+        
+        if os.path.exists(local_file_path):
+            raise AICCException("Local file already exists, please delete it first")
+        
+        if not tos_url.startswith('http'):
+            raise AICCException(f"Invalid TOS URL: {tos_url}\n URL must start with http:// or https://")
+        
+        try:
+            if not self._download_small_file_wget(tos_url, local_file_path):
+                raise AICCException("download file from TOS failed failed, please check the URL")
+
+            if encrypted_key:
+                tmp_file_path = local_file_path + ".tmp"
+                decrypt_file(encrypted_key, local_file_path, tmp_file_path, mode="b")
+                os.remove(local_file_path)
+                os.rename(tmp_file_path, local_file_path)
+
+            return True
+
+        except AICCException as e:
+            logger.error(f"Error extracting filename from TOS URL: {e}")
+            raise e
+        except Exception as e:
+            logger.error(f"Error extracting filename from TOS URL: {str(e)}")
+            raise AICCException("download filename from TOS file failed, please check the URL")
+
+    def download_files(self, tos_url: str, local_file_path: str, encrypted_key: bytes) -> bool:
+        """
+        下载多个文件并合并到,local_file_path, encrypted_key有值时需要先解密文件, 再保存
+        :return: 是否成功
+        :param tos_url: TOS文件路径
+        :param local_file_path: 本地文件路径
+        :param encrypted_key: AES加密密钥, 为空时表示文件未加密
+        """
+        if not tos_url or not local_file_path:
+            raise AICCException("TOS URL and local file path are required")
+        
+        if os.path.exists(local_file_path):
+            raise AICCException("Local file already exists, please delete it first")
+        
+        if not tos_url.startswith('http'):
+            raise AICCException(f"Invalid TOS URL: {tos_url}\n URL must start with http:// or https://")
+        
+        try:
+            meta_file_npath = tos_url + "/meta.json" 
+            resp = request.urlopen(meta_file_npath)
+            resp.raise_for_status()
+            meta_json = resp.json()
+            logger.info(f"Meta file content: {meta_json}")
+
+            file_list = meta_json.get("files", [])
+            if not file_list:
+                raise AICCException("No files found in the TOS directory")
+            
+            # sort file_list by filename
+            file_list.sort()
+
+            video_files = []
+            for file_name in file_list:
+                tos_file_path = tos_url + "/" + file_name
+                tmp_file_path = local_file_path + "." + file_name
+                if not self._download_small_file_wget(tos_file_path, tmp_file_path):
+                    raise AICCException(f"download file from TOS failed failed, please check the URL{tos_file_path}")
+
+                if encrypted_key:
+                    decrypt_file_path = tmp_file_path + ".dec"
+                    decrypt_file(encrypted_key, tmp_file_path, decrypt_file_path, mode="b")
+                    os.remove(tmp_file_path)
+                    os.rename(decrypt_file_path, tmp_file_path)
+
+                video_files.append(tmp_file_path)
+            
+            # 将ret_video_file中的文件合并输出到ret_video_file_path
+            with open(local_file_path, 'wb') as f:
+                for video_file in video_files:
+                    with open(video_file, 'rb') as video:
+                        f.write(video.read())
+
+            return True
+
+        except AICCException as e:
+            logger.error(f"Error extracting filename from TOS URL: {e}")
+            raise e
+        except Exception as e:
+            logger.error(f"Error extracting filename from TOS URL: {str(e)}")
+            raise AICCException("download filename from TOS file failed, please check the URL")

bytedance/volcengine_aicc_sdk/common/utils.py

@@ -0,0 +1,45 @@
+__all__ = [
+    "format_query",
+    "file_to_base64",
+    "base64_to_file"
+]
+
+import base64
+import json
+import urllib.parse
+from typing import Dict
+
+
+def file_to_base64(file_path):
+    with open(file_path, 'rb') as file:
+        file_data = file.read()  # 读取文件内容
+        base64_data = base64.b64encode(file_data).decode('utf-8')  # Base64 编码
+    return base64_data
+
+
+# 辅助函数：将Base64数据写入文件
+def base64_to_file(base64_data, output_file_path):
+    """将Base64编码的数据写入文件"""
+    if not base64_data:
+        print("没有Base64数据可写入文件。")
+        return False
+    try:
+        # 将Base64字符串解码为二进制数据
+        binary_data = base64.b64decode(base64_data)
+        # 写入文件
+        with open(output_file_path, 'wb') as f:
+            f.write(binary_data)
+        print(f"成功将Base64数据写入文件：{output_file_path}")
+        return True
+    except Exception as e:
+        print(f"写入文件失败：{e}")
+        return False
+
+
+def format_query(params: Dict[str, str]) -> str:
+    return "&".join(
+        urllib.parse.quote(key, safe="-_.~")
+        + "="
+        + urllib.parse.quote(value, safe="-_.~")
+        for key, value in params.items()
+    )

bytedance/volcengine_aicc_sdk/cu_inference_client.py

@@ -0,0 +1,292 @@
+"""
+CU开源模型机密推理客户端
+"""
+__all__ = ["CUClient"]
+
+import base64
+import json
+import time
+import requests
+import bytedance.jeddak_secure_channel as jsc
+from typing import IO, Union
+from .aicc_inference import AICCInference
+from .common.constant import Constants
+from .common.error import AICCException
+from .common.log import logger
+
+
+class CUClient(AICCInference):
+    """
+    AICC机密推理访问客户端.
+    """
+
+    def __init__(self,
+                 base_url: str = None,
+                 aicc_service_name: str = None,
+                 policy_id: str = None,
+                 is_openai: bool = True,
+                 is_secure: bool = True,
+                 ak: str = None,
+                 sk: str = None,
+                 region: str = "cn-beijing",
+                 **kwargs):
+
+        if not base_url:
+            raise AICCException("base_url cannot be empty")
+
+        super().__init__(base_url, is_openai, is_secure, ak, sk, region)
+
+        self.aicc_service_name = aicc_service_name
+        self.policy_id = policy_id
+        self.kwargs = kwargs
+
+        self.jsc_client = None
+        self.max_retries = 3
+
+        if is_secure:
+            jsc_config = dict()
+            jsc_config["ra_service_name"] = aicc_service_name
+            jsc_config["ra_policy_id"] = policy_id
+            jsc_config["bytedance_top_info"] = dict()
+            jsc_config["bytedance_top_info"]["ak"] = ak
+            jsc_config["bytedance_top_info"]["sk"] = sk
+            jsc_config["bytedance_top_info"]["region"] = region
+            jsc_config["bytedance_top_info"]["url"] = self.kwargs.get("top_url", "open.volcengineapi.com")
+            jsc_config["bytedance_top_info"]["service"] = self.kwargs.get("top_service", "pcc")
+            jsc_config["log_config"] = self.kwargs.get("log_config", Constants.LOG_CONFIG)
+
+            try:
+                secure_channel_config = jsc.ClientConfig.from_dict(jsc_config)
+                self.jsc_client = jsc.Client(secure_channel_config)
+            except Exception as e:
+                logger.error(f"Error init jsc: {e}")
+                raise AICCException("create cu client failed")
+
+    def completions(self, prompt: str, output_stream: IO = None) -> Union[str, None]:
+        """
+        生成聊天补全
+        :param prompt:
+        :param output_stream:
+        :return:
+        """
+        if not prompt:
+            raise AICCException("vllm completions Prompt cannot be empty")
+
+        try:
+            if self.is_secure:
+                url = f"{self.base_url}/v2/completions"
+            else:
+                url = f"{self.base_url}/v1/completions"
+
+            headers = {"Content-Type": "application/json"}
+            payload = {
+                "prompt": prompt,
+                "max_tokens": int(self.kwargs.get("max_tokens", 100)),
+                "temperature": float(self.kwargs.get("temperature", 1.0)),
+                "stream": True if output_stream else False
+            }
+
+            encrypt_key = None
+            if self.is_secure:
+                # 加密数据并发送
+                try:
+                    encrypted_prompt, encrypt_key = self.jsc_client.encrypt_with_response(json.dumps(payload["prompt"]))
+                    payload["prompt"] = encrypted_prompt  # 敏感字段
+                except Exception as e:
+                    logger.error(f"Error during encrypt: {e}")
+                    raise AICCException("error during encrypt")
+
+            if output_stream:
+                try:
+                    # 启用流式响应
+                    with requests.post(url, json=payload, headers=headers, stream=True) as response:
+                        response.raise_for_status()
+                        for chunk in response.iter_lines():
+                            if chunk:
+                                if self.is_secure and encrypt_key:
+                                    data = encrypt_key.decrypt(base64.b64decode(chunk))
+                                    output_stream.write(f"{data}")
+                                else:
+                                    output_stream.write(f"{chunk}")
+                except Exception as e:
+                    logger.error(f"call AICC VLLM service completions stream error:{e}")
+                    raise AICCException("vllm completions_stream error")
+            else:
+                for attempt in range(self.max_retries):
+                    try:
+                        response = requests.post(url, json=payload, headers=headers)
+                        response.raise_for_status()
+                        data = response.json()
+
+                        if self.is_secure and encrypt_key:
+                            for elm in data['choices']:
+                                elm['text'] = encrypt_key.decrypt(elm['text']).decode()
+
+                        return json.dumps(data)
+                    except Exception as e:
+                        if attempt == self.max_retries - 1:
+                            logger.error(f"call AICC VLLM service completions error:{e}")
+                            raise AICCException("call vllm completions error")
+                        else:
+                            logger.error(f"call AICC VLLM service completions error:{e}, try again....")
+                            time.sleep(1)
+        except AICCException as e:
+            raise e
+        except Exception as e:
+            logger.error(f"completions error,Exception={e}")
+            raise AICCException("completions error")
+
+    def completions_stream(self, prompt: str, output_stream: IO) -> Union[str, None]:
+        return self.completions(prompt, output_stream)
+
+    def chat_completions(self, messages: list, output_stream: IO = None) -> Union[str, None]:
+        """
+        生成聊天补全
+        :param messages:
+        :param output_stream:
+        :return:
+        """
+        try:
+            if not messages:
+                raise AICCException("vllm chat_completions Messages cannot be empty")
+
+            if not self.base_url:
+                raise AICCException("vllm chat_completions base_url cannot be empty")
+
+            if self.is_secure:
+                url = f"{self.base_url}/v2/chat/completions"
+            else:
+                url = f"{self.base_url}/v1/chat/completions"
+
+            headers = {"Content-Type": "application/json"}
+            payload = {
+                "messages": messages,
+                "max_tokens": int(self.kwargs.get("max_tokens", 100)),
+                "temperature": float(self.kwargs.get("temperature", 1.0)),
+                "stream": True if output_stream else False  # 启用流式响应
+            }
+
+            encrypt_key = None
+            if self.is_secure:
+                # 加密数据并发送
+                try:
+                    encrypted_messages, encrypt_key = \
+                        self.jsc_client.encrypt_with_response(json.dumps(payload["messages"]))
+                    payload["messages"] = encrypted_messages  # 敏感字段
+                except Exception as e:
+                    logger.error(f"Error during encrypt: {e}")
+                    raise AICCException("Error during encrypt")
+
+            if output_stream:
+                with requests.post(url, json=payload, headers=headers, stream=True) as response:
+                    response.raise_for_status()
+                    for chunk in response.iter_lines():
+                        if chunk:
+                            if self.is_secure and encrypt_key:
+                                data = encrypt_key.decrypt(base64.b64decode(chunk))
+                                output_stream.write(f"{data}")
+                            else:
+                                output_stream.write(f"{chunk}")
+            else:
+                for attempt in range(self.max_retries):
+                    try:
+                        response = requests.post(url, json=payload, headers=headers)
+                        response.raise_for_status()
+                        data = response.json()
+
+                        if self.is_secure and encrypt_key:
+                            for elm in data['choices']:
+                                elm['message'] = json.loads(encrypt_key.decrypt(elm['message']))
+
+                        return json.dumps(data)
+                    except Exception as e:
+                        if attempt == self.max_retries - 1:
+                            logger.error(f"call AICC VLLM service chat_completions error:{e}")
+                            raise AICCException("vllm chat_completions error")
+                        else:
+                            logger.error(f"call AICC VLLM service chat_completions error:{e}, try again....")
+                            time.sleep(1)
+        except AICCException as e:
+            raise e
+        except Exception as e:
+            logger.error(f"call AICC VLLM service chat_completions error: {e}")
+            raise AICCException("AICC VLLM service chat_completions error")
+
+    def chat_completions_stream(self, messages: list, output_stream: IO) -> Union[str, None]:
+        return self.chat_completions(messages, output_stream)
+
+    def chat_completions_stream_yield(self, messages: list):
+        """
+        生成聊天补全流式响应
+        :param messages:
+        :return:
+        """
+        try:
+            if not messages:
+                raise AICCException("vllm chat_completions Messages cannot be empty")
+
+            if not self.base_url:
+                raise AICCException("vllm chat_completions_stream_yield base_url cannot be empty")
+
+            if self.is_secure:
+                url = f"{self.base_url}/v2/chat/completions"
+            else:
+                url = f"{self.base_url}/v1/chat/completions"
+
+            headers = {"Content-Type": "application/json"}
+            payload = {
+                "messages": messages,
+                "max_tokens": int(self.kwargs.get("max_tokens", 100)),
+                "temperature": float(self.kwargs.get("temperature", 1.0)),
+                "stream": True
+            }
+
+            encrypt_key = None
+            if self.is_secure:
+                if not self.jsc_client:
+                    raise AICCException("jsc client is NULL")
+                try:
+                    encrypted_messages, encrypt_key = \
+                        self.jsc_client.encrypt_with_response(json.dumps(payload["messages"]))
+                    payload["messages"] = encrypted_messages
+                except Exception as e:
+                    logger.error(f"Error during encrypt: {e}")
+                    raise AICCException("Error during encrypt")
+
+            with requests.post(url, json=payload, headers=headers, stream=True) as response:
+                response.raise_for_status()
+                for chunk in response.iter_lines():
+                    if chunk:
+                        if self.is_secure and encrypt_key:
+                            data = encrypt_key.decrypt(base64.b64decode(chunk))
+                            if isinstance(data, bytes):
+                                data = data.decode()
+                            yield data
+                        else:
+                            if isinstance(chunk, bytes):
+                                chunk = chunk.decode()
+                            yield chunk
+        except AICCException as e:
+            raise e
+        except Exception as e:
+            logger.error(f"call AICC VLLM service chat_completions_stream_yield error: {e}")
+            raise AICCException("AICC VLLM service chat_completions_stream_yield error")
+
+    def function_call(self, messages: list, tools: list, output_stream: IO = None):
+        """
+        :param messages:
+        :param tools:
+        :param output_stream:
+        :return:
+        """
+        error_info = "Sorry, this version does not support function_call at this time."
+        logger.error(error_info)
+        raise AICCException(error_info)
+
+    def function_call_stream(self, messages: list, tools: list, output_stream: IO):
+        return self.function_call(messages, tools, output_stream)
+
+    def asr_recognize(self, file_name: str) -> Union[str, None]:
+        error_info = "Sorry, this version does not support asr at this time."
+        logger.error(error_info)
+        return error_info