Support pydantic & dataclasses in ToolNode (#1625)

Co-authored-by: Abinash Panda <abinashpanda205@gmail.com>

libs/langgraph/langgraph/prebuilt/tool_node.py

@@ -1,7 +1,10 @@
+from __future__ import annotations
+
 import asyncio
 import json
 from copy import copy
 from typing import (
+    TYPE_CHECKING,
     Any,
     Callable,
     Dict,
@@ -14,15 +17,26 @@ from typing import (
     cast,
 )
 
-from langchain_core.messages import AIMessage, AnyMessage, ToolCall, ToolMessage
+from langchain_core.messages import (
+    AIMessage,
+    AnyMessage,
+    ToolCall,
+    ToolMessage,
+)
 from langchain_core.runnables import RunnableConfig
-from langchain_core.runnables.config import get_config_list, get_executor_for_config
+from langchain_core.runnables.config import (
+    get_config_list,
+    get_executor_for_config,
+)
 from langchain_core.tools import BaseTool, InjectedToolArg
 from langchain_core.tools import tool as create_tool
-from typing_extensions import get_args
+from typing_extensions import Annotated, get_args, get_origin
 
 from langgraph.utils.runnable import RunnableCallable
 
+if TYPE_CHECKING:
+    from pydantic import BaseModel
+
 INVALID_TOOL_NAME_ERROR_TEMPLATE = (
     "Error: {requested_tool} is not a valid tool, try one of [{available_tools}]."
 )
@@ -82,21 +96,35 @@ class ToolNode(RunnableCallable):
             self.tools_by_name[tool_.name] = tool_
 
     def _func(
-        self, input: Union[list[AnyMessage], dict[str, Any]], config: RunnableConfig
+        self,
+        input: Union[
+            list[AnyMessage],
+            dict[str, Any],
+            BaseModel,
+        ],
+        config: RunnableConfig,
     ) -> Any:
         tool_calls, output_type = self._parse_input(input)
         config_list = get_config_list(config, len(tool_calls))
         with get_executor_for_config(config) as executor:
             outputs = [*executor.map(self._run_one, tool_calls, config_list)]
+        # TypedDict, pydantic, dataclass, etc. should all be able to load from dict
         return outputs if output_type == "list" else {"messages": outputs}
 
     async def _afunc(
-        self, input: Union[list[AnyMessage], dict[str, Any]], config: RunnableConfig
+        self,
+        input: Union[
+            list[AnyMessage],
+            dict[str, Any],
+            BaseModel,
+        ],
+        config: RunnableConfig,
     ) -> Any:
         tool_calls, output_type = self._parse_input(input)
         outputs = await asyncio.gather(
             *(self._arun_one(call, config) for call in tool_calls)
         )
+        # TypedDict, pydantic, dataclass, etc. should all be able to load from dict
         return outputs if output_type == "list" else {"messages": outputs}
 
     def _run_one(self, call: ToolCall, config: RunnableConfig) -> ToolMessage:
@@ -135,12 +163,21 @@ class ToolNode(RunnableCallable):
             return ToolMessage(content, name=call["name"], tool_call_id=call["id"])
 
     def _parse_input(
-        self, input: Union[list[AnyMessage], dict[str, Any]]
+        self,
+        input: Union[
+            list[AnyMessage],
+            dict[str, Any],
+            BaseModel,
+        ],
     ) -> Tuple[List[ToolCall], Literal["list", "dict"]]:
         if isinstance(input, list):
             output_type = "list"
             message: AnyMessage = input[-1]
-        elif messages := input.get("messages", []):
+        elif isinstance(input, dict) and (messages := input.get("messages", [])):
+            output_type = "dict"
+            message = messages[-1]
+        elif messages := getattr(input, "messages", None):
+            # Assume dataclass-like state that can coerce from dict
             output_type = "dict"
             message = messages[-1]
         else:
@@ -166,12 +203,18 @@ class ToolNode(RunnableCallable):
             return None
 
     def _inject_state(
-        self, tool_call: ToolCall, input: Union[list[AnyMessage], dict[str, Any]]
+        self,
+        tool_call: ToolCall,
+        input: Union[
+            list[AnyMessage],
+            dict[str, Any],
+            BaseModel,
+        ],
     ) -> ToolCall:
         if tool_call["name"] not in self.tools_by_name:
             return tool_call
         state_args = _get_state_args(self.tools_by_name[tool_call["name"]])
-        if state_args and not isinstance(input, dict):
+        if state_args and isinstance(input, list):
             required_fields = list(state_args.values())
             if (
                 len(required_fields) == 1
@@ -188,26 +231,35 @@ class ToolNode(RunnableCallable):
                     required_fields_str = ", ".join(f for f in required_fields if f)
                     err_msg += f" State should contain fields {required_fields_str}."
                 raise ValueError(err_msg)
+        if isinstance(input, dict):
+            tool_state_args = {
+                tool_arg: input[state_field] if state_field else input
+                for tool_arg, state_field in state_args.items()
+            }
+
+        else:
+            tool_state_args = {
+                tool_arg: getattr(input, state_field) if state_field else input
+                for tool_arg, state_field in state_args.items()
+            }
+
         tool_call_copy: ToolCall = copy(tool_call)
         tool_call_copy["args"] = {
             **tool_call_copy["args"],
-            **{
-                tool_arg: cast(dict, input)[state_field] if state_field else input
-                for tool_arg, state_field in state_args.items()
-            },
+            **tool_state_args,
         }
         return tool_call_copy
 
 
 def tools_condition(
-    state: Union[list[AnyMessage], dict[str, Any]],
+    state: Union[list[AnyMessage], dict[str, Any], BaseModel],
 ) -> Literal["tools", "__end__"]:
     """Use in the conditional_edge to route to the ToolNode if the last message
 
     has tool calls. Otherwise, route to the end.
 
     Args:
-        state (Union[list[AnyMessage], dict[str, Any]]): The state to check for
+        state (Union[list[AnyMessage], dict[str, Any], BaseModel]): The state to check for
             tool calls. Must have a list of messages (MessageGraph) or have the
             "messages" key (StateGraph).
 
@@ -253,7 +305,9 @@ def tools_condition(
     """
     if isinstance(state, list):
         ai_message = state[-1]
-    elif messages := state.get("messages", []):
+    elif isinstance(state, dict) and (messages := state.get("messages", [])):
+        ai_message = messages[-1]
+    elif messages := getattr(state, "messages", []):
         ai_message = messages[-1]
     else:
         raise ValueError(f"No messages found in input state to tool_edge: {state}")
@@ -328,12 +382,20 @@ class InjectedState(InjectedToolArg):
 def _get_state_args(tool: BaseTool) -> Dict[str, Optional[str]]:
     full_schema = tool.get_input_schema()
     tool_args_to_state_fields: Dict = {}
+
+    def _is_injection(type_arg: Any):
+        if isinstance(type_arg, InjectedState) or (
+            isinstance(type_arg, type) and issubclass(type_arg, InjectedState)
+        ):
+            return True
+        origin_ = get_origin(type_arg)
+        if origin_ is Union or origin_ is Annotated:
+            return any(_is_injection(ta) for ta in get_args(type_arg))
+        return False
+
     for name, type_ in full_schema.__annotations__.items():
         injections = [
-            type_arg
-            for type_arg in get_args(type_)
-            if isinstance(type_arg, InjectedState)
-            or (isinstance(type_arg, type) and issubclass(type_arg, InjectedState))
+            type_arg for type_arg in get_args(type_) if _is_injection(type_arg)
         ]
         if len(injections) > 1:
             raise ValueError(

libs/langgraph/tests/test_prebuilt.py

@@ -1,5 +1,17 @@
+import dataclasses
 import json
-from typing import Annotated, Any, Callable, Dict, List, Optional, Sequence, Type, Union
+from typing import (
+    Annotated,
+    Any,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    Sequence,
+    Type,
+    TypeVar,
+    Union,
+)
 
 import pytest
 from langchain_core.callbacks import CallbackManagerForLLMRun
@@ -19,6 +31,7 @@ from langchain_core.tools import BaseTool
 from langchain_core.tools import tool as dec_tool
 from pydantic import BaseModel
 from pydantic.v1 import BaseModel as BaseModelV1
+from typing_extensions import TypedDict
 
 from langgraph.checkpoint.base import BaseCheckpointSaver
 from langgraph.prebuilt import ToolNode, ValidationNode, create_react_agent
@@ -440,14 +453,53 @@ async def test_validation_node(tool_schema: Any, use_message_key: bool):
     check_results(result_sync)
 
 
-def test_tool_node_inject_state() -> None:
-    def tool1(some_val: int, state: Annotated[dict, InjectedState]) -> str:
-        """Tool 1 docstring."""
-        return state["foo"]
+class _InjectStateSchema(TypedDict):
+    messages: list
+    foo: str
 
-    def tool2(some_val: int, state: Annotated[dict, InjectedState()]) -> str:
+
+class _InjectedStatePydanticSchema(BaseModelV1):
+    messages: list
+    foo: str
+
+
+class _InjectedStatePydanticV2Schema(BaseModel):
+    messages: list
+    foo: str
+
+
+@dataclasses.dataclass
+class _InjectedStateDataclassSchema:
+    messages: list
+    foo: str
+
+
+T = TypeVar("T")
+
+
+@pytest.mark.parametrize(
+    "schema_",
+    [
+        _InjectStateSchema,
+        _InjectedStatePydanticSchema,
+        _InjectedStatePydanticV2Schema,
+        _InjectedStateDataclassSchema,
+    ],
+)
+def test_tool_node_inject_state(schema_: Type[T]) -> None:
+    def tool1(some_val: int, state: Annotated[T, InjectedState]) -> str:
         """Tool 1 docstring."""
-        return state["foo"]
+        if isinstance(state, dict):
+            return state["foo"]
+        else:
+            return getattr(state, "foo")
+
+    def tool2(some_val: int, state: Annotated[T, InjectedState()]) -> str:
+        """Tool 2 docstring."""
+        if isinstance(state, dict):
+            return state["foo"]
+        else:
+            return getattr(state, "foo")
 
     def tool3(
         some_val: int,
@@ -472,23 +524,36 @@ def test_tool_node_inject_state() -> None:
             "type": "tool_call",
         }
         msg = AIMessage("hi?", tool_calls=[tool_call])
-        result = node.invoke({"messages": [msg], "foo": "bar"})
+        result = node.invoke(schema_(**{"messages": [msg], "foo": "bar"}))
         tool_message = result["messages"][-1]
-        assert tool_message.content == "bar"
+        assert tool_message.content == "bar", f"Failed for tool={tool_name}"
 
         if tool_name == "tool3":
-            with pytest.raises(KeyError):
-                node.invoke({"messages": [msg], "notfoo": "bar"})
+            failure_input = None
+            try:
+                failure_input = schema_(**{"messages": [msg], "notfoo": "bar"})
+            except Exception:
+                pass
+            if failure_input is not None:
+                with pytest.raises(KeyError):
+                    node.invoke(failure_input)
 
-            with pytest.raises(ValueError):
-                node.invoke([msg])
+                with pytest.raises(ValueError):
+                    node.invoke([msg])
         else:
-            tool_message = node.invoke({"messages": [msg], "notfoo": "bar"})[
-                "messages"
-            ][-1]
-            assert "KeyError" in tool_message.content
-            tool_message = node.invoke([msg])[-1]
-            assert "KeyError" in tool_message.content
+            failure_input = None
+            try:
+                failure_input = schema_(**{"messages": [msg], "notfoo": "bar"})
+            except Exception:
+                # We'd get a validation error from pydantic state and wouldn't make it to the node
+                # anyway
+                pass
+            if failure_input is not None:
+                messages_ = node.invoke(failure_input)
+                tool_message = messages_["messages"][-1]
+                assert "KeyError" in tool_message.content
+                tool_message = node.invoke([msg])[-1]
+                assert "KeyError" in tool_message.content
 
     tool_call = {
         "name": "tool4",
@@ -497,7 +562,7 @@ def test_tool_node_inject_state() -> None:
         "type": "tool_call",
     }
     msg = AIMessage("hi?", tool_calls=[tool_call])
-    result = node.invoke({"messages": [msg]})
+    result = node.invoke(schema_(**{"messages": [msg], "foo": ""}))
     tool_message = result["messages"][-1]
     assert tool_message.content == "hi?"