Confusing behavior when using add_conditional_edges with BaseModel

[observerw]

Discussed in #2226

^{Originally posted by observerw October 30, 2024}
Let's say we have the following graph:

class State(TypedDict):
    a: int


class NodeAInput(State):
    private_b: int | None


class NodeAOutput(TypedDict):
    private_a: int


class NodeBInput(State):
    private_a: int


class NodeBOutput(TypedDict):
    private_b: int


def node_a(input: NodeAInput) -> NodeAOutput:
    return NodeAOutput(private_a=1)


def node_b(input: NodeBInput) -> NodeBOutput:
    return NodeBOutput(private_b=2)

# !!! Confusing behavior happens here
def node_b_edge(input: Any):
    print(f"input type: {type(input)}, value: {input}")

    return END


graph = StateGraph(State)
graph.add_node(node_a)
graph.add_node(node_b)

graph.add_edge("node_a", "node_b")
graph.add_conditional_edges("node_b", node_b_edge)

graph.set_entry_point("node_a")

graph = graph.compile()

graph.invoke(State(a=1))

The output is:

input type: <class '__main__.NodeBInput'>, value: a=1 private_a=1

After I added node_b_edge after node_b, I was expecting this function to accept a State with a private_b field, or at least a complete State. However, the function actually accepts a parameter of type NodeBInput instead of State! This is very confusing and results in an error if I wish to use the private_b field to determine the output of node_b_edge.

Nonetheless, when I ditched the BaseModel and used TypedDict, the run results were:

input type: <class 'dict'>, value: {'private_b': 2, 'a': 1, 'private_a': 1}

Well, at least I was able to access the private_b, So I guess the problems here associated with the use of BaseModel.

Is this behavior expected? How can I access private_b when using BaseModel?

[vbarda]

I think this is the same issue as #2504? if so, we can close this as duplicate and we'll prioritize fixing it!

[observerw]

I think this is the same issue as #2504? if so, we can close this as duplicate and we'll prioritize fixing it!

It seems so. I think langgraph is using NodeBOutput to update the input of node_b instead of the State of the graph, then pass it to node_b_edge, which is not expected.