UN-2793 [FEAT] Added exponential backoff retry mechanism for platform service connections

src/unstract/sdk/__init__.py

@@ -1,4 +1,5 @@
-__version__ = "v0.77.3"
+__version__ = "v0.78.0"
+
 
 def get_sdk_version() -> str:
     """Returns the SDK version."""

src/unstract/sdk/adapter.py

@@ -1,4 +1,5 @@
 import json
+import logging
 from typing import Any
 
 import requests
@@ -9,6 +10,9 @@
 from unstract.sdk.helper import SdkHelper
 from unstract.sdk.platform import PlatformBase
 from unstract.sdk.tool.base import BaseTool
+from unstract.sdk.utils.retry_utils import retry_platform_service_call
+
+logger = logging.getLogger(__name__)
 
 
 class ToolAdapter(PlatformBase):
@@ -24,10 +28,12 @@ def __init__(
         platform_host: str,
         platform_port: str,
     ) -> None:
-        """Args:
+        """Constructor for ToolAdapter.
+
+        Args:
             tool (AbstractTool): Instance of AbstractTool
             platform_host (str): Host of platform service
-            platform_port (str): Port of platform service
+            platform_port (str): Port of platform service.
 
         Notes:
             - PLATFORM_SERVICE_API_KEY environment variable is required.
@@ -38,14 +44,19 @@ def __init__(
             tool=tool, platform_host=platform_host, platform_port=platform_port
         )
 
+    @retry_platform_service_call
     def _get_adapter_configuration(
         self,
         adapter_instance_id: str,
     ) -> dict[str, Any]:
-        """Get Adapter
-            1. Get the adapter config from platform service
-            using the adapter_instance_id
+        """Get Adapter.
 
+               Get the adapter config from platform service
+               using the adapter_instance_id. This method automatically
+               retries on connection errors with exponential backoff.
+
+               Retry behavior is configurable via environment variables:
+        Check decorator for details
         Args:
             adapter_instance_id (str): Adapter instance ID
 
@@ -70,18 +81,14 @@ def _get_adapter_configuration(
                 f"'{adapter_type}', provider: '{provider}', name: '{adapter_name}'",
                 level=LogLevel.DEBUG,
             )
-        except ConnectionError:
-            raise SdkError(
-                "Unable to connect to platform service, please contact the admin."
-            )
         except HTTPError as e:
             default_err = (
                 "Error while calling the platform service, please contact the admin."
             )
             msg = AdapterUtils.get_msg_from_request_exc(
                 err=e, message_key="error", default_err=default_err
             )
-            raise SdkError(f"Error retrieving adapter. {msg}")
+            raise SdkError(f"Error retrieving adapter. {msg}") from e
         return adapter_data
 
     @staticmethod
@@ -121,4 +128,10 @@ def get_adapter_config(
             platform_host=platform_host,
             platform_port=platform_port,
         )
-        return tool_adapter._get_adapter_configuration(adapter_instance_id)
+
+        try:
+            return tool_adapter._get_adapter_configuration(adapter_instance_id)
+        except ConnectionError as e:
+            raise SdkError(
+                "Unable to connect to platform service, please contact the admin."
+            ) from e

src/unstract/sdk/platform.py

@@ -1,15 +1,20 @@
+import logging
 from typing import Any
 
 import requests
 from requests import ConnectionError, RequestException, Response
 from unstract.sdk.constants import (
+    LogLevel,
     MimeType,
     PromptStudioKeys,
     RequestHeader,
     ToolEnv,
 )
 from unstract.sdk.helper import SdkHelper
 from unstract.sdk.tool.base import BaseTool
+from unstract.sdk.utils.retry_utils import retry_platform_service_call
+
+logger = logging.getLogger(__name__)
 
 
 class PlatformBase:
@@ -86,6 +91,7 @@ def _get_headers(self, headers: dict[str, str] | None = None) -> dict[str, str]:
             request_headers.update(headers)
         return request_headers
 
+    @retry_platform_service_call
     def _call_service(
         self,
         url_path: str,
@@ -97,6 +103,10 @@ def _call_service(
         """Talks to platform-service to make GET / POST calls.
 
         Only GET calls are made to platform-service though functionality exists.
+        This method automatically retries on connection errors with exponential backoff.
+
+        Retry behavior is configurable via environment variables.
+        Check decorator for details
 
         Args:
             url_path (str): URL path to the service endpoint
@@ -130,9 +140,13 @@ def _call_service(
 
             response.raise_for_status()
         except ConnectionError as connect_err:
-            msg = "Unable to connect to platform service. Please contact admin."
-            msg += " \n" + str(connect_err)
-            self.tool.stream_error_and_exit(msg)
+            logger.exception("Connection error to platform service")
+            msg = (
+                "Unable to connect to platform service. Will retry with backoff, "
+                "please contact admin if retries ultimately fail."
+            )
+            self.tool.stream_log(msg, level=LogLevel.ERROR)
+            raise ConnectionError(msg) from connect_err
         except RequestException as e:
             # Extract error information from the response if available
             error_message = str(e)
@@ -200,4 +214,3 @@ def get_llm_profile(self, llm_profile_id: str) -> dict[str, Any]:
             headers=None,
             method="GET",
         )
-

src/unstract/sdk/prompt.py

@@ -15,6 +15,7 @@
 from unstract.sdk.platform import PlatformHelper
 from unstract.sdk.tool.base import BaseTool
 from unstract.sdk.utils.common_utils import log_elapsed
+from unstract.sdk.utils.retry_utils import retry_prompt_service_call
 
 logger = logging.getLogger(__name__)
 
@@ -185,6 +186,7 @@ def _get_headers(self, headers: dict[str, str] | None = None) -> dict[str, str]:
             request_headers.update(headers)
         return request_headers
 
+    @retry_prompt_service_call
     def _call_service(
         self,
         url_path: str,
@@ -196,6 +198,14 @@ def _call_service(
         """Communicates to prompt service to fetch response for the prompt.
 
         Only POST calls are made to prompt-service though functionality exists.
+        This method automatically retries on connection errors with exponential backoff.
+
+        Retry behavior is configurable via environment variables:
+          - PROMPT_SERVICE_MAX_RETRIES (default: 3)
+          - PROMPT_SERVICE_MAX_TIME (default: 60s)
+          - PROMPT_SERVICE_BASE_DELAY (default: 1.0s)
+          - PROMPT_SERVICE_MULTIPLIER (default: 2.0)
+          - PROMPT_SERVICE_JITTER (default: true)
 
         Args:
             url_path (str): URL path to the service endpoint