langgraph-checkpoint-postgres 1.0.1__tar.gz → 1.0.3__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: langgraph-checkpoint-postgres
3
- Version: 1.0.1
3
+ Version: 1.0.3
4
4
  Summary: Library with a Postgres implementation of LangGraph checkpoint saver.
5
5
  Home-page: https://www.github.com/langchain-ai/langgraph
6
6
  License: MIT
@@ -23,6 +23,12 @@ Implementation of LangGraph CheckpointSaver that uses Postgres.
23
23
 
24
24
  ## Usage
25
25
 
26
+ > [!IMPORTANT]
27
+ > When using Postgres checkpointers for the first time, make sure to call `.setup()` method on them to create required tables. See example below.
28
+
29
+ > [!IMPORTANT]
30
+ > When manually creating Postgres connections and passing them to `PostgresSaver` or `AsyncPostgresSaver`, make sure to include `autocommit=True` and `row_factory=dict_row` (`from psycopg.rows import dict_row`). See a full example in this [how-to guide](https://langchain-ai.github.io/langgraph/how-tos/persistence_postgres/).
31
+
26
32
  ```python
27
33
  from langgraph.checkpoint.postgres import PostgresSaver
28
34
 
@@ -31,6 +37,8 @@ read_config = {"configurable": {"thread_id": "1"}}
31
37
 
32
38
  DB_URI = "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable"
33
39
  with PostgresSaver.from_conn_string(DB_URI) as checkpointer:
40
+ # call .setup() the first time you're using the checkpointer
41
+ checkpointer.setup()
34
42
  checkpoint = {
35
43
  "v": 1,
36
44
  "ts": "2024-07-31T20:14:19.804150+00:00",
@@ -110,3 +118,4 @@ async with AsyncPostgresSaver.from_conn_string(DB_URI) as checkpointer:
110
118
  # list checkpoints
111
119
  [c async for c in checkpointer.alist(read_config)]
112
120
  ```
121
+
@@ -4,6 +4,12 @@ Implementation of LangGraph CheckpointSaver that uses Postgres.
4
4
 
5
5
  ## Usage
6
6
 
7
+ > [!IMPORTANT]
8
+ > When using Postgres checkpointers for the first time, make sure to call `.setup()` method on them to create required tables. See example below.
9
+
10
+ > [!IMPORTANT]
11
+ > When manually creating Postgres connections and passing them to `PostgresSaver` or `AsyncPostgresSaver`, make sure to include `autocommit=True` and `row_factory=dict_row` (`from psycopg.rows import dict_row`). See a full example in this [how-to guide](https://langchain-ai.github.io/langgraph/how-tos/persistence_postgres/).
12
+
7
13
  ```python
8
14
  from langgraph.checkpoint.postgres import PostgresSaver
9
15
 
@@ -12,6 +18,8 @@ read_config = {"configurable": {"thread_id": "1"}}
12
18
 
13
19
  DB_URI = "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable"
14
20
  with PostgresSaver.from_conn_string(DB_URI) as checkpointer:
21
+ # call .setup() the first time you're using the checkpointer
22
+ checkpointer.setup()
15
23
  checkpoint = {
16
24
  "v": 1,
17
25
  "ts": "2024-07-31T20:14:19.804150+00:00",
@@ -90,4 +98,4 @@ async with AsyncPostgresSaver.from_conn_string(DB_URI) as checkpointer:
90
98
 
91
99
  # list checkpoints
92
100
  [c async for c in checkpointer.alist(read_config)]
93
- ```
101
+ ```
@@ -61,9 +61,9 @@ class PostgresSaver(BasePostgresSaver):
61
61
  def setup(self) -> None:
62
62
  """Set up the checkpoint database asynchronously.
63
63
 
64
- This method creates the necessary tables in the SQLite database if they don't
65
- already exist. It is called automatically when needed and should not be called
66
- directly by the user.
64
+ This method creates the necessary tables in the Postgres database if they don't
65
+ already exist and runs database migrations. It MUST be called directly by the user
66
+ the first time checkpointer is used.
67
67
  """
68
68
  with self.lock:
69
69
  with self.conn.cursor(binary=True) as cur:
@@ -90,6 +90,38 @@ class PostgresSaver(BasePostgresSaver):
90
90
  before: Optional[RunnableConfig] = None,
91
91
  limit: Optional[int] = None,
92
92
  ) -> Iterator[CheckpointTuple]:
93
+ """List checkpoints from the database.
94
+
95
+ This method retrieves a list of checkpoint tuples from the Postgres database based
96
+ on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first).
97
+
98
+ Args:
99
+ config (RunnableConfig): The config to use for listing the checkpoints.
100
+ filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata. Defaults to None.
101
+ before (Optional[RunnableConfig]): If provided, only checkpoints before the specified checkpoint ID are returned. Defaults to None.
102
+ limit (Optional[int]): The maximum number of checkpoints to return. Defaults to None.
103
+
104
+ Yields:
105
+ Iterator[CheckpointTuple]: An iterator of checkpoint tuples.
106
+
107
+ Examples:
108
+ >>> from langgraph.checkpoint.postgres import PostgresSaver
109
+ >>> DB_URI = "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable"
110
+ >>> with PostgresSaver.from_conn_string(DB_URI) as memory:
111
+ ... # Run a graph, then list the checkpoints
112
+ >>> config = {"configurable": {"thread_id": "1"}}
113
+ >>> checkpoints = list(memory.list(config, limit=2))
114
+ >>> print(checkpoints)
115
+ [CheckpointTuple(...), CheckpointTuple(...)]
116
+
117
+ >>> config = {"configurable": {"thread_id": "1"}}
118
+ >>> before = {"configurable": {"checkpoint_id": "1ef4f797-8335-6428-8001-8a1503f9b875"}}
119
+ >>> with PostgresSaver.from_conn_string(DB_URI) as memory:
120
+ ... # Run a graph, then list the checkpoints
121
+ >>> checkpoints = list(memory.list(config, before=before))
122
+ >>> print(checkpoints)
123
+ [CheckpointTuple(...), ...]
124
+ """
93
125
  where, args = self._search_where(config, filter, before)
94
126
  query = self.SELECT_SQL + where + " ORDER BY checkpoint_id DESC"
95
127
  if limit:
@@ -121,6 +153,40 @@ class PostgresSaver(BasePostgresSaver):
121
153
  )
122
154
 
123
155
  def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
156
+ """Get a checkpoint tuple from the database.
157
+
158
+ This method retrieves a checkpoint tuple from the Postgres database based on the
159
+ provided config. If the config contains a "checkpoint_id" key, the checkpoint with
160
+ the matching thread ID and timestamp is retrieved. Otherwise, the latest checkpoint
161
+ for the given thread ID is retrieved.
162
+
163
+ Args:
164
+ config (RunnableConfig): The config to use for retrieving the checkpoint.
165
+
166
+ Returns:
167
+ Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found.
168
+
169
+ Examples:
170
+
171
+ Basic:
172
+ >>> config = {"configurable": {"thread_id": "1"}}
173
+ >>> checkpoint_tuple = memory.get_tuple(config)
174
+ >>> print(checkpoint_tuple)
175
+ CheckpointTuple(...)
176
+
177
+ With timestamp:
178
+
179
+ >>> config = {
180
+ ... "configurable": {
181
+ ... "thread_id": "1",
182
+ ... "checkpoint_ns": "",
183
+ ... "checkpoint_id": "1ef4f797-8335-6428-8001-8a1503f9b875",
184
+ ... }
185
+ ... }
186
+ >>> checkpoint_tuple = memory.get_tuple(config)
187
+ >>> print(checkpoint_tuple)
188
+ CheckpointTuple(...)
189
+ """ # noqa
124
190
  thread_id = config["configurable"]["thread_id"]
125
191
  checkpoint_id = get_checkpoint_id(config)
126
192
  checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
@@ -171,6 +237,31 @@ class PostgresSaver(BasePostgresSaver):
171
237
  metadata: CheckpointMetadata,
172
238
  new_versions: ChannelVersions,
173
239
  ) -> RunnableConfig:
240
+ """Save a checkpoint to the database.
241
+
242
+ This method saves a checkpoint to the Postgres database. The checkpoint is associated
243
+ with the provided config and its parent config (if any).
244
+
245
+ Args:
246
+ config (RunnableConfig): The config to associate with the checkpoint.
247
+ checkpoint (Checkpoint): The checkpoint to save.
248
+ metadata (CheckpointMetadata): Additional metadata to save with the checkpoint.
249
+ new_versions (ChannelVersions): New channel versions as of this write.
250
+
251
+ Returns:
252
+ RunnableConfig: Updated configuration after storing the checkpoint.
253
+
254
+ Examples:
255
+
256
+ >>> from langgraph.checkpoint.postgres import PostgresSaver
257
+ >>> DB_URI = "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable"
258
+ >>> with PostgresSaver.from_conn_string(DB_URI) as memory:
259
+ >>> config = {"configurable": {"thread_id": "1", "checkpoint_ns": ""}}
260
+ >>> checkpoint = {"ts": "2024-05-04T06:32:42.235444+00:00", "id": "1ef4f797-8335-6428-8001-8a1503f9b875", "data": {"key": "value"}}
261
+ >>> saved_config = memory.put(config, checkpoint, {"source": "input", "step": 1, "writes": {"key": "value"}}, {})
262
+ >>> print(saved_config)
263
+ {'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef4f797-8335-6428-8001-8a1503f9b875'}}
264
+ """
174
265
  configurable = config["configurable"].copy()
175
266
  thread_id = configurable.pop("thread_id")
176
267
  checkpoint_ns = configurable.pop("checkpoint_ns")
@@ -194,7 +285,6 @@ class PostgresSaver(BasePostgresSaver):
194
285
  thread_id,
195
286
  checkpoint_ns,
196
287
  copy.pop("channel_values"),
197
- copy["channel_versions"],
198
288
  new_versions,
199
289
  ),
200
290
  )
@@ -217,6 +307,15 @@ class PostgresSaver(BasePostgresSaver):
217
307
  writes: List[tuple[str, Any]],
218
308
  task_id: str,
219
309
  ) -> None:
310
+ """Store intermediate writes linked to a checkpoint.
311
+
312
+ This method saves intermediate writes associated with a checkpoint to the Postgres database.
313
+
314
+ Args:
315
+ config (RunnableConfig): Configuration of the related checkpoint.
316
+ writes (List[Tuple[str, Any]]): List of writes to store.
317
+ task_id (str): Identifier for the task creating the writes.
318
+ """
220
319
  with self._cursor() as cur:
221
320
  cur.executemany(
222
321
  self.UPSERT_CHECKPOINT_WRITES_SQL,
@@ -59,18 +59,17 @@ class AsyncPostgresSaver(BasePostgresSaver):
59
59
  async def setup(self) -> None:
60
60
  """Set up the checkpoint database asynchronously.
61
61
 
62
- This method creates the necessary tables in the SQLite database if they don't
63
- already exist. It is called automatically when needed and should not be called
64
- directly by the user.
62
+ This method creates the necessary tables in the Postgres database if they don't
63
+ already exist and runs database migrations. It MUST be called directly by the user
64
+ the first time checkpointer is used.
65
65
  """
66
66
  async with self.lock:
67
67
  async with self.conn.cursor(binary=True) as cur:
68
68
  try:
69
- version = (
70
- await cur.execute(
71
- "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
72
- )
73
- ).fetchone()["v"]
69
+ results = await cur.execute(
70
+ "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
71
+ )
72
+ version = (await results.fetchone())["v"]
74
73
  except UndefinedTable:
75
74
  version = -1
76
75
  for v, migration in zip(
@@ -92,6 +91,20 @@ class AsyncPostgresSaver(BasePostgresSaver):
92
91
  before: Optional[RunnableConfig] = None,
93
92
  limit: Optional[int] = None,
94
93
  ) -> AsyncIterator[CheckpointTuple]:
94
+ """List checkpoints from the database asynchronously.
95
+
96
+ This method retrieves a list of checkpoint tuples from the Postgres database based
97
+ on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first).
98
+
99
+ Args:
100
+ config (Optional[RunnableConfig]): Base configuration for filtering checkpoints.
101
+ filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata.
102
+ before (Optional[RunnableConfig]): If provided, only checkpoints before the specified checkpoint ID are returned. Defaults to None.
103
+ limit (Optional[int]): Maximum number of checkpoints to return.
104
+
105
+ Yields:
106
+ AsyncIterator[CheckpointTuple]: An asynchronous iterator of matching checkpoint tuples.
107
+ """
95
108
  where, args = self._search_where(config, filter, before)
96
109
  query = self.SELECT_SQL + where + " ORDER BY checkpoint_id DESC"
97
110
  if limit:
@@ -125,6 +138,19 @@ class AsyncPostgresSaver(BasePostgresSaver):
125
138
  )
126
139
 
127
140
  async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
141
+ """Get a checkpoint tuple from the database asynchronously.
142
+
143
+ This method retrieves a checkpoint tuple from the Postgres database based on the
144
+ provided config. If the config contains a "checkpoint_id" key, the checkpoint with
145
+ the matching thread ID and "checkpoint_id" is retrieved. Otherwise, the latest checkpoint
146
+ for the given thread ID is retrieved.
147
+
148
+ Args:
149
+ config (RunnableConfig): The config to use for retrieving the checkpoint.
150
+
151
+ Returns:
152
+ Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found.
153
+ """
128
154
  thread_id = config["configurable"]["thread_id"]
129
155
  checkpoint_id = get_checkpoint_id(config)
130
156
  checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
@@ -177,6 +203,20 @@ class AsyncPostgresSaver(BasePostgresSaver):
177
203
  metadata: CheckpointMetadata,
178
204
  new_versions: ChannelVersions,
179
205
  ) -> RunnableConfig:
206
+ """Save a checkpoint to the database asynchronously.
207
+
208
+ This method saves a checkpoint to the Postgres database. The checkpoint is associated
209
+ with the provided config and its parent config (if any).
210
+
211
+ Args:
212
+ config (RunnableConfig): The config to associate with the checkpoint.
213
+ checkpoint (Checkpoint): The checkpoint to save.
214
+ metadata (CheckpointMetadata): Additional metadata to save with the checkpoint.
215
+ new_versions (ChannelVersions): New channel versions as of this write.
216
+
217
+ Returns:
218
+ RunnableConfig: Updated configuration after storing the checkpoint.
219
+ """
180
220
  configurable = config["configurable"].copy()
181
221
  thread_id = configurable.pop("thread_id")
182
222
  checkpoint_ns = configurable.pop("checkpoint_ns")
@@ -201,7 +241,6 @@ class AsyncPostgresSaver(BasePostgresSaver):
201
241
  thread_id,
202
242
  checkpoint_ns,
203
243
  copy.pop("channel_values"),
204
- copy["channel_versions"],
205
244
  new_versions,
206
245
  ),
207
246
  )
@@ -224,6 +263,15 @@ class AsyncPostgresSaver(BasePostgresSaver):
224
263
  writes: list[tuple[str, Any]],
225
264
  task_id: str,
226
265
  ) -> None:
266
+ """Store intermediate writes linked to a checkpoint asynchronously.
267
+
268
+ This method saves intermediate writes associated with a checkpoint to the database.
269
+
270
+ Args:
271
+ config (RunnableConfig): Configuration of the related checkpoint.
272
+ writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair.
273
+ task_id (str): Identifier for the task creating the writes.
274
+ """
227
275
  async with self._cursor() as cur:
228
276
  await cur.executemany(
229
277
  self.UPSERT_CHECKPOINT_WRITES_SQL,
@@ -40,7 +40,7 @@ MIGRATIONS = [
40
40
  channel TEXT NOT NULL,
41
41
  version TEXT NOT NULL,
42
42
  type TEXT NOT NULL,
43
- blob BYTEA NOT NULL,
43
+ blob BYTEA,
44
44
  PRIMARY KEY (thread_id, checkpoint_ns, channel, version)
45
45
  );""",
46
46
  """CREATE TABLE IF NOT EXISTS checkpoint_writes (
@@ -54,6 +54,7 @@ MIGRATIONS = [
54
54
  blob BYTEA NOT NULL,
55
55
  PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx)
56
56
  );""",
57
+ "ALTER TABLE checkpoint_blobs ALTER COLUMN blob DROP not null;",
57
58
  ]
58
59
 
59
60
  SELECT_SQL = """
@@ -140,6 +141,7 @@ class BasePostgresSaver(BaseCheckpointSaver):
140
141
  return {
141
142
  k.decode(): self.serde.loads_typed((t.decode(), v))
142
143
  for k, t, v in blob_values
144
+ if t.decode() != "empty"
143
145
  }
144
146
 
145
147
  def _dump_blobs(
@@ -148,24 +150,23 @@ class BasePostgresSaver(BaseCheckpointSaver):
148
150
  checkpoint_ns: str,
149
151
  values: dict[str, Any],
150
152
  versions: dict[str, str],
151
- new_versions: Optional[dict[str, str]],
152
153
  ) -> list[tuple[str, str, str, str, str, bytes]]:
153
154
  if not versions:
154
155
  return []
155
156
 
156
- if new_versions:
157
- versions = new_versions
158
-
159
157
  return [
160
158
  (
161
159
  thread_id,
162
160
  checkpoint_ns,
163
161
  k,
164
162
  ver,
165
- *self.serde.dumps_typed(values[k]),
163
+ *(
164
+ self.serde.dumps_typed(values[k])
165
+ if k in values
166
+ else ("empty", None)
167
+ ),
166
168
  )
167
169
  for k, ver in versions.items()
168
- if k in values
169
170
  ]
170
171
 
171
172
  def _load_writes(
@@ -1,6 +1,6 @@
1
1
  [tool.poetry]
2
2
  name = "langgraph-checkpoint-postgres"
3
- version = "1.0.1"
3
+ version = "1.0.3"
4
4
  description = "Library with a Postgres implementation of LangGraph checkpoint saver."
5
5
  authors = []
6
6
  license = "MIT"