From cd3a1cbf8f8590fa432229671917d4522afc7b0b Mon Sep 17 00:00:00 2001
From: Vadym Barda <vadym@langchain.dev>
Date: Wed, 3 Jul 2024 15:12:50 -0400
Subject: [PATCH] langgraph: add tool error handling (#910)

---
 .../langgraph/langgraph/prebuilt/tool_node.py | 20 ++++++--
 libs/langgraph/tests/test_prebuilt.py         | 48 +++++++++++++++++++
 2 files changed, 64 insertions(+), 4 deletions(-)

diff --git a/libs/langgraph/langgraph/prebuilt/tool_node.py b/libs/langgraph/langgraph/prebuilt/tool_node.py
index 0dddaf584..099930fd5 100644
--- a/libs/langgraph/langgraph/prebuilt/tool_node.py
+++ b/libs/langgraph/langgraph/prebuilt/tool_node.py
@@ -52,9 +52,11 @@ class ToolNode(RunnableCallable):
         *,
         name: str = "tools",
         tags: Optional[list[str]] = None,
+        handle_tool_errors: Optional[bool] = True,
     ) -> None:
         super().__init__(self._func, self._afunc, name=name, tags=tags, trace=False)
         self.tools_by_name: Dict[str, BaseTool] = {}
+        self.handle_tool_errors = handle_tool_errors
         for tool_ in tools:
             if not isinstance(tool_, BaseTool):
                 tool_ = create_tool(tool_)
@@ -76,7 +78,12 @@ class ToolNode(RunnableCallable):
             raise ValueError("Last message is not an AIMessage")
 
         def run_one(call: ToolCall):
-            output = self.tools_by_name[call["name"]].invoke(call["args"], config)
+            try:
+                output = self.tools_by_name[call["name"]].invoke(call["args"], config)
+            except Exception as e:
+                if not self.handle_tool_errors:
+                    raise e
+                output = f"Error: {repr(e)}\n Please fix your mistakes."
             return ToolMessage(
                 content=str_output(output), name=call["name"], tool_call_id=call["id"]
             )
@@ -104,9 +111,14 @@ class ToolNode(RunnableCallable):
             raise ValueError("Last message is not an AIMessage")
 
         async def run_one(call: ToolCall):
-            output = await self.tools_by_name[call["name"]].ainvoke(
-                call["args"], config
-            )
+            try:
+                output = await self.tools_by_name[call["name"]].ainvoke(
+                    call["args"], config
+                )
+            except Exception as e:
+                if not self.handle_tool_errors:
+                    raise e
+                output = f"Error: {repr(e)}\n Please fix your mistakes."
             return ToolMessage(
                 content=str_output(output), name=call["name"], tool_call_id=call["id"]
             )
diff --git a/libs/langgraph/tests/test_prebuilt.py b/libs/langgraph/tests/test_prebuilt.py
index dd2345bb8..634f6a136 100644
--- a/libs/langgraph/tests/test_prebuilt.py
+++ b/libs/langgraph/tests/test_prebuilt.py
@@ -109,10 +109,14 @@ def test_runnable_modifier():
 async def test_tool_node():
     def tool1(some_val: int, some_other_val: str) -> str:
         """Tool 1 docstring."""
+        if some_val == 0:
+            raise ValueError("Test error")
         return f"{some_val} - {some_other_val}"
 
     async def tool2(some_val: int, some_other_val: str) -> str:
         """Tool 2 docstring."""
+        if some_val == 0:
+            raise ValueError("Test error")
         return f"tool2: {some_val} - {some_other_val}"
 
     result = ToolNode([tool1]).invoke(
@@ -131,11 +135,37 @@ async def test_tool_node():
             ]
         }
     )
+
     tool_message: ToolMessage = result["messages"][-1]
     assert tool_message.type == "tool"
     assert tool_message.content == "1 - foo"
     assert tool_message.tool_call_id == "some 0"
 
+    result_error = ToolNode([tool1]).invoke(
+        {
+            "messages": [
+                AIMessage(
+                    "hi?",
+                    tool_calls=[
+                        {
+                            "name": "tool1",
+                            "args": {"some_val": 0, "some_other_val": "foo"},
+                            "id": "some 0",
+                        }
+                    ],
+                )
+            ]
+        }
+    )
+
+    tool_message: ToolMessage = result_error["messages"][-1]
+    assert tool_message.type == "tool"
+    assert (
+        tool_message.content
+        == f"Error: {repr(ValueError('Test error'))}\n Please fix your mistakes."
+    )
+    assert tool_message.tool_call_id == "some 0"
+
     result2 = await ToolNode([tool2]).ainvoke(
         {
             "messages": [
@@ -156,6 +186,24 @@ async def test_tool_node():
     assert tool_message.type == "tool"
     assert tool_message.content == "tool2: 2 - bar"
 
+    with pytest.raises(ValueError):
+        await ToolNode([tool2], handle_tool_errors=False).ainvoke(
+            {
+                "messages": [
+                    AIMessage(
+                        "hi?",
+                        tool_calls=[
+                            {
+                                "name": "tool2",
+                                "args": {"some_val": 0, "some_other_val": "bar"},
+                                "id": "some 1",
+                            }
+                        ],
+                    )
+                ]
+            }
+        )
+
 
 def my_function(some_val: int, some_other_val: str) -> str:
     return f"{some_val} - {some_other_val}"