mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-17 21:25:46 +02:00
58 KiB
58 KiB
In [1]:
%%capture --no-stderr
%pip install -U psycopg psycopg-pool langgraphIn [2]:
"""Implementation of a langgraph checkpoint saver using Postgres."""
from contextlib import asynccontextmanager, contextmanager
from typing import (
Any,
AsyncGenerator,
AsyncIterator,
Generator,
Optional,
Union,
Tuple,
List,
Sequence,
)
import psycopg
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.serde.jsonplus import JsonPlusSerializer
from langgraph.checkpoint.base import Checkpoint, CheckpointMetadata, CheckpointTuple
from psycopg_pool import AsyncConnectionPool, ConnectionPool
class JsonAndBinarySerializer(JsonPlusSerializer):
def _default(self, obj):
if isinstance(obj, (bytes, bytearray)):
return self._encode_constructor_args(
obj.__class__, method="fromhex", args=[obj.hex()]
)
return super()._default(obj)
def dumps(self, obj: Any) -> tuple[str, bytes]:
if isinstance(obj, bytes):
return "bytes", obj
elif isinstance(obj, bytearray):
return "bytearray", obj
return "json", super().dumps(obj)
def loads(self, s: tuple[str, bytes]) -> Any:
if s[0] == "bytes":
return s[1]
elif s[0] == "bytearray":
return bytearray(s[1])
elif s[0] == "json":
return super().loads(s[1])
else:
raise NotImplementedError(f"Unknown serialization type: {s[0]}")
@contextmanager
def _get_sync_connection(
connection: Union[psycopg.Connection, ConnectionPool, None],
) -> Generator[psycopg.Connection, None, None]:
"""Get the connection to the Postgres database."""
if isinstance(connection, psycopg.Connection):
yield connection
elif isinstance(connection, ConnectionPool):
with connection.connection() as conn:
yield conn
else:
raise ValueError(
"Invalid sync connection object. Please initialize the check pointer "
f"with an appropriate sync connection object. "
f"Got {type(connection)}."
)
@asynccontextmanager
async def _get_async_connection(
connection: Union[psycopg.AsyncConnection, AsyncConnectionPool, None],
) -> AsyncGenerator[psycopg.AsyncConnection, None]:
"""Get the connection to the Postgres database."""
if isinstance(connection, psycopg.AsyncConnection):
yield connection
elif isinstance(connection, AsyncConnectionPool):
async with connection.connection() as conn:
yield conn
else:
raise ValueError(
"Invalid async connection object. Please initialize the check pointer "
f"with an appropriate async connection object. "
f"Got {type(connection)}."
)
class PostgresSaver(BaseCheckpointSaver):
sync_connection: Optional[Union[psycopg.Connection, ConnectionPool]] = None
"""The synchronous connection or pool to the Postgres database.
If providing a connection object, please ensure that the connection is open
and remember to close the connection when done.
"""
async_connection: Optional[
Union[psycopg.AsyncConnection, AsyncConnectionPool]
] = None
"""The asynchronous connection or pool to the Postgres database.
If providing a connection object, please ensure that the connection is open
and remember to close the connection when done.
"""
def __init__(
self,
sync_connection: Optional[Union[psycopg.Connection, ConnectionPool]] = None,
async_connection: Optional[
Union[psycopg.AsyncConnection, AsyncConnectionPool]
] = None,
):
super().__init__(serde=JsonPlusSerializer())
self.sync_connection = sync_connection
self.async_connection = async_connection
@contextmanager
def _get_sync_connection(self) -> Generator[psycopg.Connection, None, None]:
"""Get the connection to the Postgres database."""
with _get_sync_connection(self.sync_connection) as connection:
yield connection
@asynccontextmanager
async def _get_async_connection(
self,
) -> AsyncGenerator[psycopg.AsyncConnection, None]:
"""Get the connection to the Postgres database."""
async with _get_async_connection(self.async_connection) as connection:
yield connection
CREATE_TABLES_QUERY = """
CREATE TABLE IF NOT EXISTS checkpoints (
thread_id TEXT NOT NULL,
thread_ts TEXT NOT NULL,
parent_ts TEXT,
checkpoint BYTEA NOT NULL,
metadata BYTEA NOT NULL,
PRIMARY KEY (thread_id, thread_ts)
);
CREATE TABLE IF NOT EXISTS writes (
thread_id TEXT NOT NULL,
thread_ts TEXT NOT NULL,
task_id TEXT NOT NULL,
idx INTEGER NOT NULL,
channel TEXT NOT NULL,
value BYTEA,
PRIMARY KEY (thread_id, thread_ts, task_id, idx)
);
"""
@staticmethod
def create_tables(connection: Union[psycopg.Connection, ConnectionPool], /) -> None:
"""Create the schema for the checkpoint saver."""
with _get_sync_connection(connection) as conn:
with conn.cursor() as cur:
cur.execute(PostgresSaver.CREATE_TABLES_QUERY)
@staticmethod
async def acreate_tables(
connection: Union[psycopg.AsyncConnection, AsyncConnectionPool], /
) -> None:
"""Create the schema for the checkpoint saver."""
async with _get_async_connection(connection) as conn:
async with conn.cursor() as cur:
await cur.execute(PostgresSaver.CREATE_TABLES_QUERY)
@staticmethod
def drop_tables(connection: psycopg.Connection, /) -> None:
"""Drop the table for the checkpoint saver."""
with connection.cursor() as cur:
cur.execute("DROP TABLE IF EXISTS checkpoints, writes;")
@staticmethod
async def adrop_tables(connection: psycopg.AsyncConnection, /) -> None:
"""Drop the table for the checkpoint saver."""
async with connection.cursor() as cur:
await cur.execute("DROP TABLE IF EXISTS checkpoints, writes;")
UPSERT_CHECKPOINT_QUERY = """
INSERT INTO checkpoints
(thread_id, thread_ts, parent_ts, checkpoint, metadata)
VALUES
(%s, %s, %s, %s, %s)
ON CONFLICT (thread_id, thread_ts)
DO UPDATE SET checkpoint = EXCLUDED.checkpoint,
metadata = EXCLUDED.metadata;
"""
def put(
self,
config: RunnableConfig,
checkpoint: Checkpoint,
metadata: CheckpointMetadata,
) -> RunnableConfig:
"""Put the checkpoint for the given configuration.
Args:
config: The configuration for the checkpoint.
A dict with a `configurable` key which is a dict with
a `thread_id` key and an optional `thread_ts` key.
For example, { 'configurable': { 'thread_id': 'test_thread' } }
checkpoint: The checkpoint to persist.
Returns:
The RunnableConfig that describes the checkpoint that was just created.
It'll contain the `thread_id` and `thread_ts` of the checkpoint.
"""
thread_id = config["configurable"]["thread_id"]
parent_ts = config["configurable"].get("thread_ts")
with self._get_sync_connection() as conn:
with conn.cursor() as cur:
cur.execute(
self.UPSERT_CHECKPOINT_QUERY,
(
thread_id,
checkpoint["id"],
parent_ts if parent_ts else None,
self.serde.dumps(checkpoint),
self.serde.dumps(metadata),
),
)
return {
"configurable": {
"thread_id": thread_id,
"thread_ts": checkpoint["id"],
},
}
async def aput(
self,
config: RunnableConfig,
checkpoint: Checkpoint,
metadata: CheckpointMetadata,
) -> RunnableConfig:
"""Put the checkpoint for the given configuration.
Args:
config: The configuration for the checkpoint.
A dict with a `configurable` key which is a dict with
a `thread_id` key and an optional `thread_ts` key.
For example, { 'configurable': { 'thread_id': 'test_thread' } }
checkpoint: The checkpoint to persist.
Returns:
The RunnableConfig that describes the checkpoint that was just created.
It'll contain the `thread_id` and `thread_ts` of the checkpoint.
"""
thread_id = config["configurable"]["thread_id"]
parent_ts = config["configurable"].get("thread_ts")
async with self._get_async_connection() as conn:
async with conn.cursor() as cur:
await cur.execute(
self.UPSERT_CHECKPOINT_QUERY,
(
thread_id,
checkpoint["id"],
parent_ts if parent_ts else None,
self.serde.dumps(checkpoint),
self.serde.dumps(metadata),
),
)
return {
"configurable": {
"thread_id": thread_id,
"thread_ts": checkpoint["id"],
},
}
UPSERT_WRITES_QUERY = """
INSERT INTO writes
(thread_id, thread_ts, task_id, idx, channel, value)
VALUES
(%s, %s, %s, %s, %s, %s)
ON CONFLICT (thread_id, thread_ts, task_id, idx)
DO UPDATE SET value = EXCLUDED.value;
"""
def put_writes(
self,
config: RunnableConfig,
writes: Sequence[Tuple[str, Any]],
task_id: str,
) -> None:
with self._get_sync_connection() as conn:
with conn.cursor() as cur:
cur.executemany(
self.UPSERT_WRITES_QUERY,
[
(
str(config["configurable"]["thread_id"]),
str(config["configurable"]["thread_ts"]),
task_id,
idx,
channel,
self.serde.dumps(value),
)
for idx, (channel, value) in enumerate(writes)
],
)
conn.commit()
async def aput_writes(
self,
config: RunnableConfig,
writes: Sequence[Tuple[str, Any]],
task_id: str,
) -> None:
async with self._get_async_connection() as conn:
async with conn.cursor() as cur:
await cur.executemany(
self.UPSERT_WRITES_QUERY,
[
(
str(config["configurable"]["thread_id"]),
str(config["configurable"]["thread_ts"]),
task_id,
idx,
channel,
self.serde.dumps(value),
)
for idx, (channel, value) in enumerate(writes)
],
)
await conn.commit()
LIST_CHECKPOINTS_QUERY_STR = """
SELECT checkpoint, metadata, thread_ts, parent_ts
FROM checkpoints
{where}
ORDER BY thread_ts DESC
"""
def list(
self,
config: Optional[RunnableConfig],
*,
filter: Optional[dict[str, Any]] = None,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
) -> Generator[CheckpointTuple, None, None]:
"""Get all the checkpoints for the given configuration."""
where, args = self._search_where(config, filter, before)
query = self.LIST_CHECKPOINTS_QUERY_STR.format(where=where)
if limit:
query += f" LIMIT {limit}"
with self._get_sync_connection() as conn:
with conn.cursor() as cur:
thread_id = config["configurable"]["thread_id"]
cur.execute(query, tuple(args))
for value in cur:
checkpoint, metadata, thread_ts, parent_ts = value
yield CheckpointTuple(
config={
"configurable": {
"thread_id": thread_id,
"thread_ts": thread_ts,
}
},
checkpoint=self.serde.loads(checkpoint),
metadata=self.serde.loads(metadata),
parent_config={
"configurable": {
"thread_id": thread_id,
"thread_ts": thread_ts,
}
}
if parent_ts
else None,
)
async def alist(
self,
config: Optional[RunnableConfig],
*,
filter: Optional[dict[str, Any]] = None,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
) -> AsyncIterator[CheckpointTuple]:
"""Get all the checkpoints for the given configuration."""
where, args = self._search_where(config, filter, before)
query = self.LIST_CHECKPOINTS_QUERY_STR.format(where=where)
if limit:
query += f" LIMIT {limit}"
async with self._get_async_connection() as conn:
async with conn.cursor() as cur:
thread_id = config["configurable"]["thread_id"]
await cur.execute(query, tuple(args))
async for value in cur:
checkpoint, metadata, thread_ts, parent_ts = value
yield CheckpointTuple(
config={
"configurable": {
"thread_id": thread_id,
"thread_ts": thread_ts,
}
},
checkpoint=self.serde.loads(checkpoint),
metadata=self.serde.loads(metadata),
parent_config={
"configurable": {
"thread_id": thread_id,
"thread_ts": thread_ts,
}
}
if parent_ts
else None,
)
GET_CHECKPOINT_BY_TS_QUERY = """
SELECT checkpoint, metadata, thread_ts, parent_ts
FROM checkpoints
WHERE thread_id = %(thread_id)s AND thread_ts = %(thread_ts)s
"""
GET_CHECKPOINT_QUERY = """
SELECT checkpoint, metadata, thread_ts, parent_ts
FROM checkpoints
WHERE thread_id = %(thread_id)s
ORDER BY thread_ts DESC LIMIT 1
"""
def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
"""Get the checkpoint tuple for the given configuration.
Args:
config: The configuration for the checkpoint.
A dict with a `configurable` key which is a dict with
a `thread_id` key and an optional `thread_ts` key.
For example, { 'configurable': { 'thread_id': 'test_thread' } }
Returns:
The checkpoint tuple for the given configuration if it exists,
otherwise None.
If thread_ts is None, the latest checkpoint is returned if it exists.
"""
thread_id = config["configurable"]["thread_id"]
thread_ts = config["configurable"].get("thread_ts")
with self._get_sync_connection() as conn:
with conn.cursor() as cur:
# find the latest checkpoint for the thread_id
if thread_ts:
cur.execute(
self.GET_CHECKPOINT_BY_TS_QUERY,
{
"thread_id": thread_id,
"thread_ts": thread_ts,
},
)
else:
cur.execute(
self.GET_CHECKPOINT_QUERY,
{
"thread_id": thread_id,
},
)
# if a checkpoint is found, return it
if value := cur.fetchone():
checkpoint, metadata, thread_ts, parent_ts = value
if not config["configurable"].get("thread_ts"):
config = {
"configurable": {
"thread_id": thread_id,
"thread_ts": thread_ts,
}
}
# find any pending writes
cur.execute(
"SELECT task_id, channel, value FROM writes WHERE thread_id = %(thread_id)s AND thread_ts = %(thread_ts)s",
{
"thread_id": thread_id,
"thread_ts": thread_ts,
},
)
# deserialize the checkpoint and metadata
return CheckpointTuple(
config=config,
checkpoint=self.serde.loads(checkpoint),
metadata=self.serde.loads(metadata),
parent_config={
"configurable": {
"thread_id": thread_id,
"thread_ts": parent_ts,
}
}
if parent_ts
else None,
pending_writes=[
(task_id, channel, self.serde.loads(value))
for task_id, channel, value in cur
],
)
async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
"""Get the checkpoint tuple for the given configuration.
Args:
config: The configuration for the checkpoint.
A dict with a `configurable` key which is a dict with
a `thread_id` key and an optional `thread_ts` key.
For example, { 'configurable': { 'thread_id': 'test_thread' } }
Returns:
The checkpoint tuple for the given configuration if it exists,
otherwise None.
If thread_ts is None, the latest checkpoint is returned if it exists.
"""
thread_id = config["configurable"]["thread_id"]
thread_ts = config["configurable"].get("thread_ts")
async with self._get_async_connection() as conn:
async with conn.cursor() as cur:
# find the latest checkpoint for the thread_id
if thread_ts:
await cur.execute(
self.GET_CHECKPOINT_BY_TS_QUERY,
{
"thread_id": thread_id,
"thread_ts": thread_ts,
},
)
else:
await cur.execute(
self.GET_CHECKPOINT_QUERY,
{
"thread_id": thread_id,
},
)
# if a checkpoint is found, return it
if value := await cur.fetchone():
checkpoint, metadata, thread_ts, parent_ts = value
if not config["configurable"].get("thread_ts"):
config = {
"configurable": {
"thread_id": thread_id,
"thread_ts": thread_ts,
}
}
# find any pending writes
await cur.execute(
"SELECT task_id, channel, value FROM writes WHERE thread_id = %(thread_id)s AND thread_ts = %(thread_ts)s",
{
"thread_id": thread_id,
"thread_ts": thread_ts,
},
)
# deserialize the checkpoint and metadata
return CheckpointTuple(
config=config,
checkpoint=self.serde.loads(checkpoint),
metadata=self.serde.loads(metadata),
parent_config={
"configurable": {
"thread_id": thread_id,
"thread_ts": parent_ts,
}
}
if parent_ts
else None,
pending_writes=[
(task_id, channel, self.serde.loads(value))
async for task_id, channel, value in cur
],
)
def _search_where(
self,
config: Optional[RunnableConfig],
filter: Optional[dict[str, Any]] = None,
before: Optional[RunnableConfig] = None,
) -> Tuple[str, List[Any]]:
"""Return WHERE clause predicates for given config, filter, and before parameters.
Args:
config (Optional[RunnableConfig]): The config to use for filtering.
filter (Optional[Dict[str, Any]]): Additional filtering criteria.
before (Optional[RunnableConfig]): A config to limit results before a certain timestamp.
Returns:
Tuple[str, Sequence[Any]]: A tuple containing the WHERE clause and parameter values.
"""
wheres = []
param_values = []
# Add predicate for config
if config is not None:
wheres.append("thread_id = %s ")
param_values.append(config["configurable"]["thread_id"])
if filter:
raise NotImplementedError()
# Add predicate for limiting results before a certain timestamp
if before is not None:
wheres.append("thread_ts < %s")
param_values.append(before["configurable"]["thread_ts"])
where_clause = "WHERE " + " AND ".join(wheres) if wheres else ""
return where_clause, param_valuesIn [3]:
import getpass
import os
def _set_env(var: str):
if not os.environ.get(var):
os.environ[var] = getpass.getpass(f"{var}: ")
_set_env("OPENAI_API_KEY")OPENAI_API_KEY: ········
In [4]:
from typing import Literal
from langchain_core.runnables import ConfigurableField
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import create_react_agent
@tool
def get_weather(city: Literal["nyc", "sf"]):
"""Use this to get weather information."""
if city == "nyc":
return "It might be cloudy in nyc"
elif city == "sf":
return "It's always sunny in sf"
else:
raise AssertionError("Unknown city")
tools = [get_weather]
model = ChatOpenAI(model_name="gpt-4o", temperature=0)In [5]:
DB_URI = "postgresql://postgres:postgres@localhost:5432/postgres?sslmode=disable"In [6]:
from psycopg_pool import ConnectionPool
pool = ConnectionPool(
# Example configuration
conninfo=DB_URI,
max_size=20,
)
checkpointer = PostgresSaver(sync_connection=pool)
checkpointer.create_tables(pool)In [7]:
graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)
config = {"configurable": {"thread_id": "1"}}
res = graph.invoke({"messages": [("human", "what's the weather in sf")]}, config)In [8]:
resOut [8]:
{'messages': [HumanMessage(content="what's the weather in sf", id='bc87fac7-1da1-4818-a43b-6ba7c9b9b3e4'),
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}),
ToolMessage(content="It's always sunny in sf", name='get_weather', id='8d8f9596-a683-4644-a898-1e303b5a01ea', tool_call_id='call_MjkmibJlXeuNchL6B8qpIjOW'),
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})]}In [9]:
checkpointer.get(config)Out [9]:
{'v': 1,
'ts': '2024-07-12T15:21:51.891852+00:00',
'id': '1ef40627-6fb2-6962-8003-b74d816658c5',
'channel_values': {'messages': [HumanMessage(content="what's the weather in sf", id='bc87fac7-1da1-4818-a43b-6ba7c9b9b3e4'),
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}),
ToolMessage(content="It's always sunny in sf", name='get_weather', id='8d8f9596-a683-4644-a898-1e303b5a01ea', tool_call_id='call_MjkmibJlXeuNchL6B8qpIjOW'),
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})],
'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': []}In [10]:
from psycopg import Connection
with Connection.connect(DB_URI) as conn:
checkpointer = PostgresSaver(sync_connection=conn)
graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)
config = {"configurable": {"thread_id": "2"}}
res = graph.invoke({"messages": [("human", "what's the weather in sf")]}, config)
checkpoint_tuple = checkpointer.get_tuple(config)In [11]:
checkpoint_tupleOut [11]:
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=[])In [12]:
from psycopg_pool import AsyncConnectionPool
pool = AsyncConnectionPool(
# Example configuration
conninfo=DB_URI,
max_size=20,
)
checkpointer = PostgresSaver(async_connection=pool)
await checkpointer.acreate_tables(pool)/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: `... warnings.warn(
In [13]:
graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)
config = {"configurable": {"thread_id": "3"}}
res = await graph.ainvoke(
{"messages": [("human", "what's the weather in nyc")]}, config
)In [14]:
checkpoint_tuple = await checkpointer.aget_tuple(config)In [15]:
checkpoint_tupleOut [15]:
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=[])In [16]:
from psycopg import AsyncConnection
async with await AsyncConnection.connect(DB_URI) as conn:
checkpointer = PostgresSaver(async_connection=conn)
graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)
config = {"configurable": {"thread_id": "4"}}
res = await graph.ainvoke(
{"messages": [("human", "what's the weather in nyc")]}, config
)
checkpoint_tuples = [c async for c in checkpointer.alist(config)]In [17]:
checkpoint_tuplesOut [17]:
[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),
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),
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),
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),
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)]In [ ]: