arkruntime 0.1.0__py3-none-any.whl

arkruntime/_client.py

@@ -0,0 +1,585 @@
+# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+# License text: https://github.com/volcengine/ark-runtime-python/blob/main/LICENSE
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+import threading
+import time
+from collections import defaultdict
+from typing import Dict, Optional, Tuple
+
+import volcenginesdkark
+from httpx import URL, AsyncClient, Client, Timeout
+from volcenginesdkcore.rest import ApiException
+
+from . import resources
+from ._base_client import AsyncAPIClient, SyncAPIClient
+from ._constants import (
+    _CLOUD_PRESETS,
+    _DEFAULT_ADVISORY_REFRESH_TIMEOUT,
+    _DEFAULT_MANDATORY_REFRESH_TIMEOUT,
+    _DEFAULT_RESOURCE_TYPE,
+    _DEFAULT_STS_TIMEOUT,
+    _PRESETENDPOINT_RESOURCE_TYPE,
+    BASE_URL,
+    DEFAULT_MAX_RETRIES,
+    DEFAULT_TIMEOUT,
+    Cloud,
+)
+from ._exceptions import ArkAPIError
+from ._models import BaseModel
+from ._streaming import Stream
+from ._utils._key_agreement import get_cert_info, key_agreement_client
+from ._utils._model_breaker import ModelBreaker
+from .resources.batch import batch
+from .resources.beta import beta
+
+__all__ = ["Ark", "AsyncArk"]
+
+
+class Ark(SyncAPIClient):
+    beta: beta.Beta
+    chat: resources.Chat
+    embeddings: resources.Embeddings
+    tokenization: resources.Tokenization
+    multimodal_embeddings: resources.MultimodalEmbeddings
+    content_generation: resources.ContentGeneration
+    images: resources.Images
+    responses: resources.Responses
+    input_items: resources.InputItems
+    batch: batch.Batch
+    model_breaker_map: dict[str, ModelBreaker]
+    model_breaker_lock: threading.Lock
+    files: resources.Files
+
+    def __init__(
+        self,
+        *,
+        base_url: str | URL = BASE_URL,
+        ak: str | None = None,
+        sk: str | None = None,
+        api_key: str | None = None,
+        region: str = "cn-beijing",
+        timeout: float | Timeout | None = DEFAULT_TIMEOUT,
+        max_retries: int = DEFAULT_MAX_RETRIES,
+        http_client: Client | None = None,
+    ) -> None:
+        """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
+                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:
+                ark client
+        """
+        if ak is None:
+            ak = os.environ.get("VOLC_ACCESSKEY")
+        if sk is None:
+            sk = os.environ.get("VOLC_SECRETKEY")
+        if api_key is None:
+            api_key = os.environ.get("ARK_API_KEY")
+        self._base_url = base_url
+        self.ak = ak
+        self.sk = sk
+        self.api_key = api_key
+        self.region = region
+
+        # Reject empty api_key eagerly — passing "" would otherwise skip the
+        # `is not None` guard, produce an "Authorization: Bearer " header
+        # with no token, and surface as an opaque ArkAPIConnectionError
+        # ("Connection error") from the HTTP layer instead of a clear
+        # config-time failure.
+        if api_key is not None and not api_key.strip():
+            raise ValueError(
+                "api_key must be a non-empty string; got '' — supply a valid API "
+                "key, set ARK_API_KEY, or pass ak/sk instead"
+            )
+        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,
+            max_retries=max_retries,
+            timeout=timeout,
+            http_client=http_client,
+            custom_query=None,
+        )
+
+        self._default_stream_cls = Stream
+        self._sts_token_manager: StsTokenManager | None = None
+        self._certificate_manager: E2ECertificateManager | None = None
+
+        self.beta = beta.Beta(self)
+        self.chat = resources.Chat(self)
+        self.embeddings = resources.Embeddings(self)
+        self.tokenization = resources.Tokenization(self)
+        self.multimodal_embeddings = resources.MultimodalEmbeddings(self)
+        self.content_generation = resources.ContentGeneration(self)
+        self.images = resources.Images(self)
+        self.responses = resources.Responses(self)
+        self.input_items = resources.InputItems(self)
+        self.batch = batch.Batch(self)
+        self.model_breaker_map = defaultdict(ModelBreaker)
+        self.model_breaker_lock = threading.Lock()
+        self.files = resources.Files(self)
+        # Managed Agents (data plane /api/v3/*)
+        self.agents = resources.Agents(self)
+        self.environments = resources.Environments(self)
+        self.memory_stores = resources.MemoryStores(self)
+        self.sessions = resources.Sessions(self)
+        self.skills = resources.Skills(self)
+        self.vaults = resources.Vaults(self)
+
+    @classmethod
+    def volc(cls, **kwargs) -> "Ark":
+        """Construct an Ark client targeting the Volcengine cloud
+        (``ark.cn-beijing.volces.com``). Reads ``VOLC_ACCESSKEY`` /
+        ``VOLC_SECRETKEY`` for AK/SK and ``ARK_API_KEY`` for the API key
+        when those kwargs are not passed explicitly."""
+        return _from_cloud(cls, Cloud.VOLC, **kwargs)
+
+    @classmethod
+    def byteplus(cls, **kwargs) -> "Ark":
+        """Construct an Ark client targeting the Byteplus cloud
+        (``ark.ap-southeast.bytepluses.com``). Reads
+        ``BYTEPLUS_ACCESSKEY`` / ``BYTEPLUS_SECRETKEY`` and ``ARK_API_KEY``
+        when those kwargs are not passed explicitly."""
+        return _from_cloud(cls, Cloud.BYTEPLUS, **kwargs)
+
+    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)
+        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) -> 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
+            )
+        return self._certificate_manager.get(endpoint_id)
+
+    @property
+    def auth_headers(self) -> dict[str, str]:
+        api_key = self.api_key
+        return {"Authorization": f"Bearer {api_key}"}
+
+    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
+        # for model id, default to preset endpoint
+        return _PRESETENDPOINT_RESOURCE_TYPE
+
+
+class AsyncArk(AsyncAPIClient):
+    beta: beta.AsyncBeta
+    chat: resources.AsyncChat
+    embeddings: resources.AsyncEmbeddings
+    tokenization: resources.AsyncTokenization
+    multimodal_embeddings: resources.AsyncMultimodalEmbeddings
+    content_generation: resources.AsyncContentGeneration
+    images: resources.AsyncImages
+    responses: resources.AsyncResponses
+    input_items: resources.AsyncInputItems
+    batch: batch.AsyncBatch
+    model_breaker_map: dict[str, ModelBreaker]
+    model_breaker_lock: asyncio.Lock
+    files: resources.AsyncFiles
+
+    def __init__(
+        self,
+        *,
+        ak: str | None = None,
+        sk: str | None = None,
+        api_key: str | None = None,
+        base_url: str | URL = BASE_URL,
+        region: str = "cn-beijing",
+        timeout: float | Timeout | None = DEFAULT_TIMEOUT,
+        max_retries: int = DEFAULT_MAX_RETRIES,
+        http_client: AsyncClient | None = None,
+    ) -> 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
+        """
+
+        if ak is None:
+            ak = os.environ.get("VOLC_ACCESSKEY")
+        if sk is None:
+            sk = os.environ.get("VOLC_SECRETKEY")
+        if api_key is None:
+            api_key = os.environ.get("ARK_API_KEY")
+        self._base_url = base_url
+        self.ak = ak
+        self.sk = sk
+        self.api_key = api_key
+        self.region = region
+
+        # Reject empty api_key eagerly — passing "" would otherwise skip the
+        # `is not None` guard, produce an "Authorization: Bearer " header
+        # with no token, and surface as an opaque ArkAPIConnectionError
+        # ("Connection error") from the HTTP layer instead of a clear
+        # config-time failure.
+        if api_key is not None and not api_key.strip():
+            raise ValueError(
+                "api_key must be a non-empty string; got '' — supply a valid API "
+                "key, set ARK_API_KEY, or pass ak/sk instead"
+            )
+        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,
+            max_retries=max_retries,
+            timeout=timeout,
+            http_client=http_client,
+            custom_query=None,
+        )
+
+        self._default_stream_cls = Stream
+        self._sts_token_manager: StsTokenManager | None = None
+        self._certificate_manager: E2ECertificateManager | None = None
+
+        self.beta = beta.AsyncBeta(self)
+        self.chat = resources.AsyncChat(self)
+        self.embeddings = resources.AsyncEmbeddings(self)
+        self.tokenization = resources.AsyncTokenization(self)
+        self.multimodal_embeddings = resources.AsyncMultimodalEmbeddings(self)
+        self.content_generation = resources.AsyncContentGeneration(self)
+        self.images = resources.AsyncImages(self)
+        self.responses = resources.AsyncResponses(self)
+        self.input_items = resources.AsyncInputItems(self)
+        self.batch = batch.AsyncBatch(self)
+        self.model_breaker_map = defaultdict(ModelBreaker)
+        self.model_breaker_lock = asyncio.Lock()
+        self.files = resources.AsyncFiles(self)
+        # Managed Agents (data plane /api/v3/*)
+        self.agents = resources.AsyncAgents(self)
+        self.environments = resources.AsyncEnvironments(self)
+        self.memory_stores = resources.AsyncMemoryStores(self)
+        self.sessions = resources.AsyncSessions(self)
+        self.skills = resources.AsyncSkills(self)
+        self.vaults = resources.AsyncVaults(self)
+
+    @classmethod
+    def volc(cls, **kwargs) -> "AsyncArk":
+        """See :meth:`Ark.volc`."""
+        return _from_cloud(cls, Cloud.VOLC, **kwargs)
+
+    @classmethod
+    def byteplus(cls, **kwargs) -> "AsyncArk":
+        """See :meth:`Ark.byteplus`."""
+        return _from_cloud(cls, Cloud.BYTEPLUS, **kwargs)
+
+    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, self.region)
+        return self._sts_token_manager.get(endpoint_id)
+
+    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
+            )
+        return self._certificate_manager.get(endpoint_id)
+
+    @property
+    def auth_headers(self) -> dict[str, str]:
+        api_key = self.api_key
+        return {"Authorization": f"Bearer {api_key}"}
+
+    async def get_model_breaker(self, model_name: str) -> ModelBreaker:
+        async with self.model_breaker_lock:
+            return self.model_breaker_map[model_name]
+
+
+def _from_cloud(cls, cloud: str, **kwargs):
+    """Resolve a cloud preset, fill in env-var defaults for the cloud's
+    specific AK/SK names plus the shared ``ARK_API_KEY``, then delegate to
+    ``cls.__init__``. Explicit kwargs always win."""
+    preset = _CLOUD_PRESETS[cloud]
+    kwargs.setdefault("base_url", preset["base_url"])
+    kwargs.setdefault("region", preset["region"])
+    if "ak" not in kwargs:
+        kwargs["ak"] = os.environ.get(preset["ak_env"])
+    if "sk" not in kwargs:
+        kwargs["sk"] = os.environ.get(preset["sk_env"])
+    if "api_key" not in kwargs:
+        kwargs["api_key"] = os.environ.get("ARK_API_KEY")
+    return cls(**kwargs)
+
+
+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, region: 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 = region
+        configuration.schema = "https"
+
+        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,
+        resource_type: str = _DEFAULT_RESOURCE_TYPE,
+        project_name: str = None,
+    ):
+        if ttl < 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(ep, ttl, resource_type=resource_type, project_name=project_name)
+            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, resource_type: str = _DEFAULT_RESOURCE_TYPE, project_name: str = None):
+        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, resource_type=resource_type, project_name=project_name
+                )
+                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, resource_type=resource_type, project_name=project_name)
+
+    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)
+
+        return resp.api_key, resp.expired_time
+
+
+class CertificateResponse(BaseModel):
+    error: Optional[Dict[str, str]] = None
+    """The error information."""
+
+    Certificate: str
+    """The certificate content."""
+
+
+class E2ECertificateManager(object):
+    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
+        configuration.region = region
+        configuration.schema = "https"
+        volcenginesdkcore.Configuration.set_default(configuration)
+        self.api_instance = volcenginesdkark.ARKApi()
+
+        # global cert path prepare
+        self.cert_path = os.environ.get("E2E_CERTIFICATE_PATH")
+
+        # ark client prepare
+        self._ark_client_enabled = True
+        if api_key is None:
+            self._ark_client_enabled = False
+        else:
+            self.client = Ark(
+                base_url=base_url,
+                api_key=api_key,
+            )
+        self._e2e_uri = "/e2e/get/certificate"
+        self._x_session_token = {"X-Session-Token": self._e2e_uri}
+
+        self._aicc_enabled = False
+        if os.environ.get("VOLC_ARK_ENCRYPTION") == "AICC":
+            self._aicc_enabled = True
+
+    def _load_cert_by_cert_path(self) -> str:
+        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)
+        if self._aicc_enabled:
+            get_endpoint_certificate_request = volcenginesdkark.GetEndpointCertificateRequest(id=ep, type="AICCv0.1")
+        try:
+            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)
+
+        return resp.pca_instance_certificate
+
+    def _sync_load_cert_by_auth(self, ep: str) -> str:
+        # try to make request with session header (used for header statistic)
+        try:
+            req_body = {"model": ep}
+            if self._aicc_enabled:
+                req_body["type"] = "AICCv0.1"
+            resp = self.client.post(
+                self._e2e_uri,
+                options={"headers": self._x_session_token},
+                body=req_body,
+                cast_to=CertificateResponse,
+            )
+        except Exception as e:
+            raise ArkAPIError("Getting Certificate failed: %s\n" % e)
+        if resp.error is not None:
+            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:
+            f.write(cert_pem)
+
+    def _load_cert_locally(self, ep: str) -> str | None:
+        cert_file_path = os.path.join(self._cert_storage_path, f"{ep}.pem")
+        if os.path.exists(cert_file_path):
+            last_modified_time = os.path.getmtime(cert_file_path)
+            current_time = time.time()
+            time_difference = current_time - last_modified_time
+            if time_difference <= self._cert_expiration_seconds:
+                cert_pem = None
+                with open(cert_file_path, "r") as f:
+                    cert_pem = f.read()
+                ring, key, _ = get_cert_info(cert_pem)
+                # check cert is complement with AICC/PCA
+                if (ring == "" or key == "") and not self._aicc_enabled:
+                    return cert_pem
+                if ring != "" and key != "" and self._aicc_enabled:
+                    return cert_pem
+            os.remove(cert_file_path)
+        return None
+
+    def _init_local_cert_cache(self):
+        self._cert_storage_path = os.path.expanduser("~/.ark/certificates")
+        self._cert_expiration_seconds = 14 * 24 * 60 * 60  # 14 days
+
+        if not os.path.exists(self._cert_storage_path):
+            try:
+                os.makedirs(self._cert_storage_path)
+            except FileExistsError:
+                pass
+            except Exception as e:
+                raise ArkAPIError("failed to create certificate directory %s: %s\n" % (self._cert_storage_path, e))
+
+    def get(self, ep: str) -> key_agreement_client:
+        if ep in self._certificate_manager and not self._certificate_manager[ep].need_reload():
+            return self._certificate_manager[ep]
+        cert_pem = self._load_cert_locally(ep)
+        if cert_pem is None:
+            if self.cert_path is not None:
+                cert_pem = self._load_cert_by_cert_path()
+            elif self._ark_client_enabled:
+                cert_pem = self._sync_load_cert_by_auth(ep)
+            else:
+                cert_pem = self._load_cert_by_ak_sk(ep)
+            self._save_cert_to_file(ep, cert_pem)
+        client = key_agreement_client(certificate_pem_string=cert_pem)
+        self._certificate_manager[ep] = client
+        return client