WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-07 10:17:50 +02:00
Code Issues Packages Projects Releases Wiki Activity

Implement changes in sqlite checkpointer

Browse Source
This commit is contained in:
Nuno Campos
2024-09-10 16:17:22 -07:00
parent 98b8595a94
commit 270559b880
3 changed files with 19 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py
+7 -1
View File
@@ -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"],
libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py
+6 -1
View File
@@ -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"]),
libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py
+6 -1
View File
@@ -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"]),