feat(crypto): fit batch encryption

Author: liuhuiqi.7

volcenginesdkarkruntime/resources/batch/chat/completions.py

@@ -1,18 +1,16 @@
 from __future__ import annotations
 
-import warnings
-from typing import Callable, Dict, Iterable, List, Optional, Union
+from typing import Dict, Iterable, List, Optional, Union
 
 import httpx
 from typing_extensions import Literal
 
 from ...._base_client import make_request_options
 from ...._compat import cached_property
-from ...._constants import ARK_E2E_ENCRYPTION_HEADER
 from ...._resource import AsyncAPIResource, SyncAPIResource
 from ...._types import Body, Headers, Query
-from ...._utils import async_with_sts_token, deepcopy_minimal, with_sts_token
-from ...._utils._key_agreement import aes_gcm_decrypt_base64_string
+from ...._utils import async_with_sts_token, with_sts_token
+from ...encryption import with_e2e_encryption, async_with_e2e_encryption
 from ....types.chat import (
     ChatCompletion,
     ChatCompletionMessageParam,
@@ -29,98 +27,13 @@
 __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:
         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,
         *,
@@ -151,15 +64,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)
-
         deadline = get_request_last_time(self._client, timeout)
         breaker = self._client.get_model_breaker(model)
 
@@ -200,10 +104,6 @@ def create(
             ),
             cast_to=ChatCompletion,
         )
-
-        if is_encrypt:
-            resp = self._decrypt(e2e_key, e2e_nonce, resp)
-            return resp
         return resp
 
 
@@ -212,38 +112,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,
         *,
@@ -274,15 +144,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)
-
         deadline = get_request_last_time(self._client, timeout)
         breaker = await self._client.get_model_breaker(model)
 
@@ -324,6 +185,4 @@ async def create(
             cast_to=ChatCompletion,
         )
 
-        if is_encrypt:
-            resp = await self._decrypt(e2e_key, e2e_nonce, resp)
         return resp