Merge branch 'volc_sdk_20251204' into 'integration_2025-12-04_1090671646466'

feat: [development task] ark runtime (1900541)

See merge request iaasng/volcengine-python-sdk!956

Author: BitsAdmin

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

@@ -31,9 +31,8 @@
 from datetime import date, datetime
 from typing_extensions import TypeGuard
 
-import sniffio
-
 from .._types import NotGiven, FileTypes, NotGivenOr
+from .._constants import ARK_APIKEY_PROJECT_NAME
 
 _T = TypeVar("_T")
 _TupleT = TypeVar("_TupleT", bound=Tuple[object, ...])
@@ -382,6 +381,7 @@ def file_from_path(path: str) -> FileTypes:
 
 def get_async_library() -> str:
     try:
+        import sniffio
         return sniffio.current_async_library()
     except Exception:
         return "false"
@@ -445,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 {}
@@ -459,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}
 
 

volcenginesdkarkruntime/resources/content_generation/tasks.py

@@ -46,6 +46,8 @@ def create(
         content: Iterable[CreateTaskContentParam],
         callback_url: Optional[str] = None,
         return_last_frame: Optional[bool] = None,
+        service_tier: Optional[str] = None,
+        execution_expires_after: Optional[int] = None,
         extra_headers: Headers | None = None,
         extra_query: Query | None = None,
         extra_body: Body | None = None,
@@ -58,6 +60,8 @@ def create(
                 "content": content,
                 "callback_url": callback_url,
                 "return_last_frame": return_last_frame,
+                "service_tier": service_tier,
+                "execution_expires_after": execution_expires_after,
             },
             options=make_request_options(
                 extra_headers=extra_headers,
@@ -99,6 +103,7 @@ def list(
         status: str | None = None,
         task_ids: Union[List[str], str] | None = None,
         model: str | None = None,
+        service_tier: str | None = None,
         extra_headers: Headers | None = None,
         extra_body: Body | None = None,
         extra_query: Query | None = None,
@@ -113,6 +118,8 @@ def list(
             query_params.append(("filter.status", status))
         if model:
             query_params.append(("filter.model", model))
+        if service_tier:
+            query_params.append(("filter.service_tier", service_tier))
         if task_ids:
             if isinstance(task_ids, str):
                 task_ids = [task_ids]
@@ -167,6 +174,8 @@ async def create(
         content: Iterable[CreateTaskContentParam],
         callback_url: Optional[str] = None,
         return_last_frame: Optional[bool] = None,
+        service_tier: Optional[str] = None,
+        execution_expires_after: Optional[int] = None,
         extra_headers: Headers | None = None,
         extra_query: Query | None = None,
         extra_body: Body | None = None,
@@ -179,6 +188,8 @@ async def create(
                 "content": content,
                 "callback_url": callback_url,
                 "return_last_frame": return_last_frame,
+                "service_tier": service_tier,
+                "execution_expires_after": execution_expires_after,
             },
             options=make_request_options(
                 extra_headers=extra_headers,
@@ -220,6 +231,7 @@ async def list(
         status: str | None = None,
         task_ids: Union[List[str], str] | None = None,
         model: str | None = None,
+        service_tier: str | None = None,
         extra_headers: Headers | None = None,
         extra_body: Body | None = None,
         extra_query: Query | None = None,
@@ -234,6 +246,8 @@ async def list(
             query_params.append(("filter.status", status))
         if model:
             query_params.append(("filter.model", model))
+        if service_tier:
+            query_params.append(("filter.service_tier", service_tier))
         if task_ids:
             if isinstance(task_ids, str):
                 task_ids = [task_ids]

volcenginesdkarkruntime/types/content_generation/content_generation_task.py

@@ -80,3 +80,9 @@ class ContentGenerationTask(BaseModel):
 
     revised_prompt: str
     """The revised prompt the model uses to generate content"""
+
+    service_tier: str
+    """The service tier used to run the task (optional)."""
+
+    execution_expires_after: int
+    """The expiration time in seconds after which execution should end (optional)."""

volcenginesdkarkruntime/types/responses/easy_input_message.py

@@ -33,3 +33,6 @@ class EasyInputMessage(BaseModel):
 
     type: Optional[Literal["message"]] = None
     """The type of the message input. Always `message`."""
+
+    partial: Optional[bool] = None
+    """Whether or not to include partial assistant responses."""

volcenginesdkarkruntime/types/responses/easy_input_message_param.py

@@ -11,7 +11,7 @@
 
 from __future__ import annotations
 
-from typing import Union
+from typing import Union, Optional
 from typing_extensions import Literal, Required, TypedDict
 
 from .response_input_message_content_list_param import (
@@ -36,3 +36,6 @@ class EasyInputMessageParam(TypedDict, total=False):
 
     type: Literal["message"]
     """The type of the message input. Always `message`."""
+
+    partial: Optional[bool] = None
+    """Whether or not to include partial assistant responses."""

volcenginesdkarkruntime/types/responses/response_output_message.py

@@ -9,7 +9,7 @@
 #
 # This modified file is released under the same license.
 
-from typing import List, Union
+from typing import List, Union, Optional
 from typing_extensions import Literal, Annotated, TypeAlias
 
 from ..._utils import PropertyInfo
@@ -42,3 +42,6 @@ class ResponseOutputMessage(BaseModel):
 
     type: Literal["message"]
     """The type of the output message. Always `message`."""
+
+    partial: Optional[bool]
+    """Whether or not to include partial assistant responses."""

volcenginesdkexamples/volcenginesdkarkruntime/content_generation_tasks.py

@@ -24,13 +24,17 @@
                 # "role": "first_frame"
             }
         ],
+        service_tier="default",
+        execution_expires_after=3600,
         # callback_url="${YOUR_CALLBACK_URL}"
     )
     print(create_result)
 
     print("----- get request -----")
     get_result = client.content_generation.tasks.get(task_id=create_result.id)
     print(get_result)
+    print("ServiceTier:", getattr(get_result, "service_tier", None))
+    print("ExecutionExpiresAfter:", getattr(get_result, "execution_expires_after", None))
 
     print("----- list request -----")
     list_result = client.content_generation.tasks.list(
@@ -41,10 +45,52 @@
         # task_ids=["test-id-1", "test-id-2"] # Filter by task_ids
     )
     print(list_result)
+    if list_result.items:
+        print("List Item ServiceTier:", getattr(list_result.items[0], "service_tier", None))
+        print("List Item ExecutionExpiresAfter:", getattr(list_result.items[0], "execution_expires_after", None))
 
     print("----- delete request -----")
     try:
         client.content_generation.tasks.delete(task_id=create_result.id)
         print(create_result.id)
     except Exception as e:
         print(f"failed to delete task: {e}")
+
+    # ---- flex tier flow: create + GET + LIST + DELETE ----
+    print("----- create request (flex) -----")
+    create_result_flex = client.content_generation.tasks.create(
+        model="${YOUR_MODEL_EP}",
+        content=[
+            {
+                "type": "text",
+                "text": "使用 flex 级别进行内容生成测试，验证 service_tier 与 expire 字段"
+            }
+        ],
+        service_tier="flex",
+        execution_expires_after=3600,
+    )
+    print(create_result_flex)
+
+    print("----- get request (flex) -----")
+    get_result_flex = client.content_generation.tasks.get(task_id=create_result_flex.id)
+    print(get_result_flex)
+    print("Flex ServiceTier:", getattr(get_result_flex, "service_tier", None))
+    print("Flex ExecutionExpiresAfter:", getattr(get_result_flex, "execution_expires_after", None))
+
+    print("----- list request (flex) -----")
+    list_result_flex = client.content_generation.tasks.list(
+        page_num=1,
+        page_size=10,
+        service_tier="flex",
+    )
+    print(list_result_flex)
+    if list_result_flex.items:
+        print("Flex List Item ServiceTier:", getattr(list_result_flex.items[0], "service_tier", None))
+        print("Flex List Item ExecutionExpiresAfter:", getattr(list_result_flex.items[0], "execution_expires_after", None))
+
+    print("----- delete request (flex) -----")
+    try:
+        client.content_generation.tasks.delete(task_id=create_result_flex.id)
+        print(create_result_flex.id)
+    except Exception as e:
+        print(f"failed to delete flex task: {e}")