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

Author: songguocola

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/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,
@@ -606,10 +629,14 @@ def call(
             workspace=workspace,
             **kwargs,
         )
+        wait_timeout = -1
+        if "wait_timeout" in kwargs:
+            wait_timeout = kwargs.pop("wait_timeout")
         response = cls.wait(
             task_response,
             api_key=api_key,
             workspace=workspace,
+            wait_timeout=wait_timeout,
         )
         return response
 
@@ -767,6 +794,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 +803,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 +818,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 +829,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: