diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index 1c064415a..6d5fbac8c 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -10,6 +10,7 @@ from psycopg.types.json import Jsonb from psycopg_pool import ConnectionPool from langgraph.checkpoint.base import ( + WRITES_IDX_MAP, ChannelVersions, Checkpoint, CheckpointMetadata, @@ -337,9 +338,14 @@ class PostgresSaver(BasePostgresSaver): writes (List[Tuple[str, Any]]): List of writes to store. task_id (str): Identifier for the task creating the writes. """ + query = ( + self.UPSERT_CHECKPOINT_WRITES_SQL + if all(w[0] in WRITES_IDX_MAP for w in writes) + else self.INSERT_CHECKPOINT_WRITES_SQL + ) with self._cursor(pipeline=True) as cur: cur.executemany( - self.UPSERT_CHECKPOINT_WRITES_SQL, + query, self._dump_writes( config["configurable"]["thread_id"], config["configurable"]["checkpoint_ns"], diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py index 31bffd19c..8b6f728e9 100644 --- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py +++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py @@ -435,9 +435,14 @@ class SqliteSaver(BaseCheckpointSaver): 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. """ + query = ( + "INSERT OR REPLACE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)" + if all(w[0] in WRITES_IDX_MAP for w in writes) + else "INSERT OR IGNORE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)" + ) with self.cursor() as cur: cur.executemany( - "INSERT OR IGNORE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)", + query, [ ( str(config["configurable"]["thread_id"]), diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py index 28ad76109..3ffd715c1 100644 --- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py +++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py @@ -476,10 +476,15 @@ class AsyncSqliteSaver(BaseCheckpointSaver): 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. """ + query = ( + "INSERT OR REPLACE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)" + if all(w[0] in WRITES_IDX_MAP for w in writes) + else "INSERT OR IGNORE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)" + ) await self.setup() async with self.lock, self.conn.cursor() as cur: await cur.executemany( - "INSERT OR IGNORE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)", + query, [ ( str(config["configurable"]["thread_id"]),