feat: support thinking

Author: 潘婉宁

volcenginesdkarkruntime/_base_client.py

@@ -24,7 +24,7 @@
     _DEFAULT_MANDATORY_REFRESH_TIMEOUT,
     _DEFAULT_STS_TIMEOUT,
     _DEFAULT_RESOURCE_TYPE,
-    DEFAULT_TIMEOUT
+    DEFAULT_TIMEOUT,
 )
 from ._streaming import Stream
 
@@ -84,7 +84,9 @@ def __init__(
         self.api_key = api_key
         self.region = region
 
-        assert (api_key is not None) or (ak is not None and sk is not None), "you need to support api_key or ak&sk"
+        assert (api_key is not None) or (ak is not None and sk is not None), (
+            "you need to support api_key or ak&sk"
+        )
 
         super().__init__(
             base_url=base_url,
@@ -120,10 +122,18 @@ def _get_endpoint_sts_token(self, endpoint_id: str):
     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 and self.api_key is None:
-                raise ArkAPIError("must set (api_key) or (ak and sk) \
-                                  or (E2E_CERTIFICATE_PATH) before get endpoint token.")
-            self._certificate_manager = E2ECertificateManager(self.ak, self.sk, self.region, self._base_url, self.api_key)
+            if (
+                (self.ak is None or self.sk is None)
+                and cert_path is None
+                and self.api_key is None
+            ):
+                raise ArkAPIError(
+                    "must set (api_key) or (ak and sk) \
+                                  or (E2E_CERTIFICATE_PATH) before get endpoint token."
+                )
+            self._certificate_manager = E2ECertificateManager(
+                self.ak, self.sk, self.region, self._base_url, self.api_key
+            )
         return self._certificate_manager.get(endpoint_id)
 
     def _get_bot_sts_token(self, bot_id: str):
@@ -142,6 +152,7 @@ def get_model_breaker(self, model_name: str) -> ModelBreaker:
         with self.model_breaker_lock:
             return self.model_breaker_map[model_name]
 
+
 class AsyncArk(AsyncAPIClient):
     chat: resources.AsyncChat
     bot_chat: resources.AsyncBotChat
@@ -168,15 +179,15 @@ def __init__(
     ) -> None:
         """init async ark client, this client is thread unsafe
 
-            Args:
-                ak: access key id
-                sk: secret access key
-                api_key: api key，this api key will not be refreshed
-                timeout: timeout of client. default httpx.Timeout(timeout=60.0, connect=60.0)
-                max_retries: times of retry when request failed. default 1
-                http_client: specify customized http_client
-            Returns:
-                async ark client
+        Args:
+            ak: access key id
+            sk: secret access key
+            api_key: api key，this api key will not be refreshed
+            timeout: timeout of client. default httpx.Timeout(timeout=60.0, connect=60.0)
+            max_retries: times of retry when request failed. default 1
+            http_client: specify customized http_client
+        Returns:
+            async ark client
         """
 
         if ak is None:
@@ -191,7 +202,9 @@ def __init__(
         self.api_key = api_key
         self.region = region
 
-        assert (api_key is not None) or (ak is not None and sk is not None), "you need to support api_key or ak&sk"
+        assert (api_key is not None) or (ak is not None and sk is not None), (
+            "you need to support api_key or ak&sk"
+        )
 
         super().__init__(
             base_url=base_url,
@@ -227,10 +240,18 @@ def _get_endpoint_sts_token(self, endpoint_id: str):
     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 and self.api_key is None:
-                raise ArkAPIError("must set (api_key) or (ak and sk) \
-                                  or (E2E_CERTIFICATE_PATH) before get endpoint token.")
-            self._certificate_manager = E2ECertificateManager(self.ak, self.sk, self.region, self._base_url, self.api_key)
+            if (
+                (self.ak is None or self.sk is None)
+                and cert_path is None
+                and self.api_key is None
+            ):
+                raise ArkAPIError(
+                    "must set (api_key) or (ak and sk) \
+                                  or (E2E_CERTIFICATE_PATH) before get endpoint token."
+                )
+            self._certificate_manager = E2ECertificateManager(
+                self.ak, self.sk, self.region, self._base_url, self.api_key
+            )
         return self._certificate_manager.get(endpoint_id)
 
     @property
@@ -252,7 +273,9 @@ class StsTokenManager(object):
     _mandatory_refresh_timeout: int = _DEFAULT_MANDATORY_REFRESH_TIMEOUT
 
     def __init__(self, ak: str, sk: str, region: str):
-        self._endpoint_sts_tokens: Dict[str, Tuple[str, int]] = defaultdict(lambda: ("", 0))
+        self._endpoint_sts_tokens: Dict[str, Tuple[str, int]] = defaultdict(
+            lambda: ("", 0)
+        )
         self._refresh_lock = threading.Lock()
 
         import volcenginesdkcore
@@ -272,10 +295,19 @@ def _need_refresh(self, ep: str, refresh_in: int | None = None) -> bool:
 
         return self._endpoint_sts_tokens[ep][1] - time.time() < refresh_in
 
-    def _protected_refresh(self, ep: str, ttl: int = _DEFAULT_STS_TIMEOUT, is_mandatory: bool = False,
-                           resource_type: str = _DEFAULT_RESOURCE_TYPE):
+    def _protected_refresh(
+        self,
+        ep: str,
+        ttl: int = _DEFAULT_STS_TIMEOUT,
+        is_mandatory: bool = False,
+        resource_type: str = _DEFAULT_RESOURCE_TYPE,
+    ):
         if ttl < self._advisory_refresh_timeout * 2:
-            raise ArkAPIError("ttl should not be under {} seconds.".format(self._advisory_refresh_timeout * 2))
+            raise ArkAPIError(
+                "ttl should not be under {} seconds.".format(
+                    self._advisory_refresh_timeout * 2
+                )
+            )
 
         try:
             api_key, expired_time = self._load_api_key(
@@ -301,7 +333,9 @@ def _refresh(self, ep: str, resource_type: str = _DEFAULT_RESOURCE_TYPE):
                     ep, self._mandatory_refresh_timeout
                 )
 
-                self._protected_refresh(ep, is_mandatory=is_mandatory_refresh, resource_type=resource_type)
+                self._protected_refresh(
+                    ep, is_mandatory=is_mandatory_refresh, resource_type=resource_type
+                )
                 return
             finally:
                 self._refresh_lock.release()
@@ -310,14 +344,20 @@ def _refresh(self, ep: str, resource_type: str = _DEFAULT_RESOURCE_TYPE):
                 if not self._need_refresh(ep, self._mandatory_refresh_timeout):
                     return
 
-                self._protected_refresh(ep, is_mandatory=True, resource_type=resource_type)
+                self._protected_refresh(
+                    ep, is_mandatory=True, resource_type=resource_type
+                )
 
     def get(self, ep: str, resource_type: str = _DEFAULT_RESOURCE_TYPE) -> str:
         self._refresh(ep, resource_type=resource_type)
         return self._endpoint_sts_tokens[ep][0]
 
-    def _load_api_key(self, ep: str, duration_seconds: int,
-                      resource_type: str = _DEFAULT_RESOURCE_TYPE) -> Tuple[str, int]:
+    def _load_api_key(
+        self,
+        ep: str,
+        duration_seconds: int,
+        resource_type: str = _DEFAULT_RESOURCE_TYPE,
+    ) -> Tuple[str, int]:
         get_api_key_request = volcenginesdkark.GetApiKeyRequest(
             duration_seconds=duration_seconds,
             resource_type=resource_type,
@@ -331,19 +371,26 @@ def _load_api_key(self, ep: str, duration_seconds: int,
 
 
 class E2ECertificateManager(object):
-
-    class CertificateResponse():
+    class CertificateResponse:
         Certificate: str
         """The certificate content."""
 
-    def __init__(self, ak: str, sk: str, region: str, base_url: str | URL = BASE_URL, api_key: str | None = None):
+    def __init__(
+        self,
+        ak: str,
+        sk: str,
+        region: str,
+        base_url: str | URL = BASE_URL,
+        api_key: str | None = None,
+    ):
         self._certificate_manager: Dict[str, key_agreement_client] = {}
 
         # local cache prepare
         self._init_local_cert_cache()
 
         # api instance prepare
         import volcenginesdkcore
+
         configuration = volcenginesdkcore.Configuration()
         configuration.ak = ak
         configuration.sk = sk
@@ -365,38 +412,47 @@ def __init__(self, ak: str, sk: str, region: str, base_url: str | URL = BASE_URL
                 api_key=api_key,
             )
         self._e2e_uri = "/e2e/get/certificate"
-        self._x_session_token = {'X-Session-Token': self._e2e_uri}
+        self._x_session_token = {"X-Session-Token": self._e2e_uri}
 
     def _load_cert_by_cert_path(self) -> str:
-        with open(self.cert_path, 'r') as f:
+        with open(self.cert_path, "r") as f:
             cert_pem = f.read()
         return cert_pem
 
     def _load_cert_by_ak_sk(self, ep: str) -> str:
-        get_endpoint_certificate_request = volcenginesdkark.GetEndpointCertificateRequest(
-            id=ep
+        get_endpoint_certificate_request = (
+            volcenginesdkark.GetEndpointCertificateRequest(id=ep)
         )
         try:
-            resp: volcenginesdkark.GetEndpointCertificateResponse = self.api_instance.get_endpoint_certificate(
-                get_endpoint_certificate_request)
+            resp: volcenginesdkark.GetEndpointCertificateResponse = (
+                self.api_instance.get_endpoint_certificate(
+                    get_endpoint_certificate_request
+                )
+            )
         except ApiException as e:
-            raise ArkAPIError("Getting model vendor encryption certificate failed: %s\n" % e)
+            raise ArkAPIError(
+                "Getting model vendor encryption certificate failed: %s\n" % e
+            )
 
         return resp.pca_instance_certificate
 
     def _sync_load_cert_by_auth(self, ep: str) -> str:
         try:  # try to make request with session header (used for header statistic)
-            resp = self.client.post(self._e2e_uri, options={"headers": self._x_session_token},
-                                    body={"model": ep}, cast_to=self.CertificateResponse)
+            resp = self.client.post(
+                self._e2e_uri,
+                options={"headers": self._x_session_token},
+                body={"model": ep},
+                cast_to=self.CertificateResponse,
+            )
         except Exception as e:
             raise ArkAPIError("Getting Certificate failed: %s\n" % e)
-        if 'error' in resp:
-            raise ArkAPIError("Getting Certificate failed: %s\n" % resp['error'])
-        return resp['Certificate']
+        if "error" in resp:
+            raise ArkAPIError("Getting Certificate failed: %s\n" % resp["error"])
+        return resp["Certificate"]
 
     def _save_cert_to_file(self, ep: str, cert_pem: str):
         cert_file_path = os.path.join(self._cert_storage_path, f"{ep}.pem")
-        with open(cert_file_path, 'w') as f:
+        with open(cert_file_path, "w") as f:
             f.write(cert_pem)
 
     def _load_cert_locally(self, ep: str) -> str | None:
@@ -406,7 +462,7 @@ def _load_cert_locally(self, ep: str) -> str | None:
             current_time = time.time()
             time_difference = current_time - last_modified_time
             if time_difference <= self._cert_expiration_seconds:
-                with open(cert_file_path, 'r') as f:
+                with open(cert_file_path, "r") as f:
                     return f.read()
             else:
                 os.remove(cert_file_path)

volcenginesdkarkruntime/_client.py

@@ -24,9 +24,7 @@
     def parse_date(value: date | StrBytesIntFloat) -> date:  # noqa: ARG001
         ...
 
-    def parse_datetime(
-        value: Union[datetime, StrBytesIntFloat]
-    ) -> datetime:  # noqa: ARG001
+    def parse_datetime(value: Union[datetime, StrBytesIntFloat]) -> datetime:  # noqa: ARG001
         ...
 
     def get_args(t: type[Any]) -> tuple[Any, ...]:  # noqa: ARG001
@@ -87,9 +85,7 @@ def parse_obj(model: type[_ModelT], value: object) -> _ModelT:
     if PYDANTIC_V2:
         return model.model_validate(value)
     else:
-        return cast(
-            _ModelT, model.parse_obj(value)
-        )  # pyright: ignore[reportDeprecated, reportUnnecessaryCast]
+        return cast(_ModelT, model.parse_obj(value))  # pyright: ignore[reportDeprecated, reportUnnecessaryCast]
 
 
 def field_is_required(field: FieldInfo) -> bool:

volcenginesdkarkruntime/_compat.py

@@ -13,7 +13,9 @@
 DEFAULT_TIMEOUT_SECONDS = 600.0
 DEFAULT_CONNECT_TIMEOUT_SECONDS = 60.0
 # default timeout is 1 minutes
-DEFAULT_TIMEOUT = httpx.Timeout(timeout=DEFAULT_TIMEOUT_SECONDS, connect=DEFAULT_CONNECT_TIMEOUT_SECONDS)
+DEFAULT_TIMEOUT = httpx.Timeout(
+    timeout=DEFAULT_TIMEOUT_SECONDS, connect=DEFAULT_CONNECT_TIMEOUT_SECONDS
+)
 
 DEFAULT_MAX_RETRIES = 2
 DEFAULT_CONNECTION_LIMITS = httpx.Limits(

volcenginesdkarkruntime/_constants.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import Any, Optional, cast, Dict
+from typing import Optional
 from typing_extensions import Literal
 
 import httpx
@@ -130,45 +130,31 @@ def __init__(self, request: httpx.Request, request_id: str) -> None:
 
 
 class ArkBadRequestError(ArkAPIStatusError):
-    status_code: Literal[400] = (
-        400  # pyright: ignore[reportIncompatibleVariableOverride]
-    )
+    status_code: Literal[400] = 400  # pyright: ignore[reportIncompatibleVariableOverride]
 
 
 class ArkAuthenticationError(ArkAPIStatusError):
-    status_code: Literal[401] = (
-        401  # pyright: ignore[reportIncompatibleVariableOverride]
-    )
+    status_code: Literal[401] = 401  # pyright: ignore[reportIncompatibleVariableOverride]
 
 
 class ArkPermissionDeniedError(ArkAPIStatusError):
-    status_code: Literal[403] = (
-        403  # pyright: ignore[reportIncompatibleVariableOverride]
-    )
+    status_code: Literal[403] = 403  # pyright: ignore[reportIncompatibleVariableOverride]
 
 
 class ArkNotFoundError(ArkAPIStatusError):
-    status_code: Literal[404] = (
-        404  # pyright: ignore[reportIncompatibleVariableOverride]
-    )
+    status_code: Literal[404] = 404  # pyright: ignore[reportIncompatibleVariableOverride]
 
 
 class ArkConflictError(ArkAPIStatusError):
-    status_code: Literal[409] = (
-        409  # pyright: ignore[reportIncompatibleVariableOverride]
-    )
+    status_code: Literal[409] = 409  # pyright: ignore[reportIncompatibleVariableOverride]
 
 
 class ArkUnprocessableEntityError(ArkAPIStatusError):
-    status_code: Literal[422] = (
-        422  # pyright: ignore[reportIncompatibleVariableOverride]
-    )
+    status_code: Literal[422] = 422  # pyright: ignore[reportIncompatibleVariableOverride]
 
 
 class ArkRateLimitError(ArkAPIStatusError):
-    status_code: Literal[429] = (
-        429  # pyright: ignore[reportIncompatibleVariableOverride]
-    )
+    status_code: Literal[429] = 429  # pyright: ignore[reportIncompatibleVariableOverride]
 
 
 class ArkInternalServerError(ArkAPIStatusError):

volcenginesdkarkruntime/_exceptions.py

@@ -143,7 +143,7 @@ def to_json(
     @override
     def __str__(self) -> str:
         # mypy complains about an invalid self arg
-        return f'{self.__repr_name__()}({self.__repr_str__(", ")})'  # type: ignore[misc]
+        return f"{self.__repr_name__()}({self.__repr_str__(', ')})"  # type: ignore[misc]
 
     # Override the 'construct' method in a way that supports recursive parsing without validation.
     # Based on https://github.com/samuelcolvin/pydantic/issues/1168#issuecomment-817742836.

volcenginesdkarkruntime/_models.py

@@ -59,9 +59,7 @@ def construct(  # type: ignore
         }
         if PYDANTIC_V2:
             return super().model_construct(_fields_set, **kwargs)
-        return cast(
-            RequestOptions, super().construct(_fields_set, **kwargs)
-        )  # pyright: ignore[reportDeprecated]
+        return cast(RequestOptions, super().construct(_fields_set, **kwargs))  # pyright: ignore[reportDeprecated]
 
     if not TYPE_CHECKING:
         # type checkers incorrectly complain about this assignment