Merge pull request #22 from workflowai/yann/wor-2705-fix-errorresponse-validation-error

Fix ErrorResponse validation error

Author: guillaq

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "workflowai"
-version = "0.4.1"
+version = "0.4.2"
 description = ""
 authors = ["Guillaume Aquilina <guillaume@workflowai.com>"]
 readme = "README.md"
@@ -51,6 +51,8 @@ ignore = [
     "TD",
     "PYI051",
     "FIX002",
+    "SLF001", #reportPrivateUsage
+    "PT017", # Do not force using pytest.raises 
 ]
 
 # Allow fix for all enabled rules (when `--fix`) is provided.

workflowai/core/client/api.py

@@ -1,4 +1,3 @@
-from json import JSONDecodeError
 from typing import Any, AsyncIterator, Literal, Optional, TypeVar, Union, overload
 
 import httpx
@@ -99,7 +98,7 @@ def _extract_error(
         try:
             res = ErrorResponse.model_validate_json(data)
             return WorkflowAIError(error=res.error, task_run_id=res.task_run_id, response=response)
-        except JSONDecodeError:
+        except ValidationError:
             raise WorkflowAIError(
                 error=BaseError(
                     message="Unknown error" if exception is None else str(exception),
@@ -123,6 +122,11 @@ async def stream(
             content=data.model_dump_json(exclude_none=True),
             headers={"Content-Type": "application/json"},
         ) as response:
+            if not response.is_success:
+                # We need to read the response to get the error message
+                await response.aread()
+                response.raise_for_status()
+
             async for chunk in response.aiter_bytes():
                 payload = ""
                 try:

workflowai/core/client/api_test.py

@@ -0,0 +1,85 @@
+import httpx
+import pytest
+from pydantic import BaseModel
+from pytest_httpx import HTTPXMock
+
+from workflowai.core.client.api import APIClient
+from workflowai.core.domain.errors import WorkflowAIError
+
+
+class TestAPIClientExtractError:
+    def test_extract_error(self):
+        client = APIClient(endpoint="test_endpoint", api_key="test_api_key")
+
+        # Test valid JSON error response
+        response = httpx.Response(
+            status_code=400,
+            json={
+                "error": {
+                    "message": "Test error message",
+                    "details": {"key": "value"},
+                },
+                "task_run_id": "test_task_123",
+            },
+        )
+
+        error = client._extract_error(response, response.content)  # pyright:ignore[reportPrivateUsage]
+        assert isinstance(error, WorkflowAIError)
+        assert error.error.message == "Test error message"
+        assert error.error.details == {"key": "value"}
+        assert error.task_run_id == "test_task_123"
+        assert error.response == response
+
+    def test_extract_error_invalid_json(self):
+        client = APIClient(endpoint="test_endpoint", api_key="test_api_key")
+
+        # Test invalid JSON response
+        invalid_data = b"Invalid JSON data"
+        response = httpx.Response(status_code=400, content=invalid_data)
+
+        with pytest.raises(WorkflowAIError) as e:
+            client._extract_error(response, invalid_data)  # pyright:ignore[reportPrivateUsage]
+        assert isinstance(e.value, WorkflowAIError)
+        assert e.value.error.message == "Unknown error"
+        assert e.value.error.details == {"raw": "b'Invalid JSON data'"}
+        assert e.value.response == response
+
+    def test_extract_error_with_custom_error(self):
+        client = APIClient(endpoint="test_endpoint", api_key="test_api_key")
+
+        # Test with provided exception
+        invalid_data = "{'detail': 'Not Found'}"
+        response = httpx.Response(status_code=404, content=invalid_data)
+        exception = ValueError("Custom error")
+
+        with pytest.raises(WorkflowAIError) as e:
+            client._extract_error(response, invalid_data, exception)  # pyright:ignore[reportPrivateUsage]
+        assert isinstance(e.value, WorkflowAIError)
+        assert e.value.error.message == "Custom error"
+        assert e.value.error.details == {"raw": "{'detail': 'Not Found'}"}
+        assert e.value.response == response
+
+
+async def test_stream_404(httpx_mock: HTTPXMock):
+    class TestInputModel(BaseModel):
+        test_input: str
+
+    class TestOutputModel(BaseModel):
+        test_output: str
+
+    httpx_mock.add_response(status_code=404)
+
+    client = APIClient(endpoint="https://blabla.com", api_key="test_api_key")
+
+    try:
+        async for _ in client.stream(
+            method="GET",
+            path="test_path",
+            data=TestInputModel(test_input="test"),
+            returns=TestOutputModel,
+        ):
+            pass
+    except httpx.HTTPStatusError as e:
+        assert isinstance(e, httpx.HTTPStatusError)
+        assert e.response.status_code == 404
+        assert e.response.reason_phrase == "Not Found"