Confusing behavior when using add_conditional_edges with BaseModel

[observerw]

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?

[yguedidi]

Experiencing the same issue, would be great to fix it or at least understand why it behaves this way.
Current workaround is to add an intermediate node to be used with conditional edges, for which the input is the output of the previous node, so the conditional edge gets the right type

[ikhvostenkov]

Any updates on this issue?