arkruntime 0.1.0__py3-none-any.whl

arkruntime/resources/embeddings/embeddings.py

@@ -0,0 +1,140 @@
+# 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
+
+from typing import Optional
+
+import httpx
+
+from ..._base_client import make_request_options
+from ..._compat import cached_property
+from ..._resource import AsyncAPIResource, SyncAPIResource
+from ..._response import (
+    async_to_raw_response_wrapper,
+    to_raw_response_wrapper,
+)
+from ..._types import Body, Headers, Query
+from ..._utils import async_maybe_transform, maybe_transform
+
+# AUTOGEN-START create-imports
+from ...types.embedding.embedding_input_param import EmbeddingInputParam
+from ...types.embedding.embedding_request_param import EmbeddingRequestParam
+from ...types.embedding.embedding_response import EmbeddingResponse
+from ...types.embedding.encoding_format import EncodingFormat
+
+# AUTOGEN-END create-imports
+
+__all__ = [
+    "Embeddings",
+    "AsyncEmbeddings",
+    "EmbeddingsWithRawResponse",
+    "AsyncEmbeddingsWithRawResponse",
+]
+
+
+class Embeddings(SyncAPIResource):
+    @cached_property
+    def with_raw_response(self) -> EmbeddingsWithRawResponse:
+        return EmbeddingsWithRawResponse(self)
+
+    def create(
+        self,
+        *,
+        # AUTOGEN-START create-kwargs
+        model: str,
+        input: EmbeddingInputParam,
+        encoding_format: Optional[EncodingFormat] = None,
+        dimensions: Optional[int] = None,
+        user: Optional[str] = None,
+        # AUTOGEN-END create-kwargs
+        extra_headers: Headers | None = None,
+        extra_query: Query | None = None,
+        extra_body: Body | None = None,
+        timeout: float | httpx.Timeout | None = None,
+    ) -> EmbeddingResponse:
+        return self._post(
+            "/embeddings",
+            body=maybe_transform(
+                {
+                    # AUTOGEN-START create-body
+                    "model": model,
+                    "input": input,
+                    "encoding_format": encoding_format,
+                    "dimensions": dimensions,
+                    "user": user,
+                    # AUTOGEN-END create-body
+                },
+                EmbeddingRequestParam,
+            ),
+            options=make_request_options(
+                extra_headers=extra_headers,
+                extra_query=extra_query,
+                extra_body=extra_body,
+                timeout=timeout,
+            ),
+            cast_to=EmbeddingResponse,
+        )
+
+
+class AsyncEmbeddings(AsyncAPIResource):
+    @cached_property
+    def with_raw_response(self) -> AsyncEmbeddingsWithRawResponse:
+        return AsyncEmbeddingsWithRawResponse(self)
+
+    async def create(
+        self,
+        *,
+        # AUTOGEN-START create-kwargs
+        model: str,
+        input: EmbeddingInputParam,
+        encoding_format: Optional[EncodingFormat] = None,
+        dimensions: Optional[int] = None,
+        user: Optional[str] = None,
+        # AUTOGEN-END create-kwargs
+        extra_headers: Headers | None = None,
+        extra_query: Query | None = None,
+        extra_body: Body | None = None,
+        timeout: float | httpx.Timeout | None = None,
+    ) -> EmbeddingResponse:
+        return await self._post(
+            "/embeddings",
+            body=await async_maybe_transform(
+                {
+                    # AUTOGEN-START create-body
+                    "model": model,
+                    "input": input,
+                    "encoding_format": encoding_format,
+                    "dimensions": dimensions,
+                    "user": user,
+                    # AUTOGEN-END create-body
+                },
+                EmbeddingRequestParam,
+            ),
+            options=make_request_options(
+                extra_headers=extra_headers,
+                extra_query=extra_query,
+                extra_body=extra_body,
+                timeout=timeout,
+            ),
+            cast_to=EmbeddingResponse,
+        )
+
+
+class EmbeddingsWithRawResponse:
+    def __init__(self, embeddings: Embeddings) -> None:
+        self._embeddings = embeddings
+
+        self.create = to_raw_response_wrapper(
+            embeddings.create,
+        )
+
+
+class AsyncEmbeddingsWithRawResponse:
+    def __init__(self, embeddings: AsyncEmbeddings) -> None:
+        self._embeddings = embeddings
+
+        self.create = async_to_raw_response_wrapper(
+            embeddings.create,
+        )

arkruntime/resources/encryption.py

@@ -0,0 +1,187 @@
+# Copyright (c) 2026 ByteDance Ltd. and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import json
+import os
+import warnings
+from copy import deepcopy
+from typing import (
+    AsyncIterator,
+    Callable,
+    Iterable,
+    Iterator,
+    Union,
+)
+from urllib.parse import urlparse
+
+from .._constants import ARK_E2E_ENCRYPTION_HEADER
+from .._streaming import AsyncStream, Stream
+from .._utils._key_agreement import (
+    aes_gcm_decrypt_base64_list,
+    aes_gcm_decrypt_base64_string,
+    decrypt_validate,
+)
+from ..types.chat import (
+    ChatCompletion,
+    ChatCompletionChunk,
+)
+
+
+def _decrypt_chunk(key: bytes, nonce: bytes, resp: Stream[ChatCompletionChunk]) -> Iterator[ChatCompletionChunk]:
+    for chunk in resp:
+        if chunk.choices is not None:
+            for index, choice in enumerate(chunk.choices):
+                if (
+                    choice.delta is not None
+                    and choice.delta.content is not None
+                    and choice.finish_reason != "content_filter"
+                ):
+                    choice.delta.content = aes_gcm_decrypt_base64_string(key, nonce, choice.delta.content)
+                chunk.choices[index] = choice
+        yield chunk
+
+
+async def _async_decrypt_chunk(
+    key: bytes, nonce: bytes, resp: AsyncStream[ChatCompletionChunk]
+) -> AsyncIterator[ChatCompletionChunk]:
+    async for chunk in resp:
+        if chunk.choices is not None:
+            for index, choice in enumerate(chunk.choices):
+                if (
+                    choice.delta is not None
+                    and choice.delta.content is not None
+                    and choice.finish_reason != "content_filter"
+                ):
+                    choice.delta.content = aes_gcm_decrypt_base64_string(key, nonce, choice.delta.content)
+                chunk.choices[index] = choice
+        yield chunk
+
+
+def _decrypt(
+    key: bytes, nonce: bytes, resp: Union[ChatCompletion, Stream[ChatCompletionChunk]]
+) -> ChatCompletion | Stream[ChatCompletionChunk]:
+    if isinstance(resp, ChatCompletion):
+        if resp.choices is not None:
+            for index, choice in enumerate(resp.choices):
+                if (
+                    choice.message is not None
+                    and choice.finish_reason != "content_filter"
+                    and choice.message.content is not None
+                ):
+                    try:
+                        content = aes_gcm_decrypt_base64_string(key, nonce, choice.message.content)
+                    except Exception:
+                        content = ""
+                    if content == "" or not decrypt_validate(choice.message.content):
+                        content = aes_gcm_decrypt_base64_list(key, nonce, choice.message.content)
+                    choice.message.content = content
+                resp.choices[index] = choice
+        return resp
+    else:
+        return Stream._make_stream_from_iterator(_decrypt_chunk(key, nonce, resp))
+
+
+async def _async_decrypt(
+    key: bytes,
+    nonce: bytes,
+    resp: ChatCompletion | AsyncStream[ChatCompletionChunk],
+) -> ChatCompletion | AsyncStream[ChatCompletionChunk]:
+    if isinstance(resp, ChatCompletion):
+        if resp.choices is not None:
+            for index, choice in enumerate(resp.choices):
+                if (
+                    choice.message is not None
+                    and choice.finish_reason != "content_filter"
+                    and choice.message.content is not None
+                ):
+                    try:
+                        content = aes_gcm_decrypt_base64_string(key, nonce, choice.message.content)
+                    except Exception:
+                        content = ""
+                    if content == "" or not decrypt_validate(choice.message.content):
+                        content = aes_gcm_decrypt_base64_list(key, nonce, choice.message.content)
+                    choice.message.content = content
+                resp.choices[index] = choice
+        return resp
+    else:
+        return AsyncStream._make_stream_from_iterator(_async_decrypt_chunk(key, nonce, resp))
+
+
+def with_e2e_encryption(func):
+    def wrapper(*args, **kwargs):
+        is_encrypt, _crypto_key, _crypto_nonce = _content_encryption(args, kwargs)
+        resp = func(*args, **kwargs)
+        if is_encrypt:
+            resp = _decrypt(_crypto_key, _crypto_nonce, resp)
+        return resp
+
+    return wrapper
+
+
+def async_with_e2e_encryption(func):
+    async def wrapper(*args, **kwargs):
+        is_encrypt, _crypto_key, _crypto_nonce = _content_encryption(args, kwargs)
+        resp = await func(*args, **kwargs)
+        if is_encrypt:
+            resp = await _async_decrypt(_crypto_key, _crypto_nonce, resp)
+        return resp
+
+    return wrapper
+
+
+def _content_encryption(args, kwargs):
+    assert len(args) > 0
+    extra_headers = kwargs.get("extra_headers") if kwargs.get("extra_headers") else {}
+    if extra_headers is not None and extra_headers.get(ARK_E2E_ENCRYPTION_HEADER, None) == "true":
+        model: str = kwargs.get("model", "")
+        messages = deepcopy(kwargs["messages"])
+        ark_client = args[0]._client
+        client = ark_client._get_endpoint_certificate(model)
+        _crypto_key, _crypto_nonce, session_token = client.generate_ecies_key_pair()
+        extra_headers["X-Session-Token"] = session_token
+        _process_messages(
+            messages,
+            lambda x: client.encrypt_string_with_key(_crypto_key, _crypto_nonce, x),
+        )
+        info = {"ExpireTime": client.get_cert_expiration_time()}
+        if os.environ.get("VOLC_ARK_ENCRYPTION") == "AICC":
+            ring_id, key_id = client.get_cert_ring_key_id()
+            info["Version"] = "AICCv0.1"
+            info["KeyID"] = key_id
+            info["RingID"] = ring_id
+        extra_headers["X-Encrypt-Info"] = json.dumps(info)
+        kwargs["extra_headers"] = {**extra_headers}
+        kwargs["messages"] = messages
+        return True, _crypto_key, _crypto_nonce
+    return False, None, None
+
+
+def _process_messages(messages, f: Callable[[str], str]):
+    for message in messages:
+        if message.get("content", None) is not None:
+            current_content = message.get("content")
+            if isinstance(current_content, str):
+                message["content"] = f(current_content)
+            elif isinstance(current_content, Iterable):
+                for part in current_content:
+                    if part.get("type", None) == "text":
+                        part["text"] = f(part["text"])
+                    elif part.get("type", None) == "image_url":
+                        parse_result = urlparse(part["image_url"]["url"])
+                        if parse_result.scheme == "data":
+                            part["image_url"]["url"] = f(part["image_url"]["url"])
+                        elif parse_result.scheme == "http" or parse_result.scheme == "https":
+                            warnings.warn(
+                                "encryption is not supported for image url, "
+                                "please use base64 image if you want encryption"
+                            )
+                        else:
+                            raise TypeError("encryption is not supported for image url scheme {}".format(parse_result))
+                    elif part.get("type", None) == "video_url":
+                        pass
+                    else:
+                        raise TypeError("encryption is not supported for content type {}".format(type(part)))
+            else:
+                raise TypeError("encryption is not supported for content type {}".format(type(message.get("content"))))

arkruntime/resources/environments/__init__.py

@@ -0,0 +1,3 @@
+from arkruntime.resources.environments.environments import AsyncEnvironments, Environments
+
+__all__ = ["Environments", "AsyncEnvironments"]

arkruntime/resources/environments/environments.py

@@ -0,0 +1,216 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import httpx
+
+from ..._base_client import make_request_options
+from ..._managed_agents_serialize import dump_body
+from ..._resource import AsyncAPIResource, SyncAPIResource
+from ..._types import NOT_GIVEN, Body, Headers, NotGiven, Query
+from ...types.environment.delete_environment_response import DeleteEnvironmentResponse
+from ...types.environment.env_config import EnvConfig
+from ...types.environment.environment import Environment
+from ...types.environment.environment_scope import EnvironmentScope
+from ...types.environment.list_environments_response import ListEnvironmentsResponse
+
+__all__ = ["Environments", "AsyncEnvironments"]
+
+_ENV_PREFIX = "/environments"
+
+
+def _list_query(*, limit=NOT_GIVEN, page=NOT_GIVEN) -> dict:
+    q: dict = {}
+    if limit is not NOT_GIVEN and limit is not None:
+        q["limit"] = limit
+    if page is not NOT_GIVEN and page is not None:
+        q["page"] = page
+    return q
+
+
+class Environments(SyncAPIResource):
+    def create(
+        self,
+        *,
+        name: str,
+        config: Optional[EnvConfig] | NotGiven = NOT_GIVEN,
+        description: Optional[str] | NotGiven = NOT_GIVEN,
+        metadata: Optional[dict] | NotGiven = NOT_GIVEN,
+        scope: Optional[EnvironmentScope] | NotGiven = NOT_GIVEN,
+        extra_headers: Headers | None = None,
+        extra_query: Query | None = None,
+        extra_body: Body | None = None,
+        timeout: float | httpx.Timeout | None = None,
+    ) -> Environment:
+        return self._post(
+            _ENV_PREFIX,
+            body=dump_body(
+                {"name": name, "config": config, "description": description, "metadata": metadata, "scope": scope}
+            ),
+            options=make_request_options(
+                extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
+            ),
+            cast_to=Environment,
+        )
+
+    def retrieve(
+        self, environment_id: str, *, extra_headers=None, extra_query=None, extra_body=None, timeout=None
+    ) -> Environment:
+        if not environment_id:
+            raise ValueError("environment_id is required")
+        return self._get(
+            f"{_ENV_PREFIX}/{environment_id}",
+            options=make_request_options(
+                extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
+            ),
+            cast_to=Environment,
+        )
+
+    def list(
+        self, *, limit=NOT_GIVEN, page=NOT_GIVEN, extra_headers=None, extra_query=None, extra_body=None, timeout=None
+    ) -> ListEnvironmentsResponse:
+        return self._get(
+            _ENV_PREFIX,
+            options=make_request_options(
+                query=_list_query(limit=limit, page=page),
+                extra_headers=extra_headers,
+                extra_query=extra_query,
+                extra_body=extra_body,
+                timeout=timeout,
+            ),
+            cast_to=ListEnvironmentsResponse,
+        )
+
+    def update(
+        self,
+        environment_id: str,
+        *,
+        name=NOT_GIVEN,
+        config=NOT_GIVEN,
+        description=NOT_GIVEN,
+        metadata=NOT_GIVEN,
+        scope=NOT_GIVEN,
+        extra_headers=None,
+        extra_query=None,
+        extra_body=None,
+        timeout=None,
+    ) -> Environment:
+        if not environment_id:
+            raise ValueError("environment_id is required")
+        return self._post(
+            f"{_ENV_PREFIX}/{environment_id}",
+            body=dump_body(
+                {"name": name, "config": config, "description": description, "metadata": metadata, "scope": scope}
+            ),
+            options=make_request_options(
+                extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
+            ),
+            cast_to=Environment,
+        )
+
+    def delete(
+        self, environment_id: str, *, extra_headers=None, extra_query=None, extra_body=None, timeout=None
+    ) -> DeleteEnvironmentResponse:
+        if not environment_id:
+            raise ValueError("environment_id is required")
+        return self._delete(
+            f"{_ENV_PREFIX}/{environment_id}",
+            options=make_request_options(
+                extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
+            ),
+            cast_to=DeleteEnvironmentResponse,
+        )
+
+
+class AsyncEnvironments(AsyncAPIResource):
+    async def create(
+        self,
+        *,
+        name: str,
+        config=NOT_GIVEN,
+        description=NOT_GIVEN,
+        metadata=NOT_GIVEN,
+        scope=NOT_GIVEN,
+        extra_headers=None,
+        extra_query=None,
+        extra_body=None,
+        timeout=None,
+    ) -> Environment:
+        return await self._post(
+            _ENV_PREFIX,
+            body=dump_body(
+                {"name": name, "config": config, "description": description, "metadata": metadata, "scope": scope}
+            ),
+            options=make_request_options(
+                extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
+            ),
+            cast_to=Environment,
+        )
+
+    async def retrieve(
+        self, environment_id: str, *, extra_headers=None, extra_query=None, extra_body=None, timeout=None
+    ) -> Environment:
+        if not environment_id:
+            raise ValueError("environment_id is required")
+        return await self._get(
+            f"{_ENV_PREFIX}/{environment_id}",
+            options=make_request_options(
+                extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
+            ),
+            cast_to=Environment,
+        )
+
+    async def list(
+        self, *, limit=NOT_GIVEN, page=NOT_GIVEN, extra_headers=None, extra_query=None, extra_body=None, timeout=None
+    ) -> ListEnvironmentsResponse:
+        return await self._get(
+            _ENV_PREFIX,
+            options=make_request_options(
+                query=_list_query(limit=limit, page=page),
+                extra_headers=extra_headers,
+                extra_query=extra_query,
+                extra_body=extra_body,
+                timeout=timeout,
+            ),
+            cast_to=ListEnvironmentsResponse,
+        )
+
+    async def update(
+        self,
+        environment_id: str,
+        *,
+        name=NOT_GIVEN,
+        config=NOT_GIVEN,
+        description=NOT_GIVEN,
+        metadata=NOT_GIVEN,
+        scope=NOT_GIVEN,
+        extra_headers=None,
+        extra_query=None,
+        extra_body=None,
+        timeout=None,
+    ) -> Environment:
+        if not environment_id:
+            raise ValueError("environment_id is required")
+        return await self._post(
+            f"{_ENV_PREFIX}/{environment_id}",
+            body=dump_body(
+                {"name": name, "config": config, "description": description, "metadata": metadata, "scope": scope}
+            ),
+            options=make_request_options(
+                extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
+            ),
+            cast_to=Environment,
+        )
+
+    async def delete(
+        self, environment_id: str, *, extra_headers=None, extra_query=None, extra_body=None, timeout=None
+    ) -> DeleteEnvironmentResponse:
+        if not environment_id:
+            raise ValueError("environment_id is required")
+        return await self._delete(
+            f"{_ENV_PREFIX}/{environment_id}",
+            options=make_request_options(
+                extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
+            ),
+            cast_to=DeleteEnvironmentResponse,
+        )

arkruntime/resources/files/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) 2026 ByteDance Ltd. and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+
+from arkruntime.resources.files.files import AsyncFiles, Files
+
+__all__ = ["Files", "AsyncFiles"]