support ToolNode as input to create_supervisor (#135)

`create_react_agent` in `langgraph-prebuilt` allows to pass a `ToolNode`
for more precise management of the ToolExecutions (for example to raise
an error instead of returning a `ToolMessage`). It would be nice if
`langgraph-supervisor` could also accept a `ToolNode`.

---------

Co-authored-by: vbarda <[email protected]>

Author: jmaillefaud

langgraph_supervisor/supervisor.py

@@ -1,15 +1,17 @@
 import inspect
-from typing import Any, Callable, Literal, Optional, Type, Union, cast, get_args
+from typing import Any, Callable, Literal, Optional, Sequence, Type, Union, cast, get_args
 
 from langchain_core.language_models import BaseChatModel, LanguageModelLike
 from langchain_core.runnables import RunnableConfig
 from langchain_core.tools import BaseTool
 from langgraph.graph import END, START, StateGraph
+from langgraph.prebuilt import ToolNode
 from langgraph.prebuilt.chat_agent_executor import (
     AgentState,
     Prompt,
     StateSchemaType,
     StructuredResponseSchema,
+    _should_bind_tools,
     create_react_agent,
 )
 from langgraph.pregel import Pregel
@@ -93,7 +95,7 @@ async def acall_agent(state: dict, config: RunnableConfig) -> dict:
     return RunnableCallable(call_agent, acall_agent)
 
 
-def _get_handoff_destinations(tools: list[BaseTool | Callable]) -> list[str]:
+def _get_handoff_destinations(tools: Sequence[BaseTool | Callable]) -> list[str]:
     """Extract handoff destinations from provided tools.
     Args:
         tools: List of tools to inspect.
@@ -109,11 +111,70 @@ def _get_handoff_destinations(tools: list[BaseTool | Callable]) -> list[str]:
     ]
 
 
+def _prepare_tool_node(
+    tools: list[BaseTool | Callable] | ToolNode | None,
+    handoff_tool_prefix: Optional[str],
+    add_handoff_messages: bool,
+    agent_names: set[str],
+) -> ToolNode:
+    """Prepare the ToolNode to use in supervisor agent."""
+    if isinstance(tools, ToolNode):
+        input_tool_node = tools
+        tool_classes = list(tools.tools_by_name.values())
+    elif tools:
+        input_tool_node = ToolNode(tools)
+        # get the tool functions wrapped in a tool class from the ToolNode
+        tool_classes = list(input_tool_node.tools_by_name.values())
+    else:
+        input_tool_node = None
+        tool_classes = []
+
+    handoff_destinations = _get_handoff_destinations(tool_classes)
+    if handoff_destinations:
+        if missing_handoff_destinations := set(agent_names) - set(handoff_destinations):
+            raise ValueError(
+                "When providing custom handoff tools, you must provide them for all subagents. "
+                f"Missing handoff tools for agents '{missing_handoff_destinations}'."
+            )
+
+        # Handoff tools should be already provided here
+        tool_node = cast(ToolNode, input_tool_node)
+    else:
+        handoff_tools = [
+            create_handoff_tool(
+                agent_name=agent_name,
+                name=(
+                    None
+                    if handoff_tool_prefix is None
+                    else f"{handoff_tool_prefix}{_normalize_agent_name(agent_name)}"
+                ),
+                add_handoff_messages=add_handoff_messages,
+            )
+            for agent_name in agent_names
+        ]
+        all_tools = tool_classes + list(handoff_tools)
+
+        # re-wrap the combined tools in a ToolNode
+        # if the original input was a ToolNode, apply the same params
+        if input_tool_node is not None:
+            tool_node = ToolNode(
+                all_tools,
+                name=input_tool_node.name,
+                tags=list(input_tool_node.tags) if input_tool_node.tags else None,
+                handle_tool_errors=input_tool_node.handle_tool_errors,
+                messages_key=input_tool_node.messages_key,
+            )
+        else:
+            tool_node = ToolNode(all_tools)
+
+    return tool_node
+
+
 def create_supervisor(
     agents: list[Pregel],
     *,
     model: LanguageModelLike,
-    tools: list[BaseTool | Callable] | None = None,
+    tools: list[BaseTool | Callable] | ToolNode | None = None,
     prompt: Prompt | None = None,
     response_format: Optional[
         Union[StructuredResponseSchema, tuple[str, StructuredResponseSchema]]
@@ -256,45 +317,29 @@ def web_search(query: str) -> str:
 
         agent_names.add(agent.name)
 
-    handoff_destinations = _get_handoff_destinations(tools or [])
-    if handoff_destinations:
-        if missing_handoff_destinations := set(agent_names) - set(handoff_destinations):
-            raise ValueError(
-                "When providing custom handoff tools, you must provide them for all subagents. "
-                f"Missing handoff tools for agents '{missing_handoff_destinations}'."
-            )
+    tool_node = _prepare_tool_node(
+        tools,
+        handoff_tool_prefix,
+        add_handoff_messages,
+        agent_names,
+    )
+    all_tools = list(tool_node.tools_by_name.values())
 
-        # Handoff tools should be already provided here
-        all_tools = tools or []
-    else:
-        handoff_tools = [
-            create_handoff_tool(
-                agent_name=agent.name,
-                name=(
-                    None
-                    if handoff_tool_prefix is None
-                    else f"{handoff_tool_prefix}{_normalize_agent_name(agent.name)}"
-                ),
-                add_handoff_messages=add_handoff_messages,
+    if _should_bind_tools(model, all_tools):
+        if _supports_disable_parallel_tool_calls(model):
+            model = cast(BaseChatModel, model).bind_tools(
+                all_tools, parallel_tool_calls=parallel_tool_calls
             )
-            for agent in agents
-        ]
-        all_tools = (tools or []) + list(handoff_tools)
-
-    if _supports_disable_parallel_tool_calls(model):
-        model = cast(BaseChatModel, model).bind_tools(
-            all_tools, parallel_tool_calls=parallel_tool_calls
-        )
-    else:
-        model = cast(BaseChatModel, model).bind_tools(all_tools)
+        else:
+            model = cast(BaseChatModel, model).bind_tools(all_tools)
 
     if include_agent_name:
         model = with_agent_name(model, include_agent_name)
 
     supervisor_agent = create_react_agent(
         name=supervisor_name,
         model=model,
-        tools=all_tools,
+        tools=tool_node,
         prompt=prompt,
         state_schema=state_schema,
         response_format=response_format,