mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-07 10:17:50 +02:00
langgraph: fix edge case with string enums as node names (#1926)
This commit is contained in:
@@ -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]
|
||||
|
||||
@@ -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!"}
|
||||
|
||||
Reference in New Issue
Block a user