diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index daa942a39..98ac8576f 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -13,6 +13,7 @@ from typing import ( Protocol, Sequence, Union, + cast, overload, ) from uuid import UUID @@ -471,7 +472,7 @@ def prepare_single_task( else: return PregelTask(task_id, packet.node, task_path) elif task_path[0] == PULL: - name = str(task_path[1]) + name = cast(str, task_path[1]) if name not in processes: return proc = processes[name] diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 0bd5772d4..29e287da2 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -1,3 +1,4 @@ +import enum import json import operator import re @@ -11517,3 +11518,23 @@ def test_store_injected( "some_val": 0, } # Overwrites the whole doc assert len(the_store.search(("foo", "bar"))) == 1 # still overwriting the same one + + +def test_enum_node_names(): + class NodeName(str, enum.Enum): + BAZ = "baz" + + class State(TypedDict): + foo: str + bar: str + + def baz(state: State): + return {"bar": state["foo"] + "!"} + + graph = StateGraph(State) + graph.add_node(NodeName.BAZ, baz) + graph.add_edge(START, NodeName.BAZ) + graph.add_edge(NodeName.BAZ, END) + graph = graph.compile() + + assert graph.invoke({"foo": "hello"}) == {"foo": "hello", "bar": "hello!"}