From 1c22d5cbdded396777a4bad4a12d14881e6130f0 Mon Sep 17 00:00:00 2001
From: Rafid <rafid@cohere.com>
Date: Mon, 2 Dec 2024 12:02:43 -0800
Subject: [PATCH] [SEA-1683] Fix retry mechanism for 400 errors (#53)

We were unnecessarily retrying on all 400 errors except 401. This change
should fix the issue.
---
 compass_sdk/compass.py    | 33 ++++++++++-----------------------
 compass_sdk/exceptions.py | 29 +++++++++++++++++++++++++++++
 pyproject.toml            |  2 +-
 3 files changed, 40 insertions(+), 24 deletions(-)
 create mode 100644 compass_sdk/exceptions.py

diff --git a/compass_sdk/compass.py b/compass_sdk/compass.py
index fae4b8a..f7f5853 100644
--- a/compass_sdk/compass.py
+++ b/compass_sdk/compass.py
@@ -35,6 +35,7 @@
     DEFAULT_MAX_RETRIES,
     DEFAULT_SLEEP_RETRY_SECONDS,
 )
+from compass_sdk.exceptions import CompassAuthError, CompassClientError, CompassMaxErrorRateExceeded
 from compass_sdk.models import (
     CreateDataSource,
     DataSource,
@@ -50,28 +51,6 @@ class RetryResult:
     error: Optional[str] = None
 
 
-class CompassAuthError(Exception):
-    """Exception raised for authentication errors in the Compass client."""
-
-    def __init__(
-        self,
-        message=("CompassAuthError - check your bearer token or username and password."),
-    ):
-        self.message = message
-        super().__init__(self.message)
-
-
-class CompassMaxErrorRateExceeded(Exception):
-    """Exception raised when the error rate exceeds the maximum allowed error rate in the Compass client."""
-
-    def __init__(
-        self,
-        message="The maximum error rate was exceeded. Stopping the insertion process.",
-    ):
-        self.message = message
-        super().__init__(self.message)
-
-
 _DEFAULT_TIMEOUT = 30
 
 
@@ -636,7 +615,12 @@ def _send_request(
         @retry(
             stop=stop_after_attempt(max_retries),
             wait=wait_fixed(sleep_retry_seconds),
-            retry=retry_if_not_exception_type((CompassAuthError, InvalidSchema)),
+            retry=retry_if_not_exception_type(
+                (
+                    CompassClientError,
+                    InvalidSchema,
+                )
+            ),
         )
         def _send_request_with_retry():
             nonlocal error
@@ -667,6 +651,9 @@ def _send_request_with_retry():
                 if e.response.status_code == 401:
                     error = "Unauthorized. Please check your username and password."
                     raise CompassAuthError(message=str(e))
+                elif 400 <= e.response.status_code < 500:
+                    error = f"Client error occurred: {e.response.text}"
+                    raise CompassClientError(message=error)
                 else:
                     error = str(e) + " " + e.response.text
                     logger.error(
diff --git a/compass_sdk/exceptions.py b/compass_sdk/exceptions.py
new file mode 100644
index 0000000..01af1d6
--- /dev/null
+++ b/compass_sdk/exceptions.py
@@ -0,0 +1,29 @@
+class CompassClientError(Exception):
+    """Exception raised for all 4xx client errors in the Compass client."""
+
+    def __init__(self, message="Client error occurred."):
+        self.message = message
+        super().__init__(self.message)
+
+
+class CompassAuthError(CompassClientError):
+    """Exception raised for authentication errors in the Compass client."""
+
+    def __init__(
+        self,
+        message=("CompassAuthError - check your bearer token or username and password."),
+    ):
+        self.message = message
+        super().__init__(self.message)
+
+
+class CompassMaxErrorRateExceeded(Exception):
+    """Exception raised when the error rate exceeds the maximum allowed error rate in
+    the Compass client."""
+
+    def __init__(
+        self,
+        message="The maximum error rate was exceeded. Stopping the insertion process.",
+    ):
+        self.message = message
+        super().__init__(self.message)
diff --git a/pyproject.toml b/pyproject.toml
index cae1d67..15964da 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "compass-sdk"
-version = "0.5.0"
+version = "0.5.1"
 authors = []
 description = "Compass SDK"