Update retry logic in BasePromptDriver

griptape/drivers/prompt/base_prompt_driver.py

@@ -3,6 +3,7 @@
 import time
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING
+from tenacity import Retrying, wait_exponential, after_log
 from attr import define, field
 from griptape.tokenizers import BaseTokenizer
 
@@ -12,23 +13,27 @@
 
 @define
 class BasePromptDriver(ABC):
-    max_retries: int = field(default=8, kw_only=True)
-    retry_delay: float = field(default=1, kw_only=True)
+    min_retry_delay: float = field(default=2, kw_only=True)
+    max_retry_delay: float = field(default=10, kw_only=True)
+
     temperature: float = field(default=0.1, kw_only=True)
     model: str
     tokenizer: BaseTokenizer
 
     def run(self, **kwargs) -> TextArtifact:
-        for attempt in range(0, self.max_retries + 1):
-            try:
+        for attempt in Retrying(
+            wait=wait_exponential(
+                min=self.min_retry_delay,
+                max=self.max_retry_delay
+            ),
+            reraise=True,
+            after=after_log(
+                logger=logging.getLogger(__name__),
+                log_level=logging.ERROR
+            ),
+        ):
+            with attempt:
                 return self.try_run(**kwargs)
-            except Exception as e:
-                logging.error(f"PromptDriver.run attempt {attempt} failed: {e}\nRetrying in {self.retry_delay} seconds")
-
-                if attempt < self.max_retries:
-                    time.sleep(self.retry_delay)
-                else:
-                    raise e
 
     @abstractmethod
     def try_run(self, **kwargs) -> TextArtifact: