From b6cc30c6d4ee2b974d19042c745f03fac47a4a6b Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Thu, 13 Jun 2024 09:23:01 -0700
Subject: [PATCH] Adjust default error policy

- Add information about the last exception to retry logging
- Add a few more exceptions to the list not retried by default
---
 langgraph/pregel/retry.py  | 10 ++++++++--
 tests/test_pregel.py       |  2 +-
 tests/test_pregel_async.py |  2 +-
 3 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/langgraph/pregel/retry.py b/langgraph/pregel/retry.py
index 24d660e94..6ec224012 100644
--- a/langgraph/pregel/retry.py
+++ b/langgraph/pregel/retry.py
@@ -13,6 +13,8 @@ logger = logging.getLogger(__name__)
 
 
 def default_retry_on(exc: Exception) -> bool:
+    if isinstance(exc, ConnectionError):
+        return True
     if isinstance(
         exc,
         (
@@ -24,6 +26,10 @@ def default_retry_on(exc: Exception) -> bool:
             NameError,
             SyntaxError,
             RuntimeError,
+            ReferenceError,
+            StopIteration,
+            StopAsyncIteration,
+            OSError,
         ),
     ):
         return False
@@ -90,7 +96,7 @@ def run_with_retry(
             )
             # log the retry
             logger.info(
-                f"Retrying task {task.name} after {interval:.2f} seconds (attempt {attempts})"
+                f"Retrying task {task.name} after {interval:.2f} seconds (attempt {attempts}) after {exc.__class__.__name__} {exc}"
             )
 
 
@@ -138,5 +144,5 @@ async def arun_with_retry(
             )
             # log the retry
             logger.info(
-                f"Retrying task {task.name} after {interval:.2f} seconds (attempt {attempts})"
+                f"Retrying task {task.name} after {interval:.2f} seconds (attempt {attempts}) after {exc.__class__.__name__} {exc}"
             )
diff --git a/tests/test_pregel.py b/tests/test_pregel.py
index 712d260f5..640ba7376 100644
--- a/tests/test_pregel.py
+++ b/tests/test_pregel.py
@@ -842,7 +842,7 @@ def test_invoke_checkpoint(mocker: MockerFixture) -> None:
                 pass
             else:
                 errored_once = True
-                raise OSError("I will be retried")
+                raise ConnectionError("I will be retried")
         if input > 10:
             raise ValueError("Input is too large")
         return input
diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py
index 51a024fe3..e98aa9f90 100644
--- a/tests/test_pregel_async.py
+++ b/tests/test_pregel_async.py
@@ -758,7 +758,7 @@ async def test_invoke_checkpoint(mocker: MockerFixture) -> None:
                 pass
             else:
                 errored_once = True
-                raise OSError("I will be retried")
+                raise ConnectionError("I will be retried")
         if input > 10:
             raise ValueError("Input is too large")
         return input