refactor: introduce dedicated error types for workflow invocation polling

Author: Danite

src/runwayml/__init__.py

@@ -36,7 +36,12 @@
     UnprocessableEntityError,
     APIResponseValidationError,
 )
-from .lib.polling import TaskFailedError, TaskTimeoutError
+from .lib.polling import (
+    TaskFailedError,
+    TaskTimeoutError,
+    WorkflowInvocationFailedError,
+    WorkflowInvocationTimeoutError,
+)
 from ._base_client import DefaultHttpxClient, DefaultAioHttpClient, DefaultAsyncHttpxClient
 from ._utils._logs import setup_logging as _setup_logging
 
@@ -84,6 +89,8 @@
     "DefaultAioHttpClient",
     "TaskFailedError",
     "TaskTimeoutError",
+    "WorkflowInvocationFailedError",
+    "WorkflowInvocationTimeoutError",
 ]
 
 if not _t.TYPE_CHECKING:

src/runwayml/lib/polling.py

@@ -127,13 +127,13 @@ async def wait_for_task_output(self, timeout: Union[float, None] = 60 * 10) -> T
 
 
 class TaskFailedError(Exception):
-    def __init__(self, task_details: Union[TaskRetrieveResponse, WorkflowInvocationRetrieveResponse]):
+    def __init__(self, task_details: TaskRetrieveResponse):
         self.task_details = task_details
         super().__init__(f"Task failed")
 
 
 class TaskTimeoutError(Exception):
-    def __init__(self, task_details: Union[TaskRetrieveResponse, WorkflowInvocationRetrieveResponse]):
+    def __init__(self, task_details: TaskRetrieveResponse):
         self.task_details = task_details
         super().__init__(f"Task timed out")
 
@@ -265,12 +265,14 @@ def wait_for_task_output(self, timeout: Union[float, None] = 60 * 10) -> Workflo
         """
         When called, this will block until the workflow invocation is complete.
 
-        If the invocation fails or is cancelled, a `TaskFailedError` will be raised.
+        If the invocation fails or is cancelled, a `WorkflowInvocationFailedError`
+        will be raised.
 
         Args:
           timeout: The maximum amount of time to wait in seconds. If not specified,
-              the default timeout is 10 minutes. Will raise a `TaskTimeoutError` if
-              the invocation does not complete within the timeout.
+              the default timeout is 10 minutes. Will raise a
+              `WorkflowInvocationTimeoutError` if the invocation does not complete
+              within the timeout.
 
         Returns:
           The workflow invocation details, equivalent to calling
@@ -284,13 +286,15 @@ async def wait_for_task_output(self, timeout: Union[float, None] = 60 * 10) -> W
         """
         When called, this will wait until the workflow invocation is complete.
 
-        If the invocation fails or is cancelled, a `TaskFailedError` will be raised.
+        If the invocation fails or is cancelled, a `WorkflowInvocationFailedError`
+        will be raised.
 
         Args:
           timeout: The maximum amount of time to wait in seconds. If not specified,
-              the default timeout is 10 minutes. Will raise a `TaskTimeoutError` if
-              the invocation does not complete within the timeout. Setting this to
-              `None` will wait indefinitely (disabling the timeout).
+              the default timeout is 10 minutes. Will raise a
+              `WorkflowInvocationTimeoutError` if the invocation does not complete
+              within the timeout. Setting this to `None` will wait indefinitely
+              (disabling the timeout).
 
         Returns:
           The workflow invocation details, equivalent to awaiting
@@ -361,6 +365,18 @@ class AsyncAwaitableWISucceeded(AsyncAwaitableWorkflowInvocationResponseMixin, W
 ]
 
 
+class WorkflowInvocationFailedError(Exception):
+    def __init__(self, invocation_details: WorkflowInvocationRetrieveResponse):
+        self.invocation_details = invocation_details
+        super().__init__("Workflow invocation failed")
+
+
+class WorkflowInvocationTimeoutError(Exception):
+    def __init__(self, invocation_details: WorkflowInvocationRetrieveResponse):
+        self.invocation_details = invocation_details
+        super().__init__("Workflow invocation timed out")
+
+
 def _make_sync_wait_for_workflow_invocation_output(
     client: "RunwayML",
 ) -> Callable[["AwaitableWorkflowInvocationResponseMixin", Union[float, None]], WorkflowInvocationRetrieveResponse]:
@@ -374,9 +390,9 @@ def wait_for_task_output(
             if details.status == "SUCCEEDED":
                 return details
             if details.status == "FAILED" or details.status == "CANCELLED":
-                raise TaskFailedError(details)
+                raise WorkflowInvocationFailedError(details)
             if timeout is not None and time.time() - start_time > timeout:
-                raise TaskTimeoutError(details)
+                raise WorkflowInvocationTimeoutError(details)
 
     return wait_for_task_output
 
@@ -406,9 +422,9 @@ async def wait_for_task_output(
             if details.status == "SUCCEEDED":
                 return details
             if details.status == "FAILED" or details.status == "CANCELLED":
-                raise TaskFailedError(details)
+                raise WorkflowInvocationFailedError(details)
             if timeout is not None and anyio.current_time() - start_time > timeout:
-                raise TaskTimeoutError(details)
+                raise WorkflowInvocationTimeoutError(details)
 
     return wait_for_task_output
 

tests/test_workflow_invocation_polling.py

@@ -6,7 +6,7 @@
 
 import pytest
 
-from runwayml import RunwayML, AsyncRunwayML, TaskFailedError, TaskTimeoutError
+from runwayml import RunwayML, AsyncRunwayML, WorkflowInvocationFailedError, WorkflowInvocationTimeoutError
 from runwayml.types.workflow_invocation_retrieve_response import (
     Failed,
     Pending,
@@ -84,10 +84,10 @@ def test_raises_on_failed(self, _mock_sleep: MagicMock) -> None:
         response = _make_pending()
         patched: Any = inject_sync_workflow_invocation_wait_method(client, response)
 
-        with pytest.raises(TaskFailedError) as exc_info:
+        with pytest.raises(WorkflowInvocationFailedError) as exc_info:
             patched.wait_for_task_output()
 
-        assert exc_info.value.task_details.status == "FAILED"
+        assert exc_info.value.invocation_details.status == "FAILED"
 
     @patch("runwayml.lib.polling.time.sleep", return_value=None)
     def test_raises_on_cancelled(self, _mock_sleep: MagicMock) -> None:
@@ -100,10 +100,10 @@ def test_raises_on_cancelled(self, _mock_sleep: MagicMock) -> None:
         response = _make_pending()
         patched: Any = inject_sync_workflow_invocation_wait_method(client, response)
 
-        with pytest.raises(TaskFailedError) as exc_info:
+        with pytest.raises(WorkflowInvocationFailedError) as exc_info:
             patched.wait_for_task_output()
 
-        assert exc_info.value.task_details.status == "CANCELLED"
+        assert exc_info.value.invocation_details.status == "CANCELLED"
 
     @patch("runwayml.lib.polling.time.time")
     @patch("runwayml.lib.polling.time.sleep", return_value=None)
@@ -119,7 +119,7 @@ def test_raises_on_timeout(self, _mock_sleep: MagicMock, mock_time: MagicMock) -
         response = _make_pending()
         patched: Any = inject_sync_workflow_invocation_wait_method(client, response)
 
-        with pytest.raises(TaskTimeoutError):
+        with pytest.raises(WorkflowInvocationTimeoutError):
             patched.wait_for_task_output(timeout=600)
 
 
@@ -169,10 +169,10 @@ async def test_raises_on_failed(self, _mock_time: MagicMock, _mock_sleep: AsyncM
         response = _make_pending()
         patched: Any = inject_async_workflow_invocation_wait_method(client, response)
 
-        with pytest.raises(TaskFailedError) as exc_info:
+        with pytest.raises(WorkflowInvocationFailedError) as exc_info:
             await patched.wait_for_task_output()
 
-        assert exc_info.value.task_details.status == "FAILED"
+        assert exc_info.value.invocation_details.status == "FAILED"
 
     @patch("runwayml.lib.polling.anyio.sleep", new_callable=AsyncMock)
     @patch("runwayml.lib.polling.anyio.current_time", return_value=0.0)
@@ -186,10 +186,10 @@ async def test_raises_on_cancelled(self, _mock_time: MagicMock, _mock_sleep: Asy
         response = _make_pending()
         patched: Any = inject_async_workflow_invocation_wait_method(client, response)
 
-        with pytest.raises(TaskFailedError) as exc_info:
+        with pytest.raises(WorkflowInvocationFailedError) as exc_info:
             await patched.wait_for_task_output()
 
-        assert exc_info.value.task_details.status == "CANCELLED"
+        assert exc_info.value.invocation_details.status == "CANCELLED"
 
     @patch("runwayml.lib.polling.anyio.sleep", new_callable=AsyncMock)
     @patch("runwayml.lib.polling.anyio.current_time")
@@ -205,5 +205,5 @@ async def test_raises_on_timeout(self, mock_time: MagicMock, _mock_sleep: AsyncM
         response = _make_pending()
         patched: Any = inject_async_workflow_invocation_wait_method(client, response)
 
-        with pytest.raises(TaskTimeoutError):
+        with pytest.raises(WorkflowInvocationTimeoutError):
             await patched.wait_for_task_output(timeout=600)