feat(model/transcription): add timeout in sync wait api

Author: songguocola

dashscope/aigc/image_generation.py

@@ -177,6 +177,7 @@ def wait(
         task: Union[str, ImageGenerationResponse],  # type: ignore[override]
         api_key: str = None,
         workspace: str = None,
+        wait_timeout: int = -1,
         **kwargs,
     ) -> DashScopeAPIResponse:
         """Wait for image(s) synthesis task to complete, and return the result.
@@ -186,11 +187,18 @@ def wait(
                 ImageGenerationResponse return by async_call().
             api_key (str, optional): The api api_key. Defaults to None.
             workspace (str): The dashscope workspace id.
+            wait_timeout (int, optional): The maximum seconds to wait.
+                Default is -1 (no timeout).
 
         Returns:
             DashScopeAPIResponse: The task result.
         """
-        response = super().wait(task, api_key, workspace=workspace)
+        response = super().wait(
+            task,
+            api_key,
+            workspace=workspace,
+            wait_timeout=wait_timeout,
+        )
         return ImageGenerationResponse.from_api_response(response)
 
     @classmethod
@@ -487,6 +495,7 @@ async def wait(
         task: Union[str, ImageGenerationResponse],  # type: ignore[override]
         api_key: str = None,
         workspace: str = None,
+        wait_timeout: int = -1,
         **kwargs,
     ) -> DashScopeAPIResponse:
         """Wait for image(s) synthesis task to complete, and return the result.
@@ -496,11 +505,18 @@ async def wait(
                 ImageGenerationResponse return by async_call().
             api_key (str, optional): The api api_key. Defaults to None.
             workspace (str): The dashscope workspace id.
+            wait_timeout (int, optional): The maximum seconds to wait.
+                Default is -1 (no timeout).
 
         Returns:
             DashScopeAPIResponse: The task result.
         """
-        response = await super().wait(task, api_key, workspace=workspace)
+        response = await super().wait(
+            task,
+            api_key,
+            workspace=workspace,
+            wait_timeout=wait_timeout,
+        )
         return ImageGenerationResponse.from_api_response(response)
 
     @classmethod

dashscope/aigc/image_synthesis.py

@@ -739,6 +739,7 @@ async def wait(
         task: Union[str, ImageSynthesisResponse],  # type: ignore[override]
         api_key: str = None,
         workspace: str = None,
+        wait_timeout: int = -1,
         **kwargs,
     ) -> ImageSynthesisResponse:
         """Wait for image(s) synthesis task to complete, and return the result.
@@ -748,11 +749,18 @@ async def wait(
                 ImageSynthesisResponse return by async_call().
             api_key (str, optional): The api api_key. Defaults to None.
             workspace (str): The dashscope workspace id.
+            wait_timeout (int, optional): The maximum seconds to wait.
+                Default is -1 (no timeout).
 
         Returns:
             ImageSynthesisResponse: The task result.
         """
-        response = await super().wait(task, api_key, workspace=workspace)
+        response = await super().wait(
+            task,
+            api_key,
+            workspace=workspace,
+            wait_timeout=wait_timeout,
+        )
         return ImageSynthesisResponse.from_api_response(response)
 
     @classmethod

dashscope/aigc/video_synthesis.py

@@ -894,6 +894,7 @@ async def wait(
         task: Union[str, VideoSynthesisResponse],  # type: ignore[override]
         api_key: str = None,
         workspace: str = None,
+        wait_timeout: int = -1,
         **kwargs,
     ) -> VideoSynthesisResponse:
         """Wait for video synthesis task to complete, and return the result.
@@ -903,11 +904,18 @@ async def wait(
                 VideoSynthesisResponse return by async_call().
             api_key (str, optional): The api api_key. Defaults to None.
             workspace (str): The dashscope workspace id.
+            wait_timeout (int, optional): The maximum seconds to wait.
+                Default is -1 (no timeout).
 
         Returns:
             VideoSynthesisResponse: The task result.
         """
-        response = await super().wait(task, api_key, workspace=workspace)
+        response = await super().wait(
+            task,
+            api_key,
+            workspace=workspace,
+            wait_timeout=wait_timeout,
+        )
         return VideoSynthesisResponse.from_api_response(response)
 
     @classmethod

dashscope/audio/asr/transcription.py

@@ -167,14 +167,20 @@ def wait(
         task: Union[str, TranscriptionResponse],  # type: ignore[override]
         api_key: str = None,
         workspace: str = None,
+        wait_timeout: int = -1,
         **kwargs,
     ) -> TranscriptionResponse:
         """Poll task until the final results of transcription is obtained.
 
         Args:
             task (Union[str, TranscriptionResponse]): The task_id or
                 response including task_id returned from async_call().
+            api_key (str, optional): The api_key. Defaults to None.
             workspace (str): The dashscope workspace id.
+            wait_timeout (int, optional): The timeout for waiting.
+                Defaults to -1.That means no timeout.
+                If set to a value > 0, the task does not complete
+                within this time, a timeout error response will be returned.
 
         Returns:
             TranscriptionResponse: The result of batch transcription.
@@ -183,6 +189,7 @@ def wait(
             task,
             api_key=api_key,
             workspace=workspace,
+            wait_timeout=wait_timeout,
             **kwargs,
         )
         return TranscriptionResponse.from_api_response(response)

dashscope/audio/qwen_asr/qwen_transcription.py

@@ -127,6 +127,7 @@ def wait(
         task: Union[str, TranscriptionResponse],  # type: ignore[override]
         api_key: str = None,
         workspace: str = None,
+        wait_timeout: int = -1,
         **kwargs,
     ) -> TranscriptionResponse:
         """Poll task until the final results of transcription is obtained.
@@ -135,6 +136,8 @@ def wait(
             task (Union[str, TranscriptionResponse]): The task_id or
                 response including task_id returned from async_call().
             workspace (str): The dashscope workspace id.
+            wait_timeout (int, optional): The maximum seconds to wait.
+                Default is -1 (no timeout).
 
         Returns:
             TranscriptionResponse: The result of batch transcription.
@@ -143,6 +146,7 @@ def wait(
             task,
             api_key=api_key,
             workspace=workspace,
+            wait_timeout=wait_timeout,
             **kwargs,
         )
         return TranscriptionResponse.from_api_response(response)

dashscope/client/base_api.py

@@ -191,6 +191,7 @@ async def wait(
         task: Union[str, DashScopeAPIResponse],
         api_key: str = None,
         workspace: str = None,
+        wait_timeout: int = -1,
         **kwargs,
     ) -> DashScopeAPIResponse:
         """Wait for async task completion and return task result.
@@ -199,6 +200,12 @@ async def wait(
             task (Union[str, DashScopeAPIResponse]): The task_id, or
                 async_call response.
             api_key (str, optional): The api_key. Defaults to None.
+            workspace (str, optional): The dashscope workspace id.
+            wait_timeout (int, optional): The maximum seconds to wait
+                for the task to complete. Default is -1, which means no
+                timeout. When set to a value > 0, if the task does not
+                complete within this time, a timeout error response will
+                be returned.
 
         Returns:
             DashScopeAPIResponse: The async task information.
@@ -208,6 +215,7 @@ async def wait(
         max_wait_seconds = 5
         increment_steps = 3
         step = 0
+        start_time = time.time()
         while True:
             step += 1
             # we start by querying once every second, and double
@@ -217,6 +225,21 @@ async def wait(
             # (server side return immediately when ready)
             if wait_seconds < max_wait_seconds and step % increment_steps == 0:
                 wait_seconds = min(wait_seconds * 2, max_wait_seconds)
+            if 0 < wait_timeout <= (time.time() - start_time):
+                logger.warning(
+                    "Wait task: %s timeout after %s seconds.",
+                    task_id,
+                    wait_timeout,
+                )
+                return DashScopeAPIResponse(
+                    request_id=task_id,
+                    status_code=HTTPStatus.REQUEST_TIMEOUT,
+                    code="WaitTaskTimeout",
+                    message=(
+                        f"Wait task: {task_id} timeout after "
+                        f"{wait_timeout} seconds."
+                    ),
+                )
             rsp = await cls._get(
                 task_id,
                 api_key,
@@ -600,6 +623,10 @@ def call(
         **kwargs,
     ) -> DashScopeAPIResponse:
         """Call service and get result."""
+        wait_timeout = -1
+        if "wait_timeout" in kwargs:
+            wait_timeout = kwargs.pop("wait_timeout")
+
         task_response = cls.async_call(  # type: ignore[misc]
             *args,
             api_key=api_key,
@@ -610,6 +637,7 @@ def call(
             task_response,
             api_key=api_key,
             workspace=workspace,
+            wait_timeout=wait_timeout,
         )
         return response
 
@@ -767,6 +795,7 @@ def wait(
         task: Union[str, DashScopeAPIResponse],
         api_key: str = None,
         workspace: str = None,
+        wait_timeout: int = -1,
         **kwargs,
     ) -> DashScopeAPIResponse:
         """Wait for async task completion and return task result.
@@ -775,6 +804,12 @@ def wait(
             task (Union[str, DashScopeAPIResponse]): The task_id, or
                 async_call response.
             api_key (str, optional): The api_key. Defaults to None.
+            workspace (str, optional): The dashscope workspace id.
+            wait_timeout (int, optional): The maximum seconds to wait
+                for the task to complete. Default is -1, which means no
+                timeout. When set to a value > 0, if the task does not
+                complete within this time, a timeout error response will
+                be returned.
 
         Returns:
             DashScopeAPIResponse: The async task information.
@@ -784,6 +819,7 @@ def wait(
         max_wait_seconds = 5
         increment_steps = 3
         step = 0
+        start_time = time.time()
         while True:
             step += 1
             # we start by querying once every second, and double
@@ -794,6 +830,21 @@ def wait(
             # (server side return immediately when ready)
             if wait_seconds < max_wait_seconds and step % increment_steps == 0:
                 wait_seconds = min(wait_seconds * 2, max_wait_seconds)
+            if 0 < wait_timeout <= (time.time() - start_time):
+                logger.warning(
+                    "Wait task: %s timeout after %s seconds.",
+                    task_id,
+                    wait_timeout,
+                )
+                return DashScopeAPIResponse(
+                    request_id=task_id,
+                    status_code=HTTPStatus.REQUEST_TIMEOUT,
+                    code="WaitTaskTimeout",
+                    message=(
+                        f"Wait task: {task_id} timeout after "
+                        f"{wait_timeout} seconds."
+                    ),
+                )
             rsp = cls._get(task_id, api_key, workspace=workspace, **kwargs)
             if rsp.status_code == HTTPStatus.OK:
                 if rsp.output is None: