88from typing_extensions import Literal
99
1010import httpx
11+ import warnings
1112
1213from ..._exceptions import ArkAPITimeoutError , ArkAPIConnectionError , ArkAPIStatusError
1314from ..._types import Body , Query , Headers
3738__all__ = ["Completions" , "AsyncCompletions" ]
3839
3940
41+ def _process_messages (
42+ messages : Iterable [ChatCompletionMessageParam ], f : Callable [[str ], str ]
43+ ):
44+ for message in messages :
45+ if message .get ("content" , None ) is not None :
46+ current_content = message .get ("content" )
47+ if isinstance (current_content , str ):
48+ message ["content" ] = f (current_content )
49+ elif isinstance (current_content , Iterable ):
50+ for part in current_content :
51+ if part .get ("type" , None ) == "text" :
52+ part ["text" ] = f (part ["text" ])
53+ elif part .get ("type" , None ) == "image_url" :
54+ if part ["image_url" ]["url" ].startswith ("data:" ):
55+ part ["image_url" ]["url" ] = f (part ["image_url" ]["url" ])
56+ else :
57+ warnings .warn (
58+ "encryption is not supported for image url, "
59+ "please use base64 image if you want encryption"
60+ )
61+ else :
62+ raise TypeError (
63+ "encryption is not supported for content type {}" .format (
64+ type (part )
65+ )
66+ )
67+ else :
68+ raise TypeError (
69+ "encryption is not supported for content type {}" .format (
70+ type (message .get ("content" ))
71+ )
72+ )
73+
4074def _calculate_retry_timeout (retry_times ) -> float :
4175 nbRetries = min (retry_times , MAX_RETRY_DELAY / INITIAL_RETRY_DELAY )
4276 sleep_seconds = min (INITIAL_RETRY_DELAY * pow (2 , nbRetries ), MAX_RETRY_DELAY )
@@ -79,27 +113,6 @@ class Completions(SyncAPIResource):
79113 def with_raw_response (self ) -> CompletionsWithRawResponse :
80114 return CompletionsWithRawResponse (self )
81115
82- def _process_messages (
83- self , messages : Iterable [ChatCompletionMessageParam ], f : Callable [[str ], str ]
84- ):
85- for message in messages :
86- if message .get ("content" , None ) is not None :
87- current_content = message .get ("content" )
88- if isinstance (current_content , str ):
89- message ["content" ] = f (current_content )
90- elif isinstance (current_content , Iterable ):
91- raise TypeError (
92- "content type {} is not supported end-to-end encryption" .format (
93- type (message .get ("content" ))
94- )
95- )
96- else :
97- raise TypeError (
98- "content type {} is not supported end-to-end encryption" .format (
99- type (message .get ("content" ))
100- )
101- )
102-
103116 def _encrypt (
104117 self ,
105118 model : str ,
@@ -109,7 +122,7 @@ def _encrypt(
109122 client = self ._client ._get_endpoint_certificate (model )
110123 _crypto_key , _crypto_nonce , session_token = client .generate_ecies_key_pair ()
111124 extra_headers ["X-Session-Token" ] = session_token
112- self . _process_messages (
125+ _process_messages (
113126 messages ,
114127 lambda x : client .encrypt_string_with_key (_crypto_key , _crypto_nonce , x ),
115128 )
@@ -244,21 +257,6 @@ class AsyncCompletions(AsyncAPIResource):
244257 def with_raw_response (self ) -> AsyncCompletionsWithRawResponse :
245258 return AsyncCompletionsWithRawResponse (self )
246259
247- def _process_messages (
248- self , messages : Iterable [ChatCompletionMessageParam ], f : Callable [[str ], str ]
249- ):
250- for message in messages :
251- if message .get ("content" , None ) is not None :
252- current_content = message .get ("content" )
253- if isinstance (current_content , str ):
254- message ["content" ] = f (current_content )
255- else :
256- raise TypeError (
257- "content type {} is not supported end-to-end encryption" .format (
258- type (message .get ("content" ))
259- )
260- )
261-
262260 def _encrypt (
263261 self ,
264262 model : str ,
@@ -268,7 +266,7 @@ def _encrypt(
268266 client = self ._client ._get_endpoint_certificate (model )
269267 _crypto_key , _crypto_nonce , session_token = client .generate_ecies_key_pair ()
270268 extra_headers ["X-Session-Token" ] = session_token
271- self . _process_messages (
269+ _process_messages (
272270 messages ,
273271 lambda x : client .encrypt_string_with_key (_crypto_key , _crypto_nonce , x ),
274272 )
@@ -379,9 +377,7 @@ async def create(
379377 else :
380378 raise err
381379 if is_encrypt :
382- resp = self ._decrypt (e2e_key , e2e_nonce , resp )
383- if is_encrypt :
384- resp = await self ._decrypt (e2e_key , e2e_nonce , resp )
380+ resp = await self ._decrypt (e2e_key , e2e_nonce , resp )
385381 return resp
386382
387383 def _get_request_last_time (self , timeout ):
0 commit comments