From 78781b0a9004b0ff15ac784307fdf0312dcafe47 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Thu, 15 Aug 2024 15:23:54 -0700
Subject: [PATCH] Use exc_info

---
 libs/langgraph/langgraph/pregel/retry.py | 9 ++++-----
 libs/langgraph/tests/test_pregel.py      | 5 ++++-
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/libs/langgraph/langgraph/pregel/retry.py b/libs/langgraph/langgraph/pregel/retry.py
index ce5760bbf..88940b584 100644
--- a/libs/langgraph/langgraph/pregel/retry.py
+++ b/libs/langgraph/langgraph/pregel/retry.py
@@ -3,7 +3,6 @@ import logging
 import random
 import time
 from typing import Optional
-import traceback
 
 from langgraph.pregel.types import PregelExecutableTask, RetryPolicy
 
@@ -50,9 +49,8 @@ def run_with_retry(
             )
             # log the retry
             logger.info(
-                f"Retrying task {task.name} after {interval:.2f} seconds (attempt {attempts}) due to:\n"
-                f"{exc.__class__.__name__}: {exc}\n"
-                f"{traceback.format_exc()}"
+                f"Retrying task {task.name} after {interval:.2f} seconds (attempt {attempts}) after {exc.__class__.__name__} {exc}",
+                exc_info=exc,
             )
 
 
@@ -101,5 +99,6 @@ async def arun_with_retry(
             )
             # log the retry
             logger.info(
-                f"Retrying task {task.name} after {interval:.2f} seconds (attempt {attempts}) after {exc.__class__.__name__} {exc}"
+                f"Retrying task {task.name} after {interval:.2f} seconds (attempt {attempts}) after {exc.__class__.__name__} {exc}",
+                exc_info=exc,
             )
diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py
index d6a83e783..f09f94ca0 100644
--- a/libs/langgraph/tests/test_pregel.py
+++ b/libs/langgraph/tests/test_pregel.py
@@ -1,4 +1,5 @@
 import json
+import logging
 import operator
 import time
 import warnings
@@ -1191,7 +1192,8 @@ def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> Non
     assert app.invoke(2) == [3, 3]
 
 
-def test_invoke_checkpoint(mocker: MockerFixture) -> None:
+def test_invoke_checkpoint(mocker: MockerFixture, caplog) -> None:
+    caplog.set_level(logging.INFO)
     add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"])
     errored_once = False
 
@@ -1236,6 +1238,7 @@ def test_invoke_checkpoint(mocker: MockerFixture) -> None:
     assert checkpoint["channel_values"].get("total") == 2
     # total is now 2, so output is 2+3=5
     assert app.invoke(3, {"configurable": {"thread_id": "1"}}) == 5
+    assert 0
     assert errored_once, "errored and retried"
     checkpoint = memory.get({"configurable": {"thread_id": "1"}})
     assert checkpoint is not None