feat: add content generation task sync methods

xaviergoh-bytedance · xaviergoh-bytedance · commit 06db106fd6dc · 2025-01-09T12:01:02.000+08:00
diff --git a/volcenginesdkarkruntime/_base_client.py b/volcenginesdkarkruntime/_base_client.py
@@ -557,6 +557,44 @@ def post(
             ResponseT, self.request(cast_to, opts, stream=stream, stream_cls=stream_cls)
         )
 
+    def get(
+            self,
+            path: str,
+            *,
+            cast_to: Type[ResponseT],
+            params: list[tuple[str, str]] | None = None,
+            options: ExtraRequestOptions = {},
+            stream: bool = False,
+            stream_cls: type[_StreamT] | None = None,
+    ) -> ResponseT | _StreamT:
+        opts = RequestOptions.construct(
+            method="get",
+            url=path,
+            params=params,
+            **options,
+        )
+
+        return cast(
+            ResponseT, self.request(cast_to, opts, stream=stream, stream_cls=stream_cls)
+        )
+
+    def delete(
+            self,
+            path: str,
+            *,
+            cast_to: Type[ResponseT],
+            params: list[tuple[str, str]] | None = None,
+            options: ExtraRequestOptions = {},
+    ) -> ResponseT:
+        opts = RequestOptions.construct(  # type: ignore
+            method="delete",
+            url=path,
+            params=params,
+            **options,
+        )
+
+        return cast(ResponseT, self.request(cast_to, opts))
+
     def request(
             self,
             cast_to: Type[ResponseT],
diff --git a/volcenginesdkarkruntime/_client.py b/volcenginesdkarkruntime/_client.py
@@ -38,6 +38,7 @@ class Ark(SyncAPIClient):
     embeddings: resources.Embeddings
     tokenization: resources.Tokenization
     context: resources.Context
+    content_generation: resources.ContentGeneration
 
     def __init__(
         self,
@@ -96,6 +97,7 @@ def __init__(
         self.embeddings = resources.Embeddings(self)
         self.tokenization = resources.Tokenization(self)
         self.context = resources.Context(self)
+        self.content_generation = resources.ContentGeneration(self)
         # self.classification = resources.Classification(self)
 
     def _get_endpoint_sts_token(self, endpoint_id: str):
diff --git a/volcenginesdkarkruntime/_resource.py b/volcenginesdkarkruntime/_resource.py
@@ -10,6 +10,8 @@ class SyncAPIResource:
     def __init__(self, client: "Ark") -> None:
         self._client = client
         self._post = client.post
+        self._get = client.get
+        self._delete = client.delete
 
 
 class AsyncAPIResource:
diff --git a/volcenginesdkarkruntime/resources/__init__.py b/volcenginesdkarkruntime/resources/__init__.py
@@ -4,6 +4,7 @@
 from .classification import Classification, AsyncClassification
 from .bot import BotChat, AsyncBotChat
 from .context import Context, AsyncContext
+from .content_generation import ContentGeneration
 
 __all__ = [
     "Chat",
@@ -15,5 +16,6 @@
     "Tokenization",
     "AsyncTokenization",
     "Context",
-    "AsyncContext"
+    "AsyncContext",
+    "ContentGeneration"
 ]
diff --git a/volcenginesdkarkruntime/resources/content_generation/__init__.py b/volcenginesdkarkruntime/resources/content_generation/__init__.py
@@ -0,0 +1,3 @@
+from volcenginesdkarkruntime.resources.content_generation.content_generation import ContentGeneration
+
+__all__ = ["ContentGeneration"]
diff --git a/volcenginesdkarkruntime/resources/content_generation/content_generation.py b/volcenginesdkarkruntime/resources/content_generation/content_generation.py
@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+from ..._resource import SyncAPIResource
+from ..._compat import cached_property
+
+from .tasks import Tasks
+
+__all__ = ["ContentGeneration"]
+
+
+class ContentGeneration(SyncAPIResource):
+    @cached_property
+    def tasks(self) -> Tasks:
+        return Tasks(self._client)
diff --git a/volcenginesdkarkruntime/resources/content_generation/tasks.py b/volcenginesdkarkruntime/resources/content_generation/tasks.py
@@ -0,0 +1,122 @@
+from __future__ import annotations
+
+from typing import Iterable, Union, List
+import httpx
+
+from ..._types import Body, Query, Headers
+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.types.content_generation.create_task_content_param import CreateTaskContentParam
+from ...types.content_generation.list_content_generation_tasks_response import ListContentGenerationTasksResponse
+
+
+class Tasks(SyncAPIResource):
+    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 = 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
+
+
+    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 = 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
+
+    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 = 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
+
+    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,
+    ):
+        resp = 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
diff --git a/volcenginesdkarkruntime/types/content_generation/__init__.py b/volcenginesdkarkruntime/types/content_generation/__init__.py
diff --git a/volcenginesdkarkruntime/types/content_generation/content_generation_task.py b/volcenginesdkarkruntime/types/content_generation/content_generation_task.py
@@ -0,0 +1,38 @@
+
+__all__ = ["ContentGenerationTask"]
+
+from volcenginesdkarkruntime._models import BaseModel
+
+
+class Usage(BaseModel):
+    completion_tokens: int
+    """The number of tokens used for completion."""
+
+class Content(BaseModel):
+    video_url: str
+    """The URL of the generated video, if any."""
+
+class ContentGenerationTask(BaseModel):
+    id: str
+    """A unique identifier for the task."""
+
+    model: str
+    """The model used for the task."""
+
+    status: str
+    """The status of the task (running, failed, queued, succeeded, cancelled)."""
+
+    failure_reason: str
+    """The reason for failure, if applicable."""
+
+    content: Content
+    """The content generated by the task."""
+
+    usage: Usage
+    """The usage information for the task."""
+
+    created_at: int
+    """The Unix timestamp when the task was created."""
+
+    updated_at: int
+    """The Unix timestamp when the task was last updated."""
diff --git a/volcenginesdkarkruntime/types/content_generation/content_generation_task_id.py b/volcenginesdkarkruntime/types/content_generation/content_generation_task_id.py
@@ -0,0 +1,6 @@
+from volcenginesdkarkruntime._models import BaseModel
+
+
+class ContentGenerationTaskID(BaseModel):
+    id: str
+    """A unique identifier for the task."""
diff --git a/volcenginesdkarkruntime/types/content_generation/create_task_content_param.py b/volcenginesdkarkruntime/types/content_generation/create_task_content_param.py
@@ -0,0 +1,33 @@
+from __future__ import annotations
+from typing import Union
+from typing_extensions import Literal, Required, TypedDict
+
+__all__ = [
+    "CreateTaskContentTextParam",
+    "CreateTaskContentImageDataParam",
+    "CreateTaskContentImageParam",
+    "CreateTaskContentParam",
+]
+
+
+class CreateTaskContentTextParam(TypedDict):
+    type: Literal["text"]
+    text: str
+    """Text prompt describing the video generation job."""
+
+
+class CreateTaskContentImageDataParam(TypedDict):
+    url: str
+    """URL of the image to be used for content generation."""
+
+
+class CreateTaskContentImageParam(TypedDict):
+    type: Literal["image_url"]
+    image_url: CreateTaskContentImageDataParam
+    """Image data object containing image URL."""
+
+
+CreateTaskContentParam = Union[
+    CreateTaskContentTextParam,
+    CreateTaskContentImageParam
+]
diff --git a/volcenginesdkarkruntime/types/content_generation/list_content_generation_tasks_response.py b/volcenginesdkarkruntime/types/content_generation/list_content_generation_tasks_response.py
@@ -0,0 +1,49 @@
+from typing import List
+from pydantic import BaseModel
+
+__all__ = ["ListContentGenerationTasksResponse"]
+
+
+class Content(BaseModel):
+    video_url: str
+    """URL of the generated video."""
+
+
+class Usage(BaseModel):
+    completion_tokens: int
+    """Number of tokens used during the content generation process."""
+
+
+class Item(BaseModel):
+    id: str
+    """Unique identifier for the content generation task."""
+
+    model: str
+    """Name or identifier of the model used for generation."""
+
+    status: str
+    """Status of the content generation task (e.g., 'succeed', 'failed')."""
+
+    failure_reason: str
+    """Description of why the task failed, if it did. Otherwise empty."""
+
+    content: Content
+    """Generated content details (e.g., video URL)."""
+
+    usage: Usage
+    """Usage statistics for the content generation task."""
+
+    created_at: int
+    """Timestamp (in Unix epoch) of when the task was created."""
+
+    updated_at: int
+    """Timestamp (in Unix epoch) of when the task was last updated."""
+
+
+class ListContentGenerationTasksResponse(BaseModel):
+    total: int
+    """Total number of filtered content generation tasks."""
+
+    items: List[Item]
+    """List of content generation task items."""
+

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+from volcenginesdkarkruntime.resources.content_generation.content_generation import ContentGeneration`
	`2`	`+`
	`3`	`+__all__ = ["ContentGeneration"]`