feat(ark e2e): support decrypt

Change-Id: I7b07a4aaefb9cdadcaf6153435f5f9e0751290b5

Author: liuhuiqi.7

volcenginesdkarkruntime/_utils/_key_agreement.py

@@ -86,11 +86,17 @@ def encrypt_string(self, plaintext: str) -> tuple[bytes, bytes, str, str]:
         return key, nonce, token, ciphertext
 
     def encrypt_string_with_key(self, key: bytes, nonce: bytes, plaintext: str) -> str:
-        """encrypt_string encrypt plaintext with ECIES DH protocol
+        """encrypt_string_with_key encrypt plaintext with ECIES DH protocol
         """
         # Encrypt message using AES-GCM
         ciphertext = aes_gcm_encrypt_base64_string(key, nonce, plaintext)
         return ciphertext
+    
+    def decrypt_string_with_key(self, key: bytes, nonce: bytes, ciphertext: str) -> str:
+        """decrypt_string_with_key decrypt ciphertext with ECIES DH protocol
+        """
+        # Decrypt message using AES-GCM
+        return aes_gcm_decrypt_base64_string(key, nonce, ciphertext)
 
     def generate_ecies_key_pair(self) -> tuple[bytes, bytes, str]:
         """generate_ecies_key_pair generate ECIES key pair

volcenginesdkarkruntime/resources/chat/completions.py

@@ -66,12 +66,22 @@ def _process_messages(self, messages: Iterable[ChatCompletionMessageParam],
 
     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, completion: ChatCompletion):
+        if completion.choices is not None:
+            for choice in completion.choices:
+                if choice.message.content is not None:
+                    if isinstance(choice.message.content, str):
+                        choice.message.content = self._ka_client.decrypt_string_with_key(self._crypto_key,
+                                                                                        self._crypto_nonce,
+                                                                                        choice.message.content)
+
     @with_sts_token
     def create(
         self,
@@ -104,9 +114,8 @@ def create(
     ) -> ChatCompletion | Stream[ChatCompletionChunk]:
         if is_encrypt:
             self._encrypt(model, messages, extra_headers)
-        print(messages)
 
-        return self._post(
+        resp = self._post(
             "/chat/completions",
             body={
                 "messages": messages,
@@ -141,6 +150,10 @@ def create(
             stream_cls=Stream[ChatCompletionChunk],
         )
 
+        if is_encrypt:
+            return self._decrypt(resp)
+        return resp
+
 
 class AsyncCompletions(AsyncAPIResource):
     @cached_property