diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 1b1b88cf1..f5338a024 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -605,7 +605,7 @@ class CompiledStateGraph(CompiledGraph): def _get_root(input: Any) -> Any: if isinstance(input, Control): - return input.update_state + return input.state else: return input @@ -619,7 +619,7 @@ class CompiledStateGraph(CompiledGraph): ) return input.get(key, SKIP_WRITE) elif isinstance(input, Control): - return _get_state_key(input.update_state, key=key) + return _get_state_key(input.state, key=key) elif get_type_hints(type(input)): value = getattr(input, key, SKIP_WRITE) return value if value is not None else SKIP_WRITE @@ -797,13 +797,15 @@ def _coerce_state(schema: Type[Any], input: dict[str, Any]) -> dict[str, Any]: def _control_branch(value: Any) -> Sequence[Union[str, Send]]: + if isinstance(value, Send): + return [value] if not isinstance(value, Control): return EMPTY_SEQ rtn: list[Union[str, Send]] = [] - if isinstance(value.trigger, str): - rtn.append(value.trigger) + if isinstance(value.goto, str): + rtn.append(value.goto) else: - rtn.extend(value.trigger) + rtn.extend(value.goto) if isinstance(value.send, Send): rtn.append(value.send) else: @@ -811,14 +813,16 @@ def _control_branch(value: Any) -> Sequence[Union[str, Send]]: return rtn -async def _acontrol_branch(value: Any) -> None: +async def _acontrol_branch(value: Any) -> Sequence[Union[str, Send]]: + if isinstance(value, Send): + return [value] if not isinstance(value, Control): return EMPTY_SEQ rtn: list[Union[str, Send]] = [] - if isinstance(value.trigger, str): - rtn.append(value.trigger) + if isinstance(value.goto, str): + rtn.append(value.goto) else: - rtn.extend(value.trigger) + rtn.extend(value.goto) if isinstance(value.send, Send): rtn.append(value.send) else: diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index b1a40f1a1..88d2d06f2 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -229,17 +229,17 @@ N = TypeVar("N") class Control(Generic[N]): """A control object to update the graph's state, trigger nodes, and send messages.""" - __slots__ = ("update_state", "trigger", "send") + __slots__ = ("state", "goto", "send") def __init__( self, *, - update_state: Optional[dict[str, Any]] = None, - trigger: Union[str, Sequence[str]] = (), + state: Optional[dict[str, Any]] = None, + goto: Union[str, Sequence[str]] = (), send: Union[Send, Sequence[Send]] = (), ) -> None: - self.update_state = update_state - self.trigger = trigger + self.state = state + self.goto = goto self.send = send def __repr__(self) -> str: diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 71b410ed4..e84a35214 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -1809,7 +1809,7 @@ def test_send_sequences() -> None: else ["|".join((self.name, str(state)))] ) if isinstance(state, Control): - state.update_state = update + state.state = update return state else: return update @@ -2313,7 +2313,7 @@ def test_send_react_interrupt_control( def agent(state) -> Control[Literal["foo"]]: return Control( - update_state={"messages": ai_message}, + state={"messages": ai_message}, send=[Send(call["name"], call) for call in ai_message.tool_calls], ) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 7f9c533d3..d04c7e48f 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -2026,7 +2026,7 @@ async def test_send_sequences() -> None: else ["|".join((self.name, str(state)))] ) if isinstance(state, Control): - state.update_state = update + state.state = update return state else: return update @@ -2523,7 +2523,7 @@ async def test_send_react_interrupt_control(checkpointer_name: str) -> None: async def agent(state) -> Control[Literal["foo"]]: return Control( - update_state={"messages": ai_message}, + state={"messages": ai_message}, send=[Send(call["name"], call) for call in ai_message.tool_calls], ) @@ -2822,7 +2822,7 @@ async def test_max_concurrency(checkpointer_name: str) -> None: @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_max_concurrency_control(checkpointer_name: str) -> None: async def node1(state) -> Control[Literal["2"]]: - return Control(update_state=["1"], send=[Send("2", idx) for idx in range(100)]) + return Control(state=["1"], send=[Send("2", idx) for idx in range(100)]) node2_currently = 0 node2_max_currently = 0 @@ -2835,7 +2835,7 @@ async def test_max_concurrency_control(checkpointer_name: str) -> None: await asyncio.sleep(0.1) node2_currently -= 1 - return Control(update_state=[state], trigger="3") + return Control(state=[state], goto="3") async def node3(state) -> Literal["3"]: return ["3"]