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

bytedance/tks/consts.py

@@ -0,0 +1,122 @@
+import os
+
+
+class Conf:
+    # RELATIVE_CONST_INDEX_PATH = './templates/html/main/index.html'
+    # CONST_INDEX_PATH = './html/main/index.html'
+
+    # cookie key名称，用以保存登陆态的session token
+    SESSION_COOKIE_KEY = "session_token"
+    # 合法时间戳的最大误差，单位秒
+    TIMESTAMP_THRESHOLD = 60
+    # session token的过期时间，单位秒
+    SESSION_TOKEN_EXPIRE = 86400
+    # redis配置
+    REDIS_HOST = "localhost"
+    REDIS_PORT = 6379
+    REDIS_USERNAME = "default"
+    REDIS_PASSWORD = ""
+    # pcc后端的分流配置
+    BACKEND_MAP = {
+        "/api/tks/": ["http://localhost:10000", "appID", "AK", "SK"],
+        "/api/bgs/": ["http://localhost:10001", "UID", "AK", "SK"],
+        "/api/ras/": ["http://localhost:10002", "UID", "AK", "SK"],
+        "/api/rag/": ["http://localhost:10003", "UID", "AK", "SK"],
+    }
+    STREAM_PATH = "/api/rag/collection/query_and_generate"
+
+    UPLOAD_FOLDER = os.path.abspath("./upload/rag/doc")
+    DOWNLOAD_FOLDER = os.path.abspath("./download/attest/report")
+    ALLOWED_EXTENSIONS = {".txt", ".pdf", ".pptx", ".docx"}
+
+    # pcc后端转发配置
+    PCC_UID = "jeddak_team"
+    PCC_PASSWORD = "some_password"
+    PCC_AK = "some_AK"
+    PCC_SK = "some_SK"
+    RING_ID = "id_1234"
+    POLICY_ID = "policy_1234"
+    KEY_TYPE = "SYMMETRIC_256"
+
+    RAG_SERVICE_NAME = "rag-deployment"
+    HANDLER_SERVICE_NAME = "ai-handler"
+    # RAG_SERVICE = {
+    #     "ServiceInfo": {
+    #         "_key": "rag",
+    #         "_val": "rag-deployment"
+    #     }
+    # }
+    # HANDLER_SERVICE = {
+    #     "ServiceInfo": {
+    #         "_key": "handler",
+    #         "_val": "ai-handler"
+    #     }
+    # }
+    # SERVICE_INFO_LIST = {"rag": RAG_SERVICE, "handler": HANDLER_SERVICE}
+    SERVICE_INFO = {"ServiceInfo": {"rag": RAG_SERVICE_NAME, "handler": HANDLER_SERVICE_NAME}}
+
+    PCC_URL = "https://jeddakchain.bytedance.com"
+    PATH_MAP = {
+        "/api_mobile/doc/update": "/api/rag/collection",
+        "/api_mobile/rag/query_and_generate": "/api/rag/collection/query_and_generate",
+        # security使用同一个接口，如果改成不同，ENV处理逻辑要改
+        "/api_mobile/security/status": "/api/ras/v1/attestation_user",
+        "/api_mobile/security/status/download": "/api/ras/v1/attestation_user",
+        # TODO: 确定接口
+        "/api_mobile/key/set": "/api/tks/",
+    }
+
+    # RAG_ADDR = 'http://localhost:10003'
+    # RAG_COLLECTION_CREATE_ENDPOINT = '/api/rag/collection/create'
+    # RAS_SHIM_ADDR = "http://localhost:10002"
+
+    DOC_STATUS_MAP = {"creating": "处理中", "ready": "处理完成", "failed": "处理失败"}
+
+    SALT = "ai-handler-salt"
+
+    def update_from_env(self):
+        for k in self.BACKEND_MAP:
+            env_key = f"ENV_{k.split('/')[-2].upper()}_URL"
+            # ENV_TKS_URL, ...
+            if env_key in os.environ:
+                self.BACKEND_MAP[k][0] = os.getenv(env_key)
+        for k in self.PATH_MAP:
+            env_key = f"HANDLER_{k.split('/')[2].upper()}_URL"
+            if env_key in os.environ:
+                self.PATH_MAP[k] = os.getenv(env_key)
+        self.REDIS_USERNAME = (
+            os.getenv("ENV_REDIS_USERNAME") if "ENV_REDIS_USERNAME" in os.environ else None
+        )
+        self.REDIS_PASSWORD = (
+            os.getenv("ENV_REDIS_PASSWORD") if "ENV_REDIS_PASSWORD" in os.environ else None
+        )
+        if "ENV_REDIS_HOST" in os.environ:
+            self.REDIS_HOST = os.getenv("ENV_REDIS_HOST")
+        if "ENV_REDIS_PORT" in os.environ:
+            self.REDIS_PORT = int(os.getenv("ENV_REDIS_PORT"))
+        if "ENV_SESSION_TOKEN_EXPIRE" in os.environ:
+            self.SESSION_TOKEN_EXPIRE = int(os.getenv("ENV_SESSION_TOKEN_EXPIRE"))
+        if "RAG_SERVICE_NAME" in os.environ:
+            self.RAG_SERVICE_NAME = os.getenv("RAG_SERVICE_NAME")
+        if "HANDLER_SERVICE_NAME" in os.environ:
+            self.HANDLER_SERVICE_NAME = os.getenv("HANDLER_SERVICE_NAME")
+        if "RING_ID" in os.environ:
+            self.RING_ID = os.getenv("RING_ID")
+        if "POLICY_ID" in os.environ:
+            self.POLICY_ID = os.getenv("POLICY_ID")
+        if "KEY_TYPE" in os.environ:
+            self.KEY_TYPE = os.getenv("KEY_TYPE")
+        if "PCC_UID" in os.environ:
+            self.PCC_UID = os.getenv("PCC_UID")
+        if "PCC_PASSWORD" in os.environ:
+            self.PCC_PASSWORD = os.getenv("PCC_PASSWORD")
+        if "PCC_AK" in os.environ:
+            self.PCC_AK = os.getenv("PCC_AK")
+        if "PCC_SK" in os.environ:
+            self.PCC_SK = os.getenv("PCC_SK")
+        os.makedirs(conf.UPLOAD_FOLDER, exist_ok=True)
+        os.makedirs(conf.DOWNLOAD_FOLDER, exist_ok=True)
+        print(f"final backend_map: {self.BACKEND_MAP}")
+
+
+conf = Conf()

bytedance/tks/crypto/aes.py

@@ -0,0 +1,47 @@
+__all__ = [
+    "gcm_encrypt",
+    "gcm_decrypt",
+]
+
+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
+
+from ..consts import conf
+
+
+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 = conf.SALT.encode()
+    kdf = PBKDF2HMAC(
+        algorithm=hashes.SHA256(), length=32, salt=salt, iterations=100000, backend=backend
+    )
+    key = kdf.derive(password)
+    return key

bytedance/tks/crypto/ec.py

@@ -0,0 +1,39 @@
+from typing import Tuple
+
+from cryptography.hazmat.primitives import hashes, serialization
+from cryptography.hazmat.primitives.asymmetric import ec
+from cryptography.hazmat.primitives.kdf.hkdf import HKDF
+
+__all__ = [
+    "diffie_hellman_key_exchange",
+    "generate_key_pair",
+]
+
+
+def diffie_hellman_key_exchange(local_sk: bytes, peer_pk: bytes) -> bytes:
+    local_sk_key = serialization.load_pem_private_key(local_sk, password=None)
+    peer_pk_key = serialization.load_pem_public_key(peer_pk)
+    assert isinstance(local_sk_key, ec.EllipticCurvePrivateKey)
+    assert isinstance(peer_pk_key, ec.EllipticCurvePublicKey)
+    shared_key = local_sk_key.exchange(ec.ECDH(), peer_pk_key)
+    derived_key = HKDF(algorithm=hashes.SHA256(), length=32, salt=None, info=None).derive(
+        shared_key
+    )
+    return derived_key
+
+
+def generate_key_pair() -> Tuple[bytes, bytes]:
+    private_key = ec.generate_private_key(ec.SECP384R1())
+    public_key = private_key.public_key()
+
+    private_key_pem = private_key.private_bytes(
+        encoding=serialization.Encoding.PEM,
+        format=serialization.PrivateFormat.PKCS8,
+        encryption_algorithm=serialization.NoEncryption(),
+    )
+
+    public_key_pem = public_key.public_bytes(
+        encoding=serialization.Encoding.PEM, format=serialization.PublicFormat.SubjectPublicKeyInfo
+    )
+
+    return private_key_pem, public_key_pem

bytedance/tks/exceptions.py

@@ -0,0 +1,22 @@
+from typing import Optional
+
+__all__ = [
+    "TKSError",
+    "RAError",
+]
+
+
+class TKSError(Exception):
+    def __init__(self, message: str, http_status: Optional[int] = None):
+        super().__init__(message)
+        self.http_status = http_status
+
+    def __str__(self):
+        if self.http_status:
+            return f"{self.args[0]} (HTTP Status: {self.http_status})"
+        return self.args[0]
+
+
+class RAError(TKSError):
+    def __init__(self, message: str, http_status: Optional[int] = None):
+        super().__init__(message, http_status)

bytedance/volcengine_aicc_sdk/__init__.py

@@ -0,0 +1,14 @@
+__all__ = [
+    "AICCClient",
+    "AICCModelType",
+    "__version__",
+]
+
+__version__ = "0.0.41"
+
+from .aicc_client import AICCModelType
+from .aicc_client import AICCClient
+
+
+def __getattr__(name: str) -> object:
+    raise AttributeError(f"module {__name__} has no attribute {name}")

bytedance/volcengine_aicc_sdk/aicc_client.py

@@ -0,0 +1,277 @@
+from bytedance.volcengine_aicc_sdk.seedance_inference_client import SeedanceClient
+from bytedance.volcengine_aicc_sdk.seedream_inference_client import SeedreamClient
+from typing import Any, List, Optional, Union
+"""
+AICC访问客户端
+"""
+
+__all__ = [
+    "AICCClient",
+    "AICCModelType"
+]
+
+from bytedance.volcengine_aicc_sdk.cu_inference_client import CUClient
+from bytedance.volcengine_aicc_sdk.auc_inference_client import AUCClient
+from bytedance.volcengine_aicc_sdk.ark_inference_client import ARKClient
+from bytedance.volcengine_aicc_sdk.deepseek_inference_client import DeepSeekClient
+from bytedance.volcengine_aicc_sdk.asr_websocket_client import AsrWebSocketClient
+from .asr_inference_client import ASRClient
+from volcenginesdkarkruntime import Ark
+from .common.error import AICCException
+from .common.log import LogConfig, logger
+from bytedance.volcengine_aicc_sdk.secure_http_client import AICCConfig, SecureHttpClient
+
+
+class AICCModelType:
+    DOUBAO = "DOUBAO"  # 机密豆包
+    AICC_CU = "AICC-CU"   # AICC CU版模型
+    OSLM = "OSLM"  # 开源模型
+    ASR = "ASR"
+    ASR_WS = "ASR-WS"
+    SEEDANCE = "SEEDANCE"
+    SEEDREAM = "SEEDREAM"
+    AUC = "AUC"
+
+
+class AICCClient:
+    """
+    AICC机密推理访问客户端.
+    """
+
+    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.aicc_config = aicc_config
+        self.endpoint = endpoint
+        self.kwargs = kwargs
+
+        if not self.base_url or not self.api_key:
+            raise AICCException(f"New AICC client failed, base_url={self.base_url}, api_key={self.api_key}")
+
+        if not self.endpoint:
+            self.endpoint = self.kwargs.get("model_name")
+
+        self.secure_http_client = self._init_secure_http_client()
+        self.config = dict()
+        self._init_config()
+        self._init_log()
+
+    def get_llm_instance(self, model_type):
+        model_type = model_type.upper()
+        if AICCModelType.SEEDANCE in model_type:
+            return SeedanceClient(base_url=self.base_url,
+                                  api_key=self.api_key,
+                                  endpoint=self.endpoint,
+                                  aicc_config=self.aicc_config,
+                                  **self.kwargs)
+        elif AICCModelType.AICC_CU in model_type:
+            return CUClient(base_url=self.base_url,
+                            aicc_service_name=self.aicc_config.aicc_llm_service_name,
+                            policy_id=self.aicc_config.policy_id,
+                            is_openai=self.kwargs.get("is_openai"),
+                            is_secure=self.kwargs.get("is_secure"),
+                            ak=self.aicc_config.ak,
+                            sk=self.aicc_config.sk,
+                            region=self.aicc_config.region,
+                            **self.kwargs)
+        elif AICCModelType.ASR == model_type:
+            return ASRClient(base_url=self.base_url,
+                             aicc_service_name=self.aicc_config.aicc_llm_service_name,
+                             policy_id=self.aicc_config.policy_id,
+                             is_openai=self.kwargs.get("is_openai"),
+                             is_secure=self.kwargs.get("is_secure"),
+                             ak=self.aicc_config.ak,
+                             sk=self.aicc_config.sk,
+                             region=self.aicc_config.region,
+                             **self.kwargs)
+        elif AICCModelType.OSLM in model_type:
+            return DeepSeekClient(base_url=self.base_url,
+                                  api_key=self.api_key,
+                                  ep=self.endpoint,
+                                  aicc_service_name=self.aicc_config.aicc_llm_service_name,
+                                  policy_id=self.aicc_config.policy_id,
+                                  is_openai=self.kwargs.get("is_openai"),
+                                  is_secure=self.kwargs.get("is_secure"),
+                                  ak=self.aicc_config.ak,
+                                  sk=self.aicc_config.sk,
+                                  region=self.aicc_config.region,
+                                  **self.kwargs)
+        elif AICCModelType.DOUBAO in model_type:
+            return ARKClient(base_url=self.base_url,
+                             api_key=self.api_key,
+                             ep=self.endpoint,
+                             is_openai=self.kwargs.get("is_openai"),
+                             is_secure=self.kwargs.get("is_secure"),
+                             ak=self.aicc_config.ak,
+                             sk=self.aicc_config.sk,
+                             region=self.aicc_config.region,
+                             **self.kwargs)
+        elif AICCModelType.ASR_WS == model_type:
+            return AsrWebSocketClient(base_url=self.base_url,
+                                      api_key=self.api_key,
+                                      endpoint=self.endpoint,
+                                      aicc_config=self.aicc_config,
+                                      **self.kwargs)
+        elif AICCModelType.AUC == model_type:
+            return AUCClient(base_url=self.base_url,
+                             api_key=self.api_key,
+                             endpoint=self.endpoint,
+                             aicc_config=self.aicc_config,
+                             **self.kwargs)
+        elif AICCModelType.SEEDREAM == model_type:
+            return SeedreamClient(base_url=self.base_url,
+                                  api_key=self.api_key,
+                                  endpoint=self.endpoint,
+                                  aicc_config=self.aicc_config,
+                                  **self.kwargs)
+        else:
+            raise AICCException(f"New AICC client failed, model_type={model_type}")
+
+    @staticmethod
+    def get_inference_instance(model_type: str,
+                               base_url: str = None,
+                               api_key: str = None,
+                               ep: str = None,
+                               aicc_service_name: str = None,
+                               policy_id: str = None,
+                               is_openai: bool = False,
+                               is_secure: bool = True,
+                               ak: str = None,
+                               sk: str = None,
+                               region: str = "cn-beijing",
+                               **kwargs):
+        if model_type.upper() == AICCModelType.DOUBAO:
+            return ARKClient(base_url=base_url,
+                             api_key=api_key,
+                             ep=ep,
+                             is_openai=is_openai,
+                             is_secure=is_secure,
+                             ak=ak,
+                             sk=sk,
+                             region=region,
+                             **kwargs)
+        elif model_type.upper() == AICCModelType.AICC_CU:
+            return CUClient(base_url=base_url,
+                            aicc_service_name=aicc_service_name,
+                            policy_id=policy_id,
+                            is_openai=is_openai,
+                            is_secure=is_secure,
+                            ak=ak,
+                            sk=sk,
+                            region=region,
+                            **kwargs)
+        elif model_type.upper() == AICCModelType.OSLM:
+            return DeepSeekClient(base_url=base_url,
+                                  api_key=api_key,
+                                  ep=ep,
+                                  aicc_service_name=aicc_service_name,
+                                  policy_id=policy_id,
+                                  is_openai=is_openai,
+                                  is_secure=is_secure,
+                                  ak=ak,
+                                  sk=sk,
+                                  region=region,
+                                  **kwargs)
+        elif model_type.upper() == AICCModelType.ASR:
+            return ASRClient(base_url=base_url,
+                             aicc_service_name=aicc_service_name,
+                             policy_id=policy_id,
+                             is_openai=is_openai,
+                             is_secure=is_secure,
+                             ak=ak,
+                             sk=sk,
+                             region=region,
+                             **kwargs)
+        else:
+            err_str = f"can not support model type:{model_type}"
+            logger.error(err_str)
+            raise AICCException(err_str)
+
+    @staticmethod
+    def get_volcengine_ark(base_url: str = None, api_key: str = None, **kwargs):
+        return Ark(api_key=api_key, base_url=base_url, **kwargs)
+
+    def get_secure_http_client(self) -> SecureHttpClient:
+        return self.secure_http_client
+
+    def _init_secure_http_client(self) -> SecureHttpClient:
+        if self.aicc_config:
+            return SecureHttpClient(aicc_config=self.aicc_config)
+        return SecureHttpClient()
+
+    def get_ark_client(self, secure: bool = True, **kwargs: Any) -> Ark:
+        if secure:
+            kwargs.setdefault("http_client", self.secure_http_client)
+        return Ark(api_key=self.api_key, base_url=self.base_url, **kwargs)
+
+    def tokenization(
+            self,
+            text: Union[str, List[str]],
+            model: Optional[str] = None,
+            user: Optional[str] = None,
+            **kwargs: Any):
+        if not self.aicc_config:
+            raise AICCException("Tokenization failed, aicc_config is required")
+
+        tokenization_model = model or self.endpoint or self.kwargs.get("model_name")
+        if not tokenization_model:
+            raise AICCException("Tokenization failed, model is required")
+
+        request_url = f"{self.base_url.rstrip('/')}/tokenization"
+        request_headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json",
+        }
+        request_body = {
+            "model": str(tokenization_model),
+            "text": text,
+        }
+        if user:
+            request_body["user"] = user
+
+        try:
+            response = self.secure_http_client.post(
+                request_url,
+                headers=request_headers,
+                json=request_body,
+                **kwargs,
+            )
+            response.raise_for_status()
+            return response.json()
+        except Exception as e:
+            raise
+
+    @staticmethod
+    def hello() -> str:
+        info = "hello volcengine_aicc_sdk"
+        print(info)
+        return info
+
+    def _init_config(self):
+        pass
+
+    def _init_log(self):
+        config = LogConfig()
+
+        if log_dir := self.config.get("dir", "aicc_log"):
+            config.dir = log_dir
+
+        if filename := self.config.get("filename", "aicc.log"):
+            config.filename = filename
+
+        if rotation := self.config.get("rotation", 100):
+            config.rotation = rotation
+
+        if retention := self.config.get("retention", 30):
+            config.retention = retention
+
+        if level := self.config.get("level"):
+            config.level = level
+
+        logger.init_log_config(config)

bytedance/volcengine_aicc_sdk/aicc_inference.py

@@ -0,0 +1,64 @@
+"""
+AICC机密推理客户端BASE
+"""
+__all__ = ["AICCInference"]
+
+from abc import ABC, abstractmethod
+from typing import IO, Union
+
+
+class AICCInference(ABC):
+    """
+    AICC机密推理访问客户端.
+    """
+
+    def __init__(self,
+                 base_url: str = None,
+                 is_openai: bool = True,
+                 is_secure: bool = True,
+                 ak: str = None,
+                 sk: str = None,
+                 region: str = "cn-beijing"
+                 ):
+        self.base_url = base_url
+        self.is_openai = is_openai
+        self.is_secure = is_secure
+        self.ak = ak
+        self.sk = sk
+        self.region = region
+
+    @staticmethod
+    def hello() -> str:
+        info = "hello aicc inference"
+        print(info)
+        return info
+
+    @abstractmethod
+    def chat_completions(self, messages: list, output_stream: IO = None) -> Union[str, None]:
+        """
+        生成聊天补全
+        :param messages:
+        :param output_stream:是否启用流式响应
+        :return:
+        """
+        pass
+
+    @abstractmethod
+    def function_call(self, messages: list, tools: list, output_stream: IO = None) -> Union[str, None]:
+        """
+        生成聊天补全
+        :param messages:
+        :param tools:
+        :param output_stream: 是否开启流式
+        :return:
+        """
+        pass
+
+    @abstractmethod
+    def asr_recognize(self, file_name: str) -> Union[str, None]:
+        """
+        ASR转文字
+        :param file_name:
+        :return:
+        """
+        pass