From a4e88e930f41c091011ffdc3c760cba43461a74d Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Thu, 16 May 2024 15:59:45 -0700
Subject: [PATCH] Disable retries by default

---
 langgraph/pregel/__init__.py |  2 +-
 langgraph/pregel/retry.py    | 10 +++++++---
 tests/test_pregel.py         |  2 ++
 tests/test_pregel_async.py   |  2 ++
 4 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py
index 145b67fef..cb88512c4 100644
--- a/langgraph/pregel/__init__.py
+++ b/langgraph/pregel/__init__.py
@@ -214,7 +214,7 @@ class Pregel(
     checkpointer: Optional[BaseCheckpointSaver] = None
     """Checkpointer used to save and load graph state. Defaults to None."""
 
-    retry_policy: Optional[RetryPolicy] = RetryPolicy()
+    retry_policy: Optional[RetryPolicy] = None
     """Retry policy to use when running tasks. Set to None to disable."""
 
     config_type: Optional[Type[Any]] = None
diff --git a/langgraph/pregel/retry.py b/langgraph/pregel/retry.py
index 3c8fc1baf..de9699781 100644
--- a/langgraph/pregel/retry.py
+++ b/langgraph/pregel/retry.py
@@ -2,7 +2,7 @@ import asyncio
 import logging
 import random
 import time
-from typing import Callable, NamedTuple, Union
+from typing import Callable, NamedTuple, Optional, Union
 
 import httpx
 import requests
@@ -53,7 +53,7 @@ class RetryPolicy(NamedTuple):
 
 def run_with_retry(
     task: PregelExecutableTask,
-    retry_policy: RetryPolicy,
+    retry_policy: Optional[RetryPolicy],
 ) -> None:
     """Run a task with retries."""
     interval = retry_policy.initial_interval
@@ -67,6 +67,8 @@ def run_with_retry(
             # if successful, end
             break
         except Exception as exc:
+            if retry_policy is None:
+                raise
             # increment attempts
             attempts += 1
             # check if we should retry
@@ -94,7 +96,7 @@ def run_with_retry(
 
 async def arun_with_retry(
     task: PregelExecutableTask,
-    retry_policy: RetryPolicy,
+    retry_policy: Optional[RetryPolicy],
     stream: bool = False,
 ) -> None:
     """Run a task asynchronously with retries."""
@@ -113,6 +115,8 @@ async def arun_with_retry(
             # if successful, end
             break
         except Exception as exc:
+            if retry_policy is None:
+                raise
             # increment attempts
             attempts += 1
             # check if we should retry
diff --git a/tests/test_pregel.py b/tests/test_pregel.py
index d4772b39e..0154490f7 100644
--- a/tests/test_pregel.py
+++ b/tests/test_pregel.py
@@ -38,6 +38,7 @@ from langgraph.prebuilt.chat_agent_executor import (
 )
 from langgraph.prebuilt.tool_node import ToolNode
 from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot
+from langgraph.pregel.retry import RetryPolicy
 from tests.any_str import AnyStr
 from tests.memory_assert import MemorySaverAssertImmutable
 
@@ -762,6 +763,7 @@ def test_invoke_checkpoint(mocker: MockerFixture) -> None:
         input_channels="input",
         output_channels="output",
         checkpointer=memory,
+        retry_policy=RetryPolicy(),
     )
 
     # total starts out as 0, so output is 0+2=2
diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py
index 6ddf962c2..94b209c8e 100644
--- a/tests/test_pregel_async.py
+++ b/tests/test_pregel_async.py
@@ -38,6 +38,7 @@ from langgraph.prebuilt.chat_agent_executor import (
 from langgraph.prebuilt.tool_executor import ToolExecutor
 from langgraph.prebuilt.tool_node import ToolNode
 from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot
+from langgraph.pregel.retry import RetryPolicy
 from tests.any_str import AnyStr
 from tests.memory_assert import MemorySaverAssertImmutable
 
@@ -743,6 +744,7 @@ async def test_invoke_checkpoint(mocker: MockerFixture) -> None:
         input_channels="input",
         output_channels="output",
         checkpointer=memory,
+        retry_policy=RetryPolicy(),
     )
 
     # total starts out as 0, so output is 0+2=2