WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-10 11:47:51 +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/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()