[CM-8738] improve request exception handling

src/comet_llm/experiment_api/failed_response_handler.py

@@ -11,18 +11,26 @@
 #  This source code is licensed under the MIT license found in the
 #  LICENSE file in the root directory of this package.
 # *******************************************************
+
+import collections
 import json
-from typing import Optional
+from typing import NoReturn
 
 import requests  # type: ignore
 
-from .. import backend_error_codes, logging_messages
+from .. import backend_error_codes, exceptions, logging_messages
 
-SDK_ERROR_CODES_LOGGING_MESSAGE = {
-    backend_error_codes.UNABLE_TO_LOG_TO_NON_LLM_PROJECT: logging_messages.UNABLE_TO_LOG_TO_NON_LLM_PROJECT
-}
+_SDK_ERROR_CODES_LOGGING_MESSAGE = collections.defaultdict(
+    lambda: logging_messages.FAILED_TO_SEND_DATA_TO_SERVER,
+    {
+        backend_error_codes.UNABLE_TO_LOG_TO_NON_LLM_PROJECT: logging_messages.UNABLE_TO_LOG_TO_NON_LLM_PROJECT
+    },
+)
 
 
-def handle(response: requests.Response) -> Optional[str]:
+def handle(exception: requests.RequestException) -> NoReturn:
+    response = exception.response
     sdk_error_code = json.loads(response.text)["sdk_error_code"]
-    return SDK_ERROR_CODES_LOGGING_MESSAGE.get(sdk_error_code)
+    error_message = _SDK_ERROR_CODES_LOGGING_MESSAGE[sdk_error_code]
+
+    raise exceptions.CometLLMException(error_message) from exception

src/comet_llm/experiment_api/request_exception_wrapper.py

@@ -20,7 +20,7 @@
 
 import requests  # type: ignore
 
-from .. import config, exceptions
+from .. import config, exceptions, logging_messages
 from . import failed_response_handler
 
 LOGGER = logging.getLogger(__name__)
@@ -33,24 +33,23 @@ def wrapper(*args, **kwargs) -> Any:  # type: ignore
             try:
                 return func(*args, **kwargs)
             except requests.RequestException as exception:
-                exception_args: List[Any] = []
+                _debug_log(exception)
 
                 if check_on_prem:
                     comet_url = config.comet_url()
                     if _is_on_prem(comet_url):
-                        exception_args.append(
+                        raise exceptions.CometLLMException(
                             f"Failed to send prompt to your Comet installation at "
                             f"{comet_url}. Check that your Comet "
                             f"installation is up-to-date and check the traceback for more details."
-                        )
-                if exception.response is not None:
-                    exception_args.append(
-                        failed_response_handler.handle(exception.response)
-                    )
+                        ) from exception
 
-                _debug_log(exception)
+                if exception.response is None:
+                    raise exceptions.CometLLMException(
+                        logging_messages.FAILED_TO_SEND_DATA_TO_SERVER
+                    ) from exception
 
-                raise exceptions.CometLLMException(*exception_args) from exception
+                failed_response_handler.handle(exception)
 
         return wrapper
 
@@ -64,8 +63,13 @@ def _is_on_prem(url: str) -> bool:
 
 
 def _debug_log(exception: requests.RequestException) -> None:
-    if exception.request is not None:
-        LOGGER.debug(f"Request:\n{pformat(vars(exception.request))}")
-
-    if exception.response is not None:
-        LOGGER.debug(f"Response:\n{pformat(vars(exception.response))}")
+    try:
+        if exception.request is not None:
+            LOGGER.debug(f"Request:\n{pformat(vars(exception.request))}")
+
+        if exception.response is not None:
+            LOGGER.debug(f"Response:\n{pformat(vars(exception.response))}")
+    except Exception:
+        # Make sure we won't fail on attempt to debug.
+        # It's mainly for tests when response object can be mocked
+        pass

src/comet_llm/import_hooks/module_loader.py

@@ -53,6 +53,5 @@ def exec_module(self, module: ModuleType) -> None:
             patcher.patch(module, self._module_extension)
         except Exception:
             LOGGER.debug(
-                f"Failed to patch {self._module_name} with extension {self._module_extension.items()}",
-                exc_info=True,
+                f"Module {self._module_name} was not patched with an extension {self._module_extension.items()}",
             )

src/comet_llm/logging_messages.py

@@ -12,6 +12,7 @@
 #  LICENSE file in the root directory of this package.
 # *******************************************************
 
+FAILED_TO_SEND_DATA_TO_SERVER = "Failed to send data to server"
 
 UNABLE_TO_LOG_TO_NON_LLM_PROJECT = "Failed to send prompt to the specified project as it is not an LLM project, please specify a different project name."
 

tests/unit/experiment_api/test_failed_response_handler.py

@@ -0,0 +1,21 @@
+import json
+
+import box
+import pytest
+from testix import *
+
+from comet_llm.exceptions import exceptions
+from comet_llm.experiment_api import failed_response_handler
+
+
+def test_wrap__request_exception_non_llm_project_sdk_code__log_specifc_message_in_exception():
+    exception = Exception()
+    exception.response = box.Box(text=json.dumps({"sdk_error_code": 34323}))
+
+    expected_log_message = "Failed to send prompt to the specified project as it is not an LLM project, please specify a different project name."
+
+
+    with pytest.raises(exceptions.CometLLMException) as excinfo:
+        failed_response_handler.handle(exception)
+
+    assert excinfo.value.args == (expected_log_message, )

tests/unit/experiment_api/test_request_exception_wrapper.py

@@ -56,20 +56,13 @@ def f():
             f()
 
 
-def test_wrap__request_exception_non_llm_project_sdk_code__log_specifc_message_in_exception():
-    response = box.Box(text=json.dumps({"sdk_error_code": 34323}))
+def test_wrap__request_exception_with_not_None_response__exception_handled_by_failed_response_handler():
+    exception = requests.RequestException(response="not-None")
 
     @request_exception_wrapper.wrap()
     def f():
-        exception = requests.RequestException()
-        exception.response = response
         raise exception
 
-    expected_log_message = "Failed to send prompt to the specified project as it is not an LLM project, please specify a different project name."
-
     with Scenario() as s:
-        s.failed_response_handler.handle(response) >> expected_log_message
-        with pytest.raises(exceptions.CometLLMException) as excinfo:
-            f()
-
-    assert excinfo.value.args == (expected_log_message, )
+        s.failed_response_handler.handle(exception)
+        f()