From eed577ee2a60637d5dc453dbd59bb5bedc65a3c9 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Wed, 15 Jan 2025 10:58:46 -0800
Subject: [PATCH] Update sqlite signature

---
 .../langgraph/checkpoint/sqlite/__init__.py   |  2 ++
 .../langgraph/checkpoint/sqlite/aio.py        | 33 +++++++++++--------
 2 files changed, 22 insertions(+), 13 deletions(-)

diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py
index ea749473c..a41d9320e 100644
--- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py
+++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py
@@ -424,6 +424,7 @@ class SqliteSaver(BaseCheckpointSaver[str]):
         config: RunnableConfig,
         writes: Sequence[Tuple[str, Any]],
         task_id: str,
+        task_path: str = "",
     ) -> None:
         """Store intermediate writes linked to a checkpoint.
 
@@ -433,6 +434,7 @@ class SqliteSaver(BaseCheckpointSaver[str]):
             config (RunnableConfig): Configuration of the related checkpoint.
             writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair.
             task_id (str): Identifier for the task creating the writes.
+            task_path (str): Path of the task creating the writes.
         """
         query = (
             "INSERT OR REPLACE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)"
diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py
index 9d19b5889..2c20d6f05 100644
--- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py
+++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py
@@ -398,9 +398,11 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
         ORDER BY checkpoint_id DESC"""
         if limit:
             query += f" LIMIT {limit}"
-        async with self.lock, self.conn.execute(
-            query, params
-        ) as cur, self.conn.cursor() as wcur:
+        async with (
+            self.lock,
+            self.conn.execute(query, params) as cur,
+            self.conn.cursor() as wcur,
+        ):
             async for (
                 thread_id,
                 checkpoint_ns,
@@ -467,16 +469,19 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
         checkpoint_ns = config["configurable"]["checkpoint_ns"]
         type_, serialized_checkpoint = self.serde.dumps_typed(checkpoint)
         serialized_metadata = self.jsonplus_serde.dumps(metadata)
-        async with self.lock, self.conn.execute(
-            "INSERT OR REPLACE INTO checkpoints (thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata) VALUES (?, ?, ?, ?, ?, ?, ?)",
-            (
-                str(config["configurable"]["thread_id"]),
-                checkpoint_ns,
-                checkpoint["id"],
-                config["configurable"].get("checkpoint_id"),
-                type_,
-                serialized_checkpoint,
-                serialized_metadata,
+        async with (
+            self.lock,
+            self.conn.execute(
+                "INSERT OR REPLACE INTO checkpoints (thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata) VALUES (?, ?, ?, ?, ?, ?, ?)",
+                (
+                    str(config["configurable"]["thread_id"]),
+                    checkpoint_ns,
+                    checkpoint["id"],
+                    config["configurable"].get("checkpoint_id"),
+                    type_,
+                    serialized_checkpoint,
+                    serialized_metadata,
+                ),
             ),
         ):
             await self.conn.commit()
@@ -493,6 +498,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
         config: RunnableConfig,
         writes: Sequence[Tuple[str, Any]],
         task_id: str,
+        task_path: str = "",
     ) -> None:
         """Store intermediate writes linked to a checkpoint asynchronously.
 
@@ -502,6 +508,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
             config (RunnableConfig): Configuration of the related checkpoint.
             writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair.
             task_id (str): Identifier for the task creating the writes.
+            task_path (str): Path of the task creating the writes.
         """
         query = (
             "INSERT OR REPLACE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)"