WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-26 17:42:24 +02:00
Code Issues Packages Projects Releases Wiki Activity

ci: Enable mypy checks for checkpoint-postgres lib

Browse Source
This commit is contained in:
Nuno Campos
2024-09-19 08:40:31 -07:00
parent c793a9e36d
commit b8a8651c23
9 changed files with 103 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/checkpoint-postgres/Makefile
+2 -1
View File
@@ -41,7 +41,8 @@ lint lint_diff lint_package lint_tests:
poetry run ruff check .
[ "$(PYTHON_FILES)" = "" ] || poetry run ruff format $(PYTHON_FILES) --diff
[ "$(PYTHON_FILES)" = "" ] || poetry run ruff check --select I $(PYTHON_FILES)
[ "$(PYTHON_FILES)" = "" ] || mkdir -p $(MYPY_CACHE) || poetry run mypy $(PYTHON_FILES) --cache-dir $(MYPY_CACHE)
[ "$(PYTHON_FILES)" = "" ] || mkdir -p $(MYPY_CACHE)
[ "$(PYTHON_FILES)" = "" ] || poetry run mypy $(PYTHON_FILES) --cache-dir $(MYPY_CACHE)
format format_diff:
poetry run ruff format $(PYTHON_FILES)
libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py
+14 -8
View File
@@ -5,7 +5,7 @@ from typing import Any, Iterator, List, Optional, Union
from langchain_core.runnables import RunnableConfig
from psycopg import Connection, Cursor, Pipeline
from psycopg.errors import UndefinedTable
from psycopg.rows import dict_row
from psycopg.rows import DictRow, dict_row
from psycopg.types.json import Jsonb
from psycopg_pool import ConnectionPool
@@ -22,9 +22,11 @@ from langgraph.checkpoint.postgres.base import (
)
from langgraph.checkpoint.serde.base import SerializerProtocol
Conn = Union[Connection[DictRow], ConnectionPool[Connection[DictRow]]]
@contextmanager
def _get_connection(conn: Union[Connection, ConnectionPool]) -> Iterator[Connection]:
def _get_connection(conn: Conn) -> Iterator[Connection[DictRow]]:
if isinstance(conn, Connection):
yield conn
elif isinstance(conn, ConnectionPool):
@@ -39,7 +41,7 @@ class PostgresSaver(BasePostgresSaver):
def __init__(
self,
conn: Union[Connection, ConnectionPool],
conn: Conn,
pipe: Optional[Pipeline] = None,
serde: Optional[SerializerProtocol] = None,
) -> None:
@@ -85,9 +87,13 @@ class PostgresSaver(BasePostgresSaver):
"""
with self._cursor() as cur:
try:
version = cur.execute(
row = cur.execute(
"SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
).fetchone()["v"]
).fetchone()
if row is None:
version = -1
else:
version = row["v"]
except UndefinedTable:
version = -1
for v, migration in zip(
@@ -212,7 +218,7 @@ class PostgresSaver(BasePostgresSaver):
checkpoint_id = get_checkpoint_id(config)
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
if checkpoint_id:
args = (thread_id, checkpoint_ns, checkpoint_id)
args: tuple[Any, ...] = (thread_id, checkpoint_ns, checkpoint_id)
where = "WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s"
else:
args = (thread_id, checkpoint_ns)
@@ -306,7 +312,7 @@ class PostgresSaver(BasePostgresSaver):
self._dump_blobs(
thread_id,
checkpoint_ns,
copy.pop("channel_values"),
copy.pop("channel_values"), # type: ignore[misc]
new_versions,
),
)
@@ -356,7 +362,7 @@ class PostgresSaver(BasePostgresSaver):
)
@contextmanager
def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor]:
def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor[DictRow]]:
with _get_connection(self.conn) as conn:
if self.pipe:
# a connection in pipeline mode can be used concurrently
libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py
+16 -8
View File
@@ -5,7 +5,7 @@ from typing import Any, AsyncIterator, Iterator, List, Optional, Union
from langchain_core.runnables import RunnableConfig
from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline
from psycopg.errors import UndefinedTable
from psycopg.rows import dict_row
from psycopg.rows import DictRow, dict_row
from psycopg.types.json import Jsonb
from psycopg_pool import AsyncConnectionPool
@@ -20,11 +20,13 @@ from langgraph.checkpoint.base import (
from langgraph.checkpoint.postgres.base import BasePostgresSaver
from langgraph.checkpoint.serde.base import SerializerProtocol
Conn = Union[AsyncConnection[DictRow], AsyncConnectionPool[AsyncConnection[DictRow]]]
@asynccontextmanager
async def _get_connection(
conn: Union[AsyncConnection, AsyncConnectionPool],
) -> AsyncIterator[AsyncConnection]:
conn: Conn,
) -> AsyncIterator[AsyncConnection[DictRow]]:
if isinstance(conn, AsyncConnection):
yield conn
elif isinstance(conn, AsyncConnectionPool):
@@ -39,7 +41,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
def __init__(
self,
conn: Union[AsyncConnection, AsyncConnectionPool],
conn: Conn,
pipe: Optional[AsyncPipeline] = None,
serde: Optional[SerializerProtocol] = None,
) -> None:
@@ -93,7 +95,11 @@ class AsyncPostgresSaver(BasePostgresSaver):
results = await cur.execute(
"SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
)
version = (await results.fetchone())["v"]
row = await results.fetchone()
if row is None:
version = -1
else:
version = row["v"]
except UndefinedTable:
version = -1
for v, migration in zip(
@@ -180,7 +186,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
checkpoint_id = get_checkpoint_id(config)
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
if checkpoint_id:
args = (thread_id, checkpoint_ns, checkpoint_id)
args: tuple[Any, ...] = (thread_id, checkpoint_ns, checkpoint_id)
where = "WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s"
else:
args = (thread_id, checkpoint_ns)
@@ -265,7 +271,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
self._dump_blobs,
thread_id,
checkpoint_ns,
copy.pop("channel_values"),
copy.pop("channel_values"), # type: ignore[misc]
new_versions,
),
)
@@ -314,7 +320,9 @@ class AsyncPostgresSaver(BasePostgresSaver):
await cur.executemany(query, params)
@asynccontextmanager
async def _cursor(self, *, pipeline: bool = False) -> AsyncIterator[AsyncCursor]:
async def _cursor(
self, *, pipeline: bool = False
) -> AsyncIterator[AsyncCursor[DictRow]]:
async with _get_connection(self.conn) as conn:
if self.pipe:
# a connection in pipeline mode can be used concurrently
libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py
+10 -8
View File
@@ -1,5 +1,5 @@
import random
from typing import Any, List, Optional, Tuple
from typing import Any, List, Optional, Tuple, cast
from langchain_core.runnables import RunnableConfig
from psycopg.types.json import Jsonb
@@ -7,7 +7,9 @@ from psycopg.types.json import Jsonb
from langgraph.checkpoint.base import (
WRITES_IDX_MAP,
BaseCheckpointSaver,
ChannelVersions,
Checkpoint,
CheckpointMetadata,
get_checkpoint_id,
)
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
@@ -122,7 +124,7 @@ INSERT_CHECKPOINT_WRITES_SQL = """
"""
class BasePostgresSaver(BaseCheckpointSaver):
class BasePostgresSaver(BaseCheckpointSaver[str]):
SELECT_SQL = SELECT_SQL
MIGRATIONS = MIGRATIONS
UPSERT_CHECKPOINT_BLOBS_SQL = UPSERT_CHECKPOINT_BLOBS_SQL
@@ -165,8 +167,8 @@ class BasePostgresSaver(BaseCheckpointSaver):
thread_id: str,
checkpoint_ns: str,
values: dict[str, Any],
versions: dict[str, str],
) -> list[tuple[str, str, str, str, str, bytes]]:
versions: ChannelVersions,
) -> list[tuple[str, str, str, str, str, Optional[bytes]]]:
if not versions:
return []
@@ -175,7 +177,7 @@ class BasePostgresSaver(BaseCheckpointSaver):
thread_id,
checkpoint_ns,
k,
ver,
cast(str, ver),
*(
self.serde.dumps_typed(values[k])
if k in values
@@ -208,7 +210,7 @@ class BasePostgresSaver(BaseCheckpointSaver):
checkpoint_id: str,
task_id: str,
writes: list[tuple[str, Any]],
) -> list[tuple[str, str, str, int, str, str, bytes]]:
) -> list[tuple[str, str, str, str, int, str, str, bytes]]:
return [
(
thread_id,
@@ -222,10 +224,10 @@ class BasePostgresSaver(BaseCheckpointSaver):
for idx, (channel, value) in enumerate(writes)
]
def _load_metadata(self, metadata: dict[str, Any]) -> dict[str, Any]:
def _load_metadata(self, metadata: dict[str, Any]) -> CheckpointMetadata:
return self.jsonplus_serde.loads(self.jsonplus_serde.dumps(metadata))
def _dump_metadata(self, metadata) -> str:
def _dump_metadata(self, metadata: CheckpointMetadata) -> str:
serialized_metadata = self.jsonplus_serde.dumps(metadata)
return serialized_metadata.decode()
libs/checkpoint-postgres/poetry.lock
Generated
+28 -28
View File
@@ -434,38 +434,38 @@ files = [
[[package]]
name = "mypy"
version = "1.11.0"
version = "1.11.2"
description = "Optional static typing for Python"
optional = false
python-versions = ">=3.8"
files = [
{file = "mypy-1.11.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:a3824187c99b893f90c845bab405a585d1ced4ff55421fdf5c84cb7710995229"},
{file = "mypy-1.11.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:96f8dbc2c85046c81bcddc246232d500ad729cb720da4e20fce3b542cab91287"},
{file = "mypy-1.11.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1a5d8d8dd8613a3e2be3eae829ee891b6b2de6302f24766ff06cb2875f5be9c6"},
{file = "mypy-1.11.0-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:72596a79bbfb195fd41405cffa18210af3811beb91ff946dbcb7368240eed6be"},
{file = "mypy-1.11.0-cp310-cp310-win_amd64.whl", hash = "sha256:35ce88b8ed3a759634cb4eb646d002c4cef0a38f20565ee82b5023558eb90c00"},
{file = "mypy-1.11.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:98790025861cb2c3db8c2f5ad10fc8c336ed2a55f4daf1b8b3f877826b6ff2eb"},
{file = "mypy-1.11.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:25bcfa75b9b5a5f8d67147a54ea97ed63a653995a82798221cca2a315c0238c1"},
{file = "mypy-1.11.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0bea2a0e71c2a375c9fa0ede3d98324214d67b3cbbfcbd55ac8f750f85a414e3"},
{file = "mypy-1.11.0-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:d2b3d36baac48e40e3064d2901f2fbd2a2d6880ec6ce6358825c85031d7c0d4d"},
{file = "mypy-1.11.0-cp311-cp311-win_amd64.whl", hash = "sha256:d8e2e43977f0e09f149ea69fd0556623919f816764e26d74da0c8a7b48f3e18a"},
{file = "mypy-1.11.0-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:1d44c1e44a8be986b54b09f15f2c1a66368eb43861b4e82573026e04c48a9e20"},
{file = "mypy-1.11.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:cea3d0fb69637944dd321f41bc896e11d0fb0b0aa531d887a6da70f6e7473aba"},
{file = "mypy-1.11.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a83ec98ae12d51c252be61521aa5731f5512231d0b738b4cb2498344f0b840cd"},
{file = "mypy-1.11.0-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:c7b73a856522417beb78e0fb6d33ef89474e7a622db2653bc1285af36e2e3e3d"},
{file = "mypy-1.11.0-cp312-cp312-win_amd64.whl", hash = "sha256:f2268d9fcd9686b61ab64f077be7ffbc6fbcdfb4103e5dd0cc5eaab53a8886c2"},
{file = "mypy-1.11.0-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:940bfff7283c267ae6522ef926a7887305945f716a7704d3344d6d07f02df850"},
{file = "mypy-1.11.0-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:14f9294528b5f5cf96c721f231c9f5b2733164e02c1c018ed1a0eff8a18005ac"},
{file = "mypy-1.11.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d7b54c27783991399046837df5c7c9d325d921394757d09dbcbf96aee4649fe9"},
{file = "mypy-1.11.0-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:65f190a6349dec29c8d1a1cd4aa71284177aee5949e0502e6379b42873eddbe7"},
{file = "mypy-1.11.0-cp38-cp38-win_amd64.whl", hash = "sha256:dbe286303241fea8c2ea5466f6e0e6a046a135a7e7609167b07fd4e7baf151bf"},
{file = "mypy-1.11.0-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:104e9c1620c2675420abd1f6c44bab7dd33cc85aea751c985006e83dcd001095"},
{file = "mypy-1.11.0-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:f006e955718ecd8d159cee9932b64fba8f86ee6f7728ca3ac66c3a54b0062abe"},
{file = "mypy-1.11.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:becc9111ca572b04e7e77131bc708480cc88a911adf3d0239f974c034b78085c"},
{file = "mypy-1.11.0-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:6801319fe76c3f3a3833f2b5af7bd2c17bb93c00026a2a1b924e6762f5b19e13"},
{file = "mypy-1.11.0-cp39-cp39-win_amd64.whl", hash = "sha256:c1a184c64521dc549324ec6ef7cbaa6b351912be9cb5edb803c2808a0d7e85ac"},
{file = "mypy-1.11.0-py3-none-any.whl", hash = "sha256:56913ec8c7638b0091ef4da6fcc9136896914a9d60d54670a75880c3e5b99ace"},
{file = "mypy-1.11.0.tar.gz", hash = "sha256:93743608c7348772fdc717af4aeee1997293a1ad04bc0ea6efa15bf65385c538"},
{file = "mypy-1.11.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:d42a6dd818ffce7be66cce644f1dff482f1d97c53ca70908dff0b9ddc120b77a"},
{file = "mypy-1.11.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:801780c56d1cdb896eacd5619a83e427ce436d86a3bdf9112527f24a66618fef"},
{file = "mypy-1.11.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:41ea707d036a5307ac674ea172875f40c9d55c5394f888b168033177fce47383"},
{file = "mypy-1.11.2-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:6e658bd2d20565ea86da7d91331b0eed6d2eee22dc031579e6297f3e12c758c8"},
{file = "mypy-1.11.2-cp310-cp310-win_amd64.whl", hash = "sha256:478db5f5036817fe45adb7332d927daa62417159d49783041338921dcf646fc7"},
{file = "mypy-1.11.2-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:75746e06d5fa1e91bfd5432448d00d34593b52e7e91a187d981d08d1f33d4385"},
{file = "mypy-1.11.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:a976775ab2256aadc6add633d44f100a2517d2388906ec4f13231fafbb0eccca"},
{file = "mypy-1.11.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:cd953f221ac1379050a8a646585a29574488974f79d8082cedef62744f0a0104"},
{file = "mypy-1.11.2-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:57555a7715c0a34421013144a33d280e73c08df70f3a18a552938587ce9274f4"},
{file = "mypy-1.11.2-cp311-cp311-win_amd64.whl", hash = "sha256:36383a4fcbad95f2657642a07ba22ff797de26277158f1cc7bd234821468b1b6"},
{file = "mypy-1.11.2-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:e8960dbbbf36906c5c0b7f4fbf2f0c7ffb20f4898e6a879fcf56a41a08b0d318"},
{file = "mypy-1.11.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:06d26c277962f3fb50e13044674aa10553981ae514288cb7d0a738f495550b36"},
{file = "mypy-1.11.2-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6e7184632d89d677973a14d00ae4d03214c8bc301ceefcdaf5c474866814c987"},
{file = "mypy-1.11.2-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:3a66169b92452f72117e2da3a576087025449018afc2d8e9bfe5ffab865709ca"},
{file = "mypy-1.11.2-cp312-cp312-win_amd64.whl", hash = "sha256:969ea3ef09617aff826885a22ece0ddef69d95852cdad2f60c8bb06bf1f71f70"},
{file = "mypy-1.11.2-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:37c7fa6121c1cdfcaac97ce3d3b5588e847aa79b580c1e922bb5d5d2902df19b"},
{file = "mypy-1.11.2-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:4a8a53bc3ffbd161b5b2a4fff2f0f1e23a33b0168f1c0778ec70e1a3d66deb86"},
{file = "mypy-1.11.2-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2ff93107f01968ed834f4256bc1fc4475e2fecf6c661260066a985b52741ddce"},
{file = "mypy-1.11.2-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:edb91dded4df17eae4537668b23f0ff6baf3707683734b6a818d5b9d0c0c31a1"},
{file = "mypy-1.11.2-cp38-cp38-win_amd64.whl", hash = "sha256:ee23de8530d99b6db0573c4ef4bd8f39a2a6f9b60655bf7a1357e585a3486f2b"},
{file = "mypy-1.11.2-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:801ca29f43d5acce85f8e999b1e431fb479cb02d0e11deb7d2abb56bdaf24fd6"},
{file = "mypy-1.11.2-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:af8d155170fcf87a2afb55b35dc1a0ac21df4431e7d96717621962e4b9192e70"},
{file = "mypy-1.11.2-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f7821776e5c4286b6a13138cc935e2e9b6fde05e081bdebf5cdb2bb97c9df81d"},
{file = "mypy-1.11.2-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:539c570477a96a4e6fb718b8d5c3e0c0eba1f485df13f86d2970c91f0673148d"},
{file = "mypy-1.11.2-cp39-cp39-win_amd64.whl", hash = "sha256:3f14cd3d386ac4d05c5a39a51b84387403dadbd936e17cb35882134d4f8f0d24"},
{file = "mypy-1.11.2-py3-none-any.whl", hash = "sha256:b499bc07dbdcd3de92b0a8b29fdf592c111276f6a12fe29c30f6c417dd546d12"},
{file = "mypy-1.11.2.tar.gz", hash = "sha256:7f9993ad3e0ffdc95c2a14b66dee63729f021968bff8ad911867579c65d13a79"},
]
[package.dependencies]
libs/checkpoint-postgres/pyproject.toml
+10
View File
@@ -51,3 +51,13 @@ lint.select = [
"I", # isort
]
lint.ignore = ["E501", "B008", "UP007", "UP006"]
[tool.mypy]
# https://mypy.readthedocs.io/en/stable/config_file.html
disallow_untyped_defs = "True"
explicit_package_bases = "True"
warn_no_return = "False"
warn_unused_ignores = "True"
warn_redundant_casts = "True"
allow_redefinition = "True"
disable_error_code = "typeddict-item, return-value"
libs/checkpoint-postgres/tests/conftest.py
+5 -3
View File
@@ -1,13 +1,15 @@
from typing import AsyncIterator
import pytest
from psycopg import AsyncConnection
from psycopg.errors import UndefinedTable
from psycopg.rows import dict_row
from psycopg.rows import DictRow, dict_row
DEFAULT_URI = "postgres://postgres:postgres@localhost:5441/postgres?sslmode=disable"
@pytest.fixture(scope="function")
async def conn():
async def conn() -> AsyncIterator[AsyncConnection[DictRow]]:
async with await AsyncConnection.connect(
DEFAULT_URI, autocommit=True, prepare_threshold=0, row_factory=dict_row
) as conn:
@@ -15,7 +17,7 @@ async def conn():
@pytest.fixture(scope="function", autouse=True)
async def clear_test_db(conn):
async def clear_test_db(conn: AsyncConnection[DictRow]) -> None:
"""Delete all tables before each test."""
try:
await conn.execute("DELETE FROM checkpoints")
libs/checkpoint-postgres/tests/test_async.py
+9 -7
View File
@@ -1,5 +1,7 @@
from typing import Any
import pytest
from conftest import DEFAULT_URI
from conftest import DEFAULT_URI # type: ignore
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import (
@@ -13,7 +15,7 @@ from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
class TestAsyncPostgresSaver:
@pytest.fixture(autouse=True)
async def setup(self):
async def setup(self) -> None:
# objects for test setup
self.config_1: RunnableConfig = {
"configurable": {
@@ -58,20 +60,20 @@ class TestAsyncPostgresSaver:
async with AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver:
await saver.setup()
async def test_asearch(self):
async def test_asearch(self) -> None:
async with AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver:
await saver.aput(self.config_1, self.chkpnt_1, self.metadata_1, {})
await saver.aput(self.config_2, self.chkpnt_2, self.metadata_2, {})
await saver.aput(self.config_3, self.chkpnt_3, self.metadata_3, {})
# call method / assertions
query_1: CheckpointMetadata = {"source": "input"} # search by 1 key
query_2: CheckpointMetadata = {
query_1 = {"source": "input"} # search by 1 key
query_2 = {
"step": 1,
"writes": {"foo": "bar"},
} # search by multiple keys
query_3: CheckpointMetadata = {} # search by no keys, return all checkpoints
query_4: CheckpointMetadata = {"source": "update", "step": 1} # no match
query_3: dict[str, Any] = {} # search by no keys, return all checkpoints
query_4 = {"source": "update", "step": 1} # no match
search_results_1 = [c async for c in saver.alist(None, filter=query_1)]
assert len(search_results_1) == 1
libs/checkpoint-postgres/tests/test_sync.py
+9 -7
View File
@@ -1,5 +1,7 @@
from typing import Any
import pytest
from conftest import DEFAULT_URI
from conftest import DEFAULT_URI # type: ignore
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import (
@@ -13,7 +15,7 @@ from langgraph.checkpoint.postgres import PostgresSaver
class TestPostgresSaver:
@pytest.fixture(autouse=True)
def setup(self):
def setup(self) -> None:
# objects for test setup
self.config_1: RunnableConfig = {
"configurable": {
@@ -58,7 +60,7 @@ class TestPostgresSaver:
with PostgresSaver.from_conn_string(DEFAULT_URI) as saver:
saver.setup()
def test_search(self):
def test_search(self) -> None:
with PostgresSaver.from_conn_string(DEFAULT_URI) as saver:
# save checkpoints
saver.put(self.config_1, self.chkpnt_1, self.metadata_1, {})
@@ -66,13 +68,13 @@ class TestPostgresSaver:
saver.put(self.config_3, self.chkpnt_3, self.metadata_3, {})
# call method / assertions
query_1: CheckpointMetadata = {"source": "input"} # search by 1 key
query_2: CheckpointMetadata = {
query_1 = {"source": "input"} # search by 1 key
query_2 = {
"step": 1,
"writes": {"foo": "bar"},
} # search by multiple keys
query_3: CheckpointMetadata = {} # search by no keys, return all checkpoints
query_4: CheckpointMetadata = {"source": "update", "step": 1} # no match
query_3: dict[str, Any] = {} # search by no keys, return all checkpoints
query_4 = {"source": "update", "step": 1} # no match
search_results_1 = list(saver.list(None, filter=query_1))
assert len(search_results_1) == 1