diff --git a/dashscope/aigc/image_generation.py b/dashscope/aigc/image_generation.py index 7dbfbf0..6973108 100644 --- a/dashscope/aigc/image_generation.py +++ b/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 diff --git a/dashscope/aigc/image_synthesis.py b/dashscope/aigc/image_synthesis.py index 0409bfc..b9843b1 100644 --- a/dashscope/aigc/image_synthesis.py +++ b/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 diff --git a/dashscope/aigc/video_synthesis.py b/dashscope/aigc/video_synthesis.py index b976799..f13c155 100644 --- a/dashscope/aigc/video_synthesis.py +++ b/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 diff --git a/dashscope/api_entities/api_request_data.py b/dashscope/api_entities/api_request_data.py index 37a8593..c6c6692 100644 --- a/dashscope/api_entities/api_request_data.py +++ b/dashscope/api_entities/api_request_data.py @@ -158,6 +158,9 @@ def get_batch_binary_data(self) -> bytes: # type: ignore[return] def _only_parameters(self) -> str: obj = {"model": self.model, "parameters": self.parameters, "input": {}} + parameters = self.parameters.copy() + if "raw_input" in parameters: + obj["input"] = parameters.pop("raw_input") if self.task is not None: obj["task"] = self.task if self.task_group is not None: diff --git a/dashscope/audio/asr/transcription.py b/dashscope/audio/asr/transcription.py index 480b328..cf8ce4d 100644 --- a/dashscope/audio/asr/transcription.py +++ b/dashscope/audio/asr/transcription.py @@ -167,6 +167,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. @@ -174,7 +175,12 @@ def wait( 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) diff --git a/dashscope/audio/qwen_asr/qwen_transcription.py b/dashscope/audio/qwen_asr/qwen_transcription.py index 9ebed8d..859f0fc 100644 --- a/dashscope/audio/qwen_asr/qwen_transcription.py +++ b/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) diff --git a/dashscope/client/base_api.py b/dashscope/client/base_api.py index a609584..89094c3 100644 --- a/dashscope/client/base_api.py +++ b/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: