Update the files

Author: Spatika-Ganesh

workflowai/core/client/api.py

@@ -1,4 +1,4 @@
-import json
+from json import JSONDecodeError
 from typing import Any, AsyncIterator, Literal, Optional, TypeVar, Union, overload
 
 import httpx
@@ -91,13 +91,9 @@ async def delete(self, path: str) -> None:
 
     def _extract_error(self, data: Union[bytes, str], exception: Optional[Exception] = None) -> WorkflowAIError:
         try:
-            data_dict = json.loads(data)
-            if next(iter(data_dict.keys())) == "detail":
-                data_dict = {"error": {"detail": data_dict["detail"]}}
-                data = json.dumps(data_dict)
             res = ErrorResponse.model_validate_json(data)
             return WorkflowAIError(res.error, task_run_id=res.task_run_id)
-        except ValidationError:
+        except JSONDecodeError:
             raise WorkflowAIError(
                 error=BaseError(
                     message="Unknown error" if exception is None else str(exception),
@@ -129,5 +125,25 @@ async def stream(
                     raise self._extract_error(payload, e) from None
 
     async def raise_for_status(self, response: httpx.Response):
-        if response.status_code != 200:
-            raise self._extract_error(response.content) from None
+        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

workflowai/core/client/utils.py

@@ -4,11 +4,14 @@
 import asyncio
 import re
 from email.utils import parsedate_to_datetime
+from json import JSONDecodeError
 from time import time
 from typing import Any, Optional
 
 from httpx import HTTPStatusError
 
+from workflowai.core.domain.errors import BaseError, WorkflowAIError
+
 delimiter = re.compile(r'\}\n\ndata: \{"')
 
 
@@ -61,8 +64,29 @@ async def _wait_for_exception(e: HTTPStatusError):
         retry_after = retry_after_to_delay_seconds(e.response.headers.get("Retry-After"))
         leftover_delay = _leftover_delay()
         if not retry_after or leftover_delay < 0 or retry_count >= max_retry_count:
-            # TODO: convert error to WorkflowAIError
-            raise e
+            # Convert error to WorkflowAIError
+            try:
+                response_json = e.response.json()
+                r_err = response_json.get("error", {})
+                error_message = response_json.get("detail", {}) or r_err.get("message", "Unknown Error")
+                details = r_err.get("details", {})
+                error_code = r_err.get("code", "unknown_error")
+                status_code = r_err.get("status_code", e.response.status_code)
+            except JSONDecodeError:
+                error_message = "Unknown error"
+                details = {"raw": e.response.content.decode()}
+                error_code = "unknown_error"
+                status_code = e.response.status_code
+
+            raise WorkflowAIError(
+                error=BaseError(
+                    message=error_message,
+                    details=details,
+                    status_code=status_code,
+                    code=error_code,
+                ),
+            ) from None
+
         await asyncio.sleep(retry_after)
         retry_count += 1
 

workflowai/core/domain/errors.py

@@ -1,4 +1,3 @@
-import json
 from typing import Any, Literal, Optional, Union
 
 from pydantic import BaseModel
@@ -74,11 +73,4 @@ def __init__(self, error: BaseError, task_run_id: Optional[str] = None):
         self.task_run_id = task_run_id
 
     def __str__(self):
-        error_dict = {
-            "message": self.error.message,
-            "code": self.error.code,
-            "status_code": self.error.status_code,
-            "details": self.error.details,
-            "task_run_id": self.task_run_id,
-        }
-        return f"WorkflowAIError :  {json.dumps(error_dict, indent=4)}, Task Run ID : {self.task_run_id}"
+        return f"WorkflowAIError : [{self.error.code}] ({self.error.status_code}): [{self.error.message}]"