Transform API error type (#845)

* Transform API errors * Fix transform test error * fix tests
refuel-ai · Jun 6, 2024 · aba8336 · aba8336
1 parent d755ff0
commit aba8336
Show file tree

Hide file tree

Showing 6 changed files with 31 additions and 18 deletions.
diff --git a/src/autolabel/transforms/custom_api.py b/src/autolabel/transforms/custom_api.py
@@ -84,31 +84,37 @@ async def _get_result(
                 TransformErrorType.MAX_RETRIES_REACHED,
                 f"Max retries reached for URL: {url}",
             )
-
+        request_url = url
         try:
             client = self.client
             if not verify:
                 client = self.client_with_no_verify
             response = await client.get(url, headers=headers, params=params)
+            request_url = response.url
             response.raise_for_status()
             return response.text
         except self.httpx.ConnectTimeout as e:
-            logger.error(f"Timeout when fetching content from URL: {url}")
+            logger.error(f"Timeout when making request to URL: {request_url}")
             raise TransformError(
                 TransformErrorType.TRANSFORM_TIMEOUT,
-                "Timeout when fetching content from URL",
+                f"Timeout when making request to URL: {request_url}",
             )
         except ssl.SSLCertVerificationError as e:
             logger.warning(
-                f"SSL verification error when fetching content from URL: {url}, retrying with verify=False"
+                f"SSL verification error when making request to URL: {request_url}, retrying with verify=False"
             )
             await asyncio.sleep(BACKOFF**retry_count)
             return await self._get_result(
                 url, verify=False, headers=headers, retry_count=retry_count + 1
             )
         except Exception as e:
-            logger.error(f"Error fetching content from URL: {url}. Exception: {e}")
-            raise e
+            logger.error(
+                f"Error when making request to URL: {request_url}. Exception: {str(e)}"
+            )
+            raise TransformError(
+                TransformErrorType.TRANSFORM_API_ERROR,
+                f"Error when making request to URL: {request_url}",
+            )
 
     async def _apply(self, row: Dict[str, Any]) -> Dict[str, Any]:
         params = {}

diff --git a/src/autolabel/transforms/schema.py b/src/autolabel/transforms/schema.py
@@ -52,11 +52,10 @@ class TransformErrorType(str, Enum):
     """Transform error types"""
 
     INVALID_INPUT = "INVALID_INPUT"
-    TRANSFORM_ERROR = "TRANSFORM_ERROR"
-    TRANSFORM_TIMEOUT = "TRANSFORM_TIMEOUT"
+    TRANSFORM_ERROR = "ENRICHMENT_ERROR"
+    TRANSFORM_TIMEOUT = "ENRICHMENT_TIMEOUT"
     MAX_RETRIES_REACHED = "MAX_RETRIES_REACHED"
-    SERP_API_ERROR = "SERP_API_ERROR"
-    SERPER_API_ERROR = "SERPER_API_ERROR"
+    TRANSFORM_API_ERROR = "ENRICHMENT_API_ERROR"
 
 
 class TransformError(Exception):

diff --git a/src/autolabel/transforms/serp_api.py b/src/autolabel/transforms/serp_api.py
@@ -92,9 +92,10 @@ async def _get_result(self, query):
         try:
             search_result = await self.serp_api_wrapper.arun(query=query)
         except Exception as e:
+            logger.error(f"Error while making request to Serp API: {str(e)}")
             raise TransformError(
-                TransformErrorType.SERP_API_ERROR,
-                f"Error while making request to Serp API: {e}",
+                TransformErrorType.TRANSFORM_API_ERROR,
+                f"Error while making request with query: {query}",
             )
         return search_result
 

diff --git a/src/autolabel/transforms/serper_api.py b/src/autolabel/transforms/serper_api.py
@@ -89,9 +89,10 @@ async def _get_result(self, query):
         try:
             search_result = await self.serper_api_wrapper.arun(query=query)
         except Exception as e:
+            logger.error(f"Error while making request to Serper API: {str(e)}")
             raise TransformError(
-                TransformErrorType.SERPER_API_ERROR,
-                f"Error while making request to Serper API: {e}",
+                TransformErrorType.TRANSFORM_API_ERROR,
+                f"Error while making request with query: {query}",
             )
         return search_result
 

diff --git a/tests/unit/transforms/test_serp_api.py b/tests/unit/transforms/test_serp_api.py
@@ -84,8 +84,14 @@ async def test_error_handling():
             "organic_search_results_error",
         ]
     )
-    assert "Invalid API key" in transformed_row["knowledge_graph_results"]
-    assert "Invalid API key" in transformed_row["knowledge_graph_results_error"]
+    assert (
+        "Error while making request with query"
+        in transformed_row["knowledge_graph_results"]
+    )
+    assert (
+        "Error while making request with query"
+        in transformed_row["knowledge_graph_results_error"]
+    )
 
 
 @pytest.mark.asyncio

diff --git a/tests/unit/transforms/test_webpage_scrape.py b/tests/unit/transforms/test_webpage_scrape.py
@@ -79,9 +79,9 @@ async def test_unreachable_url():
     )
     assert (
         transformed_row["webpage_content"]
-        == "TRANSFORM_TIMEOUT: Timeout when fetching content from URL"
+        == "ENRICHMENT_TIMEOUT: Timeout when fetching content from URL"
     )
     assert (
         transformed_row["webpage_content_error"]
-        == "TRANSFORM_TIMEOUT: Timeout when fetching content from URL"
+        == "ENRICHMENT_TIMEOUT: Timeout when fetching content from URL"
     )