Merge branch 'feat/arkruntime-seedance-sdk' into 'integration_2025-10-27_1073654796290'

feat: [development task] ark (1769068)

See merge request iaasng/volcengine-python-sdk!879

Author: BitsAdmin

volcenginesdkarkruntime/_client.py

@@ -143,7 +143,7 @@ def _get_endpoint_sts_token(self, endpoint_id: str):
             self._sts_token_manager = StsTokenManager(self.ak, self.sk, self.region)
         return self._sts_token_manager.get(endpoint_id)
 
-    def _get_endpoint_certificate(self, endpoint_id: str) -> tuple[key_agreement_client, str, str]:
+    def _get_endpoint_certificate(self, endpoint_id: str) -> Tuple[key_agreement_client, str, str, float]:
         if self._certificate_manager is None:
             cert_path = os.environ.get("E2E_CERTIFICATE_PATH")
             if (
@@ -279,7 +279,7 @@ def _get_bot_sts_token(self, bot_id: str):
             self._sts_token_manager = StsTokenManager(self.ak, self.sk, self.region)
         return self._sts_token_manager.get(bot_id, resource_type="bot")
 
-    def _get_endpoint_certificate(self, endpoint_id: str) -> tuple[key_agreement_client, str, str]:
+    def _get_endpoint_certificate(self, endpoint_id: str) -> Tuple[key_agreement_client, str, str, float]:
         if self._certificate_manager is None:
             cert_path = os.environ.get("E2E_CERTIFICATE_PATH")
             if (
@@ -429,7 +429,7 @@ def __init__(
         base_url: str | URL = BASE_URL,
         api_key: str | None = None,
     ):
-        self._certificate_manager: Dict[str, Tuple[key_agreement_client, str, str]] = {}
+        self._certificate_manager: Dict[str, Tuple[key_agreement_client, str, str, float]] = {}
 
         # local cache prepare
         self._init_local_cert_cache()
@@ -522,7 +522,7 @@ def _load_cert_locally(self, ep: str) -> str | None:
                 cert_pem = None
                 with open(cert_file_path, "r") as f:
                     cert_pem = f.read()
-                ring, key = get_cert_info(cert_pem)
+                ring, key, _ = get_cert_info(cert_pem)
                 # check cert is complement with AICC/PCA
                 if (ring == "" or key == "") and not self._aicc_enabled:
                     return cert_pem
@@ -545,7 +545,7 @@ def _init_local_cert_cache(self):
                     "failed to create certificate directory %s: %s\n" % (self._cert_storage_path, e)
                 )
 
-    def get(self, ep: str) -> tuple[key_agreement_client, str, str]:
+    def get(self, ep: str) -> Tuple[key_agreement_client, str, str, float]:
         if ep not in self._certificate_manager:
             cert_pem = self._load_cert_locally(ep)
             if cert_pem is None:
@@ -556,12 +556,13 @@ def get(self, ep: str) -> tuple[key_agreement_client, str, str]:
                 else:
                     cert_pem = self._load_cert_by_ak_sk(ep)
                 self._save_cert_to_file(ep, cert_pem)
-            ring, key = get_cert_info(cert_pem)
+            ring, key, exp_time = get_cert_info(cert_pem)
             self._certificate_manager[ep] = (
                 key_agreement_client(
                     certificate_pem_string=cert_pem
                 ),
                 ring,
                 key,
+                exp_time,
             )
         return self._certificate_manager[ep]

volcenginesdkarkruntime/_utils/_key_agreement.py

@@ -17,7 +17,7 @@
 from typing import Tuple
 
 
-def get_cert_info(cert_pem: str) -> Tuple[str, str]:
+def get_cert_info(cert_pem: str) -> Tuple[str, str, float]:
     import re
     from cryptography import x509
     from cryptography.hazmat.backends import default_backend
@@ -27,10 +27,10 @@ def get_cert_info(cert_pem: str) -> Tuple[str, str]:
         dns = cert.extensions.get_extension_for_class(
             x509.SubjectAlternativeName).value.get_values_for_type(x509.DNSName)
         if dns and len(dns) > 1 and re.match(r"^ring\..*$", dns[0]) and re.match(r"^key\..*$", dns[1]):
-            return dns[0][5:], dns[1][4:]
+            return dns[0][5:], dns[1][4:], cert.not_valid_after_utc.timestamp()
     except Exception:
         pass
-    return "", ""
+    return "", "", cert.not_valid_after_utc.timestamp()
 
 
 def aes_gcm_encrypt_bytes(
@@ -92,6 +92,20 @@ def aes_gcm_decrypt_base64_string(key: bytes, nonce: bytes, ciphertext: str) ->
 base64_pattern = r'(?:[A-Za-z0-9+/]{4})*(?:[A-Za-z0-9+/]{2}==|[A-Za-z0-9+/]{3}=|[A-Za-z0-9+/]{4})'
 
 
+def decrypt_corner_case(key: bytes, nonce: bytes, data: str) -> str:
+    """decrypt_corner_case Decrypt corner case data"""
+    if len(data) < 24:
+        return ''
+    for i in range(20, len(data), 4):
+        try:
+            decrypted = aes_gcm_decrypt_base64_string(key, nonce, data[:i+4])
+            if i+4 == len(data):
+                return decrypted
+            return decrypted + decrypt_corner_case(key, nonce, data[i+4:])
+        except Exception:
+            pass
+
+
 def aes_gcm_decrypt_base64_list(key: bytes, nonce: bytes, ciphertext: str) -> str:
     # Decrypt
     base64_array = re.findall(base64_pattern, ciphertext)
@@ -100,17 +114,7 @@ def aes_gcm_decrypt_base64_list(key: bytes, nonce: bytes, ciphertext: str) -> st
         try:
             result.append(aes_gcm_decrypt_base64_string(key, nonce, b64))
         except Exception:
-            for i in range(20, len(b64), 4):
-                try:
-                    decrypted = aes_gcm_decrypt_base64_string(
-                        key, nonce, b64[:i+4])
-                    result.append(decrypted)
-                    decrypted = aes_gcm_decrypt_base64_string(
-                        key, nonce, b64[i+4:])
-                    result.append(decrypted)
-                    break
-                except Exception:
-                    pass
+            result.append(decrypt_corner_case(key, nonce, b64))
     return ''.join(result)
 
 

volcenginesdkarkruntime/resources/batch_chat/completions.py

@@ -20,12 +20,11 @@
 from typing_extensions import Literal
 
 import httpx
-import warnings
 
 from ..._exceptions import ArkAPITimeoutError, ArkAPIConnectionError, ArkAPIStatusError
 from ..._types import Body, Query, Headers
-from ..._utils import with_sts_token, async_with_sts_token, deepcopy_minimal
-from ..._utils._key_agreement import aes_gcm_decrypt_base64_string
+from ..._utils import with_sts_token, async_with_sts_token
+from ..encryption import with_e2e_encryption, async_with_e2e_encryption
 from ..._base_client import make_request_options
 from ..._resource import SyncAPIResource, AsyncAPIResource
 from ..._compat import cached_property
@@ -50,43 +49,10 @@
 __all__ = ["Completions", "AsyncCompletions"]
 
 
-def _process_messages(
-    messages: Iterable[ChatCompletionMessageParam], f: Callable[[str], str]
-):
-    for message in messages:
-        if message.get("content", None) is not None:
-            current_content = message.get("content")
-            if isinstance(current_content, str):
-                message["content"] = f(current_content)
-            elif isinstance(current_content, Iterable):
-                for part in current_content:
-                    if part.get("type", None) == "text":
-                        part["text"] = f(part["text"])
-                    elif part.get("type", None) == "image_url":
-                        if part["image_url"]["url"].startswith("data:"):
-                            part["image_url"]["url"] = f(part["image_url"]["url"])
-                        else:
-                            warnings.warn(
-                                "encryption is not supported for image url, "
-                                "please use base64 image if you want encryption"
-                            )
-                    else:
-                        raise TypeError(
-                            "encryption is not supported for content type {}".format(
-                                type(part)
-                            )
-                        )
-            else:
-                raise TypeError(
-                    "encryption is not supported for content type {}".format(
-                        type(message.get("content"))
-                    )
-                )
-
-
 def _calculate_retry_timeout(retry_times) -> float:
     nbRetries = min(retry_times, MAX_RETRY_DELAY / INITIAL_RETRY_DELAY)
-    sleep_seconds = min(INITIAL_RETRY_DELAY * pow(2, nbRetries), MAX_RETRY_DELAY)
+    sleep_seconds = min(INITIAL_RETRY_DELAY *
+                        pow(2, nbRetries), MAX_RETRY_DELAY)
     # Apply some jitter, plus-or-minus half a second.
     jitter = 1 - 0.25 * random()
     timeout = sleep_seconds * jitter
@@ -126,58 +92,8 @@ class Completions(SyncAPIResource):
     def with_raw_response(self) -> CompletionsWithRawResponse:
         return CompletionsWithRawResponse(self)
 
-    def _process_messages(
-        self, messages: Iterable[ChatCompletionMessageParam], f: Callable[[str], str]
-    ):
-        for message in messages:
-            if message.get("content", None) is not None:
-                current_content = message.get("content")
-                if isinstance(current_content, str):
-                    message["content"] = f(current_content)
-                elif isinstance(current_content, Iterable):
-                    raise TypeError(
-                        "content type {} is not supported end-to-end encryption".format(
-                            type(message.get("content"))
-                        )
-                    )
-                else:
-                    raise TypeError(
-                        "content type {} is not supported end-to-end encryption".format(
-                            type(message.get("content"))
-                        )
-                    )
-
-    def _encrypt(
-        self,
-        model: str,
-        messages: Iterable[ChatCompletionMessageParam],
-        extra_headers: Headers,
-    ) -> tuple[bytes, bytes]:
-        client = self._client._get_endpoint_certificate(model)
-        _crypto_key, _crypto_nonce, session_token = client.generate_ecies_key_pair()
-        extra_headers["X-Session-Token"] = session_token
-        _process_messages(
-            messages,
-            lambda x: client.encrypt_string_with_key(_crypto_key, _crypto_nonce, x),
-        )
-        return _crypto_key, _crypto_nonce
-
-    def _decrypt(
-        self, key: bytes, nonce: bytes, resp: ChatCompletion
-    ) -> ChatCompletion:
-        if resp.choices is not None:
-            for index, choice in enumerate(resp.choices):
-                if (
-                    choice.message is not None and choice.finish_reason != 'content_filter'
-                    and choice.message.content is not None
-                ):
-                    choice.message.content = aes_gcm_decrypt_base64_string(
-                        key, nonce, choice.message.content
-                    )
-                resp.choices[index] = choice
-        return resp
-
     @with_sts_token
+    @with_e2e_encryption
     def create(
         self,
         *,
@@ -208,14 +124,6 @@ def create(
         extra_body: Body | None = None,
         timeout: float | httpx.Timeout | None = None,
     ) -> ChatCompletion:
-        is_encrypt = False
-        if (
-            extra_headers is not None
-            and extra_headers.get(ARK_E2E_ENCRYPTION_HEADER, None) == "true"
-        ):
-            is_encrypt = True
-            messages = deepcopy_minimal(messages)
-            e2e_key, e2e_nonce = self._encrypt(model, messages, extra_headers)
         retryTimes = 0
         last_time = self._get_request_last_time(timeout)
         model_breaker = self._client.get_model_breaker(model)
@@ -273,8 +181,6 @@ def create(
                     continue
                 else:
                     raise err
-            if is_encrypt:
-                resp = self._decrypt(e2e_key, e2e_nonce, resp)
             return resp
 
     def _get_request_last_time(self, timeout):
@@ -289,7 +195,8 @@ def _get_request_last_time(self, timeout):
             timeoutSeconds = timeout
         else:
             raise TypeError(
-                "timeout type {} is not supported".format(type(self._client.timeout))
+                "timeout type {} is not supported".format(
+                    type(self._client.timeout))
             )
         return datetime.now() + timedelta(seconds=timeoutSeconds)
 
@@ -299,37 +206,8 @@ class AsyncCompletions(AsyncAPIResource):
     def with_raw_response(self) -> AsyncCompletionsWithRawResponse:
         return AsyncCompletionsWithRawResponse(self)
 
-    def _encrypt(
-        self,
-        model: str,
-        messages: Iterable[ChatCompletionMessageParam],
-        extra_headers: Headers,
-    ) -> tuple[bytes, bytes]:
-        client = self._client._get_endpoint_certificate(model)
-        _crypto_key, _crypto_nonce, session_token = client.generate_ecies_key_pair()
-        extra_headers["X-Session-Token"] = session_token
-        _process_messages(
-            messages,
-            lambda x: client.encrypt_string_with_key(_crypto_key, _crypto_nonce, x),
-        )
-        return _crypto_key, _crypto_nonce
-
-    async def _decrypt(
-        self, key: bytes, nonce: bytes, resp: ChatCompletion
-    ) -> ChatCompletion:
-        if resp.choices is not None:
-            for index, choice in enumerate(resp.choices):
-                if (
-                    choice.message is not None and choice.finish_reason != 'content_filter'
-                    and choice.message.content is not None
-                ):
-                    choice.message.content = aes_gcm_decrypt_base64_string(
-                        key, nonce, choice.message.content
-                    )
-                resp.choices[index] = choice
-        return resp
-
     @async_with_sts_token
+    @async_with_e2e_encryption
     async def create(
         self,
         *,
@@ -360,14 +238,6 @@ async def create(
         extra_body: Body | None = None,
         timeout: float | httpx.Timeout | None = None,
     ) -> ChatCompletion:
-        is_encrypt = False
-        if (
-            extra_headers is not None
-            and extra_headers.get(ARK_E2E_ENCRYPTION_HEADER, None) == "true"
-        ):
-            is_encrypt = True
-            messages = deepcopy_minimal(messages)
-            e2e_key, e2e_nonce = self._encrypt(model, messages, extra_headers)
 
         retryTimes = 0
         last_time = self._get_request_last_time(timeout)
@@ -426,8 +296,6 @@ async def create(
                     continue
                 else:
                     raise err
-            if is_encrypt:
-                resp = await self._decrypt(e2e_key, e2e_nonce, resp)
             return resp
 
     def _get_request_last_time(self, timeout):
@@ -442,7 +310,8 @@ def _get_request_last_time(self, timeout):
             timeoutSeconds = timeout
         else:
             raise TypeError(
-                "timeout type {} is not supported".format(type(self._client.timeout))
+                "timeout type {} is not supported".format(
+                    type(self._client.timeout))
             )
         return datetime.now() + timedelta(seconds=timeoutSeconds)