From 764929afd93e80ca7571d9e6fda6ea6df7450d9f Mon Sep 17 00:00:00 2001
From: Tat Dat Duong <david@duong.cz>
Date: Fri, 7 Mar 2025 19:32:24 +0100
Subject: [PATCH] Fix lint issues

---
 libs/langgraph/langgraph/pregel/__init__.py | 34 ++++++++++++---------
 1 file changed, 20 insertions(+), 14 deletions(-)

diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py
index 5b37e80fd..69d69bfab 100644
--- a/libs/langgraph/langgraph/pregel/__init__.py
+++ b/libs/langgraph/langgraph/pregel/__init__.py
@@ -1415,7 +1415,7 @@ class Pregel(PregelProtocol):
 
                     valid_updates.append((values, as_node))
 
-            tasks: list[tuple[str, PregelTaskWrites]] = []
+            run_tasks: list[tuple[str, PregelTaskWrites]] = []
             for values, as_node in valid_updates:
                 # create task to run all writers of the chosen node
                 writers = self.nodes[as_node].flat_writers
@@ -1424,9 +1424,9 @@ class Pregel(PregelProtocol):
                 writes: deque[tuple[str, Any]] = deque()
                 task = PregelTaskWrites((), as_node, writes, [INTERRUPT])
                 task_id = str(uuid5(UUID(checkpoint["id"]), INTERRUPT))
-                tasks.append((task_id, task))
+                run_tasks.append((task_id, task))
 
-            for _, task in tasks:
+            for _, task in run_tasks:
                 run = RunnableSequence(*writers) if len(writers) > 1 else writers[0]
                 # execute task
                 run.invoke(
@@ -1455,7 +1455,7 @@ class Pregel(PregelProtocol):
                 )
 
             # save task writes
-            for _, task in tasks:
+            for _, task in run_tasks:
                 channel_writes = [w for w in task.writes if w[0] != PUSH]
 
                 # channel writes are saved to current checkpoint
@@ -1464,7 +1464,10 @@ class Pregel(PregelProtocol):
 
             # apply to checkpoint and save
             mv_writes = apply_writes(
-                checkpoint, channels, tasks, checkpointer.get_next_version
+                checkpoint,
+                channels,
+                [t for _, t in run_tasks],
+                checkpointer.get_next_version,
             )
 
             assert not mv_writes, "Can't write to SharedValues from update_state"
@@ -1485,7 +1488,7 @@ class Pregel(PregelProtocol):
                 ),
             )
 
-            for task_id, task in tasks:
+            for task_id, task in run_tasks:
                 if push_writes := [w for w in task.writes if w[0] == PUSH]:
                     checkpointer.put_writes(next_config, push_writes, task_id)
 
@@ -1747,7 +1750,7 @@ class Pregel(PregelProtocol):
 
                     valid_updates.append((values, as_node))
 
-            tasks: list[tuple[str, PregelTaskWrites]] = []
+            run_tasks: list[tuple[str, PregelTaskWrites]] = []
 
             for values, as_node in valid_updates:
                 # create task to run all writers of the chosen node
@@ -1757,9 +1760,9 @@ class Pregel(PregelProtocol):
                 writes: deque[tuple[str, Any]] = deque()
                 task = PregelTaskWrites((), as_node, writes, [INTERRUPT])
                 task_id = str(uuid5(UUID(checkpoint["id"]), INTERRUPT))
-                tasks.append((task_id, task))
+                run_tasks.append((task_id, task))
 
-            for _, task in tasks:
+            for _, task in run_tasks:
                 run = RunnableSequence(*writers) if len(writers) > 1 else writers[0]
                 # execute task
                 await run.ainvoke(
@@ -1788,7 +1791,7 @@ class Pregel(PregelProtocol):
                 )
 
             # save task writes
-            for _, task in tasks:
+            for _, task in run_tasks:
                 # channel writes are saved to current checkpoint
                 channel_writes = [w for w in task.writes if w[0] != PUSH]
                 if saved and channel_writes:
@@ -1798,7 +1801,10 @@ class Pregel(PregelProtocol):
 
             # apply to checkpoint and save
             mv_writes = apply_writes(
-                checkpoint, channels, [task], checkpointer.get_next_version
+                checkpoint,
+                channels,
+                [t for _, t in run_tasks],
+                checkpointer.get_next_version,
             )
             assert not mv_writes, "Can't write to SharedValues from update_state"
             checkpoint = create_checkpoint(checkpoint, channels, step + 1)
@@ -1818,7 +1824,7 @@ class Pregel(PregelProtocol):
                 ),
             )
 
-            for task_id, task in tasks:
+            for task_id, task in run_tasks:
                 # save push writes
                 if push_writes := [w for w in task.writes if w[0] == PUSH]:
                     await checkpointer.aput_writes(next_config, push_writes, task_id)
@@ -1835,7 +1841,7 @@ class Pregel(PregelProtocol):
         node `as_node`. If `as_node` is not provided, it will be set to the last node
         that updated the state, if not ambiguous.
         """
-        return self.bulk_update_state(config, [(values, as_node)])
+        return self.bulk_update_state(config, [BulkUpdate(values, as_node)])
 
     async def aupdate_state(
         self,
@@ -1847,7 +1853,7 @@ class Pregel(PregelProtocol):
         node `as_node`. If `as_node` is not provided, it will be set to the last node
         that updated the state, if not ambiguous.
         """
-        return await self.bulk_update_state(config, [(values, as_node)])
+        return await self.abulk_update_state(config, [BulkUpdate(values, as_node)])
 
     def _defaults(
         self,