feat: Improve retry and exception handling in the SDK clients (#171)

# Overview

The retry mechanism was incorrectly configured, resulting in capturing
and retrying on errors it shouldn't retry on, e.g.
`asyncio.CancelledError` or `asyncio.TimeoutError`. Instead of a
negative filter when using `tenacity`'s `retry` decorator, i.e.
`retry_if_not_exception_type((CompassClientError,))`, which captured a
lot more exceptions than we originally intended, we now specifically
list the exceptions that we know we want to retry on, or otherwise let
the error bubble.

This PR also improves the exception hierarchy by introducing a few new
exceptions to help us better communicate errors to the clients without
having to bubble or wrap httpx exceptions. Specifically, we are
introducing:

- CompassNetworkError
- CompassServerError
- CompassTimeoutError

As part of this PR, the `utils.py` module were refactored into multiple
files to avoid having a pile of utility functions dumped in a single
file.

Author: corafid

.gitignore

@@ -172,3 +172,4 @@ cython_debug/
 
 # PyPI configuration file
 .pypirc
+typings/

cohere_compass/clients/compass.py

@@ -24,7 +24,12 @@
 import httpx
 from joblib import Parallel, delayed  # type: ignore
 from pydantic import BaseModel
-from tenacity import retry, retry_if_not_exception_type, stop_after_attempt, wait_fixed
+from tenacity import (
+    retry,
+    retry_if_exception,
+    stop_after_attempt,
+    wait_fixed,
+)
 
 # Local imports
 from cohere_compass import GroupAuthorizationInput
@@ -37,11 +42,10 @@
     URL_SAFE_STRING_PATTERN,
 )
 from cohere_compass.exceptions import (
-    CompassAuthError,
-    CompassClientError,
     CompassError,
     CompassInsertionError,
     CompassMaxErrorRateExceeded,
+    handle_httpx_exceptions,
 )
 from cohere_compass.models import (
     CompassDocument,
@@ -82,7 +86,8 @@
     RetrievedDocument,
     SortBy,
 )
-from cohere_compass.utils import partition_documents
+from cohere_compass.utils.documents import partition_documents
+from cohere_compass.utils.retry import is_retryable_httpx_exception
 
 
 @dataclass
@@ -1474,42 +1479,16 @@ def _send_request(
         @retry(
             stop=stop_after_attempt(max_retries),
             wait=wait_fixed(retry_wait),
-            reraise=True,  # re-raise last exception instead of wrapping in RetryError
-            # todo find alternative to InvalidSchema
-            retry=retry_if_not_exception_type((CompassClientError,)),
+            reraise=True,
+            retry=retry_if_exception(is_retryable_httpx_exception),
         )
         def _send_request_with_retry() -> _SendRequestResult:
-            try:
-                return self._send_http_request(
-                    http_method=http_method,
-                    target_path=target_path,
-                    data=data,
-                    timeout=timeout,
-                )
-            except httpx.HTTPStatusError as e:
-                if e.response.status_code == 401:
-                    error = "Unauthorized. Please check your bearer token."
-                    raise CompassAuthError(message=str(e))
-                elif 400 <= e.response.status_code < 500:
-                    error = f"Client error occurred: {e.response.text}"
-                    raise CompassClientError(message=error, code=e.response.status_code)
-                else:
-                    error = str(e) + " " + e.response.text
-                    logger.warning(
-                        f"Failed to send request to {api_name} {target_path}: "
-                        f"{type(e)} {error}. Going to sleep for "
-                        f"{retry_wait} seconds and retrying."
-                    )
-                    raise e
-            except Exception as e:
-                error = str(e)
-                logger.warning(
-                    f"Failed to send request to {api_name} {target_path}: {type(e)} "
-                    f"{error}. Sleeping {retry_wait} seconds and retrying..."
-                )
-                raise e
+            return self._send_http_request(
+                http_method=http_method,
+                target_path=target_path,
+                data=data,
+                timeout=timeout,
+            )
 
-        try:
+        with handle_httpx_exceptions():
             return _send_request_with_retry()
-        except Exception as e:
-            raise CompassError(f"Failed to send request for {api_name} API") from e

cohere_compass/clients/compass_async.py

@@ -19,7 +19,12 @@
 # 3rd party imports
 import httpx
 from pydantic import BaseModel
-from tenacity import retry, retry_if_not_exception_type, stop_after_attempt, wait_fixed
+from tenacity import (
+    retry,
+    retry_if_exception,
+    stop_after_attempt,
+    wait_fixed,
+)
 
 # Local imports
 from cohere_compass import GroupAuthorizationInput
@@ -36,11 +41,10 @@
     DEFAULT_RETRY_WAIT,
 )
 from cohere_compass.exceptions import (
-    CompassAuthError,
-    CompassClientError,
     CompassError,
     CompassInsertionError,
     CompassMaxErrorRateExceeded,
+    handle_httpx_exceptions,
 )
 from cohere_compass.models import (
     CompassDocument,
@@ -77,11 +81,11 @@
 )
 from cohere_compass.models.indexes import IndexDetails, ListIndexesResponse
 from cohere_compass.models.search import GetDocumentResponse, RetrievedDocument, SortBy
-from cohere_compass.utils import (
-    async_apply,
-    async_enumerate,
+from cohere_compass.utils.asyn import async_apply, async_enumerate
+from cohere_compass.utils.documents import (
     partition_documents_async,
 )
+from cohere_compass.utils.retry import is_retryable_httpx_exception
 
 
 class CompassAsyncClient:
@@ -1349,41 +1353,15 @@ async def _send_request(
         @retry(
             stop=stop_after_attempt(max_retries),
             wait=wait_fixed(retry_wait),
-            reraise=True,  # re-raise last exception instead of wrapping in RetryError
-            # todo find alternative to InvalidSchema
-            retry=retry_if_not_exception_type((CompassClientError,)),
+            reraise=True,
+            retry=retry_if_exception(is_retryable_httpx_exception),
         )
         async def _send_request_with_retry() -> _SendRequestResult:
-            try:
-                return await self._send_http_request(
-                    http_method=http_method,
-                    target_path=target_path,
-                    data=data,
-                )
-            except httpx.HTTPStatusError as e:
-                if e.response.status_code == 401:
-                    error = "Unauthorized. Please check your bearer token."
-                    raise CompassAuthError(message=str(e))
-                elif 400 <= e.response.status_code < 500:
-                    error = f"Client error occurred: {e.response.text}"
-                    raise CompassClientError(message=error, code=e.response.status_code)
-                else:
-                    error = str(e) + " " + e.response.text
-                    logger.warning(
-                        f"Failed to send request to {api_name} {target_path}: "
-                        f"{type(e)} {error}. Going to sleep for "
-                        f"{self.retry_wait} seconds and retrying."
-                    )
-                    raise e
-            except Exception as e:
-                error = str(e)
-                logger.warning(
-                    f"Failed to send request to {api_name} {target_path}: {type(e)} "
-                    f"{error}. Sleeping {self.retry_wait} seconds and retrying..."
-                )
-                raise e
+            return await self._send_http_request(
+                http_method=http_method,
+                target_path=target_path,
+                data=data,
+            )
 
-        try:
+        with handle_httpx_exceptions():
             return await _send_request_with_retry()
-        except Exception as e:
-            raise CompassError(f"Failed to send request for {api_name} API") from e

cohere_compass/clients/parser.py

@@ -18,10 +18,9 @@
 import httpx
 
 # 3rd party imports
-from pydantic import ValidationError
 from tenacity import (
     retry,
-    retry_if_not_exception_type,
+    retry_if_exception,
     stop_after_attempt,
     wait_fixed,
 )
@@ -35,13 +34,17 @@
     DEFAULT_MAX_RETRIES,
     DEFAULT_RETRY_WAIT,
 )
-from cohere_compass.exceptions import CompassClientError, CompassError
+from cohere_compass.exceptions import handle_httpx_exceptions
 from cohere_compass.models import (
     CompassDocument,
     MetadataConfig,
     ParserConfig,
 )
-from cohere_compass.utils import imap_parallel, open_document, scan_folder
+from cohere_compass.utils.fs import open_document, scan_folder
+from cohere_compass.utils.iter import imap_parallel
+from cohere_compass.utils.retry import (
+    is_retryable_compass_exception,
+)
 
 Fn_or_Dict = dict[str, Any] | Callable[[CompassDocument], dict[str, Any]]
 
@@ -244,12 +247,6 @@ def _get_metadata(
         else:
             return custom_context
 
-    @retry(
-        stop=stop_after_attempt(DEFAULT_MAX_RETRIES),
-        wait=wait_fixed(DEFAULT_RETRY_WAIT),
-        retry=retry_if_not_exception_type((CompassClientError, ValidationError)),
-        reraise=True,
-    )
     def process_file(
         self,
         *,
@@ -304,12 +301,6 @@ def process_file(
             timeout=timeout,
         )
 
-    @retry(
-        stop=stop_after_attempt(DEFAULT_MAX_RETRIES),
-        wait=wait_fixed(DEFAULT_RETRY_WAIT),
-        retry=retry_if_not_exception_type((CompassClientError, ValidationError)),
-        reraise=True,
-    )
     def process_file_bytes(
         self,
         *,
@@ -379,6 +370,12 @@ def _get_file_params(
             content_type=content_type,
         )
 
+    @retry(
+        stop=stop_after_attempt(DEFAULT_MAX_RETRIES),
+        wait=wait_fixed(DEFAULT_RETRY_WAIT),
+        retry=retry_if_exception(is_retryable_compass_exception),
+        reraise=True,
+    )
     def _process_file_bytes(
         self,
         *,
@@ -392,23 +389,15 @@ def _process_file_bytes(
         if self.bearer_token:
             headers = {"Authorization": f"Bearer {self.bearer_token}"}
 
-        res = self.httpx_client.post(
-            url=f"{self.parser_url}/v1/process_file",
-            data={"data": json.dumps(params.model_dump())},
-            files={"file": (filename, file_bytes)},
-            headers=headers,
-            timeout=(timeout or self.timeout).total_seconds(),
-        )
-
-        if res.is_error:
-            if res.status_code >= 400 and res.status_code < 500:
-                raise CompassClientError(
-                    f"Error processing file: {res.status_code} {res.text}"
-                )
-            else:
-                raise CompassError(
-                    f"Error processing file: {res.status_code} {res.text}"
-                )
+        with handle_httpx_exceptions():
+            res = self.httpx_client.post(
+                url=f"{self.parser_url}/v1/process_file",
+                data={"data": json.dumps(params.model_dump())},
+                files={"file": (filename, file_bytes)},
+                headers=headers,
+                timeout=(timeout or self.timeout).total_seconds(),
+            )
+            res.raise_for_status()
 
         docs: list[CompassDocument] = []
         for doc in res.json()["docs"]:

cohere_compass/clients/parser_async.py

@@ -17,7 +17,7 @@
 import httpx
 from tenacity import (
     retry,
-    retry_if_not_exception_type,
+    retry_if_exception,
     stop_after_attempt,
     wait_fixed,
 )
@@ -31,13 +31,17 @@
     DEFAULT_MAX_RETRIES,
     DEFAULT_RETRY_WAIT,
 )
-from cohere_compass.exceptions import CompassClientError, CompassError
+from cohere_compass.exceptions import handle_httpx_exceptions
 from cohere_compass.models import (
     CompassDocument,
     MetadataConfig,
     ParserConfig,
 )
-from cohere_compass.utils import async_map, open_document, scan_folder
+from cohere_compass.utils.asyn import async_map
+from cohere_compass.utils.fs import open_document, scan_folder
+from cohere_compass.utils.retry import (
+    is_retryable_compass_exception,
+)
 
 Fn_or_Dict = dict[str, Any] | Callable[[CompassDocument], dict[str, Any]]
 
@@ -240,13 +244,6 @@ def _get_metadata(
         else:
             return custom_context
 
-    @retry(
-        stop=stop_after_attempt(DEFAULT_MAX_RETRIES),
-        wait=wait_fixed(DEFAULT_RETRY_WAIT),
-        # todo find alternative to InvalidSchema
-        retry=retry_if_not_exception_type((CompassClientError,)),
-        reraise=True,
-    )
     async def process_file(
         self,
         *,
@@ -297,13 +294,6 @@ async def process_file(
             custom_context=custom_context,
         )
 
-    @retry(
-        stop=stop_after_attempt(DEFAULT_MAX_RETRIES),
-        wait=wait_fixed(DEFAULT_RETRY_WAIT),
-        # todo find alternative to InvalidSchema
-        retry=retry_if_not_exception_type((CompassClientError,)),
-        reraise=True,
-    )
     async def process_file_bytes(
         self,
         *,
@@ -373,6 +363,12 @@ def _get_file_params(
             content_type=content_type,
         )
 
+    @retry(
+        stop=stop_after_attempt(DEFAULT_MAX_RETRIES),
+        wait=wait_fixed(DEFAULT_RETRY_WAIT),
+        retry=retry_if_exception(is_retryable_compass_exception),
+        reraise=True,
+    )
     async def _process_file_bytes(
         self,
         *,
@@ -386,23 +382,15 @@ async def _process_file_bytes(
         if self.bearer_token:
             headers = {"Authorization": f"Bearer {self.bearer_token}"}
 
-        res = await self.httpx.post(
-            url=f"{self.parser_url}/v1/process_file",
-            data={"data": json.dumps(params.model_dump())},
-            files={"file": (filename, file_bytes)},
-            headers=headers,
-            timeout=(timeout or self.timeout).total_seconds(),
-        )
-
-        if res.is_error:
-            if res.status_code >= 400 and res.status_code < 500:
-                raise CompassClientError(
-                    f"Error processing file: {res.status_code} {res.text}"
-                )
-            else:
-                raise CompassError(
-                    f"Error processing file: {res.status_code} {res.text}"
-                )
+        with handle_httpx_exceptions():
+            res = await self.httpx.post(
+                url=f"{self.parser_url}/v1/process_file",
+                data={"data": json.dumps(params.model_dump())},
+                files={"file": (filename, file_bytes)},
+                headers=headers,
+                timeout=(timeout or self.timeout).total_seconds(),
+            )
+            res.raise_for_status()
 
         docs: list[CompassDocument] = []
         for doc in res.json()["docs"]: