Merge branch 'feat/ark_local_cert_cache' into 'integration_2025-01-02_663619381762'

feat: [development task] ark-runtime-manual-Python (951048)

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

Author: BitsAdmin

volcenginesdkarkruntime/_client.py

@@ -71,6 +71,7 @@ def __init__(
             sk = os.environ.get("VOLC_SECRETKEY")
         if api_key is None:
             api_key = os.environ.get("ARK_API_KEY")
+        self._base_url = base_url
         self.ak = ak
         self.sk = sk
         self.api_key = api_key
@@ -107,10 +108,10 @@ def _get_endpoint_sts_token(self, endpoint_id: str):
     def _get_endpoint_certificate(self, endpoint_id: str) -> key_agreement_client:
         if self._certificate_manager is None:
             cert_path = os.environ.get("E2E_CERTIFICATE_PATH")
-            if (self.ak is None or self.sk is None) and cert_path is None:
-                raise ArkAPIError("must set (ak and sk) or (E2E_CERTIFICATE_PATH) \
-                                  before get endpoint token.")
-            self._certificate_manager = E2ECertificateManager(self.ak, self.sk, self.region)
+            if (self.ak is None or self.sk is None) and cert_path is None and self.api_key is None:
+                raise ArkAPIError("must set (api_key) or (ak and sk) \
+                                  or (E2E_CERTIFICATE_PATH) before get endpoint token.")
+            self._certificate_manager = E2ECertificateManager(self.ak, self.sk, self.region, self._base_url, self.api_key)
         return self._certificate_manager.get(endpoint_id)
 
     def _get_bot_sts_token(self, bot_id: str):
@@ -164,6 +165,7 @@ def __init__(
             sk = os.environ.get("VOLC_SECRETKEY")
         if api_key is None:
             api_key = os.environ.get("ARK_API_KEY")
+        self._base_url = base_url
         self.ak = ak
         self.sk = sk
         self.api_key = api_key
@@ -200,10 +202,10 @@ def _get_endpoint_sts_token(self, endpoint_id: str):
     def _get_endpoint_certificate(self, endpoint_id: str) -> key_agreement_client:
         if self._certificate_manager is None:
             cert_path = os.environ.get("E2E_CERTIFICATE_PATH")
-            if (self.ak is None or self.sk is None) and cert_path is None:
-                raise ArkAPIError("must set (ak and sk) or (E2E_CERTIFICATE_PATH) \
-                                  before get endpoint token.")
-            self._certificate_manager = E2ECertificateManager(self.ak, self.sk, self.region)
+            if (self.ak is None or self.sk is None) and cert_path is None and self.api_key is None:
+                raise ArkAPIError("must set (api_key) or (ak and sk) \
+                                  or (E2E_CERTIFICATE_PATH) before get endpoint token.")
+            self._certificate_manager = E2ECertificateManager(self.ak, self.sk, self.region, self._base_url, self.api_key)
         return self._certificate_manager.get(endpoint_id)
 
     @property
@@ -301,22 +303,46 @@ def _load_api_key(self, ep: str, duration_seconds: int,
 
 class E2ECertificateManager(object):
 
-    def __init__(self, ak: str, sk: str, region: str):
+    class CertificateResponse():
+        Certificate: str
+        """The certificate content."""
+
+    def __init__(self, ak: str, sk: str, region: str, base_url: str | URL = BASE_URL, api_key: str | None = None):
         self._certificate_manager: Dict[str, key_agreement_client] = {}
 
-        import volcenginesdkcore
+        # local cache prepare
+        self._init_local_cert_cache()
 
+        # api instance prepare
+        import volcenginesdkcore
+        self._api_instance_enabled = True
+        if ak is None or sk is None:
+            self._api_instance_enabled = False
         configuration = volcenginesdkcore.Configuration()
         configuration.ak = ak
         configuration.sk = sk
         configuration.region = region
         configuration.schema = "https"
-
         volcenginesdkcore.Configuration.set_default(configuration)
         self.api_instance = volcenginesdkark.ARKApi()
 
+        # global cert path prepare
         self.cert_path = os.environ.get("E2E_CERTIFICATE_PATH")
 
+        # ark client prepare
+        self.client = Ark(
+            base_url=base_url,
+            api_key=api_key,
+            ak=ak, sk=sk,
+        )
+        self._e2e_uri = "/e2e/get/certificate"
+        self._x_session_token = {'X-Session-Token': self._e2e_uri}
+
+    def _load_cert_by_cert_path(self) -> str:
+        with open(self.cert_path, 'r') as f:
+            cert_pem = f.read()
+        return cert_pem
+
     def _load_cert_by_ak_sk(self, ep: str) -> str:
         get_endpoint_certificate_request = volcenginesdkark.GetEndpointCertificateRequest(
             id=ep
@@ -329,13 +355,52 @@ def _load_cert_by_ak_sk(self, ep: str) -> str:
 
         return resp.pca_instance_certificate
 
+    def _sync_load_cert_by_auth(self, ep: str) -> str:
+        try:  # try to make request with session header (used for header statistic)
+            resp = self.client.post(self._e2e_uri, options={"headers": self._x_session_token},
+                                    body={"model": ep}, cast_to=self.CertificateResponse)
+        except Exception as e:
+            raise ArkAPIError("Getting Certificate failed: %s\n" % e)
+        if 'error' in resp:
+            raise ArkAPIError("Getting Certificate failed: %s\n" % resp['error'])
+        return resp['Certificate']
+
+    def _save_cert_to_file(self, ep: str, cert_pem: str):
+        cert_file_path = os.path.join(self._cert_storage_path, f"{ep}.pem")
+        with open(cert_file_path, 'w') as f:
+            f.write(cert_pem)
+
+    def _load_cert_locally(self, ep: str) -> str | None:
+        cert_file_path = os.path.join(self._cert_storage_path, f"{ep}.pem")
+        if os.path.exists(cert_file_path):
+            last_modified_time = os.path.getmtime(cert_file_path)
+            current_time = time.time()
+            time_difference = current_time - last_modified_time
+            if time_difference <= self._cert_expiration_seconds:
+                with open(cert_file_path, 'r') as f:
+                    return f.read()
+            else:
+                os.remove(cert_file_path)
+        return None
+
+    def _init_local_cert_cache(self):
+        self._cert_storage_path = "/tmp/ark/certificates"
+        self._cert_expiration_seconds = 14 * 24 * 60 * 60  # 14 days
+
+        if not os.path.exists(self._cert_storage_path):
+            os.makedirs(self._cert_storage_path)
+
     def get(self, ep: str) -> key_agreement_client:
         if ep not in self._certificate_manager:
-            if self.cert_path is not None:
-                with open(self.cert_path, 'r') as f:
-                    cert_pem = f.read()
-            else:
-                cert_pem = self._load_cert_by_ak_sk(ep)
+            cert_pem = self._load_cert_locally(ep)
+            if cert_pem is None:
+                if self.cert_path is not None:
+                    cert_pem = self._load_cert_by_cert_path()
+                elif self._api_instance_enabled:
+                    cert_pem = self._load_cert_by_ak_sk(ep)
+                else:
+                    cert_pem = self._sync_load_cert_by_auth(ep)
+                self._save_cert_to_file(ep, cert_pem)
             self._certificate_manager[ep] = key_agreement_client(
                 certificate_pem_string=cert_pem
             )

volcenginesdkarkruntime/_utils/_key_agreement.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import base64
-
+from typing import Tuple
 
 def aes_gcm_encrypt_bytes(key: bytes, iv: bytes, plain_bytes: bytes, associated_data: bytes = b"") -> bytes:
     # aes_gcm_encrypt_bytes encrypt message using AES-GCM
@@ -85,7 +85,7 @@ def __init__(self, certificate_pem_string: str) -> None:
         self._public_key = ec.EllipticCurvePublicNumbers(
             cert_pub.x, cert_pub.y, self._curve).public_key()
 
-    def encrypt_string(self, plaintext: str) -> tuple[bytes, bytes, str, str]:
+    def encrypt_string(self, plaintext: str) -> Tuple[bytes, bytes, str, str]:
         """encrypt_string encrypt plaintext with ECIES DH protocol
         """
         key, nonce, token = self.generate_ecies_key_pair()
@@ -106,7 +106,7 @@ def decrypt_string_with_key(self, key: bytes, nonce: bytes, ciphertext: str) ->
         # Decrypt message using AES-GCM
         return aes_gcm_decrypt_base64_string(key, nonce, ciphertext)
 
-    def generate_ecies_key_pair(self) -> tuple[bytes, bytes, str]:
+    def generate_ecies_key_pair(self) -> Tuple[bytes, bytes, str]:
         """generate_ecies_key_pair generate ECIES key pair
         """
         from cryptography.hazmat.primitives import hashes

volcenginesdkarkruntime/resources/chat/completions.py

@@ -33,6 +33,31 @@
 
 __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'))))
+
 class Completions(SyncAPIResource):
     @cached_property
     def with_raw_response(self) -> CompletionsWithRawResponse:
@@ -42,38 +67,12 @@ def with_raw_response(self) -> CompletionsWithRawResponse:
     def with_streaming_response(self) -> CompletionsWithStreamingResponse:
         return CompletionsWithStreamingResponse(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):
-                    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 _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
-        self._process_messages(messages, lambda x: client.encrypt_string_with_key(_crypto_key,
-                                                                                  _crypto_nonce,
-                                                                                  x))
+        _process_messages(messages, lambda x: client.encrypt_string_with_key(_crypto_key, _crypto_nonce, x))
         return _crypto_key, _crypto_nonce
 
     def _decrypt_chunk(self, key: bytes, nonce: bytes, resp: Stream[ChatCompletionChunk]) -> Iterator[ChatCompletionChunk]:
@@ -180,38 +179,12 @@ def with_raw_response(self) -> AsyncCompletionsWithRawResponse:
     def with_streaming_response(self) -> AsyncCompletionsWithStreamingResponse:
         return AsyncCompletionsWithStreamingResponse(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):
-                    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 _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
-        self._process_messages(messages, lambda x: client.encrypt_string_with_key(_crypto_key,
-                                                                                  _crypto_nonce,
-                                                                                  x))
+        _process_messages(messages, lambda x: client.encrypt_string_with_key(_crypto_key, _crypto_nonce, x))
         return _crypto_key, _crypto_nonce
 
     async def _decrypt_chunk(self, key: bytes, nonce: bytes, resp: AsyncStream[ChatCompletionChunk]