langgraph-checkpoint-postgres 1.0.1__tar.gz → 1.0.2__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.2
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,9 @@ 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
+
26
29
  ```python
27
30
  from langgraph.checkpoint.postgres import PostgresSaver
28
31
 
@@ -31,6 +34,8 @@ read_config = {"configurable": {"thread_id": "1"}}
31
34
 
32
35
  DB_URI = "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable"
33
36
  with PostgresSaver.from_conn_string(DB_URI) as checkpointer:
37
+ # call .setup() the first time you're using the checkpointer
38
+ checkpointer.setup()
34
39
  checkpoint = {
35
40
  "v": 1,
36
41
  "ts": "2024-07-31T20:14:19.804150+00:00",
@@ -4,6 +4,9 @@ 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
+
7
10
  ```python
8
11
  from langgraph.checkpoint.postgres import PostgresSaver
9
12
 
@@ -12,6 +15,8 @@ read_config = {"configurable": {"thread_id": "1"}}
12
15
 
13
16
  DB_URI = "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable"
14
17
  with PostgresSaver.from_conn_string(DB_URI) as checkpointer:
18
+ # call .setup() the first time you're using the checkpointer
19
+ checkpointer.setup()
15
20
  checkpoint = {
16
21
  "v": 1,
17
22
  "ts": "2024-07-31T20:14:19.804150+00:00",
@@ -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,9 +59,9 @@ 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:
@@ -92,6 +92,20 @@ class AsyncPostgresSaver(BasePostgresSaver):
92
92
  before: Optional[RunnableConfig] = None,
93
93
  limit: Optional[int] = None,
94
94
  ) -> AsyncIterator[CheckpointTuple]:
95
+ """List checkpoints from the database asynchronously.
96
+
97
+ This method retrieves a list of checkpoint tuples from the Postgres database based
98
+ on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first).
99
+
100
+ Args:
101
+ config (Optional[RunnableConfig]): Base configuration for filtering checkpoints.
102
+ filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata.
103
+ before (Optional[RunnableConfig]): If provided, only checkpoints before the specified checkpoint ID are returned. Defaults to None.
104
+ limit (Optional[int]): Maximum number of checkpoints to return.
105
+
106
+ Yields:
107
+ AsyncIterator[CheckpointTuple]: An asynchronous iterator of matching checkpoint tuples.
108
+ """
95
109
  where, args = self._search_where(config, filter, before)
96
110
  query = self.SELECT_SQL + where + " ORDER BY checkpoint_id DESC"
97
111
  if limit:
@@ -125,6 +139,19 @@ class AsyncPostgresSaver(BasePostgresSaver):
125
139
  )
126
140
 
127
141
  async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
142
+ """Get a checkpoint tuple from the database asynchronously.
143
+
144
+ This method retrieves a checkpoint tuple from the Postgres database based on the
145
+ provided config. If the config contains a "checkpoint_id" key, the checkpoint with
146
+ the matching thread ID and "checkpoint_id" is retrieved. Otherwise, the latest checkpoint
147
+ for the given thread ID is retrieved.
148
+
149
+ Args:
150
+ config (RunnableConfig): The config to use for retrieving the checkpoint.
151
+
152
+ Returns:
153
+ Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found.
154
+ """
128
155
  thread_id = config["configurable"]["thread_id"]
129
156
  checkpoint_id = get_checkpoint_id(config)
130
157
  checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
@@ -177,6 +204,20 @@ class AsyncPostgresSaver(BasePostgresSaver):
177
204
  metadata: CheckpointMetadata,
178
205
  new_versions: ChannelVersions,
179
206
  ) -> RunnableConfig:
207
+ """Save a checkpoint to the database asynchronously.
208
+
209
+ This method saves a checkpoint to the Postgres database. The checkpoint is associated
210
+ with the provided config and its parent config (if any).
211
+
212
+ Args:
213
+ config (RunnableConfig): The config to associate with the checkpoint.
214
+ checkpoint (Checkpoint): The checkpoint to save.
215
+ metadata (CheckpointMetadata): Additional metadata to save with the checkpoint.
216
+ new_versions (ChannelVersions): New channel versions as of this write.
217
+
218
+ Returns:
219
+ RunnableConfig: Updated configuration after storing the checkpoint.
220
+ """
180
221
  configurable = config["configurable"].copy()
181
222
  thread_id = configurable.pop("thread_id")
182
223
  checkpoint_ns = configurable.pop("checkpoint_ns")
@@ -201,7 +242,6 @@ class AsyncPostgresSaver(BasePostgresSaver):
201
242
  thread_id,
202
243
  checkpoint_ns,
203
244
  copy.pop("channel_values"),
204
- copy["channel_versions"],
205
245
  new_versions,
206
246
  ),
207
247
  )
@@ -224,6 +264,15 @@ class AsyncPostgresSaver(BasePostgresSaver):
224
264
  writes: list[tuple[str, Any]],
225
265
  task_id: str,
226
266
  ) -> None:
267
+ """Store intermediate writes linked to a checkpoint asynchronously.
268
+
269
+ This method saves intermediate writes associated with a checkpoint to the database.
270
+
271
+ Args:
272
+ config (RunnableConfig): Configuration of the related checkpoint.
273
+ writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair.
274
+ task_id (str): Identifier for the task creating the writes.
275
+ """
227
276
  async with self._cursor() as cur:
228
277
  await cur.executemany(
229
278
  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.2"
4
4
  description = "Library with a Postgres implementation of LangGraph checkpoint saver."
5
5
  authors = []
6
6
  license = "MIT"