mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-07-19 14:33:37 -04:00
William FH
1044 lines
54 KiB
Plaintext
1044 lines
54 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "51466c8d-8ce4-4b3d-be4e-18fdbeda5f53",
|
|
"metadata": {},
|
|
"source": [
|
|
"# How to create a custom checkpointer using Postgres\n",
|
|
"\n",
|
|
"When creating LangGraph agents, you can also set them up so that they persist their state. This allows you to do things like interact with an agent multiple times and have it remember previous interactions.\n",
|
|
"\n",
|
|
"This example shows how to use `Postgres` as the backend for persisting checkpoint state.\n",
|
|
"\n",
|
|
"NOTE: this is just an example implementation. You can implement your own checkpointer using a different database or modify this one as long as it conforms to the `BaseCheckpointSaver` interface."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "ddc5cc5f-89fa-4d7a-9f2b-5b099feb6ecc",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Checkpointer implementation"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 1,
|
|
"id": "faadfb1b-cebe-4dcf-82fd-34044c380bc4",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"%%capture --no-stderr\n",
|
|
"%pip install -U psycopg psycopg-pool langgraph"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 2,
|
|
"id": "a35dba8e-5562-4803-ad80-160f53592dd7",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"\"\"\"Implementation of a langgraph checkpoint saver using Postgres.\"\"\"\n",
|
|
"from contextlib import asynccontextmanager, contextmanager\n",
|
|
"from typing import (\n",
|
|
" Any,\n",
|
|
" AsyncGenerator,\n",
|
|
" AsyncIterator,\n",
|
|
" Generator,\n",
|
|
" Optional,\n",
|
|
" Union,\n",
|
|
" Tuple,\n",
|
|
" List,\n",
|
|
" Sequence,\n",
|
|
")\n",
|
|
"\n",
|
|
"import psycopg\n",
|
|
"from langchain_core.runnables import RunnableConfig\n",
|
|
"from langgraph.checkpoint import BaseCheckpointSaver\n",
|
|
"from langgraph.serde.jsonplus import JsonPlusSerializer\n",
|
|
"from langgraph.checkpoint.base import Checkpoint, CheckpointMetadata, CheckpointTuple\n",
|
|
"from psycopg_pool import AsyncConnectionPool, ConnectionPool\n",
|
|
"\n",
|
|
"\n",
|
|
"class JsonAndBinarySerializer(JsonPlusSerializer):\n",
|
|
" def _default(self, obj):\n",
|
|
" if isinstance(obj, (bytes, bytearray)):\n",
|
|
" return self._encode_constructor_args(\n",
|
|
" obj.__class__, method=\"fromhex\", args=[obj.hex()]\n",
|
|
" )\n",
|
|
" return super()._default(obj)\n",
|
|
"\n",
|
|
" def dumps(self, obj: Any) -> tuple[str, bytes]:\n",
|
|
" if isinstance(obj, bytes):\n",
|
|
" return \"bytes\", obj\n",
|
|
" elif isinstance(obj, bytearray):\n",
|
|
" return \"bytearray\", obj\n",
|
|
"\n",
|
|
" return \"json\", super().dumps(obj)\n",
|
|
"\n",
|
|
" def loads(self, s: tuple[str, bytes]) -> Any:\n",
|
|
" if s[0] == \"bytes\":\n",
|
|
" return s[1]\n",
|
|
" elif s[0] == \"bytearray\":\n",
|
|
" return bytearray(s[1])\n",
|
|
" elif s[0] == \"json\":\n",
|
|
" return super().loads(s[1])\n",
|
|
" else:\n",
|
|
" raise NotImplementedError(f\"Unknown serialization type: {s[0]}\")\n",
|
|
"\n",
|
|
"\n",
|
|
"@contextmanager\n",
|
|
"def _get_sync_connection(\n",
|
|
" connection: Union[psycopg.Connection, ConnectionPool, None],\n",
|
|
") -> Generator[psycopg.Connection, None, None]:\n",
|
|
" \"\"\"Get the connection to the Postgres database.\"\"\"\n",
|
|
" if isinstance(connection, psycopg.Connection):\n",
|
|
" yield connection\n",
|
|
" elif isinstance(connection, ConnectionPool):\n",
|
|
" with connection.connection() as conn:\n",
|
|
" yield conn\n",
|
|
" else:\n",
|
|
" raise ValueError(\n",
|
|
" \"Invalid sync connection object. Please initialize the check pointer \"\n",
|
|
" f\"with an appropriate sync connection object. \"\n",
|
|
" f\"Got {type(connection)}.\"\n",
|
|
" )\n",
|
|
"\n",
|
|
"\n",
|
|
"@asynccontextmanager\n",
|
|
"async def _get_async_connection(\n",
|
|
" connection: Union[psycopg.AsyncConnection, AsyncConnectionPool, None],\n",
|
|
") -> AsyncGenerator[psycopg.AsyncConnection, None]:\n",
|
|
" \"\"\"Get the connection to the Postgres database.\"\"\"\n",
|
|
" if isinstance(connection, psycopg.AsyncConnection):\n",
|
|
" yield connection\n",
|
|
" elif isinstance(connection, AsyncConnectionPool):\n",
|
|
" async with connection.connection() as conn:\n",
|
|
" yield conn\n",
|
|
" else:\n",
|
|
" raise ValueError(\n",
|
|
" \"Invalid async connection object. Please initialize the check pointer \"\n",
|
|
" f\"with an appropriate async connection object. \"\n",
|
|
" f\"Got {type(connection)}.\"\n",
|
|
" )\n",
|
|
"\n",
|
|
"\n",
|
|
"class PostgresSaver(BaseCheckpointSaver):\n",
|
|
" sync_connection: Optional[Union[psycopg.Connection, ConnectionPool]] = None\n",
|
|
" \"\"\"The synchronous connection or pool to the Postgres database.\n",
|
|
" \n",
|
|
" If providing a connection object, please ensure that the connection is open\n",
|
|
" and remember to close the connection when done.\n",
|
|
" \"\"\"\n",
|
|
" async_connection: Optional[\n",
|
|
" Union[psycopg.AsyncConnection, AsyncConnectionPool]\n",
|
|
" ] = None\n",
|
|
" \"\"\"The asynchronous connection or pool to the Postgres database.\n",
|
|
" \n",
|
|
" If providing a connection object, please ensure that the connection is open\n",
|
|
" and remember to close the connection when done.\n",
|
|
" \"\"\"\n",
|
|
"\n",
|
|
" def __init__(\n",
|
|
" self,\n",
|
|
" sync_connection: Optional[Union[psycopg.Connection, ConnectionPool]] = None,\n",
|
|
" async_connection: Optional[\n",
|
|
" Union[psycopg.AsyncConnection, AsyncConnectionPool]\n",
|
|
" ] = None,\n",
|
|
" ):\n",
|
|
" super().__init__(serde=JsonPlusSerializer())\n",
|
|
" self.sync_connection = sync_connection\n",
|
|
" self.async_connection = async_connection\n",
|
|
"\n",
|
|
" @contextmanager\n",
|
|
" def _get_sync_connection(self) -> Generator[psycopg.Connection, None, None]:\n",
|
|
" \"\"\"Get the connection to the Postgres database.\"\"\"\n",
|
|
" with _get_sync_connection(self.sync_connection) as connection:\n",
|
|
" yield connection\n",
|
|
"\n",
|
|
" @asynccontextmanager\n",
|
|
" async def _get_async_connection(\n",
|
|
" self,\n",
|
|
" ) -> AsyncGenerator[psycopg.AsyncConnection, None]:\n",
|
|
" \"\"\"Get the connection to the Postgres database.\"\"\"\n",
|
|
" async with _get_async_connection(self.async_connection) as connection:\n",
|
|
" yield connection\n",
|
|
"\n",
|
|
" CREATE_TABLES_QUERY = \"\"\"\n",
|
|
" CREATE TABLE IF NOT EXISTS checkpoints (\n",
|
|
" thread_id TEXT NOT NULL,\n",
|
|
" thread_ts TEXT NOT NULL,\n",
|
|
" parent_ts TEXT,\n",
|
|
" checkpoint BYTEA NOT NULL,\n",
|
|
" metadata BYTEA NOT NULL,\n",
|
|
" PRIMARY KEY (thread_id, thread_ts)\n",
|
|
" );\n",
|
|
" CREATE TABLE IF NOT EXISTS writes (\n",
|
|
" thread_id TEXT NOT NULL,\n",
|
|
" thread_ts TEXT NOT NULL,\n",
|
|
" task_id TEXT NOT NULL,\n",
|
|
" idx INTEGER NOT NULL,\n",
|
|
" channel TEXT NOT NULL,\n",
|
|
" value BYTEA,\n",
|
|
" PRIMARY KEY (thread_id, thread_ts, task_id, idx)\n",
|
|
" );\n",
|
|
" \"\"\"\n",
|
|
"\n",
|
|
" @staticmethod\n",
|
|
" def create_tables(connection: Union[psycopg.Connection, ConnectionPool], /) -> None:\n",
|
|
" \"\"\"Create the schema for the checkpoint saver.\"\"\"\n",
|
|
" with _get_sync_connection(connection) as conn:\n",
|
|
" with conn.cursor() as cur:\n",
|
|
" cur.execute(PostgresSaver.CREATE_TABLES_QUERY)\n",
|
|
"\n",
|
|
" @staticmethod\n",
|
|
" async def acreate_tables(\n",
|
|
" connection: Union[psycopg.AsyncConnection, AsyncConnectionPool], /\n",
|
|
" ) -> None:\n",
|
|
" \"\"\"Create the schema for the checkpoint saver.\"\"\"\n",
|
|
" async with _get_async_connection(connection) as conn:\n",
|
|
" async with conn.cursor() as cur:\n",
|
|
" await cur.execute(PostgresSaver.CREATE_TABLES_QUERY)\n",
|
|
"\n",
|
|
" @staticmethod\n",
|
|
" def drop_tables(connection: psycopg.Connection, /) -> None:\n",
|
|
" \"\"\"Drop the table for the checkpoint saver.\"\"\"\n",
|
|
" with connection.cursor() as cur:\n",
|
|
" cur.execute(\"DROP TABLE IF EXISTS checkpoints, writes;\")\n",
|
|
"\n",
|
|
" @staticmethod\n",
|
|
" async def adrop_tables(connection: psycopg.AsyncConnection, /) -> None:\n",
|
|
" \"\"\"Drop the table for the checkpoint saver.\"\"\"\n",
|
|
" async with connection.cursor() as cur:\n",
|
|
" await cur.execute(\"DROP TABLE IF EXISTS checkpoints, writes;\")\n",
|
|
"\n",
|
|
" UPSERT_CHECKPOINT_QUERY = \"\"\"\n",
|
|
" INSERT INTO checkpoints \n",
|
|
" (thread_id, thread_ts, parent_ts, checkpoint, metadata)\n",
|
|
" VALUES \n",
|
|
" (%s, %s, %s, %s, %s)\n",
|
|
" ON CONFLICT (thread_id, thread_ts)\n",
|
|
" DO UPDATE SET checkpoint = EXCLUDED.checkpoint,\n",
|
|
" metadata = EXCLUDED.metadata;\n",
|
|
" \"\"\"\n",
|
|
"\n",
|
|
" def put(\n",
|
|
" self,\n",
|
|
" config: RunnableConfig,\n",
|
|
" checkpoint: Checkpoint,\n",
|
|
" metadata: CheckpointMetadata,\n",
|
|
" ) -> RunnableConfig:\n",
|
|
" \"\"\"Put the checkpoint for the given configuration.\n",
|
|
" Args:\n",
|
|
" config: The configuration for the checkpoint.\n",
|
|
" A dict with a `configurable` key which is a dict with\n",
|
|
" a `thread_id` key and an optional `thread_ts` key.\n",
|
|
" For example, { 'configurable': { 'thread_id': 'test_thread' } }\n",
|
|
" checkpoint: The checkpoint to persist.\n",
|
|
" Returns:\n",
|
|
" The RunnableConfig that describes the checkpoint that was just created.\n",
|
|
" It'll contain the `thread_id` and `thread_ts` of the checkpoint.\n",
|
|
" \"\"\"\n",
|
|
" thread_id = config[\"configurable\"][\"thread_id\"]\n",
|
|
" parent_ts = config[\"configurable\"].get(\"thread_ts\")\n",
|
|
" with self._get_sync_connection() as conn:\n",
|
|
" with conn.cursor() as cur:\n",
|
|
" cur.execute(\n",
|
|
" self.UPSERT_CHECKPOINT_QUERY,\n",
|
|
" (\n",
|
|
" thread_id,\n",
|
|
" checkpoint[\"id\"],\n",
|
|
" parent_ts if parent_ts else None,\n",
|
|
" self.serde.dumps(checkpoint),\n",
|
|
" self.serde.dumps(metadata),\n",
|
|
" ),\n",
|
|
" )\n",
|
|
"\n",
|
|
" return {\n",
|
|
" \"configurable\": {\n",
|
|
" \"thread_id\": thread_id,\n",
|
|
" \"thread_ts\": checkpoint[\"id\"],\n",
|
|
" },\n",
|
|
" }\n",
|
|
"\n",
|
|
" async def aput(\n",
|
|
" self,\n",
|
|
" config: RunnableConfig,\n",
|
|
" checkpoint: Checkpoint,\n",
|
|
" metadata: CheckpointMetadata,\n",
|
|
" ) -> RunnableConfig:\n",
|
|
" \"\"\"Put the checkpoint for the given configuration.\n",
|
|
" Args:\n",
|
|
" config: The configuration for the checkpoint.\n",
|
|
" A dict with a `configurable` key which is a dict with\n",
|
|
" a `thread_id` key and an optional `thread_ts` key.\n",
|
|
" For example, { 'configurable': { 'thread_id': 'test_thread' } }\n",
|
|
" checkpoint: The checkpoint to persist.\n",
|
|
" Returns:\n",
|
|
" The RunnableConfig that describes the checkpoint that was just created.\n",
|
|
" It'll contain the `thread_id` and `thread_ts` of the checkpoint.\n",
|
|
" \"\"\"\n",
|
|
" thread_id = config[\"configurable\"][\"thread_id\"]\n",
|
|
" parent_ts = config[\"configurable\"].get(\"thread_ts\")\n",
|
|
" async with self._get_async_connection() as conn:\n",
|
|
" async with conn.cursor() as cur:\n",
|
|
" await cur.execute(\n",
|
|
" self.UPSERT_CHECKPOINT_QUERY,\n",
|
|
" (\n",
|
|
" thread_id,\n",
|
|
" checkpoint[\"id\"],\n",
|
|
" parent_ts if parent_ts else None,\n",
|
|
" self.serde.dumps(checkpoint),\n",
|
|
" self.serde.dumps(metadata),\n",
|
|
" ),\n",
|
|
" )\n",
|
|
"\n",
|
|
" return {\n",
|
|
" \"configurable\": {\n",
|
|
" \"thread_id\": thread_id,\n",
|
|
" \"thread_ts\": checkpoint[\"id\"],\n",
|
|
" },\n",
|
|
" }\n",
|
|
"\n",
|
|
" UPSERT_WRITES_QUERY = \"\"\"\n",
|
|
" INSERT INTO writes\n",
|
|
" (thread_id, thread_ts, task_id, idx, channel, value)\n",
|
|
" VALUES\n",
|
|
" (%s, %s, %s, %s, %s, %s)\n",
|
|
" ON CONFLICT (thread_id, thread_ts, task_id, idx)\n",
|
|
" DO UPDATE SET value = EXCLUDED.value;\n",
|
|
" \"\"\"\n",
|
|
"\n",
|
|
" def put_writes(\n",
|
|
" self,\n",
|
|
" config: RunnableConfig,\n",
|
|
" writes: Sequence[Tuple[str, Any]],\n",
|
|
" task_id: str,\n",
|
|
" ) -> None:\n",
|
|
" with self._get_sync_connection() as conn:\n",
|
|
" with conn.cursor() as cur:\n",
|
|
" cur.executemany(\n",
|
|
" self.UPSERT_WRITES_QUERY,\n",
|
|
" [\n",
|
|
" (\n",
|
|
" str(config[\"configurable\"][\"thread_id\"]),\n",
|
|
" str(config[\"configurable\"][\"thread_ts\"]),\n",
|
|
" task_id,\n",
|
|
" idx,\n",
|
|
" channel,\n",
|
|
" self.serde.dumps(value),\n",
|
|
" )\n",
|
|
" for idx, (channel, value) in enumerate(writes)\n",
|
|
" ],\n",
|
|
" )\n",
|
|
" conn.commit()\n",
|
|
"\n",
|
|
" async def aput_writes(\n",
|
|
" self,\n",
|
|
" config: RunnableConfig,\n",
|
|
" writes: Sequence[Tuple[str, Any]],\n",
|
|
" task_id: str,\n",
|
|
" ) -> None:\n",
|
|
" async with self._get_async_connection() as conn:\n",
|
|
" async with conn.cursor() as cur:\n",
|
|
" await cur.executemany(\n",
|
|
" self.UPSERT_WRITES_QUERY,\n",
|
|
" [\n",
|
|
" (\n",
|
|
" str(config[\"configurable\"][\"thread_id\"]),\n",
|
|
" str(config[\"configurable\"][\"thread_ts\"]),\n",
|
|
" task_id,\n",
|
|
" idx,\n",
|
|
" channel,\n",
|
|
" self.serde.dumps(value),\n",
|
|
" )\n",
|
|
" for idx, (channel, value) in enumerate(writes)\n",
|
|
" ],\n",
|
|
" )\n",
|
|
" await conn.commit()\n",
|
|
"\n",
|
|
" LIST_CHECKPOINTS_QUERY_STR = \"\"\"\n",
|
|
" SELECT checkpoint, metadata, thread_ts, parent_ts\n",
|
|
" FROM checkpoints\n",
|
|
" {where}\n",
|
|
" ORDER BY thread_ts DESC\n",
|
|
" \"\"\"\n",
|
|
"\n",
|
|
" def list(\n",
|
|
" self,\n",
|
|
" config: Optional[RunnableConfig],\n",
|
|
" *,\n",
|
|
" filter: Optional[dict[str, Any]] = None,\n",
|
|
" before: Optional[RunnableConfig] = None,\n",
|
|
" limit: Optional[int] = None,\n",
|
|
" ) -> Generator[CheckpointTuple, None, None]:\n",
|
|
" \"\"\"Get all the checkpoints for the given configuration.\"\"\"\n",
|
|
" where, args = self._search_where(config, filter, before)\n",
|
|
" query = self.LIST_CHECKPOINTS_QUERY_STR.format(where=where)\n",
|
|
" if limit:\n",
|
|
" query += f\" LIMIT {limit}\"\n",
|
|
" with self._get_sync_connection() as conn:\n",
|
|
" with conn.cursor() as cur:\n",
|
|
" thread_id = config[\"configurable\"][\"thread_id\"]\n",
|
|
" cur.execute(query, tuple(args))\n",
|
|
" for value in cur:\n",
|
|
" checkpoint, metadata, thread_ts, parent_ts = value\n",
|
|
" yield CheckpointTuple(\n",
|
|
" config={\n",
|
|
" \"configurable\": {\n",
|
|
" \"thread_id\": thread_id,\n",
|
|
" \"thread_ts\": thread_ts,\n",
|
|
" }\n",
|
|
" },\n",
|
|
" checkpoint=self.serde.loads(checkpoint),\n",
|
|
" metadata=self.serde.loads(metadata),\n",
|
|
" parent_config={\n",
|
|
" \"configurable\": {\n",
|
|
" \"thread_id\": thread_id,\n",
|
|
" \"thread_ts\": thread_ts,\n",
|
|
" }\n",
|
|
" }\n",
|
|
" if parent_ts\n",
|
|
" else None,\n",
|
|
" )\n",
|
|
"\n",
|
|
" async def alist(\n",
|
|
" self,\n",
|
|
" config: Optional[RunnableConfig],\n",
|
|
" *,\n",
|
|
" filter: Optional[dict[str, Any]] = None,\n",
|
|
" before: Optional[RunnableConfig] = None,\n",
|
|
" limit: Optional[int] = None,\n",
|
|
" ) -> AsyncIterator[CheckpointTuple]:\n",
|
|
" \"\"\"Get all the checkpoints for the given configuration.\"\"\"\n",
|
|
" where, args = self._search_where(config, filter, before)\n",
|
|
" query = self.LIST_CHECKPOINTS_QUERY_STR.format(where=where)\n",
|
|
" if limit:\n",
|
|
" query += f\" LIMIT {limit}\"\n",
|
|
" async with self._get_async_connection() as conn:\n",
|
|
" async with conn.cursor() as cur:\n",
|
|
" thread_id = config[\"configurable\"][\"thread_id\"]\n",
|
|
" await cur.execute(query, tuple(args))\n",
|
|
" async for value in cur:\n",
|
|
" checkpoint, metadata, thread_ts, parent_ts = value\n",
|
|
" yield CheckpointTuple(\n",
|
|
" config={\n",
|
|
" \"configurable\": {\n",
|
|
" \"thread_id\": thread_id,\n",
|
|
" \"thread_ts\": thread_ts,\n",
|
|
" }\n",
|
|
" },\n",
|
|
" checkpoint=self.serde.loads(checkpoint),\n",
|
|
" metadata=self.serde.loads(metadata),\n",
|
|
" parent_config={\n",
|
|
" \"configurable\": {\n",
|
|
" \"thread_id\": thread_id,\n",
|
|
" \"thread_ts\": thread_ts,\n",
|
|
" }\n",
|
|
" }\n",
|
|
" if parent_ts\n",
|
|
" else None,\n",
|
|
" )\n",
|
|
"\n",
|
|
" GET_CHECKPOINT_BY_TS_QUERY = \"\"\"\n",
|
|
" SELECT checkpoint, metadata, thread_ts, parent_ts\n",
|
|
" FROM checkpoints\n",
|
|
" WHERE thread_id = %(thread_id)s AND thread_ts = %(thread_ts)s\n",
|
|
" \"\"\"\n",
|
|
"\n",
|
|
" GET_CHECKPOINT_QUERY = \"\"\"\n",
|
|
" SELECT checkpoint, metadata, thread_ts, parent_ts\n",
|
|
" FROM checkpoints\n",
|
|
" WHERE thread_id = %(thread_id)s\n",
|
|
" ORDER BY thread_ts DESC LIMIT 1\n",
|
|
" \"\"\"\n",
|
|
"\n",
|
|
" def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:\n",
|
|
" \"\"\"Get the checkpoint tuple for the given configuration.\n",
|
|
" Args:\n",
|
|
" config: The configuration for the checkpoint.\n",
|
|
" A dict with a `configurable` key which is a dict with\n",
|
|
" a `thread_id` key and an optional `thread_ts` key.\n",
|
|
" For example, { 'configurable': { 'thread_id': 'test_thread' } }\n",
|
|
" Returns:\n",
|
|
" The checkpoint tuple for the given configuration if it exists,\n",
|
|
" otherwise None.\n",
|
|
" If thread_ts is None, the latest checkpoint is returned if it exists.\n",
|
|
" \"\"\"\n",
|
|
" thread_id = config[\"configurable\"][\"thread_id\"]\n",
|
|
" thread_ts = config[\"configurable\"].get(\"thread_ts\")\n",
|
|
" with self._get_sync_connection() as conn:\n",
|
|
" with conn.cursor() as cur:\n",
|
|
" # find the latest checkpoint for the thread_id\n",
|
|
" if thread_ts:\n",
|
|
" cur.execute(\n",
|
|
" self.GET_CHECKPOINT_BY_TS_QUERY,\n",
|
|
" {\n",
|
|
" \"thread_id\": thread_id,\n",
|
|
" \"thread_ts\": thread_ts,\n",
|
|
" },\n",
|
|
" )\n",
|
|
" else:\n",
|
|
" cur.execute(\n",
|
|
" self.GET_CHECKPOINT_QUERY,\n",
|
|
" {\n",
|
|
" \"thread_id\": thread_id,\n",
|
|
" },\n",
|
|
" )\n",
|
|
"\n",
|
|
" # if a checkpoint is found, return it\n",
|
|
" if value := cur.fetchone():\n",
|
|
" checkpoint, metadata, thread_ts, parent_ts = value\n",
|
|
" if not config[\"configurable\"].get(\"thread_ts\"):\n",
|
|
" config = {\n",
|
|
" \"configurable\": {\n",
|
|
" \"thread_id\": thread_id,\n",
|
|
" \"thread_ts\": thread_ts,\n",
|
|
" }\n",
|
|
" }\n",
|
|
"\n",
|
|
" # find any pending writes\n",
|
|
" cur.execute(\n",
|
|
" \"SELECT task_id, channel, value FROM writes WHERE thread_id = %(thread_id)s AND thread_ts = %(thread_ts)s\",\n",
|
|
" {\n",
|
|
" \"thread_id\": thread_id,\n",
|
|
" \"thread_ts\": thread_ts,\n",
|
|
" },\n",
|
|
" )\n",
|
|
" # deserialize the checkpoint and metadata\n",
|
|
" return CheckpointTuple(\n",
|
|
" config=config,\n",
|
|
" checkpoint=self.serde.loads(checkpoint),\n",
|
|
" metadata=self.serde.loads(metadata),\n",
|
|
" parent_config={\n",
|
|
" \"configurable\": {\n",
|
|
" \"thread_id\": thread_id,\n",
|
|
" \"thread_ts\": parent_ts,\n",
|
|
" }\n",
|
|
" }\n",
|
|
" if parent_ts\n",
|
|
" else None,\n",
|
|
" pending_writes=[\n",
|
|
" (task_id, channel, self.serde.loads(value))\n",
|
|
" for task_id, channel, value in cur\n",
|
|
" ],\n",
|
|
" )\n",
|
|
"\n",
|
|
" async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:\n",
|
|
" \"\"\"Get the checkpoint tuple for the given configuration.\n",
|
|
" Args:\n",
|
|
" config: The configuration for the checkpoint.\n",
|
|
" A dict with a `configurable` key which is a dict with\n",
|
|
" a `thread_id` key and an optional `thread_ts` key.\n",
|
|
" For example, { 'configurable': { 'thread_id': 'test_thread' } }\n",
|
|
" Returns:\n",
|
|
" The checkpoint tuple for the given configuration if it exists,\n",
|
|
" otherwise None.\n",
|
|
" If thread_ts is None, the latest checkpoint is returned if it exists.\n",
|
|
" \"\"\"\n",
|
|
" thread_id = config[\"configurable\"][\"thread_id\"]\n",
|
|
" thread_ts = config[\"configurable\"].get(\"thread_ts\")\n",
|
|
" async with self._get_async_connection() as conn:\n",
|
|
" async with conn.cursor() as cur:\n",
|
|
" # find the latest checkpoint for the thread_id\n",
|
|
" if thread_ts:\n",
|
|
" await cur.execute(\n",
|
|
" self.GET_CHECKPOINT_BY_TS_QUERY,\n",
|
|
" {\n",
|
|
" \"thread_id\": thread_id,\n",
|
|
" \"thread_ts\": thread_ts,\n",
|
|
" },\n",
|
|
" )\n",
|
|
" else:\n",
|
|
" await cur.execute(\n",
|
|
" self.GET_CHECKPOINT_QUERY,\n",
|
|
" {\n",
|
|
" \"thread_id\": thread_id,\n",
|
|
" },\n",
|
|
" )\n",
|
|
" # if a checkpoint is found, return it\n",
|
|
" if value := await cur.fetchone():\n",
|
|
" checkpoint, metadata, thread_ts, parent_ts = value\n",
|
|
" if not config[\"configurable\"].get(\"thread_ts\"):\n",
|
|
" config = {\n",
|
|
" \"configurable\": {\n",
|
|
" \"thread_id\": thread_id,\n",
|
|
" \"thread_ts\": thread_ts,\n",
|
|
" }\n",
|
|
" }\n",
|
|
"\n",
|
|
" # find any pending writes\n",
|
|
" await cur.execute(\n",
|
|
" \"SELECT task_id, channel, value FROM writes WHERE thread_id = %(thread_id)s AND thread_ts = %(thread_ts)s\",\n",
|
|
" {\n",
|
|
" \"thread_id\": thread_id,\n",
|
|
" \"thread_ts\": thread_ts,\n",
|
|
" },\n",
|
|
" )\n",
|
|
" # deserialize the checkpoint and metadata\n",
|
|
" return CheckpointTuple(\n",
|
|
" config=config,\n",
|
|
" checkpoint=self.serde.loads(checkpoint),\n",
|
|
" metadata=self.serde.loads(metadata),\n",
|
|
" parent_config={\n",
|
|
" \"configurable\": {\n",
|
|
" \"thread_id\": thread_id,\n",
|
|
" \"thread_ts\": parent_ts,\n",
|
|
" }\n",
|
|
" }\n",
|
|
" if parent_ts\n",
|
|
" else None,\n",
|
|
" pending_writes=[\n",
|
|
" (task_id, channel, self.serde.loads(value))\n",
|
|
" async for task_id, channel, value in cur\n",
|
|
" ],\n",
|
|
" )\n",
|
|
"\n",
|
|
" def _search_where(\n",
|
|
" self,\n",
|
|
" config: Optional[RunnableConfig],\n",
|
|
" filter: Optional[dict[str, Any]] = None,\n",
|
|
" before: Optional[RunnableConfig] = None,\n",
|
|
" ) -> Tuple[str, List[Any]]:\n",
|
|
" \"\"\"Return WHERE clause predicates for given config, filter, and before parameters.\n",
|
|
" Args:\n",
|
|
" config (Optional[RunnableConfig]): The config to use for filtering.\n",
|
|
" filter (Optional[Dict[str, Any]]): Additional filtering criteria.\n",
|
|
" before (Optional[RunnableConfig]): A config to limit results before a certain timestamp.\n",
|
|
" Returns:\n",
|
|
" Tuple[str, Sequence[Any]]: A tuple containing the WHERE clause and parameter values.\n",
|
|
" \"\"\"\n",
|
|
" wheres = []\n",
|
|
" param_values = []\n",
|
|
"\n",
|
|
" # Add predicate for config\n",
|
|
" if config is not None:\n",
|
|
" wheres.append(\"thread_id = %s \")\n",
|
|
" param_values.append(config[\"configurable\"][\"thread_id\"])\n",
|
|
"\n",
|
|
" if filter:\n",
|
|
" raise NotImplementedError()\n",
|
|
"\n",
|
|
" # Add predicate for limiting results before a certain timestamp\n",
|
|
" if before is not None:\n",
|
|
" wheres.append(\"thread_ts < %s\")\n",
|
|
" param_values.append(before[\"configurable\"][\"thread_ts\"])\n",
|
|
"\n",
|
|
" where_clause = \"WHERE \" + \" AND \".join(wheres) if wheres else \"\"\n",
|
|
" return where_clause, param_values"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "456fa19c-93a5-4750-a410-f2d810b964ad",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Setup environment"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 3,
|
|
"id": "eca9aafb-a155-407a-8036-682a2f1297d7",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"OPENAI_API_KEY: ········\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"import getpass\n",
|
|
"import os\n",
|
|
"\n",
|
|
"\n",
|
|
"def _set_env(var: str):\n",
|
|
" if not os.environ.get(var):\n",
|
|
" os.environ[var] = getpass.getpass(f\"{var}: \")\n",
|
|
"\n",
|
|
"\n",
|
|
"_set_env(\"OPENAI_API_KEY\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "e26b3204-cca2-414c-800e-7e09032445ae",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Setup model and tools for the graph"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 4,
|
|
"id": "e5213193-5a7d-43e7-aeba-fe732bb1cd7a",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from typing import Literal\n",
|
|
"from langchain_core.runnables import ConfigurableField\n",
|
|
"from langchain_core.tools import tool\n",
|
|
"from langchain_openai import ChatOpenAI\n",
|
|
"from langgraph.prebuilt import create_react_agent\n",
|
|
"\n",
|
|
"\n",
|
|
"@tool\n",
|
|
"def get_weather(city: Literal[\"nyc\", \"sf\"]):\n",
|
|
" \"\"\"Use this to get weather information.\"\"\"\n",
|
|
" if city == \"nyc\":\n",
|
|
" return \"It might be cloudy in nyc\"\n",
|
|
" elif city == \"sf\":\n",
|
|
" return \"It's always sunny in sf\"\n",
|
|
" else:\n",
|
|
" raise AssertionError(\"Unknown city\")\n",
|
|
"\n",
|
|
"\n",
|
|
"tools = [get_weather]\n",
|
|
"model = ChatOpenAI(model_name=\"gpt-4o\", temperature=0)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "e9342c62-dbb4-40f6-9271-7393f1ca48c4",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Use sync connection"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 5,
|
|
"id": "2b9d13b1-9d72-48a0-b63a-adc062c06c29",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"DB_URI = \"postgresql://postgres:postgres@localhost:5432/postgres?sslmode=disable\""
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "e39fc712-9e1c-4831-9077-dd07b0c13594",
|
|
"metadata": {},
|
|
"source": [
|
|
"### With a connection pool"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 6,
|
|
"id": "2657c1c4-d8a5-4fe3-8f77-95415a98ed6c",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from psycopg_pool import ConnectionPool\n",
|
|
"\n",
|
|
"pool = ConnectionPool(\n",
|
|
" # Example configuration\n",
|
|
" conninfo=DB_URI,\n",
|
|
" max_size=20,\n",
|
|
")\n",
|
|
"\n",
|
|
"checkpointer = PostgresSaver(sync_connection=pool)\n",
|
|
"checkpointer.create_tables(pool)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 7,
|
|
"id": "6d388241-de57-4b4e-af7b-eb1081fb8f36",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)\n",
|
|
"config = {\"configurable\": {\"thread_id\": \"1\"}}\n",
|
|
"res = graph.invoke({\"messages\": [(\"human\", \"what's the weather in sf\")]}, config)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 8,
|
|
"id": "a7e0e7ec-a675-470b-9270-e4bdc59d4a4d",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"text/plain": [
|
|
"{'messages': [HumanMessage(content=\"what's the weather in sf\", id='bc87fac7-1da1-4818-a43b-6ba7c9b9b3e4'),\n",
|
|
" AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_MjkmibJlXeuNchL6B8qpIjOW', 'function': {'arguments': '{\"city\":\"sf\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 14, 'prompt_tokens': 57, 'total_tokens': 71}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_d33f7b429e', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-b9de0cab-f310-4f74-897e-97014072c001-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'sf'}, 'id': 'call_MjkmibJlXeuNchL6B8qpIjOW', 'type': 'tool_call'}], usage_metadata={'input_tokens': 57, 'output_tokens': 14, 'total_tokens': 71}),\n",
|
|
" ToolMessage(content=\"It's always sunny in sf\", name='get_weather', id='8d8f9596-a683-4644-a898-1e303b5a01ea', tool_call_id='call_MjkmibJlXeuNchL6B8qpIjOW'),\n",
|
|
" AIMessage(content='The weather in San Francisco is currently sunny.', response_metadata={'token_usage': {'completion_tokens': 10, 'prompt_tokens': 84, 'total_tokens': 94}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_dd932ca5d1', 'finish_reason': 'stop', 'logprobs': None}, id='run-4b5282a3-e7a6-42ee-ad0f-e6013a745a88-0', usage_metadata={'input_tokens': 84, 'output_tokens': 10, 'total_tokens': 94})]}"
|
|
]
|
|
},
|
|
"execution_count": 8,
|
|
"metadata": {},
|
|
"output_type": "execute_result"
|
|
}
|
|
],
|
|
"source": [
|
|
"res"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 9,
|
|
"id": "96efd8b2-97c9-4207-83b2-00131723a75a",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"text/plain": [
|
|
"{'v': 1,\n",
|
|
" 'ts': '2024-07-12T15:21:51.891852+00:00',\n",
|
|
" 'id': '1ef40627-6fb2-6962-8003-b74d816658c5',\n",
|
|
" 'channel_values': {'messages': [HumanMessage(content=\"what's the weather in sf\", id='bc87fac7-1da1-4818-a43b-6ba7c9b9b3e4'),\n",
|
|
" AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_MjkmibJlXeuNchL6B8qpIjOW', 'function': {'arguments': '{\"city\":\"sf\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 14, 'prompt_tokens': 57, 'total_tokens': 71}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_d33f7b429e', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-b9de0cab-f310-4f74-897e-97014072c001-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'sf'}, 'id': 'call_MjkmibJlXeuNchL6B8qpIjOW', 'type': 'tool_call'}], usage_metadata={'input_tokens': 57, 'output_tokens': 14, 'total_tokens': 71}),\n",
|
|
" ToolMessage(content=\"It's always sunny in sf\", name='get_weather', id='8d8f9596-a683-4644-a898-1e303b5a01ea', tool_call_id='call_MjkmibJlXeuNchL6B8qpIjOW'),\n",
|
|
" AIMessage(content='The weather in San Francisco is currently sunny.', response_metadata={'token_usage': {'completion_tokens': 10, 'prompt_tokens': 84, 'total_tokens': 94}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_dd932ca5d1', 'finish_reason': 'stop', 'logprobs': None}, id='run-4b5282a3-e7a6-42ee-ad0f-e6013a745a88-0', usage_metadata={'input_tokens': 84, 'output_tokens': 10, 'total_tokens': 94})],\n",
|
|
" 'agent': 'agent'},\n",
|
|
" 'channel_versions': {'__start__': 2,\n",
|
|
" 'messages': 5,\n",
|
|
" 'start:agent': 3,\n",
|
|
" 'agent': 5,\n",
|
|
" 'branch:agent:should_continue:tools': 4,\n",
|
|
" 'tools': 5},\n",
|
|
" 'versions_seen': {'__start__': {'__start__': 1},\n",
|
|
" 'agent': {'start:agent': 3, 'tools': 4},\n",
|
|
" 'tools': {'branch:agent:should_continue:tools': 3}},\n",
|
|
" 'pending_sends': []}"
|
|
]
|
|
},
|
|
"execution_count": 9,
|
|
"metadata": {},
|
|
"output_type": "execute_result"
|
|
}
|
|
],
|
|
"source": [
|
|
"checkpointer.get(config)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "967c95c7-e392-4819-bd71-f29e91c68df3",
|
|
"metadata": {},
|
|
"source": [
|
|
"### With a connection"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 10,
|
|
"id": "180d6daf-8fa7-4608-bd2e-bfbf44ed5836",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from psycopg import Connection\n",
|
|
"\n",
|
|
"with Connection.connect(DB_URI) as conn:\n",
|
|
" checkpointer = PostgresSaver(sync_connection=conn)\n",
|
|
"\n",
|
|
" graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)\n",
|
|
" config = {\"configurable\": {\"thread_id\": \"2\"}}\n",
|
|
" res = graph.invoke({\"messages\": [(\"human\", \"what's the weather in sf\")]}, config)\n",
|
|
"\n",
|
|
" checkpoint_tuple = checkpointer.get_tuple(config)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 11,
|
|
"id": "613d0bbc-0e38-45c4-aace-1f6f7ae27c7b",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"text/plain": [
|
|
"CheckpointTuple(config={'configurable': {'thread_id': '2', 'thread_ts': '1ef40627-7d58-6422-8003-de6e83a8c293'}}, checkpoint={'v': 1, 'ts': '2024-07-12T15:21:53.322868+00:00', 'id': '1ef40627-7d58-6422-8003-de6e83a8c293', 'channel_values': {'messages': [HumanMessage(content=\"what's the weather in sf\", id='8d0209ed-a8c2-42ae-8e77-cc71a9cca29d'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_BO4zHHp0JkEWtrtaEqFHkDjK', 'function': {'arguments': '{\"city\":\"sf\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 14, 'prompt_tokens': 57, 'total_tokens': 71}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_d33f7b429e', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-1f76b350-6a33-4de7-9276-59725b1ac101-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'sf'}, 'id': 'call_BO4zHHp0JkEWtrtaEqFHkDjK', 'type': 'tool_call'}], usage_metadata={'input_tokens': 57, 'output_tokens': 14, 'total_tokens': 71}), ToolMessage(content=\"It's always sunny in sf\", name='get_weather', id='c1bb1a24-62a8-4b43-b90e-b00899c112a8', tool_call_id='call_BO4zHHp0JkEWtrtaEqFHkDjK'), AIMessage(content='The weather in San Francisco is currently sunny.', response_metadata={'token_usage': {'completion_tokens': 10, 'prompt_tokens': 84, 'total_tokens': 94}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_dd932ca5d1', 'finish_reason': 'stop', 'logprobs': None}, id='run-7576d437-4938-48b9-b2cf-e4809d92742d-0', usage_metadata={'input_tokens': 84, 'output_tokens': 10, 'total_tokens': 94})], 'agent': 'agent'}, 'channel_versions': {'__start__': 2, 'messages': 5, 'start:agent': 3, 'agent': 5, 'branch:agent:should_continue:tools': 4, 'tools': 5}, 'versions_seen': {'__start__': {'__start__': 1}, 'agent': {'start:agent': 3, 'tools': 4}, 'tools': {'branch:agent:should_continue:tools': 3}}, 'pending_sends': []}, metadata={'source': 'loop', 'step': 3, 'writes': {'agent': {'messages': [AIMessage(content='The weather in San Francisco is currently sunny.', response_metadata={'token_usage': {'completion_tokens': 10, 'prompt_tokens': 84, 'total_tokens': 94}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_dd932ca5d1', 'finish_reason': 'stop', 'logprobs': None}, id='run-7576d437-4938-48b9-b2cf-e4809d92742d-0', usage_metadata={'input_tokens': 84, 'output_tokens': 10, 'total_tokens': 94})]}}}, parent_config={'configurable': {'thread_id': '2', 'thread_ts': '1ef40627-775a-6746-8002-a3967bf0eae6'}}, pending_writes=[])"
|
|
]
|
|
},
|
|
"execution_count": 11,
|
|
"metadata": {},
|
|
"output_type": "execute_result"
|
|
}
|
|
],
|
|
"source": [
|
|
"checkpoint_tuple"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "c0a47d3e-e588-48fc-a5d4-2145dff17e77",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Use async connection"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "ee6b6cf7-d8f7-4777-a48d-93b5855fe681",
|
|
"metadata": {},
|
|
"source": [
|
|
"### With a connection pool"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 12,
|
|
"id": "20cea8b7-8f13-4dc7-a3c9-825040eb4c57",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stderr",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"/Users/vadymbarda/.virtualenvs/langgraph-postgres/lib/python3.11/site-packages/psycopg_pool/pool_async.py:138: RuntimeWarning: opening the async pool AsyncConnectionPool in the constructor is deprecated and will not be supported anymore in a future release. Please use `await pool.open()`, or use the pool as context manager using: `async with AsyncConnectionPool(...) as pool: `...\n",
|
|
" warnings.warn(\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"from psycopg_pool import AsyncConnectionPool\n",
|
|
"\n",
|
|
"pool = AsyncConnectionPool(\n",
|
|
" # Example configuration\n",
|
|
" conninfo=DB_URI,\n",
|
|
" max_size=20,\n",
|
|
")\n",
|
|
"\n",
|
|
"checkpointer = PostgresSaver(async_connection=pool)\n",
|
|
"await checkpointer.acreate_tables(pool)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 13,
|
|
"id": "f889dce6-7ec1-4277-b8af-ace7811733fa",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)\n",
|
|
"config = {\"configurable\": {\"thread_id\": \"3\"}}\n",
|
|
"res = await graph.ainvoke(\n",
|
|
" {\"messages\": [(\"human\", \"what's the weather in nyc\")]}, config\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 14,
|
|
"id": "ed58c722-1662-4ae2-9bb7-4872158a5b29",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"checkpoint_tuple = await checkpointer.aget_tuple(config)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 15,
|
|
"id": "e0c42044-4de6-4742-8e00-fe295d50c95a",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"text/plain": [
|
|
"CheckpointTuple(config={'configurable': {'thread_id': '3', 'thread_ts': '1ef40627-8b0e-6b02-8003-68a7a04ea6a5'}}, checkpoint={'v': 1, 'ts': '2024-07-12T15:21:54.760751+00:00', 'id': '1ef40627-8b0e-6b02-8003-68a7a04ea6a5', 'channel_values': {'messages': [HumanMessage(content=\"what's the weather in nyc\", id='108ac72d-f658-4ae0-af57-af481adc8aa5'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_XY5TyZEwF5nbdNTWjjiqGtdS', 'function': {'arguments': '{\"city\":\"nyc\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 15, 'prompt_tokens': 58, 'total_tokens': 73}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_d33f7b429e', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-454e2142-6f18-4676-ac4b-91f89ea7a6d4-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'nyc'}, 'id': 'call_XY5TyZEwF5nbdNTWjjiqGtdS', 'type': 'tool_call'}], usage_metadata={'input_tokens': 58, 'output_tokens': 15, 'total_tokens': 73}), ToolMessage(content='It might be cloudy in nyc', name='get_weather', id='2d86514f-b8f0-439b-ab94-68c731309c63', tool_call_id='call_XY5TyZEwF5nbdNTWjjiqGtdS'), AIMessage(content='The weather in NYC might be cloudy.', response_metadata={'token_usage': {'completion_tokens': 9, 'prompt_tokens': 88, 'total_tokens': 97}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_d33f7b429e', 'finish_reason': 'stop', 'logprobs': None}, id='run-90ae3622-b480-4964-b689-9c1a572112f1-0', usage_metadata={'input_tokens': 88, 'output_tokens': 9, 'total_tokens': 97})], 'agent': 'agent'}, 'channel_versions': {'__start__': 2, 'messages': 5, 'start:agent': 3, 'agent': 5, 'branch:agent:should_continue:tools': 4, 'tools': 5}, 'versions_seen': {'__start__': {'__start__': 1}, 'agent': {'start:agent': 3, 'tools': 4}, 'tools': {'branch:agent:should_continue:tools': 3}}, 'pending_sends': []}, metadata={'source': 'loop', 'step': 3, 'writes': {'agent': {'messages': [AIMessage(content='The weather in NYC might be cloudy.', response_metadata={'token_usage': {'completion_tokens': 9, 'prompt_tokens': 88, 'total_tokens': 97}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_d33f7b429e', 'finish_reason': 'stop', 'logprobs': None}, id='run-90ae3622-b480-4964-b689-9c1a572112f1-0', usage_metadata={'input_tokens': 88, 'output_tokens': 9, 'total_tokens': 97})]}}}, parent_config={'configurable': {'thread_id': '3', 'thread_ts': '1ef40627-860c-63d4-8002-49a92ae87052'}}, pending_writes=[])"
|
|
]
|
|
},
|
|
"execution_count": 15,
|
|
"metadata": {},
|
|
"output_type": "execute_result"
|
|
}
|
|
],
|
|
"source": [
|
|
"checkpoint_tuple"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "56552584-9eb8-40df-a6a0-44151018b509",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Use connection"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 16,
|
|
"id": "386b78bc-2f73-49ba-a2a4-47bce6fc49b7",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from psycopg import AsyncConnection\n",
|
|
"\n",
|
|
"async with await AsyncConnection.connect(DB_URI) as conn:\n",
|
|
" checkpointer = PostgresSaver(async_connection=conn)\n",
|
|
" graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)\n",
|
|
" config = {\"configurable\": {\"thread_id\": \"4\"}}\n",
|
|
" res = await graph.ainvoke(\n",
|
|
" {\"messages\": [(\"human\", \"what's the weather in nyc\")]}, config\n",
|
|
" )\n",
|
|
" checkpoint_tuples = [c async for c in checkpointer.alist(config)]"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 17,
|
|
"id": "d1ed1344-c923-4a46-b04e-cc3646737d48",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"text/plain": [
|
|
"[CheckpointTuple(config={'configurable': {'thread_id': '4', 'thread_ts': '1ef40627-96b9-6682-8003-134aebfec1e9'}}, checkpoint={'v': 1, 'ts': '2024-07-12T15:21:55.984109+00:00', 'id': '1ef40627-96b9-6682-8003-134aebfec1e9', 'channel_values': {'messages': [HumanMessage(content=\"what's the weather in nyc\", id='1c1e48ba-fa25-4190-a847-459828f44579'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_pS4ybOXkIDOmS93jZ8wOYGfU', 'function': {'arguments': '{\"city\":\"nyc\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 15, 'prompt_tokens': 58, 'total_tokens': 73}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_d33f7b429e', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-07c7ee03-64f7-462a-9249-615572156216-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'nyc'}, 'id': 'call_pS4ybOXkIDOmS93jZ8wOYGfU', 'type': 'tool_call'}], usage_metadata={'input_tokens': 58, 'output_tokens': 15, 'total_tokens': 73}), ToolMessage(content='It might be cloudy in nyc', name='get_weather', id='50e612d7-c770-44dd-b128-4bfdbd7d5b7d', tool_call_id='call_pS4ybOXkIDOmS93jZ8wOYGfU'), AIMessage(content='The weather in NYC might be cloudy.', response_metadata={'token_usage': {'completion_tokens': 9, 'prompt_tokens': 88, 'total_tokens': 97}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_d33f7b429e', 'finish_reason': 'stop', 'logprobs': None}, id='run-924e10c9-7005-4cbf-a92e-3ce63b54092f-0', usage_metadata={'input_tokens': 88, 'output_tokens': 9, 'total_tokens': 97})], 'agent': 'agent'}, 'channel_versions': {'__start__': 2, 'messages': 5, 'start:agent': 3, 'agent': 5, 'branch:agent:should_continue:tools': 4, 'tools': 5}, 'versions_seen': {'__start__': {'__start__': 1}, 'agent': {'start:agent': 3, 'tools': 4}, 'tools': {'branch:agent:should_continue:tools': 3}}, 'pending_sends': []}, metadata={'source': 'loop', 'step': 3, 'writes': {'agent': {'messages': [AIMessage(content='The weather in NYC might be cloudy.', response_metadata={'token_usage': {'completion_tokens': 9, 'prompt_tokens': 88, 'total_tokens': 97}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_d33f7b429e', 'finish_reason': 'stop', 'logprobs': None}, id='run-924e10c9-7005-4cbf-a92e-3ce63b54092f-0', usage_metadata={'input_tokens': 88, 'output_tokens': 9, 'total_tokens': 97})]}}}, parent_config={'configurable': {'thread_id': '4', 'thread_ts': '1ef40627-96b9-6682-8003-134aebfec1e9'}}, pending_writes=None),\n",
|
|
" CheckpointTuple(config={'configurable': {'thread_id': '4', 'thread_ts': '1ef40627-91a0-6100-8002-e404dda477d4'}}, checkpoint={'v': 1, 'ts': '2024-07-12T15:21:55.449447+00:00', 'id': '1ef40627-91a0-6100-8002-e404dda477d4', 'channel_values': {'messages': [HumanMessage(content=\"what's the weather in nyc\", id='1c1e48ba-fa25-4190-a847-459828f44579'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_pS4ybOXkIDOmS93jZ8wOYGfU', 'function': {'arguments': '{\"city\":\"nyc\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 15, 'prompt_tokens': 58, 'total_tokens': 73}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_d33f7b429e', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-07c7ee03-64f7-462a-9249-615572156216-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'nyc'}, 'id': 'call_pS4ybOXkIDOmS93jZ8wOYGfU', 'type': 'tool_call'}], usage_metadata={'input_tokens': 58, 'output_tokens': 15, 'total_tokens': 73}), ToolMessage(content='It might be cloudy in nyc', name='get_weather', id='50e612d7-c770-44dd-b128-4bfdbd7d5b7d', tool_call_id='call_pS4ybOXkIDOmS93jZ8wOYGfU')], 'tools': 'tools'}, 'channel_versions': {'__start__': 2, 'messages': 4, 'start:agent': 3, 'agent': 4, 'branch:agent:should_continue:tools': 4, 'tools': 4}, 'versions_seen': {'__start__': {'__start__': 1}, 'agent': {'start:agent': 2}, 'tools': {'branch:agent:should_continue:tools': 3}}, 'pending_sends': []}, metadata={'source': 'loop', 'step': 2, 'writes': {'tools': {'messages': [ToolMessage(content='It might be cloudy in nyc', name='get_weather', id='50e612d7-c770-44dd-b128-4bfdbd7d5b7d', tool_call_id='call_pS4ybOXkIDOmS93jZ8wOYGfU')]}}}, parent_config={'configurable': {'thread_id': '4', 'thread_ts': '1ef40627-91a0-6100-8002-e404dda477d4'}}, pending_writes=None),\n",
|
|
" CheckpointTuple(config={'configurable': {'thread_id': '4', 'thread_ts': '1ef40627-9194-66de-8001-86c8d77c2d7c'}}, checkpoint={'v': 1, 'ts': '2024-07-12T15:21:55.444687+00:00', 'id': '1ef40627-9194-66de-8001-86c8d77c2d7c', 'channel_values': {'messages': [HumanMessage(content=\"what's the weather in nyc\", id='1c1e48ba-fa25-4190-a847-459828f44579'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_pS4ybOXkIDOmS93jZ8wOYGfU', 'function': {'arguments': '{\"city\":\"nyc\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 15, 'prompt_tokens': 58, 'total_tokens': 73}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_d33f7b429e', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-07c7ee03-64f7-462a-9249-615572156216-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'nyc'}, 'id': 'call_pS4ybOXkIDOmS93jZ8wOYGfU', 'type': 'tool_call'}], usage_metadata={'input_tokens': 58, 'output_tokens': 15, 'total_tokens': 73})], 'agent': 'agent', 'branch:agent:should_continue:tools': 'agent'}, 'channel_versions': {'__start__': 2, 'messages': 3, 'start:agent': 3, 'agent': 3, 'branch:agent:should_continue:tools': 3}, 'versions_seen': {'__start__': {'__start__': 1}, 'agent': {'start:agent': 2}, 'tools': {}}, 'pending_sends': []}, metadata={'source': 'loop', 'step': 1, 'writes': {'agent': {'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_pS4ybOXkIDOmS93jZ8wOYGfU', 'function': {'arguments': '{\"city\":\"nyc\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 15, 'prompt_tokens': 58, 'total_tokens': 73}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_d33f7b429e', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-07c7ee03-64f7-462a-9249-615572156216-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'nyc'}, 'id': 'call_pS4ybOXkIDOmS93jZ8wOYGfU', 'type': 'tool_call'}], usage_metadata={'input_tokens': 58, 'output_tokens': 15, 'total_tokens': 73})]}}}, parent_config={'configurable': {'thread_id': '4', 'thread_ts': '1ef40627-9194-66de-8001-86c8d77c2d7c'}}, pending_writes=None),\n",
|
|
" CheckpointTuple(config={'configurable': {'thread_id': '4', 'thread_ts': '1ef40627-8b8a-6b1c-8000-55b423aa733b'}}, checkpoint={'v': 1, 'ts': '2024-07-12T15:21:54.811566+00:00', 'id': '1ef40627-8b8a-6b1c-8000-55b423aa733b', 'channel_values': {'messages': [HumanMessage(content=\"what's the weather in nyc\", id='1c1e48ba-fa25-4190-a847-459828f44579')], 'start:agent': '__start__'}, 'channel_versions': {'__start__': 2, 'messages': 2, 'start:agent': 2}, 'versions_seen': {'__start__': {'__start__': 1}, 'agent': {}, 'tools': {}}, 'pending_sends': []}, metadata={'source': 'loop', 'step': 0, 'writes': None}, parent_config={'configurable': {'thread_id': '4', 'thread_ts': '1ef40627-8b8a-6b1c-8000-55b423aa733b'}}, pending_writes=None),\n",
|
|
" CheckpointTuple(config={'configurable': {'thread_id': '4', 'thread_ts': '1ef40627-8b88-62b8-bfff-9922bbf9342b'}}, checkpoint={'v': 1, 'ts': '2024-07-12T15:21:54.810527+00:00', 'id': '1ef40627-8b88-62b8-bfff-9922bbf9342b', 'channel_values': {'messages': [], '__start__': {'messages': [['human', \"what's the weather in nyc\"]]}}, 'channel_versions': {'__start__': 1}, 'versions_seen': {}, 'pending_sends': []}, metadata={'source': 'input', 'step': -1, 'writes': {'messages': [['human', \"what's the weather in nyc\"]]}}, parent_config=None, pending_writes=None)]"
|
|
]
|
|
},
|
|
"execution_count": 17,
|
|
"metadata": {},
|
|
"output_type": "execute_result"
|
|
}
|
|
],
|
|
"source": [
|
|
"checkpoint_tuples"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "6a39d1ff-ca37-4457-8b52-07d33b59c36e",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": []
|
|
}
|
|
],
|
|
"metadata": {
|
|
"kernelspec": {
|
|
"display_name": "langgraph-postgres",
|
|
"language": "python",
|
|
"name": "langgraph-postgres"
|
|
},
|
|
"language_info": {
|
|
"codemirror_mode": {
|
|
"name": "ipython",
|
|
"version": 3
|
|
},
|
|
"file_extension": ".py",
|
|
"mimetype": "text/x-python",
|
|
"name": "python",
|
|
"nbconvert_exporter": "python",
|
|
"pygments_lexer": "ipython3",
|
|
"version": "3.11.9"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 5
|
|
}
|