feat(ark e2e): fix bug

Change-Id: I262ff4a760b35a9db0977a2fdf771ef8666079c7

Author: liuhuiqi.7

volcenginesdkarkruntime/_streaming.py

@@ -168,7 +168,7 @@ def __init__(
             self._iterator = self.__stream__()
 
     @classmethod
-    def _make_stream_from_iterator(cls, iterator: Iterator[_T]) -> Stream[_T]:
+    def _make_stream_from_iterator(cls, iterator: AsyncIterator[_T]) -> AsyncStream[_T]:
         return AsyncStream(cast_to=None, response=None, client=None, iterator=iterator)
 
     async def __anext__(self) -> _T:

volcenginesdkarkruntime/resources/chat/completions.py

@@ -118,7 +118,7 @@ def create(
         timeout: float | httpx.Timeout | None = None,
     ) -> ChatCompletion | Stream[ChatCompletionChunk]:
         is_encrypt = False
-        if extra_headers is not None and extra_headers[ARK_E2E_ENCRYPTION_HEADER] == 'true':
+        if extra_headers is not None and extra_headers.get(ARK_E2E_ENCRYPTION_HEADER, None) == 'true':
             is_encrypt = True
             e2e_key, e2e_nonce = self._encrypt(model, messages, extra_headers)
 
@@ -211,7 +211,7 @@ async def _decrypt(self, key: bytes, nonce: bytes, resp: ChatCompletion | AsyncS
                     if choice.message.content is not None:
                         choice.message.content = aes_gcm_decrypt_base64_string(key, nonce, choice.message.content)
             resp.choices[0] = choice
-            return await resp
+            return resp
         else:
             return AsyncStream._make_stream_from_iterator(self._decrypt_chunk(key, nonce, resp))
 
@@ -245,7 +245,7 @@ async def create(
         timeout: float | httpx.Timeout | None = None,
     ) -> ChatCompletion | AsyncStream[ChatCompletionChunk]:
         is_encrypt = False
-        if extra_headers is not None and extra_headers[ARK_E2E_ENCRYPTION_HEADER] == 'true':
+        if extra_headers is not None and extra_headers.get(ARK_E2E_ENCRYPTION_HEADER, None) == 'true':
             is_encrypt = True
             e2e_key, e2e_nonce = self._encrypt(model, messages, extra_headers)