feat: batch-worker content generation

Author: xiangjie.meng

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)."""

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}")