Merge pull request #113 from askui/fix/multiple-retry

fix(retry): update on_exception_types type hint and add support for m…

Author: danyalxahid-askui

src/askui/retry.py

@@ -60,7 +60,7 @@ class ConfigurableRetry(Retry):
     parameters for delay and retry count.
 
     Args:
-        on_exception_types (Tuple[Type[Exception]]): Tuple of exception types that should trigger a retry
+        on_exception_types (Tuple[Type[Exception],...]): Tuple of exception types that should trigger a retry
         strategy (Literal["Exponential", "Fixed", "Linear"]): The retry strategy to use:
             - `"Exponential"`: Delay increases exponentially between retries
             - `"Fixed"`: Constant delay between retries
@@ -83,7 +83,7 @@ class ConfigurableRetry(Retry):
     @validate_call(config=ConfigDict(arbitrary_types_allowed=True))
     def __init__(
         self,
-        on_exception_types: Tuple[Type[Exception]],
+        on_exception_types: Tuple[Type[Exception], ...],
         strategy: Literal["Exponential", "Fixed", "Linear"],
         base_delay: Annotated[int, Field(gt=0)] = 1000,
         retry_count: Annotated[int, Field(gt=0)] = 3,

tests/integration/agent/test_retry.py

@@ -1,11 +1,12 @@
 from typing import Optional, Tuple, Union
 
 import pytest
+from httpx import HTTPStatusError
 
 from askui import ConfigurableRetry, LocateModel, VisionAgent
 from askui.locators.locators import Locator
 from askui.models import ModelComposition
-from askui.models.exceptions import ElementNotFoundError
+from askui.models.exceptions import ElementNotFoundError, ModelNotFoundError
 from askui.tools.toolbox import AgentToolbox
 from askui.utils.image_utils import ImageSource
 
@@ -52,6 +53,26 @@ def vision_agent_with_retry(
     )
 
 
+@pytest.fixture
+def vision_agent_with_retry_on_multiple_exceptions(
+    failing_model: FailingLocateModel, agent_toolbox_mock: AgentToolbox
+) -> VisionAgent:
+    return VisionAgent(
+        models={"failing-locate": failing_model},
+        tools=agent_toolbox_mock,
+        retry=ConfigurableRetry(
+            on_exception_types=(
+                ElementNotFoundError,
+                HTTPStatusError,
+                ModelNotFoundError,
+            ),
+            strategy="Fixed",
+            retry_count=3,
+            base_delay=1,
+        ),
+    )
+
+
 @pytest.fixture
 def vision_agent_always_fail(
     always_failing_model: FailingLocateModel, agent_toolbox_mock: AgentToolbox
@@ -78,6 +99,17 @@ def test_locate_retries_and_succeeds(
     assert failing_model.calls == 3  # 2 fails + 1 success
 
 
+def test_locate_retries_on_multiple_exceptions_and_succeeds(
+    vision_agent_with_retry_on_multiple_exceptions: VisionAgent,
+    failing_model: FailingLocateModel,
+) -> None:
+    result = vision_agent_with_retry_on_multiple_exceptions.locate(
+        "something", screenshot=None, model="failing-locate"
+    )
+    assert result == (10, 10)
+    assert failing_model.calls == 3
+
+
 def test_locate_retries_and_fails(
     vision_agent_always_fail: VisionAgent, always_failing_model: FailingLocateModel
 ) -> None: