x

libs/prebuilt/langgraph/prebuilt/tool_node.py

@@ -340,14 +340,22 @@ class ToolNode(RunnableCallable):
         self.tools_by_name: dict[str, BaseTool] = {}
         self.tool_to_state_args: dict[str, dict[str, Optional[str]]] = {}
         self.tool_to_store_arg: dict[str, Optional[str]] = {}
+        self.structured_output_tools: list[str] = []
         self.handle_tool_errors = handle_tool_errors
         self.messages_key = messages_key
         for tool_ in tools:
-            if not isinstance(tool_, BaseTool):
-                tool_ = create_tool(tool_)
-            self.tools_by_name[tool_.name] = tool_
-            self.tool_to_state_args[tool_.name] = _get_state_args(tool_)
-            self.tool_to_store_arg[tool_.name] = _get_store_arg(tool_)
+            if inspect.isclass(tool_) and issubclass(tool_, BaseModel):
+                # Handle Pydantic model classes as structured output tools
+                self.tools_by_name[tool_.__name__] = tool_
+                self.tool_to_state_args[tool_.__name__] = {}
+                self.tool_to_store_arg[tool_.__name__] = None
+                self.structured_output_tools.append(tool_.__name__)
+            else:
+                if not isinstance(tool_, BaseTool):
+                    tool_ = create_tool(tool_)
+                self.tools_by_name[tool_.name] = tool_
+                self.tool_to_state_args[tool_.name] = _get_state_args(tool_)
+                self.tool_to_store_arg[tool_.name] = _get_store_arg(tool_)
 
     def _func(
         self,
@@ -390,7 +398,7 @@ class ToolNode(RunnableCallable):
 
     def _combine_tool_outputs(
         self,
-        outputs: list[ToolMessage],
+        outputs: list[Union[ToolMessage, Command]],
         input_type: Literal["list", "dict", "tool_calls"],
     ) -> list[Union[Command, list[ToolMessage], dict[str, list[ToolMessage]]]]:
         # preserve existing behavior for non-command tool outputs for backwards
@@ -437,10 +445,27 @@ class ToolNode(RunnableCallable):
         call: ToolCall,
         input_type: Literal["list", "dict", "tool_calls"],
         config: RunnableConfig,
-    ) -> ToolMessage:
+    ) -> Union[ToolMessage, Command]:
         """Run a single tool call synchronously."""
         if invalid_tool_message := self._validate_tool_call(call):
             return invalid_tool_message
+
+        # Handle structured output tools
+        if call["name"] in self.structured_output_tools:
+            response_schema = self.tools_by_name[call["name"]]
+            return Command(
+                update={
+                    "messages": [
+                        ToolMessage(
+                            content=msg_content_output(call["args"]),
+                            name=call["name"],
+                            tool_call_id=call["id"],
+                        )
+                    ],
+                    "structured_response": response_schema(**call["args"]),
+                }
+            )
+
         try:
             call_args = {**call, **{"type": "tool_call"}}
             response = self.tools_by_name[call["name"]].invoke(call_args, config)
@@ -493,11 +518,27 @@ class ToolNode(RunnableCallable):
         call: ToolCall,
         input_type: Literal["list", "dict", "tool_calls"],
         config: RunnableConfig,
-    ) -> ToolMessage:
+    ) -> Union[ToolMessage, Command]:
         """Run a single tool call asynchronously."""
         if invalid_tool_message := self._validate_tool_call(call):
             return invalid_tool_message
 
+        # Handle structured output tools
+        if call["name"] in self.structured_output_tools:
+            response_schema = self.tools_by_name[call["name"]]
+            return Command(
+                update={
+                    "messages": [
+                        ToolMessage(
+                            content=msg_content_output(call["args"]),
+                            name=call["name"],
+                            tool_call_id=call["id"],
+                        )
+                    ],
+                    "structured_response": response_schema(**call["args"]),
+                }
+            )
+
         try:
             call_args = {**call, **{"type": "tool_call"}}
             response = await self.tools_by_name[call["name"]].ainvoke(call_args, config)

libs/prebuilt/tests/test_tool_node.py

@@ -1481,3 +1481,71 @@ def test_tool_node_stream_writer() -> None:
             },
         ),
     ]
+
+
+
+async def test_structured_output_tools():
+    """Test that ToolNode handles Pydantic model classes as structured output tools."""
+
+    class OutputSchema(BaseModel):
+        name: str
+        age: int
+        location: str
+
+    tool_node = ToolNode([OutputSchema])
+
+    # Test that the structured output tool is registered correctly
+    assert "OutputSchema" in tool_node.tools_by_name
+    assert "OutputSchema" in tool_node.structured_output_tools
+
+    # Create a tool call that matches the schema
+    tool_call = {
+        "name": "OutputSchema",
+        "args": {"name": "Alice", "age": 30, "location": "NYC"},
+        "id": "call_123",
+        "type": "tool_call",
+    }
+
+    # Test sync execution
+    result = tool_node.invoke(
+        {"messages": [AIMessage(content="", tool_calls=[tool_call])]}
+    )
+
+    # Should return a Command with structured response
+    assert isinstance(result, list)
+    assert len(result) == 1
+    command = result[0]
+    assert isinstance(command, Command)
+
+    # Check the update structure
+    assert "messages" in command.update
+    assert "structured_response" in command.update
+
+    # Check the tool message
+    tool_message = command.update["messages"][0]
+    assert isinstance(tool_message, ToolMessage)
+    assert tool_message.name == "OutputSchema"
+    assert tool_message.tool_call_id == "call_123"
+
+    # Check the structured response
+    structured_response = command.update["structured_response"]
+    assert isinstance(structured_response, OutputSchema)
+    assert structured_response.name == "Alice"
+    assert structured_response.age == 30
+    assert structured_response.location == "NYC"
+
+    # Test async execution
+    result_async = await tool_node.ainvoke(
+        {"messages": [AIMessage(content="", tool_calls=[tool_call])]}
+    )
+
+    # Should produce the same result
+    assert isinstance(result_async, list)
+    assert len(result_async) == 1
+    command_async = result_async[0]
+    assert isinstance(command_async, Command)
+    assert "structured_response" in command_async.update
+
+    structured_response_async = command_async.update["structured_response"]
+    assert isinstance(structured_response_async, OutputSchema)
+    assert structured_response_async.name == "Alice"