langgraph: fix edge case with string enums as node names (#1926)

This commit is contained in:
Vadym Barda
2024-09-30 18:00:33 -04:00
committed by GitHub
parent 6d1d705b84
commit da935e7805
2 changed files with 23 additions and 1 deletions
+2 -1
View File
@@ -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]
+21
View File
@@ -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!"}