Parallel using tools doesn't work

[JanPalasek]

I tried to create a langgraph with a parallel branches implementing the following algorithm:

1. Split processing into two branches: branch_a and branch_b. Those two branches are responsible for computing part of the response.
2. Process each branch separately using GPT and ToolNode.
3. Merge the branches.

Unfortunately, it crashes with the following error:

Error code: 400 - {'error': {'message': "An assistant message with 'tool_calls' must be followed by tool messages responding to each 'tool_call_id'. The following tool_call_ids did not have response messages: call_KsRwXkYLkZfRcXX3AWDmPgHB, call_41MdeD1uFBE81IpsiqdPdlwC", 'type': 'invalid_request_error', 'param': 'messages.[3].role', 'code': None}}
httpx.HTTPStatusError: Client error '400 Bad Request' for url 'https://api.openai.com/v1/chat/completions'
For more information check: https://developer.mozilla.org/en-US/docs/Web/HTTP/Status/400

During handling of the above exception, another exception occurred:

  File "/workspaces/beatrix/test.py", line 54, in handle_task
    return chain.invoke({"messages": messages})
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspaces/beatrix/test.py", line 117, in <module>
    result = graph.invoke("How many hungry and thirsty children are there?")

Right before it crashes on chain.invoke, the messages are following:

It seems that e.g. tool call for call with id call_KsRwXkYLkZfRcXX3AWDmPgHB was not made at all, while call for other tool (id call_8TcreVwtNmLAoQOwWSjwYz9j) was duplicated.

Any idea how to fix this? It seems like a bug - something in ToolNode is leaking. But I'm really not sure. Also the graph is really weird.

The graph

Code Repro

import functools

from langchain.prompts import (
    ChatPromptTemplate,
    SystemMessagePromptTemplate,
)
from langchain_core.messages import BaseMessage
from langchain_core.prompts.chat import MessagesPlaceholder
from langchain_core.runnables.graph import MermaidDrawMethod
from langchain_core.tools import Tool, tool
from langchain_openai import ChatOpenAI
from langgraph.graph import END, START, MessageGraph
from langgraph.prebuilt.tool_node import ToolNode


def _should_continue(messages, next_node: str, process_tool_node: str):
    last_message = messages[-1]

    # If there is no function call, then we finish
    if not last_message.tool_calls:
        return next_node
    return process_tool_node


def get_model():
    return ChatOpenAI(api_key="XXX", model="gpt-4o")


def handle_task(
    messages: list[BaseMessage],
    task_description: str,
    tools: list[Tool],
) -> BaseMessage:
    # bind tools if there are any
    model = get_model().bind_tools(tools=tools)

    chain = (
        ChatPromptTemplate.from_messages(
            [
                SystemMessagePromptTemplate.from_template(
                    template="You are a helpful assistant. You have a list of tools and try to help with the task.",
                    partial_variables={
                        "task_description": task_description,
                    },
                ),
                MessagesPlaceholder(variable_name="messages"),
            ]
        )
        | model
    )

    return chain.invoke({"messages": messages})


graph = MessageGraph()

# define branches
graph.add_conditional_edges(START, lambda _: ["branch_a", "branch_b"])


@tool
def number_of_hungry_children() -> int:
    """
    Returns the number of hungry children in the world.

    Returns
    -------
    int
        The number of hungry children.
    """
    return 1e7

branch_name = "branch_a"
branch_tools = [number_of_hungry_children]
graph.add_node(branch_name, functools.partial(handle_task, task_description="You respond to questions about number of hungry children.", tools=branch_tools))
graph.add_conditional_edges(branch_name, lambda messages, branch_name=branch_name, next_node=f"{branch_name}-final": functools.partial(
                _should_continue, next_node=next_node, process_tool_node=f"{branch_name}-action"
            )(messages))
graph.add_node(f"{branch_name}-action", ToolNode(branch_tools, name=f"{branch_name}-tools"))  # call the tool
graph.add_edge(f"{branch_name}-action", branch_name)  # go back to gpt so it postprocesses the answer
graph.add_node(f"{branch_name}-final", lambda messages: messages)
graph.add_edge(f"{branch_name}-final", "aggregate")

@tool
def thirsty_children_count() -> int:
    """
    Returns the number of thirsty children in the world.

    Returns
    -------
    int
        The number of thirsty children.
    """
    return 1e6

branch_name = "branch_b"
branch_tools = [thirsty_children_count]
graph.add_node(branch_name, functools.partial(handle_task, task_description="You respond to questions about thirsty children.", tools=branch_tools))
graph.add_conditional_edges(branch_name, lambda messages, branch_name=branch_name, next_node=f"{branch_name}-final": functools.partial(
                _should_continue, next_node=next_node, process_tool_node=f"{branch_name}-action"
            )(messages))
graph.add_node(f"{branch_name}-action", ToolNode(branch_tools, name=f"{branch_name}-tools"))  # call the tool
graph.add_edge(f"{branch_name}-action", branch_name)  # go back to gpt so it postprocesses the answer
graph.add_node(f"{branch_name}-final", lambda messages: messages)
graph.add_edge(f"{branch_name}-final", "aggregate")

# aggregate
graph.add_node("aggregate", lambda messages: messages)
graph.add_edge("aggregate", END)

graph = graph.compile()

graph.get_graph().draw_mermaid_png(draw_method=MermaidDrawMethod.API, output_file_path="graph.png")

result = graph.invoke("How many hungry and thirsty children are there?")

print(result)

Versions:

• Python: 3.11.9
• Langgraph: 0.1.9
• Langchain: 0.2.10
• Langchain OpenAI: 0.1.17

[hinthornw]

Will address pt 2 first since it's easy: for the "graph rendering looks funny" bit, it's because there are only 2 ways we can infer conditional edge connectivity at "compile" time:

1. Via type annotations in the function you pass in (not present here since you're using a lambda)
2. Via a "path map" parameter: either a dictionary (mapping function outputs -> node name) or list (stating the valid outputs. The function outputs must be node names in this case)

[hinthornw]

For pt 1 (the actual error), here's what I think is going on:

I think that you're understanding of branching is slightly different than how it works in StateGraph. I think the graph you shared assumes that the graph "forks" processes so that there are two separate states that exist in parallel over multiple steps and which would then be synchronized only at the end. How it actually works is that each set of tasks is executed logcially in parallel and then the updates are reduced into a single state after each superstep. Here "handle_task" for both branch a and b would be executed in a superstep, then the results would be reduced into a single shared state. Then the tool nodes would be executed both with that shared state (so they would both see the results of only one of the branches, since they each check the last ai message in the list). The first tool node doesn't have branch b's tools, so it errors. This results in two error message tool messages and two success tool messages that are all tackling only thte tool messages of branch b, and branch A's requests don't get processed. The tool nodes are executed in a single superstep, the graph reduces those 2 tool messages from each parallel invocation and adds them to the list, etc.

I think what you want is to frame these as subgraphs so that each branch would be executed entirely within a superstep of the parent graph without sharing state. There are other ways to handle this without using subgraphs (like having separate state channels for each branch), but they may be more complicated depending on how you want to implement it.

[hinthornw]

More discussion on related topics here: https://langchain-ai.github.io/langgraph/how-tos/branching/?h=branchin#stable-sorting

[JanPalasek]

Thank you very much for the quick response. I will try to fix it with the subgraphs, probably tomorrow, and then mark your reply as the answer.