fix: logger calls

AI21Labs · Dec 17, 2023 · 0e8c60e · 0e8c60e
1 parent 8a2966e
commit 0e8c60e
Show file tree

Hide file tree

Showing 5 changed files with 13 additions and 11 deletions.
diff --git a/ai21/__init__.py b/ai21/__init__.py
@@ -17,7 +17,7 @@
 from ai21.resources.responses.summarize_by_segment_response import SummarizeBySegmentResponse
 from ai21.resources.responses.summarize_response import SummarizeResponse
 from ai21.services.sagemaker import SageMaker
-from ai21.utils import logger
+from ai21.logger import logger
 from ai21.version import VERSION
 
 __version__ = VERSION

diff --git a/ai21/clients/bedrock/bedrock_session.py b/ai21/clients/bedrock/bedrock_session.py
@@ -5,9 +5,9 @@
 import boto3
 from botocore.exceptions import ClientError
 
+from ai21.logger import logger
 from ai21.errors import AccessDenied, NotFound, APITimeoutError
 from ai21.http_client import handle_non_success_response
-from ai21.utils import log_error
 
 _ERROR_MSG_TEMPLATE = (
     r"Received client error \((.*?)\) from primary with message \"(.*?)\". "
@@ -36,7 +36,7 @@ def invoke_model(self, model_id: str, input_json: str) -> Dict[str, Any]:
         except ClientError as client_error:
             self._handle_client_error(client_error)
         except Exception as exception:
-            log_error(f"Calling {model_id} failed with Exception: {exception}")
+            logger.error(f"Calling {model_id} failed with Exception: {exception}")
             raise exception
 
     def _handle_client_error(self, client_exception: ClientError) -> None:

diff --git a/ai21/clients/sagemaker/sagemaker_session.py b/ai21/clients/sagemaker/sagemaker_session.py
@@ -6,7 +6,7 @@
 
 from ai21.errors import BadRequest, ServiceUnavailable, ServerError, APIError
 from ai21.http_client import handle_non_success_response
-from ai21.utils import log_error
+from ai21.logger import logger
 
 _ERROR_MSG_TEMPLATE = (
     r"Received client error \((.*?)\) from primary with message \"(.*?)\". "
@@ -38,7 +38,7 @@ def invoke_endpoint(
         except ClientError as sm_client_error:
             self._handle_client_error(sm_client_error)
         except Exception as exception:
-            log_error(f"Calling {self._endpoint_name} failed with Exception: {exception}")
+            logger.error(f"Calling {self._endpoint_name} failed with Exception: {exception}")
             raise exception
 
     def _handle_client_error(self, client_exception: "ClientError"):

diff --git a/ai21/http_client.py b/ai21/http_client.py
@@ -4,6 +4,7 @@
 import requests
 from requests.adapters import HTTPAdapter, Retry, RetryError
 
+from ai21.logger import logger
 from ai21.errors import (
     BadRequest,
     Unauthorized,
@@ -13,7 +14,6 @@
     ServiceUnavailable,
     APIError,
 )
-from ai21.utils import log_info, log_error
 
 DEFAULT_TIMEOUT_SEC = 300
 DEFAULT_NUM_RETRIES = 0
@@ -77,7 +77,7 @@ def execute_http_request(
         timeout = self.timeout_sec
         headers = self.headers
         data = json.dumps(params).encode()
-        log_info(f"Calling {method} {url} {headers} {data}")
+        logger.info(f"Calling {method} {url} {headers} {data}")
         try:
             if method == "GET":
                 response = session.request(
@@ -108,17 +108,19 @@ def execute_http_request(
             else:
                 response = session.request(method, url, headers=headers, data=data, timeout=timeout, auth=auth)
         except ConnectionError as connection_error:
-            log_error(f"Calling {method} {url} failed with ConnectionError: {connection_error}")
+            logger.error(f"Calling {method} {url} failed with ConnectionError: {connection_error}")
             raise connection_error
         except RetryError as retry_error:
-            log_error(f"Calling {method} {url} failed with RetryError after {self.num_retries} attempts: {retry_error}")
+            logger.error(
+                f"Calling {method} {url} failed with RetryError after {self.num_retries} attempts: {retry_error}"
+            )
             raise retry_error
         except Exception as exception:
-            log_error(f"Calling {method} {url} failed with Exception: {exception}")
+            logger.error(f"Calling {method} {url} failed with Exception: {exception}")
             raise exception
 
         if response.status_code != 200:
-            log_error(f"Calling {method} {url} failed with a non-200 response code: {response.status_code}")
+            logger.error(f"Calling {method} {url} failed with a non-200 response code: {response.status_code}")
             handle_non_success_response(response.status_code, response.text)
 
         return response.json()
diff --git a/ai21/utils.py → ai21/logger.py b/ai21/utils.py → ai21/logger.py