From b06dbaae7a6b7e58bc4e61f46b3eceaa6073288d Mon Sep 17 00:00:00 2001 From: Eugene Yurtsev Date: Wed, 13 Aug 2025 11:17:19 -0400 Subject: [PATCH] x --- libs/prebuilt/langgraph/prebuilt/tool_node.py | 57 +++++++++++++--- libs/prebuilt/tests/test_tool_node.py | 68 +++++++++++++++++++ 2 files changed, 117 insertions(+), 8 deletions(-) diff --git a/libs/prebuilt/langgraph/prebuilt/tool_node.py b/libs/prebuilt/langgraph/prebuilt/tool_node.py index 252f88802..4a674ae67 100644 --- a/libs/prebuilt/langgraph/prebuilt/tool_node.py +++ b/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) diff --git a/libs/prebuilt/tests/test_tool_node.py b/libs/prebuilt/tests/test_tool_node.py index e9623bd5e..b952ee6a7 100644 --- a/libs/prebuilt/tests/test_tool_node.py +++ b/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"