add support for parallel handoffs (#85)

* filter out tool calls not meant for individual agents
* update `langgraph-prebuilt` to handle combining multiple `Send` +
Command.PARENT commands

Author: Vadym Barda

langgraph_supervisor/handoff.py

@@ -1,10 +1,11 @@
 import re
 import uuid
+from typing import cast
 
-from langchain_core.messages import AIMessage, ToolCall, ToolMessage
+from langchain_core.messages import AIMessage, BaseMessage, ToolCall, ToolMessage
 from langchain_core.tools import BaseTool, InjectedToolCallId, tool
 from langgraph.prebuilt import InjectedState
-from langgraph.types import Command
+from langgraph.types import Command, Send
 from typing_extensions import Annotated
 
 WHITESPACE_RE = re.compile(r"\s+")
@@ -15,6 +16,43 @@ def _normalize_agent_name(agent_name: str) -> str:
     return WHITESPACE_RE.sub("_", agent_name.strip()).lower()
 
 
+def _remove_non_handoff_tool_calls(
+    messages: list[BaseMessage], handoff_tool_name: str
+) -> list[BaseMessage]:
+    """Remove tool calls that are not meant for the agent."""
+    last_ai_message = cast(AIMessage, messages[-1])
+    # if the supervisor is calling multiple agents/tools in parallel,
+    # we need to remove tool calls that are not meant for this agent
+    # to ensure that the resulting message history is valid
+    if len(last_ai_message.tool_calls) > 1 and any(
+        tool_call["name"] == handoff_tool_name for tool_call in last_ai_message.tool_calls
+    ):
+        content = last_ai_message.content
+        if isinstance(content, list) and len(content) > 1 and isinstance(content[0], dict):
+            content = [
+                content_block
+                for content_block in content
+                if (
+                    content_block["type"] == "tool_use"
+                    and content_block["name"] == handoff_tool_name
+                )
+                or content_block["type"] != "tool_use"
+            ]
+
+        last_ai_message = AIMessage(
+            content=content,
+            tool_calls=[
+                tool_call
+                for tool_call in last_ai_message.tool_calls
+                if tool_call["name"] == handoff_tool_name
+            ],
+            name=last_ai_message.name,
+            id=str(uuid.uuid4()),
+        )
+
+    return messages[:-1] + [last_ai_message]
+
+
 def create_handoff_tool(*, agent_name: str) -> BaseTool:
     """Create a tool that can handoff control to the requested agent.
 
@@ -39,10 +77,17 @@ def handoff_to_agent(
             name=tool_name,
             tool_call_id=tool_call_id,
         )
+        handoff_messages = _remove_non_handoff_tool_calls(state["messages"], tool_name) + [
+            tool_message
+        ]
         return Command(
-            goto=agent_name,
             graph=Command.PARENT,
-            update={"messages": state["messages"] + [tool_message]},
+            # NOTE: we are using Send here to allow the ToolNode in langgraph.prebuilt
+            # to handle parallel handoffs by combining all Send commands into a single command
+            goto=[Send(agent_name, {"messages": handoff_messages})],
+            # we also propagate the update to make sure the handoff messages are applied
+            # to the parent graph's state
+            update={"messages": handoff_messages},
         )
 
     return handoff_to_agent

langgraph_supervisor/supervisor.py

@@ -1,7 +1,6 @@
-import inspect
 from typing import Any, Callable, Literal, Optional, Type, Union
 
-from langchain_core.language_models import BaseChatModel, LanguageModelLike
+from langchain_core.language_models import LanguageModelLike
 from langchain_core.tools import BaseTool
 from langgraph.graph import END, START, StateGraph
 from langgraph.prebuilt.chat_agent_executor import (
@@ -28,25 +27,6 @@
 """
 
 
-MODELS_NO_PARALLEL_TOOL_CALLS = {"o3-mini"}
-
-
-def _supports_disable_parallel_tool_calls(model: LanguageModelLike) -> bool:
-    if not isinstance(model, BaseChatModel):
-        return False
-
-    if hasattr(model, "model_name") and model.model_name in MODELS_NO_PARALLEL_TOOL_CALLS:
-        return False
-
-    if not hasattr(model, "bind_tools"):
-        return False
-
-    if "parallel_tool_calls" not in inspect.signature(model.bind_tools).parameters:
-        return False
-
-    return True
-
-
 def _make_call_agent(
     agent: Pregel,
     output_mode: OutputMode,
@@ -168,11 +148,7 @@ def create_supervisor(
 
     handoff_tools = [create_handoff_tool(agent_name=agent.name) for agent in agents]
     all_tools = (tools or []) + handoff_tools
-
-    if _supports_disable_parallel_tool_calls(model):
-        model = model.bind_tools(all_tools, parallel_tool_calls=False)
-    else:
-        model = model.bind_tools(all_tools)
+    model = model.bind_tools(all_tools)
 
     if include_agent_name:
         model = with_agent_name(model, include_agent_name)

pyproject.toml

@@ -13,7 +13,7 @@ readme = "README.md"
 requires-python = ">=3.10"
 dependencies = [
     "langgraph>=0.3.5,<0.4.0",
-    "langgraph-prebuilt>=0.1.2,<0.2.0"
+    "langgraph-prebuilt>=0.1.6,<0.2.0"
 ]
 
 [dependency-groups]