langgraph: expand handle_tool_errors in ToolNode (#1667)

This change expands error-handling functionality of the `ToolNode` by introducing more options for `handle_tool_errors`. Default behavior of the `ToolNode` is unchanged -- all errors are handled and wrapped in a `ToolMessage` to be sent back to LLM. With this change, users have flexibility to only handle the exceptions that they need to pass back to the LLM: * they can specify exceptions to handle by passing a tuple of exceptions in `handle_tool_errors` * specify `handle_tool_errors=True/str/callable` * when `handle_tool_errors` is a callable, the signature will be inspected and exceptions from the signature will be handled --------- Co-authored-by: vbarda <vadym@langchain.dev>
2026-09-05 17:27:47 +02:00 · 2024-10-24 00:58:05 +00:00
parent 62a5ec509d
commit bdc75a22d5
2 changed files with 512 additions and 112 deletions
@@ -1,6 +1,7 @@
 from __future__ import annotations

 import asyncio
+import inspect
 import json
 from copy import copy
 from typing import (
@@ -16,6 +17,7 @@ from typing import (
    Type,
    Union,
    cast,
+    get_type_hints,
 )

 from langchain_core.messages import (
@@ -67,13 +69,96 @@ def msg_content_output(output: Any) -> str | List[dict]:
            return str(output)


+def _handle_tool_error(
+    e: Exception,
+    *,
+    flag: Union[
+        bool,
+        str,
+        Callable[..., str],
+        tuple[type[Exception], ...],
+    ],
+) -> str:
+    if isinstance(flag, (bool, tuple)):
+        content = TOOL_CALL_ERROR_TEMPLATE.format(error=repr(e))
+    elif isinstance(flag, str):
+        content = flag
+    elif callable(flag):
+        content = flag(e)
+    else:
+        raise ValueError(
+            f"Got unexpected type of `handle_tool_error`. Expected bool, str "
+            f"or callable. Received: {flag}"
+        )
+    return content
+
+
+def _infer_handled_types(handler: Callable[..., str]) -> tuple[type[Exception]]:
+    sig = inspect.signature(handler)
+    params = list(sig.parameters.values())
+    if params:
+        # If it's a method, the first argument is typically 'self' or 'cls'
+        if params[0].name in ["self", "cls"] and len(params) == 2:
+            first_param = params[1]
+        else:
+            first_param = params[0]
+
+        type_hints = get_type_hints(handler)
+        if first_param.name in type_hints:
+            origin = get_origin(first_param.annotation)
+            if origin is Union:
+                args = get_args(first_param.annotation)
+                if all(issubclass(arg, Exception) for arg in args):
+                    return tuple(args)
+                else:
+                    raise ValueError(
+                        "All types in the error handler error annotation must be Exception types. "
+                        "For example, `def custom_handler(e: Union[ValueError, TypeError])`. "
+                        f"Got '{first_param.annotation}' instead."
+                    )
+
+            exception_type = type_hints[first_param.name]
+            if Exception in exception_type.__mro__:
+                return (exception_type,)
+            else:
+                raise ValueError(
+                    f"Arbitrary types are not supported in the error handler signature. "
+                    "Please annotate the error with either a specific Exception type or a union of Exception types. "
+                    "For example, `def custom_handler(e: ValueError)` or `def custom_handler(e: Union[ValueError, TypeError])`. "
+                    f"Got '{exception_type}' instead."
+                )
+
+    # If no type information is available, return (Exception,) for backwards compatibility.
+    return (Exception,)
+
+
 class ToolNode(RunnableCallable):
    """A node that runs the tools called in the last AIMessage.

-    It can be used either in StateGraph with a "messages" key (or a custom key passed via ToolNode's 'messages_key').
+    It can be used either in StateGraph with a "messages" state key (or a custom key passed via ToolNode's 'messages_key').
    If multiple tool calls are requested, they will be run in parallel. The output will be
    a list of ToolMessages, one for each tool call.

+    Args:
+        tools: A sequence of tools that can be invoked by the ToolNode.
+        name: The name of the ToolNode in the graph. Defaults to "tools".
+        tags: Optional tags to associate with the node. Defaults to None.
+        handle_tool_errors: How to handle tool errors raised by tools inside the node. Defaults to True.
+            Must be one of the following:
+
+            - True: all errors will be caught and
+                a ToolMessage with a default error message (TOOL_CALL_ERROR_TEMPLATE) will be returned.
+            - str: all errors will be caught and
+                a ToolMessage with the string value of 'handle_tool_errors' will be returned.
+            - tuple[type[Exception], ...]: exceptions in the tuple will be caught and
+                a ToolMessage with a default error message (TOOL_CALL_ERROR_TEMPLATE) will be returned.
+            - Callable[..., str]: exceptions from the signature of the callable will be caught and
+                a ToolMessage with the string value of the result of the 'handle_tool_errors' callable will be returned.
+            - False: none of the errors raised by the tools will be caught
+        messages_key: The state key in the input that contains the list of messages.
+            The same key will be used for the output from the ToolNode.
+            Defaults to "messages".
+
    The `ToolNode` is roughly analogous to:

    ```python
@@ -101,7 +186,9 @@ class ToolNode(RunnableCallable):
        *,
        name: str = "tools",
        tags: Optional[list[str]] = None,
-        handle_tool_errors: Optional[bool] = True,
+        handle_tool_errors: Union[
+            bool, str, Callable[..., str], tuple[type[Exception], ...]
+        ] = True,
        messages_key: str = "messages",
    ) -> None:
        super().__init__(self._func, self._afunc, name=name, tags=tags, trace=False)
@@ -181,14 +268,29 @@ class ToolNode(RunnableCallable):
            )
            return tool_message
        except Exception as e:
-            if not self.handle_tool_errors:
+            if isinstance(self.handle_tool_errors, tuple):
+                handled_types: tuple = self.handle_tool_errors
+            elif callable(self.handle_tool_errors):
+                handled_types = _infer_handled_types(self.handle_tool_errors)
+            else:
+                # default behavior is catching all exceptions
+                handled_types = (Exception,)
+
+            # Unhandled
+            if not self.handle_tool_errors or not isinstance(e, handled_types):
                raise e
-            content = TOOL_CALL_ERROR_TEMPLATE.format(error=repr(e))
-            return ToolMessage(content, name=call["name"], tool_call_id=call["id"])
+            # Handled
+            else:
+                content = _handle_tool_error(e, flag=self.handle_tool_errors)
+
+        return ToolMessage(
+            content=content, name=call["name"], tool_call_id=call["id"], status="error"
+        )

    async def _arun_one(self, call: ToolCall, config: RunnableConfig) -> ToolMessage:
        if invalid_tool_message := self._validate_tool_call(call):
            return invalid_tool_message
+
        try:
            input = {**call, **{"type": "tool_call"}}
            tool_message: ToolMessage = await self.tools_by_name[call["name"]].ainvoke(
@@ -199,10 +301,24 @@ class ToolNode(RunnableCallable):
            )
            return tool_message
        except Exception as e:
-            if not self.handle_tool_errors:
+            if isinstance(self.handle_tool_errors, tuple):
+                handled_types: tuple = self.handle_tool_errors
+            elif callable(self.handle_tool_errors):
+                handled_types = _infer_handled_types(self.handle_tool_errors)
+            else:
+                # default behavior is catching all exceptions
+                handled_types = (Exception,)
+
+            # Unhandled
+            if not self.handle_tool_errors or not isinstance(e, handled_types):
                raise e
-            content = TOOL_CALL_ERROR_TEMPLATE.format(error=repr(e))
-            return ToolMessage(content, name=call["name"], tool_call_id=call["id"])
+            # Handled
+            else:
+                content = _handle_tool_error(e, flag=self.handle_tool_errors)
+
+        return ToolMessage(
+            content=content, name=call["name"], tool_call_id=call["id"], status="error"
+        )

    def _parse_input(
        self,
@@ -240,7 +356,9 @@ class ToolNode(RunnableCallable):
                requested_tool=requested_tool,
                available_tools=", ".join(self.tools_by_name.keys()),
            )
-            return ToolMessage(content, name=requested_tool, tool_call_id=call["id"])
+            return ToolMessage(
+                content, name=requested_tool, tool_call_id=call["id"], status="error"
+            )
        else:
            return None

@@ -29,10 +29,11 @@ from langchain_core.messages import (
 )
 from langchain_core.outputs import ChatGeneration, ChatResult
 from langchain_core.runnables import Runnable, RunnableLambda
-from langchain_core.tools import BaseTool
+from langchain_core.tools import BaseTool, ToolException
 from langchain_core.tools import tool as dec_tool
-from pydantic import BaseModel
+from pydantic import BaseModel, ValidationError
 from pydantic.v1 import BaseModel as BaseModelV1
+from pydantic.v1 import ValidationError as ValidationErrorV1
 from typing_extensions import TypedDict

 from langgraph.checkpoint.base import BaseCheckpointSaver
@@ -43,7 +44,12 @@ from langgraph.prebuilt import (
    create_react_agent,
    tools_condition,
 )
-from langgraph.prebuilt.tool_node import InjectedState, InjectedStore
+from langgraph.prebuilt.tool_node import (
+    TOOL_CALL_ERROR_TEMPLATE,
+    InjectedState,
+    InjectedStore,
+    _infer_handled_types,
+)
 from langgraph.store.base import BaseStore
 from langgraph.store.memory import InMemoryStore
 from tests.conftest import (
@@ -368,32 +374,107 @@ def test_model_with_tools(tool_style: str):
        create_react_agent(model.bind_tools([tool1]), [tool2])


+def test__infer_handled_types() -> None:
+    def handle(e):  # type: ignore
+        return ""
+
+    def handle2(e: Exception) -> str:
+        return ""
+
+    def handle3(e: Union[ValueError, ToolException]) -> str:
+        return ""
+
+    class Handler:
+        def handle(self, e: ValueError) -> str:
+            return ""
+
+    handle4 = Handler().handle
+
+    def handle5(e: Union[Union[TypeError, ValueError], ToolException]):
+        return ""
+
+    expected: tuple = (Exception,)
+    actual = _infer_handled_types(handle)
+    assert expected == actual
+
+    expected = (Exception,)
+    actual = _infer_handled_types(handle2)
+    assert expected == actual
+
+    expected = (ValueError, ToolException)
+    actual = _infer_handled_types(handle3)
+    assert expected == actual
+
+    expected = (ValueError,)
+    actual = _infer_handled_types(handle4)
+    assert expected == actual
+
+    expected = (TypeError, ValueError, ToolException)
+    actual = _infer_handled_types(handle5)
+    assert expected == actual
+
+    with pytest.raises(ValueError):
+
+        def handler(e: str):
+            return ""
+
+        _infer_handled_types(handler)
+
+    with pytest.raises(ValueError):
+
+        def handler(e: list[Exception]):
+            return ""
+
+        _infer_handled_types(handler)
+
+    with pytest.raises(ValueError):
+
+        def handler(e: Union[str, int]):
+            return ""
+
+        _infer_handled_types(handler)
+
+
+# tools for testing Too
+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 ToolException("Test error")
+    return f"tool2: {some_val} - {some_other_val}"
+
+
+async def tool3(some_val: int, some_other_val: str) -> str:
+    """Tool 3 docstring."""
+    return [
+        {"key_1": some_val, "key_2": "foo"},
+        {"key_1": some_other_val, "key_2": "baz"},
+    ]
+
+
+async def tool4(some_val: int, some_other_val: str) -> str:
+    """Tool 4 docstring."""
+    return [
+        {"type": "image_url", "image_url": {"url": "abdc"}},
+    ]
+
+
+@dec_tool
+def tool5(some_val: int):
+    """Tool 5 docstring."""
+    raise ToolException("Test error")
+
+
+tool5.handle_tool_error = "foo"
+
+
 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}"
-
-    async def tool3(some_val: int, some_other_val: str) -> str:
-        """Tool 3 docstring."""
-        return [
-            {"key_1": some_val, "key_2": "foo"},
-            {"key_1": some_other_val, "key_2": "baz"},
-        ]
-
-    async def tool4(some_val: int, some_other_val: str) -> str:
-        """Tool 4 docstring."""
-        return [
-            {"type": "image_url", "image_url": {"url": "abdc"}},
-        ]
-
    result = ToolNode([tool1]).invoke(
        {
            "messages": [
@@ -416,31 +497,6 @@ async def test_tool_node():
    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": [
@@ -457,11 +513,232 @@ async def test_tool_node():
            ]
        }
    )
+
    tool_message: ToolMessage = result2["messages"][-1]
    assert tool_message.type == "tool"
    assert tool_message.content == "tool2: 2 - bar"

-    with pytest.raises(ValueError):
+    # list of dicts tool content
+    result3 = await ToolNode([tool3]).ainvoke(
+        {
+            "messages": [
+                AIMessage(
+                    "hi?",
+                    tool_calls=[
+                        {
+                            "name": "tool3",
+                            "args": {"some_val": 2, "some_other_val": "bar"},
+                            "id": "some 2",
+                        }
+                    ],
+                )
+            ]
+        }
+    )
+    tool_message: ToolMessage = result3["messages"][-1]
+    assert tool_message.type == "tool"
+    assert (
+        tool_message.content
+        == '[{"key_1": 2, "key_2": "foo"}, {"key_1": "bar", "key_2": "baz"}]'
+    )
+    assert tool_message.tool_call_id == "some 2"
+
+    # list of content blocks tool content
+    result4 = await ToolNode([tool4]).ainvoke(
+        {
+            "messages": [
+                AIMessage(
+                    "hi?",
+                    tool_calls=[
+                        {
+                            "name": "tool4",
+                            "args": {"some_val": 2, "some_other_val": "bar"},
+                            "id": "some 3",
+                        }
+                    ],
+                )
+            ]
+        }
+    )
+    tool_message: ToolMessage = result4["messages"][-1]
+    assert tool_message.type == "tool"
+    assert tool_message.content == [{"type": "image_url", "image_url": {"url": "abdc"}}]
+    assert tool_message.tool_call_id == "some 3"
+
+
+async def test_tool_node_error_handling():
+    def handle_all(e: Union[ValueError, ToolException, ValidationError]):
+        return TOOL_CALL_ERROR_TEMPLATE.format(error=repr(e))
+
+    # test catching all exceptions, via:
+    # - handle_tool_errors = True
+    # - passing a tuple of all exceptions
+    # - passing a callable with all exceptions in the signature
+    for handle_tool_errors in (
+        True,
+        (ValueError, ToolException, ValidationError),
+        handle_all,
+    ):
+        result_error = await ToolNode(
+            [tool1, tool2, tool3], handle_tool_errors=handle_tool_errors
+        ).ainvoke(
+            {
+                "messages": [
+                    AIMessage(
+                        "hi?",
+                        tool_calls=[
+                            {
+                                "name": "tool1",
+                                "args": {"some_val": 0, "some_other_val": "foo"},
+                                "id": "some id",
+                            },
+                            {
+                                "name": "tool2",
+                                "args": {"some_val": 0, "some_other_val": "bar"},
+                                "id": "some other id",
+                            },
+                            {
+                                "name": "tool3",
+                                "args": {"some_val": 0},
+                                "id": "another id",
+                            },
+                        ],
+                    )
+                ]
+            }
+        )
+
+        assert all(m.type == "tool" for m in result_error["messages"])
+        assert all(m.status == "error" for m in result_error["messages"])
+        assert (
+            result_error["messages"][0].content
+            == f"Error: {repr(ValueError('Test error'))}\n Please fix your mistakes."
+        )
+        assert (
+            result_error["messages"][1].content
+            == f"Error: {repr(ToolException('Test error'))}\n Please fix your mistakes."
+        )
+        assert (
+            "ValidationError" in result_error["messages"][2].content
+            or "validation error" in result_error["messages"][2].content
+        )
+
+        assert result_error["messages"][0].tool_call_id == "some id"
+        assert result_error["messages"][1].tool_call_id == "some other id"
+        assert result_error["messages"][2].tool_call_id == "another id"
+
+
+async def test_tool_node_error_handling_callable():
+    def handle_value_error(e: ValueError):
+        return "Value error"
+
+    def handle_tool_exception(e: ToolException):
+        return "Tool exception"
+
+    for handle_tool_errors in ("Value error", handle_value_error):
+        result_error = await ToolNode(
+            [tool1], handle_tool_errors=handle_tool_errors
+        ).ainvoke(
+            {
+                "messages": [
+                    AIMessage(
+                        "hi?",
+                        tool_calls=[
+                            {
+                                "name": "tool1",
+                                "args": {"some_val": 0, "some_other_val": "foo"},
+                                "id": "some id",
+                            },
+                        ],
+                    )
+                ]
+            }
+        )
+        tool_message: ToolMessage = result_error["messages"][-1]
+        assert tool_message.type == "tool"
+        assert tool_message.status == "error"
+        assert tool_message.content == "Value error"
+
+    # test raising for an unhandled exception, via:
+    # - passing a tuple of all exceptions
+    # - passing a callable with all exceptions in the signature
+    for handle_tool_errors in ((ValueError,), handle_value_error):
+        with pytest.raises(ToolException) as exc_info:
+            await ToolNode(
+                [tool1, tool2], handle_tool_errors=handle_tool_errors
+            ).ainvoke(
+                {
+                    "messages": [
+                        AIMessage(
+                            "hi?",
+                            tool_calls=[
+                                {
+                                    "name": "tool1",
+                                    "args": {"some_val": 0, "some_other_val": "foo"},
+                                    "id": "some id",
+                                },
+                                {
+                                    "name": "tool2",
+                                    "args": {"some_val": 0, "some_other_val": "bar"},
+                                    "id": "some other id",
+                                },
+                            ],
+                        )
+                    ]
+                }
+            )
+        assert str(exc_info.value) == "Test error"
+
+    for handle_tool_errors in ((ToolException,), handle_tool_exception):
+        with pytest.raises(ValueError) as exc_info:
+            await ToolNode(
+                [tool1, tool2], handle_tool_errors=handle_tool_errors
+            ).ainvoke(
+                {
+                    "messages": [
+                        AIMessage(
+                            "hi?",
+                            tool_calls=[
+                                {
+                                    "name": "tool1",
+                                    "args": {"some_val": 0, "some_other_val": "foo"},
+                                    "id": "some id",
+                                },
+                                {
+                                    "name": "tool2",
+                                    "args": {"some_val": 0, "some_other_val": "bar"},
+                                    "id": "some other id",
+                                },
+                            ],
+                        )
+                    ]
+                }
+            )
+        assert str(exc_info.value) == "Test error"
+
+
+async def test_tool_node_handle_tool_errors_false():
+    with pytest.raises(ValueError) as exc_info:
+        ToolNode([tool1], handle_tool_errors=False).invoke(
+            {
+                "messages": [
+                    AIMessage(
+                        "hi?",
+                        tool_calls=[
+                            {
+                                "name": "tool1",
+                                "args": {"some_val": 0, "some_other_val": "foo"},
+                                "id": "some id",
+                            }
+                        ],
+                    )
+                ]
+            }
+        )
+
+    assert str(exc_info.value) == "Test error"
+
+    with pytest.raises(ToolException):
        await ToolNode([tool2], handle_tool_errors=False).ainvoke(
            {
                "messages": [
@@ -471,7 +748,7 @@ async def test_tool_node():
                            {
                                "name": "tool2",
                                "args": {"some_val": 0, "some_other_val": "bar"},
-                                "id": "some 1",
+                                "id": "some id",
                            }
                        ],
                    )
@@ -479,7 +756,57 @@ async def test_tool_node():
            }
        )

-    # incorrect tool name
+    assert str(exc_info.value) == "Test error"
+
+    # test validation errors get raised if handle_tool_errors is False
+    with pytest.raises((ValidationError, ValidationErrorV1)):
+        ToolNode([tool1], handle_tool_errors=False).invoke(
+            {
+                "messages": [
+                    AIMessage(
+                        "hi?",
+                        tool_calls=[
+                            {
+                                "name": "tool1",
+                                "args": {"some_val": 0},
+                                "id": "some id",
+                            }
+                        ],
+                    )
+                ]
+            }
+        )
+
+
+def test_tool_node_individual_tool_error_handling():
+    # test error handling on individual tools (and that it overrides overall error handling!)
+    result_individual_tool_error_handler = ToolNode(
+        [tool5], handle_tool_errors="bar"
+    ).invoke(
+        {
+            "messages": [
+                AIMessage(
+                    "hi?",
+                    tool_calls=[
+                        {
+                            "name": "tool5",
+                            "args": {"some_val": 0},
+                            "id": "some 0",
+                        }
+                    ],
+                )
+            ]
+        }
+    )
+
+    tool_message: ToolMessage = result_individual_tool_error_handler["messages"][-1]
+    assert tool_message.type == "tool"
+    assert tool_message.status == "error"
+    assert tool_message.content == "foo"
+    assert tool_message.tool_call_id == "some 0"
+
+
+def test_tool_node_incorrect_tool_name():
    result_incorrect_name = ToolNode([tool1, tool2]).invoke(
        {
            "messages": [
@@ -496,61 +823,16 @@ async def test_tool_node():
            ]
        }
    )
+
    tool_message: ToolMessage = result_incorrect_name["messages"][-1]
    assert tool_message.type == "tool"
+    assert tool_message.status == "error"
    assert (
        tool_message.content
        == "Error: tool3 is not a valid tool, try one of [tool1, tool2]."
    )
    assert tool_message.tool_call_id == "some 0"

-    # list of dicts tool content
-    result3 = await ToolNode([tool3]).ainvoke(
-        {
-            "messages": [
-                AIMessage(
-                    "hi?",
-                    tool_calls=[
-                        {
-                            "name": "tool3",
-                            "args": {"some_val": 2, "some_other_val": "bar"},
-                            "id": "some 0",
-                        }
-                    ],
-                )
-            ]
-        }
-    )
-    tool_message: ToolMessage = result3["messages"][-1]
-    assert tool_message.type == "tool"
-    assert (
-        tool_message.content
-        == '[{"key_1": 2, "key_2": "foo"}, {"key_1": "bar", "key_2": "baz"}]'
-    )
-    assert tool_message.tool_call_id == "some 0"
-
-    # list of content blocks tool content
-    result4 = await ToolNode([tool4]).ainvoke(
-        {
-            "messages": [
-                AIMessage(
-                    "hi?",
-                    tool_calls=[
-                        {
-                            "name": "tool4",
-                            "args": {"some_val": 2, "some_other_val": "bar"},
-                            "id": "some 0",
-                        }
-                    ],
-                )
-            ]
-        }
-    )
-    tool_message: ToolMessage = result4["messages"][-1]
-    assert tool_message.type == "tool"
-    assert tool_message.content == [{"type": "image_url", "image_url": {"url": "abdc"}}]
-    assert tool_message.tool_call_id == "some 0"
-

 def my_function(some_val: int, some_other_val: str) -> str:
    return f"{some_val} - {some_other_val}"