feat: 增加refresh

Author: hexiaochun

volcenginesdkarkruntime/_base_client.py

@@ -1,15 +1,29 @@
 from __future__ import annotations
 
+import logging
 import os
-from typing import Dict
+import threading
+import time
+from collections import defaultdict
 
-import httpx
-from httpx import Timeout, URL
+from httpx import Timeout, URL, Client, AsyncClient
+from typing import Dict, Tuple
 
-from . import resources, _exceptions
+from volcenginesdkcore.rest import ApiException
+from ._exceptions import ArkAPIError
+
+import volcenginesdkark
+
+from . import resources
 from ._base_client import SyncAPIClient, AsyncAPIClient
-from ._constants import DEFAULT_MAX_RETRIES, BASE_URL
-from ._exceptions import ArkError, ArkAPIStatusError
+from ._constants import (
+    DEFAULT_MAX_RETRIES,
+    BASE_URL,
+    _DEFAULT_ADVISORY_REFRESH_TIMEOUT,
+    _DEFAULT_MANDATORY_REFRESH_TIMEOUT,
+    _DEFAULT_STS_TIMEOUT,
+    DEFAULT_TIMEOUT
+)
 from ._streaming import Stream
 
 __all__ = ["Ark", "AsyncArk"]
@@ -22,32 +36,51 @@ def __init__(
         self,
         *,
         base_url: str | URL = BASE_URL,
-        api_key: str | None = None,
-        timeout: float | Timeout | None = None,
+        ak: str | None = None,
+        sk: str | None = None,
+        timeout: float | Timeout | None = DEFAULT_TIMEOUT,
         max_retries: int = DEFAULT_MAX_RETRIES,
-        default_query: Dict[str, object] | None = None,
-        http_client: httpx.Client | None = None,
+        http_client: Client | None = None,
     ) -> None:
-        if api_key is None:
-            api_key = os.environ.get("ARK_API_KEY")
-        self.api_key = api_key
+        """init ark client, this client is thread unsafe. If need to use in multi thread, init a new `Ark` client in
+        each thread
+
+            Args:
+                ak: access key id
+                sk: secret access key
+                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:
+                ark client
+        """
+
+        if ak is None:
+            ak = os.environ.get("VOLC_ACCESSKEY")
+        if sk is None:
+            sk = os.environ.get("VOLC_SECRETKEY")
+        self.ak = ak
+        self.sk = sk
 
         super().__init__(
             base_url=base_url,
             max_retries=max_retries,
             timeout=timeout,
             http_client=http_client,
-            custom_query=default_query,
+            custom_query=None,
         )
 
         self._default_stream_cls = Stream
+        self._sts_token_manager: StsTokenManager | None = None
 
         self.chat = resources.Chat(self)
 
-    @property
-    def auth_headers(self) -> dict[str, str]:
-        api_key = self.api_key
-        return {"Authorization": f"Bearer {api_key}"}
+    def _get_endpoint_sts_token(self, endpoint_id: str):
+        if self._sts_token_manager is None:
+            if self.ak is None or self.sk is None:
+                raise ArkAPIError("must set ak and sk before get endpoint token.")
+            self._sts_token_manager = StsTokenManager(self.ak, self.sk)
+        return self._sts_token_manager.get(endpoint_id)
 
 
 class AsyncArk(AsyncAPIClient):
@@ -56,31 +89,133 @@ class AsyncArk(AsyncAPIClient):
     def __init__(
         self,
         *,
+        ak: str | None = None,
+        sk: str | None = None,
         base_url: str | URL = BASE_URL,
-        api_key: str | None = None,
-        timeout: float | Timeout | None = None,
+        timeout: float | Timeout | None = DEFAULT_TIMEOUT,
         max_retries: int = DEFAULT_MAX_RETRIES,
-        default_query: Dict[str, object] | None = None,
-        http_client: httpx.Client | None = None,
+        http_client: AsyncClient | None = None,
     ) -> None:
-        if api_key is None:
-            api_key = os.environ.get("ARK_API_KEY")
-        self.api_key = api_key
+        """init async ark client, this client is thread unsafe
+
+            Args:
+                ak: access key id
+                sk: secret access key
+                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:
+            ak = os.environ.get("VOLC_ACCESSKEY")
+        if sk is None:
+            sk = os.environ.get("VOLC_SECRETKEY")
+        self.ak = ak
+        self.sk = sk
 
         super().__init__(
             base_url=base_url,
             max_retries=max_retries,
             timeout=timeout,
             http_client=http_client,
-            custom_query=default_query,
+            custom_query=None,
         )
 
         self._default_stream_cls = Stream
+        self._sts_token_manager: StsTokenManager | None = None
 
         self.chat = resources.AsyncChat(self)
 
+    def _get_endpoint_sts_token(self, endpoint_id: str):
+        if self._sts_token_manager is None:
+            if self.ak is None or self.sk is None:
+                raise ArkAPIError("must set ak and sk before get endpoint token.")
+            self._sts_token_manager = StsTokenManager(self.ak, self.sk)
+        return self._sts_token_manager.get(endpoint_id)
+
+
+class StsTokenManager(object):
+
+    # The time at which we'll attempt to refresh, but not
+    # block if someone else is refreshing.
+    _advisory_refresh_timeout: int = _DEFAULT_ADVISORY_REFRESH_TIMEOUT
+    # The time at which all threads will block waiting for
+    # refreshed credentials.
+    _mandatory_refresh_timeout: int = _DEFAULT_MANDATORY_REFRESH_TIMEOUT
+
+    def __init__(self, ak: str, sk: str):
+        self._endpoint_sts_tokens: Dict[str, Tuple[str, int]] = defaultdict(lambda: ("", 0))
+        self._refresh_lock = threading.Lock()
+
+        import volcenginesdkcore
+
+        configuration = volcenginesdkcore.Configuration()
+        configuration.ak = ak
+        configuration.sk = sk
+        configuration.region = "cn-beijing"
+
+        volcenginesdkcore.Configuration.set_default(configuration)
+        self.api_instance = volcenginesdkark.ARKApi()
+
+    def _need_refresh(self, ep: str, refresh_in: int | None = None) -> bool:
+        if refresh_in is None:
+            refresh_in = self._advisory_refresh_timeout
+
+        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):
+        if ttl < _DEFAULT_ADVISORY_REFRESH_TIMEOUT * 2:
+            raise ArkAPIError("ttl should not be under {} seconds.".format(_DEFAULT_ADVISORY_REFRESH_TIMEOUT * 2))
+
+        try:
+            api_key, expired_time = self._load_api_key(
+                ep, ttl
+            )
+            self._endpoint_sts_tokens[ep] = (api_key, expired_time)
+        except ApiException as e:
+            if is_mandatory:
+                raise ArkAPIError("load api key cause error: e={}".format(e))
+            else:
+                logging.error("load api key cause error: e={}".format(e))
+
+    def _refresh(self, ep: str):
+        if not self._need_refresh(ep, self._advisory_refresh_timeout):
+            return
+
+        if self._refresh_lock.acquire(False):
+            if not self._need_refresh(ep, self._advisory_refresh_timeout):
+                return
+
+            try:
+                is_mandatory_refresh = self._need_refresh(
+                    ep, self._mandatory_refresh_timeout
+                )
+
+                self._protected_refresh(ep, is_mandatory=is_mandatory_refresh)
+                return
+            finally:
+                self._refresh_lock.release()
+        elif self._need_refresh(ep, self._mandatory_refresh_timeout):
+            with self._refresh_lock:
+                if not self._need_refresh(ep, self._mandatory_refresh_timeout):
+                    return
+
+                self._protected_refresh(ep, is_mandatory=True)
+
+    def get(self, ep: str) -> str:
+        self._refresh(ep)
+        return self._endpoint_sts_tokens[ep][0]
+
+    def _load_api_key(self, ep: str, duration_seconds: int) -> Tuple[str, int]:
+        get_api_key_request = volcenginesdkark.GetApiKeyRequest(
+            duration_seconds=duration_seconds,
+            resource_type="endpoint",
+            resource_ids=[ep],
+        )
+        resp: volcenginesdkark.GetApiKeyResponse = self.api_instance.get_api_key(
+            get_api_key_request
+        )
 
-    @property
-    def auth_headers(self) -> dict[str, str]:
-        api_key = self.api_key
-        return {"Authorization": f"Bearer {api_key}"}
+        return resp.api_key, resp.expired_time

volcenginesdkarkruntime/_client.py

@@ -18,3 +18,7 @@
 
 INITIAL_RETRY_DELAY = 0.5
 MAX_RETRY_DELAY = 8.0
+
+_DEFAULT_MANDATORY_REFRESH_TIMEOUT = 10 * 60  # 10 min
+_DEFAULT_ADVISORY_REFRESH_TIMEOUT = 30 * 60  # 30 min
+_DEFAULT_STS_TIMEOUT = 7 * 24 * 60 * 60  # 7 days

volcenginesdkarkruntime/_constants.py

@@ -25,9 +25,9 @@ class ArkError(Exception):
 
 class ArkAPIError(ArkError):
     message: str
-    request: httpx.Request
-    body: object | None
-    request_id: str
+    request: Optional[httpx.Request] = None
+    body: Optional[object] = None
+    request_id: Optional[str] = None
     """The API response body.
 
     If the API responded with a valid JSON structure then this property will be the
@@ -45,10 +45,10 @@ class ArkAPIError(ArkError):
     def __init__(
         self,
         message: str,
-        request: httpx.Request,
+        request: Optional[httpx.Request] = None,
         *,
-        body: object | None,
-        request_id: str,
+        body: Optional[object] = None,
+        request_id: Optional[str] = None,
     ) -> None:
         super().__init__(message)
         self.request = request

volcenginesdkarkruntime/_exceptions.py

@@ -8,7 +8,6 @@
 import pydantic
 import pydantic.generics
 from pydantic.fields import FieldInfo
-from pydantic.v1.typing import get_origin
 from typing_extensions import (
     Literal,
     ClassVar,
@@ -28,6 +27,7 @@
     is_literal_type,
     GenericModel as BaseGenericModel,
     parse_obj,
+    get_origin,
     ConfigDict,
 )
 from ._types import ModelT
@@ -37,7 +37,7 @@
 from ._utils._utils import is_mapping, is_list, coerce_boolean, lru_cache
 
 if TYPE_CHECKING:
-    from pydantic_core.core_schema import ModelField
+    from pydantic_core.core_schema import ModelField, ModelFieldsSchema
 
 __all__ = ["BaseModel", "GenericModel"]
 

volcenginesdkarkruntime/_models.py

@@ -19,23 +19,21 @@
 )
 
 import httpx
-import pydantic
 from typing_extensions import ParamSpec, override, get_origin
 
-from ._constants import CLIENT_REQUEST_HEADER, RAW_RESPONSE_HEADER
-from ._exceptions import ArkError, ArkAPIResponseValidationError
+from ._constants import CLIENT_REQUEST_HEADER, RAW_RESPONSE_HEADER  # type: ignore
+from ._exceptions import ArkError
 from ._streaming import Stream, AsyncStream
-from ._types import ResponseT
-from ._utils import extract_type_arg, is_annotated_type
+from ._utils import extract_type_arg, is_annotated_type  # type: ignore
 
 if TYPE_CHECKING:
     from ._base_client import BaseClient
 
 P = ParamSpec("P")
 R = TypeVar("R")
 _T = TypeVar("_T")
-_APIResponseT = TypeVar("_APIResponseT", bound="APIResponse[Any]")
-_AsyncAPIResponseT = TypeVar("_AsyncAPIResponseT", bound="AsyncAPIResponse[Any]")
+_APIResponseT = TypeVar("_APIResponseT", bound="ArkAPIResponse[Any]")
+_AsyncAPIResponseT = TypeVar("_AsyncAPIResponseT", bound="ArkAsyncAPIResponse[Any]")
 _DefaultStreamT = TypeVar("_DefaultStreamT", bound=Union[Stream[Any], AsyncStream[Any]])
 
 log: logging.Logger = logging.getLogger(__name__)
@@ -118,7 +116,7 @@ def _parse(self) -> R:
                 self._stream_cls(
                     cast_to=extract_type_arg(self._stream_cls),
                     response=self.http_response,
-                    client=self._client,
+                    client=self._client,  # type: ignore
                 ),
             )
 
@@ -166,7 +164,7 @@ def _parse(self) -> R:
 
 class ArkAPIResponse(BaseAPIResponse[R]):
     @property
-    def request_id(self) -> str | None:
+    def request_id(self) -> str:
         return self.http_response.headers.get(CLIENT_REQUEST_HEADER, "")
 
     def parse(self) -> R:
@@ -206,7 +204,7 @@ def iter_lines(self) -> Iterator[str]:
 
 class ArkAsyncAPIResponse(BaseAPIResponse[R]):
     @property
-    def request_id(self) -> str | None:
+    def request_id(self) -> str:
         return self.http_response.headers.get(CLIENT_REQUEST_HEADER, "")
 
     async def parse(self) -> R: