feat: refactor retryers

Author: andreiancuta-uipath

src/uipath_langchain/_utils/_retry_after_strategy.py

@@ -2,13 +2,11 @@
 import os
 from typing import Optional
 
+from tenacity import AsyncRetrying, Retrying
 from uipath._utils import resource_override
 from uipath.utils import EndpointManager
 
-from .._utils._retry_after_strategy import (
-    AsyncRetryAfterHeaderStrategy,
-    RetryAfterHeaderStrategy,
-)
+from .retryers.bedrock import AsyncBedrockRetryer, BedrockRetryer
 from .supported_models import BedrockModels
 from .types import APIFlavor, LLMProvider
 
@@ -43,18 +41,11 @@ def _check_bedrock_dependencies() -> None:
 
 import boto3
 import botocore.config
-import botocore.exceptions
 from langchain_aws import (
     ChatBedrock,
     ChatBedrockConverse,
 )
 
-_BEDROCK_RETRY_EXCEPTIONS = (
-    botocore.exceptions.ReadTimeoutError,
-    botocore.exceptions.ConnectTimeoutError,
-    botocore.exceptions.EndpointConnectionError,
-)
-
 
 class AwsBedrockCompletionsPassthroughClient:
     @resource_override(
@@ -145,7 +136,8 @@ class UiPathChatBedrockConverse(ChatBedrockConverse):
     llm_provider: LLMProvider = LLMProvider.BEDROCK
     api_flavor: APIFlavor = APIFlavor.AWS_BEDROCK_CONVERSE
     model: str = ""  # For tracing serialization
-    max_retries: int = 5
+    retryer: Optional[Retrying] = None
+    aretryer: Optional[AsyncRetrying] = None
 
     def __init__(
         self,
@@ -155,6 +147,8 @@ def __init__(
         model_name: str = BedrockModels.anthropic_claude_haiku_4_5,
         agenthub_config: Optional[str] = None,
         byo_connection_id: Optional[str] = None,
+        retryer: Optional[Retrying] = None,
+        aretryer: Optional[AsyncRetrying] = None,
         **kwargs,
     ):
         org_id = org_id or os.getenv("UIPATH_ORGANIZATION_ID")
@@ -187,29 +181,24 @@ def __init__(
         kwargs["model"] = model_name
         super().__init__(**kwargs)
         self.model = model_name
+        self.retryer = retryer
+        self.aretryer = aretryer
 
     def invoke(self, *args, **kwargs):
-        retryer = RetryAfterHeaderStrategy(
-            retry_on_exceptions=_BEDROCK_RETRY_EXCEPTIONS,
-            max_retries=self.max_retries,
-            logger=logger,
-        )
+        retryer = self.retryer or _get_default_retryer()
         return retryer(super().invoke, *args, **kwargs)
 
     async def ainvoke(self, *args, **kwargs):
-        retryer = AsyncRetryAfterHeaderStrategy(
-            retry_on_exceptions=_BEDROCK_RETRY_EXCEPTIONS,
-            max_retries=self.max_retries,
-            logger=logger,
-        )
+        retryer = self.aretryer or _get_default_async_retryer()
         return await retryer(super().ainvoke, *args, **kwargs)
 
 
 class UiPathChatBedrock(ChatBedrock):
     llm_provider: LLMProvider = LLMProvider.BEDROCK
     api_flavor: APIFlavor = APIFlavor.AWS_BEDROCK_INVOKE
     model: str = ""  # For tracing serialization
-    max_retries: int = 5
+    retryer: Optional[Retrying] = None
+    aretryer: Optional[AsyncRetrying] = None
 
     def __init__(
         self,
@@ -219,6 +208,8 @@ def __init__(
         model_name: str = BedrockModels.anthropic_claude_haiku_4_5,
         agenthub_config: Optional[str] = None,
         byo_connection_id: Optional[str] = None,
+        retryer: Optional[Retrying] = None,
+        aretryer: Optional[AsyncRetrying] = None,
         **kwargs,
     ):
         org_id = org_id or os.getenv("UIPATH_ORGANIZATION_ID")
@@ -251,19 +242,21 @@ def __init__(
         kwargs["model"] = model_name
         super().__init__(**kwargs)
         self.model = model_name
+        self.retryer = retryer
+        self.aretryer = aretryer
 
     def invoke(self, *args, **kwargs):
-        retryer = RetryAfterHeaderStrategy(
-            retry_on_exceptions=_BEDROCK_RETRY_EXCEPTIONS,
-            max_retries=self.max_retries,
-            logger=logger,
-        )
+        retryer = self.retryer or _get_default_retryer()
         return retryer(super().invoke, *args, **kwargs)
 
     async def ainvoke(self, *args, **kwargs):
-        retryer = AsyncRetryAfterHeaderStrategy(
-            retry_on_exceptions=_BEDROCK_RETRY_EXCEPTIONS,
-            max_retries=self.max_retries,
-            logger=logger,
-        )
+        retryer = self.aretryer or _get_default_async_retryer()
         return await retryer(super().ainvoke, *args, **kwargs)
+
+
+def _get_default_retryer() -> BedrockRetryer:
+    return BedrockRetryer(logger=logger)
+
+
+def _get_default_async_retryer() -> AsyncBedrockRetryer:
+    return AsyncBedrockRetryer(logger=logger)