diff --git a/examples/persistence_postgres.ipynb b/examples/persistence_postgres.ipynb index ef8be17cb..f43fa01d6 100644 --- a/examples/persistence_postgres.ipynb +++ b/examples/persistence_postgres.ipynb @@ -122,7 +122,7 @@ "metadata": {}, "outputs": [], "source": [ - "DB_URI = \"postgresql://postgres:postgres@localhost:5441/postgres?sslmode=disable\"" + "DB_URI = \"postgresql://postgres:postgres@localhost:5442/postgres?sslmode=disable\"" ] }, { @@ -137,7 +137,6 @@ "connection_kwargs ={\n", " \"autocommit\": True,\n", " \"prepare_threshold\": 0,\n", - " \"row_factory\": dict_row,\n", "}" ] }, @@ -551,9 +550,9 @@ ], "metadata": { "kernelspec": { - "display_name": "langgraph-postgres", + "display_name": "langgraph", "language": "python", - "name": "langgraph-postgres" + "name": "langgraph" }, "language_info": { "codemirror_mode": { diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index 433bc231d..cb4a79c35 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -66,7 +66,7 @@ class PostgresSaver(BasePostgresSaver): the first time checkpointer is used. """ with self.lock: - with self.conn.cursor(binary=True) as cur: + with self.conn.cursor(binary=True, row_factory=dict_row) as cur: try: version = cur.execute( "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1" @@ -127,31 +127,33 @@ class PostgresSaver(BasePostgresSaver): if limit: query += f" LIMIT {limit}" # if we change this to use .stream() we need to make sure to close the cursor - for value in self.conn.execute(query, args, binary=True): - yield CheckpointTuple( - { - "configurable": { - "thread_id": value["thread_id"], - "checkpoint_ns": value["checkpoint_ns"], - "checkpoint_id": value["checkpoint_id"], + with self._cursor() as cur: + cur.execute(query, args, binary=True) + for value in cur: + yield CheckpointTuple( + { + "configurable": { + "thread_id": value["thread_id"], + "checkpoint_ns": value["checkpoint_ns"], + "checkpoint_id": value["checkpoint_id"], + } + }, + { + **self._load_checkpoint(value["checkpoint"]), + "channel_values": self._load_blobs(value["channel_values"]), + }, + self._load_metadata(value["metadata"]), + { + "configurable": { + "thread_id": value["thread_id"], + "checkpoint_ns": value["checkpoint_ns"], + "checkpoint_id": value["parent_checkpoint_id"], + } } - }, - { - **self._load_checkpoint(value["checkpoint"]), - "channel_values": self._load_blobs(value["channel_values"]), - }, - self._load_metadata(value["metadata"]), - { - "configurable": { - "thread_id": value["thread_id"], - "checkpoint_ns": value["checkpoint_ns"], - "checkpoint_id": value["parent_checkpoint_id"], - } - } - if value["parent_checkpoint_id"] - else None, - self._load_writes(value["pending_writes"]), - ) + if value["parent_checkpoint_id"] + else None, + self._load_writes(value["pending_writes"]), + ) def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: """Get a checkpoint tuple from the database. @@ -199,7 +201,7 @@ class PostgresSaver(BasePostgresSaver): where = "WHERE thread_id = %s AND checkpoint_ns = %s ORDER BY checkpoint_id DESC LIMIT 1" with self._cursor() as cur: - cur = self.conn.execute( + cur.execute( self.SELECT_SQL + where, args, binary=True, @@ -336,7 +338,7 @@ class PostgresSaver(BasePostgresSaver): # in multiple threads/coroutines, but only one cursor can be # used at a time try: - with self.conn.cursor(binary=True) as cur: + with self.conn.cursor(binary=True, row_factory=dict_row) as cur: yield cur finally: if pipeline: @@ -344,8 +346,10 @@ class PostgresSaver(BasePostgresSaver): elif pipeline: # a connection not in pipeline mode can only be used by one # thread/coroutine at a time, so we acquire a lock - with self.lock, self.conn.pipeline(), self.conn.cursor(binary=True) as cur: + with self.lock, self.conn.pipeline(), self.conn.cursor( + binary=True, row_factory=dict_row + ) as cur: yield cur else: - with self.lock, self.conn.cursor(binary=True) as cur: + with self.lock, self.conn.cursor(binary=True, row_factory=dict_row) as cur: yield cur diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index 79ae1ddf7..569159d91 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -64,7 +64,7 @@ class AsyncPostgresSaver(BasePostgresSaver): the first time checkpointer is used. """ async with self.lock: - async with self.conn.cursor(binary=True) as cur: + async with self.conn.cursor(binary=True, row_factory=dict_row) as cur: try: results = await cur.execute( "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1" @@ -110,33 +110,35 @@ class AsyncPostgresSaver(BasePostgresSaver): if limit: query += f" LIMIT {limit}" # if we change this to use .stream() we need to make sure to close the cursor - async for value in await self.conn.execute(query, args, binary=True): - yield CheckpointTuple( - { - "configurable": { - "thread_id": value["thread_id"], - "checkpoint_ns": value["checkpoint_ns"], - "checkpoint_id": value["checkpoint_id"], + async with self._cursor() as cur: + await cur.execute(query, args, binary=True) + async for value in cur: + yield CheckpointTuple( + { + "configurable": { + "thread_id": value["thread_id"], + "checkpoint_ns": value["checkpoint_ns"], + "checkpoint_id": value["checkpoint_id"], + } + }, + { + **self._load_checkpoint(value["checkpoint"]), + "channel_values": await asyncio.to_thread( + self._load_blobs, value["channel_values"] + ), + }, + self._load_metadata(value["metadata"]), + { + "configurable": { + "thread_id": value["thread_id"], + "checkpoint_ns": value["checkpoint_ns"], + "checkpoint_id": value["parent_checkpoint_id"], + } } - }, - { - **self._load_checkpoint(value["checkpoint"]), - "channel_values": await asyncio.to_thread( - self._load_blobs, value["channel_values"] - ), - }, - self._load_metadata(value["metadata"]), - { - "configurable": { - "thread_id": value["thread_id"], - "checkpoint_ns": value["checkpoint_ns"], - "checkpoint_id": value["parent_checkpoint_id"], - } - } - if value["parent_checkpoint_id"] - else None, - await asyncio.to_thread(self._load_writes, value["pending_writes"]), - ) + if value["parent_checkpoint_id"] + else None, + await asyncio.to_thread(self._load_writes, value["pending_writes"]), + ) async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: """Get a checkpoint tuple from the database asynchronously. @@ -163,7 +165,7 @@ class AsyncPostgresSaver(BasePostgresSaver): where = "WHERE thread_id = %s AND checkpoint_ns = %s ORDER BY checkpoint_id DESC LIMIT 1" async with self._cursor() as cur: - cur = await self.conn.execute( + await cur.execute( self.SELECT_SQL + where, args, binary=True, @@ -293,7 +295,7 @@ class AsyncPostgresSaver(BasePostgresSaver): # in multiple threads/coroutines, but only one cursor can be # used at a time try: - async with self.conn.cursor(binary=True) as cur: + async with self.conn.cursor(binary=True, row_factory=dict_row) as cur: yield cur finally: if pipeline: @@ -302,9 +304,11 @@ class AsyncPostgresSaver(BasePostgresSaver): # a connection not in pipeline mode can only be used by one # thread/coroutine at a time, so we acquire a lock async with self.lock, self.conn.pipeline(), self.conn.cursor( - binary=True + binary=True, row_factory=dict_row ) as cur: yield cur else: - async with self.lock, self.conn.cursor(binary=True) as cur: + async with self.lock, self.conn.cursor( + binary=True, row_factory=dict_row + ) as cur: yield cur