From adf5cb0ff1666405e9fc51d144f64c1d9d979f37 Mon Sep 17 00:00:00 2001
From: Itay Etelis <itay.etelis@ibm.com>
Date: Wed, 18 Sep 2024 13:41:22 +0300
Subject: [PATCH] Fix `retry_on` handling in `run_with_retry`

- Correctly distinguish between exception classes, lists/tuples of exception classes, and callables.
- Add support for lists in `retry_on`, alongside tuples.
- Prevent exception classes from being incorrectly treated as callables.
- Raise a `TypeError` if `retry_on` is of an unsupported type.
---
 libs/langgraph/langgraph/pregel/retry.py | 28 +++++++++++++++++++-----
 1 file changed, 22 insertions(+), 6 deletions(-)

diff --git a/libs/langgraph/langgraph/pregel/retry.py b/libs/langgraph/langgraph/pregel/retry.py
index cdf7b33d3..34374962a 100644
--- a/libs/langgraph/langgraph/pregel/retry.py
+++ b/libs/langgraph/langgraph/pregel/retry.py
@@ -38,11 +38,19 @@ def run_with_retry(
             # increment attempts
             attempts += 1
             # check if we should retry
-            if callable(retry_policy.retry_on):
+            if isinstance(retry_policy.retry_on, (list, tuple)):
+                if not isinstance(exc, tuple(retry_policy.retry_on)):
+                    raise
+            elif isinstance(retry_policy.retry_on, type) and issubclass(retry_policy.retry_on, Exception):
+                if not isinstance(exc, retry_policy.retry_on):
+                    raise
+            elif callable(retry_policy.retry_on):
                 if not retry_policy.retry_on(exc):
                     raise
-            elif not isinstance(exc, retry_policy.retry_on):
-                raise
+            else:
+                raise TypeError(
+                    "retry_on must be an Exception class, a list or tuple of Exception classes, or a callable"
+                )
             # check if we should give up
             if attempts >= retry_policy.max_attempts:
                 raise
@@ -94,11 +102,19 @@ async def arun_with_retry(
             # increment attempts
             attempts += 1
             # check if we should retry
-            if callable(retry_policy.retry_on):
+            if isinstance(retry_policy.retry_on, (list, tuple)):
+                if not isinstance(exc, tuple(retry_policy.retry_on)):
+                    raise
+            elif isinstance(retry_policy.retry_on, type) and issubclass(retry_policy.retry_on, Exception):
+                if not isinstance(exc, retry_policy.retry_on):
+                    raise
+            elif callable(retry_policy.retry_on):
                 if not retry_policy.retry_on(exc):
                     raise
-            elif not isinstance(exc, retry_policy.retry_on):
-                raise
+            else:
+                raise TypeError(
+                    "retry_on must be an Exception class, a list or tuple of Exception classes, or a callable"
+                )
             # check if we should give up
             if attempts >= retry_policy.max_attempts:
                 raise