Allow ToolNode to be initialised with optional argument messages_key: str = "messages" or something similar

[tituskx]

I'm building a graph that consists of several subgraphs (using this example).

I'm having some difficulty with the shared states and the messages key. When you define subgraphs you can create private states by defining subgraph keys that are not defined in the parent state. For example:

class ParentState(TypedDict):
    messages: Annotated[list[AnyMessage], add_messages]

class SubgraphState(TypedDict): 
    response: str

However, the parent state usually has a messages key which means that if you want to use a messages key that is private to a particular subgraph you should call it something like subgraph_messages (if you don't do this the subgraph will share the messages key with the parent graph).

This is a problem because ToolNode requires a messages key to be present in the state that it receives . That means that if I want to use ToolNode somewhere in my subgraph I need to use a messages key in the subgraph state, which forces me to share it with the parent graph.

Would love to hear if I'm missing something and should implement state sharing / tool calling in my subgraph differently! If I'm not missing something then I would suggest making messages_key an optional argument to the ToolNode class that has default value messages. That way when I want to use the ToolNode in my subgraph I can do so by initialising tool nodes in this subgraph using the private (to this subgraph) subgraph_messages key.

[NewGHUser4321]

The ToolNode is roughly analogous to:

tools_by_name = {tool.name: tool for tool in tools}
def tool_node(state: dict):
    result = []
    for tool_call in state["messages"][-1].tool_calls:
        tool = tools_by_name[tool_call["name"]]
        observation = tool.invoke(tool_call["args"])
        result.append(ToolMessage(content=observation, tool_call_id=tool_call["id"]))
    return {"subgraph_messages": result}

You can define a custom tool node very easily in the subgraph