feat(ark e2e): support async

Change-Id: If3cae5914537c4346aff1aa81ef995801f4a96bc

Author: liuhuiqi.7

volcenginesdkarkruntime/_client.py

@@ -178,6 +178,7 @@ def __init__(
 
         self._default_stream_cls = Stream
         self._sts_token_manager: StsTokenManager | None = None
+        self._certificate_manager: E2ECertificateManager | None = None
 
         self.chat = resources.AsyncChat(self)
         self.bot_chat = resources.AsyncBotChat(self)
@@ -192,6 +193,15 @@ 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) -> 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)
+        return self._certificate_manager.get(endpoint_id)
+
     @property
     def auth_headers(self) -> dict[str, str]:
         api_key = self.api_key

volcenginesdkarkruntime/resources/chat/completions.py

@@ -47,21 +47,23 @@ def _process_messages(self, messages: Iterable[ChatCompletionMessageParam],
                 if isinstance(message.get("content"), str):
                     message["content"] = f(message.get("content"))
                 elif isinstance(message.get("content"), Iterable):
-                    content = message.get("content")
-                    for i, c in enumerate(content):
-                        if not isinstance(c, Dict):
-                            raise TypeError("content type {} is not supported end-to-end encryption".
-                                            format(type(c)))
-                        for key in c.keys():
-                            if key == 'type':
-                                continue
-                            if isinstance(c[key], str):
-                                content[i][key] = f(c[key])
-                            if isinstance(c[key], Dict):
-                                for k in c[key].keys():
-                                    if isinstance(c[key][k], str):
-                                        content[i][key][k] = f(c[key][k])
-                    message["content"] = content
+                    raise TypeError("content type {} is not supported end-to-end encryption".
+                                    format(type(message.get('content'))))
+                    # content = message.get("content")
+                    # for i, c in enumerate(content):
+                    #     if not isinstance(c, Dict):
+                    #         raise TypeError("content type {} is not supported end-to-end encryption".
+                    #                         format(type(c)))
+                    #     for key in c.keys():
+                    #         if key == 'type':
+                    #             continue
+                    #         if isinstance(c[key], str):
+                    #             content[i][key] = f(c[key])
+                    #         if isinstance(c[key], Dict):
+                    #             for k in c[key].keys():
+                    #                 if isinstance(c[key][k], str):
+                    #                     content[i][key][k] = f(c[key][k])
+                    # message["content"] = content
                 else:
                     raise TypeError("content type {} is not supported end-to-end encryption".
                                     format(type(message.get('content'))))
@@ -160,6 +162,28 @@ 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:
+                if isinstance(message.get("content"), str):
+                    message["content"] = f(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):
+        client = self._client._get_endpoint_certificate(model)
+        self._ka_client = client
+        self._crypto_key, self._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(self._crypto_key,
+                                                                                  self._crypto_nonce,
+                                                                                  x))
+
+    def decrypt(self, ciphertext: str) -> str:
+        return self._ka_client.decrypt_string_with_key(self._crypto_key, self._crypto_nonce, ciphertext)
+
     @async_with_sts_token
     async def create(
         self,
@@ -188,7 +212,13 @@ async def create(
         extra_query: Query | None = None,
         extra_body: Body | None = None,
         timeout: float | httpx.Timeout | None = None,
+        is_encrypt: bool | None = None,
     ) -> ChatCompletion | AsyncStream[ChatCompletionChunk]:
+        if is_encrypt:
+            if extra_headers is None:
+                extra_headers = dict()
+            self._encrypt(model, messages, extra_headers)
+
         return await self._post(
             "/chat/completions",
             body={