Private state doesn't work with instance methods as nodes

[baskaryan]

Privileged issue

• I am a LangChain maintainer, or was asked directly by a LangChain maintainer to create an issue here.

Issue Content

if you run below can see that node_2 doesn't receive the private_data key as input:

from langgraph.graph import StateGraph, START, END
from typing_extensions import TypedDict


# The overall state of the graph (this is the public state shared across nodes)
class OverallState(TypedDict):
    a: str


# Output from node_1 contains private data that is not part of the overall state
class Node1Output(TypedDict):
    private_data: str


# Node 2 input only requests the private data available after node_1
class Node2Input(TypedDict):
    private_data: str


class Nodes:
    # The private data is only shared between node_1 and node_2
    def node_1(self, state: OverallState) -> Node1Output:
        output = {"private_data": "set by node_1"}
        print(f"Entered node `node_1`:\n\tInput: {state}.\n\tReturned: {output}")
        return output


    def node_2(self, state: Node2Input) -> OverallState:
        output = {"a": "set by node_2"}
        print(f"Entered node `node_2`:\n\tInput: {state}.\n\tReturned: {output}")
        return output


    # Node 3 only has access to the overall state (no access to private data from node_1)
    def node_3(self, state: OverallState) -> OverallState:
        output = {"a": "set by node_3"}
        print(f"Entered node `node_3`:\n\tInput: {state}.\n\tReturned: {output}")
        return output

nodes = Nodes()

# Build the state graph
builder = StateGraph(OverallState)
builder.add_node(nodes.node_1)  # node_1 is the first node
builder.add_node(nodes.node_2)  # node_2 is the second node and accepts private data from node_1
builder.add_node(nodes.node_3)  # node_3 is the third node and does not see the private data
builder.add_edge(START, "node_1")  # Start the graph with node_1
builder.add_edge("node_1", "node_2")  # Pass from node_1 to node_2
builder.add_edge("node_2", "node_3")  # Pass from node_2 to node_3 (only overall state is shared)
builder.add_edge("node_3", END)  # End the graph after node_3
graph = builder.compile()
graph.invoke({"a": ""})

Entered node `node_1`:
	Input: {'a': ''}.
	Returned: {'private_data': 'set by node_1'}
Entered node `node_2`:
	Input: {'a': ''}.
	Returned: {'a': 'set by node_2'}
Entered node `node_3`:
	Input: {'a': 'set by node_2'}.
	Returned: {'a': 'set by node_3'}