WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-07 18:27:52 +02:00
Code Issues Packages Projects Releases Wiki Activity

More idiomatic sync/async checkpointer fixtures in pytest (#4624)

Browse Source
This commit is contained in:
Nuno Campos
2025-05-09 12:34:34 -07:00
committed by GitHub
parent 6777b5274c ed7f038a19
commit 302ae69e60
8 changed files with 1370 additions and 1357 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/tests/conftest.py
+113 -215
View File
@@ -1,29 +1,32 @@
import sys
from collections.abc import AsyncIterator, Iterator
from contextlib import asynccontextmanager
from typing import Optional
from uuid import UUID, uuid4
from uuid import UUID
import pytest
from langchain_core import __version__ as core_version
from packaging import version
from psycopg import AsyncConnection, Connection
from psycopg_pool import AsyncConnectionPool, ConnectionPool
from pytest_mock import MockerFixture
from langgraph.cache.base import BaseCache
from langgraph.cache.memory import InMemoryCache
from langgraph.cache.sqlite import SqliteCache
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.checkpoint.postgres import PostgresSaver, ShallowPostgresSaver
from langgraph.checkpoint.postgres.aio import (
AsyncPostgresSaver,
AsyncShallowPostgresSaver,
)
from langgraph.checkpoint.serde.encrypted import EncryptedSerializer
from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
from langgraph.store.base import BaseStore
from tests.conftest_checkpointer import (
_checkpointer_memory,
_checkpointer_postgres,
_checkpointer_postgres_aio,
_checkpointer_postgres_aio_pipe,
_checkpointer_postgres_aio_pool,
_checkpointer_postgres_aio_shallow,
_checkpointer_postgres_pipe,
_checkpointer_postgres_pool,
_checkpointer_postgres_shallow,
_checkpointer_sqlite,
_checkpointer_sqlite_aes,
_checkpointer_sqlite_aio,
)
from tests.conftest_store import (
_store_memory,
_store_postgres,
@@ -57,219 +60,54 @@ def deterministic_uuids(mocker: MockerFixture) -> MockerFixture:
return mocker.patch("uuid.uuid4", side_effect=side_effect)
# checkpointer fixtures
@pytest.fixture(params=[True, False])
def checkpoint_during(request: pytest.FixtureRequest) -> bool:
return request.param
# --- start of deprecated fixtures ---
@pytest.fixture(scope="function")
def checkpointer_memory():
from tests.memory_assert import MemorySaverAssertImmutable
yield MemorySaverAssertImmutable()
with _checkpointer_memory() as checkpointer:
yield checkpointer
@pytest.fixture(scope="function")
def checkpointer_sqlite():
with SqliteSaver.from_conn_string(":memory:") as checkpointer:
with _checkpointer_sqlite() as checkpointer:
yield checkpointer
@pytest.fixture(scope="function")
def checkpointer_sqlite_aes():
with SqliteSaver.from_conn_string(":memory:") as checkpointer:
checkpointer.serde = EncryptedSerializer.from_pycryptodome_aes(
key=b"1234567890123456"
)
yield checkpointer
@asynccontextmanager
async def _checkpointer_sqlite_aio():
async with AsyncSqliteSaver.from_conn_string(":memory:") as checkpointer:
with _checkpointer_sqlite_aes() as checkpointer:
yield checkpointer
@pytest.fixture(scope="function")
def checkpointer_postgres():
database = f"test_{uuid4().hex[:16]}"
# create unique db
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
conn.execute(f"CREATE DATABASE {database}")
try:
# yield checkpointer
with PostgresSaver.from_conn_string(
DEFAULT_POSTGRES_URI + database
) as checkpointer:
checkpointer.setup()
yield checkpointer
finally:
# drop unique db
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
conn.execute(f"DROP DATABASE {database}")
with _checkpointer_postgres() as checkpointer:
yield checkpointer
@pytest.fixture(scope="function")
def checkpointer_postgres_shallow():
database = f"test_{uuid4().hex[:16]}"
# create unique db
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
conn.execute(f"CREATE DATABASE {database}")
try:
# yield checkpointer
with ShallowPostgresSaver.from_conn_string(
DEFAULT_POSTGRES_URI + database
) as checkpointer:
checkpointer.setup()
yield checkpointer
finally:
# drop unique db
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
conn.execute(f"DROP DATABASE {database}")
with _checkpointer_postgres_shallow() as checkpointer:
yield checkpointer
@pytest.fixture(scope="function")
def checkpointer_postgres_pipe():
database = f"test_{uuid4().hex[:16]}"
# create unique db
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
conn.execute(f"CREATE DATABASE {database}")
try:
# yield checkpointer
with PostgresSaver.from_conn_string(
DEFAULT_POSTGRES_URI + database
) as checkpointer:
checkpointer.setup()
# setup can't run inside pipeline because of implicit transaction
with checkpointer.conn.pipeline() as pipe:
checkpointer.pipe = pipe
yield checkpointer
finally:
# drop unique db
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
conn.execute(f"DROP DATABASE {database}")
with _checkpointer_postgres_pipe() as checkpointer:
yield checkpointer
@pytest.fixture(scope="function")
def checkpointer_postgres_pool():
database = f"test_{uuid4().hex[:16]}"
# create unique db
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
conn.execute(f"CREATE DATABASE {database}")
try:
# yield checkpointer
with ConnectionPool(
DEFAULT_POSTGRES_URI + database, max_size=10, kwargs={"autocommit": True}
) as pool:
checkpointer = PostgresSaver(pool)
checkpointer.setup()
yield checkpointer
finally:
# drop unique db
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
conn.execute(f"DROP DATABASE {database}")
@asynccontextmanager
async def _checkpointer_postgres_aio():
if sys.version_info < (3, 10):
pytest.skip("Async Postgres tests require Python 3.10+")
database = f"test_{uuid4().hex[:16]}"
# create unique db
async with await AsyncConnection.connect(
DEFAULT_POSTGRES_URI, autocommit=True
) as conn:
await conn.execute(f"CREATE DATABASE {database}")
try:
# yield checkpointer
async with AsyncPostgresSaver.from_conn_string(
DEFAULT_POSTGRES_URI + database
) as checkpointer:
await checkpointer.setup()
yield checkpointer
finally:
# drop unique db
async with await AsyncConnection.connect(
DEFAULT_POSTGRES_URI, autocommit=True
) as conn:
await conn.execute(f"DROP DATABASE {database}")
@asynccontextmanager
async def _checkpointer_postgres_aio_shallow():
if sys.version_info < (3, 10):
pytest.skip("Async Postgres tests require Python 3.10+")
database = f"test_{uuid4().hex[:16]}"
# create unique db
async with await AsyncConnection.connect(
DEFAULT_POSTGRES_URI, autocommit=True
) as conn:
await conn.execute(f"CREATE DATABASE {database}")
try:
# yield checkpointer
async with AsyncShallowPostgresSaver.from_conn_string(
DEFAULT_POSTGRES_URI + database
) as checkpointer:
await checkpointer.setup()
yield checkpointer
finally:
# drop unique db
async with await AsyncConnection.connect(
DEFAULT_POSTGRES_URI, autocommit=True
) as conn:
await conn.execute(f"DROP DATABASE {database}")
@asynccontextmanager
async def _checkpointer_postgres_aio_pipe():
if sys.version_info < (3, 10):
pytest.skip("Async Postgres tests require Python 3.10+")
database = f"test_{uuid4().hex[:16]}"
# create unique db
async with await AsyncConnection.connect(
DEFAULT_POSTGRES_URI, autocommit=True
) as conn:
await conn.execute(f"CREATE DATABASE {database}")
try:
# yield checkpointer
async with AsyncPostgresSaver.from_conn_string(
DEFAULT_POSTGRES_URI + database
) as checkpointer:
await checkpointer.setup()
# setup can't run inside pipeline because of implicit transaction
async with checkpointer.conn.pipeline() as pipe:
checkpointer.pipe = pipe
yield checkpointer
finally:
# drop unique db
async with await AsyncConnection.connect(
DEFAULT_POSTGRES_URI, autocommit=True
) as conn:
await conn.execute(f"DROP DATABASE {database}")
@asynccontextmanager
async def _checkpointer_postgres_aio_pool():
if sys.version_info < (3, 10):
pytest.skip("Async Postgres tests require Python 3.10+")
database = f"test_{uuid4().hex[:16]}"
# create unique db
async with await AsyncConnection.connect(
DEFAULT_POSTGRES_URI, autocommit=True
) as conn:
await conn.execute(f"CREATE DATABASE {database}")
try:
# yield checkpointer
async with AsyncConnectionPool(
DEFAULT_POSTGRES_URI + database, max_size=10, kwargs={"autocommit": True}
) as pool:
checkpointer = AsyncPostgresSaver(pool)
await checkpointer.setup()
yield checkpointer
finally:
# drop unique db
async with await AsyncConnection.connect(
DEFAULT_POSTGRES_URI, autocommit=True
) as conn:
await conn.execute(f"DROP DATABASE {database}")
with _checkpointer_postgres_pool() as checkpointer:
yield checkpointer
@asynccontextmanager
@@ -279,9 +117,8 @@ async def awith_checkpointer(
if checkpointer_name is None:
yield None
elif checkpointer_name == "memory":
from tests.memory_assert import MemorySaverAssertImmutable
yield MemorySaverAssertImmutable()
with _checkpointer_memory() as checkpointer:
yield checkpointer
elif checkpointer_name == "sqlite_aio":
async with _checkpointer_sqlite_aio() as checkpointer:
yield checkpointer
@@ -301,6 +138,9 @@ async def awith_checkpointer(
raise NotImplementedError(f"Unknown checkpointer: {checkpointer_name}")
# --- end of deprecated fixtures ---
@pytest.fixture(scope="function", params=["sqlite", "memory"])
def cache(request: pytest.FixtureRequest) -> Iterator[BaseCache]:
if request.param == "sqlite":
@@ -359,32 +199,90 @@ async def async_store(request: pytest.FixtureRequest) -> AsyncIterator[BaseStore
raise NotImplementedError(f"Unknown store {store_name}")
SHALLOW_CHECKPOINTERS_SYNC = ["postgres_shallow"]
REGULAR_CHECKPOINTERS_SYNC = [
@pytest.fixture(
scope="function",
params=[
"memory",
"sqlite",
"sqlite_aes",
"postgres",
"postgres_pipe",
"postgres_pool",
],
)
def sync_checkpointer(
request: pytest.FixtureRequest,
) -> Iterator[BaseCheckpointSaver]:
checkpointer_name = request.param
if checkpointer_name == "memory":
with _checkpointer_memory() as checkpointer:
yield checkpointer
elif checkpointer_name == "sqlite":
with _checkpointer_sqlite() as checkpointer:
yield checkpointer
elif checkpointer_name == "sqlite_aes":
with _checkpointer_sqlite_aes() as checkpointer:
yield checkpointer
elif checkpointer_name == "postgres":
with _checkpointer_postgres() as checkpointer:
yield checkpointer
elif checkpointer_name == "postgres_pipe":
with _checkpointer_postgres_pipe() as checkpointer:
yield checkpointer
elif checkpointer_name == "postgres_pool":
with _checkpointer_postgres_pool() as checkpointer:
yield checkpointer
else:
raise NotImplementedError(f"Unknown checkpointer: {checkpointer_name}")
@pytest.fixture(
scope="function",
params=[
"memory",
"sqlite_aio",
"postgres_aio",
"postgres_aio_pipe",
"postgres_aio_pool",
],
)
async def async_checkpointer(
request: pytest.FixtureRequest,
) -> AsyncIterator[BaseCheckpointSaver]:
checkpointer_name = request.param
if checkpointer_name == "memory":
with _checkpointer_memory() as checkpointer:
yield checkpointer
elif checkpointer_name == "sqlite_aio":
async with _checkpointer_sqlite_aio() as checkpointer:
yield checkpointer
elif checkpointer_name == "postgres_aio":
async with _checkpointer_postgres_aio() as checkpointer:
yield checkpointer
elif checkpointer_name == "postgres_aio_pipe":
async with _checkpointer_postgres_aio_pipe() as checkpointer:
yield checkpointer
elif checkpointer_name == "postgres_aio_pool":
async with _checkpointer_postgres_aio_pool() as checkpointer:
yield checkpointer
else:
raise NotImplementedError(f"Unknown checkpointer: {checkpointer_name}")
ALL_CHECKPOINTERS_SYNC = [
"memory",
"sqlite",
"sqlite_aes",
"postgres",
"postgres_pipe",
"postgres_pool",
"sqlite_aes",
"postgres_shallow",
]
ALL_CHECKPOINTERS_SYNC = [
*REGULAR_CHECKPOINTERS_SYNC,
*SHALLOW_CHECKPOINTERS_SYNC,
]
SHALLOW_CHECKPOINTERS_ASYNC = ["postgres_aio_shallow"]
REGULAR_CHECKPOINTERS_ASYNC = [
ALL_CHECKPOINTERS_ASYNC = [
"memory",
"sqlite_aio",
"postgres_aio",
"postgres_aio_pipe",
"postgres_aio_pool",
]
ALL_CHECKPOINTERS_ASYNC = [
*REGULAR_CHECKPOINTERS_ASYNC,
*SHALLOW_CHECKPOINTERS_ASYNC,
]
ALL_CHECKPOINTERS_ASYNC_PLUS_NONE = [
*ALL_CHECKPOINTERS_ASYNC,
None,
"postgres_aio_shallow",
]
libs/langgraph/tests/conftest_checkpointer.py
+245
View File
@@ -0,0 +1,245 @@
import sys
from contextlib import asynccontextmanager, contextmanager
from uuid import uuid4
import pytest
from psycopg import AsyncConnection, Connection
from psycopg_pool import AsyncConnectionPool, ConnectionPool
from langgraph.checkpoint.postgres import PostgresSaver, ShallowPostgresSaver
from langgraph.checkpoint.postgres.aio import (
AsyncPostgresSaver,
AsyncShallowPostgresSaver,
)
from langgraph.checkpoint.serde.encrypted import EncryptedSerializer
from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
from tests.memory_assert import MemorySaverAssertImmutable
DEFAULT_POSTGRES_URI = "postgres://postgres:postgres@localhost:5442/"
@contextmanager
def _checkpointer_memory():
yield MemorySaverAssertImmutable()
@contextmanager
def _checkpointer_sqlite():
with SqliteSaver.from_conn_string(":memory:") as checkpointer:
yield checkpointer
@contextmanager
def _checkpointer_sqlite_aes():
with SqliteSaver.from_conn_string(":memory:") as checkpointer:
checkpointer.serde = EncryptedSerializer.from_pycryptodome_aes(
key=b"1234567890123456"
)
yield checkpointer
@contextmanager
def _checkpointer_postgres():
database = f"test_{uuid4().hex[:16]}"
# create unique db
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
conn.execute(f"CREATE DATABASE {database}")
try:
# yield checkpointer
with PostgresSaver.from_conn_string(
DEFAULT_POSTGRES_URI + database
) as checkpointer:
checkpointer.setup()
yield checkpointer
finally:
# drop unique db
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
conn.execute(f"DROP DATABASE {database}")
@contextmanager
def _checkpointer_postgres_shallow():
database = f"test_{uuid4().hex[:16]}"
# create unique db
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
conn.execute(f"CREATE DATABASE {database}")
try:
# yield checkpointer
with ShallowPostgresSaver.from_conn_string(
DEFAULT_POSTGRES_URI + database
) as checkpointer:
checkpointer.setup()
yield checkpointer
finally:
# drop unique db
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
conn.execute(f"DROP DATABASE {database}")
@contextmanager
def _checkpointer_postgres_pipe():
database = f"test_{uuid4().hex[:16]}"
# create unique db
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
conn.execute(f"CREATE DATABASE {database}")
try:
# yield checkpointer
with PostgresSaver.from_conn_string(
DEFAULT_POSTGRES_URI + database
) as checkpointer:
checkpointer.setup()
# setup can't run inside pipeline because of implicit transaction
with checkpointer.conn.pipeline() as pipe:
checkpointer.pipe = pipe
yield checkpointer
finally:
# drop unique db
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
conn.execute(f"DROP DATABASE {database}")
@contextmanager
def _checkpointer_postgres_pool():
database = f"test_{uuid4().hex[:16]}"
# create unique db
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
conn.execute(f"CREATE DATABASE {database}")
try:
# yield checkpointer
with ConnectionPool(
DEFAULT_POSTGRES_URI + database, max_size=10, kwargs={"autocommit": True}
) as pool:
checkpointer = PostgresSaver(pool)
checkpointer.setup()
yield checkpointer
finally:
# drop unique db
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
conn.execute(f"DROP DATABASE {database}")
@asynccontextmanager
async def _checkpointer_sqlite_aio():
async with AsyncSqliteSaver.from_conn_string(":memory:") as checkpointer:
yield checkpointer
@asynccontextmanager
async def _checkpointer_postgres_aio():
if sys.version_info < (3, 10):
pytest.skip("Async Postgres tests require Python 3.10+")
database = f"test_{uuid4().hex[:16]}"
# create unique db
async with await AsyncConnection.connect(
DEFAULT_POSTGRES_URI, autocommit=True
) as conn:
await conn.execute(f"CREATE DATABASE {database}")
try:
# yield checkpointer
async with AsyncPostgresSaver.from_conn_string(
DEFAULT_POSTGRES_URI + database
) as checkpointer:
await checkpointer.setup()
yield checkpointer
finally:
# drop unique db
async with await AsyncConnection.connect(
DEFAULT_POSTGRES_URI, autocommit=True
) as conn:
await conn.execute(f"DROP DATABASE {database}")
@asynccontextmanager
async def _checkpointer_postgres_aio_shallow():
if sys.version_info < (3, 10):
pytest.skip("Async Postgres tests require Python 3.10+")
database = f"test_{uuid4().hex[:16]}"
# create unique db
async with await AsyncConnection.connect(
DEFAULT_POSTGRES_URI, autocommit=True
) as conn:
await conn.execute(f"CREATE DATABASE {database}")
try:
# yield checkpointer
async with AsyncShallowPostgresSaver.from_conn_string(
DEFAULT_POSTGRES_URI + database
) as checkpointer:
await checkpointer.setup()
yield checkpointer
finally:
# drop unique db
async with await AsyncConnection.connect(
DEFAULT_POSTGRES_URI, autocommit=True
) as conn:
await conn.execute(f"DROP DATABASE {database}")
@asynccontextmanager
async def _checkpointer_postgres_aio_pipe():
if sys.version_info < (3, 10):
pytest.skip("Async Postgres tests require Python 3.10+")
database = f"test_{uuid4().hex[:16]}"
# create unique db
async with await AsyncConnection.connect(
DEFAULT_POSTGRES_URI, autocommit=True
) as conn:
await conn.execute(f"CREATE DATABASE {database}")
try:
# yield checkpointer
async with AsyncPostgresSaver.from_conn_string(
DEFAULT_POSTGRES_URI + database
) as checkpointer:
await checkpointer.setup()
# setup can't run inside pipeline because of implicit transaction
async with checkpointer.conn.pipeline() as pipe:
checkpointer.pipe = pipe
yield checkpointer
finally:
# drop unique db
async with await AsyncConnection.connect(
DEFAULT_POSTGRES_URI, autocommit=True
) as conn:
await conn.execute(f"DROP DATABASE {database}")
@asynccontextmanager
async def _checkpointer_postgres_aio_pool():
if sys.version_info < (3, 10):
pytest.skip("Async Postgres tests require Python 3.10+")
database = f"test_{uuid4().hex[:16]}"
# create unique db
async with await AsyncConnection.connect(
DEFAULT_POSTGRES_URI, autocommit=True
) as conn:
await conn.execute(f"CREATE DATABASE {database}")
try:
# yield checkpointer
async with AsyncConnectionPool(
DEFAULT_POSTGRES_URI + database, max_size=10, kwargs={"autocommit": True}
) as pool:
checkpointer = AsyncPostgresSaver(pool)
await checkpointer.setup()
yield checkpointer
finally:
# drop unique db
async with await AsyncConnection.connect(
DEFAULT_POSTGRES_URI, autocommit=True
) as conn:
await conn.execute(f"DROP DATABASE {database}")
__all__ = [
"_checkpointer_memory",
"_checkpointer_sqlite",
"_checkpointer_sqlite_aes",
"_checkpointer_postgres",
"_checkpointer_postgres_shallow",
"_checkpointer_postgres_pipe",
"_checkpointer_postgres_pool",
"_checkpointer_sqlite_aio",
"_checkpointer_postgres_aio",
"_checkpointer_postgres_aio_shallow",
"_checkpointer_postgres_aio_pipe",
"_checkpointer_postgres_aio_pool",
]
libs/langgraph/tests/test_checkpoint_migration.py
+98 -120
View File
@@ -17,11 +17,6 @@ from langgraph.types import Command, Interrupt, PregelTask, StateSnapshot, inter
from langgraph.utils.config import patch_configurable
from tests.any_int import AnyInt
from tests.any_str import AnyDict, AnyObject, AnyStr
from tests.conftest import (
REGULAR_CHECKPOINTERS_ASYNC,
REGULAR_CHECKPOINTERS_SYNC,
awith_checkpointer,
)
pytestmark = pytest.mark.anyio
@@ -1593,16 +1588,11 @@ def test_migrate_checkpoints(source: str, target: str) -> None:
@NEEDS_CONTEXTVARS
@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC)
def test_latest_checkpoint_state_graph(
request: pytest.FixtureRequest, checkpointer_name: str
sync_checkpointer: BaseCheckpointSaver,
) -> None:
checkpointer: BaseCheckpointSaver = request.getfixturevalue(
f"checkpointer_{checkpointer_name}"
)
builder = make_state_graph()
app = builder.compile(checkpointer=checkpointer)
app = builder.compile(checkpointer=sync_checkpointer)
config = {"configurable": {"thread_id": "1"}}
assert [*app.stream({"query": "what is weather in sf"}, config)] == [
@@ -1638,61 +1628,55 @@ def test_latest_checkpoint_state_graph(
@NEEDS_CONTEXTVARS
@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC)
async def test_latest_checkpoint_state_graph_async(checkpointer_name: str) -> None:
async with awith_checkpointer(checkpointer_name) as checkpointer:
builder = make_state_graph()
app = builder.compile(checkpointer=checkpointer)
config = {"configurable": {"thread_id": "1"}}
async def test_latest_checkpoint_state_graph_async(
async_checkpointer: BaseCheckpointSaver,
) -> None:
builder = make_state_graph()
app = builder.compile(checkpointer=async_checkpointer)
config = {"configurable": {"thread_id": "1"}}
assert [
c async for c in app.astream({"query": "what is weather in sf"}, config)
] == [
{"rewrite_query": {"query": "query: what is weather in sf"}},
{"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
{"retriever_two": {"docs": ["doc3", "doc4"]}},
{"retriever_one": {"docs": ["doc1", "doc2"]}},
{
"__interrupt__": (
Interrupt(
value="",
resumable=True,
ns=[AnyStr("qa:")],
),
)
},
]
assert [
c async for c in app.astream({"query": "what is weather in sf"}, config)
] == [
{"rewrite_query": {"query": "query: what is weather in sf"}},
{"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
{"retriever_two": {"docs": ["doc3", "doc4"]}},
{"retriever_one": {"docs": ["doc1", "doc2"]}},
{
"__interrupt__": (
Interrupt(
value="",
resumable=True,
ns=[AnyStr("qa:")],
),
)
},
]
assert [c async for c in app.astream(Command(resume=""), config)] == [
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
assert [c async for c in app.astream(Command(resume=""), config)] == [
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
# check history with current checkpoints matches expected history
history = [c async for c in app.aget_state_history(config)]
expected_history = get_expected_history()
assert len(history) == len(expected_history)
assert history[0] == expected_history[0]
assert history[1] == expected_history[1]
assert history[2] == expected_history[2]
assert history[3] == expected_history[3]
assert history[4] == expected_history[4]
assert history[5] == expected_history[5]
# check history with current checkpoints matches expected history
history = [c async for c in app.aget_state_history(config)]
expected_history = get_expected_history()
assert len(history) == len(expected_history)
assert history[0] == expected_history[0]
assert history[1] == expected_history[1]
assert history[2] == expected_history[2]
assert history[3] == expected_history[3]
assert history[4] == expected_history[4]
assert history[5] == expected_history[5]
@NEEDS_CONTEXTVARS
@pytest.mark.parametrize("checkpoint_version", ["3", "2-start:*", "2-quadratic"])
@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC)
def test_saved_checkpoint_state_graph(
request: pytest.FixtureRequest,
checkpointer_name: str,
sync_checkpointer: BaseCheckpointSaver,
checkpoint_version: str,
) -> None:
checkpointer: BaseCheckpointSaver = request.getfixturevalue(
f"checkpointer_{checkpointer_name}"
)
builder = make_state_graph()
app = builder.compile(checkpointer=checkpointer)
app = builder.compile(checkpointer=sync_checkpointer)
thread1 = "1"
config = {"configurable": {"thread_id": thread1, "checkpoint_ns": ""}}
@@ -1704,8 +1688,8 @@ def test_saved_checkpoint_state_graph(
for write in checkpoint.pending_writes:
grouped_writes[write[0]].append(write[1:])
for tid, group in grouped_writes.items():
checkpointer.put_writes(checkpoint.config, group, tid)
checkpointer.put(
sync_checkpointer.put_writes(checkpoint.config, group, tid)
sync_checkpointer.put(
patch_configurable(config, {"checkpoint_id": parent_id}),
checkpoint.checkpoint,
checkpoint.metadata,
@@ -1753,71 +1737,65 @@ def test_saved_checkpoint_state_graph(
@NEEDS_CONTEXTVARS
@pytest.mark.parametrize("checkpoint_version", ["3", "2-start:*", "2-quadratic"])
@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC)
async def test_saved_checkpoint_state_graph_async(
checkpointer_name: str,
async_checkpointer: BaseCheckpointSaver,
checkpoint_version: str,
) -> None:
async with awith_checkpointer(checkpointer_name) as checkpointer:
builder = make_state_graph()
app = builder.compile(checkpointer=checkpointer)
builder = make_state_graph()
app = builder.compile(checkpointer=async_checkpointer)
thread1 = "1"
config = {"configurable": {"thread_id": thread1, "checkpoint_ns": ""}}
thread1 = "1"
config = {"configurable": {"thread_id": thread1, "checkpoint_ns": ""}}
# save checkpoints
parent_id: Optional[str] = None
for checkpoint in reversed(SAVED_CHECKPOINTS[checkpoint_version]):
grouped_writes = defaultdict(list)
for write in checkpoint.pending_writes:
grouped_writes[write[0]].append(write[1:])
for tid, group in grouped_writes.items():
await checkpointer.aput_writes(checkpoint.config, group, tid)
await checkpointer.aput(
patch_configurable(config, {"checkpoint_id": parent_id}),
checkpoint.checkpoint,
checkpoint.metadata,
checkpoint.checkpoint["channel_versions"],
)
parent_id = checkpoint.checkpoint["id"]
# load history
history = [c async for c in app.aget_state_history(config)]
# check history with saved checkpoints matches expected history
exc_task_results: int = 0
if checkpoint_version == "2-start:*":
exc_task_results = 1
elif checkpoint_version == "2-quadratic":
exc_task_results = 2
expected_history = get_expected_history(exc_task_results=exc_task_results)
assert len(history) == len(expected_history)
assert history[0] == expected_history[0]
assert history[1] == expected_history[1]
assert history[2] == expected_history[2]
assert history[3] == expected_history[3]
assert history[4] == expected_history[4]
assert history[5] == expected_history[5]
# resume from 2nd to latest checkpoint
assert [
c async for c in app.astream(Command(resume=""), history[1].config)
] == [
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
# new checkpoint should match the latest checkpoint in history
latest_state = await app.aget_state(config)
assert (
StateSnapshot(
values=latest_state.values,
next=latest_state.next,
config=patch_configurable(
latest_state.config, {"checkpoint_id": AnyStr()}
),
metadata=AnyDict(latest_state.metadata),
created_at=AnyStr(),
parent_config=latest_state.parent_config,
tasks=latest_state.tasks,
interrupts=latest_state.interrupts,
)
== history[0]
# save checkpoints
parent_id: Optional[str] = None
for checkpoint in reversed(SAVED_CHECKPOINTS[checkpoint_version]):
grouped_writes = defaultdict(list)
for write in checkpoint.pending_writes:
grouped_writes[write[0]].append(write[1:])
for tid, group in grouped_writes.items():
await async_checkpointer.aput_writes(checkpoint.config, group, tid)
await async_checkpointer.aput(
patch_configurable(config, {"checkpoint_id": parent_id}),
checkpoint.checkpoint,
checkpoint.metadata,
checkpoint.checkpoint["channel_versions"],
)
parent_id = checkpoint.checkpoint["id"]
# load history
history = [c async for c in app.aget_state_history(config)]
# check history with saved checkpoints matches expected history
exc_task_results: int = 0
if checkpoint_version == "2-start:*":
exc_task_results = 1
elif checkpoint_version == "2-quadratic":
exc_task_results = 2
expected_history = get_expected_history(exc_task_results=exc_task_results)
assert len(history) == len(expected_history)
assert history[0] == expected_history[0]
assert history[1] == expected_history[1]
assert history[2] == expected_history[2]
assert history[3] == expected_history[3]
assert history[4] == expected_history[4]
assert history[5] == expected_history[5]
# resume from 2nd to latest checkpoint
assert [c async for c in app.astream(Command(resume=""), history[1].config)] == [
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
# new checkpoint should match the latest checkpoint in history
latest_state = await app.aget_state(config)
assert (
StateSnapshot(
values=latest_state.values,
next=latest_state.next,
config=patch_configurable(latest_state.config, {"checkpoint_id": AnyStr()}),
metadata=AnyDict(latest_state.metadata),
created_at=AnyStr(),
parent_config=latest_state.parent_config,
tasks=latest_state.tasks,
interrupts=latest_state.interrupts,
)
== history[0]
)
libs/langgraph/tests/test_interruption.py
+19 -29
View File
@@ -1,20 +1,14 @@
import pytest
from typing_extensions import TypedDict
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.graph import END, START, StateGraph
from tests.conftest import (
REGULAR_CHECKPOINTERS_ASYNC,
REGULAR_CHECKPOINTERS_SYNC,
awith_checkpointer,
)
pytestmark = pytest.mark.anyio
@pytest.mark.parametrize("checkpoint_during", [True, False])
@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC)
def test_interruption_without_state_updates(
request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool
sync_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
) -> None:
"""Test interruption without state updates. This test confirms that
interrupting doesn't require a state key having been updated in the prev step"""
@@ -34,8 +28,7 @@ def test_interruption_without_state_updates(
builder.add_edge("step_2", "step_3")
builder.add_edge("step_3", END)
checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
graph = builder.compile(checkpointer=checkpointer, interrupt_after="*")
graph = builder.compile(checkpointer=sync_checkpointer, interrupt_after="*")
initial_input = {"input": "hello world"}
thread = {"configurable": {"thread_id": "1"}}
@@ -56,10 +49,8 @@ def test_interruption_without_state_updates(
assert n_checkpoints == (5 if checkpoint_during else 3)
@pytest.mark.parametrize("checkpoint_during", [True, False])
@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC)
async def test_interruption_without_state_updates_async(
checkpointer_name: str, checkpoint_during: bool
async_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
) -> None:
"""Test interruption without state updates. This test confirms that
interrupting doesn't require a state key having been updated in the prev step"""
@@ -79,23 +70,22 @@ async def test_interruption_without_state_updates_async(
builder.add_edge("step_2", "step_3")
builder.add_edge("step_3", END)
async with awith_checkpointer(checkpointer_name) as checkpointer:
graph = builder.compile(checkpointer=checkpointer, interrupt_after="*")
graph = builder.compile(checkpointer=async_checkpointer, interrupt_after="*")
initial_input = {"input": "hello world"}
thread = {"configurable": {"thread_id": "1"}}
initial_input = {"input": "hello world"}
thread = {"configurable": {"thread_id": "1"}}
await graph.ainvoke(initial_input, thread, checkpoint_during=checkpoint_during)
assert (await graph.aget_state(thread)).next == ("step_2",)
n_checkpoints = len([c async for c in graph.aget_state_history(thread)])
assert n_checkpoints == (3 if checkpoint_during else 1)
await graph.ainvoke(initial_input, thread, checkpoint_during=checkpoint_during)
assert (await graph.aget_state(thread)).next == ("step_2",)
n_checkpoints = len([c async for c in graph.aget_state_history(thread)])
assert n_checkpoints == (3 if checkpoint_during else 1)
await graph.ainvoke(None, thread, checkpoint_during=checkpoint_during)
assert (await graph.aget_state(thread)).next == ("step_3",)
n_checkpoints = len([c async for c in graph.aget_state_history(thread)])
assert n_checkpoints == (4 if checkpoint_during else 2)
await graph.ainvoke(None, thread, checkpoint_during=checkpoint_during)
assert (await graph.aget_state(thread)).next == ("step_3",)
n_checkpoints = len([c async for c in graph.aget_state_history(thread)])
assert n_checkpoints == (4 if checkpoint_during else 2)
await graph.ainvoke(None, thread, checkpoint_during=checkpoint_during)
assert (await graph.aget_state(thread)).next == ()
n_checkpoints = len([c async for c in graph.aget_state_history(thread)])
assert n_checkpoints == (5 if checkpoint_during else 3)
await graph.ainvoke(None, thread, checkpoint_during=checkpoint_during)
assert (await graph.aget_state(thread)).next == ()
n_checkpoints = len([c async for c in graph.aget_state_history(thread)])
assert n_checkpoints == (5 if checkpoint_during else 3)
libs/langgraph/tests/test_large_cases.py
+5 -17
View File
@@ -41,11 +41,7 @@ from langgraph.types import (
from tests.agents import AgentAction, AgentFinish
from tests.any_int import AnyInt
from tests.any_str import AnyDict, AnyStr, UnsortedSequence
from tests.conftest import (
ALL_CHECKPOINTERS_SYNC,
REGULAR_CHECKPOINTERS_SYNC,
SHOULD_CHECK_SNAPSHOTS,
)
from tests.conftest import ALL_CHECKPOINTERS_SYNC, SHOULD_CHECK_SNAPSHOTS
from tests.fake_chat import FakeChatModel
from tests.fake_tracer import FakeTracer
from tests.messages import (
@@ -314,18 +310,16 @@ def test_invoke_two_processes_in_out_interrupt(
]
@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC)
def test_fork_always_re_runs_nodes(
request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture
sync_checkpointer: BaseCheckpointSaver, mocker: MockerFixture
) -> None:
checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
add_one = mocker.Mock(side_effect=lambda _: 1)
builder = StateGraph(Annotated[int, operator.add])
builder.add_node("add_one", add_one)
builder.add_edge(START, "add_one")
builder.add_conditional_edges("add_one", lambda cnt: "add_one" if cnt < 6 else END)
graph = builder.compile(checkpointer=checkpointer)
graph = builder.compile(checkpointer=sync_checkpointer)
thread1 = {"configurable": {"thread_id": "1"}}
@@ -7341,13 +7335,9 @@ def test_branch_then(
)
@pytest.mark.parametrize("checkpoint_during", [True, False])
@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC)
def test_send_dedupe_on_resume(
request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool
sync_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
) -> None:
checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
class InterruptOnce:
ticks: int = 0
@@ -7398,7 +7388,7 @@ def test_send_dedupe_on_resume(
builder.add_conditional_edges("1", send_for_fun)
builder.add_conditional_edges("2", route_to_three)
graph = builder.compile(checkpointer=checkpointer)
graph = builder.compile(checkpointer=sync_checkpointer)
thread1 = {"configurable": {"thread_id": "1"}}
assert graph.invoke(["0"], thread1, checkpoint_during=checkpoint_during) == {
"__interrupt__": [
@@ -7413,8 +7403,6 @@ def test_send_dedupe_on_resume(
assert builder.nodes["flaky"].runnable.func.ticks == 1
# check state
state = graph.get_state(thread1)
if "shallow" in checkpointer_name:
pytest.xfail("TODO: shallow checkpointer reports wrong next set")
assert state.next == ("flaky",)
# check history
history = [c for c in graph.get_state_history(thread1)]
libs/langgraph/tests/test_large_cases_async.py
+186 -194
View File
@@ -25,6 +25,7 @@ from typing_extensions import TypedDict
from langgraph.channels.context import Context
from langgraph.channels.last_value import LastValue
from langgraph.channels.untracked_value import UntrackedValue
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.constants import END, PULL, PUSH, START
from langgraph.graph.graph import Graph
from langgraph.graph.message import MessageGraph, add_messages
@@ -39,7 +40,6 @@ from tests.any_int import AnyInt
from tests.any_str import AnyDict, AnyStr, UnsortedSequence
from tests.conftest import (
ALL_CHECKPOINTERS_ASYNC,
REGULAR_CHECKPOINTERS_ASYNC,
awith_checkpointer,
)
from tests.fake_chat import FakeChatModel
@@ -327,9 +327,8 @@ async def test_invoke_two_processes_in_out_interrupt(
]
@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC)
async def test_fork_always_re_runs_nodes(
checkpointer_name: str, mocker: MockerFixture
async_checkpointer: BaseCheckpointSaver, mocker: MockerFixture
) -> None:
add_one = mocker.Mock(side_effect=lambda _: 1)
@@ -337,208 +336,201 @@ async def test_fork_always_re_runs_nodes(
builder.add_node("add_one", add_one)
builder.add_edge(START, "add_one")
builder.add_conditional_edges("add_one", lambda cnt: "add_one" if cnt < 6 else END)
async with awith_checkpointer(checkpointer_name) as checkpointer:
graph = builder.compile(checkpointer=checkpointer)
graph = builder.compile(checkpointer=async_checkpointer)
thread1 = {"configurable": {"thread_id": "1"}}
thread1 = {"configurable": {"thread_id": "1"}}
# start execution, stop at inbox
assert [
c
async for c in graph.astream(1, thread1, stream_mode=["values", "updates"])
] == [
("values", 1),
("updates", {"add_one": 1}),
("values", 2),
("updates", {"add_one": 1}),
("values", 3),
("updates", {"add_one": 1}),
("values", 4),
("updates", {"add_one": 1}),
("values", 5),
("updates", {"add_one": 1}),
("values", 6),
]
# start execution, stop at inbox
assert [
c async for c in graph.astream(1, thread1, stream_mode=["values", "updates"])
] == [
("values", 1),
("updates", {"add_one": 1}),
("values", 2),
("updates", {"add_one": 1}),
("values", 3),
("updates", {"add_one": 1}),
("values", 4),
("updates", {"add_one": 1}),
("values", 5),
("updates", {"add_one": 1}),
("values", 6),
]
# list history
history = [c async for c in graph.aget_state_history(thread1)]
assert history == [
StateSnapshot(
values=6,
next=(),
tasks=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"step": 5,
"writes": {"add_one": 1},
# list history
history = [c async for c in graph.aget_state_history(thread1)]
assert history == [
StateSnapshot(
values=6,
next=(),
tasks=(),
config={
"configurable": {
"thread_id": "1",
},
created_at=AnyStr(),
parent_config=history[1].config,
interrupts=(),
),
StateSnapshot(
values=5,
tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),),
next=("add_one",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"step": 4,
"writes": {"add_one": 1},
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"step": 5,
"writes": {"add_one": 1},
"thread_id": "1",
},
created_at=AnyStr(),
parent_config=history[1].config,
interrupts=(),
),
StateSnapshot(
values=5,
tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),),
next=("add_one",),
config={
"configurable": {
"thread_id": "1",
},
created_at=AnyStr(),
parent_config=history[2].config,
interrupts=(),
),
StateSnapshot(
values=4,
tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),),
next=("add_one",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"step": 3,
"writes": {"add_one": 1},
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"step": 4,
"writes": {"add_one": 1},
"thread_id": "1",
},
created_at=AnyStr(),
parent_config=history[2].config,
interrupts=(),
),
StateSnapshot(
values=4,
tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),),
next=("add_one",),
config={
"configurable": {
"thread_id": "1",
},
created_at=AnyStr(),
parent_config=history[3].config,
interrupts=(),
),
StateSnapshot(
values=3,
tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),),
next=("add_one",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"step": 2,
"writes": {"add_one": 1},
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"step": 3,
"writes": {"add_one": 1},
"thread_id": "1",
},
created_at=AnyStr(),
parent_config=history[3].config,
interrupts=(),
),
StateSnapshot(
values=3,
tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),),
next=("add_one",),
config={
"configurable": {
"thread_id": "1",
},
created_at=AnyStr(),
parent_config=history[4].config,
interrupts=(),
),
StateSnapshot(
values=2,
tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),),
next=("add_one",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"step": 1,
"writes": {"add_one": 1},
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"step": 2,
"writes": {"add_one": 1},
"thread_id": "1",
},
created_at=AnyStr(),
parent_config=history[4].config,
interrupts=(),
),
StateSnapshot(
values=2,
tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),),
next=("add_one",),
config={
"configurable": {
"thread_id": "1",
},
created_at=AnyStr(),
parent_config=history[5].config,
interrupts=(),
),
StateSnapshot(
values=1,
tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),),
next=("add_one",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"step": 0,
"writes": None,
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"step": 1,
"writes": {"add_one": 1},
"thread_id": "1",
},
created_at=AnyStr(),
parent_config=history[5].config,
interrupts=(),
),
StateSnapshot(
values=1,
tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),),
next=("add_one",),
config={
"configurable": {
"thread_id": "1",
},
created_at=AnyStr(),
parent_config=history[6].config,
interrupts=(),
),
StateSnapshot(
values=0,
tasks=(
PregelTask(AnyStr(), "__start__", (PULL, "__start__"), result=1),
),
next=("__start__",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "input",
"step": -1,
"writes": {"__start__": 1},
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"step": 0,
"writes": None,
"thread_id": "1",
},
created_at=AnyStr(),
parent_config=history[6].config,
interrupts=(),
),
StateSnapshot(
values=0,
tasks=(PregelTask(AnyStr(), "__start__", (PULL, "__start__"), result=1),),
next=("__start__",),
config={
"configurable": {
"thread_id": "1",
},
created_at=AnyStr(),
parent_config=None,
interrupts=(),
),
]
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "input",
"step": -1,
"writes": {"__start__": 1},
"thread_id": "1",
},
created_at=AnyStr(),
parent_config=None,
interrupts=(),
),
]
# forking from any previous checkpoint should re-run nodes
assert [
c
async for c in graph.astream(None, history[0].config, stream_mode="updates")
] == []
assert [
c
async for c in graph.astream(None, history[1].config, stream_mode="updates")
] == [
{"add_one": 1},
]
assert [
c
async for c in graph.astream(None, history[2].config, stream_mode="updates")
] == [
{"add_one": 1},
{"add_one": 1},
]
# forking from any previous checkpoint should re-run nodes
assert [
c async for c in graph.astream(None, history[0].config, stream_mode="updates")
] == []
assert [
c async for c in graph.astream(None, history[1].config, stream_mode="updates")
] == [
{"add_one": 1},
]
assert [
c async for c in graph.astream(None, history[2].config, stream_mode="updates")
] == [
{"add_one": 1},
{"add_one": 1},
]
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
libs/langgraph/tests/test_pregel.py
+15 -48
View File
@@ -69,11 +69,7 @@ from langgraph.types import (
)
from tests.agents import AgentAction, AgentFinish
from tests.any_str import AnyStr, AnyVersion, FloatBetween, UnsortedSequence
from tests.conftest import (
ALL_CHECKPOINTERS_SYNC,
REGULAR_CHECKPOINTERS_SYNC,
SHOULD_CHECK_SNAPSHOTS,
)
from tests.conftest import ALL_CHECKPOINTERS_SYNC, SHOULD_CHECK_SNAPSHOTS
from tests.messages import (
_AnyIdAIMessage,
_AnyIdAIMessageChunk,
@@ -675,12 +671,9 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
assert step == 2
@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC)
def test_run_from_checkpoint_id_retains_previous_writes(
request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture
sync_checkpointer: BaseCheckpointSaver,
) -> None:
checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
class MyState(TypedDict):
myval: Annotated[int, operator.add]
otherval: bool
@@ -713,7 +706,7 @@ def test_run_from_checkpoint_id_retains_previous_writes(
builder.add_conditional_edges("node_one", _getedge("node_one"))
builder.add_conditional_edges("node_two", _getedge("node_two"))
graph = builder.compile(checkpointer=checkpointer)
graph = builder.compile(checkpointer=sync_checkpointer)
thread_id = uuid.uuid4()
thread1 = {"configurable": {"thread_id": str(thread_id)}}
@@ -1105,7 +1098,6 @@ def test_invoke_checkpoint_two(
assert checkpoint["channel_values"].get("total") == 5
@pytest.mark.parametrize("checkpoint_during", [True, False])
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
def test_pending_writes_resume(
request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool
@@ -1504,7 +1496,6 @@ def test_send_sequences() -> None:
]
@pytest.mark.parametrize("checkpoint_during", [True, False])
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
def test_imp_task(
request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool
@@ -1602,7 +1593,6 @@ def test_imp_task(
assert mapper_calls == 2
@pytest.mark.parametrize("checkpoint_during", [True, False])
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
def test_imp_nested(
request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool
@@ -1676,7 +1666,6 @@ def test_imp_nested(
]
@pytest.mark.parametrize("checkpoint_during", [True, False])
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
def test_imp_stream_order(
request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool
@@ -3978,7 +3967,6 @@ def test_nested_graph(snapshot: SnapshotAssertion) -> None:
]
@pytest.mark.parametrize("checkpoint_during", [True, False])
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
def test_subgraph_checkpoint_true(
request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool
@@ -4047,7 +4035,6 @@ def test_subgraph_checkpoint_true(
]
@pytest.mark.parametrize("checkpoint_during", [True, False])
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
def test_subgraph_checkpoint_true_interrupt(
request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool
@@ -4229,7 +4216,6 @@ def test_stream_buffering_single_node(
]
@pytest.mark.parametrize("checkpoint_during", [True, False])
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
def test_nested_graph_interrupts_parallel(
request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool
@@ -4397,7 +4383,6 @@ def test_nested_graph_interrupts_parallel(
]
@pytest.mark.parametrize("checkpoint_during", [True, False])
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
def test_doubly_nested_graph_interrupts(
request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool
@@ -4741,9 +4726,8 @@ def test_checkpoint_metadata() -> None:
assert chkpnt_tuple.metadata["test_config_4"] == "bar"
@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC)
def test_remove_message_via_state_update(
request: pytest.FixtureRequest, checkpointer_name: str
sync_checkpointer: BaseCheckpointSaver,
) -> None:
from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage
@@ -4760,8 +4744,7 @@ def test_remove_message_via_state_update(
workflow.set_entry_point("chatbot")
workflow.add_edge("chatbot", END)
checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name)
app = workflow.compile(checkpointer=checkpointer)
app = workflow.compile(checkpointer=sync_checkpointer)
config = {"configurable": {"thread_id": "1"}}
output = app.invoke([HumanMessage(content="Hi")], config=config)
app.update_state(config, values=[RemoveMessage(id=output[-1].id)])
@@ -6202,10 +6185,8 @@ def test_concurrent_execution_thread_safety():
assert result["counter"] == 1
@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC)
def test_checkpoint_recovery(request: pytest.FixtureRequest, checkpointer_name: str):
def test_checkpoint_recovery(sync_checkpointer: BaseCheckpointSaver):
"""Test recovery from checkpoints after failures."""
checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
class State(TypedDict):
steps: Annotated[list[str], operator.add]
@@ -6226,7 +6207,7 @@ def test_checkpoint_recovery(request: pytest.FixtureRequest, checkpointer_name:
builder.add_edge(START, "node1")
builder.add_edge("node1", "node2")
graph = builder.compile(checkpointer=checkpointer)
graph = builder.compile(checkpointer=sync_checkpointer)
config = {"configurable": {"thread_id": "1"}}
# First attempt should fail
@@ -6244,9 +6225,6 @@ def test_checkpoint_recovery(request: pytest.FixtureRequest, checkpointer_name:
result = graph.invoke({"steps": [], "attempt": 2}, config)
assert result == {"steps": ["start", "node1", "node2"], "attempt": 2}
if "shallow" in checkpointer_name:
return
# Verify checkpoint history shows both attempts
history = list(graph.get_state_history(config))
assert len(history) == 6 # Initial + failed attempt + successful attempt
@@ -8345,12 +8323,9 @@ def test_pregel_loop_refcount():
gc.enable()
@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC)
def test_bulk_state_updates(
request: pytest.FixtureRequest, checkpointer_name: str
sync_checkpointer: BaseCheckpointSaver,
) -> None:
checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
class State(TypedDict):
foo: str
baz: str
@@ -8367,7 +8342,7 @@ def test_bulk_state_updates(
.add_node("node_b", node_b)
.add_edge(START, "node_a")
.add_edge("node_a", "node_b")
.compile(checkpointer=checkpointer)
.compile(checkpointer=sync_checkpointer)
)
config = {"configurable": {"thread_id": "1"}}
@@ -8397,7 +8372,7 @@ def test_bulk_state_updates(
assert state.values == {"foo": "updated", "baz": "new"}
# Check if there are only two checkpoints
checkpoints = list(checkpointer.list(config))
checkpoints = list(sync_checkpointer.list(config))
assert len(checkpoints) == 2
assert checkpoints[0].metadata["writes"] == {
"node_a": {"foo": "updated"},
@@ -8424,7 +8399,7 @@ def test_bulk_state_updates(
state = graph.get_state(config)
assert state.values == {"foo": "updated", "baz": "new"}
checkpoints = list(checkpointer.list(config))
checkpoints = list(sync_checkpointer.list(config))
assert len(checkpoints) == 2
assert checkpoints[0].metadata["writes"] == {
"node_a": {"foo": "updated"},
@@ -8489,12 +8464,7 @@ def test_pregel_node_copy() -> None:
graph.nodes["agent"].copy({})
@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC)
def test_update_as_input(
request: pytest.FixtureRequest, checkpointer_name: str
) -> None:
checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
def test_update_as_input(sync_checkpointer: BaseCheckpointSaver) -> None:
class State(TypedDict):
foo: str
@@ -8510,7 +8480,7 @@ def test_update_as_input(
.add_node("tool", tool)
.add_edge(START, "agent")
.add_edge("agent", "tool")
.compile(checkpointer=checkpointer)
.compile(checkpointer=sync_checkpointer)
)
assert graph.invoke({"foo": "input"}, {"configurable": {"thread_id": "1"}}) == {
@@ -8560,12 +8530,9 @@ def test_update_as_input(
assert new_history == history
@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC)
def test_batch_update_as_input(
request: pytest.FixtureRequest, checkpointer_name: str
sync_checkpointer: BaseCheckpointSaver,
) -> None:
checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
class State(TypedDict):
foo: str
tasks: Annotated[list[int], operator.add]
@@ -8593,7 +8560,7 @@ def test_batch_update_as_input(
.add_node("task", task)
.add_edge(START, "agent")
.add_edge("agent", "map")
.compile(checkpointer=checkpointer)
.compile(checkpointer=sync_checkpointer)
)
assert graph.invoke({"foo": "input"}, {"configurable": {"thread_id": "1"}}) == {
libs/langgraph/tests/test_pregel_async.py
+689 -734
View File
File diff suppressed because it is too large Load Diff