Merge pull request #2368 from langchain-ai/nc/7nov/control

lib: Rename args in Control object

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:

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:

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],
         )
 

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"]