feat: handle workflowai error

Author: Spatika-Ganesh

workflowai/core/client/api.py

@@ -89,10 +89,15 @@ async def delete(self, path: str) -> None:
             response = await client.delete(path)
             await self.raise_for_status(response)
 
-    def _extract_error(self, data: Union[bytes, str], exception: Optional[Exception] = None) -> WorkflowAIError:
+    def _extract_error(
+        self,
+        response: httpx.Response,
+        data: Union[bytes, str],
+        exception: Optional[Exception] = None,
+    ) -> WorkflowAIError:
         try:
             res = ErrorResponse.model_validate_json(data)
-            return WorkflowAIError(res.error, task_run_id=res.task_run_id)
+            return WorkflowAIError(error=res.error, task_run_id=res.task_run_id, response=response)
         except JSONDecodeError:
             raise WorkflowAIError(
                 error=BaseError(
@@ -101,6 +106,7 @@ def _extract_error(self, data: Union[bytes, str], exception: Optional[Exception]
                         "raw": str(data),
                     },
                 ),
+                response=response,
             ) from exception
 
     async def stream(
@@ -122,28 +128,8 @@ async def stream(
                     for payload in split_chunks(chunk):
                         yield returns.model_validate_json(payload)
                 except ValidationError as e:
-                    raise self._extract_error(payload, e) from None
+                    raise self._extract_error(response, payload, e) from None
 
     async def raise_for_status(self, response: httpx.Response):
-        if response.status_code < 200 and response.status_code >= 300:
-            try:
-                response_json = response.json()
-                r_error = response_json.get("error",{})
-                error_message = response_json.get("detail", {}) or r_error.get("message", "Unknown Error")
-                details =  r_error.get("details", {})
-                error_code = r_error.get("code", "unknown_error")
-                status_code = r_error.get("status_code", response.status_code)
-            except JSONDecodeError:
-                error_message = "Unknown error"
-                details = {"raw": response.content.decode()}
-                error_code ="unknown_error"
-                status_code = response.status_code
-
-            raise WorkflowAIError(
-                error=BaseError(
-                    message=error_message,
-                    details=details,
-                    status_code=status_code,
-                    code=error_code,
-                ),
-            ) from None
+        if response.status_code < 200 or response.status_code >= 300:
+            raise WorkflowAIError.from_response(response) from None

workflowai/core/domain/errors.py

@@ -1,5 +1,7 @@
+from json import JSONDecodeError
 from typing import Any, Literal, Optional, Union
 
+from httpx import Response
 from pydantic import BaseModel
 
 ProviderErrorCode = Literal[
@@ -68,9 +70,39 @@ class ErrorResponse(BaseModel):
 
 
 class WorkflowAIError(Exception):
-    def __init__(self, error: BaseError, task_run_id: Optional[str] = None):
+    def __init__(self, response:Response, error: BaseError, task_run_id: Optional[str] = None):
         self.error = error
         self.task_run_id = task_run_id
+        self.response = response
 
     def __str__(self):
         return f"WorkflowAIError : [{self.error.code}] ({self.error.status_code}): [{self.error.message}]"
+
+    @classmethod
+    def from_response(cls, response: Response):
+        try:
+            response_json = response.json()
+            r_error = response_json.get("error",{})
+            error_message = response_json.get("detail", {}) or r_error.get("message", "Unknown Error")
+            details =  r_error.get("details", {})
+            error_code = r_error.get("code", "unknown_error")
+            status_code = response.status_code
+            task_run_id = r_error.get("task_run_id", None)
+        except JSONDecodeError:
+            error_message = "Unknown error"
+            details = {"raw": response.content.decode()}
+            error_code ="unknown_error"
+            status_code = response.status_code
+            task_run_id=None
+
+        return cls(
+            response=response,
+            error=BaseError(
+                message=error_message,
+                details=details,
+                status_code=status_code,
+                code=error_code,
+            ),
+            task_run_id=task_run_id,
+        )
+

workflowai/core/domain/errors_test.py

@@ -0,0 +1,16 @@
+
+from unittest.mock import Mock
+
+from workflowai.core.domain.errors import WorkflowAIError
+
+
+def test_workflow_ai_error_404():
+    response = Mock()
+    response.status_code = 404
+    response.json = Mock()
+    response.json.return_value = {"error": {"message": "None", "status_code": 404, "code": "object_not_found"}}
+
+    error = WorkflowAIError.from_response(response)
+    assert error.error.message == "None"
+    assert error.error.status_code == 404
+    assert error.error.code == "object_not_found"