Merge 'feat/support_presetep' into 'volc_sdk_20251204'

feat: 支持使用短效apikey访问预置接入点

See merge request: !942

Author: volchyt2024

volcenginesdkarkruntime/_client.py

@@ -37,7 +37,9 @@
     _DEFAULT_MANDATORY_REFRESH_TIMEOUT,
     _DEFAULT_STS_TIMEOUT,
     _DEFAULT_RESOURCE_TYPE,
+    _PRESETENDPOINT_RESOURCE_TYPE,
     DEFAULT_TIMEOUT,
+    _BOT_RESOURCE_TYPE,
 )
 from ._streaming import Stream
 
@@ -137,12 +139,15 @@ def __init__(
         self.files = resources.Files(self)
         # self.classification = resources.Classification(self)
 
-    def _get_endpoint_sts_token(self, endpoint_id: str):
+    def _get_endpoint_sts_token(self, endpoint_id: str, project_name: str = None):
         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, self.region)
-        return self._sts_token_manager.get(endpoint_id)
+        resource_type: str = self.get_resource_type_by_endpoint_id(endpoint_id)
+        if resource_type == _PRESETENDPOINT_RESOURCE_TYPE and (project_name is None or project_name.strip() == ""):
+            raise ArkAPIError("must set project_name when get preset endpoint token.")
+        return self._sts_token_manager.get(endpoint_id, resource_type=resource_type, project_name=project_name)
 
     def _get_endpoint_certificate(
         self, endpoint_id: str
@@ -179,6 +184,16 @@ def get_model_breaker(self, model_name: str) -> ModelBreaker:
         with self.model_breaker_lock:
             return self.model_breaker_map[model_name]
 
+    def get_resource_type_by_endpoint_id(self, endpoint_id: str) -> str:
+        if endpoint_id.startswith("ep-m-"):
+            return _PRESETENDPOINT_RESOURCE_TYPE
+        if endpoint_id.startswith("ep-"):
+            return _DEFAULT_RESOURCE_TYPE
+        if endpoint_id.startswith("bot-"):
+            return _BOT_RESOURCE_TYPE
+        # for model id, default to preset endpoint
+        return _PRESETENDPOINT_RESOURCE_TYPE
+
 
 class AsyncArk(AsyncAPIClient):
     beta: beta.AsyncBeta
@@ -350,6 +365,7 @@ def _protected_refresh(
         ttl: int = _DEFAULT_STS_TIMEOUT,
         is_mandatory: bool = False,
         resource_type: str = _DEFAULT_RESOURCE_TYPE,
+        project_name: str = None,
     ):
         if ttl < self._advisory_refresh_timeout * 2:
             raise ArkAPIError(
@@ -360,7 +376,7 @@ def _protected_refresh(
 
         try:
             api_key, expired_time = self._load_api_key(
-                ep, ttl, resource_type=resource_type
+                ep, ttl, resource_type=resource_type, project_name=project_name
             )
             self._endpoint_sts_tokens[ep] = (api_key, expired_time)
         except ApiException as e:
@@ -369,7 +385,7 @@ def _protected_refresh(
             else:
                 logging.error("load api key cause error: e={}".format(e))
 
-    def _refresh(self, ep: str, resource_type: str = _DEFAULT_RESOURCE_TYPE):
+    def _refresh(self, ep: str, resource_type: str = _DEFAULT_RESOURCE_TYPE, project_name: str = None):
         if not self._need_refresh(ep, self._advisory_refresh_timeout):
             return
 
@@ -383,7 +399,7 @@ def _refresh(self, ep: str, resource_type: str = _DEFAULT_RESOURCE_TYPE):
                 )
 
                 self._protected_refresh(
-                    ep, is_mandatory=is_mandatory_refresh, resource_type=resource_type
+                    ep, is_mandatory=is_mandatory_refresh, resource_type=resource_type, project_name=project_name
                 )
                 return
             finally:
@@ -394,24 +410,27 @@ def _refresh(self, ep: str, resource_type: str = _DEFAULT_RESOURCE_TYPE):
                     return
 
                 self._protected_refresh(
-                    ep, is_mandatory=True, resource_type=resource_type
+                    ep, is_mandatory=True, resource_type=resource_type, project_name=project_name
                 )
 
-    def get(self, ep: str, resource_type: str = _DEFAULT_RESOURCE_TYPE) -> str:
-        self._refresh(ep, resource_type=resource_type)
+    def get(self, ep: str, resource_type: str = _DEFAULT_RESOURCE_TYPE, project_name: str = None) -> str:
+        self._refresh(ep, resource_type=resource_type, project_name=project_name)
         return self._endpoint_sts_tokens[ep][0]
 
     def _load_api_key(
         self,
         ep: str,
         duration_seconds: int,
         resource_type: str = _DEFAULT_RESOURCE_TYPE,
+        project_name: str = None,
     ) -> Tuple[str, int]:
         get_api_key_request = volcenginesdkark.GetApiKeyRequest(
             duration_seconds=duration_seconds,
             resource_type=resource_type,
             resource_ids=[ep],
         )
+        if project_name is not None and project_name.strip() != "":
+            get_api_key_request.project_name = project_name
         resp: volcenginesdkark.GetApiKeyResponse = self.api_instance.get_api_key(
             get_api_key_request
         )

volcenginesdkarkruntime/_constants.py

@@ -18,6 +18,7 @@
 CLIENT_REQUEST_HEADER = "X-Client-Request-Id"
 SERVER_REQUEST_HEADER = "X-Request-Id"
 ARK_E2E_ENCRYPTION_HEADER = "x-is-encrypted"
+ARK_APIKEY_PROJECT_NAME = "X-Project-Name"
 
 DEFAULT_TIMEOUT_SECONDS = 600.0
 DEFAULT_CONNECT_TIMEOUT_SECONDS = 60.0
@@ -39,3 +40,5 @@
 _DEFAULT_STS_TIMEOUT = 7 * 24 * 60 * 60  # 7 days
 
 _DEFAULT_RESOURCE_TYPE = "endpoint"
+_PRESETENDPOINT_RESOURCE_TYPE = "presetendpoint"
+_BOT_RESOURCE_TYPE = "bot"

volcenginesdkarkruntime/_utils/_utils.py

@@ -32,6 +32,7 @@
 from typing_extensions import TypeGuard
 
 from .._types import NotGiven, FileTypes, NotGivenOr
+from .._constants import ARK_APIKEY_PROJECT_NAME
 
 _T = TypeVar("_T")
 _TupleT = TypeVar("_TupleT", bound=Tuple[object, ...])
@@ -444,12 +445,12 @@ def _insert_sts_token(args, kwargs):
     if (
         ark_client.api_key is None
         and model
-        and model.startswith("ep-")
+        and model.startswith("bot-")
         and ark_client.ak
         and ark_client.sk
     ):
         default_auth_header = {
-            "Authorization": "Bearer " + ark_client._get_endpoint_sts_token(model)
+            "Authorization": "Bearer " + ark_client._get_bot_sts_token(model)
         }
         extra_headers = (
             kwargs.get("extra_headers") if kwargs.get("extra_headers") else {}
@@ -458,16 +459,18 @@ def _insert_sts_token(args, kwargs):
     elif (
         ark_client.api_key is None
         and model
-        and model.startswith("bot-")
         and ark_client.ak
         and ark_client.sk
     ):
-        default_auth_header = {
-            "Authorization": "Bearer " + ark_client._get_bot_sts_token(model)
-        }
         extra_headers = (
             kwargs.get("extra_headers") if kwargs.get("extra_headers") else {}
         )
+        project_name: str = None
+        if extra_headers is not None and extra_headers.get(ARK_APIKEY_PROJECT_NAME, None) is not None:
+            project_name = extra_headers[ARK_APIKEY_PROJECT_NAME]
+        default_auth_header = {
+            "Authorization": "Bearer " + ark_client._get_endpoint_sts_token(model, project_name)
+        }
         kwargs["extra_headers"] = {**default_auth_header, **extra_headers}