feat: add async tasks

Author: xaviergoh-bytedance

volcenginesdkarkruntime/_base_client.py

@@ -719,6 +719,44 @@ async def post(
 
         return await self.request(cast_to, opts, stream=stream, stream_cls=stream_cls)
 
+    async def get(
+            self,
+            path: str,
+            *,
+            cast_to: Type[ResponseT],
+            params: list[tuple[str, str]] | None = None,
+            options: ExtraRequestOptions = {},
+            stream: bool = False,
+            stream_cls: type[_AsyncStreamT] | None = None,
+    ) -> ResponseT | _AsyncStreamT:
+        opts = RequestOptions.construct(
+            method="get",
+            url=path,
+            params=params,
+            **options,
+        )
+
+        return await self.request(cast_to, opts, stream=stream, stream_cls=stream_cls)
+
+    async def delete(
+            self,
+            path: str,
+            *,
+            cast_to: Type[ResponseT],
+            params: list[tuple[str, str]] | None = None,
+            options: ExtraRequestOptions = {},
+    ) -> ResponseT:
+        opts = RequestOptions.construct(
+            method="delete",
+            url=path,
+            params=params,
+            **options,
+        )
+
+        return await self.request(cast_to, opts)
+
+        return await self.request(cast_to, opts, stream=stream, stream_cls=stream_cls)
+
     async def request(
             self,
             cast_to: Type[ResponseT],

volcenginesdkarkruntime/_client.py

@@ -135,6 +135,7 @@ class AsyncArk(AsyncAPIClient):
     embeddings: resources.AsyncEmbeddings
     tokenization: resources.AsyncTokenization
     context: resources.AsyncContext
+    content_generation: resources.AsyncContentGeneration
 
     def __init__(
         self,
@@ -192,6 +193,7 @@ def __init__(
         self.embeddings = resources.AsyncEmbeddings(self)
         self.tokenization = resources.AsyncTokenization(self)
         self.context = resources.AsyncContext(self)
+        self.content_generation = resources.AsyncContentGeneration(self)
         # self.classification = resources.AsyncClassification(self)
 
     def _get_endpoint_sts_token(self, endpoint_id: str):

volcenginesdkarkruntime/_resource.py

@@ -20,3 +20,5 @@ class AsyncAPIResource:
     def __init__(self, client: "AsyncArk") -> None:
         self._client = client
         self._post = client.post
+        self._get = client.get
+        self._delete = client.delete

volcenginesdkarkruntime/_utils/_utils.py

@@ -97,6 +97,13 @@ def wrapper(*args, **kwargs):
 
     return wrapper
 
+def async_disallow_aksk(func):
+    async def wrapper(*args, **kwargs):
+        _restrict_aksk(args, kwargs)
+        return await func(*args, **kwargs)
+
+    return wrapper
+
 def _restrict_aksk(args, kwargs):
     assert len(args) > 0
 

volcenginesdkarkruntime/resources/__init__.py

@@ -4,7 +4,7 @@
 from .classification import Classification, AsyncClassification
 from .bot import BotChat, AsyncBotChat
 from .context import Context, AsyncContext
-from .content_generation import ContentGeneration
+from .content_generation import ContentGeneration, AsyncContentGeneration
 
 __all__ = [
     "Chat",
@@ -17,5 +17,6 @@
     "AsyncTokenization",
     "Context",
     "AsyncContext",
-    "ContentGeneration"
+    "ContentGeneration",
+    "AsyncContentGeneration"
 ]

volcenginesdkarkruntime/resources/content_generation/__init__.py

@@ -1,3 +1,4 @@
-from volcenginesdkarkruntime.resources.content_generation.content_generation import ContentGeneration
+from volcenginesdkarkruntime.resources.content_generation.content_generation import ContentGeneration, \
+    AsyncContentGeneration
 
-__all__ = ["ContentGeneration"]
+__all__ = ["ContentGeneration", "AsyncContentGeneration"]

volcenginesdkarkruntime/resources/content_generation/content_generation.py

@@ -1,14 +1,17 @@
 from __future__ import annotations
 
-from ..._resource import SyncAPIResource
+from ..._resource import SyncAPIResource, AsyncAPIResource
 from ..._compat import cached_property
 
-from .tasks import Tasks
-
-__all__ = ["ContentGeneration"]
+from .tasks import Tasks, AsyncTasks
 
 
 class ContentGeneration(SyncAPIResource):
     @cached_property
     def tasks(self) -> Tasks:
-        return Tasks(self._client)
+        return Tasks(self._client)
+
+class AsyncContentGeneration(AsyncAPIResource):
+    @cached_property
+    def tasks(self) -> AsyncTasks:
+        return AsyncTasks(self._client)

volcenginesdkarkruntime/resources/content_generation/tasks.py

@@ -1,19 +1,24 @@
 from __future__ import annotations
 
+from functools import cached_property
 from typing import Iterable, Union, List
 import httpx
 
+from ..._response import to_raw_response_wrapper, async_to_raw_response_wrapper
 from ..._types import Body, Query, Headers
-from ..._utils._utils import disallow_aksk
+from ..._utils._utils import disallow_aksk, async_disallow_aksk
 from ...types.content_generation.content_generation_task import ContentGenerationTask
 from ...types.content_generation.content_generation_task_id import ContentGenerationTaskID
 from volcenginesdkarkruntime._base_client import make_request_options
-from volcenginesdkarkruntime._resource import SyncAPIResource
+from volcenginesdkarkruntime._resource import SyncAPIResource, AsyncAPIResource
 from volcenginesdkarkruntime.types.content_generation.create_task_content_param import CreateTaskContentParam
 from ...types.content_generation.list_content_generation_tasks_response import ListContentGenerationTasksResponse
 
 
 class Tasks(SyncAPIResource):
+    @cached_property
+    def with_raw_response(self) -> TasksWithRawResponse:
+        return TasksWithRawResponse(self)
 
     @disallow_aksk
     def create(
@@ -126,3 +131,140 @@ def delete(
             cast_to=ContentGenerationTask,
         )
         return resp
+
+
+class AsyncTasks(AsyncAPIResource):
+    @cached_property
+    def with_raw_response(self) -> AsyncTasksWithRawResponse:
+        return AsyncTasksWithRawResponse(self)
+
+
+    @async_disallow_aksk
+    async def create(
+            self,
+            *,
+            model: str,
+            content: Iterable[CreateTaskContentParam],
+            extra_headers: Headers | None = None,
+            extra_query: Query | None = None,
+            extra_body: Body | None = None,
+            timeout: float | httpx.Timeout | None = None,
+    ) -> ContentGenerationTaskID:
+        resp = await self._post(
+            "/contents/generations/tasks",
+            body={
+                "model": model,
+                "content": content,
+            },
+            options=make_request_options(
+                extra_headers=extra_headers,
+                extra_query=extra_query,
+                extra_body=extra_body,
+                timeout=timeout,
+            ),
+            cast_to=ContentGenerationTaskID,
+        )
+        return resp
+
+
+    @async_disallow_aksk
+    async def get(
+            self,
+            *,
+            task_id: str,
+            extra_headers: Headers | None = None,
+            extra_query: Query | None = None,
+            extra_body: Body | None = None,
+            timeout: float | httpx.Timeout | None = None,
+    ) -> ContentGenerationTask:
+        resp = await self._get(
+            path=f"/contents/generations/tasks/{task_id}",
+            options=make_request_options(
+                extra_headers=extra_headers,
+                extra_query=extra_query,
+                extra_body=extra_body,
+                timeout=timeout,
+            ),
+            cast_to=ContentGenerationTask,
+        )
+        return resp
+
+    @async_disallow_aksk
+    async def list(
+            self,
+            page_num: int | None = None,
+            page_size: int | None = None,
+            status: str | None = None,
+            task_ids : Union[List[str], str] | None = None,
+            model: str | None = None,
+            extra_headers: Headers | None = None,
+            extra_body: Body | None = None,
+            extra_query: Query | None = None,
+            timeout: float | httpx.Timeout | None = None,
+    ) -> ListContentGenerationTasksResponse:
+
+        query_params = []
+        if page_num:
+            query_params.append(("page_num", page_num))
+        if page_size:
+            query_params.append(("page_size", page_size))
+        if status:
+            query_params.append(("filter.status", status))
+        if model:
+            query_params.append(("filter.model", model))
+        if task_ids:
+            if isinstance(task_ids, str):
+                task_ids = [task_ids]
+            for task_id in task_ids:
+                query_params.append(("filter.task_ids", task_id))
+
+        resp = await self._get(
+            path="/contents/generations/tasks",
+            params=query_params,
+            options=make_request_options(
+                extra_headers=extra_headers,
+                extra_query=extra_query,
+                extra_body=extra_body,
+                timeout=timeout,
+            ),
+            cast_to=ListContentGenerationTasksResponse,
+        )
+        return resp
+
+    @async_disallow_aksk
+    async def delete(
+            self,
+            task_id: str,
+            extra_headers: Headers | None = None,
+            extra_query: Query | None = None,
+            extra_body: Body | None = None,
+            timeout: float | httpx.Timeout | None = None,
+    ) -> None:
+        resp = await self._delete(
+            path=f"/contents/generations/tasks/{task_id}",
+            options=make_request_options(
+                extra_headers=extra_headers,
+                extra_query=extra_query,
+                extra_body=extra_body,
+                timeout=timeout,
+            ),
+            cast_to=ContentGenerationTask,
+        )
+        return resp
+
+
+class TasksWithRawResponse:
+    def __init__(self, tasks: Tasks) -> None:
+        self._tasks = tasks
+
+        self.create = to_raw_response_wrapper(
+            tasks.create,
+        )
+
+class AsyncTasksWithRawResponse:
+    def __init__(self, tasks: AsyncTasks) -> None:
+        self._tasks = tasks
+
+        self.create = async_to_raw_response_wrapper(
+            tasks.create,
+        )

volcenginesdkarkruntime/types/content_generation/create_task_content_param.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 from typing import Union
-from typing_extensions import Literal, Required, TypedDict
+from typing_extensions import Literal, TypedDict
 
 __all__ = [
     "CreateTaskContentTextParam",