feat: support async images (#57)

examples/async_image_generator.py

@@ -0,0 +1,99 @@
+import asyncio
+import time
+
+from zai import ZaiClient
+
+
+class AsyncImageGenerator:
+    def __init__(self):
+        self.client = ZaiClient()
+
+    async def generate_image(
+        self,
+        prompt: str,
+        model: str = 'glm-image',
+        size: str = None,
+        quality: str = None,
+        max_wait_time: int = 300,
+    ):
+        """
+        Asynchronous image generation method
+
+        Args:
+            prompt: Image generation prompt
+            model: Model to use for image generation
+            size: Size of the generated image
+            quality: Quality level of the generated image
+            max_wait_time: Maximum wait time (seconds)
+
+        Returns:
+            Image generation result
+        """
+        try:
+            # Submit generation task
+            print('Submitting image generation task...')
+            response = self.client.images.async_generations(
+                model=model,
+                prompt=prompt,
+                size=size,
+                quality=quality,
+            )
+
+            print(f'Task submitted successfully, Task ID: {response.id}')
+            print(f'Initial response: {response}')
+
+            # Asynchronously wait for task completion
+            task_id = response.id
+            start_time = time.time()
+
+            while True:
+                # Check for timeout
+                if time.time() - start_time > max_wait_time:
+                    raise TimeoutError(f'Image generation timeout, exceeded {max_wait_time} seconds')
+
+                # Get task result
+                print(f'Querying task status... (waited {int(time.time() - start_time)} seconds)')
+                result = self.client.images.retrieve_images_result(id=task_id)
+
+                print(f'Task status: {result.task_status}')
+
+                # Check if task is completed
+                if result.task_status == 'SUCCESS':
+                    print('Image generation completed!')
+                    return result
+                elif result.task_status == 'FAIL':
+                    raise Exception(f'Image generation failed: {result}')
+                elif result.task_status in ['PROCESSING']:
+                    print(f'Task in progress, status: {result.task_status}')
+                else:
+                    print(f'Unknown status: {result.task_status}')
+
+                # Wait 3 seconds before querying again
+                await asyncio.sleep(30)
+
+        except Exception as e:
+            print(f'Error occurred during image generation: {e}')
+            raise
+
+
+# Usage example
+async def main():
+    generator = AsyncImageGenerator()
+    try:
+        result = await generator.generate_image(
+            prompt='A beautiful sunset over the ocean with colorful clouds',
+            model='glm-image',
+        )
+        print('\n=== Final Result ===')
+        print(f'Task Status: {result.task_status}')
+        print(f'Request ID: {result.request_id}')
+        if result.image_result:
+            for i, img in enumerate(result.image_result):
+                print(f'Image {i + 1} URL: {img.url}')
+
+    except Exception as e:
+        print(f'Generation failed: {e}')
+
+
+if __name__ == '__main__':
+    asyncio.run(main())

src/zai/api_resource/images/images.py

@@ -5,7 +5,7 @@ from typing import TYPE_CHECKING, Optional
 import httpx
 
 from zai.core import NOT_GIVEN, BaseAPI, Body, Headers, NotGiven, make_request_options
-from zai.types.image import ImagesResponded
+from zai.types.image import AsyncImagesResponded, ImagesResponded
 from zai.types.sensitive_word_check import SensitiveWordCheckRequest
 
 if TYPE_CHECKING:
@@ -84,3 +84,75 @@ class Images(BaseAPI):
 			cast_type=_cast_type,
 			stream=False,
 		)
+
+	def async_generations(
+		self,
+		*,
+		prompt: str,
+		model: str | NotGiven = NOT_GIVEN,
+		quality: Optional[str] | NotGiven = NOT_GIVEN,
+		size: Optional[str] | NotGiven = NOT_GIVEN,
+		request_id: Optional[str] | NotGiven = NOT_GIVEN,
+		user_id: Optional[str] | NotGiven = NOT_GIVEN,
+		extra_headers: Headers | None = None,
+		extra_body: Body | None = None,
+		timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
+		watermark_enabled: Optional[bool] | NotGiven = NOT_GIVEN,
+	) -> AsyncImagesResponded:
+		"""
+		Asynchronously generate images from text prompts. Only support glm-image model
+		Use retrieve_images_result() to poll for the result.
+
+		Arguments:
+			prompt (str): Text description of the desired image
+			model (str): The model to use for image generation
+			quality (Optional[str]): Quality level of the generated images
+			size (Optional[str]): Size of the generated images
+			request_id (Optional[str]): Unique identifier for the request
+			user_id (Optional[str]): User identifier
+			extra_headers (Headers): Additional headers to send
+			extra_body (Body): Additional body parameters
+			timeout (float | httpx.Timeout): Request timeout
+			watermark_enabled (Optional[bool]): Whether to enable watermark on generated images
+		"""
+		return self._post(
+			'/async/images/generations',
+			body={
+				'prompt': prompt,
+				'model': model,
+				'quality': quality,
+				'size': size,
+				'user_id': user_id,
+				'request_id': request_id,
+				'watermark_enabled': watermark_enabled,
+			},
+			options=make_request_options(extra_headers=extra_headers, extra_body=extra_body, timeout=timeout),
+			cast_type=AsyncImagesResponded,
+			stream=False,
+		)
+
+	def retrieve_images_result(
+		self,
+		id: str,
+		*,
+		extra_headers: Headers | None = None,
+		extra_body: Body | None = None,
+		timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
+	) -> AsyncImagesResponded:
+		"""
+		Retrieve the result of an async image generation operation
+
+		Arguments:
+			id (str): Unique identifier for the image generation task
+			extra_headers (Headers): Additional headers to send
+			extra_body (Body): Additional body parameters
+			timeout (float | httpx.Timeout): Request timeout
+		"""
+		if not id:
+			raise ValueError('`id` must be provided.')
+
+		return self._get(
+			f'/async-result/{id}',
+			options=make_request_options(extra_headers=extra_headers, extra_body=extra_body, timeout=timeout),
+			cast_type=AsyncImagesResponded,
+		)

src/zai/types/image/__init__.py

@@ -1,6 +1,7 @@
-from .image import GeneratedImage, ImagesResponded
+from .image import AsyncImagesResponded, GeneratedImage, ImagesResponded
 
 __all__ = [
 	'GeneratedImage',
 	'ImagesResponded',
+	'AsyncImagesResponded',
 ]

src/zai/types/image/image.py

@@ -31,3 +31,24 @@ class ImagesResponded(BaseModel):
 
 	created: int
 	data: List[GeneratedImage]
+
+
+class AsyncImagesResponded(BaseModel):
+	"""
+	Async image generation response
+
+	Attributes:
+		id (Optional[str]): The task order number generated by the Z.ai open platform.
+							Please use this order number when calling the request result interface.
+		model (str): Model name
+		request_id (str): The task number submitted by the user when requesting on the client
+							or the task number generated by the platform.
+		task_status (str): Processing status, PROCESSING (processing), SUCCESS (successful), FAIL (failed).
+		image_result (Optional[List[GeneratedImage]]): Image generation result, available when task_status is SUCCESS
+	"""
+
+	id: Optional[str] = None
+	model: str
+	request_id: str
+	task_status: str
+	image_result: Optional[List[GeneratedImage]] = None