From bdc75a22d5ec35c1c05d56d3e2a39ae2229268aa Mon Sep 17 00:00:00 2001 From: Isaac Francisco <78627776+isahers1@users.noreply.github.com> Date: Wed, 23 Oct 2024 17:58:05 -0700 Subject: [PATCH] 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 --- .../langgraph/langgraph/prebuilt/tool_node.py | 136 ++++- libs/langgraph/tests/test_prebuilt.py | 488 ++++++++++++++---- 2 files changed, 512 insertions(+), 112 deletions(-) diff --git a/libs/langgraph/langgraph/prebuilt/tool_node.py b/libs/langgraph/langgraph/prebuilt/tool_node.py index b7db3c1ce..2d2de56ff 100644 --- a/libs/langgraph/langgraph/prebuilt/tool_node.py +++ b/libs/langgraph/langgraph/prebuilt/tool_node.py @@ -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 diff --git a/libs/langgraph/tests/test_prebuilt.py b/libs/langgraph/tests/test_prebuilt.py index 2ebc2f6cb..f44650ede 100644 --- a/libs/langgraph/tests/test_prebuilt.py +++ b/libs/langgraph/tests/test_prebuilt.py @@ -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}"