From eb57c068960a6adb466e078f7966be433b53ef3c Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Sat, 1 Mar 2025 13:53:27 -0800 Subject: [PATCH] Remove features and dependencies - rm langchain_core dependency - replace callbacks w run tree - rm Runnable dependency - rm non-state Graph - rm managed values - rm entrypoint/task/call - rm async methods - rm shallow checkpointer - rm messages stream mode - rm debug flag - rm remote graph --- .../langgraph/checkpoint/postgres/__init__.py | 29 +- .../checkpoint/postgres/_ainternal.py | 24 - .../langgraph/checkpoint/postgres/aio.py | 485 -- .../langgraph/checkpoint/postgres/base.py | 6 +- .../langgraph/checkpoint/postgres/shallow.py | 928 -- .../langgraph/store/postgres/__init__.py | 3 +- .../langgraph/store/postgres/aio.py | 434 - .../langgraph/store/postgres/base.py | 8 +- .../langgraph/checkpoint/base/__init__.py | 189 +- .../langgraph/checkpoint/memory/__init__.py | 129 +- .../langgraph/checkpoint/serde/jsonplus.py | 12 +- .../langgraph/store/base/__init__.py | 8 +- libs/checkpoint/langgraph/store/base/embed.py | 80 +- .../langgraph/store/memory/__init__.py | 3 +- libs/checkpoint/tests/embed_test_utils.py | 2 +- libs/langgraph/Makefile | 2 +- libs/langgraph/langgraph/channels/__init__.py | 2 - libs/langgraph/langgraph/channels/context.py | 5 - libs/langgraph/langgraph/config.py | 57 +- libs/langgraph/langgraph/constants.py | 2 +- libs/langgraph/langgraph/func/__init__.py | 443 - libs/langgraph/langgraph/func/py.typed | 0 libs/langgraph/langgraph/graph/__init__.py | 7 +- libs/langgraph/langgraph/graph/graph.py | 642 -- libs/langgraph/langgraph/graph/message.py | 164 +- libs/langgraph/langgraph/graph/state.py | 531 +- libs/langgraph/langgraph/managed/__init__.py | 3 - libs/langgraph/langgraph/managed/base.py | 104 - libs/langgraph/langgraph/managed/context.py | 108 - .../langgraph/managed/is_last_step.py | 19 - libs/langgraph/langgraph/managed/py.typed | 0 .../langgraph/managed/shared_value.py | 123 - libs/langgraph/langgraph/pregel/__init__.py | 1468 +--- libs/langgraph/langgraph/pregel/algo.py | 183 +- libs/langgraph/langgraph/pregel/call.py | 233 - libs/langgraph/langgraph/pregel/debug.py | 63 +- libs/langgraph/langgraph/pregel/executor.py | 163 +- libs/langgraph/langgraph/pregel/io.py | 22 +- libs/langgraph/langgraph/pregel/loop.py | 298 +- libs/langgraph/langgraph/pregel/manager.py | 101 +- libs/langgraph/langgraph/pregel/messages.py | 185 - libs/langgraph/langgraph/pregel/protocol.py | 86 +- libs/langgraph/langgraph/pregel/read.py | 136 +- libs/langgraph/langgraph/pregel/remote.py | 841 -- libs/langgraph/langgraph/pregel/retry.py | 94 +- libs/langgraph/langgraph/pregel/runner.py | 395 +- libs/langgraph/langgraph/pregel/utils.py | 20 +- libs/langgraph/langgraph/pregel/write.py | 71 +- libs/langgraph/langgraph/types.py | 25 +- libs/langgraph/langgraph/utils/config.py | 268 +- libs/langgraph/langgraph/utils/pydantic.py | 37 - libs/langgraph/langgraph/utils/queue.py | 38 +- libs/langgraph/langgraph/utils/runnable.py | 616 +- .../tests/__snapshots__/test_pregel.ambr | 1852 ---- libs/langgraph/tests/conftest.py | 288 +- libs/langgraph/tests/messages.py | 3 +- libs/langgraph/tests/test_algo.py | 4 +- libs/langgraph/tests/test_channels.py | 2 - libs/langgraph/tests/test_interruption.py | 42 - libs/langgraph/tests/test_large_cases.py | 3400 +------- .../langgraph/tests/test_large_cases_async.py | 7395 ---------------- libs/langgraph/tests/test_pregel.py | 2896 +------ libs/langgraph/tests/test_pregel_async.py | 7706 ----------------- libs/langgraph/tests/test_remote_graph.py | 798 -- libs/langgraph/tests/test_runnable.py | 261 - libs/langgraph/tests/test_state.py | 330 - libs/langgraph/tests/test_utils.py | 10 +- 67 files changed, 1301 insertions(+), 33581 deletions(-) delete mode 100644 libs/checkpoint-postgres/langgraph/checkpoint/postgres/_ainternal.py delete mode 100644 libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py delete mode 100644 libs/checkpoint-postgres/langgraph/checkpoint/postgres/shallow.py delete mode 100644 libs/checkpoint-postgres/langgraph/store/postgres/aio.py delete mode 100644 libs/langgraph/langgraph/channels/context.py delete mode 100644 libs/langgraph/langgraph/func/__init__.py delete mode 100644 libs/langgraph/langgraph/func/py.typed delete mode 100644 libs/langgraph/langgraph/graph/graph.py delete mode 100644 libs/langgraph/langgraph/managed/__init__.py delete mode 100644 libs/langgraph/langgraph/managed/base.py delete mode 100644 libs/langgraph/langgraph/managed/context.py delete mode 100644 libs/langgraph/langgraph/managed/is_last_step.py delete mode 100644 libs/langgraph/langgraph/managed/py.typed delete mode 100644 libs/langgraph/langgraph/managed/shared_value.py delete mode 100644 libs/langgraph/langgraph/pregel/call.py delete mode 100644 libs/langgraph/langgraph/pregel/messages.py delete mode 100644 libs/langgraph/langgraph/pregel/remote.py delete mode 100644 libs/langgraph/langgraph/utils/pydantic.py delete mode 100644 libs/langgraph/tests/__snapshots__/test_pregel.ambr delete mode 100644 libs/langgraph/tests/test_large_cases_async.py delete mode 100644 libs/langgraph/tests/test_pregel_async.py delete mode 100644 libs/langgraph/tests/test_remote_graph.py delete mode 100644 libs/langgraph/tests/test_runnable.py delete mode 100644 libs/langgraph/tests/test_state.py diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index 1c3426603..e6a4e09eb 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -3,7 +3,6 @@ from collections.abc import Iterator, Sequence from contextlib import contextmanager from typing import Any, Optional -from langchain_core.runnables import RunnableConfig from psycopg import Capabilities, Connection, Cursor, Pipeline from psycopg.rows import DictRow, dict_row from psycopg.types.json import Jsonb @@ -13,6 +12,7 @@ from langgraph.checkpoint.base import ( WRITES_IDX_MAP, ChannelVersions, Checkpoint, + CheckpointConfig, CheckpointMetadata, CheckpointTuple, get_checkpoint_id, @@ -20,7 +20,6 @@ from langgraph.checkpoint.base import ( ) from langgraph.checkpoint.postgres import _internal from langgraph.checkpoint.postgres.base import BasePostgresSaver -from langgraph.checkpoint.postgres.shallow import ShallowPostgresSaver from langgraph.checkpoint.serde.base import SerializerProtocol Conn = _internal.Conn # For backward compatibility @@ -97,10 +96,10 @@ class PostgresSaver(BasePostgresSaver): def list( self, - config: Optional[RunnableConfig], + config: Optional[CheckpointConfig], *, filter: Optional[dict[str, Any]] = None, - before: Optional[RunnableConfig] = None, + before: Optional[CheckpointConfig] = None, limit: Optional[int] = None, ) -> Iterator[CheckpointTuple]: """List checkpoints from the database. @@ -109,9 +108,9 @@ class PostgresSaver(BasePostgresSaver): on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first). Args: - config (RunnableConfig): The config to use for listing the checkpoints. + config (CheckpointConfig): The config to use for listing the checkpoints. filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata. Defaults to None. - before (Optional[RunnableConfig]): If provided, only checkpoints before the specified checkpoint ID are returned. Defaults to None. + before (Optional[CheckpointConfig]): If provided, only checkpoints before the specified checkpoint ID are returned. Defaults to None. limit (Optional[int]): The maximum number of checkpoints to return. Defaults to None. Yields: @@ -171,7 +170,7 @@ class PostgresSaver(BasePostgresSaver): self._load_writes(value["pending_writes"]), ) - def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: + def get_tuple(self, config: CheckpointConfig) -> Optional[CheckpointTuple]: """Get a checkpoint tuple from the database. This method retrieves a checkpoint tuple from the Postgres database based on the @@ -180,7 +179,7 @@ class PostgresSaver(BasePostgresSaver): for the given thread ID is retrieved. Args: - config (RunnableConfig): The config to use for retrieving the checkpoint. + config (CheckpointConfig): The config to use for retrieving the checkpoint. Returns: Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found. @@ -254,24 +253,24 @@ class PostgresSaver(BasePostgresSaver): def put( self, - config: RunnableConfig, + config: CheckpointConfig, checkpoint: Checkpoint, metadata: CheckpointMetadata, new_versions: ChannelVersions, - ) -> RunnableConfig: + ) -> CheckpointConfig: """Save a checkpoint to the database. This method saves a checkpoint to the Postgres database. The checkpoint is associated with the provided config and its parent config (if any). Args: - config (RunnableConfig): The config to associate with the checkpoint. + config (CheckpointConfig): The config to associate with the checkpoint. checkpoint (Checkpoint): The checkpoint to save. metadata (CheckpointMetadata): Additional metadata to save with the checkpoint. new_versions (ChannelVersions): New channel versions as of this write. Returns: - RunnableConfig: Updated configuration after storing the checkpoint. + CheckpointConfig: Updated configuration after storing the checkpoint. Examples: @@ -325,7 +324,7 @@ class PostgresSaver(BasePostgresSaver): def put_writes( self, - config: RunnableConfig, + config: CheckpointConfig, writes: Sequence[tuple[str, Any]], task_id: str, task_path: str = "", @@ -335,7 +334,7 @@ class PostgresSaver(BasePostgresSaver): This method saves intermediate writes associated with a checkpoint to the Postgres database. Args: - config (RunnableConfig): Configuration of the related checkpoint. + config (CheckpointConfig): Configuration of the related checkpoint. writes (List[Tuple[str, Any]]): List of writes to store. task_id (str): Identifier for the task creating the writes. """ @@ -400,4 +399,4 @@ class PostgresSaver(BasePostgresSaver): yield cur -__all__ = ["PostgresSaver", "BasePostgresSaver", "ShallowPostgresSaver", "Conn"] +__all__ = ["PostgresSaver", "BasePostgresSaver", "Conn"] diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/_ainternal.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/_ainternal.py deleted file mode 100644 index 33d299029..000000000 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/_ainternal.py +++ /dev/null @@ -1,24 +0,0 @@ -"""Shared async utility functions for the Postgres checkpoint & storage classes.""" - -from collections.abc import AsyncIterator -from contextlib import asynccontextmanager -from typing import Union - -from psycopg import AsyncConnection -from psycopg.rows import DictRow -from psycopg_pool import AsyncConnectionPool - -Conn = Union[AsyncConnection[DictRow], AsyncConnectionPool[AsyncConnection[DictRow]]] - - -@asynccontextmanager -async def get_connection( - conn: Conn, -) -> AsyncIterator[AsyncConnection[DictRow]]: - if isinstance(conn, AsyncConnection): - yield conn - elif isinstance(conn, AsyncConnectionPool): - async with conn.connection() as conn: - yield conn - else: - raise TypeError(f"Invalid connection type: {type(conn)}") diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py deleted file mode 100644 index a19b7bfc1..000000000 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ /dev/null @@ -1,485 +0,0 @@ -import asyncio -from collections.abc import AsyncIterator, Iterator, Sequence -from contextlib import asynccontextmanager -from typing import Any, Optional - -from langchain_core.runnables import RunnableConfig -from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline, Capabilities -from psycopg.rows import DictRow, dict_row -from psycopg.types.json import Jsonb -from psycopg_pool import AsyncConnectionPool - -from langgraph.checkpoint.base import ( - WRITES_IDX_MAP, - ChannelVersions, - Checkpoint, - CheckpointMetadata, - CheckpointTuple, - get_checkpoint_id, - get_checkpoint_metadata, -) -from langgraph.checkpoint.postgres import _ainternal -from langgraph.checkpoint.postgres.base import BasePostgresSaver -from langgraph.checkpoint.postgres.shallow import AsyncShallowPostgresSaver -from langgraph.checkpoint.serde.base import SerializerProtocol - -Conn = _ainternal.Conn # For backward compatibility - - -class AsyncPostgresSaver(BasePostgresSaver): - lock: asyncio.Lock - - def __init__( - self, - conn: _ainternal.Conn, - pipe: Optional[AsyncPipeline] = None, - serde: Optional[SerializerProtocol] = None, - ) -> None: - super().__init__(serde=serde) - if isinstance(conn, AsyncConnectionPool) and pipe is not None: - raise ValueError( - "Pipeline should be used only with a single AsyncConnection, not AsyncConnectionPool." - ) - - self.conn = conn - self.pipe = pipe - self.lock = asyncio.Lock() - self.loop = asyncio.get_running_loop() - self.supports_pipeline = Capabilities().has_pipeline() - - @classmethod - @asynccontextmanager - async def from_conn_string( - cls, - conn_string: str, - *, - pipeline: bool = False, - serde: Optional[SerializerProtocol] = None, - ) -> AsyncIterator["AsyncPostgresSaver"]: - """Create a new AsyncPostgresSaver instance from a connection string. - - Args: - conn_string (str): The Postgres connection info string. - pipeline (bool): whether to use AsyncPipeline - - Returns: - AsyncPostgresSaver: A new AsyncPostgresSaver instance. - """ - async with await AsyncConnection.connect( - conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row - ) as conn: - if pipeline: - async with conn.pipeline() as pipe: - yield cls(conn=conn, pipe=pipe, serde=serde) - else: - yield cls(conn=conn, serde=serde) - - async def setup(self) -> None: - """Set up the checkpoint database asynchronously. - - This method creates the necessary tables in the Postgres database if they don't - already exist and runs database migrations. It MUST be called directly by the user - the first time checkpointer is used. - """ - async with self._cursor() as cur: - await cur.execute(self.MIGRATIONS[0]) - results = await cur.execute( - "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1" - ) - row = await results.fetchone() - if row is None: - version = -1 - else: - version = row["v"] - for v, migration in zip( - range(version + 1, len(self.MIGRATIONS)), - self.MIGRATIONS[version + 1 :], - ): - await cur.execute(migration) - await cur.execute(f"INSERT INTO checkpoint_migrations (v) VALUES ({v})") - if self.pipe: - await self.pipe.sync() - - async def alist( - self, - config: Optional[RunnableConfig], - *, - filter: Optional[dict[str, Any]] = None, - before: Optional[RunnableConfig] = None, - limit: Optional[int] = None, - ) -> AsyncIterator[CheckpointTuple]: - """List checkpoints from the database asynchronously. - - This method retrieves a list of checkpoint tuples from the Postgres database based - on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first). - - Args: - config (Optional[RunnableConfig]): Base configuration for filtering checkpoints. - filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata. - before (Optional[RunnableConfig]): If provided, only checkpoints before the specified checkpoint ID are returned. Defaults to None. - limit (Optional[int]): Maximum number of checkpoints to return. - - Yields: - AsyncIterator[CheckpointTuple]: An asynchronous iterator of matching checkpoint tuples. - """ - where, args = self._search_where(config, filter, before) - query = self.SELECT_SQL + where + " ORDER BY checkpoint_id DESC" - if limit: - query += f" LIMIT {limit}" - # if we change this to use .stream() we need to make sure to close the cursor - async with self._cursor() as cur: - await cur.execute(query, args, binary=True) - async for value in cur: - yield CheckpointTuple( - { - "configurable": { - "thread_id": value["thread_id"], - "checkpoint_ns": value["checkpoint_ns"], - "checkpoint_id": value["checkpoint_id"], - } - }, - await asyncio.to_thread( - self._load_checkpoint, - value["checkpoint"], - value["channel_values"], - value["pending_sends"], - ), - self._load_metadata(value["metadata"]), - ( - { - "configurable": { - "thread_id": value["thread_id"], - "checkpoint_ns": value["checkpoint_ns"], - "checkpoint_id": value["parent_checkpoint_id"], - } - } - if value["parent_checkpoint_id"] - else None - ), - await asyncio.to_thread(self._load_writes, value["pending_writes"]), - ) - - async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: - """Get a checkpoint tuple from the database asynchronously. - - This method retrieves a checkpoint tuple from the Postgres database based on the - provided config. If the config contains a "checkpoint_id" key, the checkpoint with - the matching thread ID and "checkpoint_id" is retrieved. Otherwise, the latest checkpoint - for the given thread ID is retrieved. - - Args: - config (RunnableConfig): The config to use for retrieving the checkpoint. - - Returns: - Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found. - """ - thread_id = config["configurable"]["thread_id"] - checkpoint_id = get_checkpoint_id(config) - checkpoint_ns = config["configurable"].get("checkpoint_ns", "") - if 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) - where = "WHERE thread_id = %s AND checkpoint_ns = %s ORDER BY checkpoint_id DESC LIMIT 1" - - async with self._cursor() as cur: - await cur.execute( - self.SELECT_SQL + where, - args, - binary=True, - ) - - async for value in cur: - return CheckpointTuple( - { - "configurable": { - "thread_id": thread_id, - "checkpoint_ns": checkpoint_ns, - "checkpoint_id": value["checkpoint_id"], - } - }, - await asyncio.to_thread( - self._load_checkpoint, - value["checkpoint"], - value["channel_values"], - value["pending_sends"], - ), - self._load_metadata(value["metadata"]), - ( - { - "configurable": { - "thread_id": thread_id, - "checkpoint_ns": checkpoint_ns, - "checkpoint_id": value["parent_checkpoint_id"], - } - } - if value["parent_checkpoint_id"] - else None - ), - await asyncio.to_thread(self._load_writes, value["pending_writes"]), - ) - - async def aput( - self, - config: RunnableConfig, - checkpoint: Checkpoint, - metadata: CheckpointMetadata, - new_versions: ChannelVersions, - ) -> RunnableConfig: - """Save a checkpoint to the database asynchronously. - - This method saves a checkpoint to the Postgres database. The checkpoint is associated - with the provided config and its parent config (if any). - - Args: - config (RunnableConfig): The config to associate with the checkpoint. - checkpoint (Checkpoint): The checkpoint to save. - metadata (CheckpointMetadata): Additional metadata to save with the checkpoint. - new_versions (ChannelVersions): New channel versions as of this write. - - Returns: - RunnableConfig: Updated configuration after storing the checkpoint. - """ - configurable = config["configurable"].copy() - thread_id = configurable.pop("thread_id") - checkpoint_ns = configurable.pop("checkpoint_ns") - checkpoint_id = configurable.pop( - "checkpoint_id", configurable.pop("thread_ts", None) - ) - - copy = checkpoint.copy() - next_config = { - "configurable": { - "thread_id": thread_id, - "checkpoint_ns": checkpoint_ns, - "checkpoint_id": checkpoint["id"], - } - } - - async with self._cursor(pipeline=True) as cur: - await cur.executemany( - self.UPSERT_CHECKPOINT_BLOBS_SQL, - await asyncio.to_thread( - self._dump_blobs, - thread_id, - checkpoint_ns, - copy.pop("channel_values"), # type: ignore[misc] - new_versions, - ), - ) - await cur.execute( - self.UPSERT_CHECKPOINTS_SQL, - ( - thread_id, - checkpoint_ns, - checkpoint["id"], - checkpoint_id, - Jsonb(self._dump_checkpoint(copy)), - self._dump_metadata(get_checkpoint_metadata(config, metadata)), - ), - ) - return next_config - - async def aput_writes( - self, - config: RunnableConfig, - writes: Sequence[tuple[str, Any]], - task_id: str, - task_path: str = "", - ) -> None: - """Store intermediate writes linked to a checkpoint asynchronously. - - This method saves intermediate writes associated with a checkpoint to the database. - - Args: - config (RunnableConfig): Configuration of the related checkpoint. - writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair. - task_id (str): Identifier for the task creating the writes. - """ - query = ( - self.UPSERT_CHECKPOINT_WRITES_SQL - if all(w[0] in WRITES_IDX_MAP for w in writes) - else self.INSERT_CHECKPOINT_WRITES_SQL - ) - params = await asyncio.to_thread( - self._dump_writes, - config["configurable"]["thread_id"], - config["configurable"]["checkpoint_ns"], - config["configurable"]["checkpoint_id"], - task_id, - task_path, - writes, - ) - async with self._cursor(pipeline=True) as cur: - await cur.executemany(query, params) - - @asynccontextmanager - async def _cursor( - self, *, pipeline: bool = False - ) -> AsyncIterator[AsyncCursor[DictRow]]: - """Create a database cursor as a context manager. - - Args: - pipeline (bool): whether to use pipeline for the DB operations inside the context manager. - Will be applied regardless of whether the AsyncPostgresSaver instance was initialized with a pipeline. - If pipeline mode is not supported, will fall back to using transaction context manager. - """ - async with _ainternal.get_connection(self.conn) as conn: - if self.pipe: - # a connection in pipeline mode can be used concurrently - # in multiple threads/coroutines, but only one cursor can be - # used at a time - try: - async with conn.cursor(binary=True, row_factory=dict_row) as cur: - yield cur - finally: - if pipeline: - await self.pipe.sync() - elif pipeline: - # a connection not in pipeline mode can only be used by one - # thread/coroutine at a time, so we acquire a lock - if self.supports_pipeline: - async with ( - self.lock, - conn.pipeline(), - conn.cursor(binary=True, row_factory=dict_row) as cur, - ): - yield cur - else: - # Use connection's transaction context manager when pipeline mode not supported - async with ( - self.lock, - conn.transaction(), - conn.cursor(binary=True, row_factory=dict_row) as cur, - ): - yield cur - else: - async with ( - self.lock, - conn.cursor(binary=True, row_factory=dict_row) as cur, - ): - yield cur - - def list( - self, - config: Optional[RunnableConfig], - *, - filter: Optional[dict[str, Any]] = None, - before: Optional[RunnableConfig] = None, - limit: Optional[int] = None, - ) -> Iterator[CheckpointTuple]: - """List checkpoints from the database. - - This method retrieves a list of checkpoint tuples from the Postgres database based - on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first). - - Args: - config (Optional[RunnableConfig]): Base configuration for filtering checkpoints. - filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata. - before (Optional[RunnableConfig]): If provided, only checkpoints before the specified checkpoint ID are returned. Defaults to None. - limit (Optional[int]): Maximum number of checkpoints to return. - - Yields: - Iterator[CheckpointTuple]: An iterator of matching checkpoint tuples. - """ - try: - # check if we are in the main thread, only bg threads can block - # we don't check in other methods to avoid the overhead - if asyncio.get_running_loop() is self.loop: - raise asyncio.InvalidStateError( - "Synchronous calls to AsyncPostgresSaver are only allowed from a " - "different thread. From the main thread, use the async interface. " - "For example, use `checkpointer.alist(...)` or `await " - "graph.ainvoke(...)`." - ) - except RuntimeError: - pass - aiter_ = self.alist(config, filter=filter, before=before, limit=limit) - while True: - try: - yield asyncio.run_coroutine_threadsafe( - anext(aiter_), # noqa: F821 - self.loop, - ).result() - except StopAsyncIteration: - break - - def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: - """Get a checkpoint tuple from the database. - - This method retrieves a checkpoint tuple from the Postgres database based on the - provided config. If the config contains a "checkpoint_id" key, the checkpoint with - the matching thread ID and "checkpoint_id" is retrieved. Otherwise, the latest checkpoint - for the given thread ID is retrieved. - - Args: - config (RunnableConfig): The config to use for retrieving the checkpoint. - - Returns: - Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found. - """ - try: - # check if we are in the main thread, only bg threads can block - # we don't check in other methods to avoid the overhead - if asyncio.get_running_loop() is self.loop: - raise asyncio.InvalidStateError( - "Synchronous calls to AsyncPostgresSaver are only allowed from a " - "different thread. From the main thread, use the async interface. " - "For example, use `await checkpointer.aget_tuple(...)` or `await " - "graph.ainvoke(...)`." - ) - except RuntimeError: - pass - return asyncio.run_coroutine_threadsafe( - self.aget_tuple(config), self.loop - ).result() - - def put( - self, - config: RunnableConfig, - checkpoint: Checkpoint, - metadata: CheckpointMetadata, - new_versions: ChannelVersions, - ) -> RunnableConfig: - """Save a checkpoint to the database. - - This method saves a checkpoint to the Postgres database. The checkpoint is associated - with the provided config and its parent config (if any). - - Args: - config (RunnableConfig): The config to associate with the checkpoint. - checkpoint (Checkpoint): The checkpoint to save. - metadata (CheckpointMetadata): Additional metadata to save with the checkpoint. - new_versions (ChannelVersions): New channel versions as of this write. - - Returns: - RunnableConfig: Updated configuration after storing the checkpoint. - """ - return asyncio.run_coroutine_threadsafe( - self.aput(config, checkpoint, metadata, new_versions), self.loop - ).result() - - def put_writes( - self, - config: RunnableConfig, - writes: Sequence[tuple[str, Any]], - task_id: str, - task_path: str = "", - ) -> None: - """Store intermediate writes linked to a checkpoint. - - This method saves intermediate writes associated with a checkpoint to the database. - - Args: - config (RunnableConfig): Configuration of the related checkpoint. - writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair. - task_id (str): Identifier for the task creating the writes. - task_path (str): Path of the task creating the writes. - """ - return asyncio.run_coroutine_threadsafe( - self.aput_writes(config, writes, task_id, task_path), self.loop - ).result() - - -__all__ = ["AsyncPostgresSaver", "AsyncShallowPostgresSaver", "Conn"] diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py index 0d901a78c..b04c1bd9e 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py @@ -2,7 +2,6 @@ import random from collections.abc import Sequence from typing import Any, Optional, cast -from langchain_core.runnables import RunnableConfig from psycopg.types.json import Jsonb from langgraph.checkpoint.base import ( @@ -10,6 +9,7 @@ from langgraph.checkpoint.base import ( BaseCheckpointSaver, ChannelVersions, Checkpoint, + CheckpointConfig, CheckpointMetadata, get_checkpoint_id, ) @@ -262,9 +262,9 @@ class BasePostgresSaver(BaseCheckpointSaver[str]): def _search_where( self, - config: Optional[RunnableConfig], + config: Optional[CheckpointConfig], filter: MetadataInput, - before: Optional[RunnableConfig] = None, + before: Optional[CheckpointConfig] = None, ) -> tuple[str, list[Any]]: """Return WHERE clause predicates for alist() given config, filter, before. diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/shallow.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/shallow.py deleted file mode 100644 index 90e140b2c..000000000 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/shallow.py +++ /dev/null @@ -1,928 +0,0 @@ -import asyncio -import threading -from collections.abc import AsyncIterator, Iterator, Sequence -from contextlib import asynccontextmanager, contextmanager -from typing import Any, Optional - -from langchain_core.runnables import RunnableConfig -from psycopg import ( - AsyncConnection, - AsyncCursor, - AsyncPipeline, - Capabilities, - Connection, - Cursor, - Pipeline, -) -from psycopg.rows import DictRow, dict_row -from psycopg.types.json import Jsonb -from psycopg_pool import AsyncConnectionPool, ConnectionPool - -from langgraph.checkpoint.base import ( - WRITES_IDX_MAP, - ChannelVersions, - Checkpoint, - CheckpointMetadata, - CheckpointTuple, - get_checkpoint_metadata, -) -from langgraph.checkpoint.postgres import _ainternal, _internal -from langgraph.checkpoint.postgres.base import BasePostgresSaver -from langgraph.checkpoint.serde.base import SerializerProtocol -from langgraph.checkpoint.serde.types import TASKS - -""" -To add a new migration, add a new string to the MIGRATIONS list. -The position of the migration in the list is the version number. -""" -MIGRATIONS = [ - """CREATE TABLE IF NOT EXISTS checkpoint_migrations ( - v INTEGER PRIMARY KEY -);""", - """CREATE TABLE IF NOT EXISTS checkpoints ( - thread_id TEXT NOT NULL, - checkpoint_ns TEXT NOT NULL DEFAULT '', - type TEXT, - checkpoint JSONB NOT NULL, - metadata JSONB NOT NULL DEFAULT '{}', - PRIMARY KEY (thread_id, checkpoint_ns) -);""", - """CREATE TABLE IF NOT EXISTS checkpoint_blobs ( - thread_id TEXT NOT NULL, - checkpoint_ns TEXT NOT NULL DEFAULT '', - channel TEXT NOT NULL, - type TEXT NOT NULL, - blob BYTEA, - PRIMARY KEY (thread_id, checkpoint_ns, channel) -);""", - """CREATE TABLE IF NOT EXISTS checkpoint_writes ( - thread_id TEXT NOT NULL, - checkpoint_ns TEXT NOT NULL DEFAULT '', - checkpoint_id TEXT NOT NULL, - task_id TEXT NOT NULL, - idx INTEGER NOT NULL, - channel TEXT NOT NULL, - type TEXT, - blob BYTEA NOT NULL, - PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) -);""", - """ - CREATE INDEX CONCURRENTLY IF NOT EXISTS checkpoints_thread_id_idx ON checkpoints(thread_id); - """, - """ - CREATE INDEX CONCURRENTLY IF NOT EXISTS checkpoint_blobs_thread_id_idx ON checkpoint_blobs(thread_id); - """, - """ - CREATE INDEX CONCURRENTLY IF NOT EXISTS checkpoint_writes_thread_id_idx ON checkpoint_writes(thread_id); - """, - """ - ALTER TABLE checkpoint_writes ADD COLUMN task_path TEXT NOT NULL DEFAULT ''; - """, -] - -SELECT_SQL = f""" -select - thread_id, - checkpoint, - checkpoint_ns, - metadata, - ( - select array_agg(array[bl.channel::bytea, bl.type::bytea, bl.blob]) - from jsonb_each_text(checkpoint -> 'channel_versions') - inner join checkpoint_blobs bl - on bl.thread_id = checkpoints.thread_id - and bl.checkpoint_ns = checkpoints.checkpoint_ns - and bl.channel = jsonb_each_text.key - ) as channel_values, - ( - select - array_agg(array[cw.task_id::text::bytea, cw.channel::bytea, cw.type::bytea, cw.blob] order by cw.task_id, cw.idx) - from checkpoint_writes cw - where cw.thread_id = checkpoints.thread_id - and cw.checkpoint_ns = checkpoints.checkpoint_ns - and cw.checkpoint_id = (checkpoint->>'id') - ) as pending_writes, - ( - select array_agg(array[cw.type::bytea, cw.blob] order by cw.task_path, cw.task_id, cw.idx) - from checkpoint_writes cw - where cw.thread_id = checkpoints.thread_id - and cw.checkpoint_ns = checkpoints.checkpoint_ns - and cw.channel = '{TASKS}' - ) as pending_sends -from checkpoints """ - -UPSERT_CHECKPOINT_BLOBS_SQL = """ - INSERT INTO checkpoint_blobs (thread_id, checkpoint_ns, channel, type, blob) - VALUES (%s, %s, %s, %s, %s) - ON CONFLICT (thread_id, checkpoint_ns, channel) DO UPDATE SET - type = EXCLUDED.type, - blob = EXCLUDED.blob; -""" - -UPSERT_CHECKPOINTS_SQL = """ - INSERT INTO checkpoints (thread_id, checkpoint_ns, checkpoint, metadata) - VALUES (%s, %s, %s, %s) - ON CONFLICT (thread_id, checkpoint_ns) - DO UPDATE SET - checkpoint = EXCLUDED.checkpoint, - metadata = EXCLUDED.metadata; -""" - -UPSERT_CHECKPOINT_WRITES_SQL = """ - INSERT INTO checkpoint_writes (thread_id, checkpoint_ns, checkpoint_id, task_id, task_path, idx, channel, type, blob) - VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s) - ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) DO UPDATE SET - channel = EXCLUDED.channel, - type = EXCLUDED.type, - blob = EXCLUDED.blob; -""" - -INSERT_CHECKPOINT_WRITES_SQL = """ - INSERT INTO checkpoint_writes (thread_id, checkpoint_ns, checkpoint_id, task_id, task_path, idx, channel, type, blob) - VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s) - ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) DO NOTHING -""" - - -def _dump_blobs( - serde: SerializerProtocol, - thread_id: str, - checkpoint_ns: str, - values: dict[str, Any], - versions: ChannelVersions, -) -> list[tuple[str, str, str, str, str, Optional[bytes]]]: - if not versions: - return [] - - return [ - ( - thread_id, - checkpoint_ns, - k, - *(serde.dumps_typed(values[k]) if k in values else ("empty", None)), - ) - for k in versions - ] - - -class ShallowPostgresSaver(BasePostgresSaver): - """A checkpoint saver that uses Postgres to store checkpoints. - - This checkpointer ONLY stores the most recent checkpoint and does NOT retain any history. - It is meant to be a light-weight drop-in replacement for the PostgresSaver that - supports most of the LangGraph persistence functionality with the exception of time travel. - """ - - SELECT_SQL = SELECT_SQL - MIGRATIONS = MIGRATIONS - UPSERT_CHECKPOINT_BLOBS_SQL = UPSERT_CHECKPOINT_BLOBS_SQL - UPSERT_CHECKPOINTS_SQL = UPSERT_CHECKPOINTS_SQL - UPSERT_CHECKPOINT_WRITES_SQL = UPSERT_CHECKPOINT_WRITES_SQL - INSERT_CHECKPOINT_WRITES_SQL = INSERT_CHECKPOINT_WRITES_SQL - - lock: threading.Lock - - def __init__( - self, - conn: _internal.Conn, - pipe: Optional[Pipeline] = None, - serde: Optional[SerializerProtocol] = None, - ) -> None: - super().__init__(serde=serde) - if isinstance(conn, ConnectionPool) and pipe is not None: - raise ValueError( - "Pipeline should be used only with a single Connection, not ConnectionPool." - ) - - self.conn = conn - self.pipe = pipe - self.lock = threading.Lock() - self.supports_pipeline = Capabilities().has_pipeline() - - @classmethod - @contextmanager - def from_conn_string( - cls, conn_string: str, *, pipeline: bool = False - ) -> Iterator["ShallowPostgresSaver"]: - """Create a new ShallowPostgresSaver instance from a connection string. - - Args: - conn_string (str): The Postgres connection info string. - pipeline (bool): whether to use Pipeline - - Returns: - ShallowPostgresSaver: A new ShallowPostgresSaver instance. - """ - with Connection.connect( - conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row - ) as conn: - if pipeline: - with conn.pipeline() as pipe: - yield cls(conn, pipe) - else: - yield cls(conn) - - def setup(self) -> None: - """Set up the checkpoint database asynchronously. - - This method creates the necessary tables in the Postgres database if they don't - already exist and runs database migrations. It MUST be called directly by the user - the first time checkpointer is used. - """ - with self._cursor() as cur: - cur.execute(self.MIGRATIONS[0]) - results = cur.execute( - "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1" - ) - row = results.fetchone() - if row is None: - version = -1 - else: - version = row["v"] - for v, migration in zip( - range(version + 1, len(self.MIGRATIONS)), - self.MIGRATIONS[version + 1 :], - ): - cur.execute(migration) - cur.execute(f"INSERT INTO checkpoint_migrations (v) VALUES ({v})") - if self.pipe: - self.pipe.sync() - - def list( - self, - config: Optional[RunnableConfig], - *, - filter: Optional[dict[str, Any]] = None, - before: Optional[RunnableConfig] = None, - limit: Optional[int] = None, - ) -> Iterator[CheckpointTuple]: - """List checkpoints from the database. - - This method retrieves a list of checkpoint tuples from the Postgres database based - on the provided config. For ShallowPostgresSaver, this method returns a list with - ONLY the most recent checkpoint. - """ - where, args = self._search_where(config, filter, before) - query = self.SELECT_SQL + where - if limit: - query += f" LIMIT {limit}" - with self._cursor() as cur: - cur.execute(self.SELECT_SQL + where, args, binary=True) - for value in cur: - checkpoint = self._load_checkpoint( - value["checkpoint"], - value["channel_values"], - value["pending_sends"], - ) - yield CheckpointTuple( - config={ - "configurable": { - "thread_id": value["thread_id"], - "checkpoint_ns": value["checkpoint_ns"], - "checkpoint_id": checkpoint["id"], - } - }, - checkpoint=checkpoint, - metadata=self._load_metadata(value["metadata"]), - pending_writes=self._load_writes(value["pending_writes"]), - ) - - def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: - """Get a checkpoint tuple from the database. - - This method retrieves a checkpoint tuple from the Postgres database based on the - provided config (matching the thread ID in the config). - - Args: - config (RunnableConfig): The config to use for retrieving the checkpoint. - - Returns: - Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found. - - Examples: - - Basic: - >>> config = {"configurable": {"thread_id": "1"}} - >>> checkpoint_tuple = memory.get_tuple(config) - >>> print(checkpoint_tuple) - CheckpointTuple(...) - - With timestamp: - - >>> config = { - ... "configurable": { - ... "thread_id": "1", - ... "checkpoint_ns": "", - ... "checkpoint_id": "1ef4f797-8335-6428-8001-8a1503f9b875", - ... } - ... } - >>> checkpoint_tuple = memory.get_tuple(config) - >>> print(checkpoint_tuple) - CheckpointTuple(...) - """ # noqa - thread_id = config["configurable"]["thread_id"] - checkpoint_ns = config["configurable"].get("checkpoint_ns", "") - args = (thread_id, checkpoint_ns) - where = "WHERE thread_id = %s AND checkpoint_ns = %s" - - with self._cursor() as cur: - cur.execute( - self.SELECT_SQL + where, - args, - binary=True, - ) - - for value in cur: - checkpoint = self._load_checkpoint( - value["checkpoint"], - value["channel_values"], - value["pending_sends"], - ) - return CheckpointTuple( - config={ - "configurable": { - "thread_id": thread_id, - "checkpoint_ns": checkpoint_ns, - "checkpoint_id": checkpoint["id"], - } - }, - checkpoint=checkpoint, - metadata=self._load_metadata(value["metadata"]), - pending_writes=self._load_writes(value["pending_writes"]), - ) - - def put( - self, - config: RunnableConfig, - checkpoint: Checkpoint, - metadata: CheckpointMetadata, - new_versions: ChannelVersions, - ) -> RunnableConfig: - """Save a checkpoint to the database. - - This method saves a checkpoint to the Postgres database. The checkpoint is associated - with the provided config. For ShallowPostgresSaver, this method saves ONLY the most recent - checkpoint and overwrites a previous checkpoint, if it exists. - - Args: - config (RunnableConfig): The config to associate with the checkpoint. - checkpoint (Checkpoint): The checkpoint to save. - metadata (CheckpointMetadata): Additional metadata to save with the checkpoint. - new_versions (ChannelVersions): New channel versions as of this write. - - Returns: - RunnableConfig: Updated configuration after storing the checkpoint. - - Examples: - - >>> from langgraph.checkpoint.postgres import ShallowPostgresSaver - >>> DB_URI = "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable" - >>> with ShallowPostgresSaver.from_conn_string(DB_URI) as memory: - >>> config = {"configurable": {"thread_id": "1", "checkpoint_ns": ""}} - >>> checkpoint = {"ts": "2024-05-04T06:32:42.235444+00:00", "id": "1ef4f797-8335-6428-8001-8a1503f9b875", "channel_values": {"key": "value"}} - >>> saved_config = memory.put(config, checkpoint, {"source": "input", "step": 1, "writes": {"key": "value"}}, {}) - >>> print(saved_config) - {'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef4f797-8335-6428-8001-8a1503f9b875'}} - """ - configurable = config["configurable"].copy() - thread_id = configurable.pop("thread_id") - checkpoint_ns = configurable.pop("checkpoint_ns") - - copy = checkpoint.copy() - next_config = { - "configurable": { - "thread_id": thread_id, - "checkpoint_ns": checkpoint_ns, - "checkpoint_id": checkpoint["id"], - } - } - - with self._cursor(pipeline=True) as cur: - cur.execute( - """DELETE FROM checkpoint_writes - WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id NOT IN (%s, %s)""", - ( - thread_id, - checkpoint_ns, - checkpoint["id"], - configurable.get("checkpoint_id", ""), - ), - ) - cur.executemany( - self.UPSERT_CHECKPOINT_BLOBS_SQL, - _dump_blobs( - self.serde, - thread_id, - checkpoint_ns, - copy.pop("channel_values"), # type: ignore[misc] - new_versions, - ), - ) - cur.execute( - self.UPSERT_CHECKPOINTS_SQL, - ( - thread_id, - checkpoint_ns, - Jsonb(self._dump_checkpoint(copy)), - self._dump_metadata(get_checkpoint_metadata(config, metadata)), - ), - ) - return next_config - - def put_writes( - self, - config: RunnableConfig, - writes: Sequence[tuple[str, Any]], - task_id: str, - task_path: str = "", - ) -> None: - """Store intermediate writes linked to a checkpoint. - - This method saves intermediate writes associated with a checkpoint to the Postgres database. - - Args: - config (RunnableConfig): Configuration of the related checkpoint. - writes (List[Tuple[str, Any]]): List of writes to store. - task_id (str): Identifier for the task creating the writes. - """ - query = ( - self.UPSERT_CHECKPOINT_WRITES_SQL - if all(w[0] in WRITES_IDX_MAP for w in writes) - else self.INSERT_CHECKPOINT_WRITES_SQL - ) - with self._cursor(pipeline=True) as cur: - cur.executemany( - query, - self._dump_writes( - config["configurable"]["thread_id"], - config["configurable"]["checkpoint_ns"], - config["configurable"]["checkpoint_id"], - task_id, - task_path, - writes, - ), - ) - - @contextmanager - def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor[DictRow]]: - """Create a database cursor as a context manager. - - Args: - pipeline (bool): whether to use pipeline for the DB operations inside the context manager. - Will be applied regardless of whether the ShallowPostgresSaver instance was initialized with a pipeline. - If pipeline mode is not supported, will fall back to using transaction context manager. - """ - with _internal.get_connection(self.conn) as conn: - if self.pipe: - # a connection in pipeline mode can be used concurrently - # in multiple threads/coroutines, but only one cursor can be - # used at a time - try: - with conn.cursor(binary=True, row_factory=dict_row) as cur: - yield cur - finally: - if pipeline: - self.pipe.sync() - elif pipeline: - # a connection not in pipeline mode can only be used by one - # thread/coroutine at a time, so we acquire a lock - if self.supports_pipeline: - with ( - self.lock, - conn.pipeline(), - conn.cursor(binary=True, row_factory=dict_row) as cur, - ): - yield cur - else: - # Use connection's transaction context manager when pipeline mode not supported - with ( - self.lock, - conn.transaction(), - conn.cursor(binary=True, row_factory=dict_row) as cur, - ): - yield cur - else: - with self.lock, conn.cursor(binary=True, row_factory=dict_row) as cur: - yield cur - - -class AsyncShallowPostgresSaver(BasePostgresSaver): - """A checkpoint saver that uses Postgres to store checkpoints asynchronously. - - This checkpointer ONLY stores the most recent checkpoint and does NOT retain any history. - It is meant to be a light-weight drop-in replacement for the AsyncPostgresSaver that - supports most of the LangGraph persistence functionality with the exception of time travel. - """ - - SELECT_SQL = SELECT_SQL - MIGRATIONS = MIGRATIONS - UPSERT_CHECKPOINT_BLOBS_SQL = UPSERT_CHECKPOINT_BLOBS_SQL - UPSERT_CHECKPOINTS_SQL = UPSERT_CHECKPOINTS_SQL - UPSERT_CHECKPOINT_WRITES_SQL = UPSERT_CHECKPOINT_WRITES_SQL - INSERT_CHECKPOINT_WRITES_SQL = INSERT_CHECKPOINT_WRITES_SQL - lock: asyncio.Lock - - def __init__( - self, - conn: _ainternal.Conn, - pipe: Optional[AsyncPipeline] = None, - serde: Optional[SerializerProtocol] = None, - ) -> None: - super().__init__(serde=serde) - if isinstance(conn, AsyncConnectionPool) and pipe is not None: - raise ValueError( - "Pipeline should be used only with a single AsyncConnection, not AsyncConnectionPool." - ) - - self.conn = conn - self.pipe = pipe - self.lock = asyncio.Lock() - self.loop = asyncio.get_running_loop() - self.supports_pipeline = Capabilities().has_pipeline() - - @classmethod - @asynccontextmanager - async def from_conn_string( - cls, - conn_string: str, - *, - pipeline: bool = False, - serde: Optional[SerializerProtocol] = None, - ) -> AsyncIterator["AsyncShallowPostgresSaver"]: - """Create a new AsyncShallowPostgresSaver instance from a connection string. - - Args: - conn_string (str): The Postgres connection info string. - pipeline (bool): whether to use AsyncPipeline - - Returns: - AsyncShallowPostgresSaver: A new AsyncShallowPostgresSaver instance. - """ - async with await AsyncConnection.connect( - conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row - ) as conn: - if pipeline: - async with conn.pipeline() as pipe: - yield cls(conn=conn, pipe=pipe, serde=serde) - else: - yield cls(conn=conn, serde=serde) - - async def setup(self) -> None: - """Set up the checkpoint database asynchronously. - - This method creates the necessary tables in the Postgres database if they don't - already exist and runs database migrations. It MUST be called directly by the user - the first time checkpointer is used. - """ - async with self._cursor() as cur: - await cur.execute(self.MIGRATIONS[0]) - results = await cur.execute( - "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1" - ) - row = await results.fetchone() - if row is None: - version = -1 - else: - version = row["v"] - for v, migration in zip( - range(version + 1, len(self.MIGRATIONS)), - self.MIGRATIONS[version + 1 :], - ): - await cur.execute(migration) - await cur.execute(f"INSERT INTO checkpoint_migrations (v) VALUES ({v})") - if self.pipe: - await self.pipe.sync() - - async def alist( - self, - config: Optional[RunnableConfig], - *, - filter: Optional[dict[str, Any]] = None, - before: Optional[RunnableConfig] = None, - limit: Optional[int] = None, - ) -> AsyncIterator[CheckpointTuple]: - """List checkpoints from the database asynchronously. - - This method retrieves a list of checkpoint tuples from the Postgres database based - on the provided config. For ShallowPostgresSaver, this method returns a list with - ONLY the most recent checkpoint. - """ - where, args = self._search_where(config, filter, before) - query = self.SELECT_SQL + where - if limit: - query += f" LIMIT {limit}" - async with self._cursor() as cur: - await cur.execute(self.SELECT_SQL + where, args, binary=True) - async for value in cur: - checkpoint = await asyncio.to_thread( - self._load_checkpoint, - value["checkpoint"], - value["channel_values"], - value["pending_sends"], - ) - yield CheckpointTuple( - config={ - "configurable": { - "thread_id": value["thread_id"], - "checkpoint_ns": value["checkpoint_ns"], - "checkpoint_id": checkpoint["id"], - } - }, - checkpoint=checkpoint, - metadata=self._load_metadata(value["metadata"]), - pending_writes=await asyncio.to_thread( - self._load_writes, value["pending_writes"] - ), - ) - - async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: - """Get a checkpoint tuple from the database asynchronously. - - This method retrieves a checkpoint tuple from the Postgres database based on the - provided config (matching the thread ID in the config). - - Args: - config (RunnableConfig): The config to use for retrieving the checkpoint. - - Returns: - Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found. - """ - thread_id = config["configurable"]["thread_id"] - checkpoint_ns = config["configurable"].get("checkpoint_ns", "") - args = (thread_id, checkpoint_ns) - where = "WHERE thread_id = %s AND checkpoint_ns = %s" - - async with self._cursor() as cur: - await cur.execute( - self.SELECT_SQL + where, - args, - binary=True, - ) - - async for value in cur: - checkpoint = await asyncio.to_thread( - self._load_checkpoint, - value["checkpoint"], - value["channel_values"], - value["pending_sends"], - ) - return CheckpointTuple( - config={ - "configurable": { - "thread_id": thread_id, - "checkpoint_ns": checkpoint_ns, - "checkpoint_id": checkpoint["id"], - } - }, - checkpoint=checkpoint, - metadata=self._load_metadata(value["metadata"]), - pending_writes=await asyncio.to_thread( - self._load_writes, value["pending_writes"] - ), - ) - - async def aput( - self, - config: RunnableConfig, - checkpoint: Checkpoint, - metadata: CheckpointMetadata, - new_versions: ChannelVersions, - ) -> RunnableConfig: - """Save a checkpoint to the database asynchronously. - - This method saves a checkpoint to the Postgres database. The checkpoint is associated - with the provided config. For AsyncShallowPostgresSaver, this method saves ONLY the most recent - checkpoint and overwrites a previous checkpoint, if it exists. - - Args: - config (RunnableConfig): The config to associate with the checkpoint. - checkpoint (Checkpoint): The checkpoint to save. - metadata (CheckpointMetadata): Additional metadata to save with the checkpoint. - new_versions (ChannelVersions): New channel versions as of this write. - - Returns: - RunnableConfig: Updated configuration after storing the checkpoint. - """ - configurable = config["configurable"].copy() - thread_id = configurable.pop("thread_id") - checkpoint_ns = configurable.pop("checkpoint_ns") - - copy = checkpoint.copy() - next_config = { - "configurable": { - "thread_id": thread_id, - "checkpoint_ns": checkpoint_ns, - "checkpoint_id": checkpoint["id"], - } - } - - async with self._cursor(pipeline=True) as cur: - await cur.execute( - """DELETE FROM checkpoint_writes - WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id NOT IN (%s, %s)""", - ( - thread_id, - checkpoint_ns, - checkpoint["id"], - configurable.get("checkpoint_id", ""), - ), - ) - await cur.executemany( - self.UPSERT_CHECKPOINT_BLOBS_SQL, - _dump_blobs( - self.serde, - thread_id, - checkpoint_ns, - copy.pop("channel_values"), # type: ignore[misc] - new_versions, - ), - ) - await cur.execute( - self.UPSERT_CHECKPOINTS_SQL, - ( - thread_id, - checkpoint_ns, - Jsonb(self._dump_checkpoint(copy)), - self._dump_metadata(get_checkpoint_metadata(config, metadata)), - ), - ) - return next_config - - async def aput_writes( - self, - config: RunnableConfig, - writes: Sequence[tuple[str, Any]], - task_id: str, - task_path: str = "", - ) -> None: - """Store intermediate writes linked to a checkpoint asynchronously. - - This method saves intermediate writes associated with a checkpoint to the database. - - Args: - config (RunnableConfig): Configuration of the related checkpoint. - writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair. - task_id (str): Identifier for the task creating the writes. - """ - query = ( - self.UPSERT_CHECKPOINT_WRITES_SQL - if all(w[0] in WRITES_IDX_MAP for w in writes) - else self.INSERT_CHECKPOINT_WRITES_SQL - ) - params = await asyncio.to_thread( - self._dump_writes, - config["configurable"]["thread_id"], - config["configurable"]["checkpoint_ns"], - config["configurable"]["checkpoint_id"], - task_id, - task_path, - writes, - ) - async with self._cursor(pipeline=True) as cur: - await cur.executemany(query, params) - - @asynccontextmanager - async def _cursor( - self, *, pipeline: bool = False - ) -> AsyncIterator[AsyncCursor[DictRow]]: - """Create a database cursor as a context manager. - - Args: - pipeline (bool): whether to use pipeline for the DB operations inside the context manager. - Will be applied regardless of whether the AsyncShallowPostgresSaver instance was initialized with a pipeline. - If pipeline mode is not supported, will fall back to using transaction context manager. - """ - async with _ainternal.get_connection(self.conn) as conn: - if self.pipe: - # a connection in pipeline mode can be used concurrently - # in multiple threads/coroutines, but only one cursor can be - # used at a time - try: - async with conn.cursor(binary=True, row_factory=dict_row) as cur: - yield cur - finally: - if pipeline: - await self.pipe.sync() - elif pipeline: - # a connection not in pipeline mode can only be used by one - # thread/coroutine at a time, so we acquire a lock - if self.supports_pipeline: - async with ( - self.lock, - conn.pipeline(), - conn.cursor(binary=True, row_factory=dict_row) as cur, - ): - yield cur - else: - # Use connection's transaction context manager when pipeline mode not supported - async with ( - self.lock, - conn.transaction(), - conn.cursor(binary=True, row_factory=dict_row) as cur, - ): - yield cur - else: - async with ( - self.lock, - conn.cursor(binary=True, row_factory=dict_row) as cur, - ): - yield cur - - def list( - self, - config: Optional[RunnableConfig], - *, - filter: Optional[dict[str, Any]] = None, - before: Optional[RunnableConfig] = None, - limit: Optional[int] = None, - ) -> Iterator[CheckpointTuple]: - """List checkpoints from the database. - - This method retrieves a list of checkpoint tuples from the Postgres database based - on the provided config. For ShallowPostgresSaver, this method returns a list with - ONLY the most recent checkpoint. - """ - aiter_ = self.alist(config, filter=filter, before=before, limit=limit) - while True: - try: - yield asyncio.run_coroutine_threadsafe( - anext(aiter_), # noqa: F821 - self.loop, - ).result() - except StopAsyncIteration: - break - - def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: - """Get a checkpoint tuple from the database. - - This method retrieves a checkpoint tuple from the Postgres database based on the - provided config (matching the thread ID in the config). - - Args: - config (RunnableConfig): The config to use for retrieving the checkpoint. - - Returns: - Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found. - """ - try: - # check if we are in the main thread, only bg threads can block - # we don't check in other methods to avoid the overhead - if asyncio.get_running_loop() is self.loop: - raise asyncio.InvalidStateError( - "Synchronous calls to AsyncShallowPostgresSaver are only allowed from a " - "different thread. From the main thread, use the async interface." - "For example, use `await checkpointer.aget_tuple(...)` or `await " - "graph.ainvoke(...)`." - ) - except RuntimeError: - pass - return asyncio.run_coroutine_threadsafe( - self.aget_tuple(config), self.loop - ).result() - - def put( - self, - config: RunnableConfig, - checkpoint: Checkpoint, - metadata: CheckpointMetadata, - new_versions: ChannelVersions, - ) -> RunnableConfig: - """Save a checkpoint to the database. - - This method saves a checkpoint to the Postgres database. The checkpoint is associated - with the provided config. For AsyncShallowPostgresSaver, this method saves ONLY the most recent - checkpoint and overwrites a previous checkpoint, if it exists. - - Args: - config (RunnableConfig): The config to associate with the checkpoint. - checkpoint (Checkpoint): The checkpoint to save. - metadata (CheckpointMetadata): Additional metadata to save with the checkpoint. - new_versions (ChannelVersions): New channel versions as of this write. - - Returns: - RunnableConfig: Updated configuration after storing the checkpoint. - """ - return asyncio.run_coroutine_threadsafe( - self.aput(config, checkpoint, metadata, new_versions), self.loop - ).result() - - def put_writes( - self, - config: RunnableConfig, - writes: Sequence[tuple[str, Any]], - task_id: str, - task_path: str = "", - ) -> None: - """Store intermediate writes linked to a checkpoint. - - This method saves intermediate writes associated with a checkpoint to the database. - - Args: - config (RunnableConfig): Configuration of the related checkpoint. - writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair. - task_id (str): Identifier for the task creating the writes. - task_path (str): Path of the task creating the writes. - """ - return asyncio.run_coroutine_threadsafe( - self.aput_writes(config, writes, task_id, task_path), self.loop - ).result() diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/store/postgres/__init__.py index 47b7836d2..0abdc3352 100644 --- a/libs/checkpoint-postgres/langgraph/store/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/store/postgres/__init__.py @@ -1,4 +1,3 @@ -from langgraph.store.postgres.aio import AsyncPostgresStore from langgraph.store.postgres.base import PostgresStore -__all__ = ["AsyncPostgresStore", "PostgresStore"] +__all__ = ["PostgresStore"] diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/aio.py b/libs/checkpoint-postgres/langgraph/store/postgres/aio.py deleted file mode 100644 index a8d434360..000000000 --- a/libs/checkpoint-postgres/langgraph/store/postgres/aio.py +++ /dev/null @@ -1,434 +0,0 @@ -import asyncio -import logging -from collections.abc import AsyncIterator, Iterable, Sequence -from contextlib import asynccontextmanager -from typing import Any, Callable, Optional, Union, cast - -import orjson -from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline, Capabilities -from psycopg.rows import DictRow, dict_row -from psycopg_pool import AsyncConnectionPool - -from langgraph.checkpoint.postgres import _ainternal -from langgraph.store.base import ( - GetOp, - ListNamespacesOp, - Op, - PutOp, - Result, - SearchOp, -) -from langgraph.store.base.batch import AsyncBatchedBaseStore -from langgraph.store.postgres.base import ( - _PLACEHOLDER, - BasePostgresStore, - PoolConfig, - PostgresIndexConfig, - Row, - _decode_ns_bytes, - _ensure_index_config, - _group_ops, - _row_to_item, - _row_to_search_item, -) - -logger = logging.getLogger(__name__) - - -class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Conn]): - """Asynchronous Postgres-backed store with optional vector search using pgvector. - - !!! example "Examples" - Basic setup and usage: - ```python - from langgraph.store.postgres import AsyncPostgresStore - - conn_string = "postgresql://user:pass@localhost:5432/dbname" - - async with AsyncPostgresStore.from_conn_string(conn_string) as store: - await store.setup() # Run migrations. Done once - - # Store and retrieve data - await store.aput(("users", "123"), "prefs", {"theme": "dark"}) - item = await store.aget(("users", "123"), "prefs") - ``` - - Vector search using LangChain embeddings: - ```python - from langchain.embeddings import init_embeddings - from langgraph.store.postgres import AsyncPostgresStore - - conn_string = "postgresql://user:pass@localhost:5432/dbname" - - async with AsyncPostgresStore.from_conn_string( - conn_string, - index={ - "dims": 1536, - "embed": init_embeddings("openai:text-embedding-3-small"), - "fields": ["text"] # specify which fields to embed. Default is the whole serialized value - } - ) as store: - await store.setup() # Run migrations. Done once - - # Store documents - await store.aput(("docs",), "doc1", {"text": "Python tutorial"}) - await store.aput(("docs",), "doc2", {"text": "TypeScript guide"}) - await store.aput(("docs",), "doc3", {"text": "Other guide"}, index=False) # don't index - - # Search by similarity - results = await store.asearch(("docs",), "programming guides", limit=2) - ``` - - Using connection pooling for better performance: - ```python - from langgraph.store.postgres import AsyncPostgresStore, PoolConfig - - conn_string = "postgresql://user:pass@localhost:5432/dbname" - - async with AsyncPostgresStore.from_conn_string( - conn_string, - pool_config=PoolConfig( - min_size=5, - max_size=20 - ) - ) as store: - await store.setup() # Run migrations. Done once - # Use store with connection pooling... - ``` - - Warning: - Make sure to: - 1. Call `setup()` before first use to create necessary tables and indexes - 2. Have the pgvector extension available to use vector search - 3. Use Python 3.10+ for async functionality - - Note: - Semantic search is disabled by default. You can enable it by providing an `index` configuration - when creating the store. Without this configuration, all `index` arguments passed to - `put` or `aput` will have no effect. - """ - - __slots__ = ( - "_deserializer", - "pipe", - "lock", - "supports_pipeline", - "index_config", - "embeddings", - ) - - def __init__( - self, - conn: _ainternal.Conn, - *, - pipe: Optional[AsyncPipeline] = None, - deserializer: Optional[ - Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]] - ] = None, - index: Optional[PostgresIndexConfig] = None, - ) -> None: - if isinstance(conn, AsyncConnectionPool) and pipe is not None: - raise ValueError( - "Pipeline should be used only with a single AsyncConnection, not AsyncConnectionPool." - ) - super().__init__() - self._deserializer = deserializer - self.conn = conn - self.pipe = pipe - self.lock = asyncio.Lock() - self.loop = asyncio.get_running_loop() - self.supports_pipeline = Capabilities().has_pipeline() - self.index_config = index - if self.index_config: - self.embeddings, self.index_config = _ensure_index_config(self.index_config) - - else: - self.embeddings = None - - async def abatch(self, ops: Iterable[Op]) -> list[Result]: - grouped_ops, num_ops = _group_ops(ops) - results: list[Result] = [None] * num_ops - - async with _ainternal.get_connection(self.conn) as conn: - if self.pipe: - async with self.pipe: - await self._execute_batch(grouped_ops, results, conn) - else: - await self._execute_batch(grouped_ops, results, conn) - - return results - - @classmethod - @asynccontextmanager - async def from_conn_string( - cls, - conn_string: str, - *, - pipeline: bool = False, - pool_config: Optional[PoolConfig] = None, - index: Optional[PostgresIndexConfig] = None, - ) -> AsyncIterator["AsyncPostgresStore"]: - """Create a new AsyncPostgresStore instance from a connection string. - - Args: - conn_string (str): The Postgres connection info string. - pipeline (bool): Whether to use AsyncPipeline (only for single connections) - pool_config (Optional[PoolConfig]): Configuration for the connection pool. - If provided, will create a connection pool and use it instead of a single connection. - This overrides the `pipeline` argument. - index (Optional[PostgresIndexConfig]): The embedding config. - - Returns: - AsyncPostgresStore: A new AsyncPostgresStore instance. - """ - if pool_config is not None: - pc = pool_config.copy() - async with cast( - AsyncConnectionPool[AsyncConnection[DictRow]], - AsyncConnectionPool( - conn_string, - min_size=pc.pop("min_size", 1), - max_size=pc.pop("max_size", None), - kwargs={ - "autocommit": True, - "prepare_threshold": 0, - "row_factory": dict_row, - **(pc.pop("kwargs", None) or {}), - }, - **cast(dict, pc), - ), - ) as pool: - yield cls(conn=pool, index=index) - else: - async with await AsyncConnection.connect( - conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row - ) as conn: - if pipeline: - async with conn.pipeline() as pipe: - yield cls(conn=conn, pipe=pipe, index=index) - else: - yield cls(conn=conn, index=index) - - async def setup(self) -> None: - """Set up the store database asynchronously. - - This method creates the necessary tables in the Postgres database if they don't - already exist and runs database migrations. It MUST be called directly by the user - the first time the store is used. - """ - - async def _get_version(cur: AsyncCursor[DictRow], table: str) -> int: - await cur.execute( - f""" - CREATE TABLE IF NOT EXISTS {table} ( - v INTEGER PRIMARY KEY - ) - """ - ) - await cur.execute(f"SELECT v FROM {table} ORDER BY v DESC LIMIT 1") - row = cast(dict, await cur.fetchone()) - if row is None: - version = -1 - else: - version = row["v"] - return version - - async with self._cursor() as cur: - version = await _get_version(cur, table="store_migrations") - for v, sql in enumerate(self.MIGRATIONS[version + 1 :], start=version + 1): - await cur.execute(sql) - await cur.execute("INSERT INTO store_migrations (v) VALUES (%s)", (v,)) - - if self.index_config: - version = await _get_version(cur, table="vector_migrations") - for v, migration in enumerate( - self.VECTOR_MIGRATIONS[version + 1 :], start=version + 1 - ): - sql = migration.sql - if migration.params: - params = { - k: v(self) if v is not None and callable(v) else v - for k, v in migration.params.items() - } - sql = sql % params - await cur.execute(sql) - await cur.execute( - "INSERT INTO vector_migrations (v) VALUES (%s)", (v,) - ) - - async def _execute_batch( - self, - grouped_ops: dict, - results: list[Result], - conn: AsyncConnection[DictRow], - ) -> None: - async with self._cursor(pipeline=True) as cur: - if GetOp in grouped_ops: - await self._batch_get_ops( - cast(Sequence[tuple[int, GetOp]], grouped_ops[GetOp]), - results, - cur, - ) - - if SearchOp in grouped_ops: - await self._batch_search_ops( - cast(Sequence[tuple[int, SearchOp]], grouped_ops[SearchOp]), - results, - cur, - ) - - if ListNamespacesOp in grouped_ops: - await self._batch_list_namespaces_ops( - cast( - Sequence[tuple[int, ListNamespacesOp]], - grouped_ops[ListNamespacesOp], - ), - results, - cur, - ) - - if PutOp in grouped_ops: - await self._batch_put_ops( - cast(Sequence[tuple[int, PutOp]], grouped_ops[PutOp]), - cur, - ) - - async def _batch_get_ops( - self, - get_ops: Sequence[tuple[int, GetOp]], - results: list[Result], - cur: AsyncCursor[DictRow], - ) -> None: - for query, params, namespace, items in self._get_batch_GET_ops_queries(get_ops): - await cur.execute(query, params) - rows = cast(list[Row], await cur.fetchall()) - key_to_row = {row["key"]: row for row in rows} - for idx, key in items: - row = key_to_row.get(key) - if row: - results[idx] = _row_to_item( - namespace, row, loader=self._deserializer - ) - else: - results[idx] = None - - async def _batch_put_ops( - self, - put_ops: Sequence[tuple[int, PutOp]], - cur: AsyncCursor[DictRow], - ) -> None: - queries, embedding_request = self._prepare_batch_PUT_queries(put_ops) - if embedding_request: - if self.embeddings is None: - # Should not get here since the embedding config is required - # to return an embedding_request above - raise ValueError( - "Embedding configuration is required for vector operations " - f"(for semantic search). " - f"Please provide an EmbeddingConfig when initializing the {self.__class__.__name__}." - ) - query, txt_params = embedding_request - vectors = await self.embeddings.aembed_documents( - [param[-1] for param in txt_params] - ) - queries.append( - ( - query, - [ - p - for (ns, k, pathname, _), vector in zip(txt_params, vectors) - for p in (ns, k, pathname, vector) - ], - ) - ) - - for query, params in queries: - await cur.execute(query, params) - - async def _batch_search_ops( - self, - search_ops: Sequence[tuple[int, SearchOp]], - results: list[Result], - cur: AsyncCursor[DictRow], - ) -> None: - queries, embedding_requests = self._prepare_batch_search_queries(search_ops) - - if embedding_requests and self.embeddings: - vectors = await self.embeddings.aembed_documents( - [query for _, query in embedding_requests] - ) - for (idx, _), vector in zip(embedding_requests, vectors): - _paramslist = queries[idx][1] - for i in range(len(_paramslist)): - if _paramslist[i] is _PLACEHOLDER: - _paramslist[i] = vector - - for (idx, _), (query, params) in zip(search_ops, queries): - await cur.execute(query, params) - rows = cast(list[Row], await cur.fetchall()) - items = [ - _row_to_search_item( - _decode_ns_bytes(row["prefix"]), row, loader=self._deserializer - ) - for row in rows - ] - results[idx] = items - - async def _batch_list_namespaces_ops( - self, - list_ops: Sequence[tuple[int, ListNamespacesOp]], - results: list[Result], - cur: AsyncCursor[DictRow], - ) -> None: - queries = self._get_batch_list_namespaces_queries(list_ops) - for (query, params), (idx, _) in zip(queries, list_ops): - await cur.execute(query, params) - rows = cast(list[dict], await cur.fetchall()) - namespaces = [_decode_ns_bytes(row["truncated_prefix"]) for row in rows] - results[idx] = namespaces - - @asynccontextmanager - async def _cursor( - self, *, pipeline: bool = False - ) -> AsyncIterator[AsyncCursor[DictRow]]: - """Create a database cursor as a context manager. - - Args: - pipeline: whether to use pipeline for the DB operations inside the context manager. - Will be applied regardless of whether the PostgresStore instance was initialized with a pipeline. - If pipeline mode is not supported, will fall back to using transaction context manager. - """ - async with _ainternal.get_connection(self.conn) as conn: - if self.pipe: - # a connection in pipeline mode can be used concurrently - # in multiple threads/coroutines, but only one cursor can be - # used at a time - try: - async with conn.cursor(binary=True, row_factory=dict_row) as cur: - yield cur - finally: - if pipeline: - await self.pipe.sync() - elif pipeline: - # a connection not in pipeline mode can only be used by one - # thread/coroutine at a time, so we acquire a lock - if self.supports_pipeline: - async with ( - self.lock, - conn.pipeline(), - conn.cursor(binary=True, row_factory=dict_row) as cur, - ): - yield cur - else: - async with ( - self.lock, - conn.transaction(), - conn.cursor(binary=True, row_factory=dict_row) as cur, - ): - yield cur - else: - async with ( - self.lock, - conn.cursor(binary=True) as cur, - ): - yield cur diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/base.py b/libs/checkpoint-postgres/langgraph/store/postgres/base.py index a3acc4744..40645e684 100644 --- a/libs/checkpoint-postgres/langgraph/store/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/store/postgres/base.py @@ -7,7 +7,6 @@ from collections.abc import Iterable, Iterator, Sequence from contextlib import contextmanager from datetime import datetime from typing import ( - TYPE_CHECKING, Any, Callable, Generic, @@ -26,10 +25,10 @@ from psycopg.types.json import Jsonb from psycopg_pool import ConnectionPool from typing_extensions import TypedDict -from langgraph.checkpoint.postgres import _ainternal as _ainternal from langgraph.checkpoint.postgres import _internal as _pg_internal from langgraph.store.base import ( BaseStore, + Embeddings, GetOp, IndexConfig, Item, @@ -44,9 +43,6 @@ from langgraph.store.base import ( tokenize_path, ) -if TYPE_CHECKING: - from langchain_core.embeddings import Embeddings - logger = logging.getLogger(__name__) @@ -127,7 +123,7 @@ CREATE INDEX CONCURRENTLY IF NOT EXISTS store_vectors_embedding_idx ON store_vec ] -C = TypeVar("C", bound=Union[_pg_internal.Conn, _ainternal.Conn]) +C = TypeVar("C", bound=_pg_internal.Conn) class PoolConfig(TypedDict, total=False): diff --git a/libs/checkpoint/langgraph/checkpoint/base/__init__.py b/libs/checkpoint/langgraph/checkpoint/base/__init__.py index ba9a69eed..4846fa670 100644 --- a/libs/checkpoint/langgraph/checkpoint/base/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/base/__init__.py @@ -1,4 +1,4 @@ -from collections.abc import AsyncIterator, Iterator, Mapping, Sequence +from collections.abc import Iterator, Mapping, Sequence from datetime import datetime, timezone from typing import ( # noqa: UP035 Any, @@ -14,8 +14,6 @@ from typing import ( # noqa: UP035 Union, ) -from langchain_core.runnables import ConfigurableFieldSpec, RunnableConfig - from langgraph.checkpoint.base.id import uuid6 from langgraph.checkpoint.serde.base import SerializerProtocol, maybe_add_typed_methods from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer @@ -32,6 +30,24 @@ V = TypeVar("V", int, float, str) PendingWrite = Tuple[str, str, Any] +class CheckpointConfig(TypedDict, total=False): + """Configuration for a Runnable.""" + + metadata: dict[str, Any] + """ + Metadata for this call and any sub-calls (eg. a Chain calling an LLM). + Keys should be strings, values should be JSON-serializable. + """ + + configurable: dict[str, Any] + """ + Runtime values for attributes previously made configurable on this Runnable, + or sub-Runnables, through .configurable_fields() or .configurable_alternatives(). + Check .output_schema() for a description of the attributes that have been made + configurable. + """ + + # Marked as total=False to allow for future expansion. class CheckpointMetadata(TypedDict, total=False): """Metadata associated with a checkpoint.""" @@ -157,41 +173,13 @@ def create_checkpoint( class CheckpointTuple(NamedTuple): """A tuple containing a checkpoint and its associated data.""" - config: RunnableConfig + config: CheckpointConfig checkpoint: Checkpoint metadata: CheckpointMetadata - parent_config: Optional[RunnableConfig] = None + parent_config: Optional[CheckpointConfig] = None pending_writes: Optional[List[PendingWrite]] = None -CheckpointThreadId = ConfigurableFieldSpec( - id="thread_id", - annotation=str, - name="Thread ID", - description=None, - default="", - is_shared=True, -) - -CheckpointNS = ConfigurableFieldSpec( - id="checkpoint_ns", - annotation=str, - name="Checkpoint NS", - description='Checkpoint namespace. Denotes the path to the subgraph node the checkpoint originates from, separated by `|` character, e.g. `"child|grandchild"`. Defaults to "" (root graph).', - default="", - is_shared=True, -) - -CheckpointId = ConfigurableFieldSpec( - id="checkpoint_id", - annotation=Optional[str], - name="Checkpoint ID", - description="Pass to fetch a past checkpoint. If None, fetches the latest checkpoint.", - default=None, - is_shared=True, -) - - class BaseCheckpointSaver(Generic[V]): """Base class for creating a graph checkpointer. @@ -215,20 +203,11 @@ class BaseCheckpointSaver(Generic[V]): ) -> None: self.serde = maybe_add_typed_methods(serde or self.serde) - @property - def config_specs(self) -> list[ConfigurableFieldSpec]: - """Define the configuration options for the checkpoint saver. - - Returns: - list[ConfigurableFieldSpec]: List of configuration field specs. - """ - return [CheckpointThreadId, CheckpointNS, CheckpointId] - - def get(self, config: RunnableConfig) -> Optional[Checkpoint]: + def get(self, config: CheckpointConfig) -> Optional[Checkpoint]: """Fetch a checkpoint using the given configuration. Args: - config (RunnableConfig): Configuration specifying which checkpoint to retrieve. + config (CheckpointConfig): Configuration specifying which checkpoint to retrieve. Returns: Optional[Checkpoint]: The requested checkpoint, or None if not found. @@ -236,11 +215,11 @@ class BaseCheckpointSaver(Generic[V]): if value := self.get_tuple(config): return value.checkpoint - def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: + def get_tuple(self, config: CheckpointConfig) -> Optional[CheckpointTuple]: """Fetch a checkpoint tuple using the given configuration. Args: - config (RunnableConfig): Configuration specifying which checkpoint to retrieve. + config (CheckpointConfig): Configuration specifying which checkpoint to retrieve. Returns: Optional[CheckpointTuple]: The requested checkpoint tuple, or None if not found. @@ -252,18 +231,18 @@ class BaseCheckpointSaver(Generic[V]): def list( self, - config: Optional[RunnableConfig], + config: Optional[CheckpointConfig], *, filter: Optional[Dict[str, Any]] = None, - before: Optional[RunnableConfig] = None, + before: Optional[CheckpointConfig] = None, limit: Optional[int] = None, ) -> Iterator[CheckpointTuple]: """List checkpoints that match the given criteria. Args: - config (Optional[RunnableConfig]): Base configuration for filtering checkpoints. + config (Optional[CheckpointConfig]): Base configuration for filtering checkpoints. filter (Optional[Dict[str, Any]]): Additional filtering criteria. - before (Optional[RunnableConfig]): List checkpoints created before this configuration. + before (Optional[CheckpointConfig]): List checkpoints created before this configuration. limit (Optional[int]): Maximum number of checkpoints to return. Returns: @@ -276,21 +255,21 @@ class BaseCheckpointSaver(Generic[V]): def put( self, - config: RunnableConfig, + config: CheckpointConfig, checkpoint: Checkpoint, metadata: CheckpointMetadata, new_versions: ChannelVersions, - ) -> RunnableConfig: + ) -> CheckpointConfig: """Store a checkpoint with its configuration and metadata. Args: - config (RunnableConfig): Configuration for the checkpoint. + config (CheckpointConfig): Configuration for the checkpoint. checkpoint (Checkpoint): The checkpoint to store. metadata (CheckpointMetadata): Additional metadata for the checkpoint. new_versions (ChannelVersions): New channel versions as of this write. Returns: - RunnableConfig: Updated configuration after storing the checkpoint. + CheckpointConfig: Updated configuration after storing the checkpoint. Raises: NotImplementedError: Implement this method in your custom checkpoint saver. @@ -299,7 +278,7 @@ class BaseCheckpointSaver(Generic[V]): def put_writes( self, - config: RunnableConfig, + config: CheckpointConfig, writes: Sequence[Tuple[str, Any]], task_id: str, task_path: str = "", @@ -307,101 +286,7 @@ class BaseCheckpointSaver(Generic[V]): """Store intermediate writes linked to a checkpoint. Args: - config (RunnableConfig): Configuration of the related checkpoint. - writes (List[Tuple[str, Any]]): List of writes to store. - task_id (str): Identifier for the task creating the writes. - task_path (str): Path of the task creating the writes. - - Raises: - NotImplementedError: Implement this method in your custom checkpoint saver. - """ - raise NotImplementedError - - async def aget(self, config: RunnableConfig) -> Optional[Checkpoint]: - """Asynchronously fetch a checkpoint using the given configuration. - - Args: - config (RunnableConfig): Configuration specifying which checkpoint to retrieve. - - Returns: - Optional[Checkpoint]: The requested checkpoint, or None if not found. - """ - if value := await self.aget_tuple(config): - return value.checkpoint - - async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: - """Asynchronously fetch a checkpoint tuple using the given configuration. - - Args: - config (RunnableConfig): Configuration specifying which checkpoint to retrieve. - - Returns: - Optional[CheckpointTuple]: The requested checkpoint tuple, or None if not found. - - Raises: - NotImplementedError: Implement this method in your custom checkpoint saver. - """ - raise NotImplementedError - - async def alist( - self, - config: Optional[RunnableConfig], - *, - filter: Optional[Dict[str, Any]] = None, - before: Optional[RunnableConfig] = None, - limit: Optional[int] = None, - ) -> AsyncIterator[CheckpointTuple]: - """Asynchronously list checkpoints that match the given criteria. - - Args: - config (Optional[RunnableConfig]): Base configuration for filtering checkpoints. - filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata. - before (Optional[RunnableConfig]): List checkpoints created before this configuration. - limit (Optional[int]): Maximum number of checkpoints to return. - - Returns: - AsyncIterator[CheckpointTuple]: Async iterator of matching checkpoint tuples. - - Raises: - NotImplementedError: Implement this method in your custom checkpoint saver. - """ - raise NotImplementedError - yield - - async def aput( - self, - config: RunnableConfig, - checkpoint: Checkpoint, - metadata: CheckpointMetadata, - new_versions: ChannelVersions, - ) -> RunnableConfig: - """Asynchronously store a checkpoint with its configuration and metadata. - - Args: - config (RunnableConfig): Configuration for the checkpoint. - checkpoint (Checkpoint): The checkpoint to store. - metadata (CheckpointMetadata): Additional metadata for the checkpoint. - new_versions (ChannelVersions): New channel versions as of this write. - - Returns: - RunnableConfig: Updated configuration after storing the checkpoint. - - Raises: - NotImplementedError: Implement this method in your custom checkpoint saver. - """ - raise NotImplementedError - - async def aput_writes( - self, - config: RunnableConfig, - writes: Sequence[Tuple[str, Any]], - task_id: str, - task_path: str = "", - ) -> None: - """Asynchronously store intermediate writes linked to a checkpoint. - - Args: - config (RunnableConfig): Configuration of the related checkpoint. + config (CheckpointConfig): Configuration of the related checkpoint. writes (List[Tuple[str, Any]]): List of writes to store. task_id (str): Identifier for the task creating the writes. task_path (str): Path of the task creating the writes. @@ -439,7 +324,7 @@ class EmptyChannelError(Exception): pass -def get_checkpoint_id(config: RunnableConfig) -> Optional[str]: +def get_checkpoint_id(config: CheckpointConfig) -> Optional[str]: """Get checkpoint ID in a backwards-compatible manner (fallback on thread_ts).""" return config["configurable"].get( "checkpoint_id", config["configurable"].get("thread_ts") @@ -447,7 +332,7 @@ def get_checkpoint_id(config: RunnableConfig) -> Optional[str]: def get_checkpoint_metadata( - config: RunnableConfig, metadata: CheckpointMetadata + config: CheckpointConfig, metadata: CheckpointMetadata ) -> CheckpointMetadata: """Get checkpoint metadata in a backwards-compatible manner.""" metadata = metadata.copy() diff --git a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py index 6eeb8867b..81fb77029 100644 --- a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py @@ -4,18 +4,17 @@ import pickle import random import shutil from collections import defaultdict -from collections.abc import AsyncIterator, Iterator, Sequence -from contextlib import AbstractAsyncContextManager, AbstractContextManager, ExitStack +from collections.abc import Iterator, Sequence +from contextlib import AbstractContextManager, ExitStack from types import TracebackType from typing import Any, Optional -from langchain_core.runnables import RunnableConfig - from langgraph.checkpoint.base import ( WRITES_IDX_MAP, BaseCheckpointSaver, ChannelVersions, Checkpoint, + CheckpointConfig, CheckpointMetadata, CheckpointTuple, SerializerProtocol, @@ -27,9 +26,7 @@ from langgraph.checkpoint.serde.types import TASKS, ChannelProtocol logger = logging.getLogger(__name__) -class InMemorySaver( - BaseCheckpointSaver[str], AbstractContextManager, AbstractAsyncContextManager -): +class InMemorySaver(BaseCheckpointSaver[str], AbstractContextManager): """An in-memory checkpoint saver. This checkpoint saver stores checkpoints in memory using a defaultdict. @@ -96,18 +93,7 @@ class InMemorySaver( ) -> Optional[bool]: return self.stack.__exit__(exc_type, exc_value, traceback) - async def __aenter__(self) -> "InMemorySaver": - return self.stack.__enter__() - - async def __aexit__( - self, - __exc_type: Optional[type[BaseException]], - __exc_value: Optional[BaseException], - __traceback: Optional[TracebackType], - ) -> Optional[bool]: - return self.stack.__exit__(__exc_type, __exc_value, __traceback) - - def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: + def get_tuple(self, config: CheckpointConfig) -> Optional[CheckpointTuple]: """Get a checkpoint tuple from the in-memory storage. This method retrieves a checkpoint tuple from the in-memory storage based on the @@ -116,7 +102,7 @@ class InMemorySaver( for the given thread ID is retrieved. Args: - config (RunnableConfig): The config to use for retrieving the checkpoint. + config (CheckpointConfig): The config to use for retrieving the checkpoint. Returns: Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found. @@ -211,10 +197,10 @@ class InMemorySaver( def list( self, - config: Optional[RunnableConfig], + config: Optional[CheckpointConfig], *, filter: Optional[dict[str, Any]] = None, - before: Optional[RunnableConfig] = None, + before: Optional[CheckpointConfig] = None, limit: Optional[int] = None, ) -> Iterator[CheckpointTuple]: """List checkpoints from the in-memory storage. @@ -223,9 +209,9 @@ class InMemorySaver( on the provided criteria. Args: - config (Optional[RunnableConfig]): Base configuration for filtering checkpoints. + config (Optional[CheckpointConfig]): Base configuration for filtering checkpoints. filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata. - before (Optional[RunnableConfig]): List checkpoints created before this configuration. + before (Optional[CheckpointConfig]): List checkpoints created before this configuration. limit (Optional[int]): Maximum number of checkpoints to return. Yields: @@ -330,24 +316,24 @@ class InMemorySaver( def put( self, - config: RunnableConfig, + config: CheckpointConfig, checkpoint: Checkpoint, metadata: CheckpointMetadata, new_versions: ChannelVersions, - ) -> RunnableConfig: + ) -> CheckpointConfig: """Save a checkpoint to the in-memory storage. This method saves a checkpoint to the in-memory storage. The checkpoint is associated with the provided config. Args: - config (RunnableConfig): The config to associate with the checkpoint. + config (CheckpointConfig): The config to associate with the checkpoint. checkpoint (Checkpoint): The checkpoint to save. metadata (CheckpointMetadata): Additional metadata to save with the checkpoint. new_versions (dict): New versions as of this write Returns: - RunnableConfig: The updated config containing the saved checkpoint's timestamp. + CheckpointConfig: The updated config containing the saved checkpoint's timestamp. """ c = checkpoint.copy() c.pop("pending_sends") # type: ignore[misc] @@ -372,7 +358,7 @@ class InMemorySaver( def put_writes( self, - config: RunnableConfig, + config: CheckpointConfig, writes: Sequence[tuple[str, Any]], task_id: str, task_path: str = "", @@ -383,13 +369,13 @@ class InMemorySaver( with the provided config. Args: - config (RunnableConfig): The config to associate with the writes. + config (CheckpointConfig): The config to associate with the writes. writes (list[tuple[str, Any]]): The writes to save. task_id (str): Identifier for the task creating the writes. task_path (str): Path of the task creating the writes. Returns: - RunnableConfig: The updated config containing the saved writes' timestamp. + CheckpointConfig: The updated config containing the saved writes' timestamp. """ thread_id = config["configurable"]["thread_id"] checkpoint_ns = config["configurable"].get("checkpoint_ns", "") @@ -408,85 +394,6 @@ class InMemorySaver( task_path, ) - async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: - """Asynchronous version of get_tuple. - - This method is an asynchronous wrapper around get_tuple that runs the synchronous - method in a separate thread using asyncio. - - Args: - config (RunnableConfig): The config to use for retrieving the checkpoint. - - Returns: - Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found. - """ - return self.get_tuple(config) - - async def alist( - self, - config: Optional[RunnableConfig], - *, - filter: Optional[dict[str, Any]] = None, - before: Optional[RunnableConfig] = None, - limit: Optional[int] = None, - ) -> AsyncIterator[CheckpointTuple]: - """Asynchronous version of list. - - This method is an asynchronous wrapper around list that runs the synchronous - method in a separate thread using asyncio. - - Args: - config (RunnableConfig): The config to use for listing the checkpoints. - - Yields: - AsyncIterator[CheckpointTuple]: An asynchronous iterator of checkpoint tuples. - """ - for item in self.list(config, filter=filter, before=before, limit=limit): - yield item - - async def aput( - self, - config: RunnableConfig, - checkpoint: Checkpoint, - metadata: CheckpointMetadata, - new_versions: ChannelVersions, - ) -> RunnableConfig: - """Asynchronous version of put. - - Args: - config (RunnableConfig): The config to associate with the checkpoint. - checkpoint (Checkpoint): The checkpoint to save. - metadata (CheckpointMetadata): Additional metadata to save with the checkpoint. - new_versions (dict): New versions as of this write - - Returns: - RunnableConfig: The updated config containing the saved checkpoint's timestamp. - """ - return self.put(config, checkpoint, metadata, new_versions) - - async def aput_writes( - self, - config: RunnableConfig, - writes: Sequence[tuple[str, Any]], - task_id: str, - task_path: str = "", - ) -> None: - """Asynchronous version of put_writes. - - This method is an asynchronous wrapper around put_writes that runs the synchronous - method in a separate thread using asyncio. - - Args: - config (RunnableConfig): The config to associate with the writes. - writes (List[Tuple[str, Any]]): The writes to save, each as a (channel, value) pair. - task_id (str): Identifier for the task creating the writes. - task_path (str): Path of the task creating the writes. - - Returns: - None - """ - return self.put_writes(config, writes, task_id, task_path) - def get_next_version(self, current: Optional[str], channel: ChannelProtocol) -> str: if current is None: current_v = 0 @@ -572,4 +479,4 @@ class PersistentDict(defaultdict): except Exception: logging.error(f"Failed to load file: {fileobj.name}") raise - raise ValueError("File not in a supported f ormat") + raise ValueError("File not in a supported format") diff --git a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py index 4bd0839c2..c3753cde7 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py @@ -17,20 +17,16 @@ from ipaddress import ( IPv6Interface, IPv6Network, ) -from typing import Any, Callable, Optional, Union, cast +from typing import Any, Callable, Optional, Union from uuid import UUID import msgpack # type: ignore[import-untyped] -from langchain_core.load.load import Reviver -from langchain_core.load.serializable import Serializable from zoneinfo import ZoneInfo from langgraph.checkpoint.serde.base import SerializerProtocol from langgraph.checkpoint.serde.types import SendProtocol from langgraph.store.base import Item -LC_REVIVER = Reviver() - class JsonPlusSerializer(SerializerProtocol): def _encode_constructor_args( @@ -55,9 +51,7 @@ class JsonPlusSerializer(SerializerProtocol): return out def _default(self, obj: Any) -> Union[str, dict[str, Any]]: - if isinstance(obj, Serializable): - return cast(dict[str, Any], obj.to_json()) - elif hasattr(obj, "model_dump") and callable(obj.model_dump): + if hasattr(obj, "model_dump") and callable(obj.model_dump): return self._encode_constructor_args( obj.__class__, method=(None, "model_construct"), kwargs=obj.model_dump() ) @@ -177,7 +171,7 @@ class JsonPlusSerializer(SerializerProtocol): except Exception: return None - return LC_REVIVER(value) + return value def dumps(self, obj: Any) -> bytes: return json.dumps(obj, default=self._default, ensure_ascii=False).encode( diff --git a/libs/checkpoint/langgraph/store/base/__init__.py b/libs/checkpoint/langgraph/store/base/__init__.py index 19fc62d8d..191b12bc6 100644 --- a/libs/checkpoint/langgraph/store/base/__init__.py +++ b/libs/checkpoint/langgraph/store/base/__init__.py @@ -13,10 +13,8 @@ from abc import ABC, abstractmethod from datetime import datetime from typing import Any, Iterable, Literal, NamedTuple, Optional, TypedDict, Union, cast -from langchain_core.embeddings import Embeddings - from langgraph.store.base.embed import ( - AEmbeddingsFunc, + Embeddings, EmbeddingsFunc, ensure_embeddings, get_text_at_path, @@ -493,13 +491,11 @@ class IndexConfig(TypedDict, total=False): - cohere:embed-multilingual-light-v3.0: 384 """ - embed: Union[Embeddings, EmbeddingsFunc, AEmbeddingsFunc, str] + embed: Union[EmbeddingsFunc, str] """Optional function to generate embeddings from text. Can be specified in three ways: - 1. A LangChain Embeddings instance 2. A synchronous embedding function (EmbeddingsFunc) - 3. An asynchronous embedding function (AEmbeddingsFunc) 4. A provider string (e.g., "openai:text-embedding-3-small") ???+ example "Examples" diff --git a/libs/checkpoint/langgraph/store/base/embed.py b/libs/checkpoint/langgraph/store/base/embed.py index 0fa70b407..8fdb72b34 100644 --- a/libs/checkpoint/langgraph/store/base/embed.py +++ b/libs/checkpoint/langgraph/store/base/embed.py @@ -6,12 +6,9 @@ with LangChain-compatible tools while maintaining support for both synchronous a asynchronous operations. """ -import asyncio import functools import json -from typing import Any, Awaitable, Callable, Optional, Sequence, Union - -from langchain_core.embeddings import Embeddings +from typing import Any, Callable, Optional, Sequence, Union EmbeddingsFunc = Callable[[Sequence[str]], list[list[float]]] """Type for synchronous embedding functions. @@ -21,16 +18,10 @@ where each embedding is a list of floats. The dimensionality of the embeddings should be consistent for all inputs. """ -AEmbeddingsFunc = Callable[[Sequence[str]], Awaitable[list[list[float]]]] -"""Type for asynchronous embedding functions. - -Similar to EmbeddingsFunc, but returns an awaitable that resolves to the embeddings. -""" - def ensure_embeddings( - embed: Union[Embeddings, EmbeddingsFunc, AEmbeddingsFunc, str, None], -) -> Embeddings: + embed: Union[EmbeddingsFunc, str, None], +) -> "Embeddings": """Ensure that an embedding function conforms to LangChain's Embeddings interface. This function wraps arbitrary embedding functions to make them compatible with @@ -96,10 +87,10 @@ def ensure_embeddings( if isinstance(embed, Embeddings): return embed - return EmbeddingsLambda(embed) + return Embeddings(embed) -class EmbeddingsLambda(Embeddings): +class Embeddings: """Wrapper to convert embedding functions into LangChain's Embeddings interface. This class allows arbitrary embedding functions to be used with LangChain-compatible @@ -140,12 +131,10 @@ class EmbeddingsLambda(Embeddings): def __init__( self, - func: Union[EmbeddingsFunc, AEmbeddingsFunc], + func: EmbeddingsFunc, ) -> None: if func is None: raise ValueError("func must be provided") - if _is_async_callable(func): - self.afunc = func else: self.func = func @@ -184,41 +173,6 @@ class EmbeddingsLambda(Embeddings): """ return self.embed_documents([text])[0] - async def aembed_documents(self, texts: list[str]) -> list[list[float]]: - """Asynchronously embed a list of texts into vectors. - - Args: - texts: list of texts to convert to embeddings. - - Returns: - list of embeddings, one per input text. Each embedding is a list of floats. - - Note: - If no async function was provided, this falls back to the sync implementation. - """ - afunc = getattr(self, "afunc", None) - if afunc is None: - return await super().aembed_documents(texts) - return await afunc(texts) - - async def aembed_query(self, text: str) -> list[float]: - """Asynchronously embed a single piece of text. - - Args: - text: Text to convert to an embedding. - - Returns: - Embedding vector as a list of floats. - - Note: - This is equivalent to calling aembed_documents with a single text - and taking the first result. - """ - afunc = getattr(self, "afunc", None) - if afunc is None: - return await super().aembed_query(text) - return (await afunc([text]))[0] - def get_text_at_path(obj: Any, path: Union[str, list[str]]) -> list[str]: """Extract text from an object using a path expression or pre-tokenized path. @@ -382,26 +336,6 @@ def tokenize_path(path: str) -> list[str]: return tokens -def _is_async_callable( - func: Any, -) -> bool: - """Check if a function is async. - - This includes both async def functions and classes with async __call__ methods. - - Args: - func: Function or callable object to check. - - Returns: - True if the function is async, False otherwise. - """ - return ( - asyncio.iscoroutinefunction(func) - or hasattr(func, "__call__") # noqa: B004 - and asyncio.iscoroutinefunction(func.__call__) - ) - - @functools.lru_cache def _get_init_embeddings() -> Optional[Callable[[str], Embeddings]]: try: @@ -415,5 +349,5 @@ def _get_init_embeddings() -> Optional[Callable[[str], Embeddings]]: __all__ = [ "ensure_embeddings", "EmbeddingsFunc", - "AEmbeddingsFunc", + "Embeddings", ] diff --git a/libs/checkpoint/langgraph/store/memory/__init__.py b/libs/checkpoint/langgraph/store/memory/__init__.py index 8a0d36b46..1b35f7393 100644 --- a/libs/checkpoint/langgraph/store/memory/__init__.py +++ b/libs/checkpoint/langgraph/store/memory/__init__.py @@ -108,10 +108,9 @@ from datetime import datetime, timezone from importlib import util from typing import Any, Iterable, Optional -from langchain_core.embeddings import Embeddings - from langgraph.store.base import ( BaseStore, + Embeddings, GetOp, IndexConfig, Item, diff --git a/libs/checkpoint/tests/embed_test_utils.py b/libs/checkpoint/tests/embed_test_utils.py index d28cd959f..7d2923b75 100644 --- a/libs/checkpoint/tests/embed_test_utils.py +++ b/libs/checkpoint/tests/embed_test_utils.py @@ -5,7 +5,7 @@ import random from collections import Counter, defaultdict from typing import Any -from langchain_core.embeddings import Embeddings +from langgraph.store.base.embed import Embeddings class CharacterEmbeddings(Embeddings): diff --git a/libs/langgraph/Makefile b/libs/langgraph/Makefile index 8974fcd32..e2bdbdb87 100644 --- a/libs/langgraph/Makefile +++ b/libs/langgraph/Makefile @@ -60,7 +60,7 @@ MAXFAIL ?= MAXFAIL_ARGS := $(if $(MAXFAIL),--maxfail $(MAXFAIL),) test_watch: - make start-postgres && poetry run ptw . -- --ff -vv -x $(XDIST_ARGS) $(MAXFAIL_ARGS) --snapshot-update --tb short $(TEST); \ + make start-postgres && poetry run ptw . -- --ff -vv -x $(MAXFAIL_ARGS) --snapshot-update --tb short $(TEST); \ EXIT_CODE=$$?; \ make stop-postgres; \ exit $$EXIT_CODE diff --git a/libs/langgraph/langgraph/channels/__init__.py b/libs/langgraph/langgraph/channels/__init__.py index 6f9ba2119..cdb193484 100644 --- a/libs/langgraph/langgraph/channels/__init__.py +++ b/libs/langgraph/langgraph/channels/__init__.py @@ -1,6 +1,5 @@ from langgraph.channels.any_value import AnyValue from langgraph.channels.binop import BinaryOperatorAggregate -from langgraph.channels.context import Context from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic @@ -9,7 +8,6 @@ from langgraph.channels.untracked_value import UntrackedValue __all__ = [ "LastValue", "Topic", - "Context", "BinaryOperatorAggregate", "UntrackedValue", "EphemeralValue", diff --git a/libs/langgraph/langgraph/channels/context.py b/libs/langgraph/langgraph/channels/context.py deleted file mode 100644 index 3b4e26805..000000000 --- a/libs/langgraph/langgraph/channels/context.py +++ /dev/null @@ -1,5 +0,0 @@ -from langgraph.managed.context import Context as ContextManagedValue - -Context = ContextManagedValue.of - -__all__ = ["Context"] diff --git a/libs/langgraph/langgraph/config.py b/libs/langgraph/langgraph/config.py index b2ef57cfb..42fb1890e 100644 --- a/libs/langgraph/langgraph/config.py +++ b/libs/langgraph/langgraph/config.py @@ -2,12 +2,10 @@ import asyncio import sys from typing import Any -from langchain_core.runnables import RunnableConfig -from langchain_core.runnables.config import var_child_runnable_config - from langgraph.constants import CONF, CONFIG_KEY_STORE, CONFIG_KEY_STREAM_WRITER from langgraph.store.base import BaseStore from langgraph.types import StreamWriter +from langgraph.utils.config import RunnableConfig, var_child_runnable_config def _no_op_stream_writer(c: Any) -> None: @@ -44,10 +42,6 @@ def get_store() -> BaseStore: .compile(store=store) ) - # or with entrypoint - @entrypoint(store=store) - def workflow(inputs): - ... ``` !!! warning "Async with Python < 3.11" @@ -87,32 +81,6 @@ def get_store() -> BaseStore: ```pycon {'foo': 3} ``` - - Example: Using with functional API - ```python - from langgraph.func import entrypoint, task - from langgraph.store.memory import InMemoryStore - from langgraph.config import get_store - - store = InMemoryStore() - store.put(("values",), "foo", {"bar": 2}) - - @task - def my_task(value: int): - my_store = get_store() - stored_value = my_store.get(("values",), "foo").value["bar"] - return stored_value + 1 - - @entrypoint(store=store) - def workflow(value: int): - return my_task(value).result() - - workflow.invoke(1) - ``` - - ```pycon - 3 - ``` """ config = get_config() return config[CONF][CONFIG_KEY_STORE] @@ -154,29 +122,6 @@ def get_stream_writer() -> StreamWriter: print(chunk) ``` - ```pycon - {'custom_data': 'Hello!'} - ``` - - Example: Using with functional API - ```python - from langgraph.func import entrypoint, task - from langgraph.config import get_stream_writer - - @task - def my_task(value: int): - my_stream_writer = get_stream_writer() - my_stream_writer({"custom_data": "Hello!"}) - return value + 1 - - @entrypoint(store=store) - def workflow(value: int): - return my_task(value).result() - - for chunk in workflow.stream(1, stream_mode="custom"): - print(chunk) - ``` - ```pycon {'custom_data': 'Hello!'} ``` diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index 4fde5e7dd..9b2ff462f 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -57,7 +57,7 @@ CONFIG_KEY_STREAM = sys.intern("__pregel_stream") CONFIG_KEY_STREAM_WRITER = sys.intern("__pregel_stream_writer") # holds a `StreamWriter` for stream_mode=custom CONFIG_KEY_STORE = sys.intern("__pregel_store") -# holds a `BaseStore` made available to managed values +# holds a `BaseStore` made available in context CONFIG_KEY_RESUMING = sys.intern("__pregel_resuming") # holds a boolean indicating if subgraphs should resume from a previous checkpoint CONFIG_KEY_TASK_ID = sys.intern("__pregel_task_id") diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py deleted file mode 100644 index a3d012e55..000000000 --- a/libs/langgraph/langgraph/func/__init__.py +++ /dev/null @@ -1,443 +0,0 @@ -import asyncio -import concurrent.futures -import functools -import inspect -from dataclasses import dataclass -from typing import ( - Any, - Awaitable, - Callable, - Generic, - Optional, - TypeVar, - Union, - get_args, - get_origin, - overload, -) - -from langgraph.channels.ephemeral_value import EphemeralValue -from langgraph.channels.last_value import LastValue -from langgraph.checkpoint.base import BaseCheckpointSaver -from langgraph.constants import END, PREVIOUS, START, TAG_HIDDEN -from langgraph.pregel import Pregel -from langgraph.pregel.call import ( - P, - SyncAsyncFuture, - T, - call, - get_runnable_for_entrypoint, -) -from langgraph.pregel.read import PregelNode -from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry -from langgraph.store.base import BaseStore -from langgraph.types import _DC_KWARGS, RetryPolicy, StreamMode - - -@overload -def task( - *, name: Optional[str] = None, retry: Optional[RetryPolicy] = None -) -> Callable[[Callable[P, T]], Callable[P, SyncAsyncFuture[T]]]: ... - - -@overload -def task( - __func_or_none__: Callable[P, T], -) -> Callable[P, SyncAsyncFuture[T]]: ... - - -def task( - __func_or_none__: Optional[Union[Callable[P, T], Callable[P, Awaitable[T]]]] = None, - *, - name: Optional[str] = None, - retry: Optional[RetryPolicy] = None, -) -> Union[ - Callable[[Callable[P, T]], Callable[P, SyncAsyncFuture[T]]], - Callable[P, SyncAsyncFuture[T]], -]: - """Define a LangGraph task using the `task` decorator. - - !!! important "Requires python 3.11 or higher for async functions" - The `task` decorator supports both sync and async functions. To use async - functions, ensure that you are using Python 3.11 or higher. - - Tasks can only be called from within an [entrypoint][langgraph.func.entrypoint] or - from within a StateGraph. A task can be called like a regular function with the - following differences: - - - When a checkpointer is enabled, the function inputs and outputs must be serializable. - - The decorated function can only be called from within an entrypoint or StateGraph. - - Calling the function produces a future. This makes it easy to parallelize tasks. - - Args: - retry: An optional retry policy to use for the task in case of a failure. - - Returns: - A callable function when used as a decorator. - - Example: Sync Task - ```python - from langgraph.func import entrypoint, task - - @task - def add_one(a: int) -> int: - return a + 1 - - @entrypoint() - def add_one(numbers: list[int]) -> list[int]: - futures = [add_one(n) for n in numbers] - results = [f.result() for f in futures] - return results - - # Call the entrypoint - add_one.invoke([1, 2, 3]) # Returns [2, 3, 4] - ``` - - Example: Async Task - ```python - import asyncio - from langgraph.func import entrypoint, task - - @task - async def add_one(a: int) -> int: - return a + 1 - - @entrypoint() - async def add_one(numbers: list[int]) -> list[int]: - futures = [add_one(n) for n in numbers] - return asyncio.gather(*futures) - - # Call the entrypoint - await add_one.ainvoke([1, 2, 3]) # Returns [2, 3, 4] - ``` - """ - - def decorator( - func: Union[Callable[P, Awaitable[T]], Callable[P, T]], - ) -> Union[ - Callable[P, concurrent.futures.Future[T]], Callable[P, asyncio.Future[T]] - ]: - if name is not None: - if hasattr(func, "__func__"): - # handle class methods - # NOTE: we're modifying the instance method to avoid modifying - # the original class method in case it's shared across multiple tasks - instance_method = functools.partial(func.__func__, func.__self__) # type: ignore [union-attr] - instance_method.__name__ = name # type: ignore [attr-defined] - func = instance_method - else: - # handle regular functions / partials / callable classes, etc. - func.__name__ = name - - call_func = functools.partial(call, func, retry=retry) - object.__setattr__(call_func, "_is_pregel_task", True) - return functools.update_wrapper(call_func, func) - - if __func_or_none__ is not None: - return decorator(__func_or_none__) - - return decorator - - -R = TypeVar("R") -S = TypeVar("S") - - -# The decorator was wrapped in a class to support the `final` attribute. -# In this form, the `final` attribute should play nicely with IDE autocompletion, -# and type checking tools. -# In addition, we'll be able to surface this information in the API Reference. -class entrypoint: - """Define a LangGraph workflow using the `entrypoint` decorator. - - ### Function signature - - The decorated function must accept a **single parameter**, which serves as the input - to the function. This input parameter can be of any type. Use a dictionary - to pass **multiple parameters** to the function. - - ### Injectable parameters - - The decorated function can request access to additional parameters - that will be injected automatically at run time. These parameters include: - - | Parameter | Description | - |------------------|----------------------------------------------------------------------------------------------------| - | **`store`** | An instance of [BaseStore][langgraph.store.base.BaseStore]. Useful for long-term memory. | - | **`writer`** | A [StreamWriter][langgraph.types.StreamWriter] instance for writing custom data to a stream. | - | **`config`** | A configuration object (aka RunnableConfig) that holds run-time configuration values. | - | **`previous`** | The previous return value for the given thread (available only when a checkpointer is provided). | - - The entrypoint decorator can be applied to sync functions or async functions. - - ### State management - - The **`previous`** parameter can be used to access the return value of the previous - invocation of the entrypoint on the same thread id. This value is only available - when a checkpointer is provided. - - If you want **`previous`** to be different from the return value, you can use the - `entrypoint.final` object to return a value while saving a different value to the - checkpoint. - - Args: - checkpointer: Specify a checkpointer to create a workflow that can persist - its state across runs. - store: A generalized key-value store. Some implementations may support - semantic search capabilities through an optional `index` configuration. - config_schema: Specifies the schema for the configuration object that will be - passed to the workflow. - - Example: Using entrypoint and tasks - ```python - import time - - from langgraph.func import entrypoint, task - from langgraph.types import interrupt, Command - from langgraph.checkpoint.memory import MemorySaver - - @task - def compose_essay(topic: str) -> str: - time.sleep(1.0) # Simulate slow operation - return f"An essay about {topic}" - - @entrypoint(checkpointer=MemorySaver()) - def review_workflow(topic: str) -> dict: - \"\"\"Manages the workflow for generating and reviewing an essay. - - The workflow includes: - 1. Generating an essay about the given topic. - 2. Interrupting the workflow for human review of the generated essay. - - Upon resuming the workflow, compose_essay task will not be re-executed - as its result is cached by the checkpointer. - - Args: - topic (str): The subject of the essay. - - Returns: - dict: A dictionary containing the generated essay and the human review. - \"\"\" - essay_future = compose_essay(topic) - essay = essay_future.result() - human_review = interrupt({ - \"question\": \"Please provide a review\", - \"essay\": essay - }) - return { - \"essay\": essay, - \"review\": human_review, - } - - # Example configuration for the workflow - config = { - \"configurable\": { - \"thread_id\": \"some_thread\" - } - } - - # Topic for the essay - topic = \"cats\" - - # Stream the workflow to generate the essay and await human review - for result in review_workflow.stream(topic, config): - print(result) - - # Example human review provided after the interrupt - human_review = \"This essay is great.\" - - # Resume the workflow with the provided human review - for result in review_workflow.stream(Command(resume=human_review), config): - print(result) - ``` - - Example: Accessing the previous return value - When a checkpointer is enabled the function can access the previous return value - of the previous invocation on the same thread id. - - ```python - from langgraph.checkpoint.memory import MemorySaver - from langgraph.func import entrypoint - - @entrypoint(checkpointer=MemorySaver()) - def my_workflow(input_data: str, previous: Optional[str] = None) -> str: - return "world" - - config = { - "configurable": { - "thread_id": "some_thread" - } - } - my_workflow.invoke("hello") - ``` - - Example: Using entrypoint.final to save a value - The `entrypoint.final` object allows you to return a value while saving - a different value to the checkpoint. This value will be accessible - in the next invocation of the entrypoint via the `previous` parameter, as - long as the same thread id is used. - - ```python - from langgraph.checkpoint.memory import MemorySaver - from langgraph.func import entrypoint - - @entrypoint(checkpointer=MemorySaver()) - def my_workflow(number: int, *, previous: Any = None) -> entrypoint.final[int, int]: - previous = previous or 0 - # This will return the previous value to the caller, saving - # 2 * number to the checkpoint, which will be used in the next invocation - # for the `previous` parameter. - return entrypoint.final(value=previous, save=2 * number) - - config = { - "configurable": { - "thread_id": "some_thread" - } - } - - my_workflow.invoke(3, config) # 0 (previous was None) - my_workflow.invoke(1, config) # 6 (previous was 3 * 2 from the previous invocation) - ``` - """ - - def __init__( - self, - checkpointer: Optional[BaseCheckpointSaver] = None, - store: Optional[BaseStore] = None, - config_schema: Optional[type[Any]] = None, - ) -> None: - """Initialize the entrypoint decorator.""" - self.checkpointer = checkpointer - self.store = store - self.config_schema = config_schema - - @dataclass(**_DC_KWARGS) - class final(Generic[R, S]): - """A primitive that can be returned from an entrypoint. - - This primitive allows to save a value to the checkpointer distinct from the - return value from the entrypoint. - - Example: Decoupling the return value and the save value - ```python - from langgraph.checkpoint.memory import MemorySaver - from langgraph.func import entrypoint - - @entrypoint(checkpointer=MemorySaver()) - def my_workflow(number: int, *, previous: Any = None) -> entrypoint.final[int, int]: - previous = previous or 0 - # This will return the previous value to the caller, saving - # 2 * number to the checkpoint, which will be used in the next invocation - # for the `previous` parameter. - return entrypoint.final(value=previous, save=2 * number) - - config = { - "configurable": { - "thread_id": "1" - } - } - - my_workflow.invoke(3, config) # 0 (previous was None) - my_workflow.invoke(1, config) # 6 (previous was 3 * 2 from the previous invocation) - ``` - """ - - value: R - """Value to return. A value will always be returned even if it is None.""" - save: S - """The value for the state for the next checkpoint. - - A value will always be saved even if it is None. - """ - - def __call__(self, func: Callable[..., Any]) -> Pregel: - """Convert a function into a Pregel graph. - - Args: - func: The function to convert. Support both sync and async functions. - - Returns: - A Pregel graph. - """ - # wrap generators in a function that writes to StreamWriter - if inspect.isgeneratorfunction(func) or inspect.isasyncgenfunction(func): - raise NotImplementedError( - "Generators are not supported in the Functional API." - ) - - bound = get_runnable_for_entrypoint(func) - stream_mode: StreamMode = "updates" - - # get input and output types - sig = inspect.signature(func) - first_parameter_name = next(iter(sig.parameters.keys()), None) - if not first_parameter_name: - raise ValueError("Entrypoint function must have at least one parameter") - input_type = ( - sig.parameters[first_parameter_name].annotation - if sig.parameters[first_parameter_name].annotation - is not inspect.Signature.empty - else Any - ) - - def _pluck_return_value(value: Any) -> Any: - """Extract the return_ value the entrypoint.final object or passthrough.""" - return value.value if isinstance(value, entrypoint.final) else value - - def _pluck_save_value(value: Any) -> Any: - """Get save value from the entrypoint.final object or passthrough.""" - return value.save if isinstance(value, entrypoint.final) else value - - output_type, save_type = Any, Any - if sig.return_annotation is not inspect.Signature.empty: - # User does not parameterize entrypoint.final properly - if ( - sig.return_annotation is entrypoint.final - ): # Un-parameterized entrypoint.final - output_type = save_type = Any - else: - origin = get_origin(sig.return_annotation) - if origin is entrypoint.final: - type_annotations = get_args(sig.return_annotation) - if len(type_annotations) != 2: - raise TypeError( - "Please an annotation for both the return_ and " - "the save values." - "For example, `-> entrypoint.final[int, str]` would assign a " - "return_ a type of `int` and save the type `str`." - ) - output_type, save_type = get_args(sig.return_annotation) - else: - output_type = save_type = sig.return_annotation - - return Pregel( - nodes={ - func.__name__: PregelNode( - bound=bound, - triggers=[START], - channels=[START], - writers=[ - ChannelWrite( - [ - ChannelWriteEntry(END, mapper=_pluck_return_value), - ChannelWriteEntry(PREVIOUS, mapper=_pluck_save_value), - ], - tags=[TAG_HIDDEN], - ) - ], - ) - }, - channels={ - START: EphemeralValue(input_type), - END: LastValue(output_type, END), - PREVIOUS: LastValue(save_type, PREVIOUS), - }, - input_channels=START, - output_channels=END, - stream_channels=END, - stream_mode=stream_mode, - stream_eager=True, - checkpointer=self.checkpointer, - store=self.store, - config_type=self.config_schema, - ) diff --git a/libs/langgraph/langgraph/func/py.typed b/libs/langgraph/langgraph/func/py.typed deleted file mode 100644 index e69de29bb..000000000 diff --git a/libs/langgraph/langgraph/graph/__init__.py b/libs/langgraph/langgraph/graph/__init__.py index c81ad9903..e24059901 100644 --- a/libs/langgraph/langgraph/graph/__init__.py +++ b/libs/langgraph/langgraph/graph/__init__.py @@ -1,13 +1,10 @@ -from langgraph.graph.graph import END, START, Graph -from langgraph.graph.message import MessageGraph, MessagesState, add_messages -from langgraph.graph.state import StateGraph +from langgraph.graph.message import MessagesState, add_messages +from langgraph.graph.state import END, START, StateGraph __all__ = [ "END", "START", - "Graph", "StateGraph", - "MessageGraph", "add_messages", "MessagesState", ] diff --git a/libs/langgraph/langgraph/graph/graph.py b/libs/langgraph/langgraph/graph/graph.py deleted file mode 100644 index 4bda590c3..000000000 --- a/libs/langgraph/langgraph/graph/graph.py +++ /dev/null @@ -1,642 +0,0 @@ -import asyncio -import logging -from collections import defaultdict -from typing import ( - Any, - Awaitable, - Callable, - Hashable, - Literal, - NamedTuple, - Optional, - Sequence, - Union, - cast, - get_args, - get_origin, - get_type_hints, - overload, -) - -from langchain_core.runnables import Runnable -from langchain_core.runnables.config import RunnableConfig -from langchain_core.runnables.graph import Graph as DrawableGraph -from langchain_core.runnables.graph import Node as DrawableNode -from typing_extensions import Self - -from langgraph.channels.ephemeral_value import EphemeralValue -from langgraph.constants import ( - EMPTY_SEQ, - END, - NS_END, - NS_SEP, - START, - TAG_HIDDEN, - Send, -) -from langgraph.errors import InvalidUpdateError -from langgraph.pregel import Channel, Pregel -from langgraph.pregel.read import PregelNode -from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry -from langgraph.types import All, Checkpointer -from langgraph.utils.runnable import RunnableCallable, RunnableLike, coerce_to_runnable - -logger = logging.getLogger(__name__) - - -class NodeSpec(NamedTuple): - runnable: Runnable - metadata: Optional[dict[str, Any]] = None - ends: Optional[Union[tuple[str, ...], dict[str, str]]] = EMPTY_SEQ - - -class Branch(NamedTuple): - path: Runnable[Any, Union[Hashable, list[Hashable]]] - ends: Optional[dict[Hashable, str]] - then: Optional[str] = None - - def run( - self, - writer: Callable[ - [Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite] - ], - reader: Optional[Callable[[RunnableConfig], Any]] = None, - ) -> RunnableCallable: - return ChannelWrite.register_writer( - RunnableCallable( - func=self._route, - afunc=self._aroute, - writer=writer, - reader=reader, - name=None, - trace=False, - ) - ) - - def _route( - self, - input: Any, - config: RunnableConfig, - *, - reader: Optional[Callable[[RunnableConfig], Any]], - writer: Callable[ - [Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite] - ], - ) -> Runnable: - if reader: - value = reader(config) - # passthrough additional keys from node to branch - # only doable when using dict states - if isinstance(value, dict) and isinstance(input, dict): - value = {**input, **value} - else: - value = input - result = self.path.invoke(value, config) - return self._finish(writer, input, result, config) - - async def _aroute( - self, - input: Any, - config: RunnableConfig, - *, - reader: Optional[Callable[[RunnableConfig], Any]], - writer: Callable[ - [Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite] - ], - ) -> Runnable: - if reader: - value = await asyncio.to_thread(reader, config) - # passthrough additional keys from node to branch - # only doable when using dict states - if isinstance(value, dict) and isinstance(input, dict): - value = {**input, **value} - else: - value = input - result = await self.path.ainvoke(value, config) - return self._finish(writer, input, result, config) - - def _finish( - self, - writer: Callable[ - [Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite] - ], - input: Any, - result: Any, - config: RunnableConfig, - ) -> Union[Runnable, Any]: - if not isinstance(result, (list, tuple)): - result = [result] - if self.ends: - destinations: Sequence[Union[Send, str]] = [ - r if isinstance(r, Send) else self.ends[r] for r in result - ] - else: - destinations = cast(Sequence[Union[Send, str]], result) - if any(dest is None or dest == START for dest in destinations): - raise ValueError("Branch did not return a valid destination") - if any(p.node == END for p in destinations if isinstance(p, Send)): - raise InvalidUpdateError("Cannot send a packet to the END node") - return writer(destinations, config) or input - - -class Graph: - def __init__(self) -> None: - self.nodes: dict[str, NodeSpec] = {} - self.edges = set[tuple[str, str]]() - self.branches: defaultdict[str, dict[str, Branch]] = defaultdict(dict) - self.support_multiple_edges = False - self.compiled = False - - @property - def _all_edges(self) -> set[tuple[str, str]]: - return self.edges - - @overload - def add_node( - self, - node: RunnableLike, - *, - metadata: Optional[dict[str, Any]] = None, - ) -> Self: ... - - @overload - def add_node( - self, - node: str, - action: RunnableLike, - *, - metadata: Optional[dict[str, Any]] = None, - ) -> Self: ... - - def add_node( - self, - node: Union[str, RunnableLike], - action: Optional[RunnableLike] = None, - *, - metadata: Optional[dict[str, Any]] = None, - ) -> Self: - if isinstance(node, str): - for character in (NS_SEP, NS_END): - if character in node: - raise ValueError( - f"'{character}' is a reserved character and is not allowed in the node names." - ) - - if self.compiled: - logger.warning( - "Adding a node to a graph that has already been compiled. This will " - "not be reflected in the compiled graph." - ) - if not isinstance(node, str): - action = node - node = getattr(action, "name", getattr(action, "__name__")) - if node is None: - raise ValueError( - "Node name must be provided if action is not a function" - ) - if action is None: - raise RuntimeError( - "Expected a function or Runnable action in add_node. Received None." - ) - if node in self.nodes: - raise ValueError(f"Node `{node}` already present.") - if node == END or node == START: - raise ValueError(f"Node `{node}` is reserved.") - - self.nodes[cast(str, node)] = NodeSpec( - coerce_to_runnable(action, name=cast(str, node), trace=False), metadata - ) - return self - - def add_edge(self, start_key: str, end_key: str) -> Self: - if self.compiled: - logger.warning( - "Adding an edge to a graph that has already been compiled. This will " - "not be reflected in the compiled graph." - ) - if start_key == END: - raise ValueError("END cannot be a start node") - if end_key == START: - raise ValueError("START cannot be an end node") - - # run this validation only for non-StateGraph graphs - if not hasattr(self, "channels") and start_key in set( - start for start, _ in self.edges - ): - raise ValueError( - f"Already found path for node '{start_key}'.\n" - "For multiple edges, use StateGraph with an Annotated state key." - ) - - self.edges.add((start_key, end_key)) - return self - - def add_conditional_edges( - self, - source: str, - path: Union[ - Callable[..., Union[Hashable, list[Hashable]]], - Callable[..., Awaitable[Union[Hashable, list[Hashable]]]], - Runnable[Any, Union[Hashable, list[Hashable]]], - ], - path_map: Optional[Union[dict[Hashable, str], list[str]]] = None, - then: Optional[str] = None, - ) -> Self: - """Add a conditional edge from the starting node to any number of destination nodes. - - Args: - source (str): The starting node. This conditional edge will run when - exiting this node. - path (Union[Callable, Runnable]): The callable that determines the next - node or nodes. If not specifying `path_map` it should return one or - more nodes. If it returns END, the graph will stop execution. - path_map (Optional[dict[Hashable, str]]): Optional mapping of paths to node - names. If omitted the paths returned by `path` should be node names. - then (Optional[str]): The name of a node to execute after the nodes - selected by `path`. - - Returns: - Self: The instance of the graph, allowing for method chaining. - - Note: Without typehints on the `path` function's return value (e.g., `-> Literal["foo", "__end__"]:`) - or a path_map, the graph visualization assumes the edge could transition to any node in the graph. - - """ # noqa: E501 - if self.compiled: - logger.warning( - "Adding an edge to a graph that has already been compiled. This will " - "not be reflected in the compiled graph." - ) - # coerce path_map to a dictionary - try: - if isinstance(path_map, dict): - path_map_ = path_map.copy() - elif isinstance(path_map, list): - path_map_ = {name: name for name in path_map} - elif isinstance(path, Runnable): - path_map_ = None - elif rtn_type := get_type_hints(path.__call__).get( # type: ignore[operator] - "return" - ) or get_type_hints(path).get("return"): - if get_origin(rtn_type) is Literal: - path_map_ = {name: name for name in get_args(rtn_type)} - else: - path_map_ = None - else: - path_map_ = None - except Exception: - path_map_ = None - # find a name for the condition - path = coerce_to_runnable(path, name=None, trace=True) - name = path.name or "condition" - # validate the condition - if name in self.branches[source]: - raise ValueError( - f"Branch with name `{path.name}` already exists for node " f"`{source}`" - ) - # save it - self.branches[source][name] = Branch(path, path_map_, then) - return self - - def set_entry_point(self, key: str) -> Self: - """Specifies the first node to be called in the graph. - - Equivalent to calling `add_edge(START, key)`. - - Parameters: - key (str): The key of the node to set as the entry point. - - Returns: - Self: The instance of the graph, allowing for method chaining. - """ - return self.add_edge(START, key) - - def set_conditional_entry_point( - self, - path: Union[ - Callable[..., Union[Hashable, list[Hashable]]], - Callable[..., Awaitable[Union[Hashable, list[Hashable]]]], - Runnable[Any, Union[Hashable, list[Hashable]]], - ], - path_map: Optional[Union[dict[Hashable, str], list[str]]] = None, - then: Optional[str] = None, - ) -> Self: - """Sets a conditional entry point in the graph. - - Args: - path (Union[Callable, Runnable]): The callable that determines the next - node or nodes. If not specifying `path_map` it should return one or - more nodes. If it returns END, the graph will stop execution. - path_map (Optional[dict[str, str]]): Optional mapping of paths to node - names. If omitted the paths returned by `path` should be node names. - then (Optional[str]): The name of a node to execute after the nodes - selected by `path`. - - Returns: - Self: The instance of the graph, allowing for method chaining. - """ - return self.add_conditional_edges(START, path, path_map, then) - - def set_finish_point(self, key: str) -> Self: - """Marks a node as a finish point of the graph. - - If the graph reaches this node, it will cease execution. - - Parameters: - key (str): The key of the node to set as the finish point. - - Returns: - Self: The instance of the graph, allowing for method chaining. - """ - return self.add_edge(key, END) - - def validate(self, interrupt: Optional[Sequence[str]] = None) -> Self: - # assemble sources - all_sources = {src for src, _ in self._all_edges} - for start, branches in self.branches.items(): - all_sources.add(start) - for cond, branch in branches.items(): - if branch.then is not None: - if branch.ends is not None: - for end in branch.ends.values(): - if end != END: - all_sources.add(end) - else: - for node in self.nodes: - if node != start and node != branch.then: - all_sources.add(node) - for name, spec in self.nodes.items(): - if spec.ends: - all_sources.add(name) - # validate sources - for source in all_sources: - if source not in self.nodes and source != START: - raise ValueError(f"Found edge starting at unknown node '{source}'") - - if START not in all_sources: - raise ValueError( - "Graph must have an entrypoint: add at least one edge from START to another node" - ) - - # assemble targets - all_targets = {end for _, end in self._all_edges} - for start, branches in self.branches.items(): - for cond, branch in branches.items(): - if branch.then is not None: - all_targets.add(branch.then) - if branch.ends is not None: - for end in branch.ends.values(): - if end not in self.nodes and end != END: - raise ValueError( - f"At '{start}' node, '{cond}' branch found unknown target '{end}'" - ) - all_targets.add(end) - else: - all_targets.add(END) - for node in self.nodes: - if node != start and node != branch.then: - all_targets.add(node) - for name, spec in self.nodes.items(): - if spec.ends: - all_targets.update(spec.ends) - for target in all_targets: - if target not in self.nodes and target != END: - raise ValueError(f"Found edge ending at unknown node `{target}`") - # validate interrupts - if interrupt: - for node in interrupt: - if node not in self.nodes: - raise ValueError(f"Interrupt node `{node}` not found") - - self.compiled = True - return self - - def compile( - self, - checkpointer: Checkpointer = None, - interrupt_before: Optional[Union[All, list[str]]] = None, - interrupt_after: Optional[Union[All, list[str]]] = None, - debug: bool = False, - name: Optional[str] = None, - ) -> "CompiledGraph": - # assign default values - interrupt_before = interrupt_before or [] - interrupt_after = interrupt_after or [] - - # validate the graph - self.validate( - interrupt=( - (interrupt_before if interrupt_before != "*" else []) + interrupt_after - if interrupt_after != "*" - else [] - ) - ) - - # create empty compiled graph - compiled = CompiledGraph( - builder=self, - nodes={}, - channels={START: EphemeralValue(Any), END: EphemeralValue(Any)}, - input_channels=START, - output_channels=END, - stream_mode="values", - stream_channels=[], - checkpointer=checkpointer, - interrupt_before_nodes=interrupt_before, - interrupt_after_nodes=interrupt_after, - auto_validate=False, - debug=debug, - name=name or "LangGraph", - ) - - # attach nodes, edges, and branches - for key, node in self.nodes.items(): - compiled.attach_node(key, node) - - for start, end in self.edges: - compiled.attach_edge(start, end) - - for start, branches in self.branches.items(): - for name, branch in branches.items(): - compiled.attach_branch(start, name, branch) - - # validate the compiled graph - return compiled.validate() - - -class CompiledGraph(Pregel): - builder: Graph - - def __init__(self, *, builder: Graph, **kwargs: Any) -> None: - super().__init__(**kwargs) - self.builder = builder - - def attach_node(self, key: str, node: NodeSpec) -> None: - self.channels[key] = EphemeralValue(Any) - self.nodes[key] = ( - PregelNode(channels=[], triggers=[], metadata=node.metadata) - | node.runnable - | ChannelWrite([ChannelWriteEntry(key)], tags=[TAG_HIDDEN]) - ) - cast(list[str], self.stream_channels).append(key) - - def attach_edge(self, start: str, end: str) -> None: - if end == END: - # publish to end channel - self.nodes[start].writers.append( - ChannelWrite([ChannelWriteEntry(END)], tags=[TAG_HIDDEN]) - ) - else: - # subscribe to start channel - self.nodes[end].triggers.append(start) - cast(list[str], self.nodes[end].channels).append(start) - - def attach_branch(self, start: str, name: str, branch: Branch) -> None: - def branch_writer( - packets: Sequence[Union[str, Send]], config: RunnableConfig - ) -> Optional[ChannelWrite]: - writes = [ - ( - ChannelWriteEntry(f"branch:{start}:{name}:{p}" if p != END else END) - if not isinstance(p, Send) - else p - ) - for p in packets - ] - return ChannelWrite( - cast(Sequence[Union[ChannelWriteEntry, Send]], writes), - tags=[TAG_HIDDEN], - ) - - # add hidden start node - if start == START and start not in self.nodes: - self.nodes[start] = Channel.subscribe_to(START, tags=[TAG_HIDDEN]) - - # attach branch writer - self.nodes[start] |= branch.run(branch_writer) - - # attach branch readers - ends = branch.ends.values() if branch.ends else [node for node in self.nodes] - for end in ends: - if end != END: - channel_name = f"branch:{start}:{name}:{end}" - self.channels[channel_name] = EphemeralValue(Any) - self.nodes[end].triggers.append(channel_name) - cast(list[str], self.nodes[end].channels).append(channel_name) - - async def aget_graph( - self, - config: Optional[RunnableConfig] = None, - *, - xray: Union[int, bool] = False, - ) -> DrawableGraph: - return self.get_graph(config, xray=xray) - - def get_graph( - self, - config: Optional[RunnableConfig] = None, - *, - xray: Union[int, bool] = False, - ) -> DrawableGraph: - """Returns a drawable representation of the computation graph.""" - graph = DrawableGraph() - start_nodes: dict[str, DrawableNode] = { - START: graph.add_node(self.get_input_schema(config), START) - } - end_nodes: dict[str, DrawableNode] = {} - if xray: - subgraphs = { - k: v for k, v in self.get_subgraphs() if isinstance(v, CompiledGraph) - } - else: - subgraphs = {} - - def add_edge( - start: str, - end: str, - label: Optional[Hashable] = None, - conditional: bool = False, - ) -> None: - if end == END and END not in end_nodes: - end_nodes[END] = graph.add_node(self.get_output_schema(config), END) - return graph.add_edge( - start_nodes[start], - end_nodes[end], - str(label) if label is not None else None, - conditional, - ) - - for key, n in self.builder.nodes.items(): - node = n.runnable - metadata = n.metadata or {} - if key in self.interrupt_before_nodes and key in self.interrupt_after_nodes: - metadata["__interrupt"] = "before,after" - elif key in self.interrupt_before_nodes: - metadata["__interrupt"] = "before" - elif key in self.interrupt_after_nodes: - metadata["__interrupt"] = "after" - if xray and key in subgraphs: - subgraph = subgraphs[key].get_graph( - config=config, - xray=xray - 1 - if isinstance(xray, int) and not isinstance(xray, bool) and xray > 0 - else xray, - ) - subgraph.trim_first_node() - subgraph.trim_last_node() - if len(subgraph.nodes) > 1: - e, s = graph.extend(subgraph, prefix=key) - if e is None: - raise ValueError( - f"Could not extend subgraph '{key}' due to missing entrypoint" - ) - if s is not None: - start_nodes[key] = s - end_nodes[key] = e - else: - nn = graph.add_node(node, key, metadata=metadata or None) - start_nodes[key] = nn - end_nodes[key] = nn - else: - nn = graph.add_node(node, key, metadata=metadata or None) - start_nodes[key] = nn - end_nodes[key] = nn - for start, end in sorted(self.builder._all_edges): - add_edge(start, end) - for start, branches in self.builder.branches.items(): - default_ends = { - **{k: k for k in self.builder.nodes if k != start}, - END: END, - } - for _, branch in branches.items(): - if branch.ends is not None: - ends = branch.ends - elif branch.then is not None: - ends = {k: k for k in default_ends if k not in (END, branch.then)} - else: - ends = cast(dict[Hashable, str], default_ends) - for label, end in ends.items(): - add_edge( - start, - end, - label if label != end else None, - conditional=True, - ) - if branch.then is not None: - add_edge(end, branch.then) - for key, n in self.builder.nodes.items(): - if isinstance(n.ends, dict): - for end, label in n.ends.items(): - add_edge(key, end, label, conditional=True) - elif isinstance(n.ends, tuple): - for end in n.ends: - add_edge(key, end, conditional=True) - - return graph - - def _repr_mimebundle_(self, **kwargs: Any) -> dict[str, Any]: - """Mime bundle used by Jupyter to display the graph""" - return { - "text/plain": repr(self), - "image/png": self.get_graph().draw_mermaid_png(), - } diff --git a/libs/langgraph/langgraph/graph/message.py b/libs/langgraph/langgraph/graph/message.py index b6ed16131..49aa013da 100644 --- a/libs/langgraph/langgraph/graph/message.py +++ b/libs/langgraph/langgraph/graph/message.py @@ -1,53 +1,51 @@ import uuid -import warnings -from functools import partial from typing import ( Annotated, Any, - Callable, Literal, Optional, + Protocol, Sequence, Union, - cast, + runtime_checkable, ) -from langchain_core.messages import ( - AnyMessage, - BaseMessage, - BaseMessageChunk, - MessageLikeRepresentation, - RemoveMessage, - convert_to_messages, - message_chunk_to_message, -) +from pydantic import BaseModel from typing_extensions import TypedDict -from langgraph.graph.state import StateGraph -Messages = Union[list[MessageLikeRepresentation], MessageLikeRepresentation] +@runtime_checkable +class MessageProtocol(Protocol): + content: Union[str, list] + id: Optional[str] -def _add_messages_wrapper(func: Callable) -> Callable[[Messages, Messages], Messages]: - def _add_messages( - left: Optional[Messages] = None, right: Optional[Messages] = None, **kwargs: Any - ) -> Union[Messages, Callable[[Messages, Messages], Messages]]: - if left is not None and right is not None: - return func(left, right, **kwargs) - elif left is not None or right is not None: - msg = ( - f"Must specify non-null arguments for both 'left' and 'right'. Only " - f"received: '{'left' if left else 'right'}'." - ) - raise ValueError(msg) - else: - return partial(func, **kwargs) - - _add_messages.__doc__ = func.__doc__ - return cast(Callable[[Messages, Messages], Messages], _add_messages) +class Message(BaseModel, extra="allow"): + role: str + content: Union[str, list] + id: Optional[str] = None + + +MessageLike = Union[MessageProtocol, list[str], tuple[str, str], str, dict[str, Any]] + +Messages = Union[list[MessageLike], MessageLike] + + +def convert_to_message( + message: MessageLike, +) -> MessageProtocol: + if isinstance(message, MessageProtocol): + return message + elif isinstance(message, str): + return Message(role="user", content=message) + elif isinstance(message, Sequence): + return Message(role=message[0], content=message[1]) + elif isinstance(message, dict): + return Message(**message) + else: + raise TypeError(f"Expected a message-like object, but got {type(message)}") -@_add_messages_wrapper def add_messages( left: Messages, right: Messages, @@ -164,14 +162,8 @@ def add_messages( if not isinstance(right, list): right = [right] # type: ignore[assignment] # coerce to message - left = [ - message_chunk_to_message(cast(BaseMessageChunk, m)) - for m in convert_to_messages(left) - ] - right = [ - message_chunk_to_message(cast(BaseMessageChunk, m)) - for m in convert_to_messages(right) - ] + left = [convert_to_message(m) for m in left] + right = [convert_to_message(m) for m in right] # assign missing ids for m in left: if m.id is None: @@ -185,98 +177,14 @@ def add_messages( ids_to_remove = set() for m in right: if (existing_idx := merged_by_id.get(m.id)) is not None: - if isinstance(m, RemoveMessage): - ids_to_remove.add(m.id) - else: - ids_to_remove.discard(m.id) - merged[existing_idx] = m + ids_to_remove.discard(m.id) + merged[existing_idx] = m else: - if isinstance(m, RemoveMessage): - raise ValueError( - f"Attempting to delete a message with an ID that doesn't exist ('{m.id}')" - ) - merged_by_id[m.id] = len(merged) merged.append(m) merged = [m for m in merged if m.id not in ids_to_remove] - - if format == "langchain-openai": - merged = _format_messages(merged) - elif format: - msg = f"Unrecognized {format=}. Expected one of 'langchain-openai', None." - raise ValueError(msg) - else: - pass - return merged -class MessageGraph(StateGraph): - """A StateGraph where every node receives a list of messages as input and returns one or more messages as output. - - MessageGraph is a subclass of StateGraph whose entire state is a single, append-only* list of messages. - Each node in a MessageGraph takes a list of messages as input and returns zero or more - messages as output. The `add_messages` function is used to merge the output messages from each node - into the existing list of messages in the graph's state. - - Examples: - ```pycon - >>> from langgraph.graph.message import MessageGraph - ... - >>> builder = MessageGraph() - >>> builder.add_node("chatbot", lambda state: [("assistant", "Hello!")]) - >>> builder.set_entry_point("chatbot") - >>> builder.set_finish_point("chatbot") - >>> builder.compile().invoke([("user", "Hi there.")]) - [HumanMessage(content="Hi there.", id='...'), AIMessage(content="Hello!", id='...')] - ``` - - ```pycon - >>> from langchain_core.messages import AIMessage, HumanMessage, ToolMessage - >>> from langgraph.graph.message import MessageGraph - ... - >>> builder = MessageGraph() - >>> builder.add_node( - ... "chatbot", - ... lambda state: [ - ... AIMessage( - ... content="Hello!", - ... tool_calls=[{"name": "search", "id": "123", "args": {"query": "X"}}], - ... ) - ... ], - ... ) - >>> builder.add_node( - ... "search", lambda state: [ToolMessage(content="Searching...", tool_call_id="123")] - ... ) - >>> builder.set_entry_point("chatbot") - >>> builder.add_edge("chatbot", "search") - >>> builder.set_finish_point("search") - >>> builder.compile().invoke([HumanMessage(content="Hi there. Can you search for X?")]) - {'messages': [HumanMessage(content="Hi there. Can you search for X?", id='b8b7d8f4-7f4d-4f4d-9c1d-f8b8d8f4d9c1'), - AIMessage(content="Hello!", id='f4d9c1d8-8d8f-4d9c-b8b7-d8f4f4d9c1d8'), - ToolMessage(content="Searching...", id='d8f4f4d9-c1d8-4f4d-b8b7-d8f4f4d9c1d8', tool_call_id="123")]} - ``` - """ - - def __init__(self) -> None: - super().__init__(Annotated[list[AnyMessage], add_messages]) # type: ignore[arg-type] - - class MessagesState(TypedDict): - messages: Annotated[list[AnyMessage], add_messages] - - -def _format_messages(messages: Sequence[BaseMessage]) -> list[BaseMessage]: - try: - from langchain_core.messages import convert_to_openai_messages - except ImportError: - msg = ( - "Must have langchain-core>=0.3.11 installed to use automatic message " - "formatting (format='langchain-openai'). Please update your langchain-core " - "version or remove the 'format' flag. Returning un-formatted " - "messages." - ) - warnings.warn(msg) - return list(messages) - else: - return convert_to_messages(convert_to_openai_messages(messages)) + messages: Annotated[list[MessageProtocol], add_messages] diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index a48356df5..2a05347b1 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -2,12 +2,14 @@ import inspect import logging import typing import warnings +from collections import defaultdict from functools import partial from inspect import isclass, isfunction, ismethod, signature from types import FunctionType from typing import ( Any, Callable, + Hashable, Literal, NamedTuple, Optional, @@ -21,9 +23,6 @@ from typing import ( overload, ) -from langchain_core.runnables import Runnable, RunnableConfig -from pydantic import BaseModel -from pydantic.v1 import BaseModel as BaseModelV1 from typing_extensions import Self from langgraph._api.deprecation import LangGraphDeprecationWarning @@ -33,22 +32,24 @@ from langgraph.channels.dynamic_barrier_value import DynamicBarrierValue, WaitFo from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue from langgraph.channels.named_barrier_value import NamedBarrierValue -from langgraph.constants import EMPTY_SEQ, MISSING, NS_END, NS_SEP, SELF, TAG_HIDDEN +from langgraph.constants import ( + EMPTY_SEQ, + END, + MISSING, + NS_END, + NS_SEP, + SELF, + START, + TAG_HIDDEN, +) from langgraph.errors import ( ErrorCode, InvalidUpdateError, ParentCommand, create_error_message, ) -from langgraph.graph.graph import END, START, Branch, CompiledGraph, Graph, Send -from langgraph.managed.base import ( - ChannelKeyPlaceholder, - ChannelTypePlaceholder, - ConfiguredManagedValue, - ManagedValueSpec, - is_managed_value, - is_writable_managed_value, -) +from langgraph.pregel import Pregel +from langgraph.pregel.protocol import PregelProtocol from langgraph.pregel.read import ChannelRead, PregelNode from langgraph.pregel.write import ( ChannelWrite, @@ -56,10 +57,14 @@ from langgraph.pregel.write import ( ChannelWriteTupleEntry, ) from langgraph.store.base import BaseStore -from langgraph.types import All, Checkpointer, Command, RetryPolicy -from langgraph.utils.fields import get_field_default -from langgraph.utils.pydantic import create_model -from langgraph.utils.runnable import RunnableCallable, RunnableLike, coerce_to_runnable +from langgraph.types import All, Checkpointer, Command, RetryPolicy, Send +from langgraph.utils.config import RunnableConfig +from langgraph.utils.runnable import ( + Runnable, + RunnableCallable, + RunnableLike, + coerce_to_runnable, +) logger = logging.getLogger(__name__) @@ -85,15 +90,81 @@ def _get_node_name(node: RunnableLike) -> str: raise TypeError(f"Unsupported node type: {type(node)}") +class Branch(NamedTuple): + path: Runnable[Any, Union[Hashable, list[Hashable]]] + ends: Optional[dict[Hashable, str]] + then: Optional[str] = None + + def run( + self, + writer: Callable[ + [Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite] + ], + reader: Optional[Callable[[RunnableConfig], Any]] = None, + ) -> RunnableCallable: + return RunnableCallable( + func=self._route, + writer=writer, + reader=reader, + name=None, + trace=False, + ) + + def _route( + self, + input: Any, + config: RunnableConfig, + *, + reader: Optional[Callable[[RunnableConfig], Any]], + writer: Callable[ + [Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite] + ], + ) -> Runnable: + if reader: + value = reader(config) + # passthrough additional keys from node to branch + # only doable when using dict states + if isinstance(value, dict) and isinstance(input, dict): + value = {**input, **value} + else: + value = input + result = self.path.invoke(value, config) + return self._finish(writer, input, result, config) + + def _finish( + self, + writer: Callable[ + [Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite] + ], + input: Any, + result: Any, + config: RunnableConfig, + ) -> Union[Runnable, Any]: + if not isinstance(result, (list, tuple)): + result = [result] + if self.ends: + destinations: Sequence[Union[Send, str]] = [ + r if isinstance(r, Send) else self.ends[r] for r in result + ] + else: + destinations = cast(Sequence[Union[Send, str]], result) + if any(dest is None or dest == START for dest in destinations): + raise ValueError("Branch did not return a valid destination") + if any(p.node == END for p in destinations if isinstance(p, Send)): + raise InvalidUpdateError("Cannot send a packet to the END node") + return writer(destinations, config) or input + + class StateNodeSpec(NamedTuple): runnable: Runnable metadata: Optional[dict[str, Any]] input: Type[Any] retry_policy: Optional[RetryPolicy] ends: Optional[Union[tuple[str, ...], dict[str, str]]] = EMPTY_SEQ + subgraphs: Optional[list[PregelProtocol]] = EMPTY_SEQ -class StateGraph(Graph): +class StateGraph: """A graph whose nodes communicate by reading and writing to a shared state. The signature of each node is State -> Partial. @@ -107,7 +178,7 @@ class StateGraph(Graph): Use this to expose configurable parameters in your API. Examples: - >>> from langchain_core.runnables import RunnableConfig + >>> from langgraph.utils.config import RunnableConfig >>> from typing_extensions import Annotated, TypedDict >>> from langgraph.checkpoint.memory import MemorySaver >>> from langgraph.graph import StateGraph @@ -145,8 +216,7 @@ class StateGraph(Graph): nodes: dict[str, StateNodeSpec] # type: ignore[assignment] channels: dict[str, BaseChannel] - managed: dict[str, ManagedValueSpec] - schemas: dict[Type[Any], dict[str, Union[BaseChannel, ManagedValueSpec]]] + schemas: dict[Type[Any], dict[str, BaseChannel]] def __init__( self, @@ -172,15 +242,20 @@ class StateGraph(Graph): input = state_schema if output is None: output = state_schema + + self.nodes: dict[str, StateNodeSpec] = {} + self.edges = set[tuple[str, str]]() + self.branches: defaultdict[str, dict[str, Branch]] = defaultdict(dict) + self.compiled = False + self.schemas = {} self.channels = {} - self.managed = {} self.schema = state_schema self.input = input self.output = output self._add_schema(state_schema) - self._add_schema(input, allow_managed=False) - self._add_schema(output, allow_managed=False) + self._add_schema(input) + self._add_schema(output) self.config_schema = config_schema self.waiting_edges: set[tuple[tuple[str, ...], str]] = set() @@ -190,18 +265,11 @@ class StateGraph(Graph): (start, end) for starts, end in self.waiting_edges for start in starts } - def _add_schema(self, schema: Type[Any], /, allow_managed: bool = True) -> None: + def _add_schema(self, schema: Type[Any]) -> None: if schema not in self.schemas: _warn_invalid_state_schema(schema) - channels, managed = _get_channels(schema) - if managed and not allow_managed: - names = ", ".join(managed) - schema_name = getattr(schema, "__name__", "") - raise ValueError( - f"Invalid managed channels detected in {schema_name}: {names}." - " Managed channels are not permitted in Input/Output schema." - ) - self.schemas[schema] = {**channels, **managed} + channels = _get_channels(schema) + self.schemas[schema] = channels for key, channel in channels.items(): if key in self.channels: if self.channels[key] != channel: @@ -213,14 +281,6 @@ class StateGraph(Graph): ) else: self.channels[key] = channel - for key, managed in managed.items(): - if key in self.managed: - if self.managed[key] != managed: - raise ValueError( - f"Managed value '{key}' already exists with a different type" - ) - else: - self.managed[key] = managed @overload def add_node( @@ -280,6 +340,7 @@ class StateGraph(Graph): input: Optional[Type[Any]] = None, retry: Optional[RetryPolicy] = None, destinations: Optional[Union[dict[str, str], tuple[str]]] = None, + subgraphs: list[PregelProtocol] = EMPTY_SEQ, ) -> Self: """Adds a new node to the state graph. @@ -420,6 +481,7 @@ class StateGraph(Graph): input=input or self.schema, retry_policy=retry, ends=ends, + subgraphs=subgraphs, ) return self @@ -440,8 +502,19 @@ class StateGraph(Graph): Returns: Self: The instance of the state graph, allowing for method chaining. """ + if self.compiled: + logger.warning( + "Adding an edge to a graph that has already been compiled. This will " + "not be reflected in the compiled graph." + ) if isinstance(start_key, str): - return super().add_edge(start_key, end_key) + if start_key == END: + raise ValueError("END cannot be a start node") + if end_key == START: + raise ValueError("START cannot be an end node") + + self.edges.add((start_key, end_key)) + return self if self.compiled: logger.warning( @@ -461,6 +534,72 @@ class StateGraph(Graph): self.waiting_edges.add((tuple(start_key), end_key)) return self + def add_conditional_edges( + self, + source: str, + path: Union[ + Callable[..., Union[Hashable, list[Hashable]]], + Runnable[Any, Union[Hashable, list[Hashable]]], + ], + path_map: Optional[Union[dict[Hashable, str], list[str]]] = None, + then: Optional[str] = None, + ) -> Self: + """Add a conditional edge from the starting node to any number of destination nodes. + + Args: + source (str): The starting node. This conditional edge will run when + exiting this node. + path (Union[Callable, Runnable]): The callable that determines the next + node or nodes. If not specifying `path_map` it should return one or + more nodes. If it returns END, the graph will stop execution. + path_map (Optional[dict[Hashable, str]]): Optional mapping of paths to node + names. If omitted the paths returned by `path` should be node names. + then (Optional[str]): The name of a node to execute after the nodes + selected by `path`. + + Returns: + Self: The instance of the graph, allowing for method chaining. + + Note: Without typehints on the `path` function's return value (e.g., `-> Literal["foo", "__end__"]:`) + or a path_map, the graph visualization assumes the edge could transition to any node in the graph. + + """ # noqa: E501 + if self.compiled: + logger.warning( + "Adding an edge to a graph that has already been compiled. This will " + "not be reflected in the compiled graph." + ) + # coerce path_map to a dictionary + try: + if isinstance(path_map, dict): + path_map_ = path_map.copy() + elif isinstance(path_map, list): + path_map_ = {name: name for name in path_map} + elif isinstance(path, Runnable): + path_map_ = None + elif rtn_type := get_type_hints(path.__call__).get( # type: ignore[operator] + "return" + ) or get_type_hints(path).get("return"): + if get_origin(rtn_type) is Literal: + path_map_ = {name: name for name in get_args(rtn_type)} + else: + path_map_ = None + else: + path_map_ = None + except Exception: + path_map_ = None + # find a name for the condition + path = coerce_to_runnable(path, name=None, trace=True) + name = path.name or "condition" + # validate the condition + if name in self.branches[source]: + raise ValueError( + f"Branch with name `{path.name}` already exists for node " f"`{source}`" + ) + # save it + self.branches[source][name] = Branch(path, path_map_, then) + return self + def add_sequence( self, nodes: Sequence[Union[RunnableLike, tuple[str, RunnableLike]]], @@ -503,6 +642,118 @@ class StateGraph(Graph): return self + def set_entry_point(self, key: str) -> Self: + """Specifies the first node to be called in the graph. + + Equivalent to calling `add_edge(START, key)`. + + Parameters: + key (str): The key of the node to set as the entry point. + + Returns: + Self: The instance of the graph, allowing for method chaining. + """ + return self.add_edge(START, key) + + def set_conditional_entry_point( + self, + path: Union[ + Callable[..., Union[Hashable, list[Hashable]]], + Runnable[Any, Union[Hashable, list[Hashable]]], + ], + path_map: Optional[Union[dict[Hashable, str], list[str]]] = None, + then: Optional[str] = None, + ) -> Self: + """Sets a conditional entry point in the graph. + + Args: + path (Union[Callable, Runnable]): The callable that determines the next + node or nodes. If not specifying `path_map` it should return one or + more nodes. If it returns END, the graph will stop execution. + path_map (Optional[dict[str, str]]): Optional mapping of paths to node + names. If omitted the paths returned by `path` should be node names. + then (Optional[str]): The name of a node to execute after the nodes + selected by `path`. + + Returns: + Self: The instance of the graph, allowing for method chaining. + """ + return self.add_conditional_edges(START, path, path_map, then) + + def set_finish_point(self, key: str) -> Self: + """Marks a node as a finish point of the graph. + + If the graph reaches this node, it will cease execution. + + Parameters: + key (str): The key of the node to set as the finish point. + + Returns: + Self: The instance of the graph, allowing for method chaining. + """ + return self.add_edge(key, END) + + def validate(self, interrupt: Optional[Sequence[str]] = None) -> Self: + # assemble sources + all_sources = {src for src, _ in self._all_edges} + for start, branches in self.branches.items(): + all_sources.add(start) + for cond, branch in branches.items(): + if branch.then is not None: + if branch.ends is not None: + for end in branch.ends.values(): + if end != END: + all_sources.add(end) + else: + for node in self.nodes: + if node != start and node != branch.then: + all_sources.add(node) + for name, spec in self.nodes.items(): + if spec.ends: + all_sources.add(name) + # validate sources + for source in all_sources: + if source not in self.nodes and source != START: + raise ValueError(f"Found edge starting at unknown node '{source}'") + + if START not in all_sources: + raise ValueError( + "Graph must have an entrypoint: add at least one edge from START to another node" + ) + + # assemble targets + all_targets = {end for _, end in self._all_edges} + for start, branches in self.branches.items(): + for cond, branch in branches.items(): + if branch.then is not None: + all_targets.add(branch.then) + if branch.ends is not None: + for end in branch.ends.values(): + if end not in self.nodes and end != END: + raise ValueError( + f"At '{start}' node, '{cond}' branch found unknown target '{end}'" + ) + all_targets.add(end) + else: + all_targets.add(END) + for node in self.nodes: + if node != start and node != branch.then: + all_targets.add(node) + for name, spec in self.nodes.items(): + if spec.ends: + all_targets.update(spec.ends) + for target in all_targets: + if target not in self.nodes and target != END: + raise ValueError(f"Found edge ending at unknown node `{target}`") + # validate interrupts + if interrupt: + for node in interrupt: + if node not in self.nodes: + raise ValueError(f"Interrupt node `{node}` not found") + + self.compiled = True + return self + def compile( self, checkpointer: Checkpointer = None, @@ -510,7 +761,6 @@ class StateGraph(Graph): store: Optional[BaseStore] = None, interrupt_before: Optional[Union[All, list[str]]] = None, interrupt_after: Optional[Union[All, list[str]]] = None, - debug: bool = False, name: Optional[str] = None, ) -> "CompiledStateGraph": """Compiles the state graph into a `CompiledGraph` object. @@ -549,18 +799,12 @@ class StateGraph(Graph): "__root__" if len(self.schemas[self.output]) == 1 and "__root__" in self.schemas[self.output] - else [ - key - for key, val in self.schemas[self.output].items() - if not is_managed_value(val) - ] + else [key for key, val in self.schemas[self.output].items()] ) stream_channels = ( "__root__" if len(self.channels) == 1 and "__root__" in self.channels - else [ - key for key, val in self.channels.items() if not is_managed_value(val) - ] + else [key for key, val in self.channels.items()] ) compiled = CompiledStateGraph( @@ -569,7 +813,6 @@ class StateGraph(Graph): nodes={}, channels={ **self.channels, - **self.managed, START: EphemeralValue(self.input), }, input_channels=START, @@ -580,7 +823,6 @@ class StateGraph(Graph): interrupt_before_nodes=interrupt_before, interrupt_after_nodes=interrupt_after, auto_validate=False, - debug=debug, store=store, name=name or "LangGraph", ) @@ -606,42 +848,20 @@ class StateGraph(Graph): return compiled.validate() -class CompiledStateGraph(CompiledGraph): +class CompiledStateGraph(Pregel): builder: StateGraph - def get_input_schema( - self, config: Optional[RunnableConfig] = None - ) -> type[BaseModel]: - return _get_schema( - typ=self.builder.input, - schemas=self.builder.schemas, - channels=self.builder.channels, - name=self.get_name("Input"), - ) - - def get_output_schema( - self, config: Optional[RunnableConfig] = None - ) -> type[BaseModel]: - return _get_schema( - typ=self.builder.output, - schemas=self.builder.schemas, - channels=self.builder.channels, - name=self.get_name("Output"), - ) + def __init__(self, *, builder: StateGraph, **kwargs: Any) -> None: + super().__init__(**kwargs) + self.builder = builder def attach_node(self, key: str, node: Optional[StateNodeSpec]) -> None: if key == START: output_keys = [ - k - for k, v in self.builder.schemas[self.builder.input].items() - if not is_managed_value(v) + k for k, v in self.builder.schemas[self.builder.input].items() ] else: - output_keys = list(self.builder.channels) + [ - k - for k, v in self.builder.managed.items() - if is_writable_managed_value(v) - ] + output_keys = list(self.builder.channels) def _get_root(input: Any) -> Optional[Sequence[tuple[str, Any]]]: if isinstance(input, Command): @@ -734,11 +954,7 @@ class CompiledStateGraph(CompiledGraph): triggers=[START], channels=[START], writers=[ - ChannelWrite( - write_entries, - tags=[TAG_HIDDEN], - require_at_least_one_of=output_keys, - ), + ChannelWrite(write_entries, tags=[TAG_HIDDEN]), ], ) elif node is not None: @@ -749,7 +965,7 @@ class CompiledStateGraph(CompiledGraph): self.channels[key] = EphemeralValue(Any, guard=False) self.nodes[key] = PregelNode( triggers=[], - # read state keys and managed values + # read state keys channels=(list(input_values) if is_single_input else input_values), # coerce state dict to schema class (eg. pydantic model) mapper=( @@ -766,6 +982,7 @@ class CompiledStateGraph(CompiledGraph): ], metadata=node.metadata, retry_policy=node.retry_policy, + subgraphs=node.subgraphs, bound=node.runnable, ) else: @@ -780,8 +997,10 @@ class CompiledStateGraph(CompiledGraph): # subscribe to channel self.nodes[end].triggers.append(channel_name) # publish to channel - self.nodes[START] |= ChannelWrite( - [ChannelWriteEntry(channel_name, START)], tags=[TAG_HIDDEN] + self.nodes[START].writers.append( + ChannelWrite( + [ChannelWriteEntry(channel_name, START)], tags=[TAG_HIDDEN] + ) ) elif end != END: # subscribe to start channel @@ -794,8 +1013,10 @@ class CompiledStateGraph(CompiledGraph): self.nodes[end].triggers.append(channel_name) # publish to channel for start in starts: - self.nodes[start] |= ChannelWrite( - [ChannelWriteEntry(channel_name, start)], tags=[TAG_HIDDEN] + self.nodes[start].writers.append( + ChannelWrite( + [ChannelWriteEntry(channel_name, start)], tags=[TAG_HIDDEN] + ) ) def attach_branch( @@ -832,9 +1053,11 @@ class CompiledStateGraph(CompiledGraph): if start in self.builder.nodes else self.builder.schema ) - self.nodes[start] |= branch.run( - branch_writer, - _get_state_reader(self.builder, schema) if with_reader else None, + self.nodes[start].writers.append( + branch.run( + branch_writer, + _get_state_reader(self.builder, schema) if with_reader else None, + ) ) # attach branch subscribers @@ -856,8 +1079,10 @@ class CompiledStateGraph(CompiledGraph): self.nodes[branch.then].triggers.append(channel_name) for end in ends: if end != END: - self.nodes[end] |= ChannelWrite( - [ChannelWriteEntry(channel_name, end)], tags=[TAG_HIDDEN] + self.nodes[end].writers.append( + ChannelWrite( + [ChannelWriteEntry(channel_name, end)], tags=[TAG_HIDDEN] + ) ) @@ -906,73 +1131,23 @@ def _control_branch(value: Any) -> Sequence[Union[str, Send]]: return rtn -async def _acontrol_branch(value: Any) -> Sequence[Union[str, Send]]: - if isinstance(value, Send): - return [value] - commands: list[Command] = [] - if isinstance(value, Command): - commands.append(value) - elif isinstance(value, (list, tuple)): - for cmd in value: - if isinstance(cmd, Command): - commands.append(cmd) - rtn: list[Union[str, Send]] = [] - for command in commands: - if command.graph == Command.PARENT: - raise ParentCommand(command) - if isinstance(command.goto, Send): - rtn.append(command.goto) - elif isinstance(command.goto, str): - rtn.append(command.goto) - else: - rtn.extend(command.goto) - return rtn - - -CONTROL_BRANCH_PATH = RunnableCallable( - _control_branch, _acontrol_branch, tags=[TAG_HIDDEN], trace=False, recurse=False -) +CONTROL_BRANCH_PATH = RunnableCallable(_control_branch, tags=[TAG_HIDDEN], trace=False) CONTROL_BRANCH = Branch(CONTROL_BRANCH_PATH, None) def _get_channels( schema: Type[dict], -) -> tuple[dict[str, BaseChannel], dict[str, ManagedValueSpec]]: - if not hasattr(schema, "__annotations__"): - return {"__root__": _get_channel("__root__", schema, allow_managed=False)}, {} - +) -> dict[str, BaseChannel]: all_keys = { name: _get_channel(name, typ) for name, typ in get_type_hints(schema, include_extras=True).items() if name != "__slots__" } - return ( - {k: v for k, v in all_keys.items() if isinstance(v, BaseChannel)}, - {k: v for k, v in all_keys.items() if is_managed_value(v)}, - ) + return {k: v for k, v in all_keys.items() if isinstance(v, BaseChannel)} -@overload -def _get_channel( - name: str, annotation: Any, *, allow_managed: Literal[False] -) -> BaseChannel: ... - - -@overload -def _get_channel( - name: str, annotation: Any, *, allow_managed: Literal[True] = True -) -> Union[BaseChannel, ManagedValueSpec]: ... - - -def _get_channel( - name: str, annotation: Any, *, allow_managed: bool = True -) -> Union[BaseChannel, ManagedValueSpec]: - if manager := _is_field_managed_value(name, annotation): - if allow_managed: - return manager - else: - raise ValueError(f"This {annotation} not allowed in this position") - elif channel := _is_field_channel(annotation): +def _get_channel(name: str, annotation: Any) -> BaseChannel: + if channel := _is_field_channel(annotation): channel.key = name return channel elif channel := _is_field_binop(annotation): @@ -1013,55 +1188,3 @@ def _is_field_binop(typ: Type[Any]) -> Optional[BinaryOperatorAggregate]: f"Invalid reducer signature. Expected (a, b) -> c. Got {sig}" ) return None - - -def _is_field_managed_value(name: str, typ: Type[Any]) -> Optional[ManagedValueSpec]: - if hasattr(typ, "__metadata__"): - meta = typ.__metadata__ - if len(meta) >= 1: - decoration = get_origin(meta[-1]) or meta[-1] - if is_managed_value(decoration): - if isinstance(decoration, ConfiguredManagedValue): - for k, v in decoration.kwargs.items(): - if v is ChannelKeyPlaceholder: - decoration.kwargs[k] = name - if v is ChannelTypePlaceholder: - decoration.kwargs[k] = typ.__origin__ - return decoration - - return None - - -def _get_schema( - typ: Type, - schemas: dict, - channels: dict, - name: str, -) -> type[BaseModel]: - if isclass(typ) and issubclass(typ, (BaseModel, BaseModelV1)): - return typ - else: - keys = list(schemas[typ].keys()) - if len(keys) == 1 and keys[0] == "__root__": - return create_model( - name, - root=(channels[keys[0]].UpdateType, None), - ) - else: - return create_model( - name, - field_definitions={ - k: ( - channels[k].UpdateType, - ( - get_field_default( - k, - channels[k].UpdateType, - typ, - ) - ), - ) - for k in schemas[typ] - if k in channels and isinstance(channels[k], BaseChannel) - }, - ) diff --git a/libs/langgraph/langgraph/managed/__init__.py b/libs/langgraph/langgraph/managed/__init__.py deleted file mode 100644 index 966348e6f..000000000 --- a/libs/langgraph/langgraph/managed/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from langgraph.managed.is_last_step import IsLastStep, RemainingSteps - -__all__ = ["IsLastStep", "RemainingSteps"] diff --git a/libs/langgraph/langgraph/managed/base.py b/libs/langgraph/langgraph/managed/base.py deleted file mode 100644 index 36962e156..000000000 --- a/libs/langgraph/langgraph/managed/base.py +++ /dev/null @@ -1,104 +0,0 @@ -from abc import ABC, abstractmethod -from contextlib import asynccontextmanager, contextmanager -from inspect import isclass -from typing import ( - Any, - AsyncIterator, - Generic, - Iterator, - NamedTuple, - Sequence, - Type, - TypeVar, - Union, -) - -from typing_extensions import Self, TypeGuard - -from langgraph.types import LoopProtocol - -V = TypeVar("V") -U = TypeVar("U") - - -class ManagedValue(ABC, Generic[V]): - def __init__(self, loop: LoopProtocol) -> None: - self.loop = loop - - @classmethod - @contextmanager - def enter(cls, loop: LoopProtocol, **kwargs: Any) -> Iterator[Self]: - try: - value = cls(loop, **kwargs) - yield value - finally: - # because managed value and Pregel have reference to each other - # let's make sure to break the reference on exit - try: - del value - except UnboundLocalError: - pass - - @classmethod - @asynccontextmanager - async def aenter(cls, loop: LoopProtocol, **kwargs: Any) -> AsyncIterator[Self]: - try: - value = cls(loop, **kwargs) - yield value - finally: - # because managed value and Pregel have reference to each other - # let's make sure to break the reference on exit - try: - del value - except UnboundLocalError: - pass - - @abstractmethod - def __call__(self) -> V: ... - - -class WritableManagedValue(Generic[V, U], ManagedValue[V], ABC): - @abstractmethod - def update(self, writes: Sequence[U]) -> None: ... - - @abstractmethod - async def aupdate(self, writes: Sequence[U]) -> None: ... - - -class ConfiguredManagedValue(NamedTuple): - cls: Type[ManagedValue] - kwargs: dict[str, Any] - - -ManagedValueSpec = Union[Type[ManagedValue], ConfiguredManagedValue] - - -def is_managed_value(value: Any) -> TypeGuard[ManagedValueSpec]: - return (isclass(value) and issubclass(value, ManagedValue)) or isinstance( - value, ConfiguredManagedValue - ) - - -def is_readonly_managed_value(value: Any) -> TypeGuard[Type[ManagedValue]]: - return ( - isclass(value) - and issubclass(value, ManagedValue) - and not issubclass(value, WritableManagedValue) - ) or ( - isinstance(value, ConfiguredManagedValue) - and not issubclass(value.cls, WritableManagedValue) - ) - - -def is_writable_managed_value(value: Any) -> TypeGuard[Type[WritableManagedValue]]: - return (isclass(value) and issubclass(value, WritableManagedValue)) or ( - isinstance(value, ConfiguredManagedValue) - and issubclass(value.cls, WritableManagedValue) - ) - - -ChannelKeyPlaceholder = object() -ChannelTypePlaceholder = object() - - -ManagedValueMapping = dict[str, ManagedValue] diff --git a/libs/langgraph/langgraph/managed/context.py b/libs/langgraph/langgraph/managed/context.py deleted file mode 100644 index d1713c11a..000000000 --- a/libs/langgraph/langgraph/managed/context.py +++ /dev/null @@ -1,108 +0,0 @@ -from contextlib import asynccontextmanager, contextmanager -from inspect import signature -from typing import ( - Any, - AsyncContextManager, - AsyncIterator, - Callable, - ContextManager, - Generic, - Iterator, - Optional, - Type, - Union, -) - -from typing_extensions import Self - -from langgraph.managed.base import ConfiguredManagedValue, ManagedValue, V -from langgraph.types import LoopProtocol - - -class Context(ManagedValue[V], Generic[V]): - runtime = True - - value: V - - @staticmethod - def of( - ctx: Union[ - None, - Callable[..., ContextManager[V]], - Type[ContextManager[V]], - Callable[..., AsyncContextManager[V]], - Type[AsyncContextManager[V]], - ] = None, - actx: Optional[ - Union[ - Callable[..., AsyncContextManager[V]], - Type[AsyncContextManager[V]], - ] - ] = None, - ) -> ConfiguredManagedValue: - if ctx is None and actx is None: - raise ValueError("Must provide either sync or async context manager.") - return ConfiguredManagedValue(Context, {"ctx": ctx, "actx": actx}) - - @classmethod - @contextmanager - def enter(cls, loop: LoopProtocol, **kwargs: Any) -> Iterator[Self]: - with super().enter(loop, **kwargs) as self: - if self.ctx is None: - raise ValueError( - "Synchronous context manager not found. Please initialize Context value with a sync context manager, or invoke your graph asynchronously." - ) - ctx = ( - self.ctx(loop.config) # type: ignore[call-arg] - if signature(self.ctx).parameters.get("config") - else self.ctx() - ) - with ctx as v: # type: ignore[union-attr] - self.value = v - yield self - - @classmethod - @asynccontextmanager - async def aenter(cls, loop: LoopProtocol, **kwargs: Any) -> AsyncIterator[Self]: - async with super().aenter(loop, **kwargs) as self: - if self.actx is not None: - ctx = ( - self.actx(loop.config) # type: ignore[call-arg] - if signature(self.actx).parameters.get("config") - else self.actx() - ) - elif self.ctx is not None: - ctx = ( - self.ctx(loop.config) # type: ignore - if signature(self.ctx).parameters.get("config") - else self.ctx() - ) - else: - raise ValueError( - "Asynchronous context manager not found. Please initialize Context value with an async context manager, or invoke your graph synchronously." - ) - if hasattr(ctx, "__aenter__"): - async with ctx as v: - self.value = v - yield self - elif hasattr(ctx, "__enter__") and hasattr(ctx, "__exit__"): - with ctx as v: - self.value = v - yield self - else: - raise ValueError( - "Context manager must have either __enter__ or __aenter__ method." - ) - - def __init__( - self, - loop: LoopProtocol, - *, - ctx: Union[None, Type[ContextManager[V]], Type[AsyncContextManager[V]]] = None, - actx: Optional[Type[AsyncContextManager[V]]] = None, - ) -> None: - self.ctx = ctx - self.actx = actx - - def __call__(self) -> V: - return self.value diff --git a/libs/langgraph/langgraph/managed/is_last_step.py b/libs/langgraph/langgraph/managed/is_last_step.py deleted file mode 100644 index 9f25a8121..000000000 --- a/libs/langgraph/langgraph/managed/is_last_step.py +++ /dev/null @@ -1,19 +0,0 @@ -from typing import Annotated - -from langgraph.managed.base import ManagedValue - - -class IsLastStepManager(ManagedValue[bool]): - def __call__(self) -> bool: - return self.loop.step == self.loop.stop - 1 - - -IsLastStep = Annotated[bool, IsLastStepManager] - - -class RemainingStepsManager(ManagedValue[int]): - def __call__(self) -> int: - return self.loop.stop - self.loop.step - - -RemainingSteps = Annotated[int, RemainingStepsManager] diff --git a/libs/langgraph/langgraph/managed/py.typed b/libs/langgraph/langgraph/managed/py.typed deleted file mode 100644 index e69de29bb..000000000 diff --git a/libs/langgraph/langgraph/managed/shared_value.py b/libs/langgraph/langgraph/managed/shared_value.py deleted file mode 100644 index 300d36c7d..000000000 --- a/libs/langgraph/langgraph/managed/shared_value.py +++ /dev/null @@ -1,123 +0,0 @@ -import collections.abc -from contextlib import asynccontextmanager, contextmanager -from typing import ( - Any, - AsyncIterator, - Iterator, - Optional, - Sequence, - Type, -) - -from typing_extensions import NotRequired, Required, Self - -from langgraph.constants import CONF -from langgraph.errors import InvalidUpdateError -from langgraph.managed.base import ( - ChannelKeyPlaceholder, - ChannelTypePlaceholder, - ConfiguredManagedValue, - WritableManagedValue, -) -from langgraph.store.base import PutOp -from langgraph.types import LoopProtocol - -V = dict[str, Any] - - -Value = dict[str, V] -Update = dict[str, Optional[V]] - - -# Adapted from typing_extensions -def _strip_extras(t): # type: ignore[no-untyped-def] - """Strips Annotated, Required and NotRequired from a given type.""" - if hasattr(t, "__origin__"): - return _strip_extras(t.__origin__) - if hasattr(t, "__origin__") and t.__origin__ in (Required, NotRequired): - return _strip_extras(t.__args__[0]) - - return t - - -class SharedValue(WritableManagedValue[Value, Update]): - @staticmethod - def on(scope: str) -> ConfiguredManagedValue: - return ConfiguredManagedValue( - SharedValue, - { - "scope": scope, - "key": ChannelKeyPlaceholder, - "typ": ChannelTypePlaceholder, - }, - ) - - @classmethod - @contextmanager - def enter(cls, loop: LoopProtocol, **kwargs: Any) -> Iterator[Self]: - with super().enter(loop, **kwargs) as value: - if loop.store is not None: - saved = loop.store.search(value.ns) - value.value = {it.key: it.value for it in saved} - yield value - - @classmethod - @asynccontextmanager - async def aenter(cls, loop: LoopProtocol, **kwargs: Any) -> AsyncIterator[Self]: - async with super().aenter(loop, **kwargs) as value: - if loop.store is not None: - saved = await loop.store.asearch(value.ns) - value.value = {it.key: it.value for it in saved} - yield value - - def __init__( - self, loop: LoopProtocol, *, typ: Type[Any], scope: str, key: str - ) -> None: - super().__init__(loop) - if typ := _strip_extras(typ): - if typ not in ( - dict, - collections.abc.Mapping, - collections.abc.MutableMapping, - ): - raise ValueError("SharedValue must be a dict") - self.scope = scope - self.value: Value = {} - if self.loop.store is None: - pass - elif scope_value := self.loop.config[CONF].get(self.scope): - self.ns = ("scoped", scope, key, scope_value) - else: - raise ValueError( - f"Scope {scope} for shared state key not in config.configurable" - ) - - def __call__(self) -> Value: - return self.value - - def _process_update(self, values: Sequence[Update]) -> list[PutOp]: - writes: list[PutOp] = [] - for vv in values: - for k, v in vv.items(): - if v is None: - if k in self.value: - del self.value[k] - writes.append(PutOp(self.ns, k, None)) - elif not isinstance(v, dict): - raise InvalidUpdateError("Received a non-dict value") - else: - self.value[k] = v - writes.append(PutOp(self.ns, k, v)) - return writes - - def update(self, values: Sequence[Update]) -> None: - if self.loop.store is None: - self._process_update(values) - else: - return self.loop.store.batch(self._process_update(values)) - - async def aupdate(self, writes: Sequence[Update]) -> None: - if self.loop.store is None: - self._process_update(writes) - else: - return await self.loop.store.abatch(self._process_update(writes)) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 68f4b8dc1..dd7ebc747 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -1,6 +1,5 @@ from __future__ import annotations -import asyncio import concurrent import concurrent.futures import queue @@ -8,37 +7,17 @@ from collections import deque from functools import partial from typing import ( Any, - AsyncIterator, Callable, Dict, Iterator, - Mapping, Optional, Sequence, Type, Union, cast, - overload, ) from uuid import UUID, uuid5 -from langchain_core.globals import get_debug -from langchain_core.runnables import ( - RunnableSequence, -) -from langchain_core.runnables.base import Input, Output -from langchain_core.runnables.config import ( - RunnableConfig, - get_async_callback_manager_for_config, - get_callback_manager_for_config, -) -from langchain_core.runnables.graph import Graph -from langchain_core.runnables.utils import ( - ConfigurableFieldSpec, - get_unique_config_specs, -) -from langchain_core.tracers._streaming import _StreamingCallbackHandler -from pydantic import BaseModel from typing_extensions import Self from langgraph.channels.base import ( @@ -46,6 +25,7 @@ from langgraph.channels.base import ( ) from langgraph.checkpoint.base import ( BaseCheckpointSaver, + CheckpointConfig, CheckpointTuple, copy_checkpoint, create_checkpoint, @@ -58,7 +38,6 @@ from langgraph.constants import ( CONFIG_KEY_CHECKPOINTER, CONFIG_KEY_NODE_FINISHED, CONFIG_KEY_READ, - CONFIG_KEY_RESUMING, CONFIG_KEY_RUNNER_SUBMIT, CONFIG_KEY_SEND, CONFIG_KEY_STORE, @@ -81,7 +60,6 @@ from langgraph.errors import ( InvalidUpdateError, create_error_message, ) -from langgraph.managed.base import ManagedValueSpec from langgraph.pregel.algo import ( PregelTaskWrites, apply_writes, @@ -91,11 +69,10 @@ from langgraph.pregel.algo import ( ) from langgraph.pregel.debug import tasks_w_writes from langgraph.pregel.io import read_channels -from langgraph.pregel.loop import AsyncPregelLoop, StreamProtocol, SyncPregelLoop -from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager -from langgraph.pregel.messages import StreamMessagesHandler +from langgraph.pregel.loop import StreamProtocol, SyncPregelLoop +from langgraph.pregel.manager import ChannelsManager from langgraph.pregel.protocol import PregelProtocol -from langgraph.pregel.read import PregelNode +from langgraph.pregel.read import DEFAULT_BOUND, PregelNode from langgraph.pregel.retry import RetryPolicy from langgraph.pregel.runner import PregelRunner from langgraph.pregel.utils import get_new_channel_versions @@ -105,88 +82,139 @@ from langgraph.store.base import BaseStore from langgraph.types import ( All, Checkpointer, - LoopProtocol, StateSnapshot, - StreamChunk, StreamMode, ) from langgraph.utils.config import ( + AnyConfig, + RunnableConfig, ensure_config, + get_runtree_for_config, merge_configs, patch_checkpoint_map, patch_config, patch_configurable, recast_checkpoint_ns, ) -from langgraph.utils.fields import get_enhanced_type_hints -from langgraph.utils.pydantic import create_model -from langgraph.utils.queue import AsyncQueue, SyncQueue # type: ignore[attr-defined] +from langgraph.utils.queue import SyncQueue # type: ignore[attr-defined] +from langgraph.utils.runnable import ( + Runnable, + RunnableLike, + RunnableSeq, + coerce_to_runnable, +) -WriteValue = Union[Callable[[Input], Output], Any] +WriteValue = Union[Callable[[Any], Any], Any] -class Channel: - @overload - @classmethod +class NodeBuilder: + channels: Union[list[str], dict[str, str]] + triggers: list[str] + tags: list[str] + writes: list[ChannelWriteEntry] + bound: Runnable + + def __init__( + self, + ) -> None: + self.channels = {} + self.triggers = [] + self.tags = [] + self.writes = [] + self.bound = DEFAULT_BOUND + def subscribe_to( - cls, - channels: str, - *, - key: Optional[str] = None, - tags: Optional[list[str]] = None, - ) -> PregelNode: ... - - @overload - @classmethod - def subscribe_to( - cls, - channels: Sequence[str], - *, - key: None = None, - tags: Optional[list[str]] = None, - ) -> PregelNode: ... - - @classmethod - def subscribe_to( - cls, + self, channels: Union[str, Sequence[str]], *, key: Optional[str] = None, tags: Optional[list[str]] = None, - ) -> PregelNode: - """Runs process.invoke() each time channels are updated, - with a dict of the channel values as input.""" + ) -> Self: + """Add channels to subscribe to with optional key and tags. + + Args: + channels: Channel name(s) to subscribe to + key: Optional key to use for the channel in input + tags: Optional tags to add to the node + + Returns: + Self for chaining + """ if not isinstance(channels, str) and key is not None: raise ValueError( "Can't specify a key when subscribing to multiple channels" ) - return PregelNode( - channels=cast( - Union[list[str], Mapping[str, str]], - ( - {key: channels} - if isinstance(channels, str) and key is not None - else ( - [channels] - if isinstance(channels, str) - else {chan: chan for chan in channels} - ) - ), - ), - triggers=[channels] if isinstance(channels, str) else channels, - tags=tags, - ) - @classmethod + if isinstance(channels, str) and key is not None: + self.channels = {key: channels} + elif isinstance(channels, str): + if isinstance(self.channels, list): + self.channels.append(channels) + elif not self.channels: + self.channels = [channels] + else: + self.channels = list(self.channels.values()) + self.channels.append(channels) + else: + if not self.channels: + self.channels = {chan: chan for chan in channels} + elif isinstance(self.channels, list): + self.channels = { + **{chan: chan for chan in self.channels}, + **{chan: chan for chan in channels}, + } + else: + self.channels.update({chan: chan for chan in channels}) + + if isinstance(channels, str): + self.triggers.append(channels) + else: + self.triggers.extend(channels) + + if tags: + self.tags.extend(tags) + + return self + + def read_from( + self, + *channels: str, + ) -> Self: + """Adds the specified channels to read from, without subscribing to them.""" + assert self.channels and isinstance( + self.channels, dict + ), "Channels must be specified first" + self.channels.update({c: c for c in channels}) + return self + + def add_node( + self, + node: RunnableLike, + ) -> Self: + """Adds the specified node.""" + if self.bound is not DEFAULT_BOUND: + self.bound = RunnableSeq(self.bound, coerce_to_runnable(node)) + else: + self.bound = coerce_to_runnable(node) + return self + def write_to( - cls, + self, *channels: str, **kwargs: WriteValue, - ) -> ChannelWrite: - """Writes to channels the result of the lambda, or None to skip writing.""" - return ChannelWrite( - [ChannelWriteEntry(c) for c in channels] - + [ + ) -> Self: + """Add channel writes. + + Args: + *channels: Channel names to write to + **kwargs: Channel name and value mappings + + Returns: + Self for chaining + """ + self.writes.extend([ChannelWriteEntry(c) for c in channels]) + self.writes.extend( + [ ( ChannelWriteEntry(k, mapper=v) if callable(v) @@ -196,6 +224,19 @@ class Channel: ] ) + return self + + def build(self) -> PregelNode: + """Builds the node.""" + assert self.triggers, "No channels specified" + return PregelNode( + channels=self.channels, + triggers=self.triggers, + tags=self.tags, + writers=[ChannelWrite(self.writes)], + bound=self.bound, + ) + class Pregel(PregelProtocol): """Pregel manages the runtime behavior for LangGraph applications. @@ -458,7 +499,7 @@ class Pregel(PregelProtocol): nodes: dict[str, PregelNode] - channels: dict[str, Union[BaseChannel, ManagedValueSpec]] + channels: dict[str, BaseChannel] stream_mode: StreamMode = "values" """Mode to stream output, defaults to 'values'.""" @@ -481,9 +522,6 @@ class Pregel(PregelProtocol): step_timeout: Optional[float] = None """Maximum time to wait for a step to complete, in seconds. Defaults to None.""" - debug: bool - """Whether to print debug information during execution. Defaults to False.""" - checkpointer: Checkpointer = None """Checkpointer used to save and load graph state. Defaults to None.""" @@ -503,7 +541,7 @@ class Pregel(PregelProtocol): self, *, nodes: dict[str, PregelNode], - channels: Optional[dict[str, Union[BaseChannel, ManagedValueSpec]]], + channels: Optional[dict[str, BaseChannel]], auto_validate: bool = True, stream_mode: StreamMode = "values", stream_eager: bool = False, @@ -513,7 +551,6 @@ class Pregel(PregelProtocol): interrupt_before_nodes: Union[All, Sequence[str]] = (), input_channels: Union[str, Sequence[str]], step_timeout: Optional[float] = None, - debug: Optional[bool] = None, checkpointer: Optional[BaseCheckpointSaver] = None, store: Optional[BaseStore] = None, retry_policy: Optional[RetryPolicy] = None, @@ -531,7 +568,6 @@ class Pregel(PregelProtocol): self.interrupt_before_nodes = interrupt_before_nodes self.input_channels = input_channels self.step_timeout = step_timeout - self.debug = debug if debug is not None else get_debug() self.checkpointer = checkpointer self.store = store self.retry_policy = retry_policy @@ -541,16 +577,6 @@ class Pregel(PregelProtocol): if auto_validate: self.validate() - def get_graph( - self, config: RunnableConfig | None = None, *, xray: int | bool = False - ) -> Graph: - raise NotImplementedError - - async def aget_graph( - self, config: RunnableConfig | None = None, *, xray: int | bool = False - ) -> Graph: - raise NotImplementedError - def copy(self, update: dict[str, Any] | None = None) -> Self: attrs = {**self.__dict__, **(update or {})} return self.__class__(**attrs) @@ -572,107 +598,6 @@ class Pregel(PregelProtocol): ) return self - @property - def config_specs(self) -> list[ConfigurableFieldSpec]: - return [ - spec - for spec in get_unique_config_specs( - [spec for node in self.nodes.values() for spec in node.config_specs] - + ( - self.checkpointer.config_specs - if isinstance(self.checkpointer, BaseCheckpointSaver) - else [] - ) - + ( - [ - ConfigurableFieldSpec( - id=name, - annotation=typ, - default=default, - description=description, - ) - for name, typ, default, description in get_enhanced_type_hints( - self.config_type - ) - ] - if self.config_type is not None - else [] - ) - ) - # these are provided by the Pregel class - if spec.id - not in [ - CONFIG_KEY_READ, - CONFIG_KEY_SEND, - CONFIG_KEY_CHECKPOINTER, - CONFIG_KEY_RESUMING, - ] - ] - - @property - def InputType(self) -> Any: - if isinstance(self.input_channels, str): - channel = self.channels[self.input_channels] - if isinstance(channel, BaseChannel): - return channel.UpdateType - - def get_input_schema( - self, config: Optional[RunnableConfig] = None - ) -> Type[BaseModel]: - config = merge_configs(self.config, config) - if isinstance(self.input_channels, str): - return super().get_input_schema(config) - else: - return create_model( - self.get_name("Input"), - field_definitions={ - k: (c.UpdateType, None) - for k in self.input_channels or self.channels.keys() - if (c := self.channels[k]) and isinstance(c, BaseChannel) - }, - ) - - def get_input_jsonschema( - self, config: Optional[RunnableConfig] = None - ) -> Dict[All, Any]: - schema = self.get_input_schema(config) - if hasattr(schema, "model_json_schema"): - return schema.model_json_schema() - else: - return schema.schema() - - @property - def OutputType(self) -> Any: - if isinstance(self.output_channels, str): - channel = self.channels[self.output_channels] - if isinstance(channel, BaseChannel): - return channel.ValueType - - def get_output_schema( - self, config: Optional[RunnableConfig] = None - ) -> Type[BaseModel]: - config = merge_configs(self.config, config) - if isinstance(self.output_channels, str): - return super().get_output_schema(config) - else: - return create_model( - self.get_name("Output"), - field_definitions={ - k: (c.ValueType, None) - for k in self.output_channels - if (c := self.channels[k]) and isinstance(c, BaseChannel) - }, - ) - - def get_output_jsonschema( - self, config: Optional[RunnableConfig] = None - ) -> Dict[All, Any]: - schema = self.get_output_schema(config) - if hasattr(schema, "model_json_schema"): - return schema.model_json_schema() - else: - return schema.schema() - @property def stream_channels_list(self) -> Sequence[str]: stream_channels = self.stream_channels_asis @@ -715,15 +640,9 @@ class Pregel(PregelProtocol): ) ) - async def aget_subgraphs( - self, *, namespace: Optional[str] = None, recurse: bool = False - ) -> AsyncIterator[tuple[str, PregelProtocol]]: - for name, node in self.get_subgraphs(namespace=namespace, recurse=recurse): - yield name, node - def _prepare_state_snapshot( self, - config: RunnableConfig, + config: CheckpointConfig, saved: Optional[CheckpointTuple], recurse: Optional[BaseCheckpointSaver] = None, apply_pending_writes: bool = False, @@ -742,20 +661,13 @@ class Pregel(PregelProtocol): with ChannelsManager( self.channels, saved.checkpoint, - LoopProtocol( - config=saved.config, - step=saved.metadata.get("step", -1) + 1, - stop=saved.metadata.get("step", -1) + 2, - ), - skip_context=True, - ) as (channels, managed): + ) as channels: # tasks for this checkpoint next_tasks = prepare_next_tasks( saved.checkpoint, saved.pending_writes or [], self.nodes, channels, - managed, saved.config, saved.metadata.get("step", -1) + 1, for_execution=True, @@ -763,7 +675,6 @@ class Pregel(PregelProtocol): checkpointer=self.checkpointer if isinstance(self.checkpointer, BaseCheckpointSaver) else None, - manager=None, ) # get the subgraphs subgraphs = dict(self.get_subgraphs()) @@ -832,123 +743,7 @@ class Pregel(PregelProtocol): ), ) - async def _aprepare_state_snapshot( - self, - config: RunnableConfig, - saved: Optional[CheckpointTuple], - recurse: Optional[BaseCheckpointSaver] = None, - apply_pending_writes: bool = False, - ) -> StateSnapshot: - if not saved: - return StateSnapshot( - values={}, - next=(), - config=config, - metadata=None, - created_at=None, - parent_config=None, - tasks=(), - ) - - async with AsyncChannelsManager( - self.channels, - saved.checkpoint, - LoopProtocol( - config=saved.config, - step=saved.metadata.get("step", -1) + 1, - stop=saved.metadata.get("step", -1) + 2, - ), - skip_context=True, - ) as ( - channels, - managed, - ): - # tasks for this checkpoint - next_tasks = prepare_next_tasks( - saved.checkpoint, - saved.pending_writes or [], - self.nodes, - channels, - managed, - saved.config, - saved.metadata.get("step", -1) + 1, - for_execution=True, - store=self.store, - checkpointer=self.checkpointer - if isinstance(self.checkpointer, BaseCheckpointSaver) - else None, - manager=None, - ) - # get the subgraphs - subgraphs = {n: g async for n, g in self.aget_subgraphs()} - parent_ns = saved.config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "") - task_states: dict[str, Union[RunnableConfig, StateSnapshot]] = {} - for task in next_tasks.values(): - if task.name not in subgraphs: - continue - # assemble checkpoint_ns for this task - task_ns = f"{task.name}{NS_END}{task.id}" - if parent_ns: - task_ns = f"{parent_ns}{NS_SEP}{task_ns}" - if not recurse: - # set config as signal that subgraph checkpoints exist - config = { - CONF: { - "thread_id": saved.config[CONF]["thread_id"], - CONFIG_KEY_CHECKPOINT_NS: task_ns, - } - } - task_states[task.id] = config - else: - # get the state of the subgraph - config = { - CONF: { - CONFIG_KEY_CHECKPOINTER: recurse, - "thread_id": saved.config[CONF]["thread_id"], - CONFIG_KEY_CHECKPOINT_NS: task_ns, - } - } - task_states[task.id] = await subgraphs[task.name].aget_state( - config, subgraphs=True - ) - # apply pending writes - if null_writes := [ - w[1:] for w in saved.pending_writes or [] if w[0] == NULL_TASK_ID - ]: - apply_writes( - saved.checkpoint, - channels, - [PregelTaskWrites((), INPUT, null_writes, [])], - None, - ) - if apply_pending_writes and saved.pending_writes: - for tid, k, v in saved.pending_writes: - if k in (ERROR, INTERRUPT, SCHEDULED): - continue - if tid not in next_tasks: - continue - next_tasks[tid].writes.append((k, v)) - if tasks := [t for t in next_tasks.values() if t.writes]: - apply_writes(saved.checkpoint, channels, tasks, None) - # assemble the state snapshot - return StateSnapshot( - read_channels(channels, self.stream_channels_asis), - tuple(t.name for t in next_tasks.values() if not t.writes), - patch_checkpoint_map(saved.config, saved.metadata), - saved.metadata, - saved.checkpoint["ts"], - patch_checkpoint_map(saved.parent_config, saved.metadata), - tasks_w_writes( - next_tasks.values(), - saved.pending_writes, - task_states, - self.stream_channels_asis, - ), - ) - - def get_state( - self, config: RunnableConfig, *, subgraphs: bool = False - ) -> StateSnapshot: + def get_state(self, config: AnyConfig, *, subgraphs: bool = False) -> StateSnapshot: """Get the current state of the graph.""" checkpointer: Optional[BaseCheckpointSaver] = ensure_config(config)[CONF].get( CONFIG_KEY_CHECKPOINTER, self.checkpointer @@ -985,48 +780,9 @@ class Pregel(PregelProtocol): apply_pending_writes=CONFIG_KEY_CHECKPOINT_ID not in config[CONF], ) - async def aget_state( - self, config: RunnableConfig, *, subgraphs: bool = False - ) -> StateSnapshot: - """Get the current state of the graph.""" - checkpointer: Optional[BaseCheckpointSaver] = ensure_config(config)[CONF].get( - CONFIG_KEY_CHECKPOINTER, self.checkpointer - ) - if not checkpointer: - raise ValueError("No checkpointer set") - - if ( - checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "") - ) and CONFIG_KEY_CHECKPOINTER not in config[CONF]: - # remove task_ids from checkpoint_ns - recast = recast_checkpoint_ns(checkpoint_ns) - # find the subgraph with the matching name - async for _, pregel in self.aget_subgraphs(namespace=recast, recurse=True): - return await pregel.aget_state( - patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}), - subgraphs=subgraphs, - ) - else: - raise ValueError(f"Subgraph {recast} not found") - - config = merge_configs(self.config, config) if self.config else config - if self.checkpointer is True: - ns = cast(str, config[CONF][CONFIG_KEY_CHECKPOINT_NS]) - config = merge_configs( - config, {CONF: {CONFIG_KEY_CHECKPOINT_NS: recast_checkpoint_ns(ns)}} - ) - - saved = await checkpointer.aget_tuple(config) - return await self._aprepare_state_snapshot( - config, - saved, - recurse=checkpointer if subgraphs else None, - apply_pending_writes=CONFIG_KEY_CHECKPOINT_ID not in config[CONF], - ) - def get_state_history( self, - config: RunnableConfig, + config: AnyConfig, *, filter: Optional[Dict[str, Any]] = None, before: Optional[RunnableConfig] = None, @@ -1070,62 +826,12 @@ class Pregel(PregelProtocol): checkpoint_tuple.config, checkpoint_tuple ) - async def aget_state_history( - self, - config: RunnableConfig, - *, - filter: Optional[Dict[str, Any]] = None, - before: Optional[RunnableConfig] = None, - limit: Optional[int] = None, - ) -> AsyncIterator[StateSnapshot]: - config = ensure_config(config) - """Get the history of the state of the graph.""" - checkpointer: Optional[BaseCheckpointSaver] = ensure_config(config)[CONF].get( - CONFIG_KEY_CHECKPOINTER, self.checkpointer - ) - if not checkpointer: - raise ValueError("No checkpointer set") - - if ( - checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "") - ) and CONFIG_KEY_CHECKPOINTER not in config[CONF]: - # remove task_ids from checkpoint_ns - recast = recast_checkpoint_ns(checkpoint_ns) - # find the subgraph with the matching name - async for _, pregel in self.aget_subgraphs(namespace=recast, recurse=True): - async for state in pregel.aget_state_history( - patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}), - filter=filter, - before=before, - limit=limit, - ): - yield state - return - else: - raise ValueError(f"Subgraph {recast} not found") - - config = merge_configs( - self.config, - config, - {CONF: {CONFIG_KEY_CHECKPOINT_NS: checkpoint_ns}}, - ) - # eagerly consume list() to avoid holding up the db cursor - for checkpoint_tuple in [ - c - async for c in checkpointer.alist( - config, before=before, limit=limit, filter=filter - ) - ]: - yield await self._aprepare_state_snapshot( - checkpoint_tuple.config, checkpoint_tuple - ) - def update_state( self, - config: RunnableConfig, + config: AnyConfig, values: Optional[Union[dict[str, Any], Any]], as_node: Optional[str] = None, - ) -> RunnableConfig: + ) -> AnyConfig: """Update the state of the graph with the given values, as if they came from node `as_node`. If `as_node` is not provided, it will be set to the last node that updated the state, if not ambiguous. @@ -1172,8 +878,7 @@ class Pregel(PregelProtocol): with ChannelsManager( self.channels, checkpoint, - LoopProtocol(config=config, step=step + 1, stop=step + 2), - ) as (channels, managed): + ) as channels: # no values as END, just clear all tasks if values is None and as_node == END: if saved is not None: @@ -1183,7 +888,6 @@ class Pregel(PregelProtocol): saved.pending_writes or [], self.nodes, channels, - managed, saved.config, saved.metadata.get("step", -1) + 1, for_execution=True, @@ -1191,7 +895,6 @@ class Pregel(PregelProtocol): checkpointer=self.checkpointer if isinstance(self.checkpointer, BaseCheckpointSaver) else None, - manager=None, ) # apply null writes if null_writes := [ @@ -1279,7 +982,6 @@ class Pregel(PregelProtocol): saved.pending_writes, self.nodes, channels, - managed, saved.config, saved.metadata.get("step", -1) + 1, for_execution=True, @@ -1287,7 +989,6 @@ class Pregel(PregelProtocol): checkpointer=self.checkpointer if isinstance(self.checkpointer, BaseCheckpointSaver) else None, - manager=None, ) # apply null writes if null_writes := [ @@ -1341,7 +1042,7 @@ class Pregel(PregelProtocol): writes: deque[tuple[str, Any]] = deque() task = PregelTaskWrites((), as_node, writes, [INTERRUPT]) task_id = str(uuid5(UUID(checkpoint["id"]), INTERRUPT)) - run = RunnableSequence(*writers) if len(writers) > 1 else writers[0] + run = RunnableSeq(*writers) if len(writers) > 1 else writers[0] # execute task run.invoke( values, @@ -1357,12 +1058,9 @@ class Pregel(PregelProtocol): ), CONFIG_KEY_READ: partial( local_read, - step + 1, checkpoint, channels, - managed, task, - config, ), }, ), @@ -1377,10 +1075,7 @@ class Pregel(PregelProtocol): if saved and channel_writes: checkpointer.put_writes(checkpoint_config, channel_writes, task_id) # apply to checkpoint and save - mv_writes = apply_writes( - checkpoint, channels, [task], checkpointer.get_next_version - ) - assert not mv_writes, "Can't write to SharedValues from update_state" + apply_writes(checkpoint, channels, [task], checkpointer.get_next_version) checkpoint = create_checkpoint(checkpoint, channels, step + 1) next_config = checkpointer.put( checkpoint_config, @@ -1400,301 +1095,17 @@ class Pregel(PregelProtocol): checkpointer.put_writes(next_config, push_writes, task_id) return patch_checkpoint_map(next_config, saved.metadata if saved else None) - async def aupdate_state( - self, - config: RunnableConfig, - values: dict[str, Any] | Any, - as_node: Optional[str] = None, - ) -> RunnableConfig: - """Update the state of the graph asynchronously with the given values, as if they came from - node `as_node`. If `as_node` is not provided, it will be set to the last node - that updated the state, if not ambiguous. - """ - checkpointer: Optional[BaseCheckpointSaver] = ensure_config(config)[CONF].get( - CONFIG_KEY_CHECKPOINTER, self.checkpointer - ) - if not checkpointer: - raise ValueError("No checkpointer set") - - # delegate to subgraph - if ( - checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "") - ) and CONFIG_KEY_CHECKPOINTER not in config[CONF]: - # remove task_ids from checkpoint_ns - recast = recast_checkpoint_ns(checkpoint_ns) - # find the subgraph with the matching name - async for _, pregel in self.aget_subgraphs(namespace=recast, recurse=True): - return await pregel.aupdate_state( - patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}), - values, - as_node, - ) - else: - raise ValueError(f"Subgraph {recast} not found") - - # get last checkpoint - config = ensure_config(self.config, config) - saved = await checkpointer.aget_tuple(config) - checkpoint = copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint() - checkpoint_previous_versions = ( - saved.checkpoint["channel_versions"].copy() if saved else {} - ) - step = saved.metadata.get("step", -1) if saved else -1 - # merge configurable fields with previous checkpoint config - checkpoint_config = patch_configurable( - config, - {CONFIG_KEY_CHECKPOINT_NS: config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "")}, - ) - checkpoint_metadata = config["metadata"] - if saved: - checkpoint_config = patch_configurable(config, saved.config[CONF]) - checkpoint_metadata = {**saved.metadata, **checkpoint_metadata} - async with AsyncChannelsManager( - self.channels, - checkpoint, - LoopProtocol(config=config, step=step + 1, stop=step + 2), - ) as ( - channels, - managed, - ): - # no values, just clear all tasks - if values is None and as_node == END: - if saved is not None: - # tasks for this checkpoint - next_tasks = prepare_next_tasks( - checkpoint, - saved.pending_writes or [], - self.nodes, - channels, - managed, - saved.config, - saved.metadata.get("step", -1) + 1, - for_execution=True, - store=self.store, - checkpointer=self.checkpointer - if isinstance(self.checkpointer, BaseCheckpointSaver) - else None, - manager=None, - ) - # apply null writes - if null_writes := [ - w[1:] - for w in saved.pending_writes or [] - if w[0] == NULL_TASK_ID - ]: - apply_writes( - saved.checkpoint, - channels, - [PregelTaskWrites((), INPUT, null_writes, [])], - None, - ) - # apply writes from tasks that already ran - for tid, k, v in saved.pending_writes or []: - if k in (ERROR, INTERRUPT, SCHEDULED): - continue - if tid not in next_tasks: - continue - next_tasks[tid].writes.append((k, v)) - # clear all current tasks - apply_writes(checkpoint, channels, next_tasks.values(), None) - # save checkpoint - next_config = await checkpointer.aput( - checkpoint_config, - create_checkpoint(checkpoint, None, step), - { - **checkpoint_metadata, - "source": "update", - "step": step + 1, - "writes": {}, - "parents": saved.metadata.get("parents", {}) if saved else {}, - }, - {}, - ) - return patch_checkpoint_map( - next_config, saved.metadata if saved else None - ) - # no values, empty checkpoint - if values is None and as_node is None: - next_checkpoint = create_checkpoint(checkpoint, None, step) - # copy checkpoint - next_config = await checkpointer.aput( - checkpoint_config, - next_checkpoint, - { - **checkpoint_metadata, - "source": "update", - "step": step + 1, - "writes": {}, - "parents": saved.metadata.get("parents", {}) if saved else {}, - }, - {}, - ) - return patch_checkpoint_map( - next_config, saved.metadata if saved else None - ) - # no values, copy checkpoint - if values is None and as_node == "__copy__": - next_checkpoint = create_checkpoint(checkpoint, None, step) - # copy checkpoint - next_config = await checkpointer.aput( - saved.parent_config or saved.config if saved else checkpoint_config, - next_checkpoint, - { - **checkpoint_metadata, - "source": "fork", - "step": step + 1, - "parents": saved.metadata.get("parents", {}) if saved else {}, - }, - {}, - ) - return patch_checkpoint_map( - next_config, saved.metadata if saved else None - ) - # apply pending writes, if not on specific checkpoint - if ( - CONFIG_KEY_CHECKPOINT_ID not in config[CONF] - and saved is not None - and saved.pending_writes - ): - # tasks for this checkpoint - next_tasks = prepare_next_tasks( - checkpoint, - saved.pending_writes, - self.nodes, - channels, - managed, - saved.config, - saved.metadata.get("step", -1) + 1, - for_execution=True, - store=self.store, - checkpointer=self.checkpointer - if isinstance(self.checkpointer, BaseCheckpointSaver) - else None, - manager=None, - ) - # apply null writes - if null_writes := [ - w[1:] for w in saved.pending_writes or [] if w[0] == NULL_TASK_ID - ]: - apply_writes( - saved.checkpoint, - channels, - [PregelTaskWrites((), INPUT, null_writes, [])], - None, - ) - for tid, k, v in saved.pending_writes: - if k in (ERROR, INTERRUPT, SCHEDULED): - continue - if tid not in next_tasks: - continue - next_tasks[tid].writes.append((k, v)) - if tasks := [t for t in next_tasks.values() if t.writes]: - apply_writes(checkpoint, channels, tasks, None) - # find last node that updated the state, if not provided - if as_node is None and not saved: - if ( - isinstance(self.input_channels, str) - and self.input_channels in self.nodes - ): - as_node = self.input_channels - elif as_node is None: - last_seen_by_node = sorted( - (v, n) - for n, seen in checkpoint["versions_seen"].items() - if n in self.nodes - for v in seen.values() - ) - # if two nodes updated the state at the same time, it's ambiguous - if last_seen_by_node: - if len(last_seen_by_node) == 1: - as_node = last_seen_by_node[0][1] - elif last_seen_by_node[-1][0] != last_seen_by_node[-2][0]: - as_node = last_seen_by_node[-1][1] - if as_node is None: - raise InvalidUpdateError("Ambiguous update, specify as_node") - if as_node not in self.nodes: - raise InvalidUpdateError(f"Node {as_node} does not exist") - # create task to run all writers of the chosen node - writers = self.nodes[as_node].flat_writers - if not writers: - raise InvalidUpdateError(f"Node {as_node} has no writers") - writes: deque[tuple[str, Any]] = deque() - task = PregelTaskWrites((), as_node, writes, [INTERRUPT]) - task_id = str(uuid5(UUID(checkpoint["id"]), INTERRUPT)) - run = RunnableSequence(*writers) if len(writers) > 1 else writers[0] - # execute task - await run.ainvoke( - values, - patch_config( - config, - run_name=self.name + "UpdateState", - configurable={ - # deque.extend is thread-safe - CONFIG_KEY_SEND: partial( - local_write, - writes.extend, - self.nodes.keys(), - ), - CONFIG_KEY_READ: partial( - local_read, - step + 1, - checkpoint, - channels, - managed, - task, - config, - ), - }, - ), - ) - # save task writes - # channel writes are saved to current checkpoint - # push writes are saved to next checkpoint - channel_writes, push_writes = ( - [w for w in task.writes if w[0] != PUSH], - [w for w in task.writes if w[0] == PUSH], - ) - if saved and channel_writes: - await checkpointer.aput_writes( - checkpoint_config, channel_writes, task_id - ) - # apply to checkpoint and save - mv_writes = apply_writes( - checkpoint, channels, [task], checkpointer.get_next_version - ) - assert not mv_writes, "Can't write to SharedValues from update_state" - checkpoint = create_checkpoint(checkpoint, channels, step + 1) - # save checkpoint, after applying writes - next_config = await checkpointer.aput( - checkpoint_config, - checkpoint, - { - **checkpoint_metadata, - "source": "update", - "step": step + 1, - "writes": {as_node: values}, - "parents": saved.metadata.get("parents", {}) if saved else {}, - }, - get_new_channel_versions( - checkpoint_previous_versions, checkpoint["channel_versions"] - ), - ) - # save push writes - if push_writes: - await checkpointer.aput_writes(next_config, push_writes, task_id) - return patch_checkpoint_map(next_config, saved.metadata if saved else None) - def _defaults( self, config: RunnableConfig, *, stream_mode: Optional[Union[StreamMode, list[StreamMode]]], + log_mode: Optional[Union[StreamMode, list[StreamMode]]], output_keys: Optional[Union[str, Sequence[str]]], interrupt_before: Optional[Union[All, Sequence[str]]], interrupt_after: Optional[Union[All, Sequence[str]]], - debug: Optional[bool], ) -> tuple[ - bool, + set[StreamMode], set[StreamMode], Union[str, Sequence[str]], Union[All, Sequence[str]], @@ -1704,7 +1115,6 @@ class Pregel(PregelProtocol): ]: if config["recursion_limit"] < 1: raise ValueError("recursion_limit must be at least 1") - debug = debug if debug is not None else self.debug if output_keys is None: output_keys = self.stream_channels_asis else: @@ -1717,6 +1127,11 @@ class Pregel(PregelProtocol): if CONFIG_KEY_TASK_ID in config.get(CONF, {}): # if being called as a node in another graph, always use values mode stream_mode = ["values"] + log_modes: set[StreamMode] = set() + if isinstance(log_mode, str): + log_modes.add(log_mode) + elif log_mode: + log_modes.update(log_mode) if self.checkpointer is False: checkpointer: Optional[BaseCheckpointSaver] = None elif CONFIG_KEY_CHECKPOINTER in config.get(CONF, {}): @@ -1727,15 +1142,15 @@ class Pregel(PregelProtocol): checkpointer = self.checkpointer if checkpointer and not config.get(CONF): raise ValueError( - f"Checkpointer requires one or more of the following 'configurable' keys: {[s.id for s in checkpointer.config_specs]}" + "Checkpointer requires one or more of the following 'configurable' fields: thread_id, checkpoint_id, checkpoint_ns" ) if CONFIG_KEY_STORE in config.get(CONF, {}): store: Optional[BaseStore] = config[CONF][CONFIG_KEY_STORE] else: store = self.store return ( - debug, set(stream_mode), + log_modes, output_keys, interrupt_before, interrupt_after, @@ -1746,13 +1161,13 @@ class Pregel(PregelProtocol): def stream( self, input: Union[dict[str, Any], Any], - config: Optional[RunnableConfig] = None, + config: Optional[AnyConfig] = None, *, stream_mode: Optional[Union[StreamMode, list[StreamMode]]] = None, + log_mode: Optional[Union[StreamMode, list[StreamMode]]] = None, output_keys: Optional[Union[str, Sequence[str]]] = None, interrupt_before: Optional[Union[All, Sequence[str]]] = None, interrupt_after: Optional[Union[All, Sequence[str]]] = None, - debug: Optional[bool] = None, subgraphs: bool = False, ) -> Iterator[Union[dict[str, Any], Any]]: """Stream graph steps for a single input. @@ -1768,12 +1183,11 @@ class Pregel(PregelProtocol): - `"updates"`: Emit only the node or task names and updates returned by the nodes or tasks after each step. If multiple updates are made in the same step (e.g. multiple nodes are run) then those updates are emitted separately. - `"custom"`: Emit custom data from inside nodes or tasks using `StreamWriter`. - - `"messages"`: Emit LLM messages token-by-token together with metadata for any LLM invocations inside nodes or tasks. - `"debug"`: Emit debug events with as much information as possible for each step. + log_modes: The stream modes to log, defaults to none, useful for debugging. output_keys: The keys to stream, defaults to all non-context channels. interrupt_before: Nodes to interrupt before, defaults to all nodes in the graph. interrupt_after: Nodes to interrupt after, defaults to all nodes in the graph. - debug: Whether to print debug information during execution, defaults to False. subgraphs: Whether to stream subgraphs, defaults to False. Yields: @@ -1843,38 +1257,6 @@ class Pregel(PregelProtocol): ... print(event) {'custom_data': 'foo'} ``` - - With stream_mode="messages": - - ```pycon - >>> from typing_extensions import Annotated, TypedDict - >>> from langgraph.graph import StateGraph, START - >>> from langchain_openai import ChatOpenAI - ... - >>> llm = ChatOpenAI(model="gpt-4o-mini") - ... - >>> class State(TypedDict): - ... question: str - ... answer: str - ... - >>> def node_a(state: State): - ... response = llm.invoke(state["question"]) - ... return {"answer": response.content} - ... - >>> builder = StateGraph(State) - >>> builder.add_node("a", node_a) - >>> builder.add_edge(START, "a") - >>> graph = builder.compile() - - >>> for event in graph.stream({"question": "What is the capital of France?"}, stream_mode="messages"): - ... print(event) - (AIMessageChunk(content='The', additional_kwargs={}, response_metadata={}, id='...'), {'langgraph_step': 1, 'langgraph_node': 'a', 'langgraph_triggers': ['start:a'], 'langgraph_path': ('__pregel_pull', 'a'), 'langgraph_checkpoint_ns': '...', 'checkpoint_ns': '...', 'ls_provider': 'openai', 'ls_model_name': 'gpt-4o-mini', 'ls_model_type': 'chat', 'ls_temperature': 0.7}) - (AIMessageChunk(content=' capital', additional_kwargs={}, response_metadata={}, id='...'), {'langgraph_step': 1, 'langgraph_node': 'a', 'langgraph_triggers': ['start:a'], ...}) - (AIMessageChunk(content=' of', additional_kwargs={}, response_metadata={}, id='...'), {...}) - (AIMessageChunk(content=' France', additional_kwargs={}, response_metadata={}, id='...'), {...}) - (AIMessageChunk(content=' is', additional_kwargs={}, response_metadata={}, id='...'), {...}) - (AIMessageChunk(content=' Paris', additional_kwargs={}, response_metadata={}, id='...'), {...}) - ``` """ stream = SyncQueue() @@ -1885,6 +1267,10 @@ class Pregel(PregelProtocol): ns, mode, payload = stream.get(block=False) except queue.Empty: break + if mode in log_modes: + print((ns, mode, payload)) + if mode not in stream_modes: + continue if subgraphs and isinstance(stream_mode, list): yield (ns, mode, payload) elif isinstance(stream_mode, list): @@ -1895,421 +1281,117 @@ class Pregel(PregelProtocol): yield payload config = ensure_config(self.config, config) - callback_manager = get_callback_manager_for_config(config) - run_manager = callback_manager.on_chain_start( - None, - input, - name=config.get("run_name", self.get_name()), - run_id=config.get("run_id"), + runtree = get_runtree_for_config( + config, input, name=config.get("run_name", self.get_name()) ) - try: - # assign defaults - ( - debug, - stream_modes, - output_keys, - interrupt_before_, - interrupt_after_, - checkpointer, - store, - ) = self._defaults( - config, - stream_mode=stream_mode, - output_keys=output_keys, - interrupt_before=interrupt_before, - interrupt_after=interrupt_after, - debug=debug, + config["run_tree"] = runtree + # assign defaults + ( + stream_modes, + log_modes, + output_keys, + interrupt_before_, + interrupt_after_, + checkpointer, + store, + ) = self._defaults( + config, + stream_mode=stream_mode, + log_mode=log_mode, + output_keys=output_keys, + interrupt_before=interrupt_before, + interrupt_after=interrupt_after, + ) + all_modes = stream_modes.union(log_modes) + # set up subgraph checkpointing + if self.checkpointer is True: + ns = cast(str, config[CONF][CONFIG_KEY_CHECKPOINT_NS]) + config[CONF][CONFIG_KEY_CHECKPOINT_NS] = recast_checkpoint_ns(ns) + # set up custom stream mode + if "custom" in all_modes: + config[CONF][CONFIG_KEY_STREAM_WRITER] = lambda c: stream.put( + ((), "custom", c) ) - # set up subgraph checkpointing - if self.checkpointer is True: - ns = cast(str, config[CONF][CONFIG_KEY_CHECKPOINT_NS]) - config[CONF][CONFIG_KEY_CHECKPOINT_NS] = recast_checkpoint_ns(ns) - # set up messages stream mode - if "messages" in stream_modes: - run_manager.inheritable_handlers.append( - StreamMessagesHandler(stream.put) - ) - # set up custom stream mode - if "custom" in stream_modes: - config[CONF][CONFIG_KEY_STREAM_WRITER] = lambda c: stream.put( - ((), "custom", c) - ) - with SyncPregelLoop( - input, - stream=StreamProtocol(stream.put, stream_modes), - config=config, - store=store, - checkpointer=checkpointer, - nodes=self.nodes, - specs=self.channels, - output_keys=output_keys, - stream_keys=self.stream_channels_asis, - interrupt_before=interrupt_before_, - interrupt_after=interrupt_after_, - manager=run_manager, - debug=debug, - ) as loop: - # create runner - runner = PregelRunner( - submit=config[CONF].get(CONFIG_KEY_RUNNER_SUBMIT, loop.submit), - put_writes=loop.put_writes, - schedule_task=loop.accept_push, - node_finished=config[CONF].get(CONFIG_KEY_NODE_FINISHED), - ) - # enable subgraph streaming - if subgraphs: - loop.config[CONF][CONFIG_KEY_STREAM] = loop.stream - # enable concurrent streaming - if ( - self.stream_eager - or subgraphs - or "messages" in stream_modes - or "custom" in stream_modes - ): - # we are careful to have a single waiter live at any one time - # because on exit we increment semaphore count by exactly 1 - waiter: Optional[concurrent.futures.Future] = None - # because sync futures cannot be cancelled, we instead - # release the stream semaphore on exit, which will cause - # a pending waiter to return immediately - loop.stack.callback(stream._count.release) - - def get_waiter() -> concurrent.futures.Future[None]: - nonlocal waiter - if waiter is None or waiter.done(): - waiter = loop.submit(stream.wait) - return waiter - else: - return waiter - - else: - get_waiter = None # type: ignore[assignment] - # Similarly to Bulk Synchronous Parallel / Pregel model - # computation proceeds in steps, while there are channel updates. - # Channel updates from step N are only visible in step N+1 - # channels are guaranteed to be immutable for the duration of the step, - # with channel updates applied only at the transition between steps. - while loop.tick(input_keys=self.input_channels): - for _ in runner.tick( - loop.tasks.values(), - timeout=self.step_timeout, - retry_policy=self.retry_policy, - get_waiter=get_waiter, - ): - # emit output - yield from output() - # emit output - yield from output() - # handle exit - if loop.status == "out_of_steps": - msg = create_error_message( - message=( - f"Recursion limit of {config['recursion_limit']} reached " - "without hitting a stop condition. You can increase the " - "limit by setting the `recursion_limit` config key." - ), - error_code=ErrorCode.GRAPH_RECURSION_LIMIT, - ) - raise GraphRecursionError(msg) - # set final channel values as run output - run_manager.on_chain_end(loop.output) - except BaseException as e: - run_manager.on_chain_error(e) - raise - - async def astream( - self, - input: Union[dict[str, Any], Any], - config: Optional[RunnableConfig] = None, - *, - stream_mode: Optional[Union[StreamMode, list[StreamMode]]] = None, - output_keys: Optional[Union[str, Sequence[str]]] = None, - interrupt_before: Optional[Union[All, Sequence[str]]] = None, - interrupt_after: Optional[Union[All, Sequence[str]]] = None, - debug: Optional[bool] = None, - subgraphs: bool = False, - ) -> AsyncIterator[Union[dict[str, Any], Any]]: - """Stream graph steps for a single input. - - Args: - input: The input to the graph. - config: The configuration to use for the run. - stream_mode: The mode to stream output, defaults to self.stream_mode. - Options are: - - - `"values"`: Emit all values in the state after each step. - When used with functional API, values are emitted once at the end of the workflow. - - `"updates"`: Emit only the node or task names and updates returned by the nodes or tasks after each step. - If multiple updates are made in the same step (e.g. multiple nodes are run) then those updates are emitted separately. - - `"custom"`: Emit custom data from inside nodes or tasks using `StreamWriter`. - - `"messages"`: Emit LLM messages token-by-token together with metadata for any LLM invocations inside nodes or tasks. - - `"debug"`: Emit debug events with as much information as possible for each step. - output_keys: The keys to stream, defaults to all non-context channels. - interrupt_before: Nodes to interrupt before, defaults to all nodes in the graph. - interrupt_after: Nodes to interrupt after, defaults to all nodes in the graph. - debug: Whether to print debug information during execution, defaults to False. - subgraphs: Whether to stream subgraphs, defaults to False. - - Yields: - The output of each step in the graph. The output shape depends on the stream_mode. - - Examples: - Using different stream modes with a graph: - ```pycon - >>> import operator - >>> from typing_extensions import Annotated, TypedDict - >>> from langgraph.graph import StateGraph, START - ... - >>> class State(TypedDict): - ... alist: Annotated[list, operator.add] - ... another_list: Annotated[list, operator.add] - ... - >>> builder = StateGraph(State) - >>> builder.add_node("a", lambda _state: {"another_list": ["hi"]}) - >>> builder.add_node("b", lambda _state: {"alist": ["there"]}) - >>> builder.add_edge("a", "b") - >>> builder.add_edge(START, "a") - >>> graph = builder.compile() - ``` - With stream_mode="values": - - ```pycon - >>> async for event in graph.astream({"alist": ['Ex for stream_mode="values"']}, stream_mode="values"): - ... print(event) - {'alist': ['Ex for stream_mode="values"'], 'another_list': []} - {'alist': ['Ex for stream_mode="values"'], 'another_list': ['hi']} - {'alist': ['Ex for stream_mode="values"', 'there'], 'another_list': ['hi']} - ``` - With stream_mode="updates": - - ```pycon - >>> async for event in graph.astream({"alist": ['Ex for stream_mode="updates"']}, stream_mode="updates"): - ... print(event) - {'a': {'another_list': ['hi']}} - {'b': {'alist': ['there']}} - ``` - With stream_mode="debug": - - ```pycon - >>> async for event in graph.astream({"alist": ['Ex for stream_mode="debug"']}, stream_mode="debug"): - ... print(event) - {'type': 'task', 'timestamp': '2024-06-23T...+00:00', 'step': 1, 'payload': {'id': '...', 'name': 'a', 'input': {'alist': ['Ex for stream_mode="debug"'], 'another_list': []}, 'triggers': ['start:a']}} - {'type': 'task_result', 'timestamp': '2024-06-23T...+00:00', 'step': 1, 'payload': {'id': '...', 'name': 'a', 'result': [('another_list', ['hi'])]}} - {'type': 'task', 'timestamp': '2024-06-23T...+00:00', 'step': 2, 'payload': {'id': '...', 'name': 'b', 'input': {'alist': ['Ex for stream_mode="debug"'], 'another_list': ['hi']}, 'triggers': ['a']}} - {'type': 'task_result', 'timestamp': '2024-06-23T...+00:00', 'step': 2, 'payload': {'id': '...', 'name': 'b', 'result': [('alist', ['there'])]}} - ``` - - With stream_mode="custom": - - ```pycon - >>> from langgraph.types import StreamWriter - ... - >>> async def node_a(state: State, writer: StreamWriter): - ... writer({"custom_data": "foo"}) - ... return {"alist": ["hi"]} - ... - >>> builder = StateGraph(State) - >>> builder.add_node("a", node_a) - >>> builder.add_edge(START, "a") - >>> graph = builder.compile() - ... - >>> async for event in graph.astream({"alist": ['Ex for stream_mode="custom"']}, stream_mode="custom"): - ... print(event) - {'custom_data': 'foo'} - ``` - - With stream_mode="messages": - - ```pycon - >>> from typing_extensions import Annotated, TypedDict - >>> from langgraph.graph import StateGraph, START - >>> from langchain_openai import ChatOpenAI - ... - >>> llm = ChatOpenAI(model="gpt-4o-mini") - ... - >>> class State(TypedDict): - ... question: str - ... answer: str - ... - >>> async def node_a(state: State): - ... response = await llm.ainvoke(state["question"]) - ... return {"answer": response.content} - ... - >>> builder = StateGraph(State) - >>> builder.add_node("a", node_a) - >>> builder.add_edge(START, "a") - >>> graph = builder.compile() - - >>> for event in graph.stream({"question": "What is the capital of France?"}, stream_mode="messages"): - ... print(event) - (AIMessageChunk(content='The', additional_kwargs={}, response_metadata={}, id='...'), {'langgraph_step': 1, 'langgraph_node': 'a', 'langgraph_triggers': ['start:a'], 'langgraph_path': ('__pregel_pull', 'a'), 'langgraph_checkpoint_ns': '...', 'checkpoint_ns': '...', 'ls_provider': 'openai', 'ls_model_name': 'gpt-4o-mini', 'ls_model_type': 'chat', 'ls_temperature': 0.7}) - (AIMessageChunk(content=' capital', additional_kwargs={}, response_metadata={}, id='...'), {'langgraph_step': 1, 'langgraph_node': 'a', 'langgraph_triggers': ['start:a'], ...}) - (AIMessageChunk(content=' of', additional_kwargs={}, response_metadata={}, id='...'), {...}) - (AIMessageChunk(content=' France', additional_kwargs={}, response_metadata={}, id='...'), {...}) - (AIMessageChunk(content=' is', additional_kwargs={}, response_metadata={}, id='...'), {...}) - (AIMessageChunk(content=' Paris', additional_kwargs={}, response_metadata={}, id='...'), {...}) - ``` - """ - - stream = AsyncQueue() - aioloop = asyncio.get_running_loop() - stream_put = cast( - Callable[[StreamChunk], None], - partial(aioloop.call_soon_threadsafe, stream.put_nowait), - ) - - def output() -> Iterator: - while True: - try: - ns, mode, payload = stream.get_nowait() - except asyncio.QueueEmpty: - break - if subgraphs and isinstance(stream_mode, list): - yield (ns, mode, payload) - elif isinstance(stream_mode, list): - yield (mode, payload) - elif subgraphs: - yield (ns, payload) - else: - yield payload - - config = ensure_config(self.config, config) - callback_manager = get_async_callback_manager_for_config(config) - run_manager = await callback_manager.on_chain_start( - None, + with SyncPregelLoop( input, - name=config.get("run_name", self.get_name()), - run_id=config.get("run_id"), - ) - # if running from astream_log() run each proc with streaming - do_stream = next( - ( - cast(_StreamingCallbackHandler, h) - for h in run_manager.handlers - if isinstance(h, _StreamingCallbackHandler) - ), - None, - ) - try: - # assign defaults - ( - debug, - stream_modes, - output_keys, - interrupt_before_, - interrupt_after_, - checkpointer, - store, - ) = self._defaults( - config, - stream_mode=stream_mode, - output_keys=output_keys, - interrupt_before=interrupt_before, - interrupt_after=interrupt_after, - debug=debug, + stream=StreamProtocol(stream.put, all_modes), + config=config, + store=store, + checkpointer=checkpointer, + nodes=self.nodes, + specs=self.channels, + output_keys=output_keys, + stream_keys=self.stream_channels_asis, + interrupt_before=interrupt_before_, + interrupt_after=interrupt_after_, + ) as loop: + # create runner + runner = PregelRunner( + submit=config[CONF].get(CONFIG_KEY_RUNNER_SUBMIT, loop.submit), + put_writes=loop.put_writes, + node_finished=config[CONF].get(CONFIG_KEY_NODE_FINISHED), ) - # set up subgraph checkpointing - if self.checkpointer is True: - ns = cast(str, config[CONF][CONFIG_KEY_CHECKPOINT_NS]) - config[CONF][CONFIG_KEY_CHECKPOINT_NS] = recast_checkpoint_ns(ns) - # set up messages stream mode - if "messages" in stream_modes: - run_manager.inheritable_handlers.append( - StreamMessagesHandler(stream_put) - ) - # set up custom stream mode - if "custom" in stream_modes: - config[CONF][CONFIG_KEY_STREAM_WRITER] = ( - lambda c: aioloop.call_soon_threadsafe( - stream.put_nowait, ((), "custom", c) - ) - ) - async with AsyncPregelLoop( - input, - stream=StreamProtocol(stream.put_nowait, stream_modes), - config=config, - store=store, - checkpointer=checkpointer, - nodes=self.nodes, - specs=self.channels, - output_keys=output_keys, - stream_keys=self.stream_channels_asis, - interrupt_before=interrupt_before_, - interrupt_after=interrupt_after_, - manager=run_manager, - debug=debug, - ) as loop: - # create runner - runner = PregelRunner( - submit=config[CONF].get(CONFIG_KEY_RUNNER_SUBMIT, loop.submit), - put_writes=loop.put_writes, - schedule_task=loop.accept_push, - use_astream=do_stream is not None, - node_finished=config[CONF].get(CONFIG_KEY_NODE_FINISHED), - ) - # enable subgraph streaming - if subgraphs: - loop.config[CONF][CONFIG_KEY_STREAM] = StreamProtocol( - stream_put, stream_modes - ) - # enable concurrent streaming - if ( - self.stream_eager - or subgraphs - or "messages" in stream_modes - or "custom" in stream_modes + # enable subgraph streaming + if subgraphs: + loop.config[CONF][CONFIG_KEY_STREAM] = loop.stream + # enable concurrent streaming + if self.stream_eager or subgraphs or "custom" in all_modes: + # we are careful to have a single waiter live at any one time + # because on exit we increment semaphore count by exactly 1 + waiter: Optional[concurrent.futures.Future] = None + # because sync futures cannot be cancelled, we instead + # release the stream semaphore on exit, which will cause + # a pending waiter to return immediately + loop.stack.callback(stream._count.release) + + def get_waiter() -> concurrent.futures.Future[None]: + nonlocal waiter + if waiter is None or waiter.done(): + waiter = loop.submit(stream.wait) + return waiter + else: + return waiter + + else: + get_waiter = None # type: ignore[assignment] + # Similarly to Bulk Synchronous Parallel / Pregel model + # computation proceeds in steps, while there are channel updates. + # Channel updates from step N are only visible in step N+1 + # channels are guaranteed to be immutable for the duration of the step, + # with channel updates applied only at the transition between steps. + while loop.tick(input_keys=self.input_channels): + for _ in runner.tick( + loop.tasks.values(), + timeout=self.step_timeout, + retry_policy=self.retry_policy, + get_waiter=get_waiter, ): - - def get_waiter() -> asyncio.Task[None]: - return aioloop.create_task(stream.wait()) - - else: - get_waiter = None # type: ignore[assignment] - # Similarly to Bulk Synchronous Parallel / Pregel model - # computation proceeds in steps, while there are channel updates - # channel updates from step N are only visible in step N+1 - # channels are guaranteed to be immutable for the duration of the step, - # with channel updates applied only at the transition between steps - while loop.tick(input_keys=self.input_channels): - async for _ in runner.atick( - loop.tasks.values(), - timeout=self.step_timeout, - retry_policy=self.retry_policy, - get_waiter=get_waiter, - ): - # emit output - for o in output(): - yield o - # emit output - for o in output(): - yield o - # handle exit - if loop.status == "out_of_steps": - msg = create_error_message( - message=( - f"Recursion limit of {config['recursion_limit']} reached " - "without hitting a stop condition. You can increase the " - "limit by setting the `recursion_limit` config key." - ), - error_code=ErrorCode.GRAPH_RECURSION_LIMIT, - ) - raise GraphRecursionError(msg) - # set final channel values as run output - await run_manager.on_chain_end(loop.output) - except BaseException as e: - await asyncio.shield(run_manager.on_chain_error(e)) - raise + # emit output + yield from output() + # emit output + yield from output() + # handle exit + if loop.status == "out_of_steps": + msg = create_error_message( + message=( + f"Recursion limit of {config['recursion_limit']} reached " + "without hitting a stop condition. You can increase the " + "limit by setting the `recursion_limit` config key." + ), + error_code=ErrorCode.GRAPH_RECURSION_LIMIT, + ) + raise GraphRecursionError(msg) def invoke( self, input: Union[dict[str, Any], Any], - config: Optional[RunnableConfig] = None, + config: Optional[AnyConfig] = None, *, stream_mode: StreamMode = "values", + log_mode: Optional[Union[StreamMode, list[StreamMode]]] = None, output_keys: Optional[Union[str, Sequence[str]]] = None, interrupt_before: Optional[Union[All, Sequence[str]]] = None, interrupt_after: Optional[Union[All, Sequence[str]]] = None, - debug: Optional[bool] = None, **kwargs: Any, ) -> Union[dict[str, Any], Any]: """Run the graph with a single input and config. @@ -2321,7 +1403,6 @@ class Pregel(PregelProtocol): output_keys: Optional. The output keys to retrieve from the graph run. interrupt_before: Optional. The nodes to interrupt the graph run before. interrupt_after: Optional. The nodes to interrupt the graph run after. - debug: Optional. Enable debug mode for the graph run. **kwargs: Additional keyword arguments to pass to the graph run. Returns: @@ -2336,64 +1417,11 @@ class Pregel(PregelProtocol): for chunk in self.stream( input, config, + log_mode=log_mode, stream_mode=stream_mode, output_keys=output_keys, interrupt_before=interrupt_before, interrupt_after=interrupt_after, - debug=debug, - **kwargs, - ): - if stream_mode == "values": - latest = chunk - else: - chunks.append(chunk) - if stream_mode == "values": - return latest - else: - return chunks - - async def ainvoke( - self, - input: Union[dict[str, Any], Any], - config: Optional[RunnableConfig] = None, - *, - stream_mode: StreamMode = "values", - output_keys: Optional[Union[str, Sequence[str]]] = None, - interrupt_before: Optional[Union[All, Sequence[str]]] = None, - interrupt_after: Optional[Union[All, Sequence[str]]] = None, - debug: Optional[bool] = None, - **kwargs: Any, - ) -> Union[dict[str, Any], Any]: - """Asynchronously invoke the graph on a single input. - - Args: - input: The input data for the computation. It can be a dictionary or any other type. - config: Optional. The configuration for the computation. - stream_mode: Optional. The stream mode for the computation. Default is "values". - output_keys: Optional. The output keys to include in the result. Default is None. - interrupt_before: Optional. The nodes to interrupt before. Default is None. - interrupt_after: Optional. The nodes to interrupt after. Default is None. - debug: Optional. Whether to enable debug mode. Default is None. - **kwargs: Additional keyword arguments. - - Returns: - The result of the computation. If stream_mode is "values", it returns the latest value. - If stream_mode is "chunks", it returns a list of chunks. - """ - - output_keys = output_keys if output_keys is not None else self.output_channels - if stream_mode == "values": - latest: Union[dict[str, Any], Any] = None - else: - chunks = [] - async for chunk in self.astream( - input, - config, - stream_mode=stream_mode, - output_keys=output_keys, - interrupt_before=interrupt_before, - interrupt_after=interrupt_after, - debug=debug, **kwargs, ): if stream_mode == "values": diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 03b2af6f4..4958d9b47 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -21,10 +21,6 @@ from typing import ( ) from uuid import UUID -from langchain_core.callbacks import Callbacks -from langchain_core.callbacks.manager import AsyncParentRunManager, ParentRunManager -from langchain_core.runnables.config import RunnableConfig - from langgraph.channels.base import BaseChannel from langgraph.checkpoint.base import ( BaseCheckpointSaver, @@ -63,8 +59,6 @@ from langgraph.constants import ( Send, ) from langgraph.errors import EmptyChannelError, InvalidUpdateError -from langgraph.managed.base import ManagedValueMapping -from langgraph.pregel.call import get_runnable_for_task from langgraph.pregel.io import read_channel, read_channels from langgraph.pregel.log import logger from langgraph.pregel.manager import ChannelsManager @@ -72,13 +66,11 @@ from langgraph.pregel.read import PregelNode from langgraph.store.base import BaseStore from langgraph.types import ( All, - LoopProtocol, PregelExecutableTask, PregelScratchpad, PregelTask, - RetryPolicy, ) -from langgraph.utils.config import merge_configs, patch_config +from langgraph.utils.config import AnyConfig, merge_configs, patch_config GetNextVersion = Callable[[Optional[V], BaseChannel], V] SUPPORTS_EXC_NOTES = sys.version_info >= (3, 11) @@ -111,28 +103,6 @@ class PregelTaskWrites(NamedTuple): triggers: Sequence[str] -class Call: - __slots__ = ("func", "input", "retry", "callbacks") - - func: Callable - input: Any - retry: Optional[RetryPolicy] - callbacks: Callbacks - - def __init__( - self, - func: Callable, - input: Any, - *, - retry: Optional[RetryPolicy], - callbacks: Callbacks, - ) -> None: - self.func = func - self.input = input - self.retry = retry - self.callbacks = callbacks - - def should_interrupt( checkpoint: Checkpoint, interrupt_nodes: Union[All, Sequence[str]], @@ -167,12 +137,9 @@ def should_interrupt( def local_read( - step: int, checkpoint: Checkpoint, channels: Mapping[str, BaseChannel], - managed: ManagedValueMapping, task: WritesProtocol, - config: RunnableConfig, select: Union[list[str], str], fresh: bool = False, ) -> Union[dict[str, Any], Any]: @@ -180,7 +147,6 @@ def local_read( Used by conditional edges to read a copy of the state with reflecting the writes from that node only.""" if isinstance(select, str): - managed_keys = [] for c, _ in task.writes: if c == select: updated = {c} @@ -188,22 +154,16 @@ def local_read( else: updated = set() else: - managed_keys = [k for k in select if k in managed] - select = [k for k in select if k not in managed] updated = set(select).intersection(c for c, _ in task.writes) if fresh and updated: with ChannelsManager( {k: v for k, v in channels.items() if k in updated}, checkpoint, - LoopProtocol(config=config, step=step, stop=step + 1), - skip_context=True, - ) as (local_channels, _): + ) as local_channels: apply_writes(copy_checkpoint(checkpoint), local_channels, [task], None) values = read_channels({**channels, **local_channels}, select) else: values = read_channels(channels, select) - if managed_keys: - values.update({k: managed[k]() for k in managed_keys}) return values @@ -233,10 +193,9 @@ def apply_writes( channels: Mapping[str, BaseChannel], tasks: Iterable[WritesProtocol], get_next_version: Optional[GetNextVersion], -) -> dict[str, list[Any]]: +) -> None: """Apply writes from a set of tasks (usually the tasks from a Pregel step) - to the checkpoint and channels, and return managed values writes to be applied - externally.""" + to the checkpoint and channels""" # sort tasks on path, to ensure deterministic order for update application # any path parts after the 3rd are ignored for sorting # (we use them for eg. task ids which aren't good for sorting) @@ -280,7 +239,6 @@ def apply_writes( # Group writes by channel pending_writes_by_channel: dict[str, list[Any]] = defaultdict(list) - pending_writes_by_managed: dict[str, list[Any]] = defaultdict(list) for task in tasks: for chan, val in task.writes: if chan in (NO_WRITES, PUSH, RESUME, INTERRUPT, RETURN, ERROR): @@ -289,8 +247,6 @@ def apply_writes( checkpoint["pending_sends"].append(val) elif chan in channels: pending_writes_by_channel[chan].append(val) - else: - pending_writes_by_managed[chan].append(val) # Find the highest version of all channels if checkpoint["channel_versions"]: @@ -319,9 +275,6 @@ def apply_writes( channels[chan], ) - # Return managed values writes to be applied externally - return pending_writes_by_managed - @overload def prepare_next_tasks( @@ -329,14 +282,12 @@ def prepare_next_tasks( pending_writes: list[PendingWrite], processes: Mapping[str, PregelNode], channels: Mapping[str, BaseChannel], - managed: ManagedValueMapping, - config: RunnableConfig, + config: AnyConfig, step: int, *, for_execution: Literal[False], store: Literal[None] = None, checkpointer: Literal[None] = None, - manager: Literal[None] = None, ) -> dict[str, PregelTask]: ... @@ -346,14 +297,12 @@ def prepare_next_tasks( pending_writes: list[PendingWrite], processes: Mapping[str, PregelNode], channels: Mapping[str, BaseChannel], - managed: ManagedValueMapping, - config: RunnableConfig, + config: AnyConfig, step: int, *, for_execution: Literal[True], store: Optional[BaseStore], checkpointer: Optional[BaseCheckpointSaver], - manager: Union[None, ParentRunManager, AsyncParentRunManager], ) -> dict[str, PregelExecutableTask]: ... @@ -362,14 +311,12 @@ def prepare_next_tasks( pending_writes: list[PendingWrite], processes: Mapping[str, PregelNode], channels: Mapping[str, BaseChannel], - managed: ManagedValueMapping, - config: RunnableConfig, + config: AnyConfig, step: int, *, for_execution: bool, store: Optional[BaseStore] = None, checkpointer: Optional[BaseCheckpointSaver] = None, - manager: Union[None, ParentRunManager, AsyncParentRunManager] = None, ) -> Union[dict[str, PregelTask], dict[str, PregelExecutableTask]]: """Prepare the set of tasks that will make up the next Pregel step. This is the union of all PUSH tasks (Sends) and PULL tasks (nodes triggered @@ -384,13 +331,11 @@ def prepare_next_tasks( pending_writes=pending_writes, processes=processes, channels=channels, - managed=managed, config=config, step=step, for_execution=for_execution, store=store, checkpointer=checkpointer, - manager=manager, ): tasks.append(task) # Check if any processes should be run in next step @@ -403,13 +348,11 @@ def prepare_next_tasks( pending_writes=pending_writes, processes=processes, channels=channels, - managed=managed, config=config, step=step, for_execution=for_execution, store=store, checkpointer=checkpointer, - manager=manager, ): tasks.append(task) return {t.id: t for t in tasks} @@ -423,13 +366,11 @@ def prepare_single_task( pending_writes: list[PendingWrite], processes: Mapping[str, PregelNode], channels: Mapping[str, BaseChannel], - managed: ManagedValueMapping, - config: RunnableConfig, + config: AnyConfig, step: int, for_execution: bool, store: Optional[BaseStore] = None, checkpointer: Optional[BaseCheckpointSaver] = None, - manager: Union[None, ParentRunManager, AsyncParentRunManager] = None, ) -> Union[None, PregelTask, PregelExecutableTask]: """Prepares a single task for the next Pregel step, given a task path, which uniquely identifies a PUSH or PULL task within the graph.""" @@ -437,90 +378,7 @@ def prepare_single_task( configurable = config.get(CONF, {}) parent_ns = configurable.get(CONFIG_KEY_CHECKPOINT_NS, "") - if task_path[0] == PUSH and isinstance(task_path[-1], Call): - # (PUSH, parent task path, idx of PUSH write, id of parent task, Call) - task_path_t = cast(tuple[str, tuple, int, str, Call], task_path) - call = task_path_t[-1] - proc_ = get_runnable_for_task(call.func) - name = proc_.name - if name is None: - raise ValueError("`call` functions must have a `__name__` attribute") - # create task id - triggers = [PUSH] - checkpoint_ns = f"{parent_ns}{NS_SEP}{name}" if parent_ns else name - task_id = _uuid5_str( - checkpoint_id, - checkpoint_ns, - str(step), - name, - PUSH, - task_path_str(task_path[1]), - str(task_path[2]), - ) - task_checkpoint_ns = f"{checkpoint_ns}:{task_id}" - metadata = { - "langgraph_step": step, - "langgraph_node": name, - "langgraph_triggers": triggers, - "langgraph_path": task_path[:3], - "langgraph_checkpoint_ns": task_checkpoint_ns, - } - if task_id_checksum is not None: - assert task_id == task_id_checksum, f"{task_id} != {task_id_checksum}" - if for_execution: - writes: deque[tuple[str, Any]] = deque() - return PregelExecutableTask( - name, - call.input, - proc_, - writes, - patch_config( - merge_configs(config, {"metadata": metadata}), - run_name=name, - callbacks=call.callbacks - or (manager.get_child(f"graph:step:{step}") if manager else None), - configurable={ - CONFIG_KEY_TASK_ID: task_id, - # deque.extend is thread-safe - CONFIG_KEY_SEND: partial( - local_write, - writes.extend, - processes.keys(), - ), - CONFIG_KEY_READ: partial( - local_read, - step, - checkpoint, - channels, - managed, - PregelTaskWrites(task_path[:3], name, writes, triggers), - config, - ), - CONFIG_KEY_STORE: (store or configurable.get(CONFIG_KEY_STORE)), - CONFIG_KEY_CHECKPOINTER: ( - checkpointer or configurable.get(CONFIG_KEY_CHECKPOINTER) - ), - CONFIG_KEY_CHECKPOINT_MAP: { - **configurable.get(CONFIG_KEY_CHECKPOINT_MAP, {}), - parent_ns: checkpoint["id"], - }, - CONFIG_KEY_CHECKPOINT_ID: None, - CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns, - CONFIG_KEY_SCRATCHPAD: _scratchpad( - pending_writes, - task_id, - ), - }, - ), - triggers, - call.retry, - None, - task_id, - task_path[:3], - ) - else: - return PregelTask(task_id, name, task_path[:3]) - elif task_path[0] == PUSH: + if task_path[0] == PUSH: if len(task_path) == 2: # SEND tasks, executed in superstep n+1 # (PUSH, idx of pending send) @@ -569,7 +427,7 @@ def prepare_single_task( if node := proc.node: if proc.metadata: metadata.update(proc.metadata) - writes = deque() + writes: deque[tuple[str, Any]] = deque() return PregelExecutableTask( packet.node, packet.arg, @@ -580,9 +438,6 @@ def prepare_single_task( config, {"metadata": metadata, "tags": proc.tags} ), run_name=packet.node, - callbacks=( - manager.get_child(f"graph:step:{step}") if manager else None - ), configurable={ CONFIG_KEY_TASK_ID: task_id, # deque.extend is thread-safe @@ -593,14 +448,11 @@ def prepare_single_task( ), CONFIG_KEY_READ: partial( local_read, - step, checkpoint, channels, - managed, PregelTaskWrites( task_path[:3], packet.node, writes, triggers ), - config, ), CONFIG_KEY_STORE: ( store or configurable.get(CONFIG_KEY_STORE) @@ -656,9 +508,7 @@ def prepare_single_task( > seen.get(chan, null_version) ): try: - val = next( - _proc_input(proc, managed, channels, for_execution=for_execution) - ) + val = next(_proc_input(proc, channels, for_execution=for_execution)) except StopIteration: return except Exception as exc: @@ -703,11 +553,6 @@ def prepare_single_task( config, {"metadata": metadata, "tags": proc.tags} ), run_name=name, - callbacks=( - manager.get_child(f"graph:step:{step}") - if manager - else None - ), configurable={ CONFIG_KEY_TASK_ID: task_id, # deque.extend is thread-safe @@ -718,14 +563,11 @@ def prepare_single_task( ), CONFIG_KEY_READ: partial( local_read, - step, checkpoint, channels, - managed, PregelTaskWrites( task_path[:3], name, writes, triggers ), - config, ), CONFIG_KEY_STORE: ( store or configurable.get(CONFIG_KEY_STORE) @@ -788,7 +630,6 @@ def _scratchpad( def _proc_input( proc: PregelNode, - managed: ManagedValueMapping, channels: Mapping[str, BaseChannel], *, for_execution: bool, @@ -807,8 +648,6 @@ def _proc_input( val[k] = read_channel(channels, chan, catch=False) except EmptyChannelError: continue - else: - val[k] = managed[k]() except EmptyChannelError: return elif isinstance(proc.channels, list): diff --git a/libs/langgraph/langgraph/pregel/call.py b/libs/langgraph/langgraph/pregel/call.py deleted file mode 100644 index 61a451335..000000000 --- a/libs/langgraph/langgraph/pregel/call.py +++ /dev/null @@ -1,233 +0,0 @@ -"""Utility to convert a user provided function into a Runnable with a ChannelWrite.""" - -import concurrent.futures -import functools -import inspect -import sys -import types -from typing import Any, Callable, Generator, Generic, Optional, TypeVar, cast - -from langchain_core.runnables import Runnable -from typing_extensions import ParamSpec - -from langgraph.constants import CONF, CONFIG_KEY_CALL, RETURN, TAG_HIDDEN -from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry -from langgraph.types import RetryPolicy -from langgraph.utils.config import get_config -from langgraph.utils.runnable import ( - RunnableCallable, - RunnableSeq, - is_async_callable, - run_in_executor, -) - -## -# Utilities borrowed from cloudpickle. -# https://github.com/cloudpipe/cloudpickle/blob/6220b0ce83ffee5e47e06770a1ee38ca9e47c850/cloudpickle/cloudpickle.py#L265 - - -def _getattribute(obj: Any, name: str) -> Any: - for subpath in name.split("."): - if subpath == "": - raise AttributeError( - "Can't get local attribute {!r} on {!r}".format(name, obj) - ) - try: - parent = obj - obj = getattr(obj, subpath) - except AttributeError: - raise AttributeError( - "Can't get attribute {!r} on {!r}".format(name, obj) - ) from None - return obj, parent - - -def _whichmodule(obj: Any, name: str) -> Optional[str]: - """Find the module an object belongs to. - - This function differs from ``pickle.whichmodule`` in two ways: - - it does not mangle the cases where obj's module is __main__ and obj was - not found in any module. - - Errors arising during module introspection are ignored, as those errors - are considered unwanted side effects. - """ - module_name = getattr(obj, "__module__", None) - - if module_name is not None: - return module_name - # Protect the iteration by using a copy of sys.modules against dynamic - # modules that trigger imports of other modules upon calls to getattr or - # other threads importing at the same time. - for module_name, module in sys.modules.copy().items(): - # Some modules such as coverage can inject non-module objects inside - # sys.modules - if ( - module_name == "__main__" - or module_name == "__mp_main__" - or module is None - or not isinstance(module, types.ModuleType) - ): - continue - try: - if _getattribute(module, name)[0] is obj: - return module_name - except Exception: - pass - return None - - -def _lookup_module_and_qualname( - obj: Any, name: Optional[str] = None -) -> Optional[tuple[types.ModuleType, str]]: - if name is None: - name = getattr(obj, "__qualname__", None) - if name is None: # pragma: no cover - # This used to be needed for Python 2.7 support but is probably not - # needed anymore. However we keep the __name__ introspection in case - # users of cloudpickle rely on this old behavior for unknown reasons. - name = getattr(obj, "__name__", None) - if name is None: - return None - - module_name = _whichmodule(obj, name) - - if module_name is None: - # In this case, obj.__module__ is None AND obj was not found in any - # imported module. obj is thus treated as dynamic. - return None - - if module_name == "__main__": - return None - - # Note: if module_name is in sys.modules, the corresponding module is - # assumed importable at unpickling time. See #357 - module = sys.modules.get(module_name, None) - if module is None: - # The main reason why obj's module would not be imported is that this - # module has been dynamically created, using for example - # types.ModuleType. The other possibility is that module was removed - # from sys.modules after obj was created/imported. But this case is not - # supported, as the standard pickle does not support it either. - return None - - try: - obj2, parent = _getattribute(module, name) - except AttributeError: - # obj was not found inside the module it points to - return None - if obj2 is not obj: - return None - return module, name - - -def _explode_args_trace_inputs( - sig: inspect.Signature, input: tuple[tuple[Any, ...], dict[str, Any]] -) -> dict[str, Any]: - args, kwargs = input - bound = sig.bind_partial(*args, **kwargs) - bound.apply_defaults() - arguments = dict(bound.arguments) - arguments.pop("self", None) - arguments.pop("cls", None) - for param_name, param in sig.parameters.items(): - if param.kind == inspect.Parameter.VAR_KEYWORD: - # Update with the **kwargs, and remove the original entry - # This is to help flatten out keyword arguments - if param_name in arguments: - arguments.update(arguments.pop(param_name)) - return arguments - - -def get_runnable_for_entrypoint(func: Callable[..., Any]) -> RunnableSeq: - key = (func, False) - if key in CACHE: - return CACHE[key] - else: - if is_async_callable(func): - run = RunnableCallable( - None, func, name=func.__name__, trace=False, recurse=False - ) - else: - afunc = functools.update_wrapper( - functools.partial(run_in_executor, None, func), func - ) - run = RunnableCallable( - func, - afunc, - name=func.__name__, - trace=False, - recurse=False, - ) - if not _lookup_module_and_qualname(func): - return run - return CACHE.setdefault(key, run) - - -def get_runnable_for_task(func: Callable[..., Any]) -> RunnableSeq: - key = (func, True) - if key in CACHE: - return CACHE[key] - else: - if hasattr(func, "__name__"): - name = func.__name__ - elif hasattr(func, "func"): - name = func.func.__name__ - elif hasattr(func, "__class__"): - name = func.__class__.__name__ - else: - name = str(func) - - if is_async_callable(func): - run = RunnableCallable( - None, - func, - explode_args=True, - name=name, - trace=False, - recurse=False, - ) - else: - run = RunnableCallable( - func, - functools.wraps(func)(functools.partial(run_in_executor, None, func)), - explode_args=True, - name=name, - trace=False, - recurse=False, - ) - seq = RunnableSeq( - run, - ChannelWrite([ChannelWriteEntry(RETURN)], tags=[TAG_HIDDEN]), - name=name, - trace_inputs=functools.partial( - _explode_args_trace_inputs, inspect.signature(func) - ), - ) - if not _lookup_module_and_qualname(func): - return seq - return CACHE.setdefault(key, seq) - - -CACHE: dict[tuple[Callable[..., Any], bool], Runnable] = {} - - -P = ParamSpec("P") -P1 = TypeVar("P1") -T = TypeVar("T") - - -class SyncAsyncFuture(Generic[T], concurrent.futures.Future[T]): - def __await__(self) -> Generator[T, None, T]: - yield cast(T, ...) - - -def call( - func: Callable[P, T], - *args: Any, - retry: Optional[RetryPolicy] = None, - **kwargs: Any, -) -> SyncAsyncFuture[T]: - config = get_config() - impl = config[CONF][CONFIG_KEY_CALL] - fut = impl(func, (args, kwargs), retry=retry, callbacks=config["callbacks"]) - return fut diff --git a/libs/langgraph/langgraph/pregel/debug.py b/libs/langgraph/langgraph/pregel/debug.py index 8429fd538..e3ff7a668 100644 --- a/libs/langgraph/langgraph/pregel/debug.py +++ b/libs/langgraph/langgraph/pregel/debug.py @@ -1,7 +1,5 @@ -from collections import defaultdict from dataclasses import asdict from datetime import datetime, timezone -from pprint import pformat from typing import ( Any, Iterable, @@ -14,12 +12,15 @@ from typing import ( ) from uuid import UUID -from langchain_core.runnables.config import RunnableConfig -from langchain_core.utils.input import get_bolded_text, get_colored_text from typing_extensions import TypedDict from langgraph.channels.base import BaseChannel -from langgraph.checkpoint.base import Checkpoint, CheckpointMetadata, PendingWrite +from langgraph.checkpoint.base import ( + Checkpoint, + CheckpointConfig, + CheckpointMetadata, + PendingWrite, +) from langgraph.constants import ( CONF, CONFIG_KEY_CHECKPOINT_NS, @@ -31,7 +32,7 @@ from langgraph.constants import ( ) from langgraph.pregel.io import read_channels from langgraph.types import PregelExecutableTask, PregelTask, StateSnapshot -from langgraph.utils.config import patch_checkpoint_map +from langgraph.utils.config import AnyConfig, RunnableConfig, patch_checkpoint_map class TaskPayload(TypedDict): @@ -140,14 +141,14 @@ def map_debug_task_results( def map_debug_checkpoint( step: int, - config: RunnableConfig, + config: AnyConfig, channels: Mapping[str, BaseChannel], stream_channels: Union[str, Sequence[str]], metadata: CheckpointMetadata, checkpoint: Checkpoint, tasks: Iterable[PregelExecutableTask], pending_writes: list[PendingWrite], - parent_config: Optional[RunnableConfig], + parent_config: Optional[CheckpointConfig], output_keys: Union[str, Sequence[str]], ) -> Iterator[DebugOutputCheckpoint]: """Produce "checkpoint" events for stream_mode=debug.""" @@ -210,52 +211,6 @@ def map_debug_checkpoint( } -def print_step_tasks(step: int, next_tasks: list[PregelExecutableTask]) -> None: - n_tasks = len(next_tasks) - print( - f"{get_colored_text(f'[{step}:tasks]', color='blue')} " - + get_bolded_text( - f"Starting {n_tasks} task{'s' if n_tasks != 1 else ''} for step {step}:\n" - ) - + "\n".join( - f"- {get_colored_text(task.name, 'green')} -> {pformat(task.input)}" - for task in next_tasks - ) - ) - - -def print_step_writes( - step: int, writes: Sequence[tuple[str, Any]], whitelist: Sequence[str] -) -> None: - by_channel: dict[str, list[Any]] = defaultdict(list) - for channel, value in writes: - if channel in whitelist: - by_channel[channel].append(value) - print( - f"{get_colored_text(f'[{step}:writes]', color='blue')} " - + get_bolded_text( - f"Finished step {step} with writes to {len(by_channel)} channel{'s' if len(by_channel) != 1 else ''}:\n" - ) - + "\n".join( - f"- {get_colored_text(name, 'yellow')} -> {', '.join(pformat(v) for v in vals)}" - for name, vals in by_channel.items() - ) - ) - - -def print_step_checkpoint( - metadata: CheckpointMetadata, - channels: Mapping[str, BaseChannel], - whitelist: Sequence[str], -) -> None: - step = metadata["step"] - print( - f"{get_colored_text(f'[{step}:checkpoint]', color='blue')} " - + get_bolded_text(f"State at the end of step {step}:\n") - + pformat(read_channels(channels, whitelist), depth=3) - ) - - def tasks_w_writes( tasks: Iterable[Union[PregelTask, PregelExecutableTask]], pending_writes: Optional[list[PendingWrite]], diff --git a/libs/langgraph/langgraph/pregel/executor.py b/libs/langgraph/langgraph/pregel/executor.py index 64bb6c90e..0f79558a0 100644 --- a/libs/langgraph/langgraph/pregel/executor.py +++ b/libs/langgraph/langgraph/pregel/executor.py @@ -1,27 +1,25 @@ -import asyncio import concurrent.futures import time from contextlib import ExitStack from contextvars import copy_context +from functools import partial from types import TracebackType from typing import ( - AsyncContextManager, - Awaitable, + Any, Callable, ContextManager, - Coroutine, + Iterable, + Iterator, Optional, Protocol, TypeVar, cast, ) -from langchain_core.runnables import RunnableConfig -from langchain_core.runnables.config import get_executor_for_config from typing_extensions import ParamSpec from langgraph.errors import GraphBubbleUp -from langgraph.utils.future import CONTEXT_NOT_SUPPORTED, run_coroutine_threadsafe +from langgraph.utils.config import RunnableConfig P = ParamSpec("P") T = TypeVar("T") @@ -40,6 +38,61 @@ class Submit(Protocol[P, T]): ) -> concurrent.futures.Future[T]: ... +class ContextThreadPoolExecutor(concurrent.futures.ThreadPoolExecutor): + """ThreadPoolExecutor that copies the context to the child thread.""" + + def submit( # type: ignore[override] + self, + func: Callable[P, T], + *args: P.args, + **kwargs: P.kwargs, + ) -> concurrent.futures.Future[T]: + """Submit a function to the executor. + + Args: + func (Callable[..., T]): The function to submit. + *args (Any): The positional arguments to the function. + **kwargs (Any): The keyword arguments to the function. + + Returns: + Future[T]: The future for the function. + """ + return super().submit( + cast(Callable[..., T], partial(copy_context().run, func, *args, **kwargs)) + ) + + def map( + self, + fn: Callable[..., T], + *iterables: Iterable[Any], + timeout: float | None = None, + chunksize: int = 1, + ) -> Iterator[T]: + """Map a function to multiple iterables. + + Args: + fn (Callable[..., T]): The function to map. + *iterables (Iterable[Any]): The iterables to map over. + timeout (float | None, optional): The timeout for the map. + Defaults to None. + chunksize (int, optional): The chunksize for the map. Defaults to 1. + + Returns: + Iterator[T]: The iterator for the mapped function. + """ + contexts = [copy_context() for _ in range(len(iterables[0]))] # type: ignore[arg-type] + + def _wrapped_fn(*args: Any) -> T: + return contexts.pop().run(fn, *args) + + return super().map( + _wrapped_fn, + *iterables, + timeout=timeout, + chunksize=chunksize, + ) + + class BackgroundExecutor(ContextManager): """A context manager that runs sync tasks in the background. Uses a thread pool executor to delegate tasks to separate threads. @@ -50,7 +103,9 @@ class BackgroundExecutor(ContextManager): def __init__(self, config: RunnableConfig) -> None: self.stack = ExitStack() - self.executor = self.stack.enter_context(get_executor_for_config(config)) + self.executor = self.stack.enter_context( + ContextThreadPoolExecutor(max_workers=config.get("max_concurrency")) + ) # mapping of Future to (__cancel_on_exit__, __reraise_on_exit__) flags self.tasks: dict[concurrent.futures.Future, tuple[bool, bool]] = {} @@ -122,98 +177,6 @@ class BackgroundExecutor(ContextManager): pass -class AsyncBackgroundExecutor(AsyncContextManager): - """A context manager that runs async tasks in the background. - Uses the current event loop to delegate tasks to asyncio tasks. - On exit, - - cancels any tasks with `__cancel_on_exit__=True` - - waits for all tasks to finish - - re-raises the first exception from tasks with `__reraise_on_exit__=True` - ignoring CancelledError""" - - def __init__(self, config: RunnableConfig) -> None: - self.tasks: dict[asyncio.Future, tuple[bool, bool]] = {} - self.sentinel = object() - self.loop = asyncio.get_running_loop() - if max_concurrency := config.get("max_concurrency"): - self.semaphore: Optional[asyncio.Semaphore] = asyncio.Semaphore( - max_concurrency - ) - else: - self.semaphore = None - - def submit( # type: ignore[valid-type] - self, - fn: Callable[P, Awaitable[T]], - *args: P.args, - __name__: Optional[str] = None, - __cancel_on_exit__: bool = False, - __reraise_on_exit__: bool = True, - __next_tick__: bool = False, # noop in async (always True) - **kwargs: P.kwargs, - ) -> asyncio.Future[T]: - coro = cast(Coroutine[None, None, T], fn(*args, **kwargs)) - if self.semaphore: - coro = gated(self.semaphore, coro) - if CONTEXT_NOT_SUPPORTED: - task = run_coroutine_threadsafe(coro, self.loop, name=__name__) - else: - task = run_coroutine_threadsafe( - coro, self.loop, name=__name__, context=copy_context() - ) - self.tasks[task] = (__cancel_on_exit__, __reraise_on_exit__) - task.add_done_callback(self.done) - return task - - def done(self, task: asyncio.Future) -> None: - try: - if exc := task.exception(): - # This exception is an interruption signal, not an error - # so we don't want to re-raise it on exit - if isinstance(exc, GraphBubbleUp): - self.tasks.pop(task) - else: - self.tasks.pop(task) - except asyncio.CancelledError: - self.tasks.pop(task) - - async def __aenter__(self) -> Submit: - return self.submit - - async def __aexit__( - self, - exc_type: Optional[type[BaseException]], - exc_value: Optional[BaseException], - traceback: Optional[TracebackType], - ) -> None: - # copy the tasks as done() callback may modify the dict - tasks = self.tasks.copy() - # cancel all tasks that should be cancelled - for task, (cancel, _) in tasks.items(): - if cancel: - task.cancel(self.sentinel) - # wait for all tasks to finish - if tasks: - await asyncio.wait(tasks) - # if there's already an exception being raised, don't raise another one - if exc_type is None: - # re-raise the first exception that occurred in a task - for task, (_, reraise) in tasks.items(): - if not reraise: - continue - try: - if exc := task.exception(): - raise exc - except asyncio.CancelledError: - pass - - -async def gated(semaphore: asyncio.Semaphore, coro: Coroutine[None, None, T]) -> T: - """A coroutine that waits for a semaphore before running another coroutine.""" - async with semaphore: - return await coro - - def next_tick(fn: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> T: """A function that yields control to other threads before running another function.""" time.sleep(0) diff --git a/libs/langgraph/langgraph/pregel/io.py b/libs/langgraph/langgraph/pregel/io.py index e9963f5c8..8b9148cc1 100644 --- a/libs/langgraph/langgraph/pregel/io.py +++ b/libs/langgraph/langgraph/pregel/io.py @@ -2,8 +2,6 @@ from collections import Counter from typing import Any, Iterator, Literal, Mapping, Optional, Sequence, TypeVar, Union from uuid import UUID -from langchain_core.runnables.utils import AddableDict - from langgraph.channels.base import BaseChannel, EmptyChannelError from langgraph.checkpoint.base import PendingWrite from langgraph.constants import ( @@ -123,14 +121,6 @@ def map_input( logger.warning(f"Input channel {k} not found in {input_channels}") -class AddableValuesDict(AddableDict): - def __add__(self, other: dict[str, Any]) -> "AddableValuesDict": - return self | other - - def __radd__(self, other: dict[str, Any]) -> "AddableValuesDict": - return other | self - - def map_output_values( output_channels: Union[str, Sequence[str]], pending_writes: Union[Literal[True], Sequence[tuple[str, Any]]], @@ -146,15 +136,7 @@ def map_output_values( if pending_writes is True or { c for c, _ in pending_writes if c in output_channels }: - yield AddableValuesDict(read_channels(channels, output_channels)) - - -class AddableUpdatesDict(AddableDict): - def __add__(self, other: dict[str, Any]) -> "AddableUpdatesDict": - return [self, other] - - def __radd__(self, other: dict[str, Any]) -> "AddableUpdatesDict": - raise TypeError("AddableUpdatesDict does not support right-side addition") + yield read_channels(channels, output_channels) def map_output_updates( @@ -213,7 +195,7 @@ def map_output_updates( grouped[node] = value[0] if cached: grouped["__metadata__"] = {"cached": cached} # type: ignore[assignment] - yield AddableUpdatesDict(grouped) + yield grouped T = TypeVar("T") diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 278994126..0ec50a7f4 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -1,13 +1,11 @@ -import asyncio import concurrent.futures from collections import defaultdict, deque -from contextlib import AsyncExitStack, ExitStack +from contextlib import ExitStack from dataclasses import replace from inspect import signature from types import TracebackType from typing import ( Any, - AsyncContextManager, Callable, ContextManager, Iterator, @@ -22,8 +20,6 @@ from typing import ( cast, ) -from langchain_core.callbacks import AsyncParentRunManager, ParentRunManager -from langchain_core.runnables import RunnableConfig from typing_extensions import ParamSpec, Self from langgraph.channels.base import BaseChannel @@ -57,7 +53,6 @@ from langgraph.constants import ( INTERRUPT, NS_SEP, NULL_TASK_ID, - PUSH, RESUME, SCHEDULED, TAG_HIDDEN, @@ -69,19 +64,12 @@ from langgraph.errors import ( GraphInterrupt, ParentCommand, ) -from langgraph.managed.base import ( - ManagedValueMapping, - ManagedValueSpec, - WritableManagedValue, -) from langgraph.pregel.algo import ( - Call, GetNextVersion, PregelTaskWrites, apply_writes, increment, prepare_next_tasks, - prepare_single_task, should_interrupt, task_path_str, ) @@ -89,12 +77,8 @@ from langgraph.pregel.debug import ( map_debug_checkpoint, map_debug_task_results, map_debug_tasks, - print_step_checkpoint, - print_step_tasks, - print_step_writes, ) from langgraph.pregel.executor import ( - AsyncBackgroundExecutor, BackgroundExecutor, Submit, ) @@ -106,7 +90,7 @@ from langgraph.pregel.io import ( read_channels, single, ) -from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager +from langgraph.pregel.manager import ChannelsManager from langgraph.pregel.read import PregelNode from langgraph.pregel.utils import get_new_channel_versions from langgraph.store.base import BaseStore @@ -119,7 +103,7 @@ from langgraph.types import ( StreamChunk, StreamProtocol, ) -from langgraph.utils.config import patch_configurable +from langgraph.utils.config import AnyConfig, RunnableConfig, patch_configurable V = TypeVar("V") P = ParamSpec("P") @@ -142,12 +126,11 @@ class PregelLoop(LoopProtocol): input: Optional[Any] checkpointer: Optional[BaseCheckpointSaver] nodes: Mapping[str, PregelNode] - specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]] + specs: Mapping[str, BaseChannel] output_keys: Union[str, Sequence[str]] stream_keys: Union[str, Sequence[str]] skip_done_tasks: bool is_nested: bool - manager: Union[None, AsyncParentRunManager, ParentRunManager] interrupt_after: Union[All, Sequence[str]] interrupt_before: Union[All, Sequence[str]] @@ -170,14 +153,13 @@ class PregelLoop(LoopProtocol): ] submit: Submit channels: Mapping[str, BaseChannel] - managed: ManagedValueMapping checkpoint: Checkpoint checkpoint_ns: tuple[str, ...] checkpoint_config: RunnableConfig checkpoint_metadata: CheckpointMetadata checkpoint_pending_writes: List[PendingWrite] checkpoint_previous_versions: dict[str, Union[str, float, int]] - prev_checkpoint_config: Optional[RunnableConfig] + prev_checkpoint_config: Optional[AnyConfig] status: Literal[ "pending", "done", "interrupt_before", "interrupt_after", "out_of_steps" @@ -197,13 +179,11 @@ class PregelLoop(LoopProtocol): store: Optional[BaseStore], checkpointer: Optional[BaseCheckpointSaver], nodes: Mapping[str, PregelNode], - specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]], + specs: Mapping[str, BaseChannel], output_keys: Union[str, Sequence[str]], stream_keys: Union[str, Sequence[str]], interrupt_after: Union[All, Sequence[str]] = EMPTY_SEQ, interrupt_before: Union[All, Sequence[str]] = EMPTY_SEQ, - manager: Union[None, AsyncParentRunManager, ParentRunManager] = None, - debug: bool = False, ) -> None: super().__init__( step=0, @@ -220,13 +200,11 @@ class PregelLoop(LoopProtocol): self.stream_keys = stream_keys self.interrupt_after = interrupt_after self.interrupt_before = interrupt_before - self.manager = manager self.is_nested = CONFIG_KEY_TASK_ID in self.config.get(CONF, {}) self.skip_done_tasks = ( CONFIG_KEY_CHECKPOINT_ID not in config[CONF] or CONFIG_KEY_DEDUPE_TASKS in config[CONF] ) - self.debug = debug if self.stream is not None and CONFIG_KEY_STREAM in config[CONF]: self.stream = DuplexStream(self.stream, config[CONF][CONFIG_KEY_STREAM]) scratchpad: Optional[PregelScratchpad] = config[CONF].get(CONFIG_KEY_SCRATCHPAD) @@ -333,53 +311,6 @@ class PregelLoop(LoopProtocol): if hasattr(self, "tasks"): self._output_writes(task_id, writes) - def accept_push( - self, task: PregelExecutableTask, write_idx: int, call: Optional[Call] = None - ) -> Optional[PregelExecutableTask]: - """Accept a PUSH from a task, potentially returning a new task to start.""" - # don't start if we should interrupt *after* the original task - if self.interrupt_after and should_interrupt( - self.checkpoint, self.interrupt_after, [task] - ): - self.to_interrupt.append(task) - return - if pushed := cast( - Optional[PregelExecutableTask], - prepare_single_task( - (PUSH, task.path, write_idx, task.id, call), - None, - checkpoint=self.checkpoint, - pending_writes=self.checkpoint_pending_writes, - processes=self.nodes, - channels=self.channels, - managed=self.managed, - config=task.config, - step=self.step, - for_execution=True, - store=self.store, - checkpointer=self.checkpointer, - manager=self.manager, - ), - ): - # don't start if we should interrupt *before* the new task - if self.interrupt_before and should_interrupt( - self.checkpoint, self.interrupt_before, [pushed] - ): - self.to_interrupt.append(pushed) - return - # produce debug output - self._emit("debug", map_debug_tasks, self.step, [pushed]) - # debug flag - if self.debug: - print_step_tasks(self.step, [pushed]) - # save the new task - self.tasks[pushed.id] = pushed - # match any pending writes to the new task - if self.skip_done_tasks: - self._match_writes({pushed.id: pushed}) - # return the new task, to be started if not run before - return pushed - def tick( self, *, @@ -405,27 +336,13 @@ class PregelLoop(LoopProtocol): elif all(task.writes for task in self.tasks.values()): # finish superstep writes = [w for t in self.tasks.values() for w in t.writes] - # debug flag - if self.debug: - print_step_writes( - self.step, - writes, - ( - [self.stream_keys] - if isinstance(self.stream_keys, str) - else self.stream_keys - ), - ) # all tasks have finished - mv_writes = apply_writes( + apply_writes( self.checkpoint, self.channels, self.tasks.values(), self.checkpointer_get_next_version, ) - # apply writes to managed values - for key, values in mv_writes.items(): - self._update_mv(key, values) # produce values output self._emit( "values", map_output_values, self.output_keys, writes, self.channels @@ -469,11 +386,9 @@ class PregelLoop(LoopProtocol): self.checkpoint_pending_writes, self.nodes, self.channels, - self.managed, self.config, self.step, for_execution=True, - manager=self.manager, store=self.store, checkpointer=self.checkpointer, ) @@ -531,10 +446,6 @@ class PregelLoop(LoopProtocol): # produce debug output self._emit("debug", map_debug_tasks, self.step, self.tasks.values()) - # debug flag - if self.debug: - print_step_tasks(self.step, list(self.tasks.values())) - # print output for any tasks we applied previous writes to for task in self.tasks.values(): if task.writes: @@ -598,14 +509,12 @@ class PregelLoop(LoopProtocol): if null_writes := [ w[1:] for w in self.checkpoint_pending_writes if w[0] == NULL_TASK_ID ]: - mv_writes = apply_writes( + apply_writes( self.checkpoint, self.channels, [PregelTaskWrites((), INPUT, null_writes, [])], self.checkpointer_get_next_version, ) - for key, values in mv_writes.items(): - self._update_mv(key, values) # proceed past previous checkpoint if is_resuming: self.checkpoint["versions_seen"].setdefault(INTERRUPT, {}) @@ -636,16 +545,14 @@ class PregelLoop(LoopProtocol): self.checkpoint_pending_writes, self.nodes, self.channels, - self.managed, self.config, self.step, for_execution=True, store=None, checkpointer=None, - manager=None, ) # apply input writes - mv_writes = apply_writes( + apply_writes( self.checkpoint, self.channels, [ @@ -654,7 +561,6 @@ class PregelLoop(LoopProtocol): ], self.checkpointer_get_next_version, ) - assert not mv_writes, "Can't write to SharedValues in graph input" # save input checkpoint self._put_checkpoint({"source": "input", "writes": dict(input_writes)}) elif CONFIG_KEY_RESUMING not in configurable: @@ -673,17 +579,6 @@ class PregelLoop(LoopProtocol): # assign step and parents metadata["step"] = self.step metadata["parents"] = self.config[CONF].get(CONFIG_KEY_CHECKPOINT_MAP, {}) - # debug flag - if self.debug: - print_step_checkpoint( - metadata, - self.channels, - ( - [self.stream_keys] - if isinstance(self.stream_keys, str) - else self.stream_keys - ), - ) # create new checkpoint self.checkpoint = create_checkpoint(self.checkpoint, self.channels, self.step) # bail if no checkpointer @@ -733,9 +628,6 @@ class PregelLoop(LoopProtocol): # increment step self.step += 1 - def _update_mv(self, key: str, values: Sequence[Any]) -> None: - raise NotImplementedError - def _suppress_interrupt( self, exc_type: Optional[Type[BaseException]], @@ -760,14 +652,12 @@ class PregelLoop(LoopProtocol): and self.checkpoint_pending_writes and any(task.writes for task in self.tasks.values()) ): - mv_writes = apply_writes( + apply_writes( self.checkpoint, self.channels, self.tasks.values(), self.checkpointer_get_next_version, ) - for key, values in mv_writes.items(): - self._update_mv(key, values) self._emit( "values", map_output_values, @@ -838,13 +728,11 @@ class SyncPregelLoop(PregelLoop, ContextManager): store: Optional[BaseStore], checkpointer: Optional[BaseCheckpointSaver], nodes: Mapping[str, PregelNode], - specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]], - manager: Union[None, AsyncParentRunManager, ParentRunManager] = None, + specs: Mapping[str, BaseChannel], interrupt_after: Union[All, Sequence[str]] = EMPTY_SEQ, interrupt_before: Union[All, Sequence[str]] = EMPTY_SEQ, output_keys: Union[str, Sequence[str]] = EMPTY_SEQ, stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ, - debug: bool = False, ) -> None: super().__init__( input, @@ -858,8 +746,6 @@ class SyncPregelLoop(PregelLoop, ContextManager): stream_keys=stream_keys, interrupt_after=interrupt_after, interrupt_before=interrupt_before, - manager=manager, - debug=debug, ) self.stack = ExitStack() if checkpointer: @@ -891,13 +777,6 @@ class SyncPregelLoop(PregelLoop, ContextManager): config, checkpoint, metadata, new_versions ) - def _update_mv(self, key: str, values: Sequence[Any]) -> None: - managed_value = self.managed.get(key) - if managed_value is None: - return - - return self.submit(cast(WritableManagedValue, managed_value).update, values) - # context manager def __enter__(self) -> Self: @@ -946,8 +825,8 @@ class SyncPregelLoop(PregelLoop, ContextManager): ) self.submit = self.stack.enter_context(BackgroundExecutor(self.config)) - self.channels, self.managed = self.stack.enter_context( - ChannelsManager(self.specs, self.checkpoint, self) + self.channels = self.stack.enter_context( + ChannelsManager(self.specs, self.checkpoint) ) self.stack.push(self._suppress_interrupt) self.status = "pending" @@ -965,154 +844,3 @@ class SyncPregelLoop(PregelLoop, ContextManager): ) -> Optional[bool]: # unwind stack return self.stack.__exit__(exc_type, exc_value, traceback) - - -class AsyncPregelLoop(PregelLoop, AsyncContextManager): - def __init__( - self, - input: Optional[Any], - *, - stream: Optional[StreamProtocol], - config: RunnableConfig, - store: Optional[BaseStore], - checkpointer: Optional[BaseCheckpointSaver], - nodes: Mapping[str, PregelNode], - specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]], - interrupt_after: Union[All, Sequence[str]] = EMPTY_SEQ, - interrupt_before: Union[All, Sequence[str]] = EMPTY_SEQ, - manager: Union[None, AsyncParentRunManager, ParentRunManager] = None, - output_keys: Union[str, Sequence[str]] = EMPTY_SEQ, - stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ, - debug: bool = False, - ) -> None: - super().__init__( - input, - stream=stream, - config=config, - checkpointer=checkpointer, - store=store, - nodes=nodes, - specs=specs, - output_keys=output_keys, - stream_keys=stream_keys, - interrupt_after=interrupt_after, - interrupt_before=interrupt_before, - manager=manager, - debug=debug, - ) - self.stack = AsyncExitStack() - if checkpointer: - self.checkpointer_get_next_version = checkpointer.get_next_version - self.checkpointer_put_writes = checkpointer.aput_writes - self.checkpointer_put_writes_accepts_task_path = ( - signature(checkpointer.aput_writes).parameters.get("task_path") - is not None - ) - else: - self.checkpointer_get_next_version = increment - self._checkpointer_put_after_previous = None # type: ignore[assignment] - self.checkpointer_put_writes = None - self.checkpointer_put_writes_accepts_task_path = False - - async def _checkpointer_put_after_previous( - self, - prev: Optional[asyncio.Task], - config: RunnableConfig, - checkpoint: Checkpoint, - metadata: CheckpointMetadata, - new_versions: ChannelVersions, - ) -> RunnableConfig: - try: - if prev is not None: - await prev - finally: - await cast(BaseCheckpointSaver, self.checkpointer).aput( - config, checkpoint, metadata, new_versions - ) - - def _update_mv(self, key: str, values: Sequence[Any]) -> None: - managed_value = self.managed.get(key) - if managed_value is None: - return - - return self.submit(cast(WritableManagedValue, managed_value).aupdate, values) - - # context manager - - async def __aenter__(self) -> Self: - if self.config.get(CONF, {}).get( - CONFIG_KEY_ENSURE_LATEST - ) and self.checkpoint_config[CONF].get(CONFIG_KEY_CHECKPOINT_ID): - if self.checkpointer is None: - raise RuntimeError( - "Cannot ensure latest checkpoint without checkpointer" - ) - saved = await self.checkpointer.aget_tuple( - patch_configurable( - self.checkpoint_config, {CONFIG_KEY_CHECKPOINT_ID: None} - ) - ) - if ( - saved is None - or saved.checkpoint["id"] - != self.checkpoint_config[CONF][CONFIG_KEY_CHECKPOINT_ID] - ): - raise CheckpointNotLatest - elif self.checkpointer: - saved = await self.checkpointer.aget_tuple(self.checkpoint_config) - else: - saved = None - if saved is None: - saved = CheckpointTuple( - self.config, empty_checkpoint(), {"step": -2}, None, [] - ) - self.checkpoint_config = { - **self.config, - **saved.config, - CONF: { - CONFIG_KEY_CHECKPOINT_NS: "", - **self.config.get(CONF, {}), - **saved.config.get(CONF, {}), - }, - } - self.prev_checkpoint_config = saved.parent_config - self.checkpoint = saved.checkpoint - self.checkpoint_metadata = saved.metadata - self.checkpoint_pending_writes = ( - [(str(tid), k, v) for tid, k, v in saved.pending_writes] - if saved.pending_writes is not None - else [] - ) - - self.submit = await self.stack.enter_async_context( - AsyncBackgroundExecutor(self.config) - ) - self.channels, self.managed = await self.stack.enter_async_context( - AsyncChannelsManager(self.specs, self.checkpoint, self) - ) - self.stack.push(self._suppress_interrupt) - self.status = "pending" - self.step = self.checkpoint_metadata["step"] + 1 - self.stop = self.step + self.config["recursion_limit"] + 1 - - self.checkpoint_previous_versions = self.checkpoint["channel_versions"].copy() - - return self - - async def __aexit__( - self, - exc_type: Optional[Type[BaseException]], - exc_value: Optional[BaseException], - traceback: Optional[TracebackType], - ) -> Optional[bool]: - # unwind stack - exit_task = asyncio.create_task( - self.stack.__aexit__(exc_type, exc_value, traceback) - ) - try: - return await exit_task - except asyncio.CancelledError as e: - # Bubble up the exit task upon cancellation to permit the API - # consumer to await it before e.g., reusing the DB connection. - e.args = (*e.args, exit_task) - raise diff --git a/libs/langgraph/langgraph/pregel/manager.py b/libs/langgraph/langgraph/pregel/manager.py index 641e1d8fe..af46cf7c1 100644 --- a/libs/langgraph/langgraph/pregel/manager.py +++ b/libs/langgraph/langgraph/pregel/manager.py @@ -1,103 +1,20 @@ -import asyncio -from contextlib import AsyncExitStack, ExitStack, asynccontextmanager, contextmanager -from typing import AsyncIterator, Iterator, Mapping, Union +from contextlib import contextmanager +from typing import Iterator, Mapping from langgraph.channels.base import BaseChannel from langgraph.checkpoint.base import Checkpoint -from langgraph.managed.base import ( - ConfiguredManagedValue, - ManagedValueMapping, - ManagedValueSpec, -) -from langgraph.managed.context import Context -from langgraph.types import LoopProtocol @contextmanager def ChannelsManager( - specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]], + specs: Mapping[str, BaseChannel], checkpoint: Checkpoint, - loop: LoopProtocol, - *, - skip_context: bool = False, -) -> Iterator[tuple[Mapping[str, BaseChannel], ManagedValueMapping]]: +) -> Iterator[Mapping[str, BaseChannel]]: """Manage channels for the lifetime of a Pregel invocation (multiple steps).""" channel_specs: dict[str, BaseChannel] = {} - managed_specs: dict[str, ManagedValueSpec] = {} for k, v in specs.items(): - if isinstance(v, BaseChannel): - channel_specs[k] = v - elif ( - skip_context and isinstance(v, ConfiguredManagedValue) and v.cls is Context - ): - managed_specs[k] = Context.of(noop_context) - else: - managed_specs[k] = v - with ExitStack() as stack: - yield ( - { - k: v.from_checkpoint(checkpoint["channel_values"].get(k)) - for k, v in channel_specs.items() - }, - ManagedValueMapping( - { - key: stack.enter_context( - value.cls.enter(loop, **value.kwargs) - if isinstance(value, ConfiguredManagedValue) - else value.enter(loop) - ) - for key, value in managed_specs.items() - } - ), - ) - - -@asynccontextmanager -async def AsyncChannelsManager( - specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]], - checkpoint: Checkpoint, - loop: LoopProtocol, - *, - skip_context: bool = False, -) -> AsyncIterator[tuple[Mapping[str, BaseChannel], ManagedValueMapping]]: - """Manage channels for the lifetime of a Pregel invocation (multiple steps).""" - channel_specs: dict[str, BaseChannel] = {} - managed_specs: dict[str, ManagedValueSpec] = {} - for k, v in specs.items(): - if isinstance(v, BaseChannel): - channel_specs[k] = v - elif ( - skip_context and isinstance(v, ConfiguredManagedValue) and v.cls is Context - ): - managed_specs[k] = Context.of(noop_context) - else: - managed_specs[k] = v - async with AsyncExitStack() as stack: - # managed: create enter tasks with reference to spec, await them - if tasks := { - asyncio.create_task( - stack.enter_async_context( - value.cls.aenter(loop, **value.kwargs) - if isinstance(value, ConfiguredManagedValue) - else value.aenter(loop) - ) - ): key - for key, value in managed_specs.items() - }: - done, _ = await asyncio.wait(tasks, return_when=asyncio.ALL_COMPLETED) - else: - done = set() - yield ( - # channels: enter each channel with checkpoint - { - k: v.from_checkpoint(checkpoint["channel_values"].get(k)) - for k, v in channel_specs.items() - }, - # managed: build mapping from spec to result - ManagedValueMapping({tasks[task]: task.result() for task in done}), - ) - - -@contextmanager -def noop_context() -> Iterator[None]: - yield None + channel_specs[k] = v + yield { + k: v.from_checkpoint(checkpoint["channel_values"].get(k)) + for k, v in channel_specs.items() + } diff --git a/libs/langgraph/langgraph/pregel/messages.py b/libs/langgraph/langgraph/pregel/messages.py deleted file mode 100644 index 867012fa6..000000000 --- a/libs/langgraph/langgraph/pregel/messages.py +++ /dev/null @@ -1,185 +0,0 @@ -from typing import ( - Any, - AsyncIterator, - Callable, - Dict, - Iterator, - List, - Optional, - Sequence, - Union, - cast, -) -from uuid import UUID, uuid4 - -from langchain_core.callbacks import BaseCallbackHandler -from langchain_core.messages import BaseMessage -from langchain_core.outputs import ChatGenerationChunk, LLMResult -from langchain_core.tracers._streaming import T, _StreamingCallbackHandler - -from langgraph.constants import NS_SEP, TAG_HIDDEN, TAG_NOSTREAM -from langgraph.types import StreamChunk - -Meta = tuple[tuple[str, ...], dict[str, Any]] - - -class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler): - """A callback handler that implements stream_mode=messages. - Collects messages from (1) chat model stream events and (2) node outputs.""" - - run_inline = True - """We want this callback to run in the main thread, to avoid order/locking issues.""" - - def __init__(self, stream: Callable[[StreamChunk], None]): - self.stream = stream - self.metadata: dict[UUID, Meta] = {} - self.seen: set[Union[int, str]] = set() - - def _emit(self, meta: Meta, message: BaseMessage, *, dedupe: bool = False) -> None: - if dedupe and message.id in self.seen: - return - else: - if message.id is None: - message.id = str(uuid4()) - self.seen.add(message.id) - self.stream((meta[0], "messages", (message, meta[1]))) - - def tap_output_aiter( - self, run_id: UUID, output: AsyncIterator[T] - ) -> AsyncIterator[T]: - return output - - def tap_output_iter(self, run_id: UUID, output: Iterator[T]) -> Iterator[T]: - return output - - def on_chat_model_start( - self, - serialized: dict[str, Any], - messages: list[list[BaseMessage]], - *, - run_id: UUID, - parent_run_id: Optional[UUID] = None, - tags: Optional[list[str]] = None, - metadata: Optional[dict[str, Any]] = None, - **kwargs: Any, - ) -> Any: - if metadata and (not tags or TAG_NOSTREAM not in tags): - self.metadata[run_id] = ( - tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP)), - metadata, - ) - - def on_llm_new_token( - self, - token: str, - *, - chunk: Optional[ChatGenerationChunk] = None, - run_id: UUID, - parent_run_id: Optional[UUID] = None, - tags: Optional[list[str]] = None, - **kwargs: Any, - ) -> Any: - if not isinstance(chunk, ChatGenerationChunk): - return - if meta := self.metadata.get(run_id): - filtered_tags = [t for t in (tags or []) if not t.startswith("seq:step")] - if filtered_tags: - meta[1]["tags"] = filtered_tags - self._emit(meta, chunk.message) - - def on_llm_end( - self, - response: LLMResult, - *, - run_id: UUID, - parent_run_id: Optional[UUID] = None, - **kwargs: Any, - ) -> Any: - self.metadata.pop(run_id, None) - - def on_llm_error( - self, - error: BaseException, - *, - run_id: UUID, - parent_run_id: Optional[UUID] = None, - **kwargs: Any, - ) -> Any: - self.metadata.pop(run_id, None) - - def on_chain_start( - self, - serialized: Dict[str, Any], - inputs: Dict[str, Any], - *, - run_id: UUID, - parent_run_id: Optional[UUID] = None, - tags: Optional[List[str]] = None, - metadata: Optional[Dict[str, Any]] = None, - **kwargs: Any, - ) -> Any: - if ( - metadata - and kwargs.get("name") == metadata.get("langgraph_node") - and (not tags or TAG_HIDDEN not in tags) - ): - self.metadata[run_id] = ( - tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP)), - metadata, - ) - if isinstance(inputs, dict): - for key, value in inputs.items(): - if isinstance(value, BaseMessage): - if value.id is not None: - self.seen.add(value.id) - elif isinstance(value, Sequence) and not isinstance(value, str): - for item in value: - if isinstance(item, BaseMessage): - if item.id is not None: - self.seen.add(item.id) - - def on_chain_end( - self, - response: Any, - *, - run_id: UUID, - parent_run_id: Optional[UUID] = None, - **kwargs: Any, - ) -> Any: - if meta := self.metadata.pop(run_id, None): - if isinstance(response, BaseMessage): - self._emit(meta, response, dedupe=True) - elif isinstance(response, Sequence): - for value in response: - if isinstance(value, BaseMessage): - self._emit(meta, value, dedupe=True) - elif isinstance(response, dict): - for value in response.values(): - if isinstance(value, BaseMessage): - self._emit(meta, value, dedupe=True) - elif isinstance(value, Sequence): - for item in value: - if isinstance(item, BaseMessage): - self._emit(meta, item, dedupe=True) - elif hasattr(response, "__dir__") and callable(response.__dir__): - for key in dir(response): - try: - value = getattr(response, key) - if isinstance(value, BaseMessage): - self._emit(meta, value, dedupe=True) - elif isinstance(value, Sequence): - for item in value: - if isinstance(item, BaseMessage): - self._emit(meta, item, dedupe=True) - except AttributeError: - pass - - def on_chain_error( - self, - error: BaseException, - *, - run_id: UUID, - parent_run_id: Optional[UUID] = None, - **kwargs: Any, - ) -> Any: - self.metadata.pop(run_id, None) diff --git a/libs/langgraph/langgraph/pregel/protocol.py b/libs/langgraph/langgraph/pregel/protocol.py index ac046e949..64b427fe3 100644 --- a/libs/langgraph/langgraph/pregel/protocol.py +++ b/libs/langgraph/langgraph/pregel/protocol.py @@ -1,95 +1,53 @@ from abc import ABC, abstractmethod from typing import ( Any, - AsyncIterator, Iterator, Optional, Sequence, Union, ) -from langchain_core.runnables import Runnable, RunnableConfig -from langchain_core.runnables.graph import Graph as DrawableGraph from typing_extensions import Self from langgraph.pregel.types import All, StateSnapshot, StreamMode +from langgraph.utils.config import AnyConfig +from langgraph.utils.runnable import Runnable -class PregelProtocol( - Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]], ABC -): +class PregelProtocol(Runnable, ABC): @abstractmethod def with_config( - self, config: Optional[RunnableConfig] = None, **kwargs: Any + self, config: Optional[AnyConfig] = None, **kwargs: Any ) -> Self: ... - @abstractmethod - def get_graph( - self, - config: Optional[RunnableConfig] = None, - *, - xray: Union[int, bool] = False, - ) -> DrawableGraph: ... - - @abstractmethod - async def aget_graph( - self, - config: Optional[RunnableConfig] = None, - *, - xray: Union[int, bool] = False, - ) -> DrawableGraph: ... - @abstractmethod def get_state( - self, config: RunnableConfig, *, subgraphs: bool = False - ) -> StateSnapshot: ... - - @abstractmethod - async def aget_state( - self, config: RunnableConfig, *, subgraphs: bool = False + self, config: AnyConfig, *, subgraphs: bool = False ) -> StateSnapshot: ... @abstractmethod def get_state_history( self, - config: RunnableConfig, + config: AnyConfig, *, filter: Optional[dict[str, Any]] = None, - before: Optional[RunnableConfig] = None, + before: Optional[AnyConfig] = None, limit: Optional[int] = None, ) -> Iterator[StateSnapshot]: ... - @abstractmethod - def aget_state_history( - self, - config: RunnableConfig, - *, - filter: Optional[dict[str, Any]] = None, - before: Optional[RunnableConfig] = None, - limit: Optional[int] = None, - ) -> AsyncIterator[StateSnapshot]: ... - @abstractmethod def update_state( self, - config: RunnableConfig, + config: AnyConfig, values: Optional[Union[dict[str, Any], Any]], as_node: Optional[str] = None, - ) -> RunnableConfig: ... - - @abstractmethod - async def aupdate_state( - self, - config: RunnableConfig, - values: Optional[Union[dict[str, Any], Any]], - as_node: Optional[str] = None, - ) -> RunnableConfig: ... + ) -> AnyConfig: ... @abstractmethod def stream( self, input: Union[dict[str, Any], Any], - config: Optional[RunnableConfig] = None, + config: Optional[AnyConfig] = None, *, stream_mode: Optional[Union[StreamMode, list[StreamMode]]] = None, interrupt_before: Optional[Union[All, Sequence[str]]] = None, @@ -97,33 +55,11 @@ class PregelProtocol( subgraphs: bool = False, ) -> Iterator[Union[dict[str, Any], Any]]: ... - @abstractmethod - def astream( - self, - input: Union[dict[str, Any], Any], - config: Optional[RunnableConfig] = None, - *, - stream_mode: Optional[Union[StreamMode, list[StreamMode]]] = None, - interrupt_before: Optional[Union[All, Sequence[str]]] = None, - interrupt_after: Optional[Union[All, Sequence[str]]] = None, - subgraphs: bool = False, - ) -> AsyncIterator[Union[dict[str, Any], Any]]: ... - @abstractmethod def invoke( self, input: Union[dict[str, Any], Any], - config: Optional[RunnableConfig] = None, - *, - interrupt_before: Optional[Union[All, Sequence[str]]] = None, - interrupt_after: Optional[Union[All, Sequence[str]]] = None, - ) -> Union[dict[str, Any], Any]: ... - - @abstractmethod - async def ainvoke( - self, - input: Union[dict[str, Any], Any], - config: Optional[RunnableConfig] = None, + config: Optional[AnyConfig] = None, *, interrupt_before: Optional[Union[All, Sequence[str]]] = None, interrupt_after: Optional[Union[All, Sequence[str]]] = None, diff --git a/libs/langgraph/langgraph/pregel/read.py b/libs/langgraph/langgraph/pregel/read.py index e3d5aab3d..1240909dd 100644 --- a/libs/langgraph/langgraph/pregel/read.py +++ b/libs/langgraph/langgraph/pregel/read.py @@ -3,31 +3,24 @@ from __future__ import annotations from functools import cached_property from typing import ( Any, - AsyncIterator, Callable, - Iterator, Mapping, Optional, Sequence, Union, ) -from langchain_core.runnables import ( - Runnable, - RunnableConfig, - RunnablePassthrough, - RunnableSerializable, -) -from langchain_core.runnables.base import Input, Other, coerce_to_runnable -from langchain_core.runnables.utils import ConfigurableFieldSpec - -from langgraph.constants import CONF, CONFIG_KEY_READ +from langgraph.constants import CONF, CONFIG_KEY_READ, EMPTY_SEQ from langgraph.pregel.protocol import PregelProtocol from langgraph.pregel.retry import RetryPolicy from langgraph.pregel.utils import find_subgraph_pregel from langgraph.pregel.write import ChannelWrite -from langgraph.utils.config import merge_configs -from langgraph.utils.runnable import RunnableCallable, RunnableSeq +from langgraph.utils.config import RunnableConfig, merge_configs +from langgraph.utils.runnable import ( + Runnable, + RunnableCallable, + RunnableSeq, +) READ_TYPE = Callable[[Union[str, Sequence[str]], bool], Union[Any, dict[str, Any]]] @@ -42,18 +35,6 @@ class ChannelRead(RunnableCallable): mapper: Optional[Callable[[Any], Any]] = None - @property - def config_specs(self) -> list[ConfigurableFieldSpec]: - return [ - ConfigurableFieldSpec( - id=CONFIG_KEY_READ, - name=CONFIG_KEY_READ, - description=None, - default=None, - annotation=None, - ), - ] - def __init__( self, channel: Union[str, list[str]], @@ -67,16 +48,14 @@ class ChannelRead(RunnableCallable): self.mapper = mapper self.channel = channel - def get_name( - self, suffix: Optional[str] = None, *, name: Optional[str] = None - ) -> str: + def get_name(self, *, name: Optional[str] = None) -> str: if name: pass elif isinstance(self.channel, str): name = f"ChannelRead<{self.channel}>" else: name = f"ChannelRead<{','.join(self.channel)}>" - return super().get_name(suffix, name=name) + return super().get_name(name=name) def _read(self, _: Any, config: RunnableConfig) -> Any: return self.do_read( @@ -109,7 +88,7 @@ class ChannelRead(RunnableCallable): return read(select, fresh) -DEFAULT_BOUND: RunnablePassthrough = RunnablePassthrough() +DEFAULT_BOUND = RunnableCallable(lambda input: input) class PregelNode(Runnable): @@ -161,6 +140,7 @@ class PregelNode(Runnable): metadata: Optional[Mapping[str, Any]] = None, bound: Optional[Runnable[Any, Any]] = None, retry_policy: Optional[RetryPolicy] = None, + subgraphs: Sequence[PregelProtocol] = EMPTY_SEQ, ) -> None: self.channels = channels self.triggers = list(triggers) @@ -170,7 +150,9 @@ class PregelNode(Runnable): self.retry_policy = retry_policy self.tags = tags self.metadata = metadata - if self.bound is not DEFAULT_BOUND: + if subgraphs: + self.subgraphs = list(subgraphs) + elif self.bound is not DEFAULT_BOUND: try: subgraph = find_subgraph_pregel(self.bound) except Exception: @@ -201,7 +183,6 @@ class PregelNode(Runnable): writers[-2] = ChannelWrite( writes=writers[-2].writes + writers[-1].writes, tags=writers[-2].tags, - require_at_least_one_of=writers[-2].require_at_least_one_of, ) writers.pop() return writers @@ -221,59 +202,9 @@ class PregelNode(Runnable): else: return self.bound - def join(self, channels: Sequence[str]) -> PregelNode: - assert isinstance(channels, list) or isinstance( - channels, tuple - ), "channels must be a list or tuple" - assert isinstance( - self.channels, dict - ), "all channels must be named when using .join()" - return self.copy( - update=dict( - channels={ - **self.channels, - **{chan: chan for chan in channels}, - } - ), - ) - - def __or__( - self, - other: Union[ - Runnable[Any, Other], - Callable[[Any], Other], - Mapping[str, Runnable[Any, Other] | Callable[[Any], Other]], - ], - ) -> PregelNode: - if isinstance(other, Runnable) and ChannelWrite.is_writer(other): - return self.copy(update=dict(writers=[*self.writers, other])) - elif self.bound is DEFAULT_BOUND: - return self.copy(update=dict(bound=coerce_to_runnable(other))) - else: - return self.copy(update=dict(bound=RunnableSeq(self.bound, other))) - - def pipe( - self, - *others: Runnable[Any, Other] | Callable[[Any], Other], - name: Optional[str] = None, - ) -> RunnableSerializable[Any, Other]: - for other in others: - self = self | other - return self - - def __ror__( - self, - other: Union[ - Runnable[Other, Any], - Callable[[Any], Other], - Mapping[str, Union[Runnable[Other, Any], Callable[[Other], Any]]], - ], - ) -> RunnableSerializable: - raise NotImplementedError() - def invoke( self, - input: Input, + input: dict[str, Any], config: Optional[RunnableConfig] = None, **kwargs: Optional[Any], ) -> Any: @@ -282,40 +213,3 @@ class PregelNode(Runnable): merge_configs({"metadata": self.metadata, "tags": self.tags}, config), **kwargs, ) - - async def ainvoke( - self, - input: Input, - config: Optional[RunnableConfig] = None, - **kwargs: Optional[Any], - ) -> Any: - return await self.bound.ainvoke( - input, - merge_configs({"metadata": self.metadata, "tags": self.tags}, config), - **kwargs, - ) - - def stream( - self, - input: Input, - config: Optional[RunnableConfig] = None, - **kwargs: Optional[Any], - ) -> Iterator[Any]: - yield from self.bound.stream( - input, - merge_configs({"metadata": self.metadata, "tags": self.tags}, config), - **kwargs, - ) - - async def astream( - self, - input: Input, - config: Optional[RunnableConfig] = None, - **kwargs: Optional[Any], - ) -> AsyncIterator[Any]: - async for item in self.bound.astream( - input, - merge_configs({"metadata": self.metadata, "tags": self.tags}, config), - **kwargs, - ): - yield item diff --git a/libs/langgraph/langgraph/pregel/remote.py b/libs/langgraph/langgraph/pregel/remote.py deleted file mode 100644 index 6ba547b14..000000000 --- a/libs/langgraph/langgraph/pregel/remote.py +++ /dev/null @@ -1,841 +0,0 @@ -from dataclasses import asdict -from typing import ( - Any, - AsyncIterator, - Iterator, - Literal, - Optional, - Sequence, - Union, - cast, -) - -import orjson -from langchain_core.runnables import RunnableConfig -from langchain_core.runnables.graph import ( - Edge as DrawableEdge, -) -from langchain_core.runnables.graph import ( - Graph as DrawableGraph, -) -from langchain_core.runnables.graph import ( - Node as DrawableNode, -) -from langgraph_sdk.client import ( - LangGraphClient, - SyncLangGraphClient, - get_client, - get_sync_client, -) -from langgraph_sdk.schema import Checkpoint, ThreadState -from langgraph_sdk.schema import Command as CommandSDK -from langgraph_sdk.schema import StreamMode as StreamModeSDK -from typing_extensions import Self - -from langgraph.checkpoint.base import CheckpointMetadata -from langgraph.constants import ( - CONF, - CONFIG_KEY_CHECKPOINT_NS, - CONFIG_KEY_STREAM, - INTERRUPT, - NS_SEP, -) -from langgraph.errors import GraphInterrupt -from langgraph.pregel.protocol import PregelProtocol -from langgraph.pregel.types import All, PregelTask, StateSnapshot, StreamMode -from langgraph.types import Command, Interrupt, StreamProtocol -from langgraph.utils.config import merge_configs - - -class RemoteException(Exception): - """Exception raised when an error occurs in the remote graph.""" - - pass - - -class RemoteGraph(PregelProtocol): - """The `RemoteGraph` class is a client implementation for calling remote - APIs that implement the LangGraph Server API specification. - - For example, the `RemoteGraph` class can be used to call APIs from deployments - on LangGraph Cloud. - - `RemoteGraph` behaves the same way as a `Graph` and can be used directly as - a node in another `Graph`. - """ - - name: str - - def __init__( - self, - name: str, # graph_id - /, - *, - url: Optional[str] = None, - api_key: Optional[str] = None, - headers: Optional[dict[str, str]] = None, - client: Optional[LangGraphClient] = None, - sync_client: Optional[SyncLangGraphClient] = None, - config: Optional[RunnableConfig] = None, - ): - """Specify `url`, `api_key`, and/or `headers` to create default sync and async clients. - - If `client` or `sync_client` are provided, they will be used instead of the default clients. - See `LangGraphClient` and `SyncLangGraphClient` for details on the default clients. At least - one of `url`, `client`, or `sync_client` must be provided. - - Args: - name: The name of the graph. - url: The URL of the remote API. - api_key: The API key to use for authentication. If not provided, it will be read from the environment (`LANGGRAPH_API_KEY`, `LANGSMITH_API_KEY`, or `LANGCHAIN_API_KEY`). - headers: Additional headers to include in the requests. - client: A `LangGraphClient` instance to use instead of creating a default client. - sync_client: A `SyncLangGraphClient` instance to use instead of creating a default client. - config: An optional `RunnableConfig` instance with additional configuration. - """ - self.name = name - self.config = config - - if client is None and url is not None: - client = get_client(url=url, api_key=api_key, headers=headers) - self.client = client - - if sync_client is None and url is not None: - sync_client = get_sync_client(url=url, api_key=api_key, headers=headers) - self.sync_client = sync_client - - def _validate_client(self) -> LangGraphClient: - if self.client is None: - raise ValueError( - "Async client is not initialized: please provide `url` or `client` when initializing `RemoteGraph`." - ) - return self.client - - def _validate_sync_client(self) -> SyncLangGraphClient: - if self.sync_client is None: - raise ValueError( - "Sync client is not initialized: please provide `url` or `sync_client` when initializing `RemoteGraph`." - ) - return self.sync_client - - def copy(self, update: dict[str, Any]) -> Self: - attrs = {**self.__dict__, **update} - return self.__class__(attrs.pop("name"), **attrs) - - def with_config( - self, config: Optional[RunnableConfig] = None, **kwargs: Any - ) -> Self: - return self.copy( - {"config": merge_configs(self.config, config, cast(RunnableConfig, kwargs))} - ) - - def _get_drawable_nodes( - self, graph: dict[str, list[dict[str, Any]]] - ) -> dict[str, DrawableNode]: - nodes = {} - for node in graph["nodes"]: - node_id = str(node["id"]) - node_data = node.get("data", {}) - - # Get node name from node_data if available. If not, use node_id. - node_name = node.get("name") - if node_name is None: - if isinstance(node_data, dict): - node_name = node_data.get("name", node_id) - else: - node_name = node_id - - nodes[node_id] = DrawableNode( - id=node_id, - name=node_name, - data=node_data, - metadata=node.get("metadata"), - ) - return nodes - - def get_graph( - self, - config: Optional[RunnableConfig] = None, - *, - xray: Union[int, bool] = False, - ) -> DrawableGraph: - """Get graph by graph name. - - This method calls `GET /assistants/{assistant_id}/graph`. - - Args: - config: This parameter is not used. - xray: Include graph representation of subgraphs. If an integer - value is provided, only subgraphs with a depth less than or - equal to the value will be included. - - Returns: - The graph information for the assistant in JSON format. - """ - sync_client = self._validate_sync_client() - graph = sync_client.assistants.get_graph( - assistant_id=self.name, - xray=xray, - ) - return DrawableGraph( - nodes=self._get_drawable_nodes(graph), - edges=[DrawableEdge(**edge) for edge in graph["edges"]], - ) - - async def aget_graph( - self, - config: Optional[RunnableConfig] = None, - *, - xray: Union[int, bool] = False, - ) -> DrawableGraph: - """Get graph by graph name. - - This method calls `GET /assistants/{assistant_id}/graph`. - - Args: - config: This parameter is not used. - xray: Include graph representation of subgraphs. If an integer - value is provided, only subgraphs with a depth less than or - equal to the value will be included. - - Returns: - The graph information for the assistant in JSON format. - """ - client = self._validate_client() - graph = await client.assistants.get_graph( - assistant_id=self.name, - xray=xray, - ) - return DrawableGraph( - nodes=self._get_drawable_nodes(graph), - edges=[DrawableEdge(**edge) for edge in graph["edges"]], - ) - - def _create_state_snapshot(self, state: ThreadState) -> StateSnapshot: - tasks = [] - for task in state["tasks"]: - interrupts = [] - for interrupt in task["interrupts"]: - interrupts.append(Interrupt(**interrupt)) - - tasks.append( - PregelTask( - id=task["id"], - name=task["name"], - path=tuple(), - error=Exception(task["error"]) if task["error"] else None, - interrupts=tuple(interrupts), - state=self._create_state_snapshot(task["state"]) - if task["state"] - else cast(RunnableConfig, {"configurable": task["checkpoint"]}) - if task["checkpoint"] - else None, - result=task.get("result"), - ) - ) - - return StateSnapshot( - values=state["values"], - next=tuple(state["next"]) if state["next"] else tuple(), - config={ - "configurable": { - "thread_id": state["checkpoint"]["thread_id"], - "checkpoint_ns": state["checkpoint"]["checkpoint_ns"], - "checkpoint_id": state["checkpoint"]["checkpoint_id"], - "checkpoint_map": state["checkpoint"].get("checkpoint_map", {}), - } - }, - metadata=CheckpointMetadata(**state["metadata"]), - created_at=state["created_at"], - parent_config={ - "configurable": { - "thread_id": state["parent_checkpoint"]["thread_id"], - "checkpoint_ns": state["parent_checkpoint"]["checkpoint_ns"], - "checkpoint_id": state["parent_checkpoint"]["checkpoint_id"], - "checkpoint_map": state["parent_checkpoint"].get( - "checkpoint_map", {} - ), - } - } - if state["parent_checkpoint"] - else None, - tasks=tuple(tasks), - ) - - def _get_checkpoint(self, config: Optional[RunnableConfig]) -> Optional[Checkpoint]: - if config is None: - return None - - checkpoint = {} - - if "thread_id" in config["configurable"]: - checkpoint["thread_id"] = config["configurable"]["thread_id"] - if "checkpoint_ns" in config["configurable"]: - checkpoint["checkpoint_ns"] = config["configurable"]["checkpoint_ns"] - if "checkpoint_id" in config["configurable"]: - checkpoint["checkpoint_id"] = config["configurable"]["checkpoint_id"] - if "checkpoint_map" in config["configurable"]: - checkpoint["checkpoint_map"] = config["configurable"]["checkpoint_map"] - - return checkpoint if checkpoint else None - - def _get_config(self, checkpoint: Checkpoint) -> RunnableConfig: - return { - "configurable": { - "thread_id": checkpoint["thread_id"], - "checkpoint_ns": checkpoint["checkpoint_ns"], - "checkpoint_id": checkpoint["checkpoint_id"], - "checkpoint_map": checkpoint.get("checkpoint_map", {}), - } - } - - def _sanitize_config(self, config: RunnableConfig) -> RunnableConfig: - reserved_configurable_keys = frozenset( - [ - "callbacks", - "checkpoint_map", - "checkpoint_id", - "checkpoint_ns", - ] - ) - - def _sanitize_obj(obj: Any) -> Any: - """Remove non-JSON serializable fields from the given object.""" - if isinstance(obj, dict): - return {k: _sanitize_obj(v) for k, v in obj.items()} - elif isinstance(obj, list): - return [_sanitize_obj(v) for v in obj] - else: - try: - orjson.dumps(obj) - return obj - except orjson.JSONEncodeError: - return None - - # Remove non-JSON serializable fields from the config. - config = _sanitize_obj(config) - - # Only include configurable keys that are not reserved and - # not starting with "__pregel_" prefix. - new_configurable = { - k: v - for k, v in config["configurable"].items() - if k not in reserved_configurable_keys and not k.startswith("__pregel_") - } - - sanitized: RunnableConfig = { - "tags": config.get("tags") or [], - "metadata": config.get("metadata") or {}, - "configurable": new_configurable, - } - if "recursion_limit" in config: - sanitized["recursion_limit"] = config["recursion_limit"] - - return sanitized - - def get_state( - self, config: RunnableConfig, *, subgraphs: bool = False - ) -> StateSnapshot: - """Get the state of a thread. - - This method calls `POST /threads/{thread_id}/state/checkpoint` if a - checkpoint is specified in the config or `GET /threads/{thread_id}/state` - if no checkpoint is specified. - - Args: - config: A `RunnableConfig` that includes `thread_id` in the - `configurable` field. - subgraphs: Include subgraphs in the state. - - Returns: - The latest state of the thread. - """ - sync_client = self._validate_sync_client() - merged_config = merge_configs(self.config, config) - - state = sync_client.threads.get_state( - thread_id=merged_config["configurable"]["thread_id"], - checkpoint=self._get_checkpoint(merged_config), - subgraphs=subgraphs, - ) - return self._create_state_snapshot(state) - - async def aget_state( - self, config: RunnableConfig, *, subgraphs: bool = False - ) -> StateSnapshot: - """Get the state of a thread. - - This method calls `POST /threads/{thread_id}/state/checkpoint` if a - checkpoint is specified in the config or `GET /threads/{thread_id}/state` - if no checkpoint is specified. - - Args: - config: A `RunnableConfig` that includes `thread_id` in the - `configurable` field. - subgraphs: Include subgraphs in the state. - - Returns: - The latest state of the thread. - """ - client = self._validate_client() - merged_config = merge_configs(self.config, config) - - state = await client.threads.get_state( - thread_id=merged_config["configurable"]["thread_id"], - checkpoint=self._get_checkpoint(merged_config), - subgraphs=subgraphs, - ) - return self._create_state_snapshot(state) - - def get_state_history( - self, - config: RunnableConfig, - *, - filter: Optional[dict[str, Any]] = None, - before: Optional[RunnableConfig] = None, - limit: Optional[int] = None, - ) -> Iterator[StateSnapshot]: - """Get the state history of a thread. - - This method calls `POST /threads/{thread_id}/history`. - - Args: - config: A `RunnableConfig` that includes `thread_id` in the - `configurable` field. - filter: Metadata to filter on. - before: A `RunnableConfig` that includes checkpoint metadata. - limit: Max number of states to return. - - Returns: - States of the thread. - """ - sync_client = self._validate_sync_client() - merged_config = merge_configs(self.config, config) - - states = sync_client.threads.get_history( - thread_id=merged_config["configurable"]["thread_id"], - limit=limit if limit else 10, - before=self._get_checkpoint(before), - metadata=filter, - checkpoint=self._get_checkpoint(merged_config), - ) - for state in states: - yield self._create_state_snapshot(state) - - async def aget_state_history( - self, - config: RunnableConfig, - *, - filter: Optional[dict[str, Any]] = None, - before: Optional[RunnableConfig] = None, - limit: Optional[int] = None, - ) -> AsyncIterator[StateSnapshot]: - """Get the state history of a thread. - - This method calls `POST /threads/{thread_id}/history`. - - Args: - config: A `RunnableConfig` that includes `thread_id` in the - `configurable` field. - filter: Metadata to filter on. - before: A `RunnableConfig` that includes checkpoint metadata. - limit: Max number of states to return. - - Returns: - States of the thread. - """ - client = self._validate_client() - merged_config = merge_configs(self.config, config) - - states = await client.threads.get_history( - thread_id=merged_config["configurable"]["thread_id"], - limit=limit if limit else 10, - before=self._get_checkpoint(before), - metadata=filter, - checkpoint=self._get_checkpoint(merged_config), - ) - for state in states: - yield self._create_state_snapshot(state) - - def update_state( - self, - config: RunnableConfig, - values: Optional[Union[dict[str, Any], Any]], - as_node: Optional[str] = None, - ) -> RunnableConfig: - """Update the state of a thread. - - This method calls `POST /threads/{thread_id}/state`. - - Args: - config: A `RunnableConfig` that includes `thread_id` in the - `configurable` field. - values: Values to update to the state. - as_node: Update the state as if this node had just executed. - - Returns: - `RunnableConfig` for the updated thread. - """ - sync_client = self._validate_sync_client() - merged_config = merge_configs(self.config, config) - - response: dict = sync_client.threads.update_state( # type: ignore - thread_id=merged_config["configurable"]["thread_id"], - values=values, - as_node=as_node, - checkpoint=self._get_checkpoint(merged_config), - ) - return self._get_config(response["checkpoint"]) - - async def aupdate_state( - self, - config: RunnableConfig, - values: Optional[Union[dict[str, Any], Any]], - as_node: Optional[str] = None, - ) -> RunnableConfig: - """Update the state of a thread. - - This method calls `POST /threads/{thread_id}/state`. - - Args: - config: A `RunnableConfig` that includes `thread_id` in the - `configurable` field. - values: Values to update to the state. - as_node: Update the state as if this node had just executed. - - Returns: - `RunnableConfig` for the updated thread. - """ - client = self._validate_client() - merged_config = merge_configs(self.config, config) - - response: dict = await client.threads.update_state( # type: ignore - thread_id=merged_config["configurable"]["thread_id"], - values=values, - as_node=as_node, - checkpoint=self._get_checkpoint(merged_config), - ) - return self._get_config(response["checkpoint"]) - - def _get_stream_modes( - self, - stream_mode: Optional[Union[StreamMode, list[StreamMode]]], - config: Optional[RunnableConfig], - default: StreamMode = "updates", - ) -> tuple[ - list[StreamModeSDK], list[StreamModeSDK], bool, Optional[StreamProtocol] - ]: - """Return a tuple of the final list of stream modes sent to the - remote graph and a boolean flag indicating if stream mode 'updates' - was present in the original list of stream modes. - - 'updates' mode is added to the list of stream modes so that interrupts - can be detected in the remote graph. - """ - updated_stream_modes: list[StreamModeSDK] = [] - req_single = True - # coerce to list, or add default stream mode - if stream_mode: - if isinstance(stream_mode, str): - updated_stream_modes.append(stream_mode) - else: - req_single = False - updated_stream_modes.extend(stream_mode) - else: - updated_stream_modes.append(default) - requested_stream_modes = updated_stream_modes.copy() - # add any from parent graph - stream: Optional[StreamProtocol] = ( - (config or {}).get(CONF, {}).get(CONFIG_KEY_STREAM) - ) - if stream: - updated_stream_modes.extend(stream.modes) - # map "messages" to "messages-tuple" - if "messages" in updated_stream_modes: - updated_stream_modes.remove("messages") - updated_stream_modes.append("messages-tuple") - - # if requested "messages-tuple", - # map to "messages" in requested_stream_modes - if "messages-tuple" in requested_stream_modes: - requested_stream_modes.remove("messages-tuple") - requested_stream_modes.append("messages") - - # add 'updates' mode if not present - if "updates" not in updated_stream_modes: - updated_stream_modes.append("updates") - - # remove 'events', as it's not supported in Pregel - if "events" in updated_stream_modes: - updated_stream_modes.remove("events") - return (updated_stream_modes, requested_stream_modes, req_single, stream) - - def stream( - self, - input: Union[dict[str, Any], Any], - config: Optional[RunnableConfig] = None, - *, - stream_mode: Optional[Union[StreamMode, list[StreamMode]]] = None, - interrupt_before: Optional[Union[All, Sequence[str]]] = None, - interrupt_after: Optional[Union[All, Sequence[str]]] = None, - subgraphs: bool = False, - **kwargs: Any, - ) -> Iterator[Union[dict[str, Any], Any]]: - """Create a run and stream the results. - - This method calls `POST /threads/{thread_id}/runs/stream` if a `thread_id` - is speciffed in the `configurable` field of the config or - `POST /runs/stream` otherwise. - - Args: - input: Input to the graph. - config: A `RunnableConfig` for graph invocation. - stream_mode: Stream mode(s) to use. - interrupt_before: Interrupt the graph before these nodes. - interrupt_after: Interrupt the graph after these nodes. - subgraphs: Stream from subgraphs. - **kwargs: Additional params to pass to client.runs.stream. - - Yields: - The output of the graph. - """ - sync_client = self._validate_sync_client() - merged_config = merge_configs(self.config, config) - sanitized_config = self._sanitize_config(merged_config) - stream_modes, requested, req_single, stream = self._get_stream_modes( - stream_mode, config - ) - if isinstance(input, Command): - command: Optional[CommandSDK] = cast(CommandSDK, asdict(input)) - input = None - else: - command = None - - for chunk in sync_client.runs.stream( - thread_id=sanitized_config["configurable"].get("thread_id"), - assistant_id=self.name, - input=input, - command=command, - config=sanitized_config, - stream_mode=stream_modes, - interrupt_before=interrupt_before, - interrupt_after=interrupt_after, - stream_subgraphs=subgraphs or stream is not None, - if_not_exists="create", - **kwargs, - ): - # split mode and ns - if NS_SEP in chunk.event: - mode, ns_ = chunk.event.split(NS_SEP, 1) - ns = tuple(ns_.split(NS_SEP)) - else: - mode, ns = chunk.event, () - # prepend caller ns (as it is not passed to remote graph) - if caller_ns := (config or {}).get(CONF, {}).get(CONFIG_KEY_CHECKPOINT_NS): - caller_ns = tuple(caller_ns.split(NS_SEP)) - ns = caller_ns + ns - # stream to parent stream - if stream is not None and mode in stream.modes: - stream((ns, mode, chunk.data)) - # raise interrupt or errors - if chunk.event.startswith("updates"): - if isinstance(chunk.data, dict) and INTERRUPT in chunk.data: - raise GraphInterrupt(chunk.data[INTERRUPT]) - elif chunk.event.startswith("error"): - raise RemoteException(chunk.data) - # filter for what was actually requested - if mode not in requested: - continue - # emit chunk - if subgraphs: - if NS_SEP in chunk.event: - mode, ns_ = chunk.event.split(NS_SEP, 1) - ns = tuple(ns_.split(NS_SEP)) - else: - mode, ns = chunk.event, () - if req_single: - yield ns, chunk.data - else: - yield ns, mode, chunk.data - elif req_single: - yield chunk.data - else: - yield chunk - - async def astream( - self, - input: Union[dict[str, Any], Any], - config: Optional[RunnableConfig] = None, - *, - stream_mode: Optional[Union[StreamMode, list[StreamMode]]] = None, - interrupt_before: Optional[Union[All, Sequence[str]]] = None, - interrupt_after: Optional[Union[All, Sequence[str]]] = None, - subgraphs: bool = False, - **kwargs: Any, - ) -> AsyncIterator[Union[dict[str, Any], Any]]: - """Create a run and stream the results. - - This method calls `POST /threads/{thread_id}/runs/stream` if a `thread_id` - is speciffed in the `configurable` field of the config or - `POST /runs/stream` otherwise. - - Args: - input: Input to the graph. - config: A `RunnableConfig` for graph invocation. - stream_mode: Stream mode(s) to use. - interrupt_before: Interrupt the graph before these nodes. - interrupt_after: Interrupt the graph after these nodes. - subgraphs: Stream from subgraphs. - **kwargs: Additional params to pass to client.runs.stream. - - Yields: - The output of the graph. - """ - client = self._validate_client() - merged_config = merge_configs(self.config, config) - sanitized_config = self._sanitize_config(merged_config) - stream_modes, requested, req_single, stream = self._get_stream_modes( - stream_mode, config - ) - if isinstance(input, Command): - command: Optional[CommandSDK] = cast(CommandSDK, asdict(input)) - input = None - else: - command = None - - async for chunk in client.runs.stream( - thread_id=sanitized_config["configurable"].get("thread_id"), - assistant_id=self.name, - input=input, - command=command, - config=sanitized_config, - stream_mode=stream_modes, - interrupt_before=interrupt_before, - interrupt_after=interrupt_after, - stream_subgraphs=subgraphs or stream is not None, - if_not_exists="create", - **kwargs, - ): - # split mode and ns - if NS_SEP in chunk.event: - mode, ns_ = chunk.event.split(NS_SEP, 1) - ns = tuple(ns_.split(NS_SEP)) - else: - mode, ns = chunk.event, () - # prepend caller ns (as it is not passed to remote graph) - if caller_ns := (config or {}).get(CONF, {}).get(CONFIG_KEY_CHECKPOINT_NS): - caller_ns = tuple(caller_ns.split(NS_SEP)) - ns = caller_ns + ns - # stream to parent stream - if stream is not None and mode in stream.modes: - stream((ns, mode, chunk.data)) - # raise interrupt or errors - if chunk.event.startswith("updates"): - if isinstance(chunk.data, dict) and INTERRUPT in chunk.data: - raise GraphInterrupt(chunk.data[INTERRUPT]) - elif chunk.event.startswith("error"): - raise RemoteException(chunk.data) - # filter for what was actually requested - if mode not in requested: - continue - # emit chunk - if subgraphs: - if NS_SEP in chunk.event: - mode, ns_ = chunk.event.split(NS_SEP, 1) - ns = tuple(ns_.split(NS_SEP)) - else: - mode, ns = chunk.event, () - if req_single: - yield ns, chunk.data - else: - yield ns, mode, chunk.data - elif req_single: - yield chunk.data - else: - yield chunk - - async def astream_events( - self, - input: Any, - config: Optional[RunnableConfig] = None, - *, - version: Literal["v1", "v2"], - include_names: Optional[Sequence[All]] = None, - include_types: Optional[Sequence[All]] = None, - include_tags: Optional[Sequence[All]] = None, - exclude_names: Optional[Sequence[All]] = None, - exclude_types: Optional[Sequence[All]] = None, - exclude_tags: Optional[Sequence[All]] = None, - **kwargs: Any, - ) -> AsyncIterator[dict[str, Any]]: - raise NotImplementedError - - def invoke( - self, - input: Union[dict[str, Any], Any], - config: Optional[RunnableConfig] = None, - *, - interrupt_before: Optional[Union[All, Sequence[str]]] = None, - interrupt_after: Optional[Union[All, Sequence[str]]] = None, - **kwargs: Any, - ) -> Union[dict[str, Any], Any]: - """Create a run, wait until it finishes and return the final state. - - Args: - input: Input to the graph. - config: A `RunnableConfig` for graph invocation. - interrupt_before: Interrupt the graph before these nodes. - interrupt_after: Interrupt the graph after these nodes. - **kwargs: Additional params to pass to RemoteGraph.stream. - - Returns: - The output of the graph. - """ - for chunk in self.stream( - input, - config=config, - interrupt_before=interrupt_before, - interrupt_after=interrupt_after, - stream_mode="values", - **kwargs, - ): - pass - try: - return chunk - except UnboundLocalError: - return None - - async def ainvoke( - self, - input: Union[dict[str, Any], Any], - config: Optional[RunnableConfig] = None, - *, - interrupt_before: Optional[Union[All, Sequence[str]]] = None, - interrupt_after: Optional[Union[All, Sequence[str]]] = None, - **kwargs: Any, - ) -> Union[dict[str, Any], Any]: - """Create a run, wait until it finishes and return the final state. - - Args: - input: Input to the graph. - config: A `RunnableConfig` for graph invocation. - interrupt_before: Interrupt the graph before these nodes. - interrupt_after: Interrupt the graph after these nodes. - **kwargs: Additional params to pass to RemoteGraph.astream. - - Returns: - The output of the graph. - """ - async for chunk in self.astream( - input, - config=config, - interrupt_before=interrupt_before, - interrupt_after=interrupt_after, - stream_mode="values", - **kwargs, - ): - pass - try: - return chunk - except UnboundLocalError: - return None diff --git a/libs/langgraph/langgraph/pregel/retry.py b/libs/langgraph/langgraph/pregel/retry.py index 6d0e43b54..820f88522 100644 --- a/libs/langgraph/langgraph/pregel/retry.py +++ b/libs/langgraph/langgraph/pregel/retry.py @@ -1,10 +1,9 @@ -import asyncio import logging import random import sys import time from dataclasses import replace -from typing import Any, Optional, Sequence +from typing import Optional, Sequence from langgraph.constants import ( CONF, @@ -23,15 +22,12 @@ SUPPORTS_EXC_NOTES = sys.version_info >= (3, 11) def run_with_retry( task: PregelExecutableTask, retry_policy: Optional[RetryPolicy], - configurable: Optional[dict[str, Any]] = None, ) -> None: """Run a task with retries.""" retry_policy = task.retry_policy or retry_policy interval = retry_policy.initial_interval if retry_policy else 0 attempts = 0 config = task.config - if configurable is not None: - config = patch_configurable(config, configurable) while True: try: # clear any writes from previous attempts @@ -99,91 +95,3 @@ def run_with_retry( ) # signal subgraphs to resume (if available) config = patch_configurable(config, {CONFIG_KEY_RESUMING: True}) - - -async def arun_with_retry( - task: PregelExecutableTask, - retry_policy: Optional[RetryPolicy], - stream: bool = False, - configurable: Optional[dict[str, Any]] = None, -) -> None: - """Run a task asynchronously with retries.""" - retry_policy = task.retry_policy or retry_policy - interval = retry_policy.initial_interval if retry_policy else 0 - attempts = 0 - config = task.config - if configurable is not None: - config = patch_configurable(config, configurable) - while True: - try: - # clear any writes from previous attempts - task.writes.clear() - # run the task - if stream: - async for _ in task.proc.astream(task.input, config): - pass - # if successful, end - break - else: - return await task.proc.ainvoke(task.input, config) - except ParentCommand as exc: - ns: str = config[CONF][CONFIG_KEY_CHECKPOINT_NS] - cmd = exc.args[0] - if cmd.graph == ns: - # this command is for the current graph, handle it - for w in task.writers: - w.invoke(cmd, config) - break - elif cmd.graph == Command.PARENT: - # this command is for the parent graph, assign it to the parent - parts = ns.split(NS_SEP) - if parts[-1].isdigit(): - parts.pop() - parent_ns = NS_SEP.join(parts[:-1]) - exc.args = (replace(cmd, graph=parent_ns),) - # bubble up - raise - except GraphBubbleUp: - # if interrupted, end - raise - except Exception as exc: - if SUPPORTS_EXC_NOTES: - exc.add_note(f"During task with name '{task.name}' and id '{task.id}'") - if retry_policy is None: - raise - # increment attempts - attempts += 1 - # check if we should retry - if isinstance(retry_policy.retry_on, Sequence): - if not isinstance(exc, tuple(retry_policy.retry_on)): - raise - elif isinstance(retry_policy.retry_on, type) and issubclass( - retry_policy.retry_on, Exception - ): - if not isinstance(exc, retry_policy.retry_on): - raise - elif callable(retry_policy.retry_on): - if not retry_policy.retry_on(exc): # type: ignore[call-arg] - raise - else: - raise TypeError( - "retry_on must be an Exception class, a list or tuple of Exception classes, or a callable" - ) - # check if we should give up - if attempts >= retry_policy.max_attempts: - raise - # sleep before retrying - interval = min( - retry_policy.max_interval, - interval * retry_policy.backoff_factor, - ) - await asyncio.sleep( - interval + random.uniform(0, 1) if retry_policy.jitter else interval - ) - # log the retry - logger.info( - f"Retrying task {task.name} after {interval:.2f} seconds (attempt {attempts}) after {exc.__class__.__name__} {exc}", - exc_info=exc, - ) - # signal subgraphs to resume (if available) - config = patch_configurable(config, {CONFIG_KEY_RESUMING: True}) diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 6336bc5a1..3929f1754 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -1,53 +1,34 @@ -import asyncio import concurrent.futures import threading import time from functools import partial from typing import ( Any, - AsyncIterator, - Awaitable, Callable, - Generic, Iterable, Iterator, Optional, Sequence, Type, - TypeVar, - Union, - cast, ) -from langchain_core.callbacks import Callbacks - from langgraph.constants import ( - CONF, - CONFIG_KEY_CALL, - CONFIG_KEY_SCRATCHPAD, - CONFIG_KEY_SEND, ERROR, INTERRUPT, - MISSING, NO_WRITES, - PUSH, RESUME, - RETURN, TAG_HIDDEN, ) from langgraph.errors import GraphBubbleUp, GraphInterrupt -from langgraph.pregel.algo import Call from langgraph.pregel.executor import Submit -from langgraph.pregel.retry import arun_with_retry, run_with_retry -from langgraph.types import PregelExecutableTask, PregelScratchpad, RetryPolicy -from langgraph.utils.future import chain_future +from langgraph.pregel.retry import run_with_retry +from langgraph.types import PregelExecutableTask, RetryPolicy -F = TypeVar("F", concurrent.futures.Future, asyncio.Future) -E = TypeVar("E", threading.Event, asyncio.Event) +F = concurrent.futures.Future -class FuturesDict(Generic[F, E], dict[F, Optional[PregelExecutableTask]]): - event: E +class FuturesDict(dict[concurrent.futures.Future, Optional[PregelExecutableTask]]): + event: threading.Event callback: Callable[[PregelExecutableTask, Optional[BaseException]], None] counter: int done: set[F] @@ -55,24 +36,22 @@ class FuturesDict(Generic[F, E], dict[F, Optional[PregelExecutableTask]]): def __init__( self, - event: E, + event: threading.Event, callback: Callable[[PregelExecutableTask, Optional[BaseException]], None], - future_type: Type[F], - # used for generic typing, newer py supports FutureDict[...](...) ) -> None: super().__init__() self.lock = threading.Lock() self.event = event self.callback = callback self.counter = 0 - self.done: set[F] = set() + self.done: set[concurrent.futures.Future] = set() def __setitem__( self, - key: F, + key: concurrent.futures.Future, value: Optional[PregelExecutableTask], ) -> None: - super().__setitem__(key, value) # type: ignore[index] + super().__setitem__(key, value) if value is not None: with self.lock: self.event.clear() @@ -82,7 +61,7 @@ class FuturesDict(Generic[F, E], dict[F, Optional[PregelExecutableTask]]): def on_done( self, task: PregelExecutableTask, - fut: F, + fut: concurrent.futures.Future, ) -> None: try: self.callback(task, _exception(fut)) @@ -104,9 +83,6 @@ class PregelRunner: *, submit: Submit, put_writes: Callable[[str, Sequence[tuple[str, Any]]], None], - schedule_task: Callable[ - [PregelExecutableTask, int, Optional[Call]], Optional[PregelExecutableTask] - ], use_astream: bool = False, node_finished: Optional[Callable[[str], None]] = None, ) -> None: @@ -114,7 +90,6 @@ class PregelRunner: self.put_writes = put_writes self.use_astream = use_astream self.node_finished = node_finished - self.schedule_task = schedule_task def tick( self, @@ -125,101 +100,10 @@ class PregelRunner: retry_policy: Optional[RetryPolicy] = None, get_waiter: Optional[Callable[[], concurrent.futures.Future[None]]] = None, ) -> Iterator[None]: - def writer( - task: PregelExecutableTask, - writes: Sequence[tuple[str, Any]], - *, - calls: Optional[Sequence[Call]] = None, - ) -> Sequence[Optional[concurrent.futures.Future]]: - if all(w[0] != PUSH for w in writes): - return task.config[CONF][CONFIG_KEY_SEND](writes) - - # schedule PUSH tasks, collect futures - scratchpad: PregelScratchpad = task.config[CONF][CONFIG_KEY_SCRATCHPAD] - rtn: dict[int, Optional[concurrent.futures.Future]] = {} - for idx, w in enumerate(writes): - # bail if not a PUSH write - if w[0] != PUSH: - continue - # schedule the next task, if the callback returns one - wcall = calls[idx] if calls else None - if next_task := self.schedule_task( - task, scratchpad.call_counter(), wcall - ): - if fut := next( - ( - f - for f, t in futures.items() - if t is not None and t == next_task.id - ), - None, - ): - # if the parent task was retried, - # the next task might already be running - rtn[idx] = fut - elif next_task.writes: - # if it already ran, return the result - fut = concurrent.futures.Future() - ret = next( - (v for c, v in next_task.writes if c == RETURN), MISSING - ) - if ret is not MISSING: - fut.set_result(ret) - elif exc := next( - (v for c, v in next_task.writes if c == ERROR), None - ): - fut.set_exception( - exc - if isinstance(exc, BaseException) - else Exception(exc) - ) - else: - fut.set_result(None) - rtn[idx] = fut - else: - # schedule the next task - fut = self.submit( - run_with_retry, - next_task, - retry_policy, - configurable={ - CONFIG_KEY_SEND: partial(writer, next_task), - CONFIG_KEY_CALL: partial(call, next_task), - }, - __reraise_on_exit__=reraise, - # starting a new task in the next tick ensures - # updates from this tick are committed/streamed first - __next_tick__=True, - ) - futures[fut] = next_task - rtn[idx] = fut - return [rtn.get(i) for i in range(len(writes))] - - def call( - task: PregelExecutableTask, - func: Callable[[Any], Union[Awaitable[Any], Any]], - input: Any, - *, - retry: Optional[RetryPolicy] = None, - callbacks: Callbacks = None, - ) -> concurrent.futures.Future[Any]: - if asyncio.iscoroutinefunction(func): - raise RuntimeError("In an sync context async tasks cannot be called") - (fut,) = writer( - task, - [(PUSH, None)], - calls=[Call(func, input, retry=retry, callbacks=callbacks)], - ) - assert fut is not None, "writer did not return a future for call" - # return a chained future to ensure commit() callback is called - # before the returned future is resolved, to ensure stream order etc - return chain_future(fut, concurrent.futures.Future()) - tasks = tuple(tasks) futures = FuturesDict( callback=self.commit, event=threading.Event(), - future_type=concurrent.futures.Future, ) # give control back to the caller yield @@ -227,14 +111,7 @@ class PregelRunner: if len(tasks) == 1 and timeout is None and get_waiter is None: t = tasks[0] try: - run_with_retry( - t, - retry_policy, - configurable={ - CONFIG_KEY_SEND: partial(writer, t), - CONFIG_KEY_CALL: partial(call, t), - }, - ) + run_with_retry(t, retry_policy) self.commit(t, None) except Exception as exc: self.commit(t, exc) @@ -259,10 +136,6 @@ class PregelRunner: run_with_retry, t, retry_policy, - configurable={ - CONFIG_KEY_SEND: partial(writer, t), - CONFIG_KEY_CALL: partial(call, t), - }, __reraise_on_exit__=reraise, ) futures[fut] = t @@ -304,243 +177,12 @@ class PregelRunner: panic=reraise, ) - async def atick( - self, - tasks: Iterable[PregelExecutableTask], - *, - reraise: bool = True, - timeout: Optional[float] = None, - retry_policy: Optional[RetryPolicy] = None, - get_waiter: Optional[Callable[[], asyncio.Future[None]]] = None, - ) -> AsyncIterator[None]: - def writer( - task: PregelExecutableTask, - writes: Sequence[tuple[str, Any]], - *, - calls: Optional[Sequence[Call]] = None, - ) -> Sequence[Optional[asyncio.Future]]: - if all(w[0] != PUSH for w in writes): - return task.config[CONF][CONFIG_KEY_SEND](writes) - - # schedule PUSH tasks, collect futures - scratchpad: PregelScratchpad = task.config[CONF][CONFIG_KEY_SCRATCHPAD] - rtn: dict[int, Optional[asyncio.Future]] = {} - for idx, w in enumerate(writes): - # bail if not a PUSH write - if w[0] != PUSH: - continue - # schedule the next task, if the callback returns one - wcall = calls[idx] if calls is not None else None - if next_task := self.schedule_task( - task, scratchpad.call_counter(), wcall - ): - # if the parent task was retried, - # the next task might already be running - if fut := next( - ( - f - for f, t in futures.items() - if t is not None and t == next_task.id - ), - None, - ): - # if the parent task was retried, - # the next task might already be running - rtn[idx] = fut - elif next_task.writes: - # if it already ran, return the result - fut = asyncio.Future(loop=loop) - ret = next( - (v for c, v in next_task.writes if c == RETURN), MISSING - ) - if ret is not MISSING: - fut.set_result(ret) - elif exc := next( - (v for c, v in next_task.writes if c == ERROR), None - ): - fut.set_exception( - exc - if isinstance(exc, BaseException) - else Exception(exc) - ) - else: - fut.set_result(None) - rtn[idx] = fut - else: - # schedule the next task - fut = cast( - asyncio.Future, - self.submit( - arun_with_retry, - next_task, - retry_policy, - stream=self.use_astream, - configurable={ - CONFIG_KEY_SEND: partial(writer, next_task), - CONFIG_KEY_CALL: partial(call, next_task), - }, - __name__=t.name, - __cancel_on_exit__=True, - __reraise_on_exit__=reraise, - # starting a new task in the next tick ensures - # updates from this tick are committed/streamed first - __next_tick__=True, - ), - ) - futures[fut] = next_task - rtn[idx] = fut - return [rtn.get(i) for i in range(len(writes))] - - def call( - task: PregelExecutableTask, - func: Callable[[Any], Union[Awaitable[Any], Any]], - input: Any, - *, - retry: Optional[RetryPolicy] = None, - callbacks: Callbacks = None, - ) -> Union[asyncio.Future[Any], concurrent.futures.Future[Any]]: - (fut,) = writer( - task, - [(PUSH, None)], - calls=[Call(func, input, retry=retry, callbacks=callbacks)], - ) - assert fut is not None, "writer did not return a future for call" - # return a chained future to ensure commit() callback is called - # before the returned future is resolved, to ensure stream order etc - try: - in_async = asyncio.current_task() is not None - except RuntimeError: - in_async = False - # if in async context return an async future - # otherwise return a chained sync future - if in_async: - if isinstance(fut, asyncio.Task): - sfut: Union[asyncio.Future[Any], concurrent.futures.Future[Any]] = ( - asyncio.Future(loop=loop) - ) - loop.call_soon_threadsafe(chain_future, fut, sfut) - return sfut - else: - # already wrapped in a future - return fut - else: - sfut = concurrent.futures.Future() - loop.call_soon_threadsafe(chain_future, fut, sfut) - return sfut - - loop = asyncio.get_event_loop() - tasks = tuple(tasks) - futures = FuturesDict( - callback=self.commit, - event=asyncio.Event(), - future_type=asyncio.Future, - ) - # give control back to the caller - yield - # fast path if single task with no waiter and no timeout - if len(tasks) == 1 and get_waiter is None and timeout is None: - t = tasks[0] - try: - await arun_with_retry( - t, - retry_policy, - stream=self.use_astream, - configurable={ - CONFIG_KEY_SEND: partial(writer, t), - CONFIG_KEY_CALL: partial(call, t), - }, - ) - self.commit(t, None) - except Exception as exc: - self.commit(t, exc) - if reraise and futures: - # will be re-raised after futures are done - fut: asyncio.Future = loop.create_future() - fut.set_exception(exc) - futures.done.add(fut) - elif reraise: - raise - if not futures: # maybe `t` schuduled another task - return - else: - tasks = () # don't reschedule this task - # add waiter task if requested - if get_waiter is not None: - futures[get_waiter()] = None - # schedule tasks - for t in tasks: - if not t.writes: - fut = cast( - asyncio.Future, - self.submit( - arun_with_retry, - t, - retry_policy, - stream=self.use_astream, - configurable={ - CONFIG_KEY_SEND: partial(writer, t), - CONFIG_KEY_CALL: partial(call, t), - }, - __name__=t.name, - __cancel_on_exit__=True, - __reraise_on_exit__=reraise, - ), - ) - futures[fut] = t - # execute tasks, and wait for one to fail or all to finish. - # each task is independent from all other concurrent tasks - # yield updates/debug output as each task finishes - end_time = timeout + loop.time() if timeout else None - while len(futures) > (1 if get_waiter is not None else 0): - done, inflight = await asyncio.wait( - futures, - return_when=asyncio.FIRST_COMPLETED, - timeout=(max(0, end_time - loop.time()) if end_time else None), - ) - if not done: - break # timed out - for fut in done: - task = futures.pop(fut) - if task is None: - # waiter task finished, schedule another - if inflight and get_waiter is not None: - futures[get_waiter()] = None - else: - # remove references to loop vars - del fut, task - # maybe stop other tasks - if _should_stop_others(done): - break - # give control back to the caller - yield - # wait for done callbacks - await asyncio.wait_for( - futures.event.wait(), - timeout=(max(0, end_time - loop.time()) if end_time else None), - ) - # give control back to the caller - yield - # cancel waiter task - for fut in futures: - fut.cancel() - # panic on failure or timeout - _panic_or_proceed( - futures.done.union(f for f, t in futures.items() if t is not None), - timeout_exc_cls=asyncio.TimeoutError, - panic=reraise, - ) - def commit( self, task: PregelExecutableTask, exception: Optional[BaseException], ) -> None: - if isinstance(exception, asyncio.CancelledError): - # for cancelled tasks, also save error in task, - # so loop can finish super-step - task.writes.append((ERROR, exception)) - self.put_writes(task.id, task.writes) - elif exception: + if exception: if isinstance(exception, GraphInterrupt): # save interrupt to checkpointer if interrupts := [(INTERRUPT, i) for i in exception.args[0]]: @@ -580,27 +222,24 @@ def _should_stop_others( def _exception( - fut: Union[concurrent.futures.Future[Any], asyncio.Future[Any]], + fut: concurrent.futures.Future[Any], ) -> Optional[BaseException]: """Return the exception from a future, without raising CancelledError.""" if fut.cancelled(): - if isinstance(fut, asyncio.Future): - return asyncio.CancelledError() - else: - return concurrent.futures.CancelledError() + return concurrent.futures.CancelledError() else: return fut.exception() def _panic_or_proceed( - futs: Union[set[concurrent.futures.Future], set[asyncio.Future]], + futs: set[concurrent.futures.Future], *, timeout_exc_cls: Type[Exception] = TimeoutError, panic: bool = True, ) -> None: """Cancel remaining tasks if any failed, re-raise exception if panic is True.""" - done: set[Union[concurrent.futures.Future[Any], asyncio.Future[Any]]] = set() - inflight: set[Union[concurrent.futures.Future[Any], asyncio.Future[Any]]] = set() + done: set[concurrent.futures.Future[Any]] = set() + inflight: set[concurrent.futures.Future[Any]] = set() for fut in futs: if fut.cancelled(): continue diff --git a/libs/langgraph/langgraph/pregel/utils.py b/libs/langgraph/langgraph/pregel/utils.py index f484664a0..effcd9b8d 100644 --- a/libs/langgraph/langgraph/pregel/utils.py +++ b/libs/langgraph/langgraph/pregel/utils.py @@ -1,11 +1,8 @@ from typing import Optional -from langchain_core.runnables import RunnableLambda, RunnableSequence -from langchain_core.runnables.utils import get_function_nonlocals - from langgraph.checkpoint.base import ChannelVersions from langgraph.pregel.protocol import PregelProtocol -from langgraph.utils.runnable import Runnable, RunnableCallable, RunnableSeq +from langgraph.utils.runnable import Runnable, RunnableSeq def get_new_channel_versions( @@ -38,20 +35,7 @@ def find_subgraph_pregel(candidate: Runnable) -> Optional[PregelProtocol]: and (not isinstance(c, Pregel) or c.checkpointer is not False) ): return c - elif isinstance(c, RunnableSequence) or isinstance(c, RunnableSeq): + elif isinstance(c, RunnableSeq): candidates.extend(c.steps) - elif isinstance(c, RunnableLambda): - candidates.extend(c.deps) - elif isinstance(c, RunnableCallable): - if c.func is not None: - candidates.extend( - nl.__self__ if hasattr(nl, "__self__") else nl - for nl in get_function_nonlocals(c.func) - ) - elif c.afunc is not None: - candidates.extend( - nl.__self__ if hasattr(nl, "__self__") else nl - for nl in get_function_nonlocals(c.afunc) - ) return None diff --git a/libs/langgraph/langgraph/pregel/write.py b/libs/langgraph/langgraph/pregel/write.py index 7469824e4..f3da5e160 100644 --- a/libs/langgraph/langgraph/pregel/write.py +++ b/libs/langgraph/langgraph/pregel/write.py @@ -11,12 +11,10 @@ from typing import ( cast, ) -from langchain_core.runnables import Runnable, RunnableConfig -from langchain_core.runnables.utils import ConfigurableFieldSpec - from langgraph.constants import CONF, CONFIG_KEY_SEND, TASKS, Send from langgraph.errors import InvalidUpdateError -from langgraph.utils.runnable import RunnableCallable +from langgraph.utils.config import RunnableConfig +from langgraph.utils.runnable import Runnable, RunnableCallable TYPE_SEND = Callable[[Sequence[tuple[str, Any]]], None] R = TypeVar("R", bound=Runnable) @@ -49,40 +47,22 @@ class ChannelWrite(RunnableCallable): writes: list[Union[ChannelWriteEntry, ChannelWriteTupleEntry, Send]] """Sequence of write entries or Send objects to write.""" - require_at_least_one_of: Optional[Sequence[str]] - """If defined, at least one of these channels must be written to.""" def __init__( self, writes: Sequence[Union[ChannelWriteEntry, ChannelWriteTupleEntry, Send]], *, tags: Optional[Sequence[str]] = None, - require_at_least_one_of: Optional[Sequence[str]] = None, ): - super().__init__(func=self._write, afunc=self._awrite, name=None, tags=tags) + super().__init__(func=self._write, name=None, tags=tags) self.writes = cast( list[Union[ChannelWriteEntry, ChannelWriteTupleEntry, Send]], writes ) - self.require_at_least_one_of = require_at_least_one_of - def get_name( - self, suffix: Optional[str] = None, *, name: Optional[str] = None - ) -> str: + def get_name(self, *, name: Optional[str] = None) -> str: if not name: name = f"ChannelWrite<{','.join(w.channel if isinstance(w, ChannelWriteEntry) else '...' if isinstance(w, ChannelWriteTupleEntry) else w.node for w in self.writes)}>" - return super().get_name(suffix, name=name) - - @property - def config_specs(self) -> list[ConfigurableFieldSpec]: - return [ - ConfigurableFieldSpec( - id=CONFIG_KEY_SEND, - name=CONFIG_KEY_SEND, - description=None, - default=None, - annotation=None, - ), - ] + return super().get_name(name=name) def _write(self, input: Any, config: RunnableConfig) -> None: writes = [ @@ -96,23 +76,6 @@ class ChannelWrite(RunnableCallable): self.do_write( config, writes, - self.require_at_least_one_of if input is not None else None, - ) - return input - - async def _awrite(self, input: Any, config: RunnableConfig) -> None: - writes = [ - ChannelWriteEntry(write.channel, input, write.skip_none, write.mapper) - if isinstance(write, ChannelWriteEntry) and write.value is PASSTHROUGH - else ChannelWriteTupleEntry(write.mapper, input) - if isinstance(write, ChannelWriteTupleEntry) and write.value is PASSTHROUGH - else write - for write in self.writes - ] - self.do_write( - config, - writes, - self.require_at_least_one_of if input is not None else None, ) return input @@ -120,7 +83,6 @@ class ChannelWrite(RunnableCallable): def do_write( config: RunnableConfig, writes: Sequence[Union[ChannelWriteEntry, ChannelWriteTupleEntry, Send]], - require_at_least_one_of: Optional[Sequence[str]] = None, ) -> None: # validate for w in writes: @@ -151,28 +113,5 @@ class ChannelWrite(RunnableCallable): tuples.append((w.channel, value)) else: raise ValueError(f"Invalid write entry: {w}") - # assert required channels - if require_at_least_one_of is not None: - if not {chan for chan, _ in tuples} & set(require_at_least_one_of): - raise InvalidUpdateError( - f"Must write to at least one of {require_at_least_one_of}" - ) write: TYPE_SEND = config[CONF][CONFIG_KEY_SEND] write(tuples) - - @staticmethod - def is_writer(runnable: Runnable) -> bool: - """Used by PregelNode to distinguish between writers and other runnables.""" - return ( - isinstance(runnable, ChannelWrite) - or getattr(runnable, "_is_channel_writer", False) is True - ) - - @staticmethod - def register_writer(runnable: R) -> R: - """Used to mark a runnable as a writer, so that it can be detected by is_writer. - Instances of ChannelWrite are automatically marked as writers.""" - # using object.__setattr__ to work around objects that override __setattr__ - # eg. pydantic models and dataclasses - object.__setattr__(runnable, "_is_channel_writer", True) - return runnable diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index b52949cd0..081e42a27 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -19,24 +19,21 @@ from typing import ( get_type_hints, ) -from langchain_core.runnables import Runnable, RunnableConfig from typing_extensions import Self -from langgraph.checkpoint.base import BaseCheckpointSaver, CheckpointMetadata +from langgraph.checkpoint.base import ( + BaseCheckpointSaver, + CheckpointConfig, + CheckpointMetadata, +) +from langgraph.utils.config import RunnableConfig +from langgraph.utils.runnable import Runnable if TYPE_CHECKING: from langgraph.pregel.protocol import PregelProtocol from langgraph.store.base import BaseStore -try: - from langchain_core.messages.tool import ToolOutputMixin -except ImportError: - - class ToolOutputMixin: # type: ignore[no-redef] - pass - - All = Literal["*"] """Special value to indicate that graph should interrupt on all nodes.""" @@ -166,13 +163,13 @@ class StateSnapshot(NamedTuple): """Current values of channels""" next: tuple[str, ...] """The name of the node to execute in each task for this step.""" - config: RunnableConfig + config: CheckpointConfig """Config used to fetch this snapshot""" metadata: Optional[CheckpointMetadata] """Metadata associated with this snapshot""" created_at: Optional[str] """Timestamp of snapshot creation""" - parent_config: Optional[RunnableConfig] + parent_config: Optional[CheckpointConfig] """Config used to fetch the parent snapshot, if any""" tasks: tuple[PregelTask, ...] """Tasks to execute in this step. If already attempted, may contain an error.""" @@ -253,7 +250,7 @@ N = TypeVar("N", bound=Hashable) @dataclasses.dataclass(**_DC_KWARGS) -class Command(Generic[N], ToolOutputMixin): +class Command(Generic[N]): """One or more commands to update the graph's state and send messages to nodes. Args: @@ -461,6 +458,7 @@ def interrupt(value: Any) -> Any: Raises: GraphInterrupt: On the first invocation within the node, halts execution and surfaces the provided value to the client. """ + from langgraph.config import get_config from langgraph.constants import ( CONFIG_KEY_CHECKPOINT_NS, CONFIG_KEY_SCRATCHPAD, @@ -469,7 +467,6 @@ def interrupt(value: Any) -> Any: RESUME, ) from langgraph.errors import GraphInterrupt - from langgraph.utils.config import get_config conf = get_config()["configurable"] # track interrupt index diff --git a/libs/langgraph/langgraph/utils/config.py b/libs/langgraph/langgraph/utils/config.py index baf1bb8a9..975be1eaa 100644 --- a/libs/langgraph/langgraph/utils/config.py +++ b/libs/langgraph/langgraph/utils/config.py @@ -1,30 +1,95 @@ +import uuid from collections import ChainMap -from typing import Any, Optional, Sequence, cast +from contextvars import ContextVar +from typing import Any, Optional, Sequence, Union, cast -from langchain_core.callbacks import ( - AsyncCallbackManager, - BaseCallbackManager, - CallbackManager, - Callbacks, -) -from langchain_core.runnables import RunnableConfig -from langchain_core.runnables.config import ( - CONFIG_KEYS, - COPIABLE_KEYS, - DEFAULT_RECURSION_LIMIT, - var_child_runnable_config, +from langsmith.run_trees import RunTree, get_cached_client +from typing_extensions import TypedDict + +from langgraph.checkpoint.base import CheckpointConfig, CheckpointMetadata + + +class RunnableConfig(TypedDict, total=False): + """Configuration for a Runnable.""" + + tags: list[str] + """ + Tags for this call and any sub-calls (eg. a Chain calling an LLM). + You can use these to filter calls. + """ + + metadata: dict[str, Any] + """ + Metadata for this call and any sub-calls (eg. a Chain calling an LLM). + Keys should be strings, values should be JSON-serializable. + """ + + run_name: str + """ + Name for the tracer run for this call. Defaults to the name of the class. + """ + + max_concurrency: Optional[int] + """ + Maximum number of parallel calls to make. If not provided, defaults to + ThreadPoolExecutor's default. + """ + + recursion_limit: int + """ + Maximum number of times a call can recurse. If not provided, defaults to 25. + """ + + configurable: dict[str, Any] + """ + Runtime values for attributes previously made configurable on this Runnable, + or sub-Runnables, through .configurable_fields() or .configurable_alternatives(). + Check .output_schema() for a description of the attributes that have been made + configurable. + """ + + run_id: Optional[uuid.UUID] + """ + Unique identifier for the tracer run for this call. If not provided, a new UUID + will be generated. + """ + + run_tree: RunTree + """ + The trace tree of the caller. + """ + + +AnyConfig = Union[RunnableConfig, CheckpointConfig] + +CONFIG_KEYS = [ + "tags", + "metadata", + "run_tree", + "run_name", + "max_concurrency", + "recursion_limit", + "configurable", + "run_id", +] + +COPIABLE_KEYS = [ + "tags", + "metadata", + "run_tree", + "configurable", +] + +DEFAULT_RECURSION_LIMIT = 25 + +var_child_runnable_config = ContextVar( + "child_runnable_config", default=RunnableConfig() ) -from langgraph.checkpoint.base import CheckpointMetadata -from langgraph.config import get_config, get_store, get_stream_writer # noqa -from langgraph.constants import ( - CONF, - CONFIG_KEY_CHECKPOINT_ID, - CONFIG_KEY_CHECKPOINT_MAP, - CONFIG_KEY_CHECKPOINT_NS, - NS_END, - NS_SEP, -) + +def set_config_in_context(context: AnyConfig) -> None: + """Set the context for the current thread.""" + var_child_runnable_config.set(context) def recast_checkpoint_ns(ns: str) -> str: @@ -36,14 +101,23 @@ def recast_checkpoint_ns(ns: str) -> str: Returns: str: The checkpoint namespace without task IDs. """ + from langgraph.constants import ( + NS_END, + NS_SEP, + ) + return NS_SEP.join( part.split(NS_END)[0] for part in ns.split(NS_SEP) if not part.isdigit() ) def patch_configurable( - config: Optional[RunnableConfig], patch: dict[str, Any] + config: Optional[AnyConfig], patch: dict[str, Any] ) -> RunnableConfig: + from langgraph.constants import ( + CONF, + ) + if config is None: return {CONF: patch} elif CONF not in config: @@ -53,8 +127,16 @@ def patch_configurable( def patch_checkpoint_map( - config: Optional[RunnableConfig], metadata: Optional[CheckpointMetadata] + config: Optional[AnyConfig], + metadata: Optional[CheckpointMetadata], ) -> RunnableConfig: + from langgraph.constants import ( + CONF, + CONFIG_KEY_CHECKPOINT_ID, + CONFIG_KEY_CHECKPOINT_MAP, + CONFIG_KEY_CHECKPOINT_NS, + ) + if config is None: return config elif parents := (metadata.get("parents") if metadata else None): @@ -72,7 +154,7 @@ def patch_checkpoint_map( return config -def merge_configs(*configs: Optional[RunnableConfig]) -> RunnableConfig: +def merge_configs(*configs: Optional[AnyConfig]) -> RunnableConfig: """Merge multiple configs into one. Args: @@ -81,6 +163,10 @@ def merge_configs(*configs: Optional[RunnableConfig]) -> RunnableConfig: Returns: RunnableConfig: The merged config. """ + from langgraph.constants import ( + CONF, + ) + base: RunnableConfig = {} # Even though the keys aren't literals, this is correct # because both dicts are the same type @@ -105,35 +191,6 @@ def merge_configs(*configs: Optional[RunnableConfig]) -> RunnableConfig: base[key] = {**base_value, **value} # type: ignore[dict-item] else: base[key] = value - elif key == "callbacks": - base_callbacks = base.get("callbacks") - # callbacks can be either None, list[handler] or manager - # so merging two callbacks values has 6 cases - if isinstance(value, list): - if base_callbacks is None: - base["callbacks"] = value.copy() - elif isinstance(base_callbacks, list): - base["callbacks"] = base_callbacks + value - else: - # base_callbacks is a manager - mngr = base_callbacks.copy() - for callback in value: - mngr.add_handler(callback, inherit=True) - base["callbacks"] = mngr - elif isinstance(value, BaseCallbackManager): - # value is a manager - if base_callbacks is None: - base["callbacks"] = value.copy() - elif isinstance(base_callbacks, list): - mngr = value.copy() - for callback in base_callbacks: - mngr.add_handler(callback, inherit=True) - base["callbacks"] = mngr - else: - # base_callbacks is also a manager - base["callbacks"] = base_callbacks.merge(value) - else: - raise NotImplementedError elif key == "recursion_limit": if config["recursion_limit"] != DEFAULT_RECURSION_LIMIT: base["recursion_limit"] = config["recursion_limit"] @@ -145,9 +202,9 @@ def merge_configs(*configs: Optional[RunnableConfig]) -> RunnableConfig: def patch_config( - config: Optional[RunnableConfig], + config: Optional[AnyConfig], *, - callbacks: Callbacks = None, + runtree: Optional[RunTree] = None, recursion_limit: Optional[int] = None, max_concurrency: Optional[int] = None, run_name: Optional[str] = None, @@ -157,7 +214,7 @@ def patch_config( Args: config (Optional[RunnableConfig]): The config to patch. - callbacks (Optional[BaseCallbackManager], optional): The callbacks to set. + runtree (Optional[RunTree], optional): The runtree to set. Defaults to None. recursion_limit (Optional[int], optional): The recursion limit to set. Defaults to None. @@ -170,11 +227,15 @@ def patch_config( Returns: RunnableConfig: The patched config. """ + from langgraph.constants import ( + CONF, + ) + config = config.copy() if config is not None else {} - if callbacks is not None: + if runtree is not None: # If we're replacing callbacks, we need to unset run_name # As that should apply only to the same run as the original callbacks - config["callbacks"] = callbacks + config["run_tree"] = runtree if "run_name" in config: del config["run_name"] if "run_id" in config: @@ -190,56 +251,21 @@ def patch_config( return config -def get_callback_manager_for_config( - config: RunnableConfig, tags: Optional[Sequence[str]] = None -) -> CallbackManager: - """Get a callback manager for a config. - - Args: - config (RunnableConfig): The config. - - Returns: - CallbackManager: The callback manager. - """ - from langchain_core.callbacks.manager import CallbackManager - - # merge tags - all_tags = config.get("tags") - if all_tags is not None and tags is not None: - all_tags = [*all_tags, *tags] - elif tags is not None: - all_tags = list(tags) - # use existing callbacks if they exist - if (callbacks := config.get("callbacks")) and isinstance( - callbacks, CallbackManager - ): - if all_tags: - callbacks.add_tags(all_tags) - if metadata := config.get("metadata"): - callbacks.add_metadata(metadata) - return callbacks - else: - # otherwise create a new manager - return CallbackManager.configure( - inheritable_callbacks=config.get("callbacks"), - inheritable_tags=all_tags, - inheritable_metadata=config.get("metadata"), - ) - - -def get_async_callback_manager_for_config( - config: RunnableConfig, +def get_runtree_for_config( + config: AnyConfig, + inputs: Any, + *, + name: str, tags: Optional[Sequence[str]] = None, -) -> AsyncCallbackManager: - """Get an async callback manager for a config. +) -> RunTree: + """Get a runtree for a config. Args: config (RunnableConfig): The config. Returns: - AsyncCallbackManager: The async callback manager. + RunTree: The runtree. """ - from langchain_core.callbacks.manager import AsyncCallbackManager # merge tags all_tags = config.get("tags") @@ -248,20 +274,26 @@ def get_async_callback_manager_for_config( elif tags is not None: all_tags = list(tags) # use existing callbacks if they exist - if (callbacks := config.get("callbacks")) and isinstance( - callbacks, AsyncCallbackManager - ): - if all_tags: - callbacks.add_tags(all_tags) - if metadata := config.get("metadata"): - callbacks.add_metadata(metadata) - return callbacks + if (runtree := config.get("run_tree")) and isinstance(runtree, RunTree): + # TODO why is this needed? + if not hasattr(runtree, "ls_client"): + runtree.ls_client = get_cached_client() + return runtree.create_child( + inputs=inputs if isinstance(inputs, dict) else {"input": inputs}, + tags=all_tags, + extra={"metadata": config.get("metadata")}, + name=name, + run_id=config.get("run_id", uuid.uuid4()), + ) else: # otherwise create a new manager - return AsyncCallbackManager.configure( - inheritable_callbacks=config.get("callbacks"), - inheritable_tags=config.get("tags"), - inheritable_metadata=config.get("metadata"), + return RunTree( + id=config.get("run_id", uuid.uuid4()), + name=name, + extra={"metadata": config.get("metadata")}, + tags=all_tags, + inputs=inputs if isinstance(inputs, dict) else {"input": inputs}, + ls_client=get_cached_client(), ) @@ -272,7 +304,7 @@ def _is_not_empty(value: Any) -> bool: return value is not None -def ensure_config(*configs: Optional[RunnableConfig]) -> RunnableConfig: +def ensure_config(*configs: Optional[AnyConfig]) -> RunnableConfig: """Ensure that a config is a dict with all keys present. Args: @@ -282,6 +314,10 @@ def ensure_config(*configs: Optional[RunnableConfig]) -> RunnableConfig: Returns: RunnableConfig: The ensured config. """ + from langgraph.constants import ( + CONF, + ) + empty = RunnableConfig( tags=[], metadata=ChainMap(), diff --git a/libs/langgraph/langgraph/utils/pydantic.py b/libs/langgraph/langgraph/utils/pydantic.py deleted file mode 100644 index cd0984202..000000000 --- a/libs/langgraph/langgraph/utils/pydantic.py +++ /dev/null @@ -1,37 +0,0 @@ -from typing import Any, Dict, Optional, Union - -from pydantic import BaseModel -from pydantic.v1 import BaseModel as BaseModelV1 - - -def create_model( - model_name: str, - *, - field_definitions: Optional[Dict[str, Any]] = None, - root: Optional[Any] = None, -) -> Union[BaseModel, BaseModelV1]: - """Create a pydantic model with the given field definitions. - - Args: - model_name: The name of the model. - field_definitions: The field definitions for the model. - root: Type for a root model (RootModel) - """ - try: - # for langchain-core >= 0.3.0 - from langchain_core.utils.pydantic import create_model_v2 - - return create_model_v2( - model_name, - field_definitions=field_definitions, - root=root, - ) - except ImportError: - # for langchain-core < 0.3.0 - from langchain_core.runnables.utils import create_model - - v1_kwargs = {} - if root is not None: - v1_kwargs["__root__"] = root - - return create_model(model_name, **v1_kwargs, **(field_definitions or {})) diff --git a/libs/langgraph/langgraph/utils/queue.py b/libs/langgraph/langgraph/utils/queue.py index b68e15322..ed0408039 100644 --- a/libs/langgraph/langgraph/utils/queue.py +++ b/libs/langgraph/langgraph/utils/queue.py @@ -1,6 +1,5 @@ # type: ignore -import asyncio import queue import sys import threading @@ -12,41 +11,6 @@ from typing import Optional PY_310 = sys.version_info >= (3, 10) -class AsyncQueue(asyncio.Queue): - """Async unbounded FIFO queue with a wait() method. - - Subclassed from asyncio.Queue, adding a wait() method.""" - - async def wait(self) -> None: - """If queue is empty, wait until an item is available. - - Copied from Queue.get(), removing the call to .get_nowait(), - ie. this doesn't consume the item, just waits for it. - """ - while self.empty(): - if PY_310: - getter = self._get_loop().create_future() - else: - getter = self._loop.create_future() - self._getters.append(getter) - try: - await getter - except: - getter.cancel() # Just in case getter is not done yet. - try: - # Clean self._getters from canceled getters. - self._getters.remove(getter) - except ValueError: - # The getter could be removed from self._getters by a - # previous put_nowait call. - pass - if not self.empty() and not getter.cancelled(): - # We were woken up by put_nowait(), but can't take - # the call. Wake up the next in line. - self._wakeup_next(self._getters) - raise - - class Semaphore(threading.Semaphore): """Semaphore subclass with a wait() method.""" @@ -130,4 +94,4 @@ class SyncQueue: __class_getitem__ = classmethod(types.GenericAlias) -__all__ = ["AsyncQueue", "SyncQueue"] +__all__ = ["SyncQueue"] diff --git a/libs/langgraph/langgraph/utils/runnable.py b/libs/langgraph/langgraph/utils/runnable.py index cb8e2f447..656652f6f 100644 --- a/libs/langgraph/langgraph/utils/runnable.py +++ b/libs/langgraph/langgraph/utils/runnable.py @@ -1,193 +1,84 @@ -import asyncio import enum import inspect -import sys -from contextlib import AsyncExitStack +from abc import ABC, abstractmethod from contextvars import copy_context -from functools import partial, wraps from typing import ( Any, - AsyncIterator, - Awaitable, Callable, - Coroutine, - Iterator, + Generic, Optional, - Protocol, Sequence, - Tuple, + TypeVar, Union, cast, ) -from langchain_core.runnables.base import ( - Runnable, - RunnableConfig, - RunnableLambda, - RunnableParallel, - RunnableSequence, -) -from langchain_core.runnables.base import ( - RunnableLike as LCRunnableLike, -) -from langchain_core.runnables.config import ( - run_in_executor, - var_child_runnable_config, -) -from langchain_core.runnables.utils import Input, Output -from langchain_core.tracers._streaming import _StreamingCallbackHandler -from typing_extensions import TypeGuard - -from langgraph.constants import ( - CONF, - CONFIG_KEY_PREVIOUS, - CONFIG_KEY_STORE, - CONFIG_KEY_STREAM_WRITER, -) -from langgraph.store.base import BaseStore -from langgraph.types import StreamWriter from langgraph.utils.config import ( + AnyConfig, ensure_config, - get_async_callback_manager_for_config, - get_callback_manager_for_config, + get_runtree_for_config, patch_config, + set_config_in_context, ) -try: - from langchain_core.runnables.config import _set_config_context -except ImportError: - # For forwards compatibility - def _set_config_context(context: RunnableConfig) -> None: # type: ignore - """Set the context for the current thread.""" - var_child_runnable_config.set(context) - # Before Python 3.11 native StrEnum is not available class StrEnum(str, enum.Enum): """A string enum.""" -# Special type to denote any type is accepted -ANY_TYPE = object() - -ASYNCIO_ACCEPTS_CONTEXT = sys.version_info >= (3, 11) - -# List of keyword arguments that can be injected at runtime from the config object. -# A named argument may appear multiple times if it appears with distinct types. -KWARGS_CONFIG_KEYS: tuple[tuple[str, tuple[Any, ...], str, Any], ...] = ( - ( - sys.intern("writer"), - (StreamWriter, "StreamWriter", inspect.Parameter.empty), - CONFIG_KEY_STREAM_WRITER, - lambda _: None, - ), - ( - # Covers store that is not optional (will raise an error if a store - # cannot be injected). - sys.intern("store"), - ( - BaseStore, - "BaseStore", - inspect.Parameter.empty, - ), - CONFIG_KEY_STORE, - inspect.Parameter.empty, - ), - ( - # Covers store that is optional. Will set to None if not found in config. - sys.intern("store"), - ( - Optional[BaseStore], - # Best effort to catch some forward references. - # This will not work for cases like `"Union[None, BaseStore]"`, - # we'll need to re-write logic to use get_type_hints() - # to resolve forward references. - "Optional[BaseStore]", - ), - CONFIG_KEY_STORE, - None, - ), - ( - sys.intern("previous"), - (ANY_TYPE,), - CONFIG_KEY_PREVIOUS, - inspect.Parameter.empty, - ), -) -"""List of kwargs that can be passed to functions, and their corresponding -config keys, default values and type annotations. - -Used to configure keyword arguments that can be injected at runtime -from the config object as kwargs to `invoke`, `ainvoke`, `stream` and `astream`. - -For a keyword to be injected from the config object, the function signature -must contain a kwarg with the same name and a matching type annotation. - -Each tuple contains: -- the name of the kwarg in the function signature -- the type annotation(s) for the kwarg -- the config key to look for the value in -- the default value for the kwarg -""" - -VALID_KINDS = (inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.KEYWORD_ONLY) +Input = TypeVar("Input", contravariant=True) +Output = TypeVar("Output", covariant=True) -class _RunnableWithWriter(Protocol[Input, Output]): - def __call__(self, state: Input, *, writer: StreamWriter) -> Output: ... +class Runnable(Generic[Input, Output], ABC): + """A unit of work that can be invoked, batched, streamed, transformed and composed.""" # noqa: E501 + name: Optional[str] + """The name of the Runnable. Used for debugging and tracing.""" -class _RunnableWithStore(Protocol[Input, Output]): - def __call__(self, state: Input, *, store: BaseStore) -> Output: ... + """ --- Public API --- """ + def get_name(self, name: Optional[str] = None) -> str: + """Get the name of the Runnable. -class _RunnableWithWriterStore(Protocol[Input, Output]): - def __call__( - self, state: Input, *, writer: StreamWriter, store: BaseStore - ) -> Output: ... + Returns: + The name of the Runnable. + """ + return name or self.name or self.__class__.__name__ + @abstractmethod + def invoke( + self, input: Input, config: Optional[AnyConfig] = None, **kwargs: Any + ) -> Output: + """Transform a single input into an output. Override to implement. -class _RunnableWithConfigWriter(Protocol[Input, Output]): - def __call__( - self, state: Input, *, config: RunnableConfig, writer: StreamWriter - ) -> Output: ... + Args: + input: The input to the Runnable. + config: A config to use when invoking the Runnable. + The config supports standard keys like 'tags', 'metadata' for tracing + purposes, 'max_concurrency' for controlling how much work to do + in parallel, and other keys. Please refer to the RunnableConfig + for more details. - -class _RunnableWithConfigStore(Protocol[Input, Output]): - def __call__( - self, state: Input, *, config: RunnableConfig, store: BaseStore - ) -> Output: ... - - -class _RunnableWithConfigWriterStore(Protocol[Input, Output]): - def __call__( - self, - state: Input, - *, - config: RunnableConfig, - writer: StreamWriter, - store: BaseStore, - ) -> Output: ... + Returns: + The output of the Runnable. + """ RunnableLike = Union[ - LCRunnableLike, - _RunnableWithWriter[Input, Output], - _RunnableWithStore[Input, Output], - _RunnableWithWriterStore[Input, Output], - _RunnableWithConfigWriter[Input, Output], - _RunnableWithConfigStore[Input, Output], - _RunnableWithConfigWriterStore[Input, Output], + Runnable[Input, Output], + Callable[[Input], Output], + Callable[[Input, AnyConfig], Output], ] class RunnableCallable(Runnable): - """A much simpler version of RunnableLambda that requires sync and async functions.""" + """Wraps a callable as a Runnable.""" def __init__( self, func: Optional[Callable[..., Union[Any, Runnable]]], - afunc: Optional[Callable[..., Awaitable[Union[Any, Runnable]]]] = None, *, name: Optional[str] = None, tags: Optional[Sequence[str]] = None, @@ -204,42 +95,13 @@ class RunnableCallable(Runnable): self.name = func.__name__ except AttributeError: pass - elif afunc: - try: - self.name = afunc.__name__ - except AttributeError: - pass self.func = func - self.afunc = afunc self.tags = tags self.kwargs = kwargs self.trace = trace - self.recurse = recurse self.explode_args = explode_args - # check signature - if func is None and afunc is None: - raise ValueError("At least one of func or afunc must be provided.") - params = inspect.signature(cast(Callable, func or afunc)).parameters - - self.func_accepts_config = "config" in params - # Mapping from kwarg name to (config key, default value) to be used. - # The default value is used if the config key is not found in the config. - self.func_accepts: dict[str, Tuple[str, Any]] = {} - - for kw, typ, config_key, default in KWARGS_CONFIG_KEYS: - p = params.get(kw) - - if p is None or p.kind not in VALID_KINDS: - # If parameter is not found or is not a valid kind, skip - continue - - if typ != (ANY_TYPE,) and p.annotation not in typ: - # A specific type is required, but the function annotation does - # not match the expected type. - continue - - # If the kwarg is accepted by the function, store the default value - self.func_accepts[kw] = (config_key, default) + self.func_accepts_config = "config" in inspect.signature(func).parameters + self.recurse = recurse def __repr__(self) -> str: repr_args = { @@ -250,7 +112,7 @@ class RunnableCallable(Runnable): return f"{self.get_name()}({', '.join(f'{k}={v!r}' for k, v in repr_args.items())})" def invoke( - self, input: Any, config: Optional[RunnableConfig] = None, **kwargs: Any + self, input: Any, config: Optional[AnyConfig] = None, **kwargs: Any ) -> Any: if self.func is None: raise TypeError( @@ -268,138 +130,35 @@ class RunnableCallable(Runnable): kwargs = {**self.kwargs, **kwargs} if self.func_accepts_config: kwargs["config"] = config - _conf = config[CONF] - - for kw, (config_key, default_value) in self.func_accepts.items(): - # If the kwarg is already set, use the set value - if kw in kwargs: - continue - - if ( - # If the kwarg is requested, but isn't in the config AND has no - # default value, raise an error - config_key not in _conf and default_value is inspect.Parameter.empty - ): - raise ValueError( - f"Missing required config key '{config_key}' for '{self.name}'." - ) - - kwargs[kw] = _conf.get(config_key, default_value) context = copy_context() if self.trace: - callback_manager = get_callback_manager_for_config(config, self.tags) - run_manager = callback_manager.on_chain_start( - None, + runtree = get_runtree_for_config( + config, input, - name=config.get("run_name") or self.get_name(), - run_id=config.pop("run_id", None), + name=cast(str, config.get("run_name", self.get_name())), + tags=self.tags, ) try: - child_config = patch_config(config, callbacks=run_manager.get_child()) + child_config = patch_config(config, runtree=runtree) context = copy_context() - context.run(_set_config_context, child_config) + context.run(set_config_in_context, child_config) ret = context.run(self.func, *args, **kwargs) except BaseException as e: - run_manager.on_chain_error(e) + runtree.end(error=str(e)) raise else: - run_manager.on_chain_end(ret) + runtree.end(outputs=ret if isinstance(ret, dict) else {"output": ret}) else: - context.run(_set_config_context, config) + context.run(set_config_in_context, config) ret = context.run(self.func, *args, **kwargs) if isinstance(ret, Runnable) and self.recurse: return ret.invoke(input, config) return ret - async def ainvoke( - self, input: Any, config: Optional[RunnableConfig] = None, **kwargs: Any - ) -> Any: - if not self.afunc: - return self.invoke(input, config) - if config is None: - config = ensure_config() - if self.explode_args: - args, _kwargs = input - kwargs = {**self.kwargs, **_kwargs, **kwargs} - else: - args = (input,) - kwargs = {**self.kwargs, **kwargs} - if self.func_accepts_config: - kwargs["config"] = config - _conf = config[CONF] - for kw, (config_key, default_value) in self.func_accepts.items(): - # If the kwarg has already been set, use the set value - if kw in kwargs: - continue - - if ( - # If the kwarg is requested, but isn't in the config AND has no - # default value, raise an error - config_key not in _conf and default_value is inspect.Parameter.empty - ): - raise ValueError( - f"Missing required config key '{config_key}' for '{self.name}'." - ) - kwargs[kw] = _conf.get(config_key, default_value) - context = copy_context() - if self.trace: - callback_manager = get_async_callback_manager_for_config(config, self.tags) - run_manager = await callback_manager.on_chain_start( - None, - input, - name=config.get("run_name") or self.name, - run_id=config.pop("run_id", None), - ) - try: - child_config = patch_config(config, callbacks=run_manager.get_child()) - context.run(_set_config_context, child_config) - coro = cast(Coroutine[None, None, Any], self.afunc(*args, **kwargs)) - if ASYNCIO_ACCEPTS_CONTEXT: - ret = await asyncio.create_task(coro, context=context) - else: - ret = await coro - except BaseException as e: - await run_manager.on_chain_error(e) - raise - else: - await run_manager.on_chain_end(ret) - else: - context.run(_set_config_context, config) - if ASYNCIO_ACCEPTS_CONTEXT: - coro = cast(Coroutine[None, None, Any], self.afunc(*args, **kwargs)) - ret = await asyncio.create_task(coro, context=context) - else: - ret = await self.afunc(*args, **kwargs) - if isinstance(ret, Runnable) and self.recurse: - return await ret.ainvoke(input, config) - return ret - - -def is_async_callable( - func: Any, -) -> TypeGuard[Callable[..., Awaitable]]: - """Check if a function is async.""" - return ( - asyncio.iscoroutinefunction(func) - or hasattr(func, "__call__") - and asyncio.iscoroutinefunction(func.__call__) - ) - - -def is_async_generator( - func: Any, -) -> TypeGuard[Callable[..., AsyncIterator]]: - """Check if a function is an async generator.""" - return ( - inspect.isasyncgenfunction(func) - or hasattr(func, "__call__") - and inspect.isasyncgenfunction(func.__call__) - ) - def coerce_to_runnable( - thing: RunnableLike, *, name: Optional[str], trace: bool + thing: RunnableLike, *, name: Optional[str] = None, trace: bool = True ) -> Runnable: """Coerce a runnable-like object into a Runnable. @@ -409,22 +168,20 @@ def coerce_to_runnable( Returns: A Runnable. """ - if isinstance(thing, Runnable): + try: + from langchain_core.runnables import Runnable as LC_Runnable + except ImportError: + LC_Runnable = None + if LC_Runnable and isinstance(thing, LC_Runnable): + return thing + elif isinstance(thing, Runnable): return thing - elif is_async_generator(thing) or inspect.isgeneratorfunction(thing): - return RunnableLambda(thing, name=name) elif callable(thing): - if is_async_callable(thing): - return RunnableCallable(None, thing, name=name, trace=trace) - else: - return RunnableCallable( - thing, - wraps(thing)(partial(run_in_executor, None, thing)), # type: ignore[arg-type] - name=name, - trace=trace, - ) - elif isinstance(thing, dict): - return RunnableParallel(thing) + return RunnableCallable( + thing, + name=name, + trace=trace, + ) else: raise TypeError( f"Expected a Runnable, callable or dict." @@ -433,11 +190,7 @@ def coerce_to_runnable( class RunnableSeq(Runnable): - """Sequence of Runnables, where the output of each is the input of the next. - - RunnableSeq is a simpler version of RunnableSequence that is internal to - LangGraph. - """ + """Sequence of Runnables, where the output of each is the input of the next.""" def __init__( self, @@ -456,9 +209,7 @@ class RunnableSeq(Runnable): """ steps_flat: list[Runnable] = [] for step in steps: - if isinstance(step, RunnableSequence): - steps_flat.extend(step.steps) - elif isinstance(step, RunnableSeq): + if isinstance(step, RunnableSeq): steps_flat.extend(step.steps) else: steps_flat.append(coerce_to_runnable(step, name=None, trace=True)) @@ -470,252 +221,31 @@ class RunnableSeq(Runnable): self.name = name self.trace_inputs = trace_inputs - def __or__( - self, - other: Any, - ) -> Runnable: - if isinstance(other, RunnableSequence): - return RunnableSeq( - *self.steps, - other.first, - *other.middle, - other.last, - name=self.name or other.name, - ) - elif isinstance(other, RunnableSeq): - return RunnableSeq( - *self.steps, - *other.steps, - name=self.name or other.name, - ) - else: - return RunnableSeq( - *self.steps, - coerce_to_runnable(other, name=None, trace=True), - name=self.name, - ) - - def __ror__( - self, - other: Any, - ) -> Runnable: - if isinstance(other, RunnableSequence): - return RunnableSequence( - other.first, - *other.middle, - other.last, - *self.steps, - name=other.name or self.name, - ) - elif isinstance(other, RunnableSeq): - return RunnableSeq( - *other.steps, - *self.steps, - name=other.name or self.name, - ) - else: - return RunnableSequence( - coerce_to_runnable(other, name=None, trace=True), - *self.steps, - name=self.name, - ) - def invoke( - self, input: Input, config: Optional[RunnableConfig] = None, **kwargs: Any + self, input: Input, config: Optional[AnyConfig] = None, **kwargs: Any ) -> Any: if config is None: config = ensure_config() - # setup callbacks and context - callback_manager = get_callback_manager_for_config(config) - # start the root run - run_manager = callback_manager.on_chain_start( - None, + # setup runtree and context + runtree = get_runtree_for_config( + config, self.trace_inputs(input) if self.trace_inputs is not None else input, - name=config.get("run_name") or self.get_name(), - run_id=config.pop("run_id", None), + name=cast(str, config.get("run_name", self.get_name())), ) # invoke all steps in sequence try: for i, step in enumerate(self.steps): # mark each step as a child run - config = patch_config( - config, callbacks=run_manager.get_child(f"seq:step:{i + 1}") - ) + config = patch_config(config, runtree=runtree) if i == 0: input = step.invoke(input, config, **kwargs) else: input = step.invoke(input, config) # finish the root run except BaseException as e: - run_manager.on_chain_error(e) + runtree.end(error=str(e)) raise else: - run_manager.on_chain_end(input) + runtree.end(outputs=input if isinstance(input, dict) else {"output": input}) return input - - async def ainvoke( - self, - input: Input, - config: Optional[RunnableConfig] = None, - **kwargs: Optional[Any], - ) -> Any: - if config is None: - config = ensure_config() - # setup callbacks - callback_manager = get_async_callback_manager_for_config(config) - # start the root run - run_manager = await callback_manager.on_chain_start( - None, - self.trace_inputs(input) if self.trace_inputs is not None else input, - name=config.get("run_name") or self.get_name(), - run_id=config.pop("run_id", None), - ) - - # invoke all steps in sequence - try: - for i, step in enumerate(self.steps): - # mark each step as a child run - config = patch_config( - config, callbacks=run_manager.get_child(f"seq:step:{i + 1}") - ) - if i == 0: - input = await step.ainvoke(input, config, **kwargs) - else: - input = await step.ainvoke(input, config) - # finish the root run - except BaseException as e: - await run_manager.on_chain_error(e) - raise - else: - await run_manager.on_chain_end(input) - return input - - def stream( - self, - input: Input, - config: Optional[RunnableConfig] = None, - **kwargs: Optional[Any], - ) -> Iterator[Any]: - if config is None: - config = ensure_config() - # setup callbacks - callback_manager = get_callback_manager_for_config(config) - # start the root run - run_manager = callback_manager.on_chain_start( - None, - self.trace_inputs(input) if self.trace_inputs is not None else input, - name=config.get("run_name") or self.get_name(), - run_id=config.pop("run_id", None), - ) - - try: - # stream the last steps - # transform the input stream of each step with the next - # steps that don't natively support transforming an input stream will - # buffer input in memory until all available, and then start emitting output - for idx, step in enumerate(self.steps): - config = patch_config( - config, - callbacks=run_manager.get_child(f"seq:step:{idx + 1}"), - ) - if idx == 0: - iterator = step.stream(input, config, **kwargs) - else: - iterator = step.transform(iterator, config) - if stream_handler := next( - ( - cast(_StreamingCallbackHandler, h) - for h in run_manager.handlers - if isinstance(h, _StreamingCallbackHandler) - ), - None, - ): - # populates streamed_output in astream_log() output if needed - iterator = stream_handler.tap_output_iter(run_manager.run_id, iterator) - output: Any = None - add_supported = False - for chunk in iterator: - yield chunk - # collect final output - if output is None: - output = chunk - elif add_supported: - try: - output = output + chunk - except TypeError: - output = chunk - add_supported = False - else: - output = chunk - except BaseException as e: - run_manager.on_chain_error(e) - raise - else: - run_manager.on_chain_end(output) - - async def astream( - self, - input: Input, - config: Optional[RunnableConfig] = None, - **kwargs: Optional[Any], - ) -> AsyncIterator[Any]: - if config is None: - config = ensure_config() - # setup callbacks - callback_manager = get_async_callback_manager_for_config(config) - # start the root run - run_manager = await callback_manager.on_chain_start( - None, - self.trace_inputs(input) if self.trace_inputs is not None else input, - name=config.get("run_name") or self.get_name(), - run_id=config.pop("run_id", None), - ) - - try: - async with AsyncExitStack() as stack: - # stream the last steps - # transform the input stream of each step with the next - # steps that don't natively support transforming an input stream will - # buffer input in memory until all available, and then start emitting output - for idx, step in enumerate(self.steps): - config = patch_config( - config, - callbacks=run_manager.get_child(f"seq:step:{idx + 1}"), - ) - if idx == 0: - aiterator = step.astream(input, config, **kwargs) - else: - aiterator = step.atransform(aiterator, config) - if hasattr(aiterator, "aclose"): - stack.push_async_callback(aiterator.aclose) - if stream_handler := next( - ( - cast(_StreamingCallbackHandler, h) - for h in run_manager.handlers - if isinstance(h, _StreamingCallbackHandler) - ), - None, - ): - # populates streamed_output in astream_log() output if needed - aiterator = stream_handler.tap_output_aiter( - run_manager.run_id, aiterator - ) - output: Any = None - add_supported = False - async for chunk in aiterator: - yield chunk - # collect final output - if add_supported: - try: - output = output + chunk - except TypeError: - output = chunk - add_supported = False - else: - output = chunk - except BaseException as e: - await run_manager.on_chain_error(e) - raise - else: - await run_manager.on_chain_end(output) diff --git a/libs/langgraph/tests/__snapshots__/test_pregel.ambr b/libs/langgraph/tests/__snapshots__/test_pregel.ambr deleted file mode 100644 index 4d9622955..000000000 --- a/libs/langgraph/tests/__snapshots__/test_pregel.ambr +++ /dev/null @@ -1,1852 +0,0 @@ -# serializer version: 1 -# name: test_conditional_entrypoint_graph - '{"title": "LangGraphInput"}' -# --- -# name: test_conditional_entrypoint_graph.1 - '{"title": "LangGraphOutput"}' -# --- -# name: test_conditional_entrypoint_graph.2 - ''' - { - "nodes": [ - { - "id": "__start__", - "type": "schema", - "data": "__start__" - }, - { - "id": "left", - "type": "runnable", - "data": { - "id": [ - "langgraph", - "utils", - "runnable", - "RunnableCallable" - ], - "name": "left" - } - }, - { - "id": "right", - "type": "runnable", - "data": { - "id": [ - "langgraph", - "utils", - "runnable", - "RunnableCallable" - ], - "name": "right" - } - }, - { - "id": "__end__", - "type": "schema", - "data": "__end__" - } - ], - "edges": [ - { - "source": "right", - "target": "__end__" - }, - { - "source": "__start__", - "target": "left", - "data": "go-left", - "conditional": true - }, - { - "source": "__start__", - "target": "right", - "data": "go-right", - "conditional": true - }, - { - "source": "left", - "target": "__end__", - "conditional": true - } - ] - } - ''' -# --- -# name: test_conditional_entrypoint_graph.3 - ''' - graph TD; - right --> __end__; - __start__ -.  go-left  .-> left; - __start__ -.  go-right  .-> right; - left -.-> __end__; - - ''' -# --- -# name: test_conditional_entrypoint_graph_state - '{"properties": {"input": {"default": null, "title": "Input", "type": "string"}, "output": {"default": null, "title": "Output", "type": "string"}, "steps": {"default": null, "items": {"type": "string"}, "title": "Steps", "type": "array"}}, "title": "LangGraphInput", "type": "object"}' -# --- -# name: test_conditional_entrypoint_graph_state.1 - '{"properties": {"input": {"default": null, "title": "Input", "type": "string"}, "output": {"default": null, "title": "Output", "type": "string"}, "steps": {"default": null, "items": {"type": "string"}, "title": "Steps", "type": "array"}}, "title": "LangGraphOutput", "type": "object"}' -# --- -# name: test_conditional_entrypoint_graph_state.2 - ''' - { - "nodes": [ - { - "id": "__start__", - "type": "schema", - "data": "__start__" - }, - { - "id": "left", - "type": "runnable", - "data": { - "id": [ - "langgraph", - "utils", - "runnable", - "RunnableCallable" - ], - "name": "left" - } - }, - { - "id": "right", - "type": "runnable", - "data": { - "id": [ - "langgraph", - "utils", - "runnable", - "RunnableCallable" - ], - "name": "right" - } - }, - { - "id": "__end__", - "type": "schema", - "data": "__end__" - } - ], - "edges": [ - { - "source": "right", - "target": "__end__" - }, - { - "source": "__start__", - "target": "left", - "data": "go-left", - "conditional": true - }, - { - "source": "__start__", - "target": "right", - "data": "go-right", - "conditional": true - }, - { - "source": "left", - "target": "__end__", - "conditional": true - } - ] - } - ''' -# --- -# name: test_conditional_entrypoint_graph_state.3 - ''' - graph TD; - right --> __end__; - __start__ -.  go-left  .-> left; - __start__ -.  go-right  .-> right; - left -.-> __end__; - - ''' -# --- -# name: test_conditional_entrypoint_to_multiple_state_graph - '{"properties": {"locations": {"items": {"type": "string"}, "title": "Locations", "type": "array"}, "results": {"items": {"type": "string"}, "title": "Results", "type": "array"}}, "required": ["locations", "results"], "title": "LangGraphInput", "type": "object"}' -# --- -# name: test_conditional_entrypoint_to_multiple_state_graph.1 - '{"properties": {"locations": {"items": {"type": "string"}, "title": "Locations", "type": "array"}, "results": {"items": {"type": "string"}, "title": "Results", "type": "array"}}, "required": ["locations", "results"], "title": "LangGraphOutput", "type": "object"}' -# --- -# name: test_conditional_entrypoint_to_multiple_state_graph.2 - ''' - { - "nodes": [ - { - "id": "__start__", - "type": "schema", - "data": "__start__" - }, - { - "id": "get_weather", - "type": "runnable", - "data": { - "id": [ - "langgraph", - "utils", - "runnable", - "RunnableCallable" - ], - "name": "get_weather" - } - }, - { - "id": "__end__", - "type": "schema", - "data": "__end__" - } - ], - "edges": [ - { - "source": "get_weather", - "target": "__end__" - }, - { - "source": "__start__", - "target": "get_weather", - "conditional": true - }, - { - "source": "__start__", - "target": "__end__", - "conditional": true - } - ] - } - ''' -# --- -# name: test_conditional_entrypoint_to_multiple_state_graph.3 - ''' - graph TD; - get_weather --> __end__; - __start__ -.-> get_weather; - __start__ -.-> __end__; - - ''' -# --- -# name: test_conditional_state_graph_with_list_edge_inputs - ''' - { - "nodes": [ - { - "id": "__start__", - "type": "schema", - "data": "__start__" - }, - { - "id": "A", - "type": "runnable", - "data": { - "id": [ - "langgraph", - "utils", - "runnable", - "RunnableCallable" - ], - "name": "A" - } - }, - { - "id": "B", - "type": "runnable", - "data": { - "id": [ - "langgraph", - "utils", - "runnable", - "RunnableCallable" - ], - "name": "B" - } - }, - { - "id": "__end__", - "type": "schema", - "data": "__end__" - } - ], - "edges": [ - { - "source": "A", - "target": "__end__" - }, - { - "source": "B", - "target": "__end__" - }, - { - "source": "__start__", - "target": "A" - }, - { - "source": "__start__", - "target": "B" - } - ] - } - ''' -# --- -# name: test_conditional_state_graph_with_list_edge_inputs.1 - ''' - graph TD; - A --> __end__; - B --> __end__; - __start__ --> A; - __start__ --> B; - - ''' -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge[memory] - ''' - graph TD; - __start__ --> rewrite_query; - analyzer_one --> retriever_one; - qa --> __end__; - retriever_one --> qa; - retriever_two --> qa; - rewrite_query --> analyzer_one; - rewrite_query --> retriever_two; - - ''' -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge[postgres] - ''' - graph TD; - __start__ --> rewrite_query; - analyzer_one --> retriever_one; - qa --> __end__; - retriever_one --> qa; - retriever_two --> qa; - rewrite_query --> analyzer_one; - rewrite_query --> retriever_two; - - ''' -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge[postgres_pipe] - ''' - graph TD; - __start__ --> rewrite_query; - analyzer_one --> retriever_one; - qa --> __end__; - retriever_one --> qa; - retriever_two --> qa; - rewrite_query --> analyzer_one; - rewrite_query --> retriever_two; - - ''' -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge[postgres_pool] - ''' - graph TD; - __start__ --> rewrite_query; - analyzer_one --> retriever_one; - qa --> __end__; - retriever_one --> qa; - retriever_two --> qa; - rewrite_query --> analyzer_one; - rewrite_query --> retriever_two; - - ''' -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge[postgres_shallow] - ''' - graph TD; - __start__ --> rewrite_query; - analyzer_one --> retriever_one; - qa --> __end__; - retriever_one --> qa; - retriever_two --> qa; - rewrite_query --> analyzer_one; - rewrite_query --> retriever_two; - - ''' -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge[sqlite] - ''' - graph TD; - __start__ --> rewrite_query; - analyzer_one --> retriever_one; - qa --> __end__; - retriever_one --> qa; - retriever_two --> qa; - rewrite_query --> analyzer_one; - rewrite_query --> retriever_two; - - ''' -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[memory] - ''' - graph TD; - __start__ --> rewrite_query; - analyzer_one --> retriever_one; - qa --> __end__; - retriever_one --> qa; - retriever_two --> qa; - rewrite_query --> analyzer_one; - rewrite_query -.-> retriever_two; - - ''' -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[memory].1 - dict({ - 'definitions': dict({ - 'InnerObject': dict({ - 'properties': dict({ - 'yo': dict({ - 'title': 'Yo', - 'type': 'integer', - }), - }), - 'required': list([ - 'yo', - ]), - 'title': 'InnerObject', - 'type': 'object', - }), - }), - 'properties': dict({ - 'inner': dict({ - '$ref': '#/definitions/InnerObject', - }), - 'query': dict({ - 'title': 'Query', - 'type': 'string', - }), - }), - 'required': list([ - 'query', - 'inner', - ]), - 'title': 'Input', - 'type': 'object', - }) -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[memory].2 - dict({ - 'properties': dict({ - 'answer': dict({ - 'title': 'Answer', - 'type': 'string', - }), - 'docs': dict({ - 'items': dict({ - 'type': 'string', - }), - 'title': 'Docs', - 'type': 'array', - }), - }), - 'required': list([ - 'answer', - 'docs', - ]), - 'title': 'Output', - 'type': 'object', - }) -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[postgres] - ''' - graph TD; - __start__ --> rewrite_query; - analyzer_one --> retriever_one; - qa --> __end__; - retriever_one --> qa; - retriever_two --> qa; - rewrite_query --> analyzer_one; - rewrite_query -.-> retriever_two; - - ''' -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[postgres].1 - dict({ - 'definitions': dict({ - 'InnerObject': dict({ - 'properties': dict({ - 'yo': dict({ - 'title': 'Yo', - 'type': 'integer', - }), - }), - 'required': list([ - 'yo', - ]), - 'title': 'InnerObject', - 'type': 'object', - }), - }), - 'properties': dict({ - 'inner': dict({ - '$ref': '#/definitions/InnerObject', - }), - 'query': dict({ - 'title': 'Query', - 'type': 'string', - }), - }), - 'required': list([ - 'query', - 'inner', - ]), - 'title': 'Input', - 'type': 'object', - }) -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[postgres].2 - dict({ - 'properties': dict({ - 'answer': dict({ - 'title': 'Answer', - 'type': 'string', - }), - 'docs': dict({ - 'items': dict({ - 'type': 'string', - }), - 'title': 'Docs', - 'type': 'array', - }), - }), - 'required': list([ - 'answer', - 'docs', - ]), - 'title': 'Output', - 'type': 'object', - }) -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[postgres_pipe] - ''' - graph TD; - __start__ --> rewrite_query; - analyzer_one --> retriever_one; - qa --> __end__; - retriever_one --> qa; - retriever_two --> qa; - rewrite_query --> analyzer_one; - rewrite_query -.-> retriever_two; - - ''' -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[postgres_pipe].1 - dict({ - 'definitions': dict({ - 'InnerObject': dict({ - 'properties': dict({ - 'yo': dict({ - 'title': 'Yo', - 'type': 'integer', - }), - }), - 'required': list([ - 'yo', - ]), - 'title': 'InnerObject', - 'type': 'object', - }), - }), - 'properties': dict({ - 'inner': dict({ - '$ref': '#/definitions/InnerObject', - }), - 'query': dict({ - 'title': 'Query', - 'type': 'string', - }), - }), - 'required': list([ - 'query', - 'inner', - ]), - 'title': 'Input', - 'type': 'object', - }) -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[postgres_pipe].2 - dict({ - 'properties': dict({ - 'answer': dict({ - 'title': 'Answer', - 'type': 'string', - }), - 'docs': dict({ - 'items': dict({ - 'type': 'string', - }), - 'title': 'Docs', - 'type': 'array', - }), - }), - 'required': list([ - 'answer', - 'docs', - ]), - 'title': 'Output', - 'type': 'object', - }) -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[postgres_pool] - ''' - graph TD; - __start__ --> rewrite_query; - analyzer_one --> retriever_one; - qa --> __end__; - retriever_one --> qa; - retriever_two --> qa; - rewrite_query --> analyzer_one; - rewrite_query -.-> retriever_two; - - ''' -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[postgres_pool].1 - dict({ - 'definitions': dict({ - 'InnerObject': dict({ - 'properties': dict({ - 'yo': dict({ - 'title': 'Yo', - 'type': 'integer', - }), - }), - 'required': list([ - 'yo', - ]), - 'title': 'InnerObject', - 'type': 'object', - }), - }), - 'properties': dict({ - 'inner': dict({ - '$ref': '#/definitions/InnerObject', - }), - 'query': dict({ - 'title': 'Query', - 'type': 'string', - }), - }), - 'required': list([ - 'query', - 'inner', - ]), - 'title': 'Input', - 'type': 'object', - }) -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[postgres_pool].2 - dict({ - 'properties': dict({ - 'answer': dict({ - 'title': 'Answer', - 'type': 'string', - }), - 'docs': dict({ - 'items': dict({ - 'type': 'string', - }), - 'title': 'Docs', - 'type': 'array', - }), - }), - 'required': list([ - 'answer', - 'docs', - ]), - 'title': 'Output', - 'type': 'object', - }) -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[postgres_shallow] - ''' - graph TD; - __start__ --> rewrite_query; - analyzer_one --> retriever_one; - qa --> __end__; - retriever_one --> qa; - retriever_two --> qa; - rewrite_query --> analyzer_one; - rewrite_query -.-> retriever_two; - - ''' -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[postgres_shallow].1 - dict({ - 'definitions': dict({ - 'InnerObject': dict({ - 'properties': dict({ - 'yo': dict({ - 'title': 'Yo', - 'type': 'integer', - }), - }), - 'required': list([ - 'yo', - ]), - 'title': 'InnerObject', - 'type': 'object', - }), - }), - 'properties': dict({ - 'inner': dict({ - '$ref': '#/definitions/InnerObject', - }), - 'query': dict({ - 'title': 'Query', - 'type': 'string', - }), - }), - 'required': list([ - 'query', - 'inner', - ]), - 'title': 'Input', - 'type': 'object', - }) -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[postgres_shallow].2 - dict({ - 'properties': dict({ - 'answer': dict({ - 'title': 'Answer', - 'type': 'string', - }), - 'docs': dict({ - 'items': dict({ - 'type': 'string', - }), - 'title': 'Docs', - 'type': 'array', - }), - }), - 'required': list([ - 'answer', - 'docs', - ]), - 'title': 'Output', - 'type': 'object', - }) -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[sqlite] - ''' - graph TD; - __start__ --> rewrite_query; - analyzer_one --> retriever_one; - qa --> __end__; - retriever_one --> qa; - retriever_two --> qa; - rewrite_query --> analyzer_one; - rewrite_query -.-> retriever_two; - - ''' -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[sqlite].1 - dict({ - 'definitions': dict({ - 'InnerObject': dict({ - 'properties': dict({ - 'yo': dict({ - 'title': 'Yo', - 'type': 'integer', - }), - }), - 'required': list([ - 'yo', - ]), - 'title': 'InnerObject', - 'type': 'object', - }), - }), - 'properties': dict({ - 'inner': dict({ - '$ref': '#/definitions/InnerObject', - }), - 'query': dict({ - 'title': 'Query', - 'type': 'string', - }), - }), - 'required': list([ - 'query', - 'inner', - ]), - 'title': 'Input', - 'type': 'object', - }) -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[sqlite].2 - dict({ - 'properties': dict({ - 'answer': dict({ - 'title': 'Answer', - 'type': 'string', - }), - 'docs': dict({ - 'items': dict({ - 'type': 'string', - }), - 'title': 'Docs', - 'type': 'array', - }), - }), - 'required': list([ - 'answer', - 'docs', - ]), - 'title': 'Output', - 'type': 'object', - }) -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[memory] - ''' - graph TD; - __start__ --> rewrite_query; - analyzer_one --> retriever_one; - qa --> __end__; - retriever_one --> qa; - retriever_two --> qa; - rewrite_query --> analyzer_one; - rewrite_query -.-> retriever_two; - - ''' -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[memory].1 - dict({ - '$defs': dict({ - 'InnerObject': dict({ - 'properties': dict({ - 'yo': dict({ - 'title': 'Yo', - 'type': 'integer', - }), - }), - 'required': list([ - 'yo', - ]), - 'title': 'InnerObject', - 'type': 'object', - }), - }), - 'properties': dict({ - 'inner': dict({ - '$ref': '#/$defs/InnerObject', - }), - 'query': dict({ - 'title': 'Query', - 'type': 'string', - }), - }), - 'required': list([ - 'query', - 'inner', - ]), - 'title': 'Input', - 'type': 'object', - }) -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[memory].2 - dict({ - 'properties': dict({ - 'answer': dict({ - 'title': 'Answer', - 'type': 'string', - }), - 'docs': dict({ - 'items': dict({ - 'type': 'string', - }), - 'title': 'Docs', - 'type': 'array', - }), - }), - 'required': list([ - 'answer', - 'docs', - ]), - 'title': 'Output', - 'type': 'object', - }) -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres] - ''' - graph TD; - __start__ --> rewrite_query; - analyzer_one --> retriever_one; - qa --> __end__; - retriever_one --> qa; - retriever_two --> qa; - rewrite_query --> analyzer_one; - rewrite_query -.-> retriever_two; - - ''' -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres].1 - dict({ - '$defs': dict({ - 'InnerObject': dict({ - 'properties': dict({ - 'yo': dict({ - 'title': 'Yo', - 'type': 'integer', - }), - }), - 'required': list([ - 'yo', - ]), - 'title': 'InnerObject', - 'type': 'object', - }), - }), - 'properties': dict({ - 'inner': dict({ - '$ref': '#/$defs/InnerObject', - }), - 'query': dict({ - 'title': 'Query', - 'type': 'string', - }), - }), - 'required': list([ - 'query', - 'inner', - ]), - 'title': 'Input', - 'type': 'object', - }) -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres].2 - dict({ - 'properties': dict({ - 'answer': dict({ - 'title': 'Answer', - 'type': 'string', - }), - 'docs': dict({ - 'items': dict({ - 'type': 'string', - }), - 'title': 'Docs', - 'type': 'array', - }), - }), - 'required': list([ - 'answer', - 'docs', - ]), - 'title': 'Output', - 'type': 'object', - }) -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_pipe] - ''' - graph TD; - __start__ --> rewrite_query; - analyzer_one --> retriever_one; - qa --> __end__; - retriever_one --> qa; - retriever_two --> qa; - rewrite_query --> analyzer_one; - rewrite_query -.-> retriever_two; - - ''' -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_pipe].1 - dict({ - '$defs': dict({ - 'InnerObject': dict({ - 'properties': dict({ - 'yo': dict({ - 'title': 'Yo', - 'type': 'integer', - }), - }), - 'required': list([ - 'yo', - ]), - 'title': 'InnerObject', - 'type': 'object', - }), - }), - 'properties': dict({ - 'inner': dict({ - '$ref': '#/$defs/InnerObject', - }), - 'query': dict({ - 'title': 'Query', - 'type': 'string', - }), - }), - 'required': list([ - 'query', - 'inner', - ]), - 'title': 'Input', - 'type': 'object', - }) -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_pipe].2 - dict({ - 'properties': dict({ - 'answer': dict({ - 'title': 'Answer', - 'type': 'string', - }), - 'docs': dict({ - 'items': dict({ - 'type': 'string', - }), - 'title': 'Docs', - 'type': 'array', - }), - }), - 'required': list([ - 'answer', - 'docs', - ]), - 'title': 'Output', - 'type': 'object', - }) -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_pool] - ''' - graph TD; - __start__ --> rewrite_query; - analyzer_one --> retriever_one; - qa --> __end__; - retriever_one --> qa; - retriever_two --> qa; - rewrite_query --> analyzer_one; - rewrite_query -.-> retriever_two; - - ''' -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_pool].1 - dict({ - '$defs': dict({ - 'InnerObject': dict({ - 'properties': dict({ - 'yo': dict({ - 'title': 'Yo', - 'type': 'integer', - }), - }), - 'required': list([ - 'yo', - ]), - 'title': 'InnerObject', - 'type': 'object', - }), - }), - 'properties': dict({ - 'inner': dict({ - '$ref': '#/$defs/InnerObject', - }), - 'query': dict({ - 'title': 'Query', - 'type': 'string', - }), - }), - 'required': list([ - 'query', - 'inner', - ]), - 'title': 'Input', - 'type': 'object', - }) -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_pool].2 - dict({ - 'properties': dict({ - 'answer': dict({ - 'title': 'Answer', - 'type': 'string', - }), - 'docs': dict({ - 'items': dict({ - 'type': 'string', - }), - 'title': 'Docs', - 'type': 'array', - }), - }), - 'required': list([ - 'answer', - 'docs', - ]), - 'title': 'Output', - 'type': 'object', - }) -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_shallow] - ''' - graph TD; - __start__ --> rewrite_query; - analyzer_one --> retriever_one; - qa --> __end__; - retriever_one --> qa; - retriever_two --> qa; - rewrite_query --> analyzer_one; - rewrite_query -.-> retriever_two; - - ''' -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_shallow].1 - dict({ - '$defs': dict({ - 'InnerObject': dict({ - 'properties': dict({ - 'yo': dict({ - 'title': 'Yo', - 'type': 'integer', - }), - }), - 'required': list([ - 'yo', - ]), - 'title': 'InnerObject', - 'type': 'object', - }), - }), - 'properties': dict({ - 'inner': dict({ - '$ref': '#/$defs/InnerObject', - }), - 'query': dict({ - 'title': 'Query', - 'type': 'string', - }), - }), - 'required': list([ - 'query', - 'inner', - ]), - 'title': 'Input', - 'type': 'object', - }) -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_shallow].2 - dict({ - 'properties': dict({ - 'answer': dict({ - 'title': 'Answer', - 'type': 'string', - }), - 'docs': dict({ - 'items': dict({ - 'type': 'string', - }), - 'title': 'Docs', - 'type': 'array', - }), - }), - 'required': list([ - 'answer', - 'docs', - ]), - 'title': 'Output', - 'type': 'object', - }) -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[sqlite] - ''' - graph TD; - __start__ --> rewrite_query; - analyzer_one --> retriever_one; - qa --> __end__; - retriever_one --> qa; - retriever_two --> qa; - rewrite_query --> analyzer_one; - rewrite_query -.-> retriever_two; - - ''' -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[sqlite].1 - dict({ - '$defs': dict({ - 'InnerObject': dict({ - 'properties': dict({ - 'yo': dict({ - 'title': 'Yo', - 'type': 'integer', - }), - }), - 'required': list([ - 'yo', - ]), - 'title': 'InnerObject', - 'type': 'object', - }), - }), - 'properties': dict({ - 'inner': dict({ - '$ref': '#/$defs/InnerObject', - }), - 'query': dict({ - 'title': 'Query', - 'type': 'string', - }), - }), - 'required': list([ - 'query', - 'inner', - ]), - 'title': 'Input', - 'type': 'object', - }) -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[sqlite].2 - dict({ - 'properties': dict({ - 'answer': dict({ - 'title': 'Answer', - 'type': 'string', - }), - 'docs': dict({ - 'items': dict({ - 'type': 'string', - }), - 'title': 'Docs', - 'type': 'array', - }), - }), - 'required': list([ - 'answer', - 'docs', - ]), - 'title': 'Output', - 'type': 'object', - }) -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[memory] - ''' - graph TD; - __start__ --> rewrite_query; - analyzer_one --> retriever_one; - qa --> __end__; - retriever_one --> qa; - retriever_two --> qa; - rewrite_query --> analyzer_one; - rewrite_query -.-> retriever_two; - - ''' -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[postgres] - ''' - graph TD; - __start__ --> rewrite_query; - analyzer_one --> retriever_one; - qa --> __end__; - retriever_one --> qa; - retriever_two --> qa; - rewrite_query --> analyzer_one; - rewrite_query -.-> retriever_two; - - ''' -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[postgres_pipe] - ''' - graph TD; - __start__ --> rewrite_query; - analyzer_one --> retriever_one; - qa --> __end__; - retriever_one --> qa; - retriever_two --> qa; - rewrite_query --> analyzer_one; - rewrite_query -.-> retriever_two; - - ''' -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[postgres_pool] - ''' - graph TD; - __start__ --> rewrite_query; - analyzer_one --> retriever_one; - qa --> __end__; - retriever_one --> qa; - retriever_two --> qa; - rewrite_query --> analyzer_one; - rewrite_query -.-> retriever_two; - - ''' -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[postgres_shallow] - ''' - graph TD; - __start__ --> rewrite_query; - analyzer_one --> retriever_one; - qa --> __end__; - retriever_one --> qa; - retriever_two --> qa; - rewrite_query --> analyzer_one; - rewrite_query -.-> retriever_two; - - ''' -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[sqlite] - ''' - graph TD; - __start__ --> rewrite_query; - analyzer_one --> retriever_one; - qa --> __end__; - retriever_one --> qa; - retriever_two --> qa; - rewrite_query --> analyzer_one; - rewrite_query -.-> retriever_two; - - ''' -# --- -# name: test_multiple_sinks_subgraphs - ''' - %%{init: {'flowchart': {'curve': 'linear'}}}%% - graph TD; - __start__([

__start__

]):::first - uno(uno) - dos(dos) - subgraph_one(one) - subgraph_two(two) - subgraph_three(three) - __start__ --> uno; - uno -.-> dos; - uno -.-> subgraph_one; - subgraph subgraph - subgraph_one -.-> subgraph_two; - subgraph_one -.-> subgraph_three; - end - classDef default fill:#f2f0ff,line-height:1.2 - classDef first fill-opacity:0 - classDef last fill:#bfb6fc - - ''' -# --- -# name: test_nested_graph - ''' - graph TD; - __start__ --> inner; - inner --> side; - side --> __end__; - - ''' -# --- -# name: test_nested_graph.1 - ''' - %%{init: {'flowchart': {'curve': 'linear'}}}%% - graph TD; - __start__([

__start__

]):::first - inner(inner) - side(side) - __end__([

__end__

]):::last - __start__ --> inner; - inner --> side; - side --> __end__; - classDef default fill:#f2f0ff,line-height:1.2 - classDef first fill-opacity:0 - classDef last fill:#bfb6fc - - ''' -# --- -# name: test_nested_graph_xray - dict({ - 'edges': list([ - dict({ - 'conditional': True, - 'source': 'tool_two:__start__', - 'target': 'tool_two:tool_two_slow', - }), - dict({ - 'source': 'tool_two:tool_two_slow', - 'target': 'tool_two:__end__', - }), - dict({ - 'conditional': True, - 'source': 'tool_two:__start__', - 'target': 'tool_two:tool_two_fast', - }), - dict({ - 'source': 'tool_two:tool_two_fast', - 'target': 'tool_two:__end__', - }), - dict({ - 'conditional': True, - 'source': '__start__', - 'target': 'tool_one', - }), - dict({ - 'source': 'tool_one', - 'target': '__end__', - }), - dict({ - 'conditional': True, - 'source': '__start__', - 'target': 'tool_two:__start__', - }), - dict({ - 'source': 'tool_two:__end__', - 'target': '__end__', - }), - dict({ - 'conditional': True, - 'source': '__start__', - 'target': 'tool_three', - }), - dict({ - 'source': 'tool_three', - 'target': '__end__', - }), - ]), - 'nodes': list([ - dict({ - 'data': '__start__', - 'id': '__start__', - 'type': 'schema', - }), - dict({ - 'data': dict({ - 'id': list([ - 'langgraph', - 'utils', - 'runnable', - 'RunnableCallable', - ]), - 'name': 'tool_one', - }), - 'id': 'tool_one', - 'type': 'runnable', - }), - dict({ - 'data': 'tool_two:__start__', - 'id': 'tool_two:__start__', - 'type': 'schema', - }), - dict({ - 'data': dict({ - 'id': list([ - 'langgraph', - 'utils', - 'runnable', - 'RunnableCallable', - ]), - 'name': 'tool_two:tool_two_slow', - }), - 'id': 'tool_two:tool_two_slow', - 'type': 'runnable', - }), - dict({ - 'data': dict({ - 'id': list([ - 'langgraph', - 'utils', - 'runnable', - 'RunnableCallable', - ]), - 'name': 'tool_two:tool_two_fast', - }), - 'id': 'tool_two:tool_two_fast', - 'type': 'runnable', - }), - dict({ - 'data': 'tool_two:__end__', - 'id': 'tool_two:__end__', - 'type': 'schema', - }), - dict({ - 'data': dict({ - 'id': list([ - 'langgraph', - 'utils', - 'runnable', - 'RunnableCallable', - ]), - 'name': 'tool_three', - }), - 'id': 'tool_three', - 'type': 'runnable', - }), - dict({ - 'data': '__end__', - 'id': '__end__', - 'type': 'schema', - }), - ]), - }) -# --- -# name: test_nested_graph_xray.1 - ''' - %%{init: {'flowchart': {'curve': 'linear'}}}%% - graph TD; - __start__([

__start__

]):::first - tool_one(tool_one) - tool_two___start__(

__start__

) - tool_two_tool_two_slow(tool_two_slow) - tool_two_tool_two_fast(tool_two_fast) - tool_two___end__(

__end__

) - tool_three(tool_three) - __end__([

__end__

]):::last - __start__ -.-> tool_one; - tool_one --> __end__; - __start__ -.-> tool_two___start__; - tool_two___end__ --> __end__; - __start__ -.-> tool_three; - tool_three --> __end__; - subgraph tool_two - tool_two___start__ -.-> tool_two_tool_two_slow; - tool_two_tool_two_slow --> tool_two___end__; - tool_two___start__ -.-> tool_two_tool_two_fast; - tool_two_tool_two_fast --> tool_two___end__; - end - classDef default fill:#f2f0ff,line-height:1.2 - classDef first fill-opacity:0 - classDef last fill:#bfb6fc - - ''' -# --- -# name: test_repeat_condition - ''' - graph TD; - __start__ --> Researcher; - Researcher -.  continue  .-> Chart_Generator; - Researcher -.  call_tool  .-> Call_Tool; - Researcher -.  end  .-> __end__; - Chart_Generator -.  continue  .-> Researcher; - Chart_Generator -.  call_tool  .-> Call_Tool; - Chart_Generator -.  end  .-> __end__; - Call_Tool -.-> Researcher; - Call_Tool -.-> Chart_Generator; - Researcher -.  redo  .-> Researcher; - - ''' -# --- -# name: test_simple_multi_edge - ''' - graph TD; - __start__ --> up; - down --> __end__; - side --> down; - up --> down; - up --> other; - up --> side; - - ''' -# --- -# name: test_state_graph_w_config_inherited_state_keys - '{"$defs": {"Configurable": {"properties": {"tools": {"default": null, "items": {"type": "string"}, "title": "Tools", "type": "array"}}, "title": "Configurable", "type": "object"}}, "properties": {"configurable": {"$ref": "#/$defs/Configurable", "default": null}}, "title": "LangGraphConfig", "type": "object"}' -# --- -# name: test_state_graph_w_config_inherited_state_keys.1 - '{"$defs": {"AgentAction": {"description": "Represents a request to execute an action by an agent.\\n\\nThe action consists of the name of the tool to execute and the input to pass\\nto the tool. The log is used to pass along extra information about the action.", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"anyOf": [{"type": "string"}, {"type": "object"}], "title": "Tool Input"}, "log": {"title": "Log", "type": "string"}, "type": {"const": "AgentAction", "default": "AgentAction", "enum": ["AgentAction"], "title": "Type", "type": "string"}}, "required": ["tool", "tool_input", "log"], "title": "AgentAction", "type": "object"}, "AgentFinish": {"description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"const": "AgentFinish", "default": "AgentFinish", "enum": ["AgentFinish"], "title": "Type", "type": "string"}}, "required": ["return_values", "log"], "title": "AgentFinish", "type": "object"}}, "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"anyOf": [{"$ref": "#/$defs/AgentAction"}, {"$ref": "#/$defs/AgentFinish"}, {"type": "null"}], "default": null, "title": "Agent Outcome"}, "intermediate_steps": {"default": null, "items": {"maxItems": 2, "minItems": 2, "prefixItems": [{"$ref": "#/$defs/AgentAction"}, {"type": "string"}], "type": "array"}, "title": "Intermediate Steps", "type": "array"}}, "required": ["input"], "title": "LangGraphInput", "type": "object"}' -# --- -# name: test_state_graph_w_config_inherited_state_keys.2 - '{"$defs": {"AgentAction": {"description": "Represents a request to execute an action by an agent.\\n\\nThe action consists of the name of the tool to execute and the input to pass\\nto the tool. The log is used to pass along extra information about the action.", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"anyOf": [{"type": "string"}, {"type": "object"}], "title": "Tool Input"}, "log": {"title": "Log", "type": "string"}, "type": {"const": "AgentAction", "default": "AgentAction", "enum": ["AgentAction"], "title": "Type", "type": "string"}}, "required": ["tool", "tool_input", "log"], "title": "AgentAction", "type": "object"}, "AgentFinish": {"description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"const": "AgentFinish", "default": "AgentFinish", "enum": ["AgentFinish"], "title": "Type", "type": "string"}}, "required": ["return_values", "log"], "title": "AgentFinish", "type": "object"}}, "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"anyOf": [{"$ref": "#/$defs/AgentAction"}, {"$ref": "#/$defs/AgentFinish"}, {"type": "null"}], "default": null, "title": "Agent Outcome"}, "intermediate_steps": {"default": null, "items": {"maxItems": 2, "minItems": 2, "prefixItems": [{"$ref": "#/$defs/AgentAction"}, {"type": "string"}], "type": "array"}, "title": "Intermediate Steps", "type": "array"}}, "required": ["input"], "title": "LangGraphOutput", "type": "object"}' -# --- -# name: test_xray_bool - ''' - %%{init: {'flowchart': {'curve': 'linear'}}}%% - graph TD; - __start__([

__start__

]):::first - gp_one(gp_one) - gp_two___start__(

__start__

) - gp_two_p_one(p_one) - gp_two_p_two___start__(

__start__

) - gp_two_p_two_c_one(c_one) - gp_two_p_two_c_two(c_two) - gp_two_p_two___end__(

__end__

) - gp_two___end__(

__end__

) - __end__([

__end__

]):::last - __start__ --> gp_one; - gp_two___end__ --> gp_one; - gp_one -.  0  .-> gp_two___start__; - gp_one -.  1  .-> __end__; - subgraph gp_two - gp_two___start__ --> gp_two_p_one; - gp_two_p_two___end__ --> gp_two_p_one; - gp_two_p_one -.  0  .-> gp_two_p_two___start__; - gp_two_p_one -.  1  .-> gp_two___end__; - subgraph p_two - gp_two_p_two___start__ --> gp_two_p_two_c_one; - gp_two_p_two_c_two --> gp_two_p_two_c_one; - gp_two_p_two_c_one -.  0  .-> gp_two_p_two_c_two; - gp_two_p_two_c_one -.  1  .-> gp_two_p_two___end__; - end - end - classDef default fill:#f2f0ff,line-height:1.2 - classDef first fill-opacity:0 - classDef last fill:#bfb6fc - - ''' -# --- -# name: test_xray_issue - ''' - %%{init: {'flowchart': {'curve': 'linear'}}}%% - graph TD; - __start__([

__start__

]):::first - p_one(p_one) - p_two___start__(

__start__

) - p_two_c_one(c_one) - p_two_c_two(c_two) - p_two___end__(

__end__

) - __end__([

__end__

]):::last - __start__ --> p_one; - p_two___end__ --> p_one; - p_one -.  0  .-> p_two___start__; - p_one -.  1  .-> __end__; - subgraph p_two - p_two___start__ --> p_two_c_one; - p_two_c_two --> p_two_c_one; - p_two_c_one -.  0  .-> p_two_c_two; - p_two_c_one -.  1  .-> p_two___end__; - end - classDef default fill:#f2f0ff,line-height:1.2 - classDef first fill-opacity:0 - classDef last fill:#bfb6fc - - ''' -# --- -# name: test_xray_lance - dict({ - 'edges': list([ - dict({ - 'source': '__start__', - 'target': 'ask_question', - }), - dict({ - 'source': 'ask_question', - 'target': 'answer_question', - }), - dict({ - 'conditional': True, - 'source': 'answer_question', - 'target': 'ask_question', - }), - dict({ - 'conditional': True, - 'source': 'answer_question', - 'target': '__end__', - }), - ]), - 'nodes': list([ - dict({ - 'data': '__start__', - 'id': '__start__', - 'type': 'schema', - }), - dict({ - 'data': dict({ - 'id': list([ - 'langgraph', - 'utils', - 'runnable', - 'RunnableCallable', - ]), - 'name': 'ask_question', - }), - 'id': 'ask_question', - 'type': 'runnable', - }), - dict({ - 'data': dict({ - 'id': list([ - 'langgraph', - 'utils', - 'runnable', - 'RunnableCallable', - ]), - 'name': 'answer_question', - }), - 'id': 'answer_question', - 'type': 'runnable', - }), - dict({ - 'data': '__end__', - 'id': '__end__', - 'type': 'schema', - }), - ]), - }) -# --- -# name: test_xray_lance.1 - dict({ - 'edges': list([ - dict({ - 'source': '__start__', - 'target': 'generate_analysts', - }), - dict({ - 'source': 'conduct_interview', - 'target': 'generate_sections', - }), - dict({ - 'source': 'generate_sections', - 'target': '__end__', - }), - dict({ - 'conditional': True, - 'source': 'generate_analysts', - 'target': 'conduct_interview', - }), - ]), - 'nodes': list([ - dict({ - 'data': '__start__', - 'id': '__start__', - 'type': 'schema', - }), - dict({ - 'data': dict({ - 'id': list([ - 'langgraph', - 'utils', - 'runnable', - 'RunnableCallable', - ]), - 'name': 'generate_analysts', - }), - 'id': 'generate_analysts', - 'type': 'runnable', - }), - dict({ - 'data': dict({ - 'id': list([ - 'langgraph', - 'graph', - 'state', - 'CompiledStateGraph', - ]), - 'name': 'conduct_interview', - }), - 'id': 'conduct_interview', - 'type': 'runnable', - }), - dict({ - 'data': dict({ - 'id': list([ - 'langgraph', - 'utils', - 'runnable', - 'RunnableCallable', - ]), - 'name': 'generate_sections', - }), - 'id': 'generate_sections', - 'type': 'runnable', - }), - dict({ - 'data': '__end__', - 'id': '__end__', - 'type': 'schema', - }), - ]), - }) -# --- -# name: test_xray_lance.2 - dict({ - 'edges': list([ - dict({ - 'source': 'conduct_interview:__start__', - 'target': 'conduct_interview:ask_question', - }), - dict({ - 'source': 'conduct_interview:ask_question', - 'target': 'conduct_interview:answer_question', - }), - dict({ - 'conditional': True, - 'source': 'conduct_interview:answer_question', - 'target': 'conduct_interview:ask_question', - }), - dict({ - 'conditional': True, - 'source': 'conduct_interview:answer_question', - 'target': 'conduct_interview:__end__', - }), - dict({ - 'source': '__start__', - 'target': 'generate_analysts', - }), - dict({ - 'source': 'conduct_interview:__end__', - 'target': 'generate_sections', - }), - dict({ - 'source': 'generate_sections', - 'target': '__end__', - }), - dict({ - 'conditional': True, - 'source': 'generate_analysts', - 'target': 'conduct_interview:__start__', - }), - ]), - 'nodes': list([ - dict({ - 'data': '__start__', - 'id': '__start__', - 'type': 'schema', - }), - dict({ - 'data': dict({ - 'id': list([ - 'langgraph', - 'utils', - 'runnable', - 'RunnableCallable', - ]), - 'name': 'generate_analysts', - }), - 'id': 'generate_analysts', - 'type': 'runnable', - }), - dict({ - 'data': 'conduct_interview:__start__', - 'id': 'conduct_interview:__start__', - 'type': 'schema', - }), - dict({ - 'data': dict({ - 'id': list([ - 'langgraph', - 'utils', - 'runnable', - 'RunnableCallable', - ]), - 'name': 'conduct_interview:ask_question', - }), - 'id': 'conduct_interview:ask_question', - 'type': 'runnable', - }), - dict({ - 'data': dict({ - 'id': list([ - 'langgraph', - 'utils', - 'runnable', - 'RunnableCallable', - ]), - 'name': 'conduct_interview:answer_question', - }), - 'id': 'conduct_interview:answer_question', - 'type': 'runnable', - }), - dict({ - 'data': 'conduct_interview:__end__', - 'id': 'conduct_interview:__end__', - 'type': 'schema', - }), - dict({ - 'data': dict({ - 'id': list([ - 'langgraph', - 'utils', - 'runnable', - 'RunnableCallable', - ]), - 'name': 'generate_sections', - }), - 'id': 'generate_sections', - 'type': 'runnable', - }), - dict({ - 'data': '__end__', - 'id': '__end__', - 'type': 'schema', - }), - ]), - }) -# --- diff --git a/libs/langgraph/tests/conftest.py b/libs/langgraph/tests/conftest.py index 510eddf62..1b2a18789 100644 --- a/libs/langgraph/tests/conftest.py +++ b/libs/langgraph/tests/conftest.py @@ -1,4 +1,3 @@ -import sys from contextlib import asynccontextmanager from typing import AsyncIterator, Optional from uuid import UUID, uuid4 @@ -6,21 +5,14 @@ from uuid import UUID, uuid4 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 psycopg import Connection +from psycopg_pool import ConnectionPool from pytest_mock import MockerFixture -from langgraph.checkpoint.base import BaseCheckpointSaver -from langgraph.checkpoint.postgres import PostgresSaver, ShallowPostgresSaver -from langgraph.checkpoint.postgres.aio import ( - AsyncPostgresSaver, - AsyncShallowPostgresSaver, -) -from langgraph.checkpoint.sqlite import SqliteSaver -from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver +from langgraph.checkpoint.postgres import PostgresSaver from langgraph.store.base import BaseStore from langgraph.store.memory import InMemoryStore -from langgraph.store.postgres import AsyncPostgresStore, PostgresStore +from langgraph.store.postgres import PostgresStore pytest.register_assert_rewrite("tests.memory_assert") @@ -55,18 +47,6 @@ def checkpointer_memory(): yield MemorySaverAssertImmutable() -@pytest.fixture(scope="function") -def checkpointer_sqlite(): - with SqliteSaver.from_conn_string(":memory:") as checkpointer: - yield checkpointer - - -@asynccontextmanager -async def _checkpointer_sqlite_aio(): - async with AsyncSqliteSaver.from_conn_string(":memory:") as checkpointer: - yield checkpointer - - @pytest.fixture(scope="function") def checkpointer_postgres(): database = f"test_{uuid4().hex[:16]}" @@ -86,25 +66,6 @@ def checkpointer_postgres(): conn.execute(f"DROP DATABASE {database}") -@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}") - - @pytest.fixture(scope="function") def checkpointer_postgres_pipe(): database = f"test_{uuid4().hex[:16]}" @@ -147,209 +108,6 @@ def checkpointer_postgres_pool(): 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}") - - -@asynccontextmanager -async def awith_checkpointer( - checkpointer_name: Optional[str], -) -> AsyncIterator[BaseCheckpointSaver]: - if checkpointer_name is None: - yield None - elif checkpointer_name == "memory": - from tests.memory_assert import MemorySaverAssertImmutable - - yield MemorySaverAssertImmutable() - 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_shallow": - async with _checkpointer_postgres_aio_shallow() 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}") - - -@asynccontextmanager -async def _store_postgres_aio(): - if sys.version_info < (3, 10): - pytest.skip("Async Postgres tests require Python 3.10+") - database = f"test_{uuid4().hex[:16]}" - async with await AsyncConnection.connect( - DEFAULT_POSTGRES_URI, autocommit=True - ) as conn: - await conn.execute(f"CREATE DATABASE {database}") - try: - async with AsyncPostgresStore.from_conn_string( - DEFAULT_POSTGRES_URI + database - ) as store: - await store.setup() - yield store - finally: - async with await AsyncConnection.connect( - DEFAULT_POSTGRES_URI, autocommit=True - ) as conn: - await conn.execute(f"DROP DATABASE {database}") - - -@asynccontextmanager -async def _store_postgres_aio_pipe(): - if sys.version_info < (3, 10): - pytest.skip("Async Postgres tests require Python 3.10+") - database = f"test_{uuid4().hex[:16]}" - async with await AsyncConnection.connect( - DEFAULT_POSTGRES_URI, autocommit=True - ) as conn: - await conn.execute(f"CREATE DATABASE {database}") - try: - async with AsyncPostgresStore.from_conn_string( - DEFAULT_POSTGRES_URI + database - ) as store: - await store.setup() # Run in its own transaction - async with AsyncPostgresStore.from_conn_string( - DEFAULT_POSTGRES_URI + database, pipeline=True - ) as store: - yield store - finally: - async with await AsyncConnection.connect( - DEFAULT_POSTGRES_URI, autocommit=True - ) as conn: - await conn.execute(f"DROP DATABASE {database}") - - -@asynccontextmanager -async def _store_postgres_aio_pool(): - if sys.version_info < (3, 10): - pytest.skip("Async Postgres tests require Python 3.10+") - database = f"test_{uuid4().hex[:16]}" - async with await AsyncConnection.connect( - DEFAULT_POSTGRES_URI, autocommit=True - ) as conn: - await conn.execute(f"CREATE DATABASE {database}") - try: - async with AsyncPostgresStore.from_conn_string( - DEFAULT_POSTGRES_URI + database, - pool_config={"max_size": 10}, - ) as store: - await store.setup() - yield store - finally: - async with await AsyncConnection.connect( - DEFAULT_POSTGRES_URI, autocommit=True - ) as conn: - await conn.execute(f"DROP DATABASE {database}") - - @pytest.fixture(scope="function") def store_postgres(): database = f"test_{uuid4().hex[:16]}" @@ -417,56 +175,20 @@ async def awith_store(store_name: Optional[str]) -> AsyncIterator[BaseStore]: yield None elif store_name == "in_memory": yield InMemoryStore() - elif store_name == "postgres_aio": - async with _store_postgres_aio() as store: - yield store - elif store_name == "postgres_aio_pipe": - async with _store_postgres_aio_pipe() as store: - yield store - elif store_name == "postgres_aio_pool": - async with _store_postgres_aio_pool() as store: - yield store else: raise NotImplementedError(f"Unknown store {store_name}") -SHALLOW_CHECKPOINTERS_SYNC = ["postgres_shallow"] +SHALLOW_CHECKPOINTERS_SYNC = [] REGULAR_CHECKPOINTERS_SYNC = [ "memory", - "sqlite", "postgres", - "postgres_pipe", - "postgres_pool", ] ALL_CHECKPOINTERS_SYNC = [ *REGULAR_CHECKPOINTERS_SYNC, *SHALLOW_CHECKPOINTERS_SYNC, ] -SHALLOW_CHECKPOINTERS_ASYNC = ["postgres_aio_shallow"] -REGULAR_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, -] ALL_STORES_SYNC = [ "in_memory", "postgres", - "postgres_pipe", - "postgres_pool", -] -ALL_STORES_ASYNC = [ - "in_memory", - "postgres_aio", - "postgres_aio_pipe", - "postgres_aio_pool", ] diff --git a/libs/langgraph/tests/messages.py b/libs/langgraph/tests/messages.py index ecc657a36..c067510d8 100644 --- a/libs/langgraph/tests/messages.py +++ b/libs/langgraph/tests/messages.py @@ -12,6 +12,7 @@ from typing import Any from langchain_core.documents import Document from langchain_core.messages import AIMessage, AIMessageChunk, HumanMessage, ToolMessage +from langgraph.graph.message import Message from tests.any_str import AnyStr @@ -38,7 +39,7 @@ def _AnyIdAIMessageChunk(**kwargs: Any) -> AIMessageChunk: def _AnyIdHumanMessage(**kwargs: Any) -> HumanMessage: """Create a human message with an any id field.""" - message = HumanMessage(**kwargs) + message = Message(role="user", **kwargs) message.id = AnyStr() return message diff --git a/libs/langgraph/tests/test_algo.py b/libs/langgraph/tests/test_algo.py index a3a588f6b..c7e1bcaf6 100644 --- a/libs/langgraph/tests/test_algo.py +++ b/libs/langgraph/tests/test_algo.py @@ -9,14 +9,13 @@ def test_prepare_next_tasks() -> None: processes = {} checkpoint = empty_checkpoint() - with ChannelsManager({}, checkpoint, config) as (channels, managed): + with ChannelsManager({}, checkpoint, config) as channels: assert ( prepare_next_tasks( checkpoint, {}, processes, channels, - managed, config, 0, for_execution=False, @@ -29,7 +28,6 @@ def test_prepare_next_tasks() -> None: {}, processes, channels, - managed, config, 0, for_execution=True, diff --git a/libs/langgraph/tests/test_channels.py b/libs/langgraph/tests/test_channels.py index 7c6fb162b..7f430a11a 100644 --- a/libs/langgraph/tests/test_channels.py +++ b/libs/langgraph/tests/test_channels.py @@ -8,8 +8,6 @@ from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic from langgraph.errors import EmptyChannelError, InvalidUpdateError -pytestmark = pytest.mark.anyio - def test_last_value() -> None: channel = LastValue(int).from_checkpoint(None) diff --git a/libs/langgraph/tests/test_interruption.py b/libs/langgraph/tests/test_interruption.py index aa543cf00..f268f9e3b 100644 --- a/libs/langgraph/tests/test_interruption.py +++ b/libs/langgraph/tests/test_interruption.py @@ -4,13 +4,9 @@ from typing_extensions import TypedDict from langgraph.graph import END, START, StateGraph from tests.conftest import ( - ALL_CHECKPOINTERS_ASYNC, ALL_CHECKPOINTERS_SYNC, - awith_checkpointer, ) -pytestmark = pytest.mark.anyio - @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_interruption_without_state_updates( @@ -48,41 +44,3 @@ def test_interruption_without_state_updates( graph.invoke(None, thread, debug=True) assert graph.get_state(thread).next == () - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_interruption_without_state_updates_async( - checkpointer_name: str, mocker: MockerFixture -): - """Test interruption without state updates. This test confirms that - interrupting doesn't require a state key having been updated in the prev step""" - - class State(TypedDict): - input: str - - async def noop(_state): - pass - - builder = StateGraph(State) - builder.add_node("step_1", noop) - builder.add_node("step_2", noop) - builder.add_node("step_3", noop) - builder.add_edge(START, "step_1") - builder.add_edge("step_1", "step_2") - 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="*") - - initial_input = {"input": "hello world"} - thread = {"configurable": {"thread_id": "1"}} - - await graph.ainvoke(initial_input, thread, debug=True) - assert (await graph.aget_state(thread)).next == ("step_2",) - - await graph.ainvoke(None, thread, debug=True) - assert (await graph.aget_state(thread)).next == ("step_3",) - - await graph.ainvoke(None, thread, debug=True) - assert (await graph.aget_state(thread)).next == () diff --git a/libs/langgraph/tests/test_large_cases.py b/libs/langgraph/tests/test_large_cases.py index f562720b8..3ef199e0b 100644 --- a/libs/langgraph/tests/test_large_cases.py +++ b/libs/langgraph/tests/test_large_cases.py @@ -4,29 +4,22 @@ import re import time from contextlib import contextmanager from dataclasses import replace -from typing import Annotated, Any, Iterator, Literal, Optional, Union, cast +from typing import Annotated, Iterator, Literal, Optional, Union, cast import httpx import pytest -from langchain_core.runnables import RunnableConfig, RunnableMap, RunnablePick from pytest_mock import MockerFixture from syrupy import SnapshotAssertion 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.errors import NodeInterrupt from langgraph.graph import StateGraph -from langgraph.graph.graph import Graph -from langgraph.graph.message import MessageGraph, MessagesState, add_messages -from langgraph.managed.shared_value import SharedValue -from langgraph.prebuilt.chat_agent_executor import create_react_agent -from langgraph.prebuilt.tool_node import ToolNode +from langgraph.graph.message import MessagesState, add_messages from langgraph.pregel import Channel, Pregel -from langgraph.store.memory import InMemoryStore from langgraph.types import ( Command, Interrupt, @@ -37,18 +30,17 @@ from langgraph.types import ( StreamWriter, interrupt, ) +from langgraph.utils.config import RunnableConfig from tests.agents import AgentAction, AgentFinish -from tests.any_str import AnyDict, AnyStr, UnsortedSequence +from tests.any_str import AnyDict, AnyStr from tests.conftest import ( ALL_CHECKPOINTERS_SYNC, REGULAR_CHECKPOINTERS_SYNC, SHOULD_CHECK_SNAPSHOTS, ) -from tests.fake_chat import FakeChatModel from tests.fake_tracer import FakeTracer from tests.messages import ( _AnyIdAIMessage, - _AnyIdAIMessageChunk, _AnyIdHumanMessage, _AnyIdToolMessage, ) @@ -503,917 +495,6 @@ def test_fork_always_re_runs_nodes( ] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_conditional_graph( - snapshot: SnapshotAssertion, request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - from langchain_core.language_models.fake import FakeStreamingListLLM - from langchain_core.prompts import PromptTemplate - from langchain_core.runnables import RunnablePassthrough - from langchain_core.tools import tool - - checkpointer: BaseCheckpointSaver = request.getfixturevalue( - f"checkpointer_{checkpointer_name}" - ) - - # Assemble the tools - @tool() - def search_api(query: str) -> str: - """Searches the API for the query.""" - return f"result for {query}" - - tools = [search_api] - - # Construct the agent - prompt = PromptTemplate.from_template("Hello!") - - llm = FakeStreamingListLLM( - responses=[ - "tool:search_api:query", - "tool:search_api:another", - "finish:answer", - ] - ) - - def agent_parser(input: str) -> Union[AgentAction, AgentFinish]: - if input.startswith("finish"): - _, answer = input.split(":") - return AgentFinish(return_values={"answer": answer}, log=input) - else: - _, tool_name, tool_input = input.split(":") - return AgentAction(tool=tool_name, tool_input=tool_input, log=input) - - agent = RunnablePassthrough.assign(agent_outcome=prompt | llm | agent_parser) - - # Define tool execution logic - def execute_tools(data: dict) -> dict: - data = data.copy() - agent_action: AgentAction = data.pop("agent_outcome") - observation = {t.name: t for t in tools}[agent_action.tool].invoke( - agent_action.tool_input - ) - if data.get("intermediate_steps") is None: - data["intermediate_steps"] = [] - else: - data["intermediate_steps"] = data["intermediate_steps"].copy() - data["intermediate_steps"].append([agent_action, observation]) - return data - - # Define decision-making logic - def should_continue(data: dict) -> str: - # Logic to decide whether to continue in the loop or exit - if isinstance(data["agent_outcome"], AgentFinish): - return "exit" - else: - return "continue" - - # Define a new graph - workflow = Graph() - - workflow.add_node("agent", agent) - workflow.add_node( - "tools", - execute_tools, - metadata={"parents": {}, "version": 2, "variant": "b"}, - ) - - workflow.set_entry_point("agent") - - workflow.add_conditional_edges( - "agent", should_continue, {"continue": "tools", "exit": END} - ) - - workflow.add_edge("tools", "agent") - - app = workflow.compile() - - if SHOULD_CHECK_SNAPSHOTS: - assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot - assert app.get_graph().draw_mermaid(with_styles=False) == snapshot - assert app.get_graph().draw_mermaid() == snapshot - assert json.dumps(app.get_graph(xray=True).to_json(), indent=2) == snapshot - assert app.get_graph(xray=True).draw_mermaid(with_styles=False) == snapshot - - assert app.invoke({"input": "what is weather in sf"}) == { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ], - [ - AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - "result for another", - ], - ], - "agent_outcome": AgentFinish( - return_values={"answer": "answer"}, log="finish:answer" - ), - } - - assert [c for c in app.stream({"input": "what is weather in sf"})] == [ - { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ), - } - }, - { - "tools": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ] - ], - } - }, - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ] - ], - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - { - "tools": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ], - [ - AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - "result for another", - ], - ], - } - }, - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ], - [ - AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - "result for another", - ], - ], - "agent_outcome": AgentFinish( - return_values={"answer": "answer"}, log="finish:answer" - ), - } - }, - ] - - # test state get/update methods with interrupt_after - - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_after=["agent"], - ) - config = {"configurable": {"thread_id": "1"}} - - if SHOULD_CHECK_SNAPSHOTS: - assert app_w_interrupt.get_graph().to_json() == snapshot - assert app_w_interrupt.get_graph().draw_mermaid() == snapshot - - assert [ - c for c in app_w_interrupt.stream({"input": "what is weather in sf"}, config) - ] == [ - { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ), - } - } - ] - - assert app_w_interrupt.get_state(config) == StateSnapshot( - values={ - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ), - }, - }, - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - created_at=AnyStr(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": { - "agent": { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - }, - }, - "thread_id": "1", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ), - ) - assert ( - app_w_interrupt.checkpointer.get_tuple(config).config["configurable"][ - "checkpoint_id" - ] - is not None - ) - - app_w_interrupt.update_state( - config, - { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "input": "what is weather in sf", - }, - ) - - assert app_w_interrupt.get_state(config) == StateSnapshot( - values={ - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "input": "what is weather in sf", - }, - }, - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 1, - "writes": { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "input": "what is weather in sf", - }, - }, - "thread_id": "1", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ), - ) - - assert [c for c in app_w_interrupt.stream(None, config)] == [ - { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "input": "what is weather in sf", - } - }, - { - "tools": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - } - }, - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - ] - - app_w_interrupt.update_state( - config, - { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ), - }, - ) - - assert app_w_interrupt.get_state(config) == StateSnapshot( - values={ - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ), - }, - }, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 4, - "writes": { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ), - } - }, - "thread_id": "1", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ), - ) - - # test state get/update methods with interrupt_before - - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_before=["tools"], - ) - config = {"configurable": {"thread_id": "2"}} - llm.i = 0 # reset the llm - - assert [ - c for c in app_w_interrupt.stream({"input": "what is weather in sf"}, config) - ] == [ - { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ), - } - } - ] - - assert app_w_interrupt.get_state(config) == StateSnapshot( - values={ - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ), - }, - }, - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": { - "agent": { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - } - }, - "thread_id": "2", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ), - ) - - app_w_interrupt.update_state( - config, - { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "input": "what is weather in sf", - }, - ) - - assert app_w_interrupt.get_state(config) == StateSnapshot( - values={ - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "input": "what is weather in sf", - }, - }, - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 1, - "writes": { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "input": "what is weather in sf", - } - }, - "thread_id": "2", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ), - ) - - assert [c for c in app_w_interrupt.stream(None, config)] == [ - { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "input": "what is weather in sf", - }, - }, - { - "tools": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - } - }, - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - ] - - app_w_interrupt.update_state( - config, - { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ), - }, - ) - - assert app_w_interrupt.get_state(config) == StateSnapshot( - values={ - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ), - }, - }, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 4, - "writes": { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ), - } - }, - "thread_id": "2", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ), - ) - - # test re-invoke to continue with interrupt_before - - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_before=["tools"], - ) - config = {"configurable": {"thread_id": "3"}} - llm.i = 0 # reset the llm - - assert [ - c for c in app_w_interrupt.stream({"input": "what is weather in sf"}, config) - ] == [ - { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ), - } - } - ] - - assert app_w_interrupt.get_state(config) == StateSnapshot( - values={ - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ), - }, - }, - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config={ - "configurable": { - "thread_id": "3", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": { - "agent": { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - } - }, - "thread_id": "3", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ), - ) - - assert [c for c in app_w_interrupt.stream(None, config)] == [ - { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ), - }, - }, - { - "tools": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ] - ], - } - }, - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ] - ], - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - ] - - assert [c for c in app_w_interrupt.stream(None, config)] == [ - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ] - ], - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - { - "tools": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ], - [ - AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - "result for another", - ], - ], - } - }, - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ], - [ - AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - "result for another", - ], - ], - "agent_outcome": AgentFinish( - return_values={"answer": "answer"}, log="finish:answer" - ), - } - }, - ] - - @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_conditional_state_graph( snapshot: SnapshotAssertion, @@ -1456,11 +537,9 @@ def test_conditional_state_graph( input: Annotated[str, UntrackedValue] agent_outcome: Optional[Union[AgentAction, AgentFinish]] intermediate_steps: Annotated[list[tuple[AgentAction, str]], operator.add] - session: Annotated[httpx.Client, Context(make_httpx_client)] class ToolState(TypedDict, total=False): agent_outcome: Union[AgentAction, AgentFinish] - session: Annotated[httpx.Client, Context(make_httpx_client)] # Assemble the tools @tool() @@ -2352,407 +1431,6 @@ def test_conditional_state_graph( ] -def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: - from langchain_core.messages import AIMessage, HumanMessage - from langchain_core.tools import tool - - @tool() - def search_api(query: str) -> str: - """Searches the API for the query.""" - return f"result for {query}" - - tools = [search_api] - - model = FakeChatModel( - messages=[ - AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - }, - ], - ), - AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call234", - "name": "search_api", - "args": {"query": "another"}, - }, - { - "id": "tool_call567", - "name": "search_api", - "args": {"query": "a third one"}, - }, - ], - ), - AIMessage(content="answer"), - ] - ) - - app = create_react_agent(model, tools) - - if SHOULD_CHECK_SNAPSHOTS: - assert json.dumps(app.get_input_schema().model_json_schema()) == snapshot - assert json.dumps(app.get_output_schema().model_json_schema()) == snapshot - assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot - assert app.get_graph().draw_mermaid(with_styles=False) == snapshot - - assert app.invoke( - {"messages": [HumanMessage(content="what is weather in sf")]} - ) == { - "messages": [ - _AnyIdHumanMessage(content="what is weather in sf"), - _AnyIdAIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - }, - ], - ), - _AnyIdToolMessage( - content="result for query", - name="search_api", - tool_call_id="tool_call123", - ), - _AnyIdAIMessage( - content="", - tool_calls=[ - { - "id": "tool_call234", - "name": "search_api", - "args": {"query": "another"}, - }, - { - "id": "tool_call567", - "name": "search_api", - "args": {"query": "a third one"}, - }, - ], - ), - _AnyIdToolMessage( - content="result for another", - name="search_api", - tool_call_id="tool_call234", - ), - _AnyIdToolMessage( - content="result for a third one", - name="search_api", - tool_call_id="tool_call567", - id=AnyStr(), - ), - _AnyIdAIMessage(content="answer"), - ] - } - - assert [ - c - for c in app.stream( - {"messages": [HumanMessage(content="what is weather in sf")]}, - stream_mode="messages", - ) - ] == [ - ( - _AnyIdAIMessageChunk( - content="", - tool_calls=[ - { - "name": "search_api", - "args": {"query": "query"}, - "id": "tool_call123", - "type": "tool_call", - } - ], - tool_call_chunks=[ - { - "name": "search_api", - "args": '{"query": "query"}', - "id": "tool_call123", - "index": None, - "type": "tool_call_chunk", - } - ], - ), - { - "langgraph_step": 1, - "langgraph_node": "agent", - "langgraph_triggers": ["start:agent"], - "langgraph_path": (PULL, "agent"), - "langgraph_checkpoint_ns": AnyStr("agent:"), - "checkpoint_ns": AnyStr("agent:"), - "ls_provider": "fakechatmodel", - "ls_model_type": "chat", - }, - ), - ( - _AnyIdToolMessage( - content="result for query", - name="search_api", - tool_call_id="tool_call123", - ), - { - "langgraph_step": 2, - "langgraph_node": "tools", - "langgraph_triggers": ["branch:agent:should_continue:tools"], - "langgraph_path": (PULL, "tools"), - "langgraph_checkpoint_ns": AnyStr("tools:"), - }, - ), - ( - _AnyIdAIMessageChunk( - content="", - tool_calls=[ - { - "name": "search_api", - "args": {"query": "another"}, - "id": "tool_call234", - "type": "tool_call", - }, - { - "name": "search_api", - "args": {"query": "a third one"}, - "id": "tool_call567", - "type": "tool_call", - }, - ], - tool_call_chunks=[ - { - "name": "search_api", - "args": '{"query": "another"}', - "id": "tool_call234", - "index": None, - "type": "tool_call_chunk", - }, - { - "name": "search_api", - "args": '{"query": "a third one"}', - "id": "tool_call567", - "index": None, - "type": "tool_call_chunk", - }, - ], - ), - { - "langgraph_step": 3, - "langgraph_node": "agent", - "langgraph_triggers": ["tools"], - "langgraph_path": (PULL, "agent"), - "langgraph_checkpoint_ns": AnyStr("agent:"), - "checkpoint_ns": AnyStr("agent:"), - "ls_provider": "fakechatmodel", - "ls_model_type": "chat", - }, - ), - ( - _AnyIdToolMessage( - content="result for another", - name="search_api", - tool_call_id="tool_call234", - ), - { - "langgraph_step": 4, - "langgraph_node": "tools", - "langgraph_triggers": ["branch:agent:should_continue:tools"], - "langgraph_path": (PULL, "tools"), - "langgraph_checkpoint_ns": AnyStr("tools:"), - }, - ), - ( - _AnyIdToolMessage( - content="result for a third one", - name="search_api", - tool_call_id="tool_call567", - ), - { - "langgraph_step": 4, - "langgraph_node": "tools", - "langgraph_triggers": ["branch:agent:should_continue:tools"], - "langgraph_path": (PULL, "tools"), - "langgraph_checkpoint_ns": AnyStr("tools:"), - }, - ), - ( - _AnyIdAIMessageChunk( - content="answer", - ), - { - "langgraph_step": 5, - "langgraph_node": "agent", - "langgraph_triggers": ["tools"], - "langgraph_path": (PULL, "agent"), - "langgraph_checkpoint_ns": AnyStr("agent:"), - "checkpoint_ns": AnyStr("agent:"), - "ls_provider": "fakechatmodel", - "ls_model_type": "chat", - }, - ), - ] - - assert app.invoke( - {"messages": [HumanMessage(content="what is weather in sf")]}, - {"recursion_limit": 2}, - debug=True, - ) == { - "messages": [ - _AnyIdHumanMessage(content="what is weather in sf"), - _AnyIdAIMessage(content="Sorry, need more steps to process this request."), - ] - } - - model.i = 0 # reset the model - - assert ( - app.invoke( - {"messages": [HumanMessage(content="what is weather in sf")]}, - stream_mode="updates", - )[0]["agent"]["messages"] - == [ - { - "agent": { - "messages": [ - _AnyIdAIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - }, - ], - ) - ] - } - }, - { - "tools": { - "messages": [ - _AnyIdToolMessage( - content="result for query", - name="search_api", - tool_call_id="tool_call123", - ) - ] - } - }, - { - "agent": { - "messages": [ - _AnyIdAIMessage( - content="", - tool_calls=[ - { - "id": "tool_call234", - "name": "search_api", - "args": {"query": "another"}, - }, - { - "id": "tool_call567", - "name": "search_api", - "args": {"query": "a third one"}, - }, - ], - ) - ] - } - }, - { - "tools": { - "messages": [ - _AnyIdToolMessage( - content="result for another", - name="search_api", - tool_call_id="tool_call234", - ), - _AnyIdToolMessage( - content="result for a third one", - name="search_api", - tool_call_id="tool_call567", - ), - ] - } - }, - {"agent": {"messages": [_AnyIdAIMessage(content="answer")]}}, - ][0]["agent"]["messages"] - ) - - assert [ - *app.stream({"messages": [HumanMessage(content="what is weather in sf")]}) - ] == [ - { - "agent": { - "messages": [ - _AnyIdAIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - }, - ], - ) - ] - } - }, - { - "tools": { - "messages": [ - _AnyIdToolMessage( - content="result for query", - name="search_api", - tool_call_id="tool_call123", - ) - ] - } - }, - { - "agent": { - "messages": [ - _AnyIdAIMessage( - content="", - tool_calls=[ - { - "id": "tool_call234", - "name": "search_api", - "args": {"query": "another"}, - }, - { - "id": "tool_call567", - "name": "search_api", - "args": {"query": "a third one"}, - }, - ], - ) - ] - } - }, - { - "tools": { - "messages": [ - _AnyIdToolMessage( - content="result for another", - name="search_api", - tool_call_id="tool_call234", - ), - _AnyIdToolMessage( - content="result for a third one", - name="search_api", - tool_call_id="tool_call567", - ), - ] - } - }, - {"agent": {"messages": [_AnyIdAIMessage(content="answer")]}}, - ] - - @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_state_graph_packets( request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture @@ -2775,7 +1453,6 @@ def test_state_graph_packets( class AgentState(TypedDict): messages: Annotated[list[BaseMessage], add_messages] - session: Annotated[httpx.Client, Context(httpx.Client)] @tool() def search_api(query: str) -> str: @@ -3648,1773 +2325,6 @@ def test_state_graph_packets( ) -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_message_graph( - snapshot: SnapshotAssertion, - deterministic_uuids: MockerFixture, - request: pytest.FixtureRequest, - checkpointer_name: str, -) -> None: - from copy import deepcopy - - from langchain_core.callbacks import CallbackManagerForLLMRun - from langchain_core.language_models.fake_chat_models import ( - FakeMessagesListChatModel, - ) - from langchain_core.messages import AIMessage, BaseMessage, HumanMessage - from langchain_core.outputs import ChatGeneration, ChatResult - from langchain_core.tools import tool - - checkpointer: BaseCheckpointSaver = request.getfixturevalue( - f"checkpointer_{checkpointer_name}" - ) - - class FakeFuntionChatModel(FakeMessagesListChatModel): - def bind_functions(self, functions: list): - return self - - def _generate( - self, - messages: list[BaseMessage], - stop: Optional[list[str]] = None, - run_manager: Optional[CallbackManagerForLLMRun] = None, - **kwargs: Any, - ) -> ChatResult: - response = deepcopy(self.responses[self.i]) - if self.i < len(self.responses) - 1: - self.i += 1 - else: - self.i = 0 - generation = ChatGeneration(message=response) - return ChatResult(generations=[generation]) - - @tool() - def search_api(query: str) -> str: - """Searches the API for the query.""" - return f"result for {query}" - - tools = [search_api] - - model = FakeFuntionChatModel( - responses=[ - AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - } - ], - id="ai1", - ), - AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call456", - "name": "search_api", - "args": {"query": "another"}, - } - ], - id="ai2", - ), - AIMessage(content="answer", id="ai3"), - ] - ) - - # Define the function that determines whether to continue or not - def should_continue(messages): - last_message = messages[-1] - # If there is no function call, then we finish - if not last_message.tool_calls: - return "end" - # Otherwise if there is, we continue - else: - return "continue" - - # Define a new graph - workflow = MessageGraph() - - # Define the two nodes we will cycle between - workflow.add_node("agent", model) - workflow.add_node("tools", ToolNode(tools)) - - # Set the entrypoint as `agent` - # This means that this node is the first one called - workflow.set_entry_point("agent") - - # We now add a conditional edge - workflow.add_conditional_edges( - # First, we define the start node. We use `agent`. - # This means these are the edges taken after the `agent` node is called. - "agent", - # Next, we pass in the function that will determine which node is called next. - should_continue, - # Finally we pass in a mapping. - # The keys are strings, and the values are other nodes. - # END is a special node marking that the graph should finish. - # What will happen is we will call `should_continue`, and then the output of that - # will be matched against the keys in this mapping. - # Based on which one it matches, that node will then be called. - { - # If `tools`, then we call the tool node. - "continue": "tools", - # Otherwise we finish. - "end": END, - }, - ) - - # We now add a normal edge from `tools` to `agent`. - # This means that after `tools` is called, `agent` node is called next. - workflow.add_edge("tools", "agent") - - # Finally, we compile it! - # This compiles it into a LangChain Runnable, - # meaning you can use it as you would any other runnable - app = workflow.compile() - - if SHOULD_CHECK_SNAPSHOTS: - assert json.dumps(app.get_input_schema().model_json_schema()) == snapshot - assert json.dumps(app.get_output_schema().model_json_schema()) == snapshot - assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot - assert app.get_graph().draw_mermaid(with_styles=False) == snapshot - - assert app.invoke(HumanMessage(content="what is weather in sf")) == [ - _AnyIdHumanMessage( - content="what is weather in sf", - ), - AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - } - ], - id="ai1", # respects ids passed in - ), - _AnyIdToolMessage( - content="result for query", - name="search_api", - tool_call_id="tool_call123", - ), - AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call456", - "name": "search_api", - "args": {"query": "another"}, - } - ], - id="ai2", - ), - _AnyIdToolMessage( - content="result for another", - name="search_api", - tool_call_id="tool_call456", - ), - AIMessage(content="answer", id="ai3"), - ] - - assert [*app.stream([HumanMessage(content="what is weather in sf")])] == [ - { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - } - ], - id="ai1", - ) - }, - { - "tools": [ - _AnyIdToolMessage( - content="result for query", - name="search_api", - tool_call_id="tool_call123", - ) - ] - }, - { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call456", - "name": "search_api", - "args": {"query": "another"}, - } - ], - id="ai2", - ) - }, - { - "tools": [ - _AnyIdToolMessage( - content="result for another", - name="search_api", - tool_call_id="tool_call456", - ) - ] - }, - {"agent": AIMessage(content="answer", id="ai3")}, - ] - - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_after=["agent"], - ) - config = {"configurable": {"thread_id": "1"}} - - assert [ - c for c in app_w_interrupt.stream(("human", "what is weather in sf"), config) - ] == [ - { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - } - ], - id="ai1", - ) - }, - {"__interrupt__": ()}, - ] - - assert app_w_interrupt.get_state(config) == StateSnapshot( - values=[ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - } - ], - id="ai1", - ), - ], - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - } - ], - id="ai1", - ) - }, - "thread_id": "1", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ), - ) - - # modify ai message - last_message = app_w_interrupt.get_state(config).values[-1] - last_message.tool_calls[0]["args"] = {"query": "a different query"} - next_config = app_w_interrupt.update_state(config, last_message) - - # message was replaced instead of appended - assert app_w_interrupt.get_state(config) == StateSnapshot( - values=[ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - content="", - id="ai1", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - } - ], - ), - ], - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config=next_config, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 2, - "writes": { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - } - ], - id="ai1", - ) - }, - "thread_id": "1", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ), - ) - - assert [c for c in app_w_interrupt.stream(None, config)] == [ - { - "tools": [ - _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", - ) - ] - }, - { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call456", - "name": "search_api", - "args": {"query": "another"}, - } - ], - id="ai2", - ) - }, - {"__interrupt__": ()}, - ] - - assert app_w_interrupt.get_state(config) == StateSnapshot( - values=[ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - content="", - id="ai1", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - } - ], - ), - _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", - ), - AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call456", - "name": "search_api", - "args": {"query": "another"}, - } - ], - id="ai2", - ), - ], - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 4, - "writes": { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call456", - "name": "search_api", - "args": {"query": "another"}, - } - ], - id="ai2", - ) - }, - "thread_id": "1", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ), - ) - - app_w_interrupt.update_state( - config, - AIMessage(content="answer", id="ai2"), # replace existing message - ) - - # replaces message even if object identity is different, as long as id is the same - assert app_w_interrupt.get_state(config) == StateSnapshot( - values=[ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - content="", - id="ai1", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - } - ], - ), - _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", - ), - AIMessage(content="answer", id="ai2"), - ], - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 5, - "writes": {"agent": AIMessage(content="answer", id="ai2")}, - "thread_id": "1", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ), - ) - - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_before=["tools"], - ) - config = {"configurable": {"thread_id": "2"}} - model.i = 0 # reset the llm - - assert [c for c in app_w_interrupt.stream("what is weather in sf", config)] == [ - { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - } - ], - id="ai1", - ) - }, - {"__interrupt__": ()}, - ] - - assert app_w_interrupt.get_state(config) == StateSnapshot( - values=[ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - } - ], - id="ai1", - ), - ], - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - } - ], - id="ai1", - ) - }, - "thread_id": "2", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ), - ) - - # modify ai message - last_message = app_w_interrupt.get_state(config).values[-1] - last_message.tool_calls[0]["args"] = {"query": "a different query"} - app_w_interrupt.update_state(config, last_message) - - # message was replaced instead of appended - assert app_w_interrupt.get_state(config) == StateSnapshot( - values=[ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - content="", - id="ai1", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - } - ], - ), - ], - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 2, - "writes": { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - } - ], - id="ai1", - ) - }, - "thread_id": "2", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ), - ) - - assert [c for c in app_w_interrupt.stream(None, config)] == [ - { - "tools": [ - _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", - ) - ] - }, - { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call456", - "name": "search_api", - "args": {"query": "another"}, - } - ], - id="ai2", - ) - }, - {"__interrupt__": ()}, - ] - - assert app_w_interrupt.get_state(config) == StateSnapshot( - values=[ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - content="", - id="ai1", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - } - ], - ), - _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", - ), - AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call456", - "name": "search_api", - "args": {"query": "another"}, - } - ], - id="ai2", - ), - ], - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 4, - "writes": { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call456", - "name": "search_api", - "args": {"query": "another"}, - } - ], - id="ai2", - ) - }, - "thread_id": "2", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ), - ) - - app_w_interrupt.update_state( - config, - AIMessage(content="answer", id="ai2"), - ) - - # replaces message even if object identity is different, as long as id is the same - assert app_w_interrupt.get_state(config) == StateSnapshot( - values=[ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - content="", - id="ai1", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - } - ], - ), - _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", - id=AnyStr(), - ), - AIMessage(content="answer", id="ai2"), - ], - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 5, - "writes": {"agent": AIMessage(content="answer", id="ai2")}, - "thread_id": "2", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ), - ) - - # add an extra message as if it came from "tools" node - app_w_interrupt.update_state(config, ("ai", "an extra message"), as_node="tools") - - # extra message is coerced BaseMessge and appended - # now the next node is "agent" per the graph edges - assert app_w_interrupt.get_state(config) == StateSnapshot( - values=[ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - content="", - id="ai1", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - } - ], - ), - _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", - id=AnyStr(), - ), - AIMessage(content="answer", id="ai2"), - _AnyIdAIMessage(content="an extra message"), - ], - tasks=(PregelTask(AnyStr(), "agent", (PULL, "agent")),), - next=("agent",), - config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 6, - "writes": {"tools": UnsortedSequence("ai", "an extra message")}, - "thread_id": "2", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ), - ) - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_root_graph( - deterministic_uuids: MockerFixture, - request: pytest.FixtureRequest, - checkpointer_name: str, -) -> None: - from copy import deepcopy - - from langchain_core.callbacks import CallbackManagerForLLMRun - from langchain_core.language_models.fake_chat_models import ( - FakeMessagesListChatModel, - ) - from langchain_core.messages import ( - AIMessage, - BaseMessage, - HumanMessage, - ToolMessage, - ) - from langchain_core.outputs import ChatGeneration, ChatResult - from langchain_core.tools import tool - - checkpointer: BaseCheckpointSaver = request.getfixturevalue( - f"checkpointer_{checkpointer_name}" - ) - - class FakeFuntionChatModel(FakeMessagesListChatModel): - def bind_functions(self, functions: list): - return self - - def _generate( - self, - messages: list[BaseMessage], - stop: Optional[list[str]] = None, - run_manager: Optional[CallbackManagerForLLMRun] = None, - **kwargs: Any, - ) -> ChatResult: - response = deepcopy(self.responses[self.i]) - if self.i < len(self.responses) - 1: - self.i += 1 - else: - self.i = 0 - generation = ChatGeneration(message=response) - return ChatResult(generations=[generation]) - - @tool() - def search_api(query: str) -> str: - """Searches the API for the query.""" - return f"result for {query}" - - tools = [search_api] - - model = FakeFuntionChatModel( - responses=[ - AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - } - ], - id="ai1", - ), - AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call456", - "name": "search_api", - "args": {"query": "another"}, - } - ], - id="ai2", - ), - AIMessage(content="answer", id="ai3"), - ] - ) - - # Define the function that determines whether to continue or not - def should_continue(messages): - last_message = messages[-1] - # If there is no function call, then we finish - if not last_message.tool_calls: - return "end" - # Otherwise if there is, we continue - else: - return "continue" - - class State(TypedDict): - __root__: Annotated[list[BaseMessage], add_messages] - - # Define a new graph - workflow = StateGraph(State) - - # Define the two nodes we will cycle between - workflow.add_node("agent", model) - workflow.add_node("tools", ToolNode(tools)) - - # Set the entrypoint as `agent` - # This means that this node is the first one called - workflow.set_entry_point("agent") - - # We now add a conditional edge - workflow.add_conditional_edges( - # First, we define the start node. We use `agent`. - # This means these are the edges taken after the `agent` node is called. - "agent", - # Next, we pass in the function that will determine which node is called next. - should_continue, - # Finally we pass in a mapping. - # The keys are strings, and the values are other nodes. - # END is a special node marking that the graph should finish. - # What will happen is we will call `should_continue`, and then the output of that - # will be matched against the keys in this mapping. - # Based on which one it matches, that node will then be called. - { - # If `tools`, then we call the tool node. - "continue": "tools", - # Otherwise we finish. - "end": END, - }, - ) - - # We now add a normal edge from `tools` to `agent`. - # This means that after `tools` is called, `agent` node is called next. - workflow.add_edge("tools", "agent") - - # Finally, we compile it! - # This compiles it into a LangChain Runnable, - # meaning you can use it as you would any other runnable - app = workflow.compile() - - assert app.invoke(HumanMessage(content="what is weather in sf")) == [ - _AnyIdHumanMessage( - content="what is weather in sf", - ), - AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - } - ], - id="ai1", # respects ids passed in - ), - _AnyIdToolMessage( - content="result for query", - name="search_api", - tool_call_id="tool_call123", - ), - AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call456", - "name": "search_api", - "args": {"query": "another"}, - } - ], - id="ai2", - ), - _AnyIdToolMessage( - content="result for another", - name="search_api", - tool_call_id="tool_call456", - ), - AIMessage(content="answer", id="ai3"), - ] - - assert [*app.stream([HumanMessage(content="what is weather in sf")])] == [ - { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - } - ], - id="ai1", - ) - }, - { - "tools": [ - ToolMessage( - content="result for query", - name="search_api", - tool_call_id="tool_call123", - id="00000000-0000-4000-8000-000000000037", - ) - ] - }, - { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call456", - "name": "search_api", - "args": {"query": "another"}, - } - ], - id="ai2", - ) - }, - { - "tools": [ - ToolMessage( - content="result for another", - name="search_api", - tool_call_id="tool_call456", - id="00000000-0000-4000-8000-000000000045", - ) - ] - }, - {"agent": AIMessage(content="answer", id="ai3")}, - ] - - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_after=["agent"], - ) - config = {"configurable": {"thread_id": "1"}} - - assert [ - c for c in app_w_interrupt.stream(("human", "what is weather in sf"), config) - ] == [ - { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - } - ], - id="ai1", - ) - }, - {"__interrupt__": ()}, - ] - - assert app_w_interrupt.get_state(config) == StateSnapshot( - values=[ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - } - ], - id="ai1", - ), - ], - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - } - ], - id="ai1", - ) - }, - "thread_id": "1", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ), - ) - - # modify ai message - last_message = app_w_interrupt.get_state(config).values[-1] - last_message.tool_calls[0]["args"] = {"query": "a different query"} - next_config = app_w_interrupt.update_state(config, last_message) - - # message was replaced instead of appended - assert app_w_interrupt.get_state(config) == StateSnapshot( - values=[ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - content="", - id="ai1", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - } - ], - ), - ], - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config=next_config, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 2, - "writes": { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - } - ], - id="ai1", - ) - }, - "thread_id": "1", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ), - ) - - assert [c for c in app_w_interrupt.stream(None, config)] == [ - { - "tools": [ - _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", - ) - ] - }, - { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call456", - "name": "search_api", - "args": {"query": "another"}, - } - ], - id="ai2", - ) - }, - {"__interrupt__": ()}, - ] - - assert app_w_interrupt.get_state(config) == StateSnapshot( - values=[ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - content="", - id="ai1", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - } - ], - ), - _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", - id=AnyStr(), - ), - AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call456", - "name": "search_api", - "args": {"query": "another"}, - } - ], - id="ai2", - ), - ], - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 4, - "writes": { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call456", - "name": "search_api", - "args": {"query": "another"}, - } - ], - id="ai2", - ) - }, - "thread_id": "1", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ), - ) - - app_w_interrupt.update_state( - config, - AIMessage(content="answer", id="ai2"), # replace existing message - ) - - # replaces message even if object identity is different, as long as id is the same - assert app_w_interrupt.get_state(config) == StateSnapshot( - values=[ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - content="", - id="ai1", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - } - ], - ), - _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", - id=AnyStr(), - ), - AIMessage(content="answer", id="ai2"), - ], - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 5, - "writes": {"agent": AIMessage(content="answer", id="ai2")}, - "thread_id": "1", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ), - ) - - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_before=["tools"], - ) - config = {"configurable": {"thread_id": "2"}} - model.i = 0 # reset the llm - - assert [c for c in app_w_interrupt.stream("what is weather in sf", config)] == [ - { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - } - ], - id="ai1", - ) - }, - {"__interrupt__": ()}, - ] - - assert app_w_interrupt.get_state(config) == StateSnapshot( - values=[ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - } - ], - id="ai1", - ), - ], - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - } - ], - id="ai1", - ) - }, - "thread_id": "2", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ), - ) - - # modify ai message - last_message = app_w_interrupt.get_state(config).values[-1] - last_message.tool_calls[0]["args"] = {"query": "a different query"} - app_w_interrupt.update_state(config, last_message) - - # message was replaced instead of appended - assert app_w_interrupt.get_state(config) == StateSnapshot( - values=[ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - content="", - id="ai1", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - } - ], - ), - ], - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 2, - "writes": { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - } - ], - id="ai1", - ) - }, - "thread_id": "2", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ), - ) - - assert [c for c in app_w_interrupt.stream(None, config)] == [ - { - "tools": [ - _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", - ) - ] - }, - { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call456", - "name": "search_api", - "args": {"query": "another"}, - } - ], - id="ai2", - ) - }, - {"__interrupt__": ()}, - ] - - assert app_w_interrupt.get_state(config) == StateSnapshot( - values=[ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - content="", - id="ai1", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - } - ], - ), - _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", - id=AnyStr(), - ), - AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call456", - "name": "search_api", - "args": {"query": "another"}, - } - ], - id="ai2", - ), - ], - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 4, - "writes": { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call456", - "name": "search_api", - "args": {"query": "another"}, - } - ], - id="ai2", - ) - }, - "thread_id": "2", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ), - ) - - app_w_interrupt.update_state( - config, - AIMessage(content="answer", id="ai2"), - ) - - # replaces message even if object identity is different, as long as id is the same - assert app_w_interrupt.get_state(config) == StateSnapshot( - values=[ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - content="", - id="ai1", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - } - ], - ), - _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", - ), - AIMessage(content="answer", id="ai2"), - ], - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 5, - "writes": {"agent": AIMessage(content="answer", id="ai2")}, - "thread_id": "2", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ), - ) - - # add an extra message as if it came from "tools" node - app_w_interrupt.update_state(config, ("ai", "an extra message"), as_node="tools") - - # extra message is coerced BaseMessge and appended - # now the next node is "agent" per the graph edges - assert app_w_interrupt.get_state(config) == StateSnapshot( - values=[ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - content="", - id="ai1", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - } - ], - ), - _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", - id=AnyStr(), - ), - AIMessage(content="answer", id="ai2"), - _AnyIdAIMessage(content="an extra message"), - ], - tasks=(PregelTask(AnyStr(), "agent", (PULL, "agent")),), - next=("agent",), - config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 6, - "writes": {"tools": UnsortedSequence("ai", "an extra message")}, - "thread_id": "2", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ), - ) - - # create new graph with one more state key, reuse previous thread history - - def simple_add(left, right): - if not isinstance(right, list): - right = [right] - return left + right - - class MoreState(TypedDict): - __root__: Annotated[list[BaseMessage], simple_add] - something_else: str - - # Define a new graph - new_workflow = StateGraph(MoreState) - new_workflow.add_node( - "agent", RunnableMap(__root__=RunnablePick("__root__") | model) - ) - new_workflow.add_node( - "tools", RunnableMap(__root__=RunnablePick("__root__") | ToolNode(tools)) - ) - new_workflow.set_entry_point("agent") - new_workflow.add_conditional_edges( - "agent", - RunnablePick("__root__") | should_continue, - { - # If `tools`, then we call the tool node. - "continue": "tools", - # Otherwise we finish. - "end": END, - }, - ) - new_workflow.add_edge("tools", "agent") - new_app = new_workflow.compile(checkpointer=checkpointer) - model.i = 0 # reset the llm - - # previous state is converted to new schema - assert new_app.get_state(config) == StateSnapshot( - values={ - "__root__": [ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - content="", - id="ai1", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - } - ], - ), - _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", - ), - AIMessage(content="answer", id="ai2"), - _AnyIdAIMessage(content="an extra message"), - ] - }, - tasks=(PregelTask(AnyStr(), "agent", (PULL, "agent")),), - next=("agent",), - config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 6, - "writes": {"tools": UnsortedSequence("ai", "an extra message")}, - "thread_id": "2", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(new_app.checkpointer.list(config, limit=2))[-1].config - ), - ) - - # new input is merged to old state - assert new_app.invoke( - { - "__root__": [HumanMessage(content="what is weather in la")], - "something_else": "value", - }, - config, - interrupt_before=["agent"], - ) == { - "__root__": [ - HumanMessage( - content="what is weather in sf", - id="00000000-0000-4000-8000-000000000078", - ), - AIMessage( - content="", - id="ai1", - tool_calls=[ - { - "name": "search_api", - "args": {"query": "a different query"}, - "id": "tool_call123", - } - ], - ), - _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", - ), - AIMessage(content="answer", id="ai2"), - AIMessage( - content="an extra message", id="00000000-0000-4000-8000-000000000100" - ), - HumanMessage(content="what is weather in la"), - ], - "something_else": "value", - } - - def test_in_one_fan_out_out_one_graph_state() -> None: def sorted_add(x: list[str], y: list[str]) -> list[str]: return sorted(operator.add(x, y)) @@ -6198,308 +3108,6 @@ def test_dynamic_interrupt_subgraph( ) -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_start_branch_then( - snapshot: SnapshotAssertion, request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - - class State(TypedDict): - my_key: Annotated[str, operator.add] - market: str - shared: Annotated[dict[str, dict[str, Any]], SharedValue.on("assistant_id")] - - def assert_shared_value(data: State, config: RunnableConfig) -> State: - assert "shared" in data - if thread_id := config["configurable"].get("thread_id"): - if thread_id == "1": - # this is the first thread, so should not see a value - assert data["shared"] == {} - return {"shared": {"1": {"hello": "world"}}} - elif thread_id == "2": - # this should get value saved by thread 1 - assert data["shared"] == {"1": {"hello": "world"}} - elif thread_id == "3": - # this is a different assistant, so should not see previous value - assert data["shared"] == {} - return {} - - def tool_two_slow(data: State, config: RunnableConfig) -> State: - return {"my_key": " slow", **assert_shared_value(data, config)} - - def tool_two_fast(data: State, config: RunnableConfig) -> State: - return {"my_key": " fast", **assert_shared_value(data, config)} - - tool_two_graph = StateGraph(State) - tool_two_graph.add_node("tool_two_slow", tool_two_slow) - tool_two_graph.add_node("tool_two_fast", tool_two_fast) - tool_two_graph.set_conditional_entry_point( - lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast", then=END - ) - tool_two = tool_two_graph.compile() - assert tool_two.get_graph().draw_mermaid() == snapshot - - assert tool_two.invoke({"my_key": "value", "market": "DE"}) == { - "my_key": "value slow", - "market": "DE", - } - assert tool_two.invoke({"my_key": "value", "market": "US"}) == { - "my_key": "value fast", - "market": "US", - } - - tool_two = tool_two_graph.compile( - store=InMemoryStore(), - checkpointer=checkpointer, - interrupt_before=["tool_two_fast", "tool_two_slow"], - ) - - # missing thread_id - with pytest.raises(ValueError, match="thread_id"): - tool_two.invoke({"my_key": "value", "market": "DE"}) - - thread1 = {"configurable": {"thread_id": "1", "assistant_id": "a"}} - # stop when about to enter node - assert tool_two.invoke({"my_key": "value ⛰️", "market": "DE"}, thread1) == { - "my_key": "value ⛰️", - "market": "DE", - } - - if "shallow" not in checkpointer_name: - assert [c.metadata for c in tool_two.checkpointer.list(thread1)] == [ - { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "assistant_id": "a", - "thread_id": "1", - }, - { - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, - "assistant_id": "a", - "thread_id": "1", - }, - ] - - assert tool_two.get_state(thread1) == StateSnapshot( - values={"my_key": "value ⛰️", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "tool_two_slow", (PULL, "tool_two_slow")),), - next=("tool_two_slow",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "assistant_id": "a", - "thread_id": "1", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(tool_two.checkpointer.list(thread1, limit=2))[-1].config - ), - ) - # resume, for same result as above - assert tool_two.invoke(None, thread1, debug=1) == { - "my_key": "value ⛰️ slow", - "market": "DE", - } - assert tool_two.get_state(thread1) == StateSnapshot( - values={"my_key": "value ⛰️ slow", "market": "DE"}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"tool_two_slow": {"my_key": " slow"}}, - "assistant_id": "a", - "thread_id": "1", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(tool_two.checkpointer.list(thread1, limit=2))[-1].config - ), - ) - - thread2 = {"configurable": {"thread_id": "2", "assistant_id": "a"}} - # stop when about to enter node - assert tool_two.invoke({"my_key": "value", "market": "US"}, thread2) == { - "my_key": "value", - "market": "US", - } - assert tool_two.get_state(thread2) == StateSnapshot( - values={"my_key": "value", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),), - next=("tool_two_fast",), - config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "assistant_id": "a", - "thread_id": "2", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(tool_two.checkpointer.list(thread2, limit=2))[-1].config - ), - ) - # resume, for same result as above - assert tool_two.invoke(None, thread2, debug=1) == { - "my_key": "value fast", - "market": "US", - } - assert tool_two.get_state(thread2) == StateSnapshot( - values={"my_key": "value fast", "market": "US"}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"tool_two_fast": {"my_key": " fast"}}, - "assistant_id": "a", - "thread_id": "2", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(tool_two.checkpointer.list(thread2, limit=2))[-1].config - ), - ) - - thread3 = {"configurable": {"thread_id": "3", "assistant_id": "b"}} - # stop when about to enter node - assert tool_two.invoke({"my_key": "value", "market": "US"}, thread3) == { - "my_key": "value", - "market": "US", - } - assert tool_two.get_state(thread3) == StateSnapshot( - values={"my_key": "value", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),), - next=("tool_two_fast",), - config={ - "configurable": { - "thread_id": "3", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "assistant_id": "b", - "thread_id": "3", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(tool_two.checkpointer.list(thread3, limit=2))[-1].config - ), - ) - # update state - tool_two.update_state(thread3, {"my_key": "key"}) # appends to my_key - assert tool_two.get_state(thread3) == StateSnapshot( - values={"my_key": "valuekey", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),), - next=("tool_two_fast",), - config={ - "configurable": { - "thread_id": "3", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 1, - "writes": {START: {"my_key": "key"}}, - "assistant_id": "b", - "thread_id": "3", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(tool_two.checkpointer.list(thread3, limit=2))[-1].config - ), - ) - # resume, for same result as above - assert tool_two.invoke(None, thread3, debug=1) == { - "my_key": "valuekey fast", - "market": "US", - } - assert tool_two.get_state(thread3) == StateSnapshot( - values={"my_key": "valuekey fast", "market": "US"}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "3", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 2, - "writes": {"tool_two_fast": {"my_key": " fast"}}, - "assistant_id": "b", - "thread_id": "3", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(tool_two.checkpointer.list(thread3, limit=2))[-1].config - ), - ) - - @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_branch_then( snapshot: SnapshotAssertion, request: pytest.FixtureRequest, checkpointer_name: str diff --git a/libs/langgraph/tests/test_large_cases_async.py b/libs/langgraph/tests/test_large_cases_async.py deleted file mode 100644 index 250cbfe72..000000000 --- a/libs/langgraph/tests/test_large_cases_async.py +++ /dev/null @@ -1,7395 +0,0 @@ -import asyncio -import operator -import re -import sys -from contextlib import asynccontextmanager -from typing import ( - Annotated, - Any, - AsyncIterator, - Literal, - Optional, - Union, - cast, -) - -import httpx -import pytest -from langchain_core.messages import ToolCall -from langchain_core.runnables import RunnableConfig, RunnablePick -from pydantic import BaseModel -from pytest_mock import MockerFixture -from syrupy import SnapshotAssertion -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.constants import END, PULL, PUSH, START -from langgraph.graph.graph import Graph -from langgraph.graph.message import MessageGraph, add_messages -from langgraph.graph.state import StateGraph -from langgraph.managed.shared_value import SharedValue -from langgraph.prebuilt.chat_agent_executor import create_react_agent -from langgraph.prebuilt.tool_node import ToolNode -from langgraph.pregel import Channel, Pregel -from langgraph.store.memory import InMemoryStore -from langgraph.types import PregelTask, Send, StateSnapshot, StreamWriter -from tests.any_str import AnyDict, AnyStr -from tests.conftest import ( - ALL_CHECKPOINTERS_ASYNC, - REGULAR_CHECKPOINTERS_ASYNC, - awith_checkpointer, -) -from tests.fake_chat import FakeChatModel -from tests.fake_tracer import FakeTracer -from tests.messages import ( - _AnyIdAIMessage, - _AnyIdAIMessageChunk, - _AnyIdHumanMessage, - _AnyIdToolMessage, -) - -pytestmark = pytest.mark.anyio - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_invoke_two_processes_in_out_interrupt( - checkpointer_name: str, mocker: MockerFixture -) -> None: - add_one = mocker.Mock(side_effect=lambda x: x + 1) - one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") - two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output") - async with awith_checkpointer(checkpointer_name) as checkpointer: - app = Pregel( - nodes={"one": one, "two": two}, - channels={ - "inbox": LastValue(int), - "output": LastValue(int), - "input": LastValue(int), - }, - input_channels="input", - output_channels="output", - checkpointer=checkpointer, - interrupt_after_nodes=["one"], - ) - thread1 = {"configurable": {"thread_id": "1"}} - thread2 = {"configurable": {"thread_id": "2"}} - - # start execution, stop at inbox - assert await app.ainvoke(2, thread1) is None - - # inbox == 3 - checkpoint = await checkpointer.aget(thread1) - assert checkpoint is not None - assert checkpoint["channel_values"]["inbox"] == 3 - - # resume execution, finish - assert await app.ainvoke(None, thread1) == 4 - - # start execution again, stop at inbox - assert await app.ainvoke(20, thread1) is None - - # inbox == 21 - checkpoint = await checkpointer.aget(thread1) - assert checkpoint is not None - assert checkpoint["channel_values"]["inbox"] == 21 - - # send a new value in, interrupting the previous execution - assert await app.ainvoke(3, thread1) is None - assert await app.ainvoke(None, thread1) == 5 - - # start execution again, stopping at inbox - assert await app.ainvoke(20, thread2) is None - - # inbox == 21 - snapshot = await app.aget_state(thread2) - assert snapshot.values["inbox"] == 21 - assert snapshot.next == ("two",) - - # update the state, resume - await app.aupdate_state(thread2, 25, as_node="one") - assert await app.ainvoke(None, thread2) == 26 - - # no pending tasks - snapshot = await app.aget_state(thread2) - assert snapshot.next == () - - if "shallow" in checkpointer_name: - return - - # list history - history = [c async for c in app.aget_state_history(thread1)] - assert history == [ - StateSnapshot( - values={"inbox": 4, "output": 5, "input": 3}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 6, - "writes": {"two": 5}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[1].config, - ), - StateSnapshot( - values={"inbox": 4, "output": 4, "input": 3}, - tasks=( - PregelTask(AnyStr(), "two", (PULL, "two"), result={"output": 5}), - ), - next=("two",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 5, - "writes": {"one": None}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[2].config, - ), - StateSnapshot( - values={"inbox": 21, "output": 4, "input": 3}, - tasks=( - PregelTask(AnyStr(), "one", (PULL, "one"), result={"inbox": 4}), - ), - next=("one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "step": 4, - "writes": {"input": 3}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[3].config, - ), - StateSnapshot( - values={"inbox": 21, "output": 4, "input": 20}, - tasks=(PregelTask(AnyStr(), "two", (PULL, "two")),), - next=("two",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"one": None}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[4].config, - ), - StateSnapshot( - values={"inbox": 3, "output": 4, "input": 20}, - tasks=( - PregelTask(AnyStr(), "one", (PULL, "one"), result={"inbox": 21}), - ), - next=("one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "step": 2, - "writes": {"input": 20}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[5].config, - ), - StateSnapshot( - values={"inbox": 3, "output": 4, "input": 2}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"two": 4}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[6].config, - ), - StateSnapshot( - values={"inbox": 3, "input": 2}, - tasks=( - PregelTask(AnyStr(), "two", (PULL, "two"), result={"output": 4}), - ), - next=("two",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": {"one": None}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[7].config, - ), - StateSnapshot( - values={"input": 2}, - tasks=( - PregelTask(AnyStr(), "one", (PULL, "one"), result={"inbox": 3}), - ), - next=("one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "step": -1, - "writes": {"input": 2}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - - # forking from any previous checkpoint should re-run nodes - assert [ - c async for c in app.astream(None, history[0].config, stream_mode="updates") - ] == [] - assert [ - c async for c in app.astream(None, history[1].config, stream_mode="updates") - ] == [ - {"two": {"output": 5}}, - ] - assert [ - c async for c in app.astream(None, history[2].config, stream_mode="updates") - ] == [ - {"one": {"inbox": 4}}, - {"__interrupt__": ()}, - ] - - -@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC) -async def test_fork_always_re_runs_nodes( - checkpointer_name: str, mocker: MockerFixture -) -> None: - 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) - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer) - - 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), - ] - - # 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}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[1].config, - ), - 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}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[2].config, - ), - 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}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[3].config, - ), - 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}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[4].config, - ), - 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}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[5].config, - ), - 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, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[6].config, - ), - 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}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - - # 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) -async def test_conditional_graph(checkpointer_name: str) -> None: - from langchain_core.agents import AgentAction, AgentFinish - from langchain_core.language_models.fake import FakeStreamingListLLM - from langchain_core.prompts import PromptTemplate - from langchain_core.runnables import RunnablePassthrough - from langchain_core.tools import tool - - # Assemble the tools - @tool() - def search_api(query: str) -> str: - """Searches the API for the query.""" - return f"result for {query}" - - tools = [search_api] - - # Construct the agent - prompt = PromptTemplate.from_template("Hello!") - - llm = FakeStreamingListLLM( - responses=[ - "tool:search_api:query", - "tool:search_api:another", - "finish:answer", - ] - ) - - async def agent_parser(input: str) -> Union[AgentAction, AgentFinish]: - if input.startswith("finish"): - _, answer = input.split(":") - return AgentFinish(return_values={"answer": answer}, log=input) - else: - _, tool_name, tool_input = input.split(":") - return AgentAction(tool=tool_name, tool_input=tool_input, log=input) - - agent = RunnablePassthrough.assign(agent_outcome=prompt | llm | agent_parser) - - # Define tool execution logic - async def execute_tools(data: dict) -> dict: - data = data.copy() - agent_action: AgentAction = data.pop("agent_outcome") - observation = await {t.name: t for t in tools}[agent_action.tool].ainvoke( - agent_action.tool_input - ) - if data.get("intermediate_steps") is None: - data["intermediate_steps"] = [] - else: - data["intermediate_steps"] = data["intermediate_steps"].copy() - data["intermediate_steps"].append([agent_action, observation]) - return data - - # Define decision-making logic - async def should_continue(data: dict, config: RunnableConfig) -> str: - # Logic to decide whether to continue in the loop or exit - if isinstance(data["agent_outcome"], AgentFinish): - return "exit" - else: - return "continue" - - # Define a new graph - workflow = Graph() - - workflow.add_node("agent", agent) - workflow.add_node("tools", execute_tools) - - workflow.set_entry_point("agent") - - workflow.add_conditional_edges( - "agent", should_continue, {"continue": "tools", "exit": END} - ) - - workflow.add_edge("tools", "agent") - - app = workflow.compile() - - assert await app.ainvoke({"input": "what is weather in sf"}) == { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ], - [ - AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - "result for another", - ], - ], - "agent_outcome": AgentFinish( - return_values={"answer": "answer"}, log="finish:answer" - ), - } - - assert [c async for c in app.astream({"input": "what is weather in sf"})] == [ - { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ), - } - }, - { - "tools": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ] - ], - } - }, - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ] - ], - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - { - "tools": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ], - [ - AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - "result for another", - ], - ], - } - }, - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ], - [ - AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - "result for another", - ], - ], - "agent_outcome": AgentFinish( - return_values={"answer": "answer"}, log="finish:answer" - ), - } - }, - ] - - patches = [c async for c in app.astream_log({"input": "what is weather in sf"})] - patch_paths = {op["path"] for log in patches for op in log.ops} - - # Check that agent (one of the nodes) has its output streamed to the logs - assert "/logs/agent/streamed_output/-" in patch_paths - assert "/logs/agent:2/streamed_output/-" in patch_paths - assert "/logs/agent:3/streamed_output/-" in patch_paths - # Check that agent (one of the nodes) has its final output set in the logs - assert "/logs/agent/final_output" in patch_paths - assert "/logs/agent:2/final_output" in patch_paths - assert "/logs/agent:3/final_output" in patch_paths - assert [ - p["value"] - for log in patches - for p in log.ops - if p["path"] == "/logs/agent/final_output" - or p["path"] == "/logs/agent:2/final_output" - or p["path"] == "/logs/agent:3/final_output" - ] == [ - { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ), - }, - { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ] - ], - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - }, - { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ], - [ - AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - "result for another", - ], - ], - "agent_outcome": AgentFinish( - return_values={"answer": "answer"}, log="finish:answer" - ), - }, - ] - - async with awith_checkpointer(checkpointer_name) as checkpointer: - # test state get/update methods with interrupt_after - - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_after=["agent"], - ) - config = {"configurable": {"thread_id": "1"}} - - assert [ - c - async for c in app_w_interrupt.astream( - {"input": "what is weather in sf"}, config - ) - ] == [ - { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - } - ] - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - }, - }, - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=( - await app_w_interrupt.checkpointer.aget_tuple(config) - ).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": { - "agent": { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - } - }, - "thread_id": "1", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), - ) - - await app_w_interrupt.aupdate_state( - config, - { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "input": "what is weather in sf", - }, - ) - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "input": "what is weather in sf", - }, - }, - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 1, - "writes": { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "input": "what is weather in sf", - } - }, - "thread_id": "1", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), - ) - - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "input": "what is weather in sf", - }, - }, - { - "tools": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - } - }, - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - ] - - await app_w_interrupt.aupdate_state( - config, - { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ), - }, - ) - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ), - }, - }, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 4, - "writes": { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ), - } - }, - "thread_id": "1", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), - ) - - # test state get/update methods with interrupt_before - - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_before=["tools"], - ) - config = {"configurable": {"thread_id": "2"}} - llm.i = 0 - - assert [ - c - async for c in app_w_interrupt.astream( - {"input": "what is weather in sf"}, config - ) - ] == [ - { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - } - ] - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - }, - }, - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": { - "agent": { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - } - }, - "thread_id": "2", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), - ) - - await app_w_interrupt.aupdate_state( - config, - { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "input": "what is weather in sf", - }, - ) - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "input": "what is weather in sf", - }, - }, - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 1, - "writes": { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "input": "what is weather in sf", - } - }, - "thread_id": "2", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), - ) - - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "input": "what is weather in sf", - }, - }, - { - "tools": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - } - }, - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - ] - - await app_w_interrupt.aupdate_state( - config, - { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ), - }, - ) - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ), - }, - }, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 4, - "writes": { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ), - } - }, - "thread_id": "2", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), - ) - - # test re-invoke to continue with interrupt_before - - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_before=["tools"], - ) - config = {"configurable": {"thread_id": "3"}} - llm.i = 0 # reset the llm - - assert [ - c - async for c in app_w_interrupt.astream( - {"input": "what is weather in sf"}, config - ) - ] == [ - { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - } - ] - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - }, - }, - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config={ - "configurable": { - "thread_id": "3", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": { - "agent": { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - } - }, - "thread_id": "3", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), - ) - - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - }, - }, - { - "tools": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ] - ], - } - }, - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ] - ], - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - ] - - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ] - ], - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - { - "tools": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ], - [ - AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - "result for another", - ], - ], - } - }, - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ], - [ - AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - "result for another", - ], - ], - "agent_outcome": AgentFinish( - return_values={"answer": "answer"}, log="finish:answer" - ), - } - }, - ] - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_conditional_graph_state( - mocker: MockerFixture, checkpointer_name: str -) -> None: - from langchain_core.agents import AgentAction, AgentFinish - from langchain_core.language_models.fake import FakeStreamingListLLM - from langchain_core.prompts import PromptTemplate - from langchain_core.tools import tool - - setup = mocker.Mock() - teardown = mocker.Mock() - - @asynccontextmanager - async def assert_ctx_once() -> AsyncIterator[None]: - assert setup.call_count == 0 - assert teardown.call_count == 0 - try: - yield - finally: - assert setup.call_count == 1 - assert teardown.call_count == 1 - setup.reset_mock() - teardown.reset_mock() - - class MyPydanticContextModel(BaseModel, arbitrary_types_allowed=True): - session: httpx.AsyncClient - something_else: str - - @asynccontextmanager - async def make_context( - config: RunnableConfig, - ) -> AsyncIterator[MyPydanticContextModel]: - assert isinstance(config, dict) - setup() - session = httpx.AsyncClient() - try: - yield MyPydanticContextModel(session=session, something_else="hello") - finally: - await session.aclose() - teardown() - - class AgentState(TypedDict): - input: Annotated[str, UntrackedValue] - agent_outcome: Optional[Union[AgentAction, AgentFinish]] - intermediate_steps: Annotated[list[tuple[AgentAction, str]], operator.add] - context: Annotated[MyPydanticContextModel, Context(make_context)] - - # Assemble the tools - @tool() - def search_api(query: str) -> str: - """Searches the API for the query.""" - return f"result for {query}" - - tools = [search_api] - - # Construct the agent - prompt = PromptTemplate.from_template("Hello!") - - llm = FakeStreamingListLLM( - responses=[ - "tool:search_api:query", - "tool:search_api:another", - "finish:answer", - ] - ) - - def agent_parser(input: str) -> dict[str, Union[AgentAction, AgentFinish]]: - if input.startswith("finish"): - _, answer = input.split(":") - return { - "agent_outcome": AgentFinish( - return_values={"answer": answer}, log=input - ) - } - else: - _, tool_name, tool_input = input.split(":") - return { - "agent_outcome": AgentAction( - tool=tool_name, tool_input=tool_input, log=input - ) - } - - agent = prompt | llm | agent_parser - - # Define tool execution logic - def execute_tools(data: AgentState) -> dict: - # check we have httpx session in AgentState - assert isinstance(data["context"], MyPydanticContextModel) - # execute the tool - agent_action: AgentAction = data.pop("agent_outcome") - observation = {t.name: t for t in tools}[agent_action.tool].invoke( - agent_action.tool_input - ) - return {"intermediate_steps": [[agent_action, observation]]} - - # Define decision-making logic - def should_continue(data: AgentState) -> str: - # check we have httpx session in AgentState - assert isinstance(data["context"], MyPydanticContextModel) - # Logic to decide whether to continue in the loop or exit - if isinstance(data["agent_outcome"], AgentFinish): - return "exit" - else: - return "continue" - - # Define a new graph - workflow = StateGraph(AgentState) - - workflow.add_node("agent", agent) - workflow.add_node("tools", execute_tools) - - workflow.set_entry_point("agent") - - workflow.add_conditional_edges( - "agent", should_continue, {"continue": "tools", "exit": END} - ) - - workflow.add_edge("tools", "agent") - - app = workflow.compile() - - async with assert_ctx_once(): - assert await app.ainvoke({"input": "what is weather in sf"}) == { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ], - [ - AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - "result for another", - ], - ], - "agent_outcome": AgentFinish( - return_values={"answer": "answer"}, log="finish:answer" - ), - } - - async with assert_ctx_once(): - assert [c async for c in app.astream({"input": "what is weather in sf"})] == [ - { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - }, - { - "tools": { - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ] - ], - } - }, - { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - { - "tools": { - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - "result for another", - ], - ], - } - }, - { - "agent": { - "agent_outcome": AgentFinish( - return_values={"answer": "answer"}, log="finish:answer" - ), - } - }, - ] - - async with assert_ctx_once(): - patches = [c async for c in app.astream_log({"input": "what is weather in sf"})] - patch_paths = {op["path"] for log in patches for op in log.ops} - - # Check that agent (one of the nodes) has its output streamed to the logs - assert "/logs/agent/streamed_output/-" in patch_paths - # Check that agent (one of the nodes) has its final output set in the logs - assert "/logs/agent/final_output" in patch_paths - assert [ - p["value"] - for log in patches - for p in log.ops - if p["path"] == "/logs/agent/final_output" - or p["path"] == "/logs/agent:2/final_output" - or p["path"] == "/logs/agent:3/final_output" - ] == [ - { - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ) - }, - { - "agent_outcome": AgentAction( - tool="search_api", tool_input="another", log="tool:search_api:another" - ) - }, - { - "agent_outcome": AgentFinish( - return_values={"answer": "answer"}, log="finish:answer" - ), - }, - ] - - async with awith_checkpointer(checkpointer_name) as checkpointer: - # test state get/update methods with interrupt_after - - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_after=["agent"], - ) - config = {"configurable": {"thread_id": "1"}} - - async with assert_ctx_once(): - assert [ - c - async for c in app_w_interrupt.astream( - {"input": "what is weather in sf"}, config - ) - ] == [ - { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - }, - {"__interrupt__": ()}, - ] - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "intermediate_steps": [], - }, - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - }, - "thread_id": "1", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), - ) - - async with assert_ctx_once(): - await app_w_interrupt.aupdate_state( - config, - { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ) - }, - ) - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "intermediate_steps": [], - }, - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 2, - "writes": { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ) - } - }, - "thread_id": "1", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), - ) - - async with assert_ctx_once(): - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - { - "tools": { - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - } - }, - { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - {"__interrupt__": ()}, - ] - - async with assert_ctx_once(): - await app_w_interrupt.aupdate_state( - config, - { - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ) - }, - ) - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ), - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - }, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 5, - "writes": { - "agent": { - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ) - } - }, - "thread_id": "1", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), - ) - - # test state get/update methods with interrupt_before - - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_before=["tools"], - ) - config = {"configurable": {"thread_id": "2"}} - llm.i = 0 # reset the llm - - assert [ - c - async for c in app_w_interrupt.astream( - {"input": "what is weather in sf"}, config - ) - ] == [ - { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - }, - {"__interrupt__": ()}, - ] - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ), - "intermediate_steps": [], - }, - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - }, - "thread_id": "2", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), - ) - - await app_w_interrupt.aupdate_state( - config, - { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ) - }, - ) - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "intermediate_steps": [], - }, - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 2, - "writes": { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ) - } - }, - "thread_id": "2", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), - ) - - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - { - "tools": { - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - } - }, - { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - {"__interrupt__": ()}, - ] - - await app_w_interrupt.aupdate_state( - config, - { - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ) - }, - ) - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ), - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - }, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 5, - "writes": { - "agent": { - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ) - } - }, - "thread_id": "2", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), - ) - - -async def test_prebuilt_tool_chat() -> None: - from langchain_core.messages import AIMessage, HumanMessage - from langchain_core.tools import tool - - model = FakeChatModel( - messages=[ - AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - }, - ], - ), - AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call234", - "name": "search_api", - "args": {"query": "another"}, - }, - { - "id": "tool_call567", - "name": "search_api", - "args": {"query": "a third one"}, - }, - ], - ), - AIMessage(content="answer"), - ] - ) - - @tool() - def search_api(query: str) -> str: - """Searches the API for the query.""" - return f"result for {query}" - - tools = [search_api] - - app = create_react_agent(model, tools) - - assert await app.ainvoke( - {"messages": [HumanMessage(content="what is weather in sf")]} - ) == { - "messages": [ - _AnyIdHumanMessage(content="what is weather in sf"), - _AnyIdAIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - }, - ], - ), - _AnyIdToolMessage( - content="result for query", - name="search_api", - tool_call_id="tool_call123", - ), - _AnyIdAIMessage( - content="", - tool_calls=[ - { - "id": "tool_call234", - "name": "search_api", - "args": {"query": "another"}, - }, - { - "id": "tool_call567", - "name": "search_api", - "args": {"query": "a third one"}, - }, - ], - ), - _AnyIdToolMessage( - content="result for another", - name="search_api", - tool_call_id="tool_call234", - ), - _AnyIdToolMessage( - content="result for a third one", - name="search_api", - tool_call_id="tool_call567", - id=AnyStr(), - ), - _AnyIdAIMessage(content="answer"), - ] - } - - assert [ - c - async for c in app.astream( - {"messages": [HumanMessage(content="what is weather in sf")]}, - stream_mode="messages", - ) - ] == [ - ( - _AnyIdAIMessageChunk( - content="", - tool_calls=[ - { - "name": "search_api", - "args": {"query": "query"}, - "id": "tool_call123", - "type": "tool_call", - } - ], - tool_call_chunks=[ - { - "name": "search_api", - "args": '{"query": "query"}', - "id": "tool_call123", - "index": None, - "type": "tool_call_chunk", - } - ], - ), - { - "langgraph_step": 1, - "langgraph_node": "agent", - "langgraph_triggers": ["start:agent"], - "langgraph_path": ("__pregel_pull", "agent"), - "langgraph_checkpoint_ns": AnyStr("agent:"), - "checkpoint_ns": AnyStr("agent:"), - "ls_provider": "fakechatmodel", - "ls_model_type": "chat", - }, - ), - ( - _AnyIdToolMessage( - content="result for query", - name="search_api", - tool_call_id="tool_call123", - ), - { - "langgraph_step": 2, - "langgraph_node": "tools", - "langgraph_triggers": ["branch:agent:should_continue:tools"], - "langgraph_path": ("__pregel_pull", "tools"), - "langgraph_checkpoint_ns": AnyStr("tools:"), - }, - ), - ( - _AnyIdAIMessageChunk( - content="", - tool_calls=[ - { - "name": "search_api", - "args": {"query": "another"}, - "id": "tool_call234", - "type": "tool_call", - }, - { - "name": "search_api", - "args": {"query": "a third one"}, - "id": "tool_call567", - "type": "tool_call", - }, - ], - tool_call_chunks=[ - { - "name": "search_api", - "args": '{"query": "another"}', - "id": "tool_call234", - "index": None, - "type": "tool_call_chunk", - }, - { - "name": "search_api", - "args": '{"query": "a third one"}', - "id": "tool_call567", - "index": None, - "type": "tool_call_chunk", - }, - ], - ), - { - "langgraph_step": 3, - "langgraph_node": "agent", - "langgraph_triggers": ["tools"], - "langgraph_path": ("__pregel_pull", "agent"), - "langgraph_checkpoint_ns": AnyStr("agent:"), - "checkpoint_ns": AnyStr("agent:"), - "ls_provider": "fakechatmodel", - "ls_model_type": "chat", - }, - ), - ( - _AnyIdToolMessage( - content="result for another", - name="search_api", - tool_call_id="tool_call234", - ), - { - "langgraph_step": 4, - "langgraph_node": "tools", - "langgraph_triggers": ["branch:agent:should_continue:tools"], - "langgraph_path": ("__pregel_pull", "tools"), - "langgraph_checkpoint_ns": AnyStr("tools:"), - }, - ), - ( - _AnyIdToolMessage( - content="result for a third one", - name="search_api", - tool_call_id="tool_call567", - ), - { - "langgraph_step": 4, - "langgraph_node": "tools", - "langgraph_triggers": ["branch:agent:should_continue:tools"], - "langgraph_path": ("__pregel_pull", "tools"), - "langgraph_checkpoint_ns": AnyStr("tools:"), - }, - ), - ( - _AnyIdAIMessageChunk( - content="answer", - ), - { - "langgraph_step": 5, - "langgraph_node": "agent", - "langgraph_triggers": ["tools"], - "langgraph_path": ("__pregel_pull", "agent"), - "langgraph_checkpoint_ns": AnyStr("agent:"), - "checkpoint_ns": AnyStr("agent:"), - "ls_provider": "fakechatmodel", - "ls_model_type": "chat", - }, - ), - ] - - assert [ - c - async for c in app.astream( - {"messages": [HumanMessage(content="what is weather in sf")]} - ) - ] == [ - { - "agent": { - "messages": [ - _AnyIdAIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - }, - ], - ) - ] - } - }, - { - "tools": { - "messages": [ - _AnyIdToolMessage( - content="result for query", - name="search_api", - tool_call_id="tool_call123", - ) - ] - } - }, - { - "agent": { - "messages": [ - _AnyIdAIMessage( - content="", - tool_calls=[ - { - "id": "tool_call234", - "name": "search_api", - "args": {"query": "another"}, - }, - { - "id": "tool_call567", - "name": "search_api", - "args": {"query": "a third one"}, - }, - ], - ) - ] - } - }, - { - "tools": { - "messages": [ - _AnyIdToolMessage( - content="result for another", - name="search_api", - tool_call_id="tool_call234", - ), - _AnyIdToolMessage( - content="result for a third one", - name="search_api", - tool_call_id="tool_call567", - ), - ] - } - }, - {"agent": {"messages": [_AnyIdAIMessage(content="answer")]}}, - ] - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_state_graph_packets(checkpointer_name: str) -> None: - from langchain_core.language_models.fake_chat_models import ( - FakeMessagesListChatModel, - ) - from langchain_core.messages import ( - AIMessage, - BaseMessage, - HumanMessage, - ToolMessage, - ) - from langchain_core.tools import tool - - class AgentState(TypedDict): - messages: Annotated[list[BaseMessage], add_messages] - session: Annotated[httpx.AsyncClient, Context(httpx.AsyncClient)] - - @tool() - def search_api(query: str) -> str: - """Searches the API for the query.""" - return f"result for {query}" - - tools = [search_api] - tools_by_name = {t.name: t for t in tools} - - model = FakeMessagesListChatModel( - responses=[ - AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - }, - ], - ), - AIMessage( - id="ai2", - content="", - tool_calls=[ - { - "id": "tool_call234", - "name": "search_api", - "args": {"query": "another", "idx": 0}, - }, - { - "id": "tool_call567", - "name": "search_api", - "args": {"query": "a third one", "idx": 1}, - }, - ], - ), - AIMessage(id="ai3", content="answer"), - ] - ) - - # Define decision-making logic - def should_continue(data: AgentState) -> str: - assert isinstance(data["session"], httpx.AsyncClient) - # Logic to decide whether to continue in the loop or exit - if tool_calls := data["messages"][-1].tool_calls: - return [Send("tools", tool_call) for tool_call in tool_calls] - else: - return END - - async def tools_node(input: ToolCall, config: RunnableConfig) -> AgentState: - await asyncio.sleep(input["args"].get("idx", 0) / 10) - output = await tools_by_name[input["name"]].ainvoke(input["args"], config) - return { - "messages": ToolMessage( - content=output, name=input["name"], tool_call_id=input["id"] - ) - } - - # Define a new graph - workflow = StateGraph(AgentState) - - # Define the two nodes we will cycle between - workflow.add_node("agent", {"messages": RunnablePick("messages") | model}) - workflow.add_node("tools", tools_node) - - # Set the entrypoint as `agent` - # This means that this node is the first one called - workflow.set_entry_point("agent") - - # We now add a conditional edge - workflow.add_conditional_edges("agent", should_continue) - - # We now add a normal edge from `tools` to `agent`. - # This means that after `tools` is called, `agent` node is called next. - workflow.add_edge("tools", "agent") - - # Finally, we compile it! - # This compiles it into a LangChain Runnable, - # meaning you can use it as you would any other runnable - app = workflow.compile() - - assert await app.ainvoke( - {"messages": HumanMessage(content="what is weather in sf")} - ) == { - "messages": [ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - }, - ], - ), - _AnyIdToolMessage( - content="result for query", - name="search_api", - tool_call_id="tool_call123", - ), - AIMessage( - id="ai2", - content="", - tool_calls=[ - { - "id": "tool_call234", - "name": "search_api", - "args": {"query": "another", "idx": 0}, - }, - { - "id": "tool_call567", - "name": "search_api", - "args": {"query": "a third one", "idx": 1}, - }, - ], - ), - _AnyIdToolMessage( - content="result for another", - name="search_api", - tool_call_id="tool_call234", - ), - _AnyIdToolMessage( - content="result for a third one", - name="search_api", - tool_call_id="tool_call567", - ), - AIMessage(content="answer", id="ai3"), - ] - } - - assert [ - c - async for c in app.astream( - {"messages": [HumanMessage(content="what is weather in sf")]} - ) - ] == [ - { - "agent": { - "messages": AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - }, - ], - ) - }, - }, - { - "tools": { - "messages": _AnyIdToolMessage( - content="result for query", - name="search_api", - tool_call_id="tool_call123", - ) - } - }, - { - "agent": { - "messages": AIMessage( - id="ai2", - content="", - tool_calls=[ - { - "id": "tool_call234", - "name": "search_api", - "args": {"query": "another", "idx": 0}, - }, - { - "id": "tool_call567", - "name": "search_api", - "args": {"query": "a third one", "idx": 1}, - }, - ], - ) - } - }, - { - "tools": { - "messages": _AnyIdToolMessage( - content="result for another", - name="search_api", - tool_call_id="tool_call234", - ) - }, - }, - { - "tools": { - "messages": _AnyIdToolMessage( - content="result for a third one", - name="search_api", - tool_call_id="tool_call567", - ), - }, - }, - {"agent": {"messages": AIMessage(content="answer", id="ai3")}}, - ] - - async with awith_checkpointer(checkpointer_name) as checkpointer: - # interrupt after agent - - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_after=["agent"], - ) - config = {"configurable": {"thread_id": "1"}} - - assert [ - c - async for c in app_w_interrupt.astream( - {"messages": HumanMessage(content="what is weather in sf")}, config - ) - ] == [ - { - "agent": { - "messages": AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - }, - ], - ) - } - }, - {"__interrupt__": ()}, - ] - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - }, - ], - ), - ] - }, - tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),), - next=("tools",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": { - "agent": { - "messages": AIMessage( - content="", - id="ai1", - tool_calls=[ - { - "name": "search_api", - "args": {"query": "query"}, - "id": "tool_call123", - "type": "tool_call", - } - ], - ) - } - }, - "thread_id": "1", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), - ) - - # modify ai message - last_message = (await app_w_interrupt.aget_state(config)).values["messages"][-1] - last_message.tool_calls[0]["args"]["query"] = "a different query" - await app_w_interrupt.aupdate_state(config, {"messages": last_message}) - - # message was replaced instead of appended - tup = await app_w_interrupt.checkpointer.aget_tuple(config) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - }, - ], - ), - ] - }, - tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),), - next=("tools",), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 2, - "writes": { - "agent": { - "messages": AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - }, - ], - ) - } - }, - "thread_id": "1", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), - ) - - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - { - "tools": { - "messages": _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", - ) - } - }, - { - "agent": { - "messages": AIMessage( - id="ai2", - content="", - tool_calls=[ - { - "id": "tool_call234", - "name": "search_api", - "args": {"query": "another", "idx": 0}, - }, - { - "id": "tool_call567", - "name": "search_api", - "args": {"query": "a third one", "idx": 1}, - }, - ], - ) - }, - }, - {"__interrupt__": ()}, - ] - - tup = await app_w_interrupt.checkpointer.aget_tuple(config) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - }, - ], - ), - _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", - ), - AIMessage( - id="ai2", - content="", - tool_calls=[ - { - "id": "tool_call234", - "name": "search_api", - "args": {"query": "another", "idx": 0}, - }, - { - "id": "tool_call567", - "name": "search_api", - "args": {"query": "a third one", "idx": 1}, - }, - ], - ), - ] - }, - tasks=( - PregelTask(AnyStr(), "tools", (PUSH, 0)), - PregelTask(AnyStr(), "tools", (PUSH, 1)), - ), - next=("tools", "tools"), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 4, - "writes": { - "agent": { - "messages": AIMessage( - id="ai2", - content="", - tool_calls=[ - { - "id": "tool_call234", - "name": "search_api", - "args": {"query": "another", "idx": 0}, - }, - { - "id": "tool_call567", - "name": "search_api", - "args": {"query": "a third one", "idx": 1}, - }, - ], - ), - }, - }, - "thread_id": "1", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), - ) - - await app_w_interrupt.aupdate_state( - config, - {"messages": AIMessage(content="answer", id="ai2")}, - ) - - # replaces message even if object identity is different, as long as id is the same - tup = await app_w_interrupt.checkpointer.aget_tuple(config) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - }, - ], - ), - _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", - ), - AIMessage(content="answer", id="ai2"), - ] - }, - tasks=(), - next=(), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 5, - "writes": { - "agent": { - "messages": AIMessage(content="answer", id="ai2"), - } - }, - "thread_id": "1", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), - ) - - # interrupt before tools - - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_before=["tools"], - ) - config = {"configurable": {"thread_id": "2"}} - model.i = 0 - - assert [ - c - async for c in app_w_interrupt.astream( - {"messages": HumanMessage(content="what is weather in sf")}, config - ) - ] == [ - { - "agent": { - "messages": AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - }, - ], - ) - } - }, - {"__interrupt__": ()}, - ] - tup = await app_w_interrupt.checkpointer.aget_tuple(config) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - }, - ], - ), - ] - }, - tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),), - next=("tools",), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": { - "agent": { - "messages": AIMessage( - content="", - additional_kwargs={}, - response_metadata={}, - id="ai1", - tool_calls=[ - { - "name": "search_api", - "args": {"query": "query"}, - "id": "tool_call123", - "type": "tool_call", - } - ], - ) - } - }, - "thread_id": "2", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), - ) - - # modify ai message - last_message = (await app_w_interrupt.aget_state(config)).values["messages"][-1] - last_message.tool_calls[0]["args"]["query"] = "a different query" - await app_w_interrupt.aupdate_state(config, {"messages": last_message}) - - # message was replaced instead of appended - tup = await app_w_interrupt.checkpointer.aget_tuple(config) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - }, - ], - ), - ] - }, - tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),), - next=("tools",), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 2, - "writes": { - "agent": { - "messages": AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - }, - ], - ) - } - }, - "thread_id": "2", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), - ) - - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - { - "tools": { - "messages": _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", - ) - } - }, - { - "agent": { - "messages": AIMessage( - id="ai2", - content="", - tool_calls=[ - { - "id": "tool_call234", - "name": "search_api", - "args": {"query": "another", "idx": 0}, - }, - { - "id": "tool_call567", - "name": "search_api", - "args": {"query": "a third one", "idx": 1}, - }, - ], - ) - }, - }, - {"__interrupt__": ()}, - ] - - tup = await app_w_interrupt.checkpointer.aget_tuple(config) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - }, - ], - ), - _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", - ), - AIMessage( - id="ai2", - content="", - tool_calls=[ - { - "id": "tool_call234", - "name": "search_api", - "args": {"query": "another", "idx": 0}, - }, - { - "id": "tool_call567", - "name": "search_api", - "args": {"query": "a third one", "idx": 1}, - }, - ], - ), - ] - }, - tasks=( - PregelTask(AnyStr(), "tools", (PUSH, 0)), - PregelTask(AnyStr(), "tools", (PUSH, 1)), - ), - next=("tools", "tools"), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 4, - "writes": { - "agent": { - "messages": AIMessage( - id="ai2", - content="", - tool_calls=[ - { - "id": "tool_call234", - "name": "search_api", - "args": {"query": "another", "idx": 0}, - }, - { - "id": "tool_call567", - "name": "search_api", - "args": {"query": "a third one", "idx": 1}, - }, - ], - ), - }, - }, - "thread_id": "2", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), - ) - - await app_w_interrupt.aupdate_state( - config, - {"messages": AIMessage(content="answer", id="ai2")}, - ) - - # replaces message even if object identity is different, as long as id is the same - tup = await app_w_interrupt.checkpointer.aget_tuple(config) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - }, - ], - ), - _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", - ), - AIMessage(content="answer", id="ai2"), - ] - }, - tasks=(), - next=(), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 5, - "writes": { - "agent": { - "messages": AIMessage(content="answer", id="ai2"), - } - }, - "thread_id": "2", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), - ) - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_message_graph(checkpointer_name: str) -> None: - from langchain_core.language_models.fake_chat_models import ( - FakeMessagesListChatModel, - ) - from langchain_core.messages import AIMessage, HumanMessage - from langchain_core.tools import tool - - class FakeFuntionChatModel(FakeMessagesListChatModel): - def bind_functions(self, functions: list): - return self - - @tool() - def search_api(query: str) -> str: - """Searches the API for the query.""" - return f"result for {query}" - - tools = [search_api] - - model = FakeFuntionChatModel( - responses=[ - AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - } - ], - id="ai1", - ), - AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call456", - "name": "search_api", - "args": {"query": "another"}, - } - ], - id="ai2", - ), - AIMessage(content="answer", id="ai3"), - ] - ) - - # Define the function that determines whether to continue or not - def should_continue(messages): - last_message = messages[-1] - # If there is no function call, then we finish - if not last_message.tool_calls: - return "end" - # Otherwise if there is, we continue - else: - return "continue" - - # Define a new graph - workflow = MessageGraph() - - # Define the two nodes we will cycle between - workflow.add_node("agent", model) - workflow.add_node("tools", ToolNode(tools)) - - # Set the entrypoint as `agent` - # This means that this node is the first one called - workflow.set_entry_point("agent") - - # We now add a conditional edge - workflow.add_conditional_edges( - # First, we define the start node. We use `agent`. - # This means these are the edges taken after the `agent` node is called. - "agent", - # Next, we pass in the function that will determine which node is called next. - should_continue, - # Finally we pass in a mapping. - # The keys are strings, and the values are other nodes. - # END is a special node marking that the graph should finish. - # What will happen is we will call `should_continue`, and then the output of that - # will be matched against the keys in this mapping. - # Based on which one it matches, that node will then be called. - { - # If `tools`, then we call the tool node. - "continue": "tools", - # Otherwise we finish. - "end": END, - }, - ) - - # We now add a normal edge from `tools` to `agent`. - # This means that after `tools` is called, `agent` node is called next. - workflow.add_edge("tools", "agent") - - # Finally, we compile it! - # This compiles it into a LangChain Runnable, - # meaning you can use it as you would any other runnable - app = workflow.compile() - - assert await app.ainvoke(HumanMessage(content="what is weather in sf")) == [ - _AnyIdHumanMessage( - content="what is weather in sf", - ), - AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - } - ], - id="ai1", # respects ids passed in - ), - _AnyIdToolMessage( - content="result for query", - name="search_api", - tool_call_id="tool_call123", - ), - AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call456", - "name": "search_api", - "args": {"query": "another"}, - } - ], - id="ai2", - ), - _AnyIdToolMessage( - content="result for another", - name="search_api", - tool_call_id="tool_call456", - ), - AIMessage(content="answer", id="ai3"), - ] - - assert [ - c async for c in app.astream([HumanMessage(content="what is weather in sf")]) - ] == [ - { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - } - ], - id="ai1", - ) - }, - { - "tools": [ - _AnyIdToolMessage( - content="result for query", - name="search_api", - tool_call_id="tool_call123", - ) - ] - }, - { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call456", - "name": "search_api", - "args": {"query": "another"}, - } - ], - id="ai2", - ) - }, - { - "tools": [ - _AnyIdToolMessage( - content="result for another", - name="search_api", - tool_call_id="tool_call456", - ) - ] - }, - {"agent": AIMessage(content="answer", id="ai3")}, - ] - - async with awith_checkpointer(checkpointer_name) as checkpointer: - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_after=["agent"], - ) - config = {"configurable": {"thread_id": "1"}} - - assert [ - c - async for c in app_w_interrupt.astream( - HumanMessage(content="what is weather in sf"), config - ) - ] == [ - { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - } - ], - id="ai1", - ) - }, - {"__interrupt__": ()}, - ] - - tup = await app_w_interrupt.checkpointer.aget_tuple(config) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values=[ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - } - ], - id="ai1", - ), - ], - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - } - ], - id="ai1", - ) - }, - "thread_id": "1", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), - ) - - # modify ai message - last_message = (await app_w_interrupt.aget_state(config)).values[-1] - last_message.tool_calls[0]["args"] = {"query": "a different query"} - await app_w_interrupt.aupdate_state(config, last_message) - - # message was replaced instead of appended - tup = await app_w_interrupt.checkpointer.aget_tuple(config) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values=[ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - content="", - id="ai1", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - } - ], - ), - ], - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 2, - "writes": { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - } - ], - id="ai1", - ) - }, - "thread_id": "1", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), - ) - - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - { - "tools": [ - _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", - ) - ] - }, - { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call456", - "name": "search_api", - "args": {"query": "another"}, - } - ], - id="ai2", - ) - }, - {"__interrupt__": ()}, - ] - - tup = await app_w_interrupt.checkpointer.aget_tuple(config) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values=[ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - content="", - id="ai1", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - } - ], - ), - _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", - ), - AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call456", - "name": "search_api", - "args": {"query": "another"}, - } - ], - id="ai2", - ), - ], - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 4, - "writes": { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call456", - "name": "search_api", - "args": {"query": "another"}, - } - ], - id="ai2", - ) - }, - "thread_id": "1", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), - ) - - await app_w_interrupt.aupdate_state( - config, - AIMessage(content="answer", id="ai2"), - ) - - # replaces message even if object identity is different, as long as id is the same - tup = await app_w_interrupt.checkpointer.aget_tuple(config) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values=[ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - content="", - id="ai1", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - } - ], - ), - _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", - ), - AIMessage(content="answer", id="ai2"), - ], - tasks=(), - next=(), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 5, - "writes": {"agent": AIMessage(content="answer", id="ai2")}, - "thread_id": "1", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), - ) - - -async def test_in_one_fan_out_out_one_graph_state() -> None: - def sorted_add(x: list[str], y: list[str]) -> list[str]: - return sorted(operator.add(x, y)) - - class State(TypedDict, total=False): - query: str - answer: str - docs: Annotated[list[str], operator.add] - - async def rewrite_query(data: State) -> State: - return {"query": f'query: {data["query"]}'} - - async def retriever_one(data: State) -> State: - await asyncio.sleep(0.1) - return {"docs": ["doc1", "doc2"]} - - async def retriever_two(data: State) -> State: - return {"docs": ["doc3", "doc4"]} - - async def qa(data: State) -> State: - return {"answer": ",".join(data["docs"])} - - workflow = StateGraph(State) - - workflow.add_node("rewrite_query", rewrite_query) - workflow.add_node("retriever_one", retriever_one) - workflow.add_node("retriever_two", retriever_two) - workflow.add_node("qa", qa) - - workflow.set_entry_point("rewrite_query") - workflow.add_edge("rewrite_query", "retriever_one") - workflow.add_edge("rewrite_query", "retriever_two") - workflow.add_edge("retriever_one", "qa") - workflow.add_edge("retriever_two", "qa") - workflow.set_finish_point("qa") - - app = workflow.compile() - - assert await app.ainvoke({"query": "what is weather in sf"}) == { - "query": "query: what is weather in sf", - "docs": ["doc1", "doc2", "doc3", "doc4"], - "answer": "doc1,doc2,doc3,doc4", - } - - assert [c async for c in app.astream({"query": "what is weather in sf"})] == [ - {"rewrite_query": {"query": "query: what is weather in sf"}}, - {"retriever_two": {"docs": ["doc3", "doc4"]}}, - {"retriever_one": {"docs": ["doc1", "doc2"]}}, - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] - - assert [ - c - async for c in app.astream( - {"query": "what is weather in sf"}, stream_mode="values" - ) - ] == [ - {"query": "what is weather in sf", "docs": []}, - {"query": "query: what is weather in sf", "docs": []}, - { - "query": "query: what is weather in sf", - "docs": ["doc1", "doc2", "doc3", "doc4"], - }, - { - "query": "query: what is weather in sf", - "docs": ["doc1", "doc2", "doc3", "doc4"], - "answer": "doc1,doc2,doc3,doc4", - }, - ] - - assert [ - c - async for c in app.astream( - {"query": "what is weather in sf"}, - stream_mode=["values", "updates", "debug"], - ) - ] == [ - ("values", {"query": "what is weather in sf", "docs": []}), - ( - "debug", - { - "type": "task", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": AnyStr(), - "name": "rewrite_query", - "input": {"query": "what is weather in sf", "docs": []}, - "triggers": ["start:rewrite_query"], - }, - }, - ), - ("updates", {"rewrite_query": {"query": "query: what is weather in sf"}}), - ( - "debug", - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": AnyStr(), - "name": "rewrite_query", - "result": [("query", "query: what is weather in sf")], - "error": None, - "interrupts": [], - }, - }, - ), - ("values", {"query": "query: what is weather in sf", "docs": []}), - ( - "debug", - { - "type": "task", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "id": AnyStr(), - "name": "retriever_one", - "input": {"query": "query: what is weather in sf", "docs": []}, - "triggers": ["rewrite_query"], - }, - }, - ), - ( - "debug", - { - "type": "task", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "id": AnyStr(), - "name": "retriever_two", - "input": {"query": "query: what is weather in sf", "docs": []}, - "triggers": ["rewrite_query"], - }, - }, - ), - ( - "updates", - {"retriever_two": {"docs": ["doc3", "doc4"]}}, - ), - ( - "debug", - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "id": AnyStr(), - "name": "retriever_two", - "result": [("docs", ["doc3", "doc4"])], - "error": None, - "interrupts": [], - }, - }, - ), - ( - "updates", - {"retriever_one": {"docs": ["doc1", "doc2"]}}, - ), - ( - "debug", - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "id": AnyStr(), - "name": "retriever_one", - "result": [("docs", ["doc1", "doc2"])], - "error": None, - "interrupts": [], - }, - }, - ), - ( - "values", - { - "query": "query: what is weather in sf", - "docs": ["doc1", "doc2", "doc3", "doc4"], - }, - ), - ( - "debug", - { - "type": "task", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "id": AnyStr(), - "name": "qa", - "input": { - "query": "query: what is weather in sf", - "docs": ["doc1", "doc2", "doc3", "doc4"], - }, - "triggers": ["retriever_one", "retriever_two"], - }, - }, - ), - ("updates", {"qa": {"answer": "doc1,doc2,doc3,doc4"}}), - ( - "debug", - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "id": AnyStr(), - "name": "qa", - "result": [("answer", "doc1,doc2,doc3,doc4")], - "error": None, - "interrupts": [], - }, - }, - ), - ( - "values", - { - "query": "query: what is weather in sf", - "answer": "doc1,doc2,doc3,doc4", - "docs": ["doc1", "doc2", "doc3", "doc4"], - }, - ), - ] - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_start_branch_then(checkpointer_name: str) -> None: - class State(TypedDict): - my_key: Annotated[str, operator.add] - market: str - shared: Annotated[dict[str, dict[str, Any]], SharedValue.on("assistant_id")] - other: Annotated[dict[str, dict[str, Any]], SharedValue.on("assistant_id")] - - def assert_shared_value(data: State, config: RunnableConfig) -> State: - assert "shared" in data - if thread_id := config["configurable"].get("thread_id"): - if thread_id == "1": - # this is the first thread, so should not see a value - assert data["shared"] == {} - return {"shared": {"1": {"hello": "world"}}, "other": {"2": {1: 2}}} - elif thread_id == "2": - # this should get value saved by thread 1 - assert data["shared"] == {"1": {"hello": "world"}} - elif thread_id == "3": - # this is a different assistant, so should not see previous value - assert data["shared"] == {} - return {} - - def tool_two_slow(data: State, config: RunnableConfig) -> State: - return {"my_key": " slow", **assert_shared_value(data, config)} - - def tool_two_fast(data: State, config: RunnableConfig) -> State: - return {"my_key": " fast", **assert_shared_value(data, config)} - - tool_two_graph = StateGraph(State) - tool_two_graph.add_node("tool_two_slow", tool_two_slow) - tool_two_graph.add_node("tool_two_fast", tool_two_fast) - tool_two_graph.set_conditional_entry_point( - lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast", then=END - ) - tool_two = tool_two_graph.compile() - - assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}) == { - "my_key": "value slow", - "market": "DE", - } - assert await tool_two.ainvoke({"my_key": "value", "market": "US"}) == { - "my_key": "value fast", - "market": "US", - } - - async with awith_checkpointer(checkpointer_name) as checkpointer: - tool_two = tool_two_graph.compile( - store=InMemoryStore(), - checkpointer=checkpointer, - interrupt_before=["tool_two_fast", "tool_two_slow"], - ) - - # missing thread_id - with pytest.raises(ValueError, match="thread_id"): - await tool_two.ainvoke({"my_key": "value", "market": "DE"}) - - thread1 = {"configurable": {"thread_id": "1", "assistant_id": "a"}} - # stop when about to enter node - assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}, thread1) == { - "my_key": "value", - "market": "DE", - } - if "shallow" not in checkpointer_name: - assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [ - { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "assistant_id": "a", - "thread_id": "1", - }, - { - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": {"my_key": "value", "market": "DE"}}, - "assistant_id": "a", - "thread_id": "1", - }, - ] - - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "tool_two_slow", (PULL, "tool_two_slow")),), - next=("tool_two_slow",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "assistant_id": "a", - "thread_id": "1", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ - -1 - ].config - ), - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread1, debug=1) == { - "my_key": "value slow", - "market": "DE", - } - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value slow", "market": "DE"}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"tool_two_slow": {"my_key": " slow"}}, - "assistant_id": "a", - "thread_id": "1", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ - -1 - ].config - ), - ) - - thread2 = {"configurable": {"thread_id": "2", "assistant_id": "a"}} - # stop when about to enter node - assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread2) == { - "my_key": "value", - "market": "US", - } - assert await tool_two.aget_state(thread2) == StateSnapshot( - values={"my_key": "value", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),), - next=("tool_two_fast",), - config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "assistant_id": "a", - "thread_id": "2", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ - -1 - ].config - ), - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread2, debug=1) == { - "my_key": "value fast", - "market": "US", - } - assert await tool_two.aget_state(thread2) == StateSnapshot( - values={"my_key": "value fast", "market": "US"}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"tool_two_fast": {"my_key": " fast"}}, - "assistant_id": "a", - "thread_id": "2", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ - -1 - ].config - ), - ) - - thread3 = {"configurable": {"thread_id": "3", "assistant_id": "b"}} - # stop when about to enter node - assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread3) == { - "my_key": "value", - "market": "US", - } - assert await tool_two.aget_state(thread3) == StateSnapshot( - values={"my_key": "value", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),), - next=("tool_two_fast",), - config={ - "configurable": { - "thread_id": "3", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "assistant_id": "b", - "thread_id": "3", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread3, limit=2)][ - -1 - ].config - ), - ) - # update state - await tool_two.aupdate_state(thread3, {"my_key": "key"}) # appends to my_key - assert await tool_two.aget_state(thread3) == StateSnapshot( - values={"my_key": "valuekey", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),), - next=("tool_two_fast",), - config={ - "configurable": { - "thread_id": "3", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 1, - "writes": {START: {"my_key": "key"}}, - "assistant_id": "b", - "thread_id": "3", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread3, limit=2)][ - -1 - ].config - ), - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread3, debug=1) == { - "my_key": "valuekey fast", - "market": "US", - } - assert await tool_two.aget_state(thread3) == StateSnapshot( - values={"my_key": "valuekey fast", "market": "US"}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "3", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 2, - "writes": {"tool_two_fast": {"my_key": " fast"}}, - "assistant_id": "b", - "thread_id": "3", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread3, limit=2)][ - -1 - ].config - ), - ) - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_branch_then(checkpointer_name: str) -> None: - class State(TypedDict): - my_key: Annotated[str, operator.add] - market: str - - tool_two_graph = StateGraph(State) - tool_two_graph.set_entry_point("prepare") - tool_two_graph.set_finish_point("finish") - tool_two_graph.add_conditional_edges( - source="prepare", - path=lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast", - then="finish", - ) - tool_two_graph.add_node("prepare", lambda s: {"my_key": " prepared"}) - tool_two_graph.add_node("tool_two_slow", lambda s: {"my_key": " slow"}) - tool_two_graph.add_node("tool_two_fast", lambda s: {"my_key": " fast"}) - tool_two_graph.add_node("finish", lambda s: {"my_key": " finished"}) - tool_two = tool_two_graph.compile() - - assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}, debug=1) == { - "my_key": "value prepared slow finished", - "market": "DE", - } - assert await tool_two.ainvoke({"my_key": "value", "market": "US"}) == { - "my_key": "value prepared fast finished", - "market": "US", - } - - async with awith_checkpointer(checkpointer_name) as checkpointer: - # test stream_mode=debug - tool_two = tool_two_graph.compile(checkpointer=checkpointer) - thread10 = {"configurable": {"thread_id": "10"}} - assert [ - c - async for c in tool_two.astream( - {"my_key": "value", "market": "DE"}, thread10, stream_mode="debug" - ) - ] == [ - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": -1, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": {"my_key": ""}, - "metadata": { - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": {"my_key": "value", "market": "DE"}}, - "thread_id": "10", - }, - "parent_config": None, - "next": ["__start__"], - "tasks": [ - { - "id": AnyStr(), - "name": "__start__", - "interrupts": (), - "state": None, - } - ], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": { - "my_key": "value", - "market": "DE", - }, - "metadata": { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "thread_id": "10", - }, - "parent_config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "next": ["prepare"], - "tasks": [ - { - "id": AnyStr(), - "name": "prepare", - "interrupts": (), - "state": None, - } - ], - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": AnyStr(), - "name": "prepare", - "input": {"my_key": "value", "market": "DE"}, - "triggers": ["start:prepare"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": AnyStr(), - "name": "prepare", - "result": [("my_key", " prepared")], - "error": None, - "interrupts": [], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": { - "my_key": "value prepared", - "market": "DE", - }, - "metadata": { - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - "thread_id": "10", - }, - "parent_config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "next": ["tool_two_slow"], - "tasks": [ - { - "id": AnyStr(), - "name": "tool_two_slow", - "interrupts": (), - "state": None, - } - ], - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "id": AnyStr(), - "name": "tool_two_slow", - "input": {"my_key": "value prepared", "market": "DE"}, - "triggers": ["branch:prepare:condition:tool_two_slow"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "id": AnyStr(), - "name": "tool_two_slow", - "result": [("my_key", " slow")], - "error": None, - "interrupts": [], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": { - "my_key": "value prepared slow", - "market": "DE", - }, - "metadata": { - "parents": {}, - "source": "loop", - "step": 2, - "writes": {"tool_two_slow": {"my_key": " slow"}}, - "thread_id": "10", - }, - "parent_config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "next": ["finish"], - "tasks": [ - { - "id": AnyStr(), - "name": "finish", - "interrupts": (), - "state": None, - } - ], - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "id": AnyStr(), - "name": "finish", - "input": {"my_key": "value prepared slow", "market": "DE"}, - "triggers": ["branch:prepare:condition::then"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "id": AnyStr(), - "name": "finish", - "result": [("my_key", " finished")], - "error": None, - "interrupts": [], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": { - "my_key": "value prepared slow finished", - "market": "DE", - }, - "metadata": { - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - "thread_id": "10", - }, - "parent_config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "next": [], - "tasks": [], - }, - }, - ] - - tool_two = tool_two_graph.compile( - checkpointer=checkpointer, - interrupt_before=["tool_two_fast", "tool_two_slow"], - ) - - # missing thread_id - with pytest.raises(ValueError, match="thread_id"): - await tool_two.ainvoke({"my_key": "value", "market": "DE"}) - - thread1 = {"configurable": {"thread_id": "11"}} - # stop when about to enter node - assert [ - c - async for c in tool_two.astream( - {"my_key": "value", "market": "DE"}, thread1, stream_mode="debug" - ) - ] == [ - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": -1, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "11"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "11", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": {"my_key": ""}, - "metadata": { - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": {"my_key": "value", "market": "DE"}}, - "thread_id": "11", - }, - "parent_config": None, - "next": ["__start__"], - "tasks": [ - { - "id": AnyStr(), - "name": "__start__", - "interrupts": (), - "state": None, - } - ], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "11"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "11", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": { - "my_key": "value", - "market": "DE", - }, - "metadata": { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "thread_id": "11", - }, - "parent_config": { - "tags": [], - "metadata": {"thread_id": "11"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "11", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "next": ["prepare"], - "tasks": [ - { - "id": AnyStr(), - "name": "prepare", - "interrupts": (), - "state": None, - } - ], - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": AnyStr(), - "name": "prepare", - "input": {"my_key": "value", "market": "DE"}, - "triggers": ["start:prepare"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": AnyStr(), - "name": "prepare", - "result": [("my_key", " prepared")], - "error": None, - "interrupts": [], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "11"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "11", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": { - "my_key": "value prepared", - "market": "DE", - }, - "metadata": { - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - "thread_id": "11", - }, - "parent_config": { - "tags": [], - "metadata": {"thread_id": "11"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "11", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "next": ["tool_two_slow"], - "tasks": [ - { - "id": AnyStr(), - "name": "tool_two_slow", - "interrupts": (), - "state": None, - } - ], - }, - }, - ] - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value prepared", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "tool_two_slow", (PULL, "tool_two_slow")),), - next=("tool_two_slow",), - config={ - "configurable": { - "thread_id": "11", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - "thread_id": "11", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ - -1 - ].config - ), - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread1, debug=1) == { - "my_key": "value prepared slow finished", - "market": "DE", - } - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value prepared slow finished", "market": "DE"}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "11", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - "thread_id": "11", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ - -1 - ].config - ), - ) - - thread2 = {"configurable": {"thread_id": "12"}} - # stop when about to enter node - assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread2) == { - "my_key": "value prepared", - "market": "US", - } - assert await tool_two.aget_state(thread2) == StateSnapshot( - values={"my_key": "value prepared", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),), - next=("tool_two_fast",), - config={ - "configurable": { - "thread_id": "12", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - "thread_id": "12", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ - -1 - ].config - ), - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread2, debug=1) == { - "my_key": "value prepared fast finished", - "market": "US", - } - assert await tool_two.aget_state(thread2) == StateSnapshot( - values={"my_key": "value prepared fast finished", "market": "US"}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "12", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - "thread_id": "12", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ - -1 - ].config - ), - ) - - tool_two = tool_two_graph.compile( - checkpointer=checkpointer, interrupt_after=["prepare"] - ) - - # missing thread_id - with pytest.raises(ValueError, match="thread_id"): - await tool_two.ainvoke({"my_key": "value", "market": "DE"}) - - thread1 = {"configurable": {"thread_id": "21"}} - # stop when about to enter node - assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}, thread1) == { - "my_key": "value prepared", - "market": "DE", - } - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value prepared", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "tool_two_slow", (PULL, "tool_two_slow")),), - next=("tool_two_slow",), - config={ - "configurable": { - "thread_id": "21", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - "thread_id": "21", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ - -1 - ].config - ), - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread1, debug=1) == { - "my_key": "value prepared slow finished", - "market": "DE", - } - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value prepared slow finished", "market": "DE"}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "21", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - "thread_id": "21", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ - -1 - ].config - ), - ) - - thread2 = {"configurable": {"thread_id": "22"}} - # stop when about to enter node - assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread2) == { - "my_key": "value prepared", - "market": "US", - } - assert await tool_two.aget_state(thread2) == StateSnapshot( - values={"my_key": "value prepared", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),), - next=("tool_two_fast",), - config={ - "configurable": { - "thread_id": "22", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - "thread_id": "22", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ - -1 - ].config - ), - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread2, debug=1) == { - "my_key": "value prepared fast finished", - "market": "US", - } - assert await tool_two.aget_state(thread2) == StateSnapshot( - values={"my_key": "value prepared fast finished", "market": "US"}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "22", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - "thread_id": "22", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ - -1 - ].config - ), - ) - - thread3 = {"configurable": {"thread_id": "23"}} - # update an empty thread before first run - uconfig = await tool_two.aupdate_state( - thread3, {"my_key": "key", "market": "DE"} - ) - # check current state - assert await tool_two.aget_state(thread3) == StateSnapshot( - values={"my_key": "key", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "prepare", (PULL, "prepare")),), - next=("prepare",), - config=uconfig, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 0, - "writes": {START: {"my_key": "key", "market": "DE"}}, - "thread_id": "23", - }, - parent_config=None, - ) - # run from this point - assert await tool_two.ainvoke(None, thread3) == { - "my_key": "key prepared", - "market": "DE", - } - # get state after first node - assert await tool_two.aget_state(thread3) == StateSnapshot( - values={"my_key": "key prepared", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "tool_two_slow", (PULL, "tool_two_slow")),), - next=("tool_two_slow",), - config={ - "configurable": { - "thread_id": "23", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - "thread_id": "23", - }, - parent_config=(None if "shallow" in checkpointer_name else uconfig), - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread3, debug=1) == { - "my_key": "key prepared slow finished", - "market": "DE", - } - assert await tool_two.aget_state(thread3) == StateSnapshot( - values={"my_key": "key prepared slow finished", "market": "DE"}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "23", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - "thread_id": "23", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread3, limit=2)][ - -1 - ].config - ), - ) - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_nested_graph_state(checkpointer_name: str) -> None: - class InnerState(TypedDict): - my_key: str - my_other_key: str - - def inner_1(state: InnerState): - return { - "my_key": state["my_key"] + " here", - "my_other_key": state["my_key"], - } - - def inner_2(state: InnerState): - return { - "my_key": state["my_key"] + " and there", - "my_other_key": state["my_key"], - } - - inner = StateGraph(InnerState) - inner.add_node("inner_1", inner_1) - inner.add_node("inner_2", inner_2) - inner.add_edge("inner_1", "inner_2") - inner.set_entry_point("inner_1") - inner.set_finish_point("inner_2") - - class State(TypedDict): - my_key: str - other_parent_key: str - - def outer_1(state: State): - return {"my_key": "hi " + state["my_key"]} - - def outer_2(state: State): - return {"my_key": state["my_key"] + " and back again"} - - graph = StateGraph(State) - graph.add_node("outer_1", outer_1) - graph.add_node( - "inner", - inner.compile(interrupt_before=["inner_2"]), - ) - graph.add_node("outer_2", outer_2) - graph.set_entry_point("outer_1") - graph.add_edge("outer_1", "inner") - graph.add_edge("inner", "outer_2") - graph.set_finish_point("outer_2") - - async with awith_checkpointer(checkpointer_name) as checkpointer: - app = graph.compile(checkpointer=checkpointer) - - config = {"configurable": {"thread_id": "1"}} - await app.ainvoke({"my_key": "my value"}, config, debug=True) - # test state w/ nested subgraph state (right after interrupt) - # first get_state without subgraph state - assert await app.aget_state(config) == StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "inner", - (PULL, "inner"), - state={ - "configurable": {"thread_id": "1", "checkpoint_ns": AnyStr()} - }, - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), - ) - # now, get_state with subgraphs state - assert await app.aget_state(config, subgraphs=True) == StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "inner", - (PULL, "inner"), - state=StateSnapshot( - values={ - "my_key": "hi my value here", - "my_other_key": "hi my value", - }, - tasks=( - PregelTask( - AnyStr(), - "inner_2", - (PULL, "inner_2"), - ), - ), - next=("inner_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "parents": { - "": AnyStr(), - }, - "source": "loop", - "writes": { - "inner_1": { - "my_key": "hi my value here", - "my_other_key": "hi my value", - } - }, - "step": 1, - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "langgraph_node": "inner", - "langgraph_path": [PULL, "inner"], - "langgraph_step": 2, - "langgraph_triggers": ["outer_1"], - "langgraph_checkpoint_ns": AnyStr("inner:"), - }, - created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - } - ), - ), - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), - ) - # get_state_history returns outer graph checkpoints - history = [c async for c in app.aget_state_history(config)] - expected_history = [ - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "inner", - (PULL, "inner"), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - } - }, - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=( - PregelTask( - AnyStr(), - "outer_1", - (PULL, "outer_1"), - result={"my_key": "hi my value"}, - ), - ), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": None, - "step": 0, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=( - PregelTask( - AnyStr(), - "__start__", - (PULL, "__start__"), - result={"my_key": "my value"}, - ), - ), - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"__start__": {"my_key": "my value"}}, - "step": -1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - - if "shallow" in checkpointer_name: - expected_history = expected_history[:1] - - assert history == expected_history - - # get_state_history for a subgraph returns its checkpoints - child_history = [ - c async for c in app.aget_state_history(history[0].tasks[0].state) - ] - expected_child_history = [ - StateSnapshot( - values={"my_key": "hi my value here", "my_other_key": "hi my value"}, - next=("inner_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "source": "loop", - "writes": { - "inner_1": { - "my_key": "hi my value here", - "my_other_key": "hi my value", - } - }, - "step": 1, - "parents": {"": AnyStr()}, - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "langgraph_node": "inner", - "langgraph_path": [PULL, "inner"], - "langgraph_step": 2, - "langgraph_triggers": ["outer_1"], - "langgraph_checkpoint_ns": AnyStr("inner:"), - }, - created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - } - ), - tasks=(PregelTask(AnyStr(), "inner_2", (PULL, "inner_2")),), - ), - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "source": "loop", - "writes": None, - "step": 0, - "parents": {"": AnyStr()}, - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "langgraph_node": "inner", - "langgraph_path": [PULL, "inner"], - "langgraph_step": 2, - "langgraph_triggers": ["outer_1"], - "langgraph_checkpoint_ns": AnyStr("inner:"), - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - tasks=( - PregelTask( - AnyStr(), - "inner_1", - (PULL, "inner_1"), - result={ - "my_key": "hi my value here", - "my_other_key": "hi my value", - }, - ), - ), - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "source": "input", - "writes": {"__start__": {"my_key": "hi my value"}}, - "step": -1, - "parents": {"": AnyStr()}, - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "langgraph_node": "inner", - "langgraph_path": [PULL, "inner"], - "langgraph_step": 2, - "langgraph_triggers": ["outer_1"], - "langgraph_checkpoint_ns": AnyStr("inner:"), - }, - created_at=AnyStr(), - parent_config=None, - tasks=( - PregelTask( - AnyStr(), - "__start__", - (PULL, "__start__"), - result={"my_key": "hi my value"}, - ), - ), - ), - ] - - if "shallow" in checkpointer_name: - expected_child_history = expected_child_history[:1] - - assert child_history == expected_child_history - - # resume - await app.ainvoke(None, config, debug=True) - # test state w/ nested subgraph state (after resuming from interrupt) - assert await app.aget_state(config) == StateSnapshot( - values={"my_key": "hi my value here and there and back again"}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "outer_2": {"my_key": "hi my value here and there and back again"} - }, - "step": 3, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), - ) - # test full history at the end - actual_history = [c async for c in app.aget_state_history(config)] - expected_history = [ - StateSnapshot( - values={"my_key": "hi my value here and there and back again"}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "outer_2": { - "my_key": "hi my value here and there and back again" - } - }, - "step": 3, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), - ), - StateSnapshot( - values={"my_key": "hi my value here and there"}, - tasks=( - PregelTask( - AnyStr(), - "outer_2", - (PULL, "outer_2"), - result={"my_key": "hi my value here and there and back again"}, - ), - ), - next=("outer_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"inner": {"my_key": "hi my value here and there"}}, - "step": 2, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "inner", - (PULL, "inner"), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - } - }, - result={"my_key": "hi my value here and there"}, - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=( - PregelTask( - AnyStr(), - "outer_1", - (PULL, "outer_1"), - result={"my_key": "hi my value"}, - ), - ), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": None, - "step": 0, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=( - PregelTask( - AnyStr(), - "__start__", - (PULL, "__start__"), - result={"my_key": "my value"}, - ), - ), - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"__start__": {"my_key": "my value"}}, - "step": -1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - if "shallow" in checkpointer_name: - expected_history = expected_history[:1] - - assert actual_history == expected_history - # test looking up parent state by checkpoint ID - for actual_snapshot, expected_snapshot in zip(actual_history, expected_history): - assert await app.aget_state(actual_snapshot.config) == expected_snapshot - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: - class State(TypedDict): - my_key: str - - class ChildState(TypedDict): - my_key: str - - class GrandChildState(TypedDict): - my_key: str - - def grandchild_1(state: ChildState): - return {"my_key": state["my_key"] + " here"} - - def grandchild_2(state: ChildState): - return { - "my_key": state["my_key"] + " and there", - } - - grandchild = StateGraph(GrandChildState) - grandchild.add_node("grandchild_1", grandchild_1) - grandchild.add_node("grandchild_2", grandchild_2) - grandchild.add_edge("grandchild_1", "grandchild_2") - grandchild.set_entry_point("grandchild_1") - grandchild.set_finish_point("grandchild_2") - - child = StateGraph(ChildState) - child.add_node( - "child_1", - grandchild.compile(interrupt_before=["grandchild_2"]), - ) - child.set_entry_point("child_1") - child.set_finish_point("child_1") - - def parent_1(state: State): - return {"my_key": "hi " + state["my_key"]} - - def parent_2(state: State): - return {"my_key": state["my_key"] + " and back again"} - - graph = StateGraph(State) - graph.add_node("parent_1", parent_1) - graph.add_node("child", child.compile()) - graph.add_node("parent_2", parent_2) - graph.set_entry_point("parent_1") - graph.add_edge("parent_1", "child") - graph.add_edge("child", "parent_2") - graph.set_finish_point("parent_2") - - async with awith_checkpointer(checkpointer_name) as checkpointer: - app = graph.compile(checkpointer=checkpointer) - - # test invoke w/ nested interrupt - config = {"configurable": {"thread_id": "1"}} - assert [ - c async for c in app.astream({"my_key": "my value"}, config, subgraphs=True) - ] == [ - ((), {"parent_1": {"my_key": "hi my value"}}), - ( - (AnyStr("child:"), AnyStr("child_1:")), - {"grandchild_1": {"my_key": "hi my value here"}}, - ), - ((), {"__interrupt__": ()}), - ] - # get state without subgraphs - outer_state = await app.aget_state(config) - assert outer_state == StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "child", - (PULL, "child"), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child"), - } - }, - ), - ), - next=("child",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"parent_1": {"my_key": "hi my value"}}, - "step": 1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), - ) - child_state = await app.aget_state(outer_state.tasks[0].state) - assert ( - child_state.tasks[0] - == StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "child_1", - (PULL, "child_1"), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - } - }, - ), - ), - next=("child_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {"": AnyStr()}, - "source": "loop", - "writes": None, - "step": 0, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - } - } - ), - ).tasks[0] - ) - grandchild_state = await app.aget_state(child_state.tasks[0].state) - assert grandchild_state == StateSnapshot( - values={"my_key": "hi my value here"}, - tasks=( - PregelTask( - AnyStr(), - "grandchild_2", - (PULL, "grandchild_2"), - ), - ), - next=("grandchild_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), - } - ), - } - }, - metadata={ - "parents": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - } - ), - "source": "loop", - "writes": {"grandchild_1": {"my_key": "hi my value here"}}, - "step": 1, - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "langgraph_checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child_1", - "langgraph_path": [PULL, AnyStr("child_1")], - "langgraph_step": 1, - "langgraph_triggers": [AnyStr("start:child_1")], - }, - created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), - } - ), - } - } - ), - ) - # get state with subgraphs - assert await app.aget_state(config, subgraphs=True) == StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "child", - (PULL, "child"), - state=StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "child_1", - (PULL, "child_1"), - state=StateSnapshot( - values={"my_key": "hi my value here"}, - tasks=( - PregelTask( - AnyStr(), - "grandchild_2", - (PULL, "grandchild_2"), - ), - ), - next=("grandchild_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr( - re.compile(r"child:.+|child1:") - ): AnyStr(), - } - ), - } - }, - metadata={ - "parents": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - } - ), - "source": "loop", - "writes": { - "grandchild_1": { - "my_key": "hi my value here" - } - }, - "step": 1, - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "langgraph_checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child_1", - "langgraph_path": [ - PULL, - AnyStr("child_1"), - ], - "langgraph_step": 1, - "langgraph_triggers": [AnyStr("start:child_1")], - }, - created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr( - re.compile( - r"child:.+|child1:" - ) - ): AnyStr(), - } - ), - } - } - ), - ), - ), - ), - next=("child_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "parents": {"": AnyStr()}, - "source": "loop", - "writes": None, - "step": 0, - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child", - "langgraph_path": [PULL, AnyStr("child")], - "langgraph_step": 2, - "langgraph_triggers": [AnyStr("parent_1")], - "langgraph_checkpoint_ns": AnyStr("child:"), - }, - created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - } - ), - ), - ), - ), - next=("child",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"parent_1": {"my_key": "hi my value"}}, - "step": 1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), - ) - # resume - assert [c async for c in app.astream(None, config, subgraphs=True)] == [ - ( - (AnyStr("child:"), AnyStr("child_1:")), - {"grandchild_2": {"my_key": "hi my value here and there"}}, - ), - ( - (AnyStr("child:"),), - {"child_1": {"my_key": "hi my value here and there"}}, - ), - ((), {"child": {"my_key": "hi my value here and there"}}), - ((), {"parent_2": {"my_key": "hi my value here and there and back again"}}), - ] - # get state with and without subgraphs - assert ( - await app.aget_state(config) - == await app.aget_state(config, subgraphs=True) - == StateSnapshot( - values={"my_key": "hi my value here and there and back again"}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "parent_2": { - "my_key": "hi my value here and there and back again" - } - }, - "step": 3, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), - ) - ) - - if "shallow" in checkpointer_name: - return - - # get outer graph history - outer_history = [c async for c in app.aget_state_history(config)] - assert ( - outer_history[0] - == [ - StateSnapshot( - values={"my_key": "hi my value here and there and back again"}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "parent_2": { - "my_key": "hi my value here and there and back again" - } - }, - "step": 3, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value here and there"}, - next=("parent_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"child": {"my_key": "hi my value here and there"}}, - "step": 2, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - tasks=( - PregelTask( - id=AnyStr(), name="parent_2", path=(PULL, "parent_2") - ), - ), - ), - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "child", - (PULL, "child"), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child"), - } - }, - ), - ), - next=("child",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"parent_1": {"my_key": "hi my value"}}, - "step": 1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - next=("parent_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": None, - "step": 0, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - tasks=( - PregelTask( - id=AnyStr(), name="parent_1", path=(PULL, "parent_1") - ), - ), - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=None, - tasks=( - PregelTask( - id=AnyStr(), name="__start__", path=(PULL, "__start__") - ), - ), - ), - ][0] - ) - # get child graph history - child_history = [ - c async for c in app.aget_state_history(outer_history[2].tasks[0].state) - ] - assert child_history == [ - StateSnapshot( - values={"my_key": "hi my value here and there"}, - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "source": "loop", - "writes": {"child_1": {"my_key": "hi my value here and there"}}, - "step": 1, - "parents": {"": AnyStr()}, - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child", - "langgraph_path": [PULL, AnyStr("child")], - "langgraph_step": 2, - "langgraph_triggers": [AnyStr("parent_1")], - "langgraph_checkpoint_ns": AnyStr("child:"), - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - tasks=(), - ), - StateSnapshot( - values={"my_key": "hi my value"}, - next=("child_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "source": "loop", - "writes": None, - "step": 0, - "parents": {"": AnyStr()}, - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child", - "langgraph_path": [PULL, AnyStr("child")], - "langgraph_step": 2, - "langgraph_triggers": [AnyStr("parent_1")], - "langgraph_checkpoint_ns": AnyStr("child:"), - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - tasks=( - PregelTask( - id=AnyStr(), - name="child_1", - path=(PULL, "child_1"), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - } - }, - result={"my_key": "hi my value here and there"}, - ), - ), - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "source": "input", - "writes": {"__start__": {"my_key": "hi my value"}}, - "step": -1, - "parents": {"": AnyStr()}, - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child", - "langgraph_path": [PULL, AnyStr("child")], - "langgraph_step": 2, - "langgraph_triggers": [AnyStr("parent_1")], - "langgraph_checkpoint_ns": AnyStr("child:"), - }, - created_at=AnyStr(), - parent_config=None, - tasks=( - PregelTask( - id=AnyStr(), - name="__start__", - path=(PULL, "__start__"), - result={"my_key": "hi my value"}, - ), - ), - ), - ] - # get grandchild graph history - grandchild_history = [ - c async for c in app.aget_state_history(child_history[1].tasks[0].state) - ] - assert grandchild_history == [ - StateSnapshot( - values={"my_key": "hi my value here and there"}, - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), - } - ), - } - }, - metadata={ - "source": "loop", - "writes": { - "grandchild_2": {"my_key": "hi my value here and there"} - }, - "step": 2, - "parents": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - } - ), - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "langgraph_checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child_1", - "langgraph_path": [ - PULL, - AnyStr("child_1"), - ], - "langgraph_step": 1, - "langgraph_triggers": [AnyStr("start:child_1")], - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), - } - ), - } - }, - tasks=(), - ), - StateSnapshot( - values={"my_key": "hi my value here"}, - next=("grandchild_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), - } - ), - } - }, - metadata={ - "source": "loop", - "writes": {"grandchild_1": {"my_key": "hi my value here"}}, - "step": 1, - "parents": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - } - ), - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "langgraph_checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child_1", - "langgraph_path": [ - PULL, - AnyStr("child_1"), - ], - "langgraph_step": 1, - "langgraph_triggers": [AnyStr("start:child_1")], - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), - } - ), - } - }, - tasks=( - PregelTask( - id=AnyStr(), - name="grandchild_2", - path=(PULL, "grandchild_2"), - result={"my_key": "hi my value here and there"}, - ), - ), - ), - StateSnapshot( - values={"my_key": "hi my value"}, - next=("grandchild_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), - } - ), - } - }, - metadata={ - "source": "loop", - "writes": None, - "step": 0, - "parents": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - } - ), - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "langgraph_checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child_1", - "langgraph_path": [ - PULL, - AnyStr("child_1"), - ], - "langgraph_step": 1, - "langgraph_triggers": [AnyStr("start:child_1")], - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), - } - ), - } - }, - tasks=( - PregelTask( - id=AnyStr(), - name="grandchild_1", - path=(PULL, "grandchild_1"), - result={"my_key": "hi my value here"}, - ), - ), - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), - } - ), - } - }, - metadata={ - "source": "input", - "writes": {"__start__": {"my_key": "hi my value"}}, - "step": -1, - "parents": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - } - ), - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "langgraph_checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child_1", - "langgraph_path": [ - PULL, - AnyStr("child_1"), - ], - "langgraph_step": 1, - "langgraph_triggers": [AnyStr("start:child_1")], - }, - created_at=AnyStr(), - parent_config=None, - tasks=( - PregelTask( - id=AnyStr(), - name="__start__", - path=(PULL, "__start__"), - result={"my_key": "hi my value"}, - ), - ), - ), - ] - - # replay grandchild checkpoint - assert [ - c - async for c in app.astream( - None, grandchild_history[2].config, subgraphs=True - ) - ] == [ - ( - (AnyStr("child:"), AnyStr("child_1:")), - {"grandchild_1": {"my_key": "hi my value here"}}, - ), - ((), {"__interrupt__": ()}), - ] - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_send_to_nested_graphs(checkpointer_name: str) -> None: - class OverallState(TypedDict): - subjects: list[str] - jokes: Annotated[list[str], operator.add] - - async def continue_to_jokes(state: OverallState): - return [Send("generate_joke", {"subject": s}) for s in state["subjects"]] - - class JokeState(TypedDict): - subject: str - - async def edit(state: JokeState): - subject = state["subject"] - return {"subject": f"{subject} - hohoho"} - - # subgraph - subgraph = StateGraph(JokeState, output=OverallState) - subgraph.add_node("edit", edit) - subgraph.add_node( - "generate", lambda state: {"jokes": [f"Joke about {state['subject']}"]} - ) - subgraph.set_entry_point("edit") - subgraph.add_edge("edit", "generate") - subgraph.set_finish_point("generate") - - # parent graph - builder = StateGraph(OverallState) - builder.add_node( - "generate_joke", - subgraph.compile(interrupt_before=["generate"]), - ) - builder.add_conditional_edges(START, continue_to_jokes) - builder.add_edge("generate_joke", END) - - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer) - config = {"configurable": {"thread_id": "1"}} - tracer = FakeTracer() - - # invoke and pause at nested interrupt - assert await graph.ainvoke( - {"subjects": ["cats", "dogs"]}, - config={**config, "callbacks": [tracer]}, - ) == { - "subjects": ["cats", "dogs"], - "jokes": [], - } - assert len(tracer.runs) == 1, "Should produce exactly 1 root run" - - # check state - outer_state = await graph.aget_state(config) - - # update state of dogs joke graph - await graph.aupdate_state( - outer_state.tasks[1].state, {"subject": "turtles - hohoho"} - ) - - # continue past interrupt - assert await graph.ainvoke(None, config=config) == { - "subjects": ["cats", "dogs"], - "jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"], - } - - -@pytest.mark.skipif( - sys.version_info < (3, 11), - reason="Python 3.11+ is required for async contextvars support", -) -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_weather_subgraph( - checkpointer_name: str, snapshot: SnapshotAssertion -) -> None: - from langchain_core.language_models.fake_chat_models import ( - FakeMessagesListChatModel, - ) - from langchain_core.messages import AIMessage, ToolCall - from langchain_core.tools import tool - - from langgraph.graph import MessagesState - - # setup subgraph - - @tool - def get_weather(city: str): - """Get the weather for a specific city""" - return f"I'ts sunny in {city}!" - - weather_model = FakeMessagesListChatModel( - responses=[ - AIMessage( - content="", - tool_calls=[ - ToolCall( - id="tool_call123", - name="get_weather", - args={"city": "San Francisco"}, - ) - ], - ) - ] - ) - - class SubGraphState(MessagesState): - city: str - - def model_node(state: SubGraphState, writer: StreamWriter): - writer(" very") - result = weather_model.invoke(state["messages"]) - return {"city": cast(AIMessage, result).tool_calls[0]["args"]["city"]} - - def weather_node(state: SubGraphState, writer: StreamWriter): - writer(" good") - result = get_weather.invoke({"city": state["city"]}) - return {"messages": [{"role": "assistant", "content": result}]} - - subgraph = StateGraph(SubGraphState) - subgraph.add_node(model_node) - subgraph.add_node(weather_node) - subgraph.add_edge(START, "model_node") - subgraph.add_edge("model_node", "weather_node") - subgraph.add_edge("weather_node", END) - subgraph = subgraph.compile(interrupt_before=["weather_node"]) - - # setup main graph - - class RouterState(MessagesState): - route: Literal["weather", "other"] - - class Router(TypedDict): - route: Literal["weather", "other"] - - router_model = FakeMessagesListChatModel( - responses=[ - AIMessage( - content="", - tool_calls=[ - ToolCall( - id="tool_call123", - name="router", - args={"dest": "weather"}, - ) - ], - ) - ] - ) - - def router_node(state: RouterState, writer: StreamWriter): - writer("I'm") - system_message = "Classify the incoming query as either about weather or not." - messages = [{"role": "system", "content": system_message}] + state["messages"] - route = router_model.invoke(messages) - return {"route": cast(AIMessage, route).tool_calls[0]["args"]["dest"]} - - def normal_llm_node(state: RouterState): - return {"messages": [AIMessage("Hello!")]} - - def route_after_prediction(state: RouterState): - if state["route"] == "weather": - return "weather_graph" - else: - return "normal_llm_node" - - def weather_graph(state: RouterState): - # this tests that all async checkpointers tested also implement sync methods - # as the subgraph called with sync invoke will use sync checkpointer methods - return subgraph.invoke(state) - - graph = StateGraph(RouterState) - graph.add_node(router_node) - graph.add_node(normal_llm_node) - graph.add_node("weather_graph", weather_graph) - graph.add_edge(START, "router_node") - graph.add_conditional_edges("router_node", route_after_prediction) - graph.add_edge("normal_llm_node", END) - graph.add_edge("weather_graph", END) - - def get_first_in_list(): - return [*graph.get_state_history(config, limit=1)][0] - - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = graph.compile(checkpointer=checkpointer) - - assert graph.get_graph(xray=1).draw_mermaid() == snapshot - - config = {"configurable": {"thread_id": "1"}} - thread2 = {"configurable": {"thread_id": "2"}} - inputs = {"messages": [{"role": "user", "content": "what's the weather in sf"}]} - - # run with custom output - assert [ - c async for c in graph.astream(inputs, thread2, stream_mode="custom") - ] == [ - "I'm", - " very", - ] - assert [ - c async for c in graph.astream(None, thread2, stream_mode="custom") - ] == [ - " good", - ] - - # run until interrupt - assert [ - c - async for c in graph.astream( - inputs, config=config, stream_mode="updates", subgraphs=True - ) - ] == [ - ((), {"router_node": {"route": "weather"}}), - ((AnyStr("weather_graph:"),), {"model_node": {"city": "San Francisco"}}), - ((), {"__interrupt__": ()}), - ] - - # check current state - state = await graph.aget_state(config) - assert state == StateSnapshot( - values={ - "messages": [_AnyIdHumanMessage(content="what's the weather in sf")], - "route": "weather", - }, - next=("weather_graph",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"router_node": {"route": "weather"}}, - "step": 1, - "parents": {}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), - tasks=( - PregelTask( - id=AnyStr(), - name="weather_graph", - path=(PULL, "weather_graph"), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("weather_graph:"), - } - }, - ), - ), - ) - # confirm that list() delegates to alist() correctly - assert await asyncio.to_thread(get_first_in_list) == state - - # update - await graph.aupdate_state(state.tasks[0].state, {"city": "la"}) - - # run after update - assert [ - c - async for c in graph.astream( - None, config=config, stream_mode="updates", subgraphs=True - ) - ] == [ - ( - (AnyStr("weather_graph:"),), - { - "weather_node": { - "messages": [ - {"role": "assistant", "content": "I'ts sunny in la!"} - ] - } - }, - ), - ( - (), - { - "weather_graph": { - "messages": [ - _AnyIdHumanMessage(content="what's the weather in sf"), - _AnyIdAIMessage(content="I'ts sunny in la!"), - ] - } - }, - ), - ] - - # try updating acting as weather node - config = {"configurable": {"thread_id": "14"}} - inputs = {"messages": [{"role": "user", "content": "what's the weather in sf"}]} - assert [ - c - async for c in graph.astream( - inputs, config=config, stream_mode="updates", subgraphs=True - ) - ] == [ - ((), {"router_node": {"route": "weather"}}), - ((AnyStr("weather_graph:"),), {"model_node": {"city": "San Francisco"}}), - ((), {"__interrupt__": ()}), - ] - state = await graph.aget_state(config, subgraphs=True) - assert state == StateSnapshot( - values={ - "messages": [_AnyIdHumanMessage(content="what's the weather in sf")], - "route": "weather", - }, - next=("weather_graph",), - config={ - "configurable": { - "thread_id": "14", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"router_node": {"route": "weather"}}, - "step": 1, - "parents": {}, - "thread_id": "14", - }, - created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { - "configurable": { - "thread_id": "14", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), - tasks=( - PregelTask( - id=AnyStr(), - name="weather_graph", - path=(PULL, "weather_graph"), - state=StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="what's the weather in sf") - ], - "city": "San Francisco", - }, - next=("weather_node",), - config={ - "configurable": { - "thread_id": "14", - "checkpoint_ns": AnyStr("weather_graph:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("weather_graph:"): AnyStr(), - } - ), - } - }, - metadata={ - "source": "loop", - "writes": {"model_node": {"city": "San Francisco"}}, - "step": 1, - "parents": {"": AnyStr()}, - "thread_id": "14", - "checkpoint_ns": AnyStr("weather_graph:"), - "langgraph_node": "weather_graph", - "langgraph_path": [PULL, "weather_graph"], - "langgraph_step": 2, - "langgraph_triggers": [ - "branch:router_node:route_after_prediction:weather_graph" - ], - "langgraph_checkpoint_ns": AnyStr("weather_graph:"), - }, - created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { - "configurable": { - "thread_id": "14", - "checkpoint_ns": AnyStr("weather_graph:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("weather_graph:"): AnyStr(), - } - ), - } - } - ), - tasks=( - PregelTask( - id=AnyStr(), - name="weather_node", - path=(PULL, "weather_node"), - ), - ), - ), - ), - ), - ) - await graph.aupdate_state( - state.tasks[0].state.config, - {"messages": [{"role": "assistant", "content": "rainy"}]}, - as_node="weather_node", - ) - state = await graph.aget_state(config, subgraphs=True) - assert state == StateSnapshot( - values={ - "messages": [_AnyIdHumanMessage(content="what's the weather in sf")], - "route": "weather", - }, - next=("weather_graph",), - config={ - "configurable": { - "thread_id": "14", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"router_node": {"route": "weather"}}, - "step": 1, - "parents": {}, - "thread_id": "14", - }, - created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { - "configurable": { - "thread_id": "14", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), - tasks=( - PregelTask( - id=AnyStr(), - name="weather_graph", - path=(PULL, "weather_graph"), - state=StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="what's the weather in sf"), - _AnyIdAIMessage(content="rainy"), - ], - "city": "San Francisco", - }, - next=(), - config={ - "configurable": { - "thread_id": "14", - "checkpoint_ns": AnyStr("weather_graph:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("weather_graph:"): AnyStr(), - } - ), - } - }, - metadata={ - "step": 2, - "source": "update", - "writes": { - "weather_node": { - "messages": [ - {"role": "assistant", "content": "rainy"} - ] - } - }, - "parents": {"": AnyStr()}, - "thread_id": "14", - "checkpoint_id": AnyStr(), - "checkpoint_ns": AnyStr("weather_graph:"), - "langgraph_node": "weather_graph", - "langgraph_path": [PULL, "weather_graph"], - "langgraph_step": 2, - "langgraph_triggers": [ - "branch:router_node:route_after_prediction:weather_graph" - ], - "langgraph_checkpoint_ns": AnyStr("weather_graph:"), - }, - created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { - "configurable": { - "thread_id": "14", - "checkpoint_ns": AnyStr("weather_graph:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("weather_graph:"): AnyStr(), - } - ), - } - } - ), - tasks=(), - ), - ), - ), - ) - assert [ - c - async for c in graph.astream( - None, config=config, stream_mode="updates", subgraphs=True - ) - ] == [ - ( - (), - { - "weather_graph": { - "messages": [ - _AnyIdHumanMessage(content="what's the weather in sf"), - _AnyIdAIMessage(content="rainy"), - ] - } - }, - ), - ] diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 862c859b5..12bd1fca3 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -1,91 +1,61 @@ import enum -import functools -import json import logging import operator import threading import time import uuid -import warnings from collections import Counter, deque from concurrent.futures import ThreadPoolExecutor -from contextlib import contextmanager from dataclasses import dataclass -from random import randrange from typing import ( Annotated, Any, Dict, - Generator, - Iterator, List, Literal, Optional, Sequence, - Tuple, Union, get_type_hints, ) -import httpx import pytest -from langchain_core.language_models import GenericFakeChatModel -from langchain_core.runnables import ( - RunnableConfig, - RunnableLambda, - RunnablePassthrough, -) -from langchain_core.runnables.graph import Edge -from langsmith import traceable +from langchain_core.runnables import RunnableConfig, RunnableLambda from pytest_mock import MockerFixture from syrupy import SnapshotAssertion from typing_extensions import TypedDict -from langgraph.channels.base import BaseChannel from langgraph.channels.binop import BinaryOperatorAggregate -from langgraph.channels.context import Context -from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic from langgraph.checkpoint.base import ( BaseCheckpointSaver, - Checkpoint, - CheckpointMetadata, CheckpointTuple, ) from langgraph.checkpoint.memory import InMemorySaver, MemorySaver from langgraph.config import get_stream_writer from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, START from langgraph.errors import InvalidUpdateError -from langgraph.func import entrypoint, task -from langgraph.graph import END, Graph, StateGraph -from langgraph.graph.message import MessageGraph, MessagesState, add_messages -from langgraph.prebuilt.tool_node import ToolNode -from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot +from langgraph.graph import END, StateGraph +from langgraph.graph.message import MessagesState, add_messages +from langgraph.pregel import GraphRecursionError, NodeBuilder, Pregel, StateSnapshot from langgraph.pregel.retry import RetryPolicy -from langgraph.store.base import BaseStore from langgraph.types import ( Command, Interrupt, PregelTask, Send, - StreamWriter, interrupt, ) +from langgraph.utils.runnable import RunnableSeq from tests.agents import AgentAction, AgentFinish from tests.any_str import AnyStr, AnyVersion, FloatBetween, UnsortedSequence from tests.conftest import ( ALL_CHECKPOINTERS_SYNC, - ALL_STORES_SYNC, REGULAR_CHECKPOINTERS_SYNC, - SHOULD_CHECK_SNAPSHOTS, ) -from tests.memory_assert import MemorySaverAssertCheckpointMetadata from tests.messages import ( - _AnyIdAIMessage, - _AnyIdAIMessageChunk, _AnyIdHumanMessage, - _AnyIdToolMessage, ) pytestmark = pytest.mark.anyio @@ -94,28 +64,31 @@ logger = logging.getLogger(__name__) def test_graph_validation() -> None: + class State(TypedDict): + hello: str + def logic(inp: str) -> str: return "" - workflow = Graph() + workflow = StateGraph(State) workflow.add_node("agent", logic) workflow.set_entry_point("agent") workflow.set_finish_point("agent") assert workflow.compile(), "valid graph" # Accept a dead-end - workflow = Graph() + workflow = StateGraph(State) workflow.add_node("agent", logic) workflow.set_entry_point("agent") workflow.compile() - workflow = Graph() + workflow = StateGraph(State) workflow.add_node("agent", logic) workflow.set_finish_point("agent") with pytest.raises(ValueError, match="must have an entrypoint"): workflow.compile() - workflow = Graph() + workflow = StateGraph(State) workflow.add_node("agent", logic) workflow.add_node("tools", logic) workflow.set_entry_point("agent") @@ -123,7 +96,7 @@ def test_graph_validation() -> None: workflow.add_edge("tools", "agent") assert workflow.compile(), "valid graph" - workflow = Graph() + workflow = StateGraph(State) workflow.add_node("agent", logic) workflow.add_node("tools", logic) workflow.set_entry_point("tools") @@ -131,7 +104,7 @@ def test_graph_validation() -> None: workflow.add_edge("tools", "agent") assert workflow.compile(), "valid graph" - workflow = Graph() + workflow = StateGraph(State) workflow.set_entry_point("tools") workflow.add_conditional_edges("agent", logic, {"continue": "tools", "exit": END}) workflow.add_edge("tools", "agent") @@ -139,7 +112,7 @@ def test_graph_validation() -> None: workflow.add_node("tools", logic) assert workflow.compile(), "valid graph" - workflow = Graph() + workflow = StateGraph(State) workflow.set_entry_point("tools") workflow.add_conditional_edges( "agent", logic, {"continue": "tools", "exit": END, "hmm": "extra"} @@ -150,7 +123,7 @@ def test_graph_validation() -> None: with pytest.raises(ValueError, match="unknown"): # extra is not defined workflow.compile() - workflow = Graph() + workflow = StateGraph(State) workflow.set_entry_point("agent") workflow.add_conditional_edges("agent", logic, {"continue": "tools", "exit": END}) workflow.add_edge("tools", "extra") @@ -159,7 +132,7 @@ def test_graph_validation() -> None: with pytest.raises(ValueError, match="unknown"): # extra is not defined workflow.compile() - workflow = Graph() + workflow = StateGraph(State) workflow.add_node("agent", logic) workflow.add_node("tools", logic) workflow.add_node("extra", logic) @@ -169,9 +142,6 @@ def test_graph_validation() -> None: # Accept, even though extra is dead-end workflow.compile() - class State(TypedDict): - hello: str - graph = StateGraph(State) graph.add_node("start", lambda x: x) graph.add_edge("__start__", "start") @@ -223,58 +193,6 @@ def test_graph_validation_with_command() -> None: assert graph.invoke({"foo": ""}) == {"foo": "bar", "bar": "baz"} -def test_checkpoint_errors() -> None: - class FaultyGetCheckpointer(InMemorySaver): - def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: - raise ValueError("Faulty get_tuple") - - class FaultyPutCheckpointer(MemorySaver): - def put( - self, - config: RunnableConfig, - checkpoint: Checkpoint, - metadata: CheckpointMetadata, - new_versions: Optional[dict[str, Union[str, int, float]]] = None, - ) -> RunnableConfig: - raise ValueError("Faulty put") - - class FaultyPutWritesCheckpointer(InMemorySaver): - def put_writes( - self, config: RunnableConfig, writes: List[Tuple[str, Any]], task_id: str - ) -> RunnableConfig: - raise ValueError("Faulty put_writes") - - class FaultyVersionCheckpointer(InMemorySaver): - def get_next_version(self, current: Optional[int], channel: BaseChannel) -> int: - raise ValueError("Faulty get_next_version") - - def logic(inp: str) -> str: - return "" - - builder = StateGraph(Annotated[str, operator.add]) - builder.add_node("agent", logic) - builder.add_edge(START, "agent") - - graph = builder.compile(checkpointer=FaultyGetCheckpointer()) - with pytest.raises(ValueError, match="Faulty get_tuple"): - graph.invoke("", {"configurable": {"thread_id": "thread-1"}}) - - graph = builder.compile(checkpointer=FaultyPutCheckpointer()) - with pytest.raises(ValueError, match="Faulty put"): - graph.invoke("", {"configurable": {"thread_id": "thread-1"}}) - - graph = builder.compile(checkpointer=FaultyVersionCheckpointer()) - with pytest.raises(ValueError, match="Faulty get_next_version"): - graph.invoke("", {"configurable": {"thread_id": "thread-1"}}) - - # add parallel node - builder.add_node("parallel", logic) - builder.add_edge(START, "parallel") - graph = builder.compile(checkpointer=FaultyPutWritesCheckpointer()) - with pytest.raises(ValueError, match="Faulty put_writes"): - graph.invoke("", {"configurable": {"thread_id": "thread-1"}}) - - def test_node_schemas_custom_output() -> None: class State(TypedDict): hello: str @@ -431,7 +349,9 @@ def test_reducer_before_first_node() -> None: def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") + chain = ( + NodeBuilder().subscribe_to("input").add_node(add_one).write_to("output").build() + ) app = Pregel( nodes={ @@ -444,55 +364,20 @@ def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: input_channels="input", output_channels="output", ) - graph = Graph() - graph.add_node("add_one", add_one) - graph.set_entry_point("add_one") - graph.set_finish_point("add_one") - gapp = graph.compile() - - if SHOULD_CHECK_SNAPSHOTS: - assert app.input_schema.model_json_schema() == { - "title": "LangGraphInput", - "type": "integer", - } - assert app.output_schema.model_json_schema() == { - "title": "LangGraphOutput", - "type": "integer", - } - with warnings.catch_warnings(): - warnings.simplefilter("error") # raise warnings as errors - assert app.config_schema().model_json_schema() == { - "properties": {}, - "title": "LangGraphConfig", - "type": "object", - } assert app.invoke(2) == 3 assert app.invoke(2, output_keys=["output"]) == {"output": 3} assert repr(app), "does not raise recursion error" - assert gapp.invoke(2, debug=True) == 3 - - -@pytest.mark.parametrize( - "falsy_value", - [None, False, 0, "", [], {}, set(), frozenset(), 0.0, 0j], -) -def test_invoke_single_process_in_out_falsy_values(falsy_value: Any) -> None: - graph = Graph() - graph.add_node("return_falsy_const", lambda *args, **kwargs: falsy_value) - graph.set_entry_point("return_falsy_const") - graph.set_finish_point("return_falsy_const") - gapp = graph.compile() - assert gapp.invoke(1) == falsy_value - def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) chain = ( - Channel.subscribe_to("input") - | add_one - | Channel.write_to("output", fixed=5, output_plus_one=lambda x: x + 1) + NodeBuilder() + .subscribe_to("input") + .add_node(add_one) + .write_to("output", fixed=5, output_plus_one=lambda x: x + 1) + .build() ) app = Pregel( @@ -507,30 +392,14 @@ def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> None: input_channels="input", ) - if SHOULD_CHECK_SNAPSHOTS: - assert app.input_schema.model_json_schema() == { - "title": "LangGraphInput", - "type": "integer", - } - assert app.output_schema.model_json_schema() == { - "title": "LangGraphOutput", - "type": "object", - "properties": { - "output": {"title": "Output", "type": "integer", "default": None}, - "fixed": {"title": "Fixed", "type": "integer", "default": None}, - "output_plus_one": { - "title": "Output Plus One", - "type": "integer", - "default": None, - }, - }, - } assert app.invoke(2) == {"output": 3, "fixed": 5, "output_plus_one": 4} def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") + chain = ( + NodeBuilder().subscribe_to("input").add_node(add_one).write_to("output").build() + ) app = Pregel( nodes={"one": chain}, @@ -539,24 +408,14 @@ def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None: output_channels=["output"], ) - if SHOULD_CHECK_SNAPSHOTS: - assert app.input_schema.model_json_schema() == { - "title": "LangGraphInput", - "type": "integer", - } - assert app.output_schema.model_json_schema() == { - "title": "LangGraphOutput", - "type": "object", - "properties": { - "output": {"title": "Output", "type": "integer", "default": None} - }, - } assert app.invoke(2) == {"output": 3} def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") + chain = ( + NodeBuilder().subscribe_to("input").add_node(add_one).write_to("output").build() + ) app = Pregel( nodes={"one": chain}, @@ -564,28 +423,18 @@ def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None: input_channels=["input"], output_channels=["output"], ) - if SHOULD_CHECK_SNAPSHOTS: - assert app.input_schema.model_json_schema() == { - "title": "LangGraphInput", - "type": "object", - "properties": { - "input": {"title": "Input", "type": "integer", "default": None} - }, - } - assert app.output_schema.model_json_schema() == { - "title": "LangGraphOutput", - "type": "object", - "properties": { - "output": {"title": "Output", "type": "integer", "default": None} - }, - } + assert app.invoke({"input": 2}) == {"output": 3} def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") - two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output") + one = ( + NodeBuilder().subscribe_to("input").add_node(add_one).write_to("inbox").build() + ) + two = ( + NodeBuilder().subscribe_to("inbox").add_node(add_one).write_to("output").build() + ) app = Pregel( nodes={"one": one, "two": two}, @@ -601,30 +450,7 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: assert app.invoke(2) == 4 with pytest.raises(GraphRecursionError): - app.invoke(2, {"recursion_limit": 1}, debug=1) - - graph = Graph() - graph.add_node("add_one", add_one) - graph.add_node("add_one_more", add_one) - graph.set_entry_point("add_one") - graph.set_finish_point("add_one_more") - graph.add_edge("add_one", "add_one_more") - gapp = graph.compile() - - assert gapp.invoke(2) == 4 - - for step, values in enumerate(gapp.stream(2, debug=1), start=1): - if step == 1: - assert values == { - "add_one": 3, - } - elif step == 2: - assert values == { - "add_one_more": 4, - } - else: - assert 0, f"{step}:{values}" - assert step == 2 + app.invoke(2, {"recursion_limit": 1}) @pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC) @@ -643,6 +469,7 @@ def test_run_from_checkpoint_id_retains_previous_writes( def __call__(self, state: MyState): self.switch = not self.switch + print("node", state) return {"myval": 2 if self.switch else 1, "otherval": self.switch} builder = StateGraph(MyState) @@ -655,6 +482,7 @@ def test_run_from_checkpoint_id_retains_previous_writes( swap = "node_one" if src == "node_two" else "node_two" def _edge(st: MyState) -> Literal["__end__", "node_one", "node_two"]: + print("edge", src, st) if st["myval"] > 3: return END if st["otherval"]: @@ -670,7 +498,7 @@ def test_run_from_checkpoint_id_retains_previous_writes( thread_id = uuid.uuid4() thread1 = {"configurable": {"thread_id": str(thread_id)}} - result = graph.invoke({"myval": 1}, thread1) + result = graph.invoke({"myval": 1}, thread1, log_mode=["updates", "values"]) assert result["myval"] == 4 history = [c for c in graph.get_state_history(thread1)] @@ -707,174 +535,28 @@ def test_run_from_checkpoint_id_retains_previous_writes( assert _get_tasks(new_history, 1) == _get_tasks(history, 0) -def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: - add_one = mocker.Mock(side_effect=lambda x: x + 1) - one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") - two = ( - Channel.subscribe_to("inbox") - | RunnableLambda(add_one).batch - | RunnablePassthrough(lambda _: time.sleep(0.1)) - | Channel.write_to("output").batch - ) - - app = Pregel( - nodes={"one": one, "two": two}, - channels={ - "inbox": Topic(int), - "output": LastValue(int), - "input": LastValue(int), - }, - input_channels=["input", "inbox"], - stream_channels=["output", "inbox"], - output_channels=["output"], - ) - - # [12 + 1, 2 + 1 + 1] - assert [ - *app.stream( - {"input": 2, "inbox": 12}, output_keys="output", stream_mode="updates" - ) - ] == [ - {"one": None}, - {"two": 13}, - {"two": 4}, - ] - assert [*app.stream({"input": 2, "inbox": 12}, output_keys="output")] == [ - 13, - 4, - ] - - assert [*app.stream({"input": 2, "inbox": 12}, stream_mode="updates")] == [ - {"one": {"inbox": 3}}, - {"two": {"output": 13}}, - {"two": {"output": 4}}, - ] - assert [*app.stream({"input": 2, "inbox": 12})] == [ - {"inbox": [3], "output": 13}, - {"output": 4}, - ] - assert [*app.stream({"input": 2, "inbox": 12}, stream_mode="debug")] == [ - { - "type": "task", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "id": AnyStr(), - "name": "one", - "input": 2, - "triggers": ["input"], - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "id": AnyStr(), - "name": "two", - "input": [12], - "triggers": ["inbox"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "id": AnyStr(), - "name": "one", - "result": [("inbox", 3)], - "error": None, - "interrupts": [], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "id": AnyStr(), - "name": "two", - "result": [("output", 13)], - "error": None, - "interrupts": [], - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": AnyStr(), - "name": "two", - "input": [3], - "triggers": ["inbox"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": AnyStr(), - "name": "two", - "result": [("output", 4)], - "error": None, - "interrupts": [], - }, - }, - ] - - -def test_batch_two_processes_in_out() -> None: - def add_one_with_delay(inp: int) -> int: - time.sleep(inp / 10) - return inp + 1 - - one = Channel.subscribe_to("input") | add_one_with_delay | Channel.write_to("one") - two = Channel.subscribe_to("one") | add_one_with_delay | Channel.write_to("output") - - app = Pregel( - nodes={"one": one, "two": two}, - channels={ - "one": LastValue(int), - "output": LastValue(int), - "input": LastValue(int), - }, - input_channels="input", - output_channels="output", - ) - - assert app.batch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7] - assert app.batch([3, 2, 1, 3, 5], output_keys=["output"]) == [ - {"output": 5}, - {"output": 4}, - {"output": 3}, - {"output": 5}, - {"output": 7}, - ] - - graph = Graph() - graph.add_node("add_one", add_one_with_delay) - graph.add_node("add_one_more", add_one_with_delay) - graph.set_entry_point("add_one") - graph.set_finish_point("add_one_more") - graph.add_edge("add_one", "add_one_more") - gapp = graph.compile() - - assert gapp.batch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7] - - def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None: test_size = 100 add_one = mocker.Mock(side_effect=lambda x: x + 1) - nodes = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")} + nodes = { + "-1": NodeBuilder() + .subscribe_to("input") + .add_node(add_one) + .write_to("-1") + .build() + } for i in range(test_size - 2): nodes[str(i)] = ( - Channel.subscribe_to(str(i - 1)) | add_one | Channel.write_to(str(i)) + NodeBuilder() + .subscribe_to(str(i - 1)) + .add_node(add_one) + .write_to(str(i)) + .build() ) - nodes["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output") + nodes["last"] = ( + NodeBuilder().subscribe_to(str(i)).add_node(add_one).write_to("output").build() + ) app = Pregel( nodes=nodes, @@ -893,49 +575,15 @@ def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None: ] == [2 + test_size] * 10 -def test_batch_many_processes_in_out(mocker: MockerFixture) -> None: - test_size = 100 - add_one = mocker.Mock(side_effect=lambda x: x + 1) - - nodes = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")} - for i in range(test_size - 2): - nodes[str(i)] = ( - Channel.subscribe_to(str(i - 1)) | add_one | Channel.write_to(str(i)) - ) - nodes["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output") - - app = Pregel( - nodes=nodes, - channels={str(i): LastValue(int) for i in range(-1, test_size - 2)} - | {"input": LastValue(int), "output": LastValue(int)}, - input_channels="input", - output_channels="output", - ) - - for _ in range(3): - assert app.batch([2, 1, 3, 4, 5], {"recursion_limit": test_size}) == [ - 2 + test_size, - 1 + test_size, - 3 + test_size, - 4 + test_size, - 5 + test_size, - ] - - with ThreadPoolExecutor() as executor: - assert [ - *executor.map( - app.batch, [[2, 1, 3, 4, 5]] * 3, [{"recursion_limit": test_size}] * 3 - ) - ] == [ - [2 + test_size, 1 + test_size, 3 + test_size, 4 + test_size, 5 + test_size] - ] * 3 - - def test_invoke_two_processes_two_in_two_out_invalid(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - one = Channel.subscribe_to("input") | add_one | Channel.write_to("output") - two = Channel.subscribe_to("input") | add_one | Channel.write_to("output") + one = ( + NodeBuilder().subscribe_to("input").add_node(add_one).write_to("output").build() + ) + two = ( + NodeBuilder().subscribe_to("input").add_node(add_one).write_to("output").build() + ) app = Pregel( nodes={"one": one, "two": two}, @@ -961,14 +609,18 @@ def test_invoke_two_processes_two_in_two_out_invalid(mocker: MockerFixture) -> N graph = builder.compile() with pytest.raises(InvalidUpdateError, match="At key 'hello'"): - graph.invoke({"hello": "there"}, debug=True) + graph.invoke({"hello": "there"}) def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - one = Channel.subscribe_to("input") | add_one | Channel.write_to("output") - two = Channel.subscribe_to("input") | add_one | Channel.write_to("output") + one = ( + NodeBuilder().subscribe_to("input").add_node(add_one).write_to("output").build() + ) + two = ( + NodeBuilder().subscribe_to("input").add_node(add_one).write_to("output").build() + ) app = Pregel( nodes={"one": one, "two": two}, @@ -1007,10 +659,13 @@ def test_invoke_checkpoint_two( return input one = ( - Channel.subscribe_to(["input"]).join(["total"]) - | add_one - | Channel.write_to("output", "total") - | raise_if_above_10 + NodeBuilder() + .subscribe_to(["input"]) + .read_from("total") + .add_node(add_one) + .write_to("output", "total") + .add_node(raise_if_above_10) + .build() ) app = Pregel( @@ -1317,327 +972,6 @@ def test_pending_writes_resume( ) -def test_cond_edge_after_send() -> None: - class Node: - def __init__(self, name: str): - self.name = name - setattr(self, "__name__", name) - - def __call__(self, state): - return [self.name] - - def send_for_fun(state): - return [Send("2", state), Send("2", state)] - - def route_to_three(state) -> Literal["3"]: - return "3" - - builder = StateGraph(Annotated[list, operator.add]) - builder.add_node(Node("1")) - builder.add_node(Node("2")) - builder.add_node(Node("3")) - builder.add_edge(START, "1") - builder.add_conditional_edges("1", send_for_fun) - builder.add_conditional_edges("2", route_to_three) - graph = builder.compile() - assert graph.invoke(["0"]) == ["0", "1", "2", "2", "3"] - - -def test_concurrent_emit_sends() -> None: - class Node: - def __init__(self, name: str): - self.name = name - setattr(self, "__name__", name) - - def __call__(self, state): - return ( - [self.name] - if isinstance(state, list) - else ["|".join((self.name, str(state)))] - ) - - def send_for_fun(state): - return [Send("2", 1), Send("2", 2), "3.1"] - - def send_for_profit(state): - return [Send("2", 3), Send("2", 4)] - - def route_to_three(state) -> Literal["3"]: - return "3" - - builder = StateGraph(Annotated[list, operator.add]) - builder.add_node(Node("1")) - builder.add_node(Node("1.1")) - builder.add_node(Node("2")) - builder.add_node(Node("3")) - builder.add_node(Node("3.1")) - builder.add_edge(START, "1") - builder.add_edge(START, "1.1") - builder.add_conditional_edges("1", send_for_fun) - builder.add_conditional_edges("1.1", send_for_profit) - builder.add_conditional_edges("2", route_to_three) - graph = builder.compile() - assert graph.invoke(["0"]) == [ - "0", - "1", - "1.1", - "3.1", - "2|1", - "2|2", - "2|3", - "2|4", - "3", - ] - - -def test_send_sequences() -> None: - class Node: - def __init__(self, name: str): - self.name = name - setattr(self, "__name__", name) - - def __call__(self, state): - update = ( - [self.name] - if isinstance(state, list) - else ["|".join((self.name, str(state)))] - ) - if isinstance(state, Command): - return [state, Command(update=update)] - else: - return update - - def send_for_fun(state): - return [ - Send("2", Command(goto=Send("2", 3))), - Send("2", Command(goto=Send("2", 4))), - "3.1", - ] - - def route_to_three(state) -> Literal["3"]: - return "3" - - builder = StateGraph(Annotated[list, operator.add]) - builder.add_node(Node("1")) - builder.add_node(Node("2")) - builder.add_node(Node("3")) - builder.add_node(Node("3.1")) - builder.add_edge(START, "1") - builder.add_conditional_edges("1", send_for_fun) - builder.add_conditional_edges("2", route_to_three) - graph = builder.compile() - assert graph.invoke(["0"]) == [ - "0", - "1", - "3.1", - "2|Command(goto=Send(node='2', arg=3))", - "2|Command(goto=Send(node='2', arg=4))", - "3", - "2|3", - "2|4", - "3", - ] - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_imp_task(request: pytest.FixtureRequest, checkpointer_name: str) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - mapper_calls = 0 - - class Config: - model: str - - @task() - def mapper(input: int) -> str: - nonlocal mapper_calls - mapper_calls += 1 - time.sleep(input / 100) - return str(input) * 2 - - @entrypoint(checkpointer=checkpointer, config_schema=Config) - def graph(input: list[int]) -> list[str]: - futures = [mapper(i) for i in input] - mapped = [f.result() for f in futures] - answer = interrupt("question") - return [m + answer for m in mapped] - - assert graph.get_input_jsonschema() == { - "type": "array", - "items": {"type": "integer"}, - "title": "LangGraphInput", - } - assert graph.get_output_jsonschema() == { - "type": "array", - "items": {"type": "string"}, - "title": "LangGraphOutput", - } - assert graph.get_config_jsonschema() == { - "$defs": { - "Configurable": { - "properties": { - "model": {"default": None, "title": "Model", "type": "string"}, - "checkpoint_id": { - "anyOf": [{"type": "string"}, {"type": "null"}], - "default": None, - "description": "Pass to fetch a past checkpoint. If None, fetches the latest checkpoint.", - "title": "Checkpoint ID", - }, - "checkpoint_ns": { - "default": "", - "description": 'Checkpoint namespace. Denotes the path to the subgraph node the checkpoint originates from, separated by `|` character, e.g. `"child|grandchild"`. Defaults to "" (root graph).', - "title": "Checkpoint NS", - "type": "string", - }, - "thread_id": { - "default": "", - "title": "Thread ID", - "type": "string", - }, - }, - "title": "Configurable", - "type": "object", - } - }, - "properties": { - "configurable": {"$ref": "#/$defs/Configurable", "default": None} - }, - "title": "LangGraphConfig", - "type": "object", - } - - thread1 = {"configurable": {"thread_id": "1"}} - assert [*graph.stream([0, 1], thread1)] == [ - {"mapper": "00"}, - {"mapper": "11"}, - { - "__interrupt__": ( - Interrupt( - value="question", - resumable=True, - ns=[AnyStr("graph:")], - when="during", - ), - ) - }, - ] - assert mapper_calls == 2 - - assert graph.invoke(Command(resume="answer"), thread1) == [ - "00answer", - "11answer", - ] - assert mapper_calls == 2 - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_imp_nested( - request: pytest.FixtureRequest, checkpointer_name: str, snapshot: SnapshotAssertion -) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - - def mynode(input: list[str]) -> list[str]: - return [it + "a" for it in input] - - builder = StateGraph(list[str]) - builder.add_node(mynode) - builder.add_edge(START, "mynode") - add_a = builder.compile() - - @task - def submapper(input: int) -> str: - time.sleep(input / 100) - return str(input) - - @task() - def mapper(input: int) -> str: - sub = submapper(input) - time.sleep(input / 100) - return sub.result() * 2 - - @entrypoint(checkpointer=checkpointer) - def graph(input: list[int]) -> list[str]: - futures = [mapper(i) for i in input] - mapped = [f.result() for f in futures] - answer = interrupt("question") - final = [m + answer for m in mapped] - return add_a.invoke(final) - - assert graph.get_input_jsonschema() == { - "type": "array", - "items": {"type": "integer"}, - "title": "LangGraphInput", - } - assert graph.get_output_jsonschema() == { - "type": "array", - "items": {"type": "string"}, - "title": "LangGraphOutput", - } - - thread1 = {"configurable": {"thread_id": "1"}} - assert [*graph.stream([0, 1], thread1)] == [ - {"submapper": "0"}, - {"mapper": "00"}, - {"submapper": "1"}, - {"mapper": "11"}, - { - "__interrupt__": ( - Interrupt( - value="question", - resumable=True, - ns=[AnyStr("graph:")], - when="during", - ), - ) - }, - ] - - assert graph.invoke(Command(resume="answer"), thread1) == [ - "00answera", - "11answera", - ] - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_imp_stream_order( - request: pytest.FixtureRequest, checkpointer_name: str, snapshot: SnapshotAssertion -) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - - @task() - def foo(state: dict) -> tuple: - return state["a"] + "foo", "bar" - - @task - def bar(a: str, b: str, c: Optional[str] = None) -> dict: - return {"a": a + b, "c": (c or "") + "bark"} - - @task - def baz(state: dict) -> dict: - return {"a": state["a"] + "baz", "c": "something else"} - - @entrypoint(checkpointer=checkpointer) - def graph(state: dict) -> dict: - fut_foo = foo(state) - fut_bar = bar(*fut_foo.result()) - fut_baz = baz(fut_bar.result()) - return fut_baz.result() - - thread1 = {"configurable": {"thread_id": "1"}} - assert [c for c in graph.stream({"a": "0"}, thread1)] == [ - { - "foo": ( - "0foo", - "bar", - ) - }, - {"bar": {"a": "0foobar", "c": "bark"}}, - {"baz": {"a": "0foobarbaz", "c": "something else"}}, - {"graph": {"a": "0foobarbaz", "c": "something else"}}, - ] - - assert graph.get_state(thread1).values == {"a": "0foobarbaz", "c": "something else"} - - @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_invoke_checkpoint_three( mocker: MockerFixture, request: pytest.FixtureRequest, checkpointer_name: str @@ -1651,10 +985,13 @@ def test_invoke_checkpoint_three( return input one = ( - Channel.subscribe_to(["input"]).join(["total"]) - | adder - | Channel.write_to("output", "total") - | raise_if_above_10 + NodeBuilder() + .subscribe_to(["input"]) + .read_from("total") + .add_node(adder) + .write_to("output", "total") + .add_node(raise_if_above_10) + .build() ) app = Pregel( @@ -1671,7 +1008,7 @@ def test_invoke_checkpoint_three( thread_1 = {"configurable": {"thread_id": "1"}} # total starts out as 0, so output is 0+2=2 - assert app.invoke(2, thread_1, debug=1) == 2 + assert app.invoke(2, thread_1) == 2 state = app.get_state(thread_1) assert state is not None assert state.values.get("total") == 2 @@ -1707,7 +1044,7 @@ def test_invoke_checkpoint_three( thread_2 = {"configurable": {"thread_id": "2"}} # on a new thread, total starts out as 0, so output is 0+5=5 - assert app.invoke(5, thread_2, debug=True) == 5 + assert app.invoke(5, thread_2) == 5 state = app.get_state({"configurable": {"thread_id": "1"}}) assert state is not None assert state.values.get("total") == 16 @@ -1779,10 +1116,18 @@ def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None add_one = mocker.Mock(side_effect=lambda x: x + 1) add_10_each = mocker.Mock(side_effect=lambda x: sorted(y + 10 for y in x)) - one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") - chain_three = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") + one = ( + NodeBuilder().subscribe_to("input").add_node(add_one).write_to("inbox").build() + ) + chain_three = ( + NodeBuilder().subscribe_to("input").add_node(add_one).write_to("inbox").build() + ) chain_four = ( - Channel.subscribe_to("inbox") | add_10_each | Channel.write_to("output") + NodeBuilder() + .subscribe_to("inbox") + .add_node(add_10_each) + .write_to("output") + .build() ) app = Pregel( @@ -1810,80 +1155,23 @@ def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None assert [*executor.map(app.invoke, [2] * 100)] == [[13, 13]] * 100 -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_invoke_join_then_call_other_pregel( - mocker: MockerFixture, request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - - add_one = mocker.Mock(side_effect=lambda x: x + 1) - add_10_each = mocker.Mock(side_effect=lambda x: [y + 10 for y in x]) - - inner_app = Pregel( - nodes={ - "one": Channel.subscribe_to("input") | add_one | Channel.write_to("output") - }, - channels={ - "output": LastValue(int), - "input": LastValue(int), - }, - input_channels="input", - output_channels="output", - ) - - one = ( - Channel.subscribe_to("input") - | add_10_each - | Channel.write_to("inbox_one").map() - ) - two = ( - Channel.subscribe_to("inbox_one") - | inner_app.map() - | sorted - | Channel.write_to("outbox_one") - ) - chain_three = Channel.subscribe_to("outbox_one") | sum | Channel.write_to("output") - - app = Pregel( - nodes={ - "one": one, - "two": two, - "chain_three": chain_three, - }, - channels={ - "inbox_one": Topic(int), - "outbox_one": LastValue(int), - "output": LastValue(int), - "input": LastValue(int), - }, - input_channels="input", - output_channels="output", - ) - - for _ in range(10): - assert app.invoke([2, 3]) == 27 - - with ThreadPoolExecutor() as executor: - assert [*executor.map(app.invoke, [[2, 3]] * 10)] == [27] * 10 - - # add checkpointer - app.checkpointer = checkpointer - # subgraph is called twice in the same node, but that works - assert app.invoke([2, 3], {"configurable": {"thread_id": "1"}}) == 27 - - # set inner graph checkpointer NeverCheckpoint - inner_app.checkpointer = False - # subgraph still called twice, but checkpointing for inner graph is disabled - assert app.invoke([2, 3], {"configurable": {"thread_id": "1"}}) == 27 - - def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) one = ( - Channel.subscribe_to("input") | add_one | Channel.write_to("output", "between") + NodeBuilder() + .subscribe_to("input") + .add_node(add_one) + .write_to("output", "between") + .build() + ) + two = ( + NodeBuilder() + .subscribe_to("between") + .add_node(add_one) + .write_to("output") + .build() ) - two = Channel.subscribe_to("between") | add_one | Channel.write_to("output") app = Pregel( nodes={"one": one, "two": two}, @@ -1909,8 +1197,14 @@ def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None: def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - one = Channel.subscribe_to("input") | add_one | Channel.write_to("between") - two = Channel.subscribe_to("between") | add_one + one = ( + NodeBuilder() + .subscribe_to("input") + .add_node(add_one) + .write_to("between") + .build() + ) + two = NodeBuilder().subscribe_to("between").add_node(add_one).build() app = Pregel( nodes={"one": one, "two": two}, @@ -1931,104 +1225,19 @@ def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None: def test_invoke_two_processes_no_in(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - one = Channel.subscribe_to("between") | add_one | Channel.write_to("output") - two = Channel.subscribe_to("between") | add_one + one = ( + NodeBuilder() + .subscribe_to("between") + .add_node(add_one) + .write_to("output") + .build() + ) + two = NodeBuilder().subscribe_to("between").add_node(add_one).build() with pytest.raises(TypeError): Pregel(nodes={"one": one, "two": two}) -def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: - setup = mocker.Mock() - cleanup = mocker.Mock() - - @contextmanager - def an_int() -> Generator[int, None, None]: - setup() - try: - yield 5 - finally: - cleanup() - - add_one = mocker.Mock(side_effect=lambda x: x + 1) - one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") - two = ( - Channel.subscribe_to("inbox") - | RunnableLambda(add_one).batch - | Channel.write_to("output").batch - ) - - app = Pregel( - nodes={"one": one, "two": two}, - channels={ - "inbox": Topic(int), - "ctx": Context(an_int), - "output": LastValue(int), - "input": LastValue(int), - }, - input_channels="input", - output_channels=["inbox", "output"], - stream_channels=["inbox", "output"], - ) - - assert setup.call_count == 0 - assert cleanup.call_count == 0 - for i, chunk in enumerate(app.stream(2)): - assert setup.call_count == 1, "Expected setup to be called once" - if i == 0: - assert chunk == {"inbox": [3]} - elif i == 1: - assert chunk == {"output": 4} - else: - assert False, "Expected only two chunks" - assert cleanup.call_count == 1, "Expected cleanup to be called once" - - -def test_conditional_entrypoint_graph(snapshot: SnapshotAssertion) -> None: - def left(data: str) -> str: - return data + "->left" - - def right(data: str) -> str: - return data + "->right" - - def should_start(data: str) -> str: - # Logic to decide where to start - if len(data) > 10: - return "go-right" - else: - return "go-left" - - # Define a new graph - workflow = Graph() - - workflow.add_node("left", left) - workflow.add_node("right", right) - - workflow.set_conditional_entry_point( - should_start, {"go-left": "left", "go-right": "right"} - ) - - workflow.add_conditional_edges("left", lambda data: END, {END: END}) - workflow.add_edge("right", END) - - app = workflow.compile() - - if SHOULD_CHECK_SNAPSHOTS: - assert json.dumps(app.get_input_schema().model_json_schema()) == snapshot - assert json.dumps(app.get_output_schema().model_json_schema()) == snapshot - assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot - assert app.get_graph().draw_mermaid(with_styles=False) == snapshot - - assert ( - app.invoke("what is weather in sf", debug=True) - == "what is weather in sf->right" - ) - - assert [*app.stream("what is weather in sf")] == [ - {"right": "what is weather in sf->right"}, - ] - - def test_conditional_entrypoint_to_multiple_state_graph( snapshot: SnapshotAssertion, ) -> None: @@ -2055,13 +1264,7 @@ def test_conditional_entrypoint_to_multiple_state_graph( app = workflow.compile() - if SHOULD_CHECK_SNAPSHOTS: - assert json.dumps(app.get_input_schema().model_json_schema()) == snapshot - assert json.dumps(app.get_output_schema().model_json_schema()) == snapshot - assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot - assert app.get_graph().draw_mermaid(with_styles=False) == snapshot - - assert app.invoke({"locations": ["sf", "nyc"]}, debug=True) == { + assert app.invoke({"locations": ["sf", "nyc"]}) == { "locations": ["sf", "nyc"], "results": ["It's cloudy in sf", "It's sunny in nyc"], } @@ -2086,9 +1289,6 @@ def test_conditional_state_graph_with_list_edge_inputs(snapshot: SnapshotAsserti app = graph_builder.compile() assert app.invoke({"foo": []}) == {"foo": ["A", "B"]} - assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot - assert app.get_graph().draw_mermaid(with_styles=False) == snapshot - def test_state_graph_w_config_inherited_state_keys(snapshot: SnapshotAssertion) -> None: from langchain_core.language_models.fake import FakeStreamingListLLM @@ -2146,7 +1346,7 @@ def test_state_graph_w_config_inherited_state_keys(snapshot: SnapshotAssertion) ) } - agent = prompt | llm | agent_parser + agent = RunnableSeq(prompt, llm, agent_parser) # Define tool execution logic def execute_tools(data: AgentState) -> dict: @@ -2180,11 +1380,6 @@ def test_state_graph_w_config_inherited_state_keys(snapshot: SnapshotAssertion) app = builder.compile() - if SHOULD_CHECK_SNAPSHOTS: - assert json.dumps(app.config_schema().model_json_schema()) == snapshot - assert json.dumps(app.get_input_schema().model_json_schema()) == snapshot - assert json.dumps(app.get_output_schema().model_json_schema()) == snapshot - assert builder.channels.keys() == {"input", "agent_outcome", "intermediate_steps"} assert app.invoke({"input": "what is weather in sf"}) == { @@ -2246,12 +1441,6 @@ def test_conditional_entrypoint_graph_state(snapshot: SnapshotAssertion) -> None app = workflow.compile() - if SHOULD_CHECK_SNAPSHOTS: - assert json.dumps(app.get_input_schema().model_json_schema()) == snapshot - assert json.dumps(app.get_output_schema().model_json_schema()) == snapshot - assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot - assert app.get_graph().draw_mermaid(with_styles=False) == snapshot - assert app.invoke({"input": "what is weather in sf"}) == { "input": "what is weather in sf", "output": "what is weather in sf->right", @@ -2318,8 +1507,6 @@ def test_in_one_fan_out_state_graph_waiting_edge( app = workflow.compile() - assert app.get_graph().draw_mermaid(with_styles=False) == snapshot - assert app.invoke({"query": "what is weather in sf"}) == { "query": "analyzed: query: what is weather in sf", "docs": ["doc1", "doc2", "doc3", "doc4"], @@ -2401,7 +1588,7 @@ def test_in_one_fan_out_state_graph_waiting_edge( parent_config=expected_parent_config, ) - assert [c for c in app_w_interrupt.stream(None, config, debug=1)] == [ + assert [c for c in app_w_interrupt.stream(None, config)] == [ {"qa": {"answer": "doc1,doc2,doc3,doc4,doc5"}}, ] @@ -2464,9 +1651,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_via_branch( app = workflow.compile() - assert app.get_graph().draw_mermaid(with_styles=False) == snapshot - - assert app.invoke({"query": "what is weather in sf"}, debug=True) == { + assert app.invoke({"query": "what is weather in sf"}) == { "query": "analyzed: query: what is weather in sf", "docs": ["doc1", "doc2", "doc3", "doc4"], "answer": "doc1,doc2,doc3,doc4", @@ -2511,29 +1696,6 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1( from pydantic.v1 import BaseModel, ValidationError checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - setup = mocker.Mock() - teardown = mocker.Mock() - - @contextmanager - def assert_ctx_once() -> Iterator[None]: - assert setup.call_count == 0 - assert teardown.call_count == 0 - try: - yield - finally: - assert setup.call_count == 1 - assert teardown.call_count == 1 - setup.reset_mock() - teardown.reset_mock() - - @contextmanager - def make_httpx_client() -> Iterator[httpx.Client]: - setup() - with httpx.Client() as client: - try: - yield client - finally: - teardown() def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] @@ -2555,7 +1717,6 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1( inner: InnerObject answer: Optional[str] = None docs: Annotated[list[str], sorted_add] - client: Annotated[httpx.Client, Context(make_httpx_client)] class Input(BaseModel): query: str @@ -2609,29 +1770,21 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1( app = workflow.compile() - assert app.get_graph().draw_mermaid(with_styles=False) == snapshot - assert app.get_input_jsonschema() == snapshot - assert app.get_output_jsonschema() == snapshot - - with pytest.raises(ValidationError), assert_ctx_once(): + with pytest.raises(ValidationError): app.invoke({"query": {}}) - with assert_ctx_once(): - assert app.invoke({"query": "what is weather in sf", "inner": {"yo": 1}}) == { - "docs": ["doc1", "doc2", "doc3", "doc4"], - "answer": "doc1,doc2,doc3,doc4", - } + assert app.invoke({"query": "what is weather in sf", "inner": {"yo": 1}}) == { + "docs": ["doc1", "doc2", "doc3", "doc4"], + "answer": "doc1,doc2,doc3,doc4", + } - with assert_ctx_once(): - assert [ - *app.stream({"query": "what is weather in sf", "inner": {"yo": 1}}) - ] == [ - {"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"]}}, - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] + assert [*app.stream({"query": "what is weather in sf", "inner": {"yo": 1}})] == [ + {"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"]}}, + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] app_w_interrupt = workflow.compile( checkpointer=checkpointer, @@ -2639,204 +1792,32 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1( ) config = {"configurable": {"thread_id": "1"}} - with assert_ctx_once(): - assert [ - c - for c in app_w_interrupt.stream( - {"query": "what is weather in sf", "inner": {"yo": 1}}, 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__": ()}, - ] + assert [ + c + for c in app_w_interrupt.stream( + {"query": "what is weather in sf", "inner": {"yo": 1}}, 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__": ()}, + ] - with assert_ctx_once(): - assert [c for c in app_w_interrupt.stream(None, config)] == [ - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] + assert [c for c in app_w_interrupt.stream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] - with assert_ctx_once(): - assert app_w_interrupt.update_state( - config, {"docs": ["doc5"]}, as_node="rewrite_query" - ) == { - "configurable": { - "thread_id": "1", - "checkpoint_id": AnyStr(), - "checkpoint_ns": "", - } - } - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2( - snapshot: SnapshotAssertion, - mocker: MockerFixture, - request: pytest.FixtureRequest, - checkpointer_name: str, -) -> None: - from pydantic import BaseModel, ConfigDict, ValidationError - - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - setup = mocker.Mock() - teardown = mocker.Mock() - - @contextmanager - def assert_ctx_once() -> Iterator[None]: - assert setup.call_count == 0 - assert teardown.call_count == 0 - try: - yield - finally: - assert setup.call_count == 1 - assert teardown.call_count == 1 - setup.reset_mock() - teardown.reset_mock() - - @contextmanager - def make_httpx_client() -> Iterator[httpx.Client]: - setup() - with httpx.Client() as client: - try: - yield client - finally: - teardown() - - def sorted_add( - x: list[str], y: Union[list[str], list[tuple[str, str]]] - ) -> list[str]: - if isinstance(y[0], tuple): - for rem, _ in y: - x.remove(rem) - y = [t[1] for t in y] - return sorted(operator.add(x, y)) - - class InnerObject(BaseModel): - yo: int - - class State(BaseModel): - model_config = ConfigDict(arbitrary_types_allowed=True) - - query: str - inner: InnerObject - answer: Optional[str] = None - docs: Annotated[list[str], sorted_add] - client: Annotated[httpx.Client, Context(make_httpx_client)] - - class StateUpdate(BaseModel): - query: Optional[str] = None - answer: Optional[str] = None - docs: Optional[list[str]] = None - - class Input(BaseModel): - query: str - inner: InnerObject - - class Output(BaseModel): - answer: str - docs: list[str] - - def rewrite_query(data: State) -> State: - return {"query": f"query: {data.query}"} - - def analyzer_one(data: State) -> State: - return StateUpdate(query=f"analyzed: {data.query}") - - def retriever_one(data: State) -> State: - return {"docs": ["doc1", "doc2"]} - - def retriever_two(data: State) -> State: - time.sleep(0.1) - return {"docs": ["doc3", "doc4"]} - - def qa(data: State) -> State: - return {"answer": ",".join(data.docs)} - - def decider(data: State) -> str: - assert isinstance(data, State) - return "retriever_two" - - workflow = StateGraph(State, input=Input, output=Output) - - workflow.add_node("rewrite_query", rewrite_query) - workflow.add_node("analyzer_one", analyzer_one) - workflow.add_node("retriever_one", retriever_one) - workflow.add_node("retriever_two", retriever_two) - workflow.add_node("qa", qa) - - workflow.set_entry_point("rewrite_query") - workflow.add_edge("rewrite_query", "analyzer_one") - workflow.add_edge("analyzer_one", "retriever_one") - workflow.add_conditional_edges( - "rewrite_query", decider, {"retriever_two": "retriever_two"} - ) - workflow.add_edge(["retriever_one", "retriever_two"], "qa") - workflow.set_finish_point("qa") - - app = workflow.compile() - - if SHOULD_CHECK_SNAPSHOTS: - assert app.get_graph().draw_mermaid(with_styles=False) == snapshot - assert app.get_input_schema().model_json_schema() == snapshot - assert app.get_output_schema().model_json_schema() == snapshot - - with pytest.raises(ValidationError), assert_ctx_once(): - app.invoke({"query": {}}) - - with assert_ctx_once(): - assert app.invoke({"query": "what is weather in sf", "inner": {"yo": 1}}) == { - "docs": ["doc1", "doc2", "doc3", "doc4"], - "answer": "doc1,doc2,doc3,doc4", - } - - with assert_ctx_once(): - assert [ - *app.stream({"query": "what is weather in sf", "inner": {"yo": 1}}) - ] == [ - {"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"]}}, - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] - - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_after=["retriever_one"], - ) - config = {"configurable": {"thread_id": "1"}} - - with assert_ctx_once(): - assert [ - c - for c in app_w_interrupt.stream( - {"query": "what is weather in sf", "inner": {"yo": 1}}, 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__": ()}, - ] - - with assert_ctx_once(): - assert [c for c in app_w_interrupt.stream(None, config)] == [ - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] - - with assert_ctx_once(): - assert app_w_interrupt.update_state( - config, {"docs": ["doc5"]}, as_node="rewrite_query" - ) == { - "configurable": { - "thread_id": "1", - "checkpoint_id": AnyStr(), - "checkpoint_ns": "", - } + assert app_w_interrupt.update_state( + config, {"docs": ["doc5"]}, as_node="rewrite_query" + ) == { + "configurable": { + "thread_id": "1", + "checkpoint_id": AnyStr(), + "checkpoint_ns": "", } + } @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) @@ -3187,8 +2168,6 @@ def test_simple_multi_edge(snapshot: SnapshotAssertion) -> None: graph.set_finish_point("down") app = graph.compile() - - assert app.get_graph().draw_mermaid(with_styles=False) == snapshot assert app.invoke({"my_key": "my_value"}) == {"my_key": "my_value_more"} assert [*app.stream({"my_key": "my_value"})] in ( [ @@ -3206,34 +2185,6 @@ def test_simple_multi_edge(snapshot: SnapshotAssertion) -> None: ) -def test_nested_graph_xray(snapshot: SnapshotAssertion) -> None: - class State(TypedDict): - my_key: Annotated[str, operator.add] - market: str - - def logic(state: State): - pass - - tool_two_graph = StateGraph(State) - tool_two_graph.add_node("tool_two_slow", logic) - tool_two_graph.add_node("tool_two_fast", logic) - tool_two_graph.set_conditional_entry_point( - lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast", - then=END, - ) - tool_two = tool_two_graph.compile() - - graph = StateGraph(State) - graph.add_node("tool_one", logic) - graph.add_node("tool_two", tool_two) - graph.add_node("tool_three", logic) - graph.set_conditional_entry_point(lambda s: "tool_one", then=END) - app = graph.compile() - - assert app.get_graph(xray=True).to_json() == snapshot - assert app.get_graph(xray=True).draw_mermaid() == snapshot - - def test_nested_graph(snapshot: SnapshotAssertion) -> None: def never_called_fn(state: Any): assert 0, "This function should never be called" @@ -3267,12 +2218,7 @@ def test_nested_graph(snapshot: SnapshotAssertion) -> None: graph.set_finish_point("side") app = graph.compile() - - assert app.get_graph().draw_mermaid(with_styles=False) == snapshot - assert app.get_graph(xray=True).draw_mermaid() == snapshot - assert app.invoke( - {"my_key": "my value", "never_called": never_called}, debug=True - ) == { + assert app.invoke({"my_key": "my value", "never_called": never_called}) == { "my_key": "my value there and back again", "never_called": never_called, } @@ -3299,17 +2245,6 @@ def test_nested_graph(snapshot: SnapshotAssertion) -> None: }, ] - chain = app | RunnablePassthrough() - - assert chain.invoke({"my_key": "my value", "never_called": never_called}) == { - "my_key": "my value there and back again", - "never_called": never_called, - } - assert [*chain.stream({"my_key": "my value", "never_called": never_called})] == [ - {"inner": {"my_key": "my value there"}}, - {"side": {"my_key": "my value there and back again"}}, - ] - @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_subgraph_checkpoint_true( @@ -3409,7 +2344,7 @@ def test_subgraph_checkpoint_true_interrupt( builder = StateGraph(ParentState) builder.add_node("node_1", node_1) - builder.add_node("node_2", node_2) + builder.add_node("node_2", node_2, subgraphs=[subgraph]) builder.add_edge(START, "node_1") builder.add_edge("node_1", "node_2") @@ -3418,6 +2353,7 @@ def test_subgraph_checkpoint_true_interrupt( config = {"configurable": {"thread_id": "1"}} assert graph.invoke({"foo": "foo"}, config) == {"foo": "hi! foo"} + print(graph.get_state(config, subgraphs=True).tasks) assert graph.get_state(config, subgraphs=True).tasks[0].state.values == { "bar": "hi! foo" } @@ -3513,7 +2449,8 @@ def test_stream_buffering_single_node( class State(TypedDict): my_key: Annotated[str, operator.add] - def node(state: State, writer: StreamWriter): + def node(state: State): + writer = get_stream_writer() writer("Before sleep") time.sleep(0.2) writer("After sleep") @@ -3587,11 +2524,11 @@ def test_nested_graph_interrupts_parallel( # test invoke w/ nested interrupt config = {"configurable": {"thread_id": "1"}} - assert app.invoke({"my_key": ""}, config, debug=True) == { + assert app.invoke({"my_key": ""}, config) == { "my_key": " and parallel", } - assert app.invoke(None, config, debug=True) == { + assert app.invoke(None, config) == { "my_key": "got here and there and parallel and back again", } @@ -3715,11 +2652,11 @@ def test_doubly_nested_graph_interrupts( # test invoke w/ nested interrupt config = {"configurable": {"thread_id": "1"}} - assert app.invoke({"my_key": "my value"}, config, debug=True) == { + assert app.invoke({"my_key": "my value"}, config) == { "my_key": "hi my value", } - assert app.invoke(None, config, debug=True) == { + assert app.invoke(None, config) == { "my_key": "hi my value here and there and back again", } @@ -3759,543 +2696,6 @@ def test_doubly_nested_graph_interrupts( ] -def test_repeat_condition(snapshot: SnapshotAssertion) -> None: - class AgentState(TypedDict): - hello: str - - def router(state: AgentState) -> str: - return "hmm" - - workflow = StateGraph(AgentState) - workflow.add_node("Researcher", lambda x: x) - workflow.add_node("Chart Generator", lambda x: x) - workflow.add_node("Call Tool", lambda x: x) - workflow.add_conditional_edges( - "Researcher", - router, - { - "redo": "Researcher", - "continue": "Chart Generator", - "call_tool": "Call Tool", - "end": END, - }, - ) - workflow.add_conditional_edges( - "Chart Generator", - router, - {"continue": "Researcher", "call_tool": "Call Tool", "end": END}, - ) - workflow.add_conditional_edges( - "Call Tool", - # Each agent node updates the 'sender' field - # the tool calling node does not, meaning - # this edge will route back to the original agent - # who invoked the tool - lambda x: x["sender"], - { - "Researcher": "Researcher", - "Chart Generator": "Chart Generator", - }, - ) - workflow.set_entry_point("Researcher") - - app = workflow.compile() - assert app.get_graph().draw_mermaid(with_styles=False) == snapshot - - -def test_checkpoint_metadata() -> None: - """This test verifies that a run's configurable fields are merged with the - previous checkpoint config for each step in the run. - """ - # set up test - from langchain_core.language_models.fake_chat_models import ( - FakeMessagesListChatModel, - ) - from langchain_core.messages import AIMessage, AnyMessage - from langchain_core.prompts import ChatPromptTemplate - from langchain_core.tools import tool - - # graph state - class BaseState(TypedDict): - messages: Annotated[list[AnyMessage], add_messages] - - # initialize graph nodes - @tool() - def search_api(query: str) -> str: - """Searches the API for the query.""" - return f"result for {query}" - - tools = [search_api] - - prompt = ChatPromptTemplate.from_messages( - [ - ("system", "You are a nice assistant."), - ("placeholder", "{messages}"), - ] - ) - - model = FakeMessagesListChatModel( - responses=[ - AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - }, - ], - ), - AIMessage(content="answer"), - ] - ) - - @traceable(run_type="llm") - def agent(state: BaseState) -> BaseState: - formatted = prompt.invoke(state) - response = model.invoke(formatted) - return {"messages": response, "usage_metadata": {"total_tokens": 123}} - - def should_continue(data: BaseState) -> str: - # Logic to decide whether to continue in the loop or exit - if not data["messages"][-1].tool_calls: - return "exit" - else: - return "continue" - - # define graphs w/ and w/o interrupt - workflow = StateGraph(BaseState) - workflow.add_node("agent", agent) - workflow.add_node("tools", ToolNode(tools)) - workflow.set_entry_point("agent") - workflow.add_conditional_edges( - "agent", should_continue, {"continue": "tools", "exit": END} - ) - workflow.add_edge("tools", "agent") - - # graph w/o interrupt - checkpointer_1 = MemorySaverAssertCheckpointMetadata() - app = workflow.compile(checkpointer=checkpointer_1) - - # graph w/ interrupt - checkpointer_2 = MemorySaverAssertCheckpointMetadata() - app_w_interrupt = workflow.compile( - checkpointer=checkpointer_2, interrupt_before=["tools"] - ) - - # assertions - - # invoke graph w/o interrupt - assert app.invoke( - {"messages": ["what is weather in sf"]}, - { - "configurable": { - "thread_id": "1", - "test_config_1": "foo", - "test_config_2": "bar", - }, - }, - ) == { - "messages": [ - _AnyIdHumanMessage(content="what is weather in sf"), - _AnyIdAIMessage( - content="", - tool_calls=[ - { - "name": "search_api", - "args": {"query": "query"}, - "id": "tool_call123", - "type": "tool_call", - } - ], - ), - _AnyIdToolMessage( - content="result for query", - name="search_api", - tool_call_id="tool_call123", - ), - _AnyIdAIMessage(content="answer"), - ] - } - - config = {"configurable": {"thread_id": "1"}} - - # assert that checkpoint metadata contains the run's configurable fields - chkpnt_metadata_1 = checkpointer_1.get_tuple(config).metadata - assert chkpnt_metadata_1["thread_id"] == "1" - assert chkpnt_metadata_1["test_config_1"] == "foo" - assert chkpnt_metadata_1["test_config_2"] == "bar" - - # Verify that all checkpoint metadata have the expected keys. This check - # is needed because a run may have an arbitrary number of steps depending - # on how the graph is constructed. - chkpnt_tuples_1 = checkpointer_1.list(config) - for chkpnt_tuple in chkpnt_tuples_1: - assert chkpnt_tuple.metadata["thread_id"] == "1" - assert chkpnt_tuple.metadata["test_config_1"] == "foo" - assert chkpnt_tuple.metadata["test_config_2"] == "bar" - - # invoke graph, but interrupt before tool call - app_w_interrupt.invoke( - {"messages": ["what is weather in sf"]}, - { - "configurable": { - "thread_id": "2", - "test_config_3": "foo", - "test_config_4": "bar", - }, - }, - ) - - config = {"configurable": {"thread_id": "2"}} - - # assert that checkpoint metadata contains the run's configurable fields - chkpnt_metadata_2 = checkpointer_2.get_tuple(config).metadata - assert chkpnt_metadata_2["thread_id"] == "2" - assert chkpnt_metadata_2["test_config_3"] == "foo" - assert chkpnt_metadata_2["test_config_4"] == "bar" - - # resume graph execution - app_w_interrupt.invoke( - input=None, - config={ - "configurable": { - "thread_id": "2", - "test_config_3": "foo", - "test_config_4": "bar", - } - }, - ) - - # assert that checkpoint metadata contains the run's configurable fields - chkpnt_metadata_3 = checkpointer_2.get_tuple(config).metadata - assert chkpnt_metadata_3["thread_id"] == "2" - assert chkpnt_metadata_3["test_config_3"] == "foo" - assert chkpnt_metadata_3["test_config_4"] == "bar" - - # Verify that all checkpoint metadata have the expected keys. This check - # is needed because a run may have an arbitrary number of steps depending - # on how the graph is constructed. - chkpnt_tuples_2 = checkpointer_2.list(config) - for chkpnt_tuple in chkpnt_tuples_2: - assert chkpnt_tuple.metadata["thread_id"] == "2" - assert chkpnt_tuple.metadata["test_config_3"] == "foo" - assert chkpnt_tuple.metadata["test_config_4"] == "bar" - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_remove_message_via_state_update( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage - - workflow = MessageGraph() - workflow.add_node( - "chatbot", - lambda state: [ - AIMessage( - content="Hello! How can I help you", - ) - ], - ) - - workflow.set_entry_point("chatbot") - workflow.add_edge("chatbot", END) - - checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) - app = workflow.compile(checkpointer=checkpointer) - config = {"configurable": {"thread_id": "1"}} - output = app.invoke([HumanMessage(content="Hi")], config=config) - app.update_state(config, values=[RemoveMessage(id=output[-1].id)]) - - updated_state = app.get_state(config) - - assert len(updated_state.values) == 1 - assert updated_state.values[-1].content == "Hi" - - -def test_remove_message_from_node(): - from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage - - workflow = MessageGraph() - workflow.add_node( - "chatbot", - lambda state: [ - AIMessage( - content="Hello!", - ), - AIMessage( - content="How can I help you?", - ), - ], - ) - workflow.add_node("delete_messages", lambda state: [RemoveMessage(id=state[-2].id)]) - workflow.set_entry_point("chatbot") - workflow.add_edge("chatbot", "delete_messages") - workflow.add_edge("delete_messages", END) - - app = workflow.compile() - output = app.invoke([HumanMessage(content="Hi")]) - assert len(output) == 2 - assert output[-1].content == "How can I help you?" - - -def test_xray_lance(snapshot: SnapshotAssertion): - from langchain_core.messages import AnyMessage, HumanMessage - from pydantic import BaseModel, Field - - class Analyst(BaseModel): - affiliation: str = Field( - description="Primary affiliation of the investment analyst.", - ) - name: str = Field( - description="Name of the investment analyst.", - pattern=r"^[a-zA-Z0-9_-]{1,64}$", - ) - role: str = Field( - description="Role of the investment analyst in the context of the topic.", - ) - description: str = Field( - description="Description of the investment analyst focus, concerns, and motives.", - ) - - @property - def persona(self) -> str: - return f"Name: {self.name}\nRole: {self.role}\nAffiliation: {self.affiliation}\nDescription: {self.description}\n" - - class Perspectives(BaseModel): - analysts: List[Analyst] = Field( - description="Comprehensive list of investment analysts with their roles and affiliations.", - ) - - class Section(BaseModel): - section_title: str = Field(..., title="Title of the section") - context: str = Field( - ..., title="Provide a clear summary of the focus area that you researched." - ) - findings: str = Field( - ..., - title="Give a clear and detailed overview of your findings based upon the expert interview.", - ) - thesis: str = Field( - ..., - title="Give a clear and specific investment thesis based upon these findings.", - ) - - class InterviewState(TypedDict): - messages: Annotated[List[AnyMessage], add_messages] - analyst: Analyst - section: Section - - class ResearchGraphState(TypedDict): - analysts: List[Analyst] - topic: str - max_analysts: int - sections: List[Section] - interviews: Annotated[list, operator.add] - - # Conditional edge - def route_messages(state): - return "ask_question" - - def generate_question(state): - return ... - - def generate_answer(state): - return ... - - # Add nodes and edges - interview_builder = StateGraph(InterviewState) - interview_builder.add_node("ask_question", generate_question) - interview_builder.add_node("answer_question", generate_answer) - - # Flow - interview_builder.add_edge(START, "ask_question") - interview_builder.add_edge("ask_question", "answer_question") - interview_builder.add_conditional_edges("answer_question", route_messages) - - # Set up memory - memory = InMemorySaver() - - # Interview - interview_graph = interview_builder.compile(checkpointer=memory).with_config( - run_name="Conduct Interviews" - ) - - # View - assert interview_graph.get_graph().to_json() == snapshot - - def run_all_interviews(state: ResearchGraphState): - """Edge to run the interview sub-graph using Send""" - return [ - Send( - "conduct_interview", - { - "analyst": Analyst(), - "messages": [ - HumanMessage( - content="So you said you were writing an article on ...?" - ) - ], - }, - ) - for s in state["analysts"] - ] - - def generate_sections(state: ResearchGraphState): - return ... - - def generate_analysts(state: ResearchGraphState): - return ... - - builder = StateGraph(ResearchGraphState) - builder.add_node("generate_analysts", generate_analysts) - builder.add_node("conduct_interview", interview_builder.compile()) - builder.add_node("generate_sections", generate_sections) - - builder.add_edge(START, "generate_analysts") - builder.add_conditional_edges( - "generate_analysts", run_all_interviews, ["conduct_interview"] - ) - builder.add_edge("conduct_interview", "generate_sections") - builder.add_edge("generate_sections", END) - - graph = builder.compile() - - # View - assert graph.get_graph().to_json() == snapshot - assert graph.get_graph(xray=1).to_json() == snapshot - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_channel_values(request: pytest.FixtureRequest, checkpointer_name: str) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - - config = {"configurable": {"thread_id": "1"}} - chain = Channel.subscribe_to("input") | Channel.write_to("output") - app = Pregel( - nodes={ - "one": chain, - }, - channels={ - "ephemeral": EphemeralValue(Any), - "input": LastValue(int), - "output": LastValue(int), - }, - input_channels=["input", "ephemeral"], - output_channels="output", - checkpointer=checkpointer, - ) - app.invoke({"input": 1, "ephemeral": "meow"}, config) - assert checkpointer.get(config)["channel_values"] == {"input": 1, "output": 1} - - -def test_xray_issue(snapshot: SnapshotAssertion) -> None: - class State(TypedDict): - messages: Annotated[list, add_messages] - - def node(name): - def _node(state: State): - return {"messages": [("human", f"entered {name} node")]} - - return _node - - parent = StateGraph(State) - child = StateGraph(State) - - child.add_node("c_one", node("c_one")) - child.add_node("c_two", node("c_two")) - - child.add_edge("__start__", "c_one") - child.add_edge("c_two", "c_one") - - child.add_conditional_edges( - "c_one", lambda x: str(randrange(0, 2)), {"0": "c_two", "1": "__end__"} - ) - - parent.add_node("p_one", node("p_one")) - parent.add_node("p_two", child.compile()) - - parent.add_edge("__start__", "p_one") - parent.add_edge("p_two", "p_one") - - parent.add_conditional_edges( - "p_one", lambda x: str(randrange(0, 2)), {"0": "p_two", "1": "__end__"} - ) - - app = parent.compile() - - assert app.get_graph(xray=True).draw_mermaid() == snapshot - - -def test_xray_bool(snapshot: SnapshotAssertion) -> None: - class State(TypedDict): - messages: Annotated[list, add_messages] - - def node(name): - def _node(state: State): - return {"messages": [("human", f"entered {name} node")]} - - return _node - - grand_parent = StateGraph(State) - - child = StateGraph(State) - - child.add_node("c_one", node("c_one")) - child.add_node("c_two", node("c_two")) - - child.add_edge("__start__", "c_one") - child.add_edge("c_two", "c_one") - - child.add_conditional_edges( - "c_one", lambda x: str(randrange(0, 2)), {"0": "c_two", "1": "__end__"} - ) - - parent = StateGraph(State) - parent.add_node("p_one", node("p_one")) - parent.add_node("p_two", child.compile()) - parent.add_edge("__start__", "p_one") - parent.add_edge("p_two", "p_one") - parent.add_conditional_edges( - "p_one", lambda x: str(randrange(0, 2)), {"0": "p_two", "1": "__end__"} - ) - - grand_parent.add_node("gp_one", node("gp_one")) - grand_parent.add_node("gp_two", parent.compile()) - grand_parent.add_edge("__start__", "gp_one") - grand_parent.add_edge("gp_two", "gp_one") - grand_parent.add_conditional_edges( - "gp_one", lambda x: str(randrange(0, 2)), {"0": "gp_two", "1": "__end__"} - ) - - app = grand_parent.compile() - assert app.get_graph(xray=True).draw_mermaid() == snapshot - - -def test_multiple_sinks_subgraphs(snapshot: SnapshotAssertion) -> None: - class State(TypedDict): - messages: Annotated[list, add_messages] - - subgraph_builder = StateGraph(State) - subgraph_builder.add_node("one", lambda x: x) - subgraph_builder.add_node("two", lambda x: x) - subgraph_builder.add_node("three", lambda x: x) - subgraph_builder.add_edge("__start__", "one") - subgraph_builder.add_conditional_edges("one", lambda x: "two", ["two", "three"]) - subgraph = subgraph_builder.compile() - - builder = StateGraph(State) - builder.add_node("uno", lambda x: x) - builder.add_node("dos", lambda x: x) - builder.add_node("subgraph", subgraph) - builder.add_edge("__start__", "uno") - builder.add_conditional_edges("uno", lambda x: "dos", ["dos", "subgraph"]) - - app = builder.compile() - assert app.get_graph(xray=True).draw_mermaid() == snapshot - - def test_subgraph_retries(): class State(TypedDict): count: int @@ -4349,86 +2749,6 @@ def test_subgraph_retries(): app.invoke({"count": 0}, {"configurable": {"thread_id": "foo"}}) -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -@pytest.mark.parametrize("store_name", ALL_STORES_SYNC) -def test_store_injected( - request: pytest.FixtureRequest, checkpointer_name: str, store_name: str -) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - the_store = request.getfixturevalue(f"store_{store_name}") - - class State(TypedDict): - count: Annotated[int, operator.add] - - doc_id = str(uuid.uuid4()) - doc = {"some-key": "this-is-a-val"} - uid = uuid.uuid4().hex - namespace = (f"foo-{uid}", "bar") - thread_1 = str(uuid.uuid4()) - thread_2 = str(uuid.uuid4()) - - class Node: - def __init__(self, i: Optional[int] = None): - self.i = i - - def __call__(self, inputs: State, config: RunnableConfig, store: BaseStore): - assert isinstance(store, BaseStore) - store.put( - ( - namespace - if self.i is not None - and config["configurable"]["thread_id"] in (thread_1, thread_2) - else (f"foo_{self.i}", "bar") - ), - doc_id, - { - **doc, - "from_thread": config["configurable"]["thread_id"], - "some_val": inputs["count"], - }, - ) - return {"count": 1} - - builder = StateGraph(State) - builder.add_node("node", Node()) - builder.add_edge("__start__", "node") - N = 500 - M = 1 - - for i in range(N): - builder.add_node(f"node_{i}", Node(i)) - builder.add_edge("__start__", f"node_{i}") - - graph = builder.compile(store=the_store, checkpointer=checkpointer) - - results = graph.batch( - [{"count": 0}] * M, - ([{"configurable": {"thread_id": str(uuid.uuid4())}}] * (M - 1)) - + [{"configurable": {"thread_id": thread_1}}], - ) - result = results[-1] - assert result == {"count": N + 1} - returned_doc = the_store.get(namespace, doc_id).value - assert returned_doc == {**doc, "from_thread": thread_1, "some_val": 0} - assert len(the_store.search(namespace)) == 1 - # Check results after another turn of the same thread - result = graph.invoke({"count": 0}, {"configurable": {"thread_id": thread_1}}) - assert result == {"count": (N + 1) * 2} - returned_doc = the_store.get(namespace, doc_id).value - assert returned_doc == {**doc, "from_thread": thread_1, "some_val": N + 1} - assert len(the_store.search(namespace)) == 1 - - result = graph.invoke({"count": 0}, {"configurable": {"thread_id": thread_2}}) - assert result == {"count": N + 1} - returned_doc = the_store.get(namespace, doc_id).value - assert returned_doc == { - **doc, - "from_thread": thread_2, - "some_val": 0, - } # Overwrites the whole doc - assert len(the_store.search(namespace)) == 1 # still overwriting the same one - - def test_enum_node_names(): class NodeName(str, enum.Enum): BAZ = "baz" @@ -4807,25 +3127,6 @@ def test_add_sequence(): ] -def test_runnable_passthrough_node_graph() -> None: - class State(TypedDict): - changeme: str - - async def dummy(state): - return state - - agent = dummy | RunnablePassthrough.assign(prediction=RunnableLambda(lambda x: x)) - - graph_builder = StateGraph(State) - - graph_builder.add_node("agent", agent) - graph_builder.add_edge(START, "agent") - - graph = graph_builder.compile() - - assert graph.get_graph(xray=True).to_json() == graph.get_graph(xray=False).to_json() - - @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_parent_command(request: pytest.FixtureRequest, checkpointer_name: str) -> None: from langchain_core.messages import BaseMessage @@ -4856,18 +3157,14 @@ def test_parent_command(request: pytest.FixtureRequest, checkpointer_name: str) assert graph.invoke({"messages": [("user", "get user name")]}, config) == { "messages": [ - _AnyIdHumanMessage( - content="get user name", additional_kwargs={}, response_metadata={} - ), + _AnyIdHumanMessage(content="get user name"), ], "user_name": "Meow", } assert graph.get_state(config) == StateSnapshot( values={ "messages": [ - _AnyIdHumanMessage( - content="get user name", additional_kwargs={}, response_metadata={} - ), + _AnyIdHumanMessage(content="get user name"), ], "user_name": "Meow", }, @@ -4884,11 +3181,7 @@ def test_parent_command(request: pytest.FixtureRequest, checkpointer_name: str) "writes": { "alice": { "messages": [ - _AnyIdHumanMessage( - content="get user name", - additional_kwargs={}, - response_metadata={}, - ), + _AnyIdHumanMessage(content="get user name"), ], "user_name": "Meow", } @@ -5085,98 +3378,6 @@ def test_interrupt_loop(request: pytest.FixtureRequest, checkpointer_name: str): ] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_interrupt_functional( - request: pytest.FixtureRequest, checkpointer_name: str, snapshot: SnapshotAssertion -) -> None: - checkpointer: BaseCheckpointSaver = request.getfixturevalue( - f"checkpointer_{checkpointer_name}" - ) - - @task - def foo(state: dict) -> dict: - return {"a": state["a"] + "foo"} - - @task - def bar(state: dict) -> dict: - return {"a": state["a"] + "bar", "b": state["b"]} - - @entrypoint(checkpointer=checkpointer) - def graph(inputs: dict) -> dict: - fut_foo = foo(inputs) - value = interrupt("Provide value for bar:") - bar_input = {**fut_foo.result(), "b": value} - fut_bar = bar(bar_input) - return fut_bar.result() - - config = {"configurable": {"thread_id": "1"}} - # First run, interrupted at bar - graph.invoke({"a": ""}, config) - # Resume with an answer - res = graph.invoke(Command(resume="bar"), config) - assert res == {"a": "foobar", "b": "bar"} - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_interrupt_task_functional( - request: pytest.FixtureRequest, checkpointer_name: str, snapshot: SnapshotAssertion -) -> None: - checkpointer: BaseCheckpointSaver = request.getfixturevalue( - f"checkpointer_{checkpointer_name}" - ) - - @task - def foo(state: dict) -> dict: - return {"a": state["a"] + "foo"} - - @task - def bar(state: dict) -> dict: - value = interrupt("Provide value for bar:") - return {"a": state["a"] + value} - - @entrypoint(checkpointer=checkpointer) - def graph(inputs: dict) -> dict: - fut_foo = foo(inputs) - fut_bar = bar(fut_foo.result()) - return fut_bar.result() - - config = {"configurable": {"thread_id": "1"}} - # First run, interrupted at bar - assert not graph.invoke({"a": ""}, config) - # Resume with an answer - res = graph.invoke(Command(resume="bar"), config) - assert res == {"a": "foobar"} - - # Test that we can interrupt the same task multiple times - config = {"configurable": {"thread_id": "2"}} - - @entrypoint(checkpointer=checkpointer) - def graph(inputs: dict) -> dict: - foo_result = foo(inputs).result() - bar_result = bar(foo_result).result() - baz_result = bar(bar_result).result() - return baz_result - - # First run, interrupted at bar - assert not graph.invoke({"a": ""}, config) - # Provide resumes - assert not graph.invoke(Command(resume="bar"), config) - assert graph.invoke(Command(resume="baz"), config) == {"a": "foobarbaz"} - - -def test_root_mixed_return() -> None: - def my_node(state: list[str]): - return [Command(update=["a"]), ["b"]] - - graph = StateGraph(Annotated[list[str], operator.add]) - - graph.add_node(my_node) - graph.add_edge(START, "my_node") - graph = graph.compile() - - assert graph.invoke([]) == ["a", "b"] - - def test_dict_mixed_return() -> None: class State(TypedDict): foo: Annotated[str, operator.add] @@ -5286,16 +3487,16 @@ def test_multistep_plan(request: pytest.FixtureRequest, checkpointer_name: str): pass def step1(state: State): - return Command(goto="planner", update={"messages": [("human", "step1")]}) + return Command(goto="planner", update={"messages": [("user", "step1")]}) def step2(state: State): - return Command(goto="planner", update={"messages": [("human", "step2")]}) + return Command(goto="planner", update={"messages": [("user", "step2")]}) def step3(state: State): - return Command(goto="planner", update={"messages": [("human", "step3")]}) + return Command(goto="planner", update={"messages": [("user", "step3")]}) def step4(state: State): - return Command(goto="planner", update={"messages": [("human", "step4")]}) + return Command(goto="planner", update={"messages": [("user", "step4")]}) builder = StateGraph(State) builder.add_node(planner) @@ -5308,7 +3509,7 @@ def test_multistep_plan(request: pytest.FixtureRequest, checkpointer_name: str): config = {"configurable": {"thread_id": "1"}} - assert graph.invoke({"messages": [("human", "start")]}, config) == { + assert graph.invoke({"messages": [("user", "start")]}, config) == { "messages": [ _AnyIdHumanMessage(content="start"), _AnyIdHumanMessage(content="step1"), @@ -5359,37 +3560,6 @@ def test_command_goto_with_static_breakpoints( assert result == {"foo": "abc|node-1|node-2|node-2"} -def test_nested_graph_state_error_handling(): - """Test error handling when updating state in nested graphs.""" - - class State(TypedDict): - count: int - - def child_node(state: State): - return {"count": state["count"] + 1} - - child = StateGraph(State) - child.add_node("child", child_node) - child.add_edge(START, "child") - - parent = StateGraph(State) - parent.add_node("child_graph", child.compile()) - parent.add_edge(START, "child_graph") - - app = parent.compile(checkpointer=MemorySaver()) - - # Test invalid state update on parent - with pytest.raises(InvalidUpdateError): - app.update_state({"configurable": {"thread_id": "1"}}, {"invalid_key": "value"}) - - # Test invalid state update on child - with pytest.raises(InvalidUpdateError): - app.update_state( - {"configurable": {"thread_id": "1", "checkpoint_ns": "child_graph"}}, - {"invalid_key": "value"}, - ) - - def test_parallel_node_execution(): """Test that parallel nodes execute concurrently.""" @@ -5561,29 +3731,6 @@ def test_checkpoint_recovery(request: pytest.FixtureRequest, checkpointer_name: assert "RuntimeError('Simulated failure')" in failed_checkpoint.tasks[0].error -def test_multiple_updates_root() -> None: - def node_a(state): - return [Command(update="a1"), Command(update="a2")] - - def node_b(state): - return "b" - - graph = ( - StateGraph(Annotated[str, operator.add]) - .add_sequence([node_a, node_b]) - .add_edge(START, "node_a") - .compile() - ) - - assert graph.invoke("") == "a1a2b" - - # only streams the last update from node_a - assert [c for c in graph.stream("", stream_mode="updates")] == [ - {"node_a": ["a1", "a2"]}, - {"node_b": "b"}, - ] - - def test_multiple_updates() -> None: class State(TypedDict): foo: Annotated[str, operator.add] @@ -5612,67 +3759,6 @@ def test_multiple_updates() -> None: ] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_falsy_return_from_task( - request: pytest.FixtureRequest, checkpointer_name: str, snapshot: SnapshotAssertion -): - """Test with a falsy return from a task.""" - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - - @task - def falsy_task() -> bool: - return False - - @entrypoint(checkpointer=checkpointer) - def graph(state: dict) -> dict: - """React tool.""" - falsy_task().result() - interrupt("test") - - configurable = {"configurable": {"thread_id": str(uuid.uuid4())}} - graph.invoke({"a": 5}, configurable) - graph.invoke(Command(resume="123"), configurable) - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_multiple_interrupts_functional( - request: pytest.FixtureRequest, checkpointer_name: str, snapshot: SnapshotAssertion -): - """Test multiple interrupts with functional API.""" - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - - counter = 0 - - @task - def double(x: int) -> int: - """Increment the counter.""" - nonlocal counter - counter += 1 - return 2 * x - - @entrypoint(checkpointer=checkpointer) - def graph(state: dict) -> dict: - """React tool.""" - - values = [] - - for idx in [1, 2, 3]: - values.extend([double(idx).result(), interrupt({"a": "boo"})]) - - return {"values": values} - - configurable = {"configurable": {"thread_id": str(uuid.uuid4())}} - graph.invoke({}, configurable) - graph.invoke(Command(resume="a"), configurable) - graph.invoke(Command(resume="b"), configurable) - result = graph.invoke(Command(resume="c"), configurable) - # `double` value should be cached appropriately when used w/ `interrupt` - assert result == { - "values": [2, "a", 4, "b", 6, "c"], - } - assert counter == 3 - - @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_double_interrupt_subgraph( request: pytest.FixtureRequest, checkpointer_name: str @@ -5785,118 +3871,6 @@ def test_double_interrupt_subgraph( ] -def test_sync_streaming_with_functional_api() -> None: - """Test streaming with functional API. - - This test verifies that we're able to stream results as they're being generated - rather than have all the results arrive at once after the graph has completed. - - The time of arrival between the two updates corresponding to the two `slow` tasks - should be greater than the time delay between the two tasks. - """ - - time_delay = 0.01 - - @task() - def slow() -> dict: - time.sleep(time_delay) # Simulate a delay of 10 ms - return {"tic": time.time()} - - @entrypoint() - def graph(inputs: dict) -> list: - first = slow().result() - second = slow().result() - return [first, second] - - arrival_times = [] - - for chunk in graph.stream({}): - if "slow" not in chunk: # We'll just look at the updates from `slow` - continue - arrival_times.append(time.time()) - - assert len(arrival_times) == 2 - delta = arrival_times[1] - arrival_times[0] - # Delta cannot be less than 10 ms if it is streaming as results are generated. - assert delta > time_delay - - -def test_entrypoint_without_checkpointer() -> None: - """Test no checkpointer.""" - states = [] - config = {"configurable": {"thread_id": "1"}} - - # Test without previous - @entrypoint() - def foo(inputs: Any) -> Any: - states.append(inputs) - return inputs - - assert foo.invoke({"a": "1"}, config) == {"a": "1"} - - @entrypoint() - def foo(inputs: Any, *, previous: Any) -> Any: - states.append(previous) - return {"previous": previous, "current": inputs} - - assert foo.invoke({"a": "1"}, config) == {"current": {"a": "1"}, "previous": None} - assert foo.invoke({"a": "1"}, config) == {"current": {"a": "1"}, "previous": None} - - -def test_entrypoint_stateful() -> None: - """Test stateful entrypoint invoke.""" - - # Test invoke - states = [] - - @entrypoint(checkpointer=MemorySaver()) - def foo(inputs, *, previous: Any) -> Any: - states.append(previous) - return {"previous": previous, "current": inputs} - - config = {"configurable": {"thread_id": "1"}} - - assert foo.invoke({"a": "1"}, config) == {"current": {"a": "1"}, "previous": None} - assert foo.invoke({"a": "2"}, config) == { - "current": {"a": "2"}, - "previous": {"current": {"a": "1"}, "previous": None}, - } - assert foo.invoke({"a": "3"}, config) == { - "current": {"a": "3"}, - "previous": { - "current": {"a": "2"}, - "previous": {"current": {"a": "1"}, "previous": None}, - }, - } - assert states == [ - None, - {"current": {"a": "1"}, "previous": None}, - {"current": {"a": "2"}, "previous": {"current": {"a": "1"}, "previous": None}}, - ] - - # Test stream - @entrypoint(checkpointer=MemorySaver()) - def foo(inputs, *, previous: Any) -> Any: - return {"previous": previous, "current": inputs} - - config = {"configurable": {"thread_id": "1"}} - items = [item for item in foo.stream({"a": "1"}, config)] - assert items == [{"foo": {"current": {"a": "1"}, "previous": None}}] - - -def test_entrypoint_from_sync_generator() -> None: - """@entrypoint does not support sync generators.""" - previous_return_values = [] - - with pytest.raises(NotImplementedError): - - @entrypoint(checkpointer=MemorySaver()) - def foo(inputs, previous=None) -> Any: - previous_return_values.append(previous) - yield "a" - yield "b" - - @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_multiple_subgraphs( request: pytest.FixtureRequest, checkpointer_name: str @@ -5969,187 +3943,6 @@ def test_multiple_subgraphs( } -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_multiple_subgraphs_functional( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - - # Define addition subgraph - @entrypoint() - def add(inputs: tuple[int, int]): - a, b = inputs - return a + b - - # Define multiplication subgraph using tasks - @task - def multiply_task(a, b): - return a * b - - @entrypoint() - def multiply(inputs: tuple[int, int]): - return multiply_task(*inputs).result() - - # Test calling the same subgraph multiple times - @task - def call_same_subgraph(a, b): - result = add.invoke([a, b]) - another_result = add.invoke([result, 10]) - return another_result - - @entrypoint(checkpointer=checkpointer) - def parent_call_same_subgraph(inputs): - return call_same_subgraph(*inputs).result() - - config = {"configurable": {"thread_id": "1"}} - assert parent_call_same_subgraph.invoke([2, 3], config) == 15 - - # Test calling multiple subgraphs - @task - def call_multiple_subgraphs(a, b): - add_result = add.invoke([a, b]) - multiply_result = multiply.invoke([a, b]) - return [add_result, multiply_result] - - @entrypoint(checkpointer=checkpointer) - def parent_call_multiple_subgraphs(inputs): - return call_multiple_subgraphs(*inputs).result() - - config = {"configurable": {"thread_id": "2"}} - assert parent_call_multiple_subgraphs.invoke([2, 3], config) == [5, 6] - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_multiple_subgraphs_mixed_entrypoint( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - """Test calling multiple StateGraph subgraphs from an entrypoint.""" - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - - class State(TypedDict): - a: int - b: int - - class Output(TypedDict): - result: int - - # Define the subgraphs - def add(state): - return {"result": state["a"] + state["b"]} - - add_subgraph = ( - StateGraph(State, output=Output).add_node(add).add_edge(START, "add").compile() - ) - - def multiply(state): - return {"result": state["a"] * state["b"]} - - multiply_subgraph = ( - StateGraph(State, output=Output) - .add_node(multiply) - .add_edge(START, "multiply") - .compile() - ) - - # Test calling the same subgraph multiple times - @task - def call_same_subgraph(a, b): - result = add_subgraph.invoke({"a": a, "b": b})["result"] - another_result = add_subgraph.invoke({"a": result, "b": 10})["result"] - return another_result - - @entrypoint(checkpointer=checkpointer) - def parent_call_same_subgraph(inputs): - return call_same_subgraph(*inputs).result() - - config = {"configurable": {"thread_id": "1"}} - assert parent_call_same_subgraph.invoke([2, 3], config) == 15 - - # Test calling multiple subgraphs - @task - def call_multiple_subgraphs(a, b): - add_result = add_subgraph.invoke({"a": a, "b": b})["result"] - multiply_result = multiply_subgraph.invoke({"a": a, "b": b})["result"] - return [add_result, multiply_result] - - @entrypoint(checkpointer=checkpointer) - def parent_call_multiple_subgraphs(inputs): - return call_multiple_subgraphs(*inputs).result() - - config = {"configurable": {"thread_id": "2"}} - assert parent_call_multiple_subgraphs.invoke([2, 3], config) == [5, 6] - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_multiple_subgraphs_mixed_state_graph( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - """Test calling multiple entrypoint "subgraphs" from a StateGraph.""" - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - - class State(TypedDict): - a: int - b: int - - class Output(TypedDict): - result: int - - # Define addition subgraph - @entrypoint() - def add(inputs: tuple[int, int]): - a, b = inputs - return a + b - - # Define multiplication subgraph using tasks - @task - def multiply_task(a, b): - return a * b - - @entrypoint() - def multiply(inputs: tuple[int, int]): - return multiply_task(*inputs).result() - - # Test calling the same subgraph multiple times - def call_same_subgraph(state): - result = add.invoke([state["a"], state["b"]]) - another_result = add.invoke([result, 10]) - return {"result": another_result} - - parent_call_same_subgraph = ( - StateGraph(State, output=Output) - .add_node(call_same_subgraph) - .add_edge(START, "call_same_subgraph") - .compile(checkpointer=checkpointer) - ) - config = {"configurable": {"thread_id": "1"}} - assert parent_call_same_subgraph.invoke({"a": 2, "b": 3}, config) == {"result": 15} - - # Test calling multiple subgraphs - class Output(TypedDict): - add_result: int - multiply_result: int - - def call_multiple_subgraphs(state): - add_result = add.invoke([state["a"], state["b"]]) - multiply_result = multiply.invoke([state["a"], state["b"]]) - return { - "add_result": add_result, - "multiply_result": multiply_result, - } - - parent_call_multiple_subgraphs = ( - StateGraph(State, output=Output) - .add_node(call_multiple_subgraphs) - .add_edge(START, "call_multiple_subgraphs") - .compile(checkpointer=checkpointer) - ) - config = {"configurable": {"thread_id": "2"}} - assert parent_call_multiple_subgraphs.invoke({"a": 2, "b": 3}, config) == { - "add_result": 5, - "multiply_result": 6, - } - - @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_multiple_subgraphs_checkpointer( request: pytest.FixtureRequest, checkpointer_name: str @@ -6273,7 +4066,7 @@ def test_merging_updates_command_parent(): update={"bar": ["node_1"]}, ) - def node_3(state: State, store): + def node_3(state: State): return Command( update={"bar": ["node_3"]}, ) @@ -6356,7 +4149,7 @@ def test_merging_non_overlapping_updates_command_parent(): update={"foo": ["foo"]}, ) - def node_3(state: State, store): + def node_3(state: State): return Command( update={"foo": ["baz"]}, ) @@ -6374,231 +4167,6 @@ def test_merging_non_overlapping_updates_command_parent(): } -def test_entrypoint_output_schema_with_return_and_save() -> None: - """Test output schema inference with entrypoint.final.""" - - # Un-parameterized entrypoint.final is interpreted as entrypoint.final[Any, Any] - @entrypoint() - def foo2(inputs, *, previous: Any) -> entrypoint.final: - return entrypoint.final(value="foo", save=1) - - assert foo2.get_output_schema().model_json_schema() == { - "title": "LangGraphOutput", - } - - @entrypoint() - def foo(inputs, *, previous: Any) -> entrypoint.final[str, int]: - return entrypoint.final(value="foo", save=1) - - assert foo.get_output_schema().model_json_schema() == { - "title": "LangGraphOutput", - "type": "string", - } - - with pytest.raises(TypeError): - # Raise an exception on an improperly parameterized entrypoint.final - # User is attempting to parameterize in this case, so we'll offer - # a bit of help if it's not done correctly. - @entrypoint() - def foo(inputs, *, previous: Any) -> entrypoint.final[int]: - return entrypoint.final(value=1, save=1) # type: ignore - - -def test_entrypoint_with_return_and_save() -> None: - """Test entrypoint with return and save.""" - previous_ = None - - @entrypoint(checkpointer=MemorySaver()) - def foo(msg: str, *, previous: Any) -> entrypoint.final[int, list[str]]: - nonlocal previous_ - previous_ = previous - previous = previous or [] - return entrypoint.final(value=len(previous), save=previous + [msg]) - - assert foo.get_output_schema().model_json_schema() == { - "title": "LangGraphOutput", - "type": "integer", - } - - config = {"configurable": {"thread_id": "1"}} - assert foo.invoke("hello", config) == 0 - assert previous_ is None - assert foo.invoke("goodbye", config) == 1 - assert previous_ == ["hello"] - assert foo.invoke("definitely", config) == 2 - assert previous_ == ["hello", "goodbye"] - - -def test_overriding_injectable_args_with_tasks() -> None: - """Test overriding injectable args in tasks.""" - from langgraph.store.memory import InMemoryStore - - @task - def foo(store: BaseStore, writer: StreamWriter, value: Any) -> None: - assert store is value - assert writer is value - - @entrypoint(store=InMemoryStore()) - def main(inputs, store: BaseStore) -> str: - assert store is not None - foo(store=None, writer=None, value=None).result() - foo(store="hello", writer="hello", value="hello").result() - return "OK" - - assert main.invoke({}) == "OK" - - -def test_named_tasks_functional() -> None: - class Foo: - def foo(self, value: str) -> dict: - return value + "foo" - - f = Foo() - - # class method task - foo = task(f.foo, name="custom_foo") - other_foo = task(f.foo, name="other_foo") - - # regular function task - @task(name="custom_bar") - def bar(value: str) -> dict: - return value + "|bar" - - def baz(update: str, value: str) -> dict: - return value + f"|{update}" - - # partial function task (unnamed) - baz_task = task(functools.partial(baz, "baz")) - # partial function task (named_) - custom_baz_task = task(functools.partial(baz, "custom_baz"), name="custom_baz") - - class Qux: - def __call__(self, value: str) -> dict: - return value + "|qux" - - qux_task = task(Qux(), name="qux") - - @entrypoint() - def workflow(inputs: dict) -> dict: - foo_result = foo(inputs).result() - other_foo(inputs).result() - fut_bar = bar(foo_result) - fut_baz = baz_task(fut_bar.result()) - fut_custom_baz = custom_baz_task(fut_baz.result()) - fut_qux = qux_task(fut_custom_baz.result()) - return fut_qux.result() - - assert list(workflow.stream("", stream_mode="updates")) == [ - {"custom_foo": "foo"}, - {"other_foo": "foo"}, - {"custom_bar": "foo|bar"}, - {"baz": "foo|bar|baz"}, - {"custom_baz": "foo|bar|baz|custom_baz"}, - {"qux": "foo|bar|baz|custom_baz|qux"}, - {"workflow": "foo|bar|baz|custom_baz|qux"}, - ] - - -def test_tags_stream_mode_messages() -> None: - model = GenericFakeChatModel(messages=iter(["foo"]), tags=["meow"]) - graph = ( - StateGraph(MessagesState) - .add_node( - "call_model", lambda state: {"messages": model.invoke(state["messages"])} - ) - .add_edge(START, "call_model") - .compile() - ) - assert list( - graph.stream( - { - "messages": "hi", - }, - stream_mode="messages", - ) - ) == [ - ( - _AnyIdAIMessageChunk(content="foo"), - { - "langgraph_step": 1, - "langgraph_node": "call_model", - "langgraph_triggers": ["start:call_model"], - "langgraph_path": ("__pregel_pull", "call_model"), - "langgraph_checkpoint_ns": AnyStr("call_model:"), - "checkpoint_ns": AnyStr("call_model:"), - "ls_provider": "genericfakechatmodel", - "ls_model_type": "chat", - "tags": ["meow"], - }, - ) - ] - - -def test_node_destinations() -> None: - class State(TypedDict): - foo: Annotated[str, operator.add] - - def node_a(state: State): - value = state["foo"] - if value == "a": - goto = "node_b" - else: - goto = "node_c" - - return Command( - update={"foo": value}, - goto=goto, - graph=Command.PARENT, - ) - - subgraph = StateGraph(State).add_node(node_a).add_edge(START, "node_a").compile() - - # test calling subgraph inside a node function - def call_subgraph(state: State): - return subgraph.invoke(state) - - def node_b(state: State): - return {"foo": "b"} - - def node_c(state: State): - return {"foo": "c"} - - for subgraph_node in (subgraph, call_subgraph): - # destinations w/ tuples - builder = StateGraph(State) - builder.add_edge(START, "child") - builder.add_node("child", subgraph_node, destinations=("node_b", "node_c")) - builder.add_node(node_b) - builder.add_node(node_c) - compiled_graph = builder.compile() - assert compiled_graph.invoke({"foo": ""}) == {"foo": "c"} - - graph = compiled_graph.get_graph() - assert [ - Edge(source="__start__", target="child", data=None, conditional=False), - Edge(source="child", target="node_b", data=None, conditional=True), - Edge(source="child", target="node_c", data=None, conditional=True), - ] == graph.edges - - # destinations w/ dicts - builder = StateGraph(State) - builder.add_edge(START, "child") - builder.add_node( - "child", subgraph_node, destinations={"node_b": "foo", "node_c": "bar"} - ) - builder.add_node(node_b) - builder.add_node(node_c) - compiled_graph = builder.compile() - assert compiled_graph.invoke({"foo": ""}) == {"foo": "c"} - - graph = compiled_graph.get_graph() - assert [ - Edge(source="__start__", target="child", data=None, conditional=False), - Edge(source="child", target="node_b", data="foo", conditional=True), - Edge(source="child", target="node_c", data="bar", conditional=True), - ] == graph.edges - - def test_pydantic_none_state_update() -> None: from pydantic import BaseModel @@ -6643,104 +4211,6 @@ def test_get_stream_writer() -> None: ] -def test_stream_messages_dedupe_inputs() -> None: - from langchain_core.messages import AIMessage - - def call_model(state): - return {"messages": AIMessage("hi", id="1")} - - def route(state): - return Command(goto="node_2", graph=Command.PARENT) - - subgraph = ( - StateGraph(MessagesState) - .add_node(call_model) - .add_node(route) - .add_edge(START, "call_model") - .add_edge("call_model", "route") - .compile() - ) - - graph = ( - StateGraph(MessagesState) - .add_node("node_1", subgraph) - .add_node("node_2", lambda state: state) - .add_edge(START, "node_1") - .compile() - ) - - chunks = [ - chunk - for ns, chunk in graph.stream( - {"messages": "hi"}, stream_mode="messages", subgraphs=True - ) - ] - - assert len(chunks) == 1 - assert chunks[0][0] == AIMessage("hi", id="1") - assert chunks[0][1]["langgraph_node"] == "call_model" - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_stream_messages_dedupe_state( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - from langchain_core.messages import AIMessage - - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - to_emit = [AIMessage("bye", id="1"), AIMessage("bye again", id="2")] - - def call_model(state): - return {"messages": to_emit.pop(0)} - - def route(state): - return Command(goto="node_2", graph=Command.PARENT) - - subgraph = ( - StateGraph(MessagesState) - .add_node(call_model) - .add_node(route) - .add_edge(START, "call_model") - .add_edge("call_model", "route") - .compile() - ) - - graph = ( - StateGraph(MessagesState) - .add_node("node_1", subgraph) - .add_node("node_2", lambda state: state) - .add_edge(START, "node_1") - .compile(checkpointer=checkpointer) - ) - - thread1 = {"configurable": {"thread_id": "1"}} - - chunks = [ - chunk - for ns, chunk in graph.stream( - {"messages": "hi"}, thread1, stream_mode="messages", subgraphs=True - ) - ] - - assert len(chunks) == 1 - assert chunks[0][0] == AIMessage("bye", id="1") - assert chunks[0][1]["langgraph_node"] == "call_model" - - chunks = [ - chunk - for ns, chunk in graph.stream( - {"messages": "hi again"}, - thread1, - stream_mode="messages", - subgraphs=True, - ) - ] - - assert len(chunks) == 1 - assert chunks[0][0] == AIMessage("bye again", id="2") - assert chunks[0][1]["langgraph_node"] == "call_model" - - @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_interrupt_subgraph_reenter_checkpointer_true( request: pytest.FixtureRequest, checkpointer_name: str diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py deleted file mode 100644 index ce54bb44c..000000000 --- a/libs/langgraph/tests/test_pregel_async.py +++ /dev/null @@ -1,7706 +0,0 @@ -import asyncio -import functools -import logging -import operator -import random -import sys -import uuid -from collections import Counter, deque -from contextlib import asynccontextmanager, contextmanager -from dataclasses import replace -from time import perf_counter -from typing import ( - Annotated, - Any, - AsyncGenerator, - AsyncIterator, - Dict, - Generator, - List, - Literal, - Optional, - Tuple, - Union, -) -from uuid import UUID - -import httpx -import pytest -from langchain_core.language_models import GenericFakeChatModel -from langchain_core.runnables import ( - RunnableConfig, - RunnableLambda, - RunnablePassthrough, -) -from langchain_core.utils.aiter import aclosing -from pytest_mock import MockerFixture -from syrupy import SnapshotAssertion -from typing_extensions import TypedDict - -from langgraph.channels.base import BaseChannel -from langgraph.channels.binop import BinaryOperatorAggregate -from langgraph.channels.context import Context -from langgraph.channels.last_value import LastValue -from langgraph.channels.topic import Topic -from langgraph.checkpoint.base import ( - ChannelVersions, - Checkpoint, - CheckpointMetadata, - CheckpointTuple, -) -from langgraph.checkpoint.memory import InMemorySaver, MemorySaver -from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, PUSH, START -from langgraph.errors import InvalidUpdateError, NodeInterrupt -from langgraph.func import entrypoint, task -from langgraph.graph import END, Graph, StateGraph -from langgraph.graph.message import MessagesState, add_messages -from langgraph.prebuilt.tool_node import ToolNode -from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot -from langgraph.pregel.retry import RetryPolicy -from langgraph.store.base import BaseStore -from langgraph.types import ( - Command, - Interrupt, - PregelTask, - Send, - StreamWriter, - interrupt, -) -from tests.any_str import AnyStr, AnyVersion, FloatBetween, UnsortedSequence -from tests.conftest import ( - ALL_CHECKPOINTERS_ASYNC, - ALL_CHECKPOINTERS_ASYNC_PLUS_NONE, - ALL_STORES_ASYNC, - REGULAR_CHECKPOINTERS_ASYNC, - SHOULD_CHECK_SNAPSHOTS, - awith_checkpointer, - awith_store, -) -from tests.fake_tracer import FakeTracer -from tests.memory_assert import ( - MemorySaverAssertCheckpointMetadata, - MemorySaverNoPending, -) -from tests.messages import ( - _AnyIdAIMessage, - _AnyIdAIMessageChunk, - _AnyIdHumanMessage, - _AnyIdToolMessage, -) - -logger = logging.getLogger(__name__) - -pytestmark = pytest.mark.anyio - -NEEDS_CONTEXTVARS = pytest.mark.skipif( - sys.version_info < (3, 11), - reason="Python 3.11+ is required for async contextvars support", -) - - -async def test_checkpoint_errors() -> None: - class FaultyGetCheckpointer(InMemorySaver): - async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: - raise ValueError("Faulty get_tuple") - - class FaultyPutCheckpointer(InMemorySaver): - async def aput( - self, - config: RunnableConfig, - checkpoint: Checkpoint, - metadata: CheckpointMetadata, - new_versions: ChannelVersions, - ) -> RunnableConfig: - raise ValueError("Faulty put") - - class FaultyPutWritesCheckpointer(InMemorySaver): - async def aput_writes( - self, config: RunnableConfig, writes: List[Tuple[str, Any]], task_id: str - ) -> RunnableConfig: - raise ValueError("Faulty put_writes") - - class FaultyVersionCheckpointer(InMemorySaver): - def get_next_version(self, current: Optional[int], channel: BaseChannel) -> int: - raise ValueError("Faulty get_next_version") - - def logic(inp: str) -> str: - return "" - - builder = StateGraph(Annotated[str, operator.add]) - builder.add_node("agent", logic) - builder.add_edge(START, "agent") - - graph = builder.compile(checkpointer=FaultyGetCheckpointer()) - with pytest.raises(ValueError, match="Faulty get_tuple"): - await graph.ainvoke("", {"configurable": {"thread_id": "thread-1"}}) - with pytest.raises(ValueError, match="Faulty get_tuple"): - async for _ in graph.astream("", {"configurable": {"thread_id": "thread-2"}}): - pass - with pytest.raises(ValueError, match="Faulty get_tuple"): - async for _ in graph.astream_events( - "", {"configurable": {"thread_id": "thread-3"}}, version="v2" - ): - pass - - graph = builder.compile(checkpointer=FaultyPutCheckpointer()) - with pytest.raises(ValueError, match="Faulty put"): - await graph.ainvoke("", {"configurable": {"thread_id": "thread-1"}}) - with pytest.raises(ValueError, match="Faulty put"): - async for _ in graph.astream("", {"configurable": {"thread_id": "thread-2"}}): - pass - with pytest.raises(ValueError, match="Faulty put"): - async for _ in graph.astream_events( - "", {"configurable": {"thread_id": "thread-3"}}, version="v2" - ): - pass - - graph = builder.compile(checkpointer=FaultyVersionCheckpointer()) - with pytest.raises(ValueError, match="Faulty get_next_version"): - await graph.ainvoke("", {"configurable": {"thread_id": "thread-1"}}) - with pytest.raises(ValueError, match="Faulty get_next_version"): - async for _ in graph.astream("", {"configurable": {"thread_id": "thread-2"}}): - pass - with pytest.raises(ValueError, match="Faulty get_next_version"): - async for _ in graph.astream_events( - "", {"configurable": {"thread_id": "thread-3"}}, version="v2" - ): - pass - - # add a parallel node - builder.add_node("parallel", logic) - builder.add_edge(START, "parallel") - graph = builder.compile(checkpointer=FaultyPutWritesCheckpointer()) - with pytest.raises(ValueError, match="Faulty put_writes"): - await graph.ainvoke("", {"configurable": {"thread_id": "thread-1"}}) - with pytest.raises(ValueError, match="Faulty put_writes"): - async for _ in graph.astream("", {"configurable": {"thread_id": "thread-2"}}): - pass - with pytest.raises(ValueError, match="Faulty put_writes"): - async for _ in graph.astream_events( - "", {"configurable": {"thread_id": "thread-3"}}, version="v2" - ): - pass - - -async def test_py_async_with_cancel_behavior() -> None: - """This test confirms that in all versions of Python we support, __aexit__ - is not cancelled when the coroutine containing the async with block is cancelled.""" - - logs: list[str] = [] - - class MyContextManager: - async def __aenter__(self): - logs.append("Entering") - return self - - async def __aexit__(self, exc_type, exc_val, exc_tb): - logs.append("Starting exit") - try: - # Simulate some cleanup work - await asyncio.sleep(2) - logs.append("Cleanup completed") - except asyncio.CancelledError: - logs.append("Cleanup was cancelled!") - raise - logs.append("Exit finished") - - async def main(): - try: - async with MyContextManager(): - logs.append("In context") - await asyncio.sleep(1) - logs.append("This won't print if cancelled") - except asyncio.CancelledError: - logs.append("Context was cancelled") - raise - - # create task - t = asyncio.create_task(main()) - # cancel after 0.2 seconds - await asyncio.sleep(0.2) - t.cancel() - # check logs before cancellation is handled - assert logs == [ - "Entering", - "In context", - ], "Cancelled before cleanup started" - # wait for task to finish - try: - await t - except asyncio.CancelledError: - # check logs after cancellation is handled - assert logs == [ - "Entering", - "In context", - "Starting exit", - "Cleanup completed", - "Exit finished", - "Context was cancelled", - ], "Cleanup started and finished after cancellation" - else: - assert False, "Task should be cancelled" - - -async def test_checkpoint_put_after_cancellation() -> None: - logs: list[str] = [] - - class LongPutCheckpointer(MemorySaver): - async def aput( - self, - config: RunnableConfig, - checkpoint: Checkpoint, - metadata: CheckpointMetadata, - new_versions: ChannelVersions, - ) -> RunnableConfig: - logs.append("checkpoint.aput.start") - try: - await asyncio.sleep(1) - return await super().aput(config, checkpoint, metadata, new_versions) - finally: - logs.append("checkpoint.aput.end") - - inner_task_cancelled = False - - async def awhile(input: Any) -> None: - logs.append("awhile.start") - try: - await asyncio.sleep(1) - except asyncio.CancelledError: - nonlocal inner_task_cancelled - inner_task_cancelled = True - raise - finally: - logs.append("awhile.end") - - builder = Graph() - builder.add_node("agent", awhile) - builder.set_entry_point("agent") - builder.set_finish_point("agent") - - graph = builder.compile(checkpointer=LongPutCheckpointer()) - thread1 = {"configurable": {"thread_id": "1"}} - - # start the task - t = asyncio.create_task(graph.ainvoke(1, thread1)) - # cancel after 0.2 seconds - await asyncio.sleep(0.2) - t.cancel() - # check logs before cancellation is handled - assert sorted(logs) == [ - "awhile.start", - "checkpoint.aput.start", - ], "Cancelled before checkpoint put started" - # wait for task to finish - try: - await t - except asyncio.CancelledError: - # check logs after cancellation is handled - assert sorted(logs) == [ - "awhile.end", - "awhile.start", - "checkpoint.aput.end", - "checkpoint.aput.start", - ], "Checkpoint put is not cancelled" - else: - assert False, "Task should be cancelled" - - -async def test_checkpoint_put_after_cancellation_stream_anext() -> None: - logs: list[str] = [] - - class LongPutCheckpointer(MemorySaver): - async def aput( - self, - config: RunnableConfig, - checkpoint: Checkpoint, - metadata: CheckpointMetadata, - new_versions: ChannelVersions, - ) -> RunnableConfig: - logs.append("checkpoint.aput.start") - try: - await asyncio.sleep(1) - return await super().aput(config, checkpoint, metadata, new_versions) - finally: - logs.append("checkpoint.aput.end") - - inner_task_cancelled = False - - async def awhile(input: Any) -> None: - logs.append("awhile.start") - try: - await asyncio.sleep(1) - except asyncio.CancelledError: - nonlocal inner_task_cancelled - inner_task_cancelled = True - raise - finally: - logs.append("awhile.end") - - builder = Graph() - builder.add_node("agent", awhile) - builder.set_entry_point("agent") - builder.set_finish_point("agent") - - graph = builder.compile(checkpointer=LongPutCheckpointer()) - thread1 = {"configurable": {"thread_id": "1"}} - - # start the task - s = graph.astream(1, thread1) - t = asyncio.create_task(s.__anext__()) - # cancel after 0.2 seconds - await asyncio.sleep(0.2) - t.cancel() - # check logs before cancellation is handled - assert sorted(logs) == [ - "awhile.start", - "checkpoint.aput.start", - ], "Cancelled before checkpoint put started" - # wait for task to finish - try: - await t - except asyncio.CancelledError: - # check logs after cancellation is handled - assert sorted(logs) == [ - "awhile.end", - "awhile.start", - "checkpoint.aput.end", - "checkpoint.aput.start", - ], "Checkpoint put is not cancelled" - else: - assert False, "Task should be cancelled" - - -async def test_checkpoint_put_after_cancellation_stream_events_anext() -> None: - logs: list[str] = [] - - class LongPutCheckpointer(MemorySaver): - async def aput( - self, - config: RunnableConfig, - checkpoint: Checkpoint, - metadata: CheckpointMetadata, - new_versions: ChannelVersions, - ) -> RunnableConfig: - logs.append("checkpoint.aput.start") - try: - await asyncio.sleep(1) - return await super().aput(config, checkpoint, metadata, new_versions) - finally: - logs.append("checkpoint.aput.end") - - inner_task_cancelled = False - - async def awhile(input: Any) -> None: - logs.append("awhile.start") - try: - await asyncio.sleep(1) - except asyncio.CancelledError: - nonlocal inner_task_cancelled - inner_task_cancelled = True - raise - finally: - logs.append("awhile.end") - - builder = Graph() - builder.add_node("agent", awhile) - builder.set_entry_point("agent") - builder.set_finish_point("agent") - - graph = builder.compile(checkpointer=LongPutCheckpointer()) - thread1 = {"configurable": {"thread_id": "1"}} - - # start the task - s = graph.astream_events(1, thread1, version="v2", include_names=["LangGraph"]) - # skip first event (happens right away) - await s.__anext__() - # start the task for 2nd event - t = asyncio.create_task(s.__anext__()) - # cancel after 0.2 seconds - await asyncio.sleep(0.2) - t.cancel() - # check logs before cancellation is handled - assert logs == [ - "checkpoint.aput.start", - "awhile.start", - ], "Cancelled before checkpoint put started" - # wait for task to finish - try: - await t - except asyncio.CancelledError: - # check logs after cancellation is handled - assert logs == [ - "checkpoint.aput.start", - "awhile.start", - "awhile.end", - "checkpoint.aput.end", - ], "Checkpoint put is not cancelled" - else: - assert False, "Task should be cancelled" - - -async def test_node_cancellation_on_external_cancel() -> None: - inner_task_cancelled = False - - async def awhile(input: Any) -> None: - try: - await asyncio.sleep(1) - except asyncio.CancelledError: - nonlocal inner_task_cancelled - inner_task_cancelled = True - raise - - builder = Graph() - builder.add_node("agent", awhile) - builder.set_entry_point("agent") - builder.set_finish_point("agent") - - graph = builder.compile() - - with pytest.raises(asyncio.TimeoutError): - await asyncio.wait_for(graph.ainvoke(1), 0.5) - - assert inner_task_cancelled - - -async def test_node_cancellation_on_other_node_exception() -> None: - inner_task_cancelled = False - - async def awhile(input: Any) -> None: - try: - await asyncio.sleep(1) - except asyncio.CancelledError: - nonlocal inner_task_cancelled - inner_task_cancelled = True - raise - - async def iambad(input: Any) -> None: - raise ValueError("I am bad") - - builder = Graph() - builder.add_node("agent", awhile) - builder.add_node("bad", iambad) - builder.set_conditional_entry_point(lambda _: ["agent", "bad"], then=END) - - graph = builder.compile() - - with pytest.raises(ValueError, match="I am bad"): - # This will raise ValueError, not TimeoutError - await asyncio.wait_for(graph.ainvoke(1), 0.5) - - assert inner_task_cancelled - - -async def test_node_cancellation_on_other_node_exception_two() -> None: - async def awhile(input: Any) -> None: - await asyncio.sleep(1) - - async def iambad(input: Any) -> None: - raise ValueError("I am bad") - - builder = Graph() - builder.add_node("agent", awhile) - builder.add_node("bad", iambad) - builder.set_conditional_entry_point(lambda _: ["agent", "bad"], then=END) - - graph = builder.compile() - - with pytest.raises(ValueError, match="I am bad"): - # This will raise ValueError, not CancelledError - await graph.ainvoke(1) - - -@NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_dynamic_interrupt(checkpointer_name: str) -> None: - class State(TypedDict): - my_key: Annotated[str, operator.add] - market: str - - tool_two_node_count = 0 - - async def tool_two_node(s: State) -> State: - nonlocal tool_two_node_count - tool_two_node_count += 1 - if s["market"] == "DE": - answer = interrupt("Just because...") - else: - answer = " all good" - return {"my_key": answer} - - tool_two_graph = StateGraph(State) - tool_two_graph.add_node("tool_two", tool_two_node, retry=RetryPolicy()) - tool_two_graph.add_edge(START, "tool_two") - tool_two = tool_two_graph.compile() - - tracer = FakeTracer() - assert await tool_two.ainvoke( - {"my_key": "value", "market": "DE"}, {"callbacks": [tracer]} - ) == { - "my_key": "value", - "market": "DE", - } - assert tool_two_node_count == 1, "interrupts aren't retried" - assert len(tracer.runs) == 1 - run = tracer.runs[0] - assert run.end_time is not None - assert run.error is None - assert run.outputs == {"market": "DE", "my_key": "value"} - - assert await tool_two.ainvoke({"my_key": "value", "market": "US"}) == { - "my_key": "value all good", - "market": "US", - } - - async with awith_checkpointer(checkpointer_name) as checkpointer: - tool_two = tool_two_graph.compile(checkpointer=checkpointer) - - # missing thread_id - with pytest.raises(ValueError, match="thread_id"): - await tool_two.ainvoke({"my_key": "value", "market": "DE"}) - - # flow: interrupt -> resume with answer - thread2 = {"configurable": {"thread_id": "2"}} - # stop when about to enter node - assert [ - c - async for c in tool_two.astream( - {"my_key": "value ⛰️", "market": "DE"}, thread2 - ) - ] == [ - { - "__interrupt__": ( - Interrupt( - value="Just because...", - resumable=True, - ns=[AnyStr("tool_two:")], - ), - ) - }, - ] - # resume with answer - assert [ - c async for c in tool_two.astream(Command(resume=" my answer"), thread2) - ] == [ - {"tool_two": {"my_key": " my answer"}}, - ] - - # flow: interrupt -> clear - thread1 = {"configurable": {"thread_id": "1"}} - # stop when about to enter node - assert [ - c - async for c in tool_two.astream( - {"my_key": "value ⛰️", "market": "DE"}, thread1 - ) - ] == [ - { - "__interrupt__": ( - Interrupt( - value="Just because...", - resumable=True, - ns=[AnyStr("tool_two:")], - ), - ) - }, - ] - if "shallow" not in checkpointer_name: - assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [ - { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "thread_id": "1", - }, - { - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, - "thread_id": "1", - }, - ] - tup = await tool_two.checkpointer.aget_tuple(thread1) - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value ⛰️", "market": "DE"}, - next=("tool_two",), - tasks=( - PregelTask( - AnyStr(), - "tool_two", - (PULL, "tool_two"), - interrupts=( - Interrupt( - value="Just because...", - resumable=True, - ns=[AnyStr("tool_two:")], - ), - ), - ), - ), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "thread_id": "1", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ - -1 - ].config - ), - ) - - # clear the interrupt and next tasks - await tool_two.aupdate_state(thread1, None, as_node=END) - # interrupt is cleared, as well as the next tasks - tup = await tool_two.checkpointer.aget_tuple(thread1) - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value ⛰️", "market": "DE"}, - next=(), - tasks=(), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 1, - "writes": {}, - "thread_id": "1", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ - -1 - ].config - ), - ) - - -@NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_dynamic_interrupt_subgraph(checkpointer_name: str) -> None: - class SubgraphState(TypedDict): - my_key: str - market: str - - tool_two_node_count = 0 - - def tool_two_node(s: SubgraphState) -> SubgraphState: - nonlocal tool_two_node_count - tool_two_node_count += 1 - if s["market"] == "DE": - answer = interrupt("Just because...") - else: - answer = " all good" - return {"my_key": answer} - - subgraph = StateGraph(SubgraphState) - subgraph.add_node("do", tool_two_node, retry=RetryPolicy()) - subgraph.add_edge(START, "do") - - class State(TypedDict): - my_key: Annotated[str, operator.add] - market: str - - tool_two_graph = StateGraph(State) - tool_two_graph.add_node("tool_two", subgraph.compile()) - tool_two_graph.add_edge(START, "tool_two") - tool_two = tool_two_graph.compile() - - tracer = FakeTracer() - assert await tool_two.ainvoke( - {"my_key": "value", "market": "DE"}, {"callbacks": [tracer]} - ) == { - "my_key": "value", - "market": "DE", - } - assert tool_two_node_count == 1, "interrupts aren't retried" - assert len(tracer.runs) == 1 - run = tracer.runs[0] - assert run.end_time is not None - assert run.error is None - assert run.outputs == {"market": "DE", "my_key": "value"} - - assert await tool_two.ainvoke({"my_key": "value", "market": "US"}) == { - "my_key": "value all good", - "market": "US", - } - - async with awith_checkpointer(checkpointer_name) as checkpointer: - tool_two = tool_two_graph.compile(checkpointer=checkpointer) - - # missing thread_id - with pytest.raises(ValueError, match="thread_id"): - await tool_two.ainvoke({"my_key": "value", "market": "DE"}) - - # flow: interrupt -> resume with answer - thread2 = {"configurable": {"thread_id": "2"}} - # stop when about to enter node - assert [ - c - async for c in tool_two.astream( - {"my_key": "value ⛰️", "market": "DE"}, thread2 - ) - ] == [ - { - "__interrupt__": ( - Interrupt( - value="Just because...", - resumable=True, - ns=[AnyStr("tool_two:"), AnyStr("do:")], - ), - ) - }, - ] - # resume with answer - assert [ - c async for c in tool_two.astream(Command(resume=" my answer"), thread2) - ] == [ - {"tool_two": {"my_key": " my answer", "market": "DE"}}, - ] - - # flow: interrupt -> clear - thread1 = {"configurable": {"thread_id": "1"}} - thread1root = {"configurable": {"thread_id": "1", "checkpoint_ns": ""}} - # stop when about to enter node - assert [ - c - async for c in tool_two.astream( - {"my_key": "value ⛰️", "market": "DE"}, thread1 - ) - ] == [ - { - "__interrupt__": ( - Interrupt( - value="Just because...", - resumable=True, - ns=[AnyStr("tool_two:"), AnyStr("do:")], - ), - ) - }, - ] - if "shallow" not in checkpointer_name: - assert [ - c.metadata async for c in tool_two.checkpointer.alist(thread1root) - ] == [ - { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "thread_id": "1", - }, - { - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, - "thread_id": "1", - }, - ] - tup = await tool_two.checkpointer.aget_tuple(thread1) - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value ⛰️", "market": "DE"}, - next=("tool_two",), - tasks=( - PregelTask( - AnyStr(), - "tool_two", - (PULL, "tool_two"), - interrupts=( - Interrupt( - value="Just because...", - resumable=True, - ns=[AnyStr("tool_two:"), AnyStr("do:")], - ), - ), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("tool_two:"), - } - }, - ), - ), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "thread_id": "1", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in tool_two.checkpointer.alist(thread1root, limit=2) - ][-1].config - ), - ) - - # clear the interrupt and next tasks - await tool_two.aupdate_state(thread1, None, as_node=END) - # interrupt is cleared, as well as the next tasks - tup = await tool_two.checkpointer.aget_tuple(thread1) - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value ⛰️", "market": "DE"}, - next=(), - tasks=(), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 1, - "writes": {}, - "thread_id": "1", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in tool_two.checkpointer.alist(thread1root, limit=2) - ][-1].config - ), - ) - - -@NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_copy_checkpoint(checkpointer_name: str) -> None: - class State(TypedDict): - my_key: Annotated[str, operator.add] - market: str - - def tool_one(s: State) -> State: - return {"my_key": " one"} - - tool_two_node_count = 0 - - def tool_two_node(s: State) -> State: - nonlocal tool_two_node_count - tool_two_node_count += 1 - if s["market"] == "DE": - answer = interrupt("Just because...") - else: - answer = " all good" - return {"my_key": answer} - - def start(state: State) -> list[Union[Send, str]]: - return ["tool_two", Send("tool_one", state)] - - tool_two_graph = StateGraph(State) - tool_two_graph.add_node("tool_two", tool_two_node, retry=RetryPolicy()) - tool_two_graph.add_node("tool_one", tool_one) - tool_two_graph.set_conditional_entry_point(start) - tool_two = tool_two_graph.compile() - - tracer = FakeTracer() - assert await tool_two.ainvoke( - {"my_key": "value", "market": "DE"}, {"callbacks": [tracer]}, debug=True - ) == { - "my_key": "value one", - "market": "DE", - } - assert tool_two_node_count == 1, "interrupts aren't retried" - assert len(tracer.runs) == 1 - run = tracer.runs[0] - assert run.end_time is not None - assert run.error is None - assert run.outputs == {"market": "DE", "my_key": "value one"} - - assert await tool_two.ainvoke({"my_key": "value", "market": "US"}) == { - "my_key": "value all good one", - "market": "US", - } - - async with awith_checkpointer(checkpointer_name) as checkpointer: - tool_two = tool_two_graph.compile(checkpointer=checkpointer) - - # missing thread_id - with pytest.raises(ValueError, match="thread_id"): - await tool_two.ainvoke({"my_key": "value", "market": "DE"}) - - # flow: interrupt -> resume with answer - thread2 = {"configurable": {"thread_id": "2"}} - # stop when about to enter node - assert [ - c - async for c in tool_two.astream( - {"my_key": "value ⛰️", "market": "DE"}, thread2 - ) - ] == [ - { - "tool_one": {"my_key": " one"}, - }, - { - "__interrupt__": ( - Interrupt( - value="Just because...", - resumable=True, - ns=[AnyStr("tool_two:")], - ), - ) - }, - ] - # resume with answer - assert [ - c async for c in tool_two.astream(Command(resume=" my answer"), thread2) - ] == [ - { - "__metadata__": {"cached": True}, - "tool_one": {"my_key": " one"}, - }, - {"tool_two": {"my_key": " my answer"}}, - ] - - # flow: interrupt -> clear tasks - thread1 = {"configurable": {"thread_id": "1"}} - # stop when about to enter node - assert await tool_two.ainvoke( - {"my_key": "value ⛰️", "market": "DE"}, thread1 - ) == { - "my_key": "value ⛰️ one", - "market": "DE", - } - - if "shallow" not in checkpointer_name: - assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [ - { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "thread_id": "1", - }, - { - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, - "thread_id": "1", - }, - ] - - tup = await tool_two.checkpointer.aget_tuple(thread1) - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value ⛰️ one", "market": "DE"}, - next=("tool_two",), - tasks=( - PregelTask( - AnyStr(), - name="tool_one", - path=("__pregel_push", 0), - error=None, - interrupts=(), - state=None, - result={"my_key": " one"}, - ), - PregelTask( - AnyStr(), - "tool_two", - (PULL, "tool_two"), - interrupts=( - Interrupt( - value="Just because...", - resumable=True, - ns=[AnyStr("tool_two:")], - ), - ), - ), - ), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "thread_id": "1", - }, - parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ - -1 - ].config - ), - ) - - if "shallow" in checkpointer_name: - # shallow checkpointer doesn't support copy - return - - # clear the interrupt and next tasks - await tool_two.aupdate_state(thread1, None, as_node="__copy__") - # interrupt is cleared, next task is kept - tup = await tool_two.checkpointer.aget_tuple(thread1) - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value ⛰️", "market": "DE"}, - next=("tool_one", "tool_two"), - tasks=( - PregelTask( - AnyStr(), - "tool_one", - (PUSH, 0), - result=None, - ), - PregelTask( - AnyStr(), - "tool_two", - (PULL, "tool_two"), - interrupts=(), - ), - ), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "fork", - "step": 1, - "writes": None, - "thread_id": "1", - }, - parent_config=( - [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ - -1 - ].parent_config - ), - ) - - -@NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_node_not_cancelled_on_other_node_interrupted( - checkpointer_name: str, -) -> None: - class State(TypedDict): - hello: Annotated[str, operator.add] - - awhiles = 0 - inner_task_cancelled = False - - async def awhile(input: State) -> None: - nonlocal awhiles - - awhiles += 1 - try: - await asyncio.sleep(1) - return {"hello": " again"} - except asyncio.CancelledError: - nonlocal inner_task_cancelled - inner_task_cancelled = True - raise - - async def iambad(input: State) -> None: - return {"hello": interrupt("I am bad")} - - builder = StateGraph(State) - builder.add_node("agent", awhile) - builder.add_node("bad", iambad) - builder.set_conditional_entry_point(lambda _: ["agent", "bad"], then=END) - - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer) - thread = {"configurable": {"thread_id": "1"}} - - # writes from "awhile" are applied to last chunk - assert await graph.ainvoke({"hello": "world"}, thread) == { - "hello": "world again" - } - - assert not inner_task_cancelled - assert awhiles == 1 - - assert await graph.ainvoke(None, thread, debug=True) == {"hello": "world again"} - - assert not inner_task_cancelled - assert awhiles == 1 - - # resume with answer - assert await graph.ainvoke(Command(resume=" okay"), thread) == { - "hello": "world again okay" - } - - assert not inner_task_cancelled - assert awhiles == 1 - - -@pytest.mark.parametrize("stream_hang_s", [0.3, 0.6]) -async def test_step_timeout_on_stream_hang(stream_hang_s: float) -> None: - inner_task_cancelled = False - - async def awhile(input: Any) -> None: - try: - await asyncio.sleep(1.5) - except asyncio.CancelledError: - nonlocal inner_task_cancelled - inner_task_cancelled = True - raise - - async def alittlewhile(input: Any) -> None: - await asyncio.sleep(0.6) - return "1" - - builder = Graph() - builder.add_node(awhile) - builder.add_node(alittlewhile) - builder.set_conditional_entry_point(lambda _: ["awhile", "alittlewhile"], then=END) - graph = builder.compile() - graph.step_timeout = 1 - - with pytest.raises(asyncio.TimeoutError): - async for chunk in graph.astream(1, stream_mode="updates"): - assert chunk == {"alittlewhile": {"alittlewhile": "1"}} - await asyncio.sleep(stream_hang_s) - - assert inner_task_cancelled - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC_PLUS_NONE) -async def test_cancel_graph_astream(checkpointer_name: str) -> None: - class State(TypedDict): - value: Annotated[int, operator.add] - - class AwhileMaker: - def __init__(self) -> None: - self.reset() - - async def __call__(self, input: State) -> Any: - self.started = True - try: - await asyncio.sleep(1.5) - except asyncio.CancelledError: - self.cancelled = True - raise - - def reset(self): - self.started = False - self.cancelled = False - - async def alittlewhile(input: State) -> None: - await asyncio.sleep(0.6) - return {"value": 2} - - awhile = AwhileMaker() - aparallelwhile = AwhileMaker() - builder = StateGraph(State) - builder.add_node("awhile", awhile) - builder.add_node("aparallelwhile", aparallelwhile) - builder.add_node(alittlewhile) - builder.add_edge(START, "alittlewhile") - builder.add_edge(START, "aparallelwhile") - builder.add_edge("alittlewhile", "awhile") - - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer) - - # test interrupting astream - got_event = False - thread1: RunnableConfig = {"configurable": {"thread_id": "1"}} - async with aclosing(graph.astream({"value": 1}, thread1)) as stream: - async for chunk in stream: - assert chunk == {"alittlewhile": {"value": 2}} - got_event = True - break - - assert got_event - - # node aparallelwhile should start, but be cancelled - assert aparallelwhile.started is True - assert aparallelwhile.cancelled is True - - # node "awhile" should never start - assert awhile.started is False - - # checkpoint with output of "alittlewhile" should not be saved - # but we should have applied pending writes - if checkpointer is not None: - state = await graph.aget_state(thread1) - assert state is not None - assert state.values == {"value": 3} # 1 + 2 - assert state.next == ("aparallelwhile",) - assert state.metadata == { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "thread_id": "1", - } - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC_PLUS_NONE) -async def test_cancel_graph_astream_events_v2(checkpointer_name: Optional[str]) -> None: - class State(TypedDict): - value: int - - class AwhileMaker: - def __init__(self) -> None: - self.reset() - - async def __call__(self, input: State) -> Any: - self.started = True - try: - await asyncio.sleep(1.5) - except asyncio.CancelledError: - self.cancelled = True - raise - - def reset(self): - self.started = False - self.cancelled = False - - async def alittlewhile(input: State) -> None: - await asyncio.sleep(0.6) - return {"value": 2} - - awhile = AwhileMaker() - anotherwhile = AwhileMaker() - builder = StateGraph(State) - builder.add_node(alittlewhile) - builder.add_node("awhile", awhile) - builder.add_node("anotherwhile", anotherwhile) - builder.add_edge(START, "alittlewhile") - builder.add_edge("alittlewhile", "awhile") - builder.add_edge("awhile", "anotherwhile") - - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer) - - # test interrupting astream_events v2 - got_event = False - thread2: RunnableConfig = {"configurable": {"thread_id": "2"}} - async with aclosing( - graph.astream_events({"value": 1}, thread2, version="v2") - ) as stream: - async for chunk in stream: - if chunk["event"] == "on_chain_stream" and not chunk["parent_ids"]: - got_event = True - assert chunk["data"]["chunk"] == {"alittlewhile": {"value": 2}} - await asyncio.sleep(0.1) - break - - # did break - assert got_event - - # node "awhile" maybe starts (impl detail of astream_events) - # if it does start, it must be cancelled - if awhile.started: - assert awhile.cancelled is True - - # node "anotherwhile" should never start - assert anotherwhile.started is False - - # checkpoint with output of "alittlewhile" should not be saved - if checkpointer is not None: - state = await graph.aget_state(thread2) - assert state is not None - assert state.values == {"value": 2} - assert state.next == ("awhile",) - assert state.metadata == { - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"alittlewhile": {"value": 2}}, - "thread_id": "2", - } - - -async def test_node_schemas_custom_output() -> None: - class State(TypedDict): - hello: str - bye: str - messages: Annotated[list[str], add_messages] - - class Output(TypedDict): - messages: list[str] - - class StateForA(TypedDict): - hello: str - messages: Annotated[list[str], add_messages] - - async def node_a(state: StateForA): - assert state == { - "hello": "there", - "messages": [_AnyIdHumanMessage(content="hello")], - } - - class StateForB(TypedDict): - bye: str - now: int - - async def node_b(state: StateForB): - assert state == { - "bye": "world", - } - return { - "now": 123, - "hello": "again", - } - - class StateForC(TypedDict): - hello: str - now: int - - async def node_c(state: StateForC): - assert state == { - "hello": "again", - "now": 123, - } - - builder = StateGraph(State, output=Output) - builder.add_node("a", node_a) - builder.add_node("b", node_b) - builder.add_node("c", node_c) - builder.add_edge(START, "a") - builder.add_edge("a", "b") - builder.add_edge("b", "c") - graph = builder.compile() - - assert await graph.ainvoke( - {"hello": "there", "bye": "world", "messages": "hello"} - ) == { - "messages": [_AnyIdHumanMessage(content="hello")], - } - - builder = StateGraph(State, output=Output) - builder.add_node("a", node_a) - builder.add_node("b", node_b) - builder.add_node("c", node_c) - builder.add_edge(START, "a") - builder.add_edge("a", "b") - builder.add_edge("b", "c") - graph = builder.compile() - - assert await graph.ainvoke( - { - "hello": "there", - "bye": "world", - "messages": "hello", - "now": 345, # ignored because not in input schema - } - ) == { - "messages": [_AnyIdHumanMessage(content="hello")], - } - - assert [ - c - async for c in graph.astream( - { - "hello": "there", - "bye": "world", - "messages": "hello", - "now": 345, # ignored because not in input schema - } - ) - ] == [ - {"a": None}, - {"b": {"hello": "again", "now": 123}}, - {"c": None}, - ] - - -async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: - add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") - - app = Pregel( - nodes={ - "one": chain, - }, - channels={ - "input": LastValue(int), - "output": LastValue(int), - }, - input_channels="input", - output_channels="output", - ) - graph = Graph() - graph.add_node("add_one", add_one) - graph.set_entry_point("add_one") - graph.set_finish_point("add_one") - gapp = graph.compile() - - if SHOULD_CHECK_SNAPSHOTS: - assert app.input_schema.model_json_schema() == { - "title": "LangGraphInput", - "type": "integer", - } - assert app.output_schema.model_json_schema() == { - "title": "LangGraphOutput", - "type": "integer", - } - assert await app.ainvoke(2) == 3 - assert await app.ainvoke(2, output_keys=["output"]) == {"output": 3} - - assert await gapp.ainvoke(2) == 3 - - -@pytest.mark.parametrize( - "falsy_value", - [None, False, 0, "", [], {}, set(), frozenset(), 0.0, 0j], -) -async def test_invoke_single_process_in_out_falsy_values(falsy_value: Any) -> None: - graph = Graph() - graph.add_node("return_falsy_const", lambda *args, **kwargs: falsy_value) - graph.set_entry_point("return_falsy_const") - graph.set_finish_point("return_falsy_const") - gapp = graph.compile() - assert falsy_value == await gapp.ainvoke(1) - - -async def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> None: - add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain = ( - Channel.subscribe_to("input") - | add_one - | Channel.write_to("output", fixed=5, output_plus_one=lambda x: x + 1) - ) - - app = Pregel( - nodes={"one": chain}, - channels={ - "input": LastValue(int), - "output": LastValue(int), - "fixed": LastValue(int), - "output_plus_one": LastValue(int), - }, - output_channels=["output", "fixed", "output_plus_one"], - input_channels="input", - ) - - if SHOULD_CHECK_SNAPSHOTS: - assert app.input_schema.model_json_schema() == { - "title": "LangGraphInput", - "type": "integer", - } - assert app.output_schema.model_json_schema() == { - "title": "LangGraphOutput", - "type": "object", - "properties": { - "output": {"title": "Output", "type": "integer", "default": None}, - "fixed": {"title": "Fixed", "type": "integer", "default": None}, - "output_plus_one": { - "title": "Output Plus One", - "type": "integer", - "default": None, - }, - }, - } - assert await app.ainvoke(2) == {"output": 3, "fixed": 5, "output_plus_one": 4} - - -async def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None: - add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") - - app = Pregel( - nodes={"one": chain}, - channels={"input": LastValue(int), "output": LastValue(int)}, - input_channels="input", - output_channels=["output"], - ) - - if SHOULD_CHECK_SNAPSHOTS: - assert app.input_schema.model_json_schema() == { - "title": "LangGraphInput", - "type": "integer", - } - assert app.output_schema.model_json_schema() == { - "title": "LangGraphOutput", - "type": "object", - "properties": { - "output": {"title": "Output", "type": "integer", "default": None} - }, - } - assert await app.ainvoke(2) == {"output": 3} - - -async def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None: - add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") - - app = Pregel( - nodes={"one": chain}, - channels={"input": LastValue(int), "output": LastValue(int)}, - input_channels=["input"], - output_channels=["output"], - ) - - if SHOULD_CHECK_SNAPSHOTS: - assert app.input_schema.model_json_schema() == { - "title": "LangGraphInput", - "type": "object", - "properties": { - "input": {"title": "Input", "type": "integer", "default": None} - }, - } - assert app.output_schema.model_json_schema() == { - "title": "LangGraphOutput", - "type": "object", - "properties": { - "output": {"title": "Output", "type": "integer", "default": None} - }, - } - assert await app.ainvoke({"input": 2}) == {"output": 3} - - -async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: - add_one = mocker.Mock(side_effect=lambda x: x + 1) - one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") - two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output") - - app = Pregel( - nodes={"one": one, "two": two}, - channels={ - "inbox": LastValue(int), - "output": LastValue(int), - "input": LastValue(int), - }, - input_channels="input", - output_channels="output", - stream_channels=["inbox", "output"], - ) - - assert await app.ainvoke(2) == 4 - - with pytest.raises(GraphRecursionError): - await app.ainvoke(2, {"recursion_limit": 1}) - - step = 0 - async for values in app.astream(2): - step += 1 - if step == 1: - assert values == { - "inbox": 3, - } - elif step == 2: - assert values == { - "inbox": 3, - "output": 4, - } - assert step == 2 - - graph = Graph() - graph.add_node("add_one", add_one) - graph.add_node("add_one_more", add_one) - graph.set_entry_point("add_one") - graph.set_finish_point("add_one_more") - graph.add_edge("add_one", "add_one_more") - gapp = graph.compile() - - assert await gapp.ainvoke(2) == 4 - - step = 0 - async for values in gapp.astream(2): - step += 1 - if step == 1: - assert values == { - "add_one": 3, - } - elif step == 2: - assert values == { - "add_one_more": 4, - } - assert step == 2 - - -async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: - add_one = mocker.Mock(side_effect=lambda x: x + 1) - one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") - two = ( - Channel.subscribe_to("inbox") - | RunnableLambda(add_one).abatch - | Channel.write_to("output").abatch - ) - - app = Pregel( - nodes={"one": one, "two": two}, - channels={ - "inbox": Topic(int), - "output": LastValue(int), - "input": LastValue(int), - }, - input_channels=["input", "inbox"], - stream_channels=["output", "inbox"], - output_channels=["output"], - ) - - # [12 + 1, 2 + 1 + 1] - assert [ - c - async for c in app.astream( - {"input": 2, "inbox": 12}, output_keys="output", stream_mode="updates" - ) - ] == [ - {"one": None}, - {"two": 13}, - {"two": 4}, - ] - assert [ - c async for c in app.astream({"input": 2, "inbox": 12}, output_keys="output") - ] == [13, 4] - - assert [ - c async for c in app.astream({"input": 2, "inbox": 12}, stream_mode="updates") - ] == [ - {"one": {"inbox": 3}}, - {"two": {"output": 13}}, - {"two": {"output": 4}}, - ] - assert [c async for c in app.astream({"input": 2, "inbox": 12})] == [ - {"inbox": [3], "output": 13}, - {"output": 4}, - ] - assert [ - c async for c in app.astream({"input": 2, "inbox": 12}, stream_mode="debug") - ] == [ - { - "type": "task", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "id": AnyStr(), - "name": "one", - "input": 2, - "triggers": ["input"], - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "id": AnyStr(), - "name": "two", - "input": [12], - "triggers": ["inbox"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "id": AnyStr(), - "name": "one", - "result": [("inbox", 3)], - "error": None, - "interrupts": [], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "id": AnyStr(), - "name": "two", - "result": [("output", 13)], - "error": None, - "interrupts": [], - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": AnyStr(), - "name": "two", - "input": [3], - "triggers": ["inbox"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": AnyStr(), - "name": "two", - "result": [("output", 4)], - "error": None, - "interrupts": [], - }, - }, - ] - - -async def test_batch_two_processes_in_out() -> None: - async def add_one_with_delay(inp: int) -> int: - await asyncio.sleep(inp / 10) - return inp + 1 - - one = Channel.subscribe_to("input") | add_one_with_delay | Channel.write_to("one") - two = Channel.subscribe_to("one") | add_one_with_delay | Channel.write_to("output") - - app = Pregel( - nodes={"one": one, "two": two}, - channels={ - "one": LastValue(int), - "output": LastValue(int), - "input": LastValue(int), - }, - input_channels="input", - output_channels="output", - ) - - assert await app.abatch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7] - assert await app.abatch([3, 2, 1, 3, 5], output_keys=["output"]) == [ - {"output": 5}, - {"output": 4}, - {"output": 3}, - {"output": 5}, - {"output": 7}, - ] - - graph = Graph() - graph.add_node("add_one", add_one_with_delay) - graph.add_node("add_one_more", add_one_with_delay) - graph.set_entry_point("add_one") - graph.set_finish_point("add_one_more") - graph.add_edge("add_one", "add_one_more") - gapp = graph.compile() - - assert await gapp.abatch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7] - - -async def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None: - test_size = 100 - add_one = mocker.Mock(side_effect=lambda x: x + 1) - - nodes = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")} - for i in range(test_size - 2): - nodes[str(i)] = ( - Channel.subscribe_to(str(i - 1)) | add_one | Channel.write_to(str(i)) - ) - nodes["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output") - - app = Pregel( - nodes=nodes, - channels={str(i): LastValue(int) for i in range(-1, test_size - 2)} - | {"input": LastValue(int), "output": LastValue(int)}, - input_channels="input", - output_channels="output", - ) - - # No state is left over from previous invocations - for _ in range(10): - assert await app.ainvoke(2, {"recursion_limit": test_size}) == 2 + test_size - - # Concurrent invocations do not interfere with each other - assert await asyncio.gather( - *(app.ainvoke(2, {"recursion_limit": test_size}) for _ in range(10)) - ) == [2 + test_size for _ in range(10)] - - -async def test_batch_many_processes_in_out(mocker: MockerFixture) -> None: - test_size = 100 - add_one = mocker.Mock(side_effect=lambda x: x + 1) - - nodes = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")} - for i in range(test_size - 2): - nodes[str(i)] = ( - Channel.subscribe_to(str(i - 1)) | add_one | Channel.write_to(str(i)) - ) - nodes["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output") - - app = Pregel( - nodes=nodes, - channels={str(i): LastValue(int) for i in range(-1, test_size - 2)} - | {"input": LastValue(int), "output": LastValue(int)}, - input_channels="input", - output_channels="output", - ) - - # No state is left over from previous invocations - for _ in range(3): - # Then invoke pubsub - assert await app.abatch([2, 1, 3, 4, 5], {"recursion_limit": test_size}) == [ - 2 + test_size, - 1 + test_size, - 3 + test_size, - 4 + test_size, - 5 + test_size, - ] - - # Concurrent invocations do not interfere with each other - assert await asyncio.gather( - *(app.abatch([2, 1, 3, 4, 5], {"recursion_limit": test_size}) for _ in range(3)) - ) == [ - [2 + test_size, 1 + test_size, 3 + test_size, 4 + test_size, 5 + test_size] - for _ in range(3) - ] - - -async def test_invoke_two_processes_two_in_two_out_invalid( - mocker: MockerFixture, -) -> None: - add_one = mocker.Mock(side_effect=lambda x: x + 1) - - one = Channel.subscribe_to("input") | add_one | Channel.write_to("output") - two = Channel.subscribe_to("input") | add_one | Channel.write_to("output") - - app = Pregel( - nodes={"one": one, "two": two}, - channels={"output": LastValue(int), "input": LastValue(int)}, - input_channels="input", - output_channels="output", - ) - - with pytest.raises(InvalidUpdateError): - # LastValue channels can only be updated once per iteration - await app.ainvoke(2) - - -async def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> None: - add_one = mocker.Mock(side_effect=lambda x: x + 1) - - one = Channel.subscribe_to("input") | add_one | Channel.write_to("output") - two = Channel.subscribe_to("input") | add_one | Channel.write_to("output") - - app = Pregel( - nodes={"one": one, "two": two}, - channels={ - "input": LastValue(int), - "output": Topic(int), - }, - input_channels="input", - output_channels="output", - ) - - # An Topic channel accumulates updates into a sequence - assert await app.ainvoke(2) == [3, 3] - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_invoke_checkpoint(mocker: MockerFixture, checkpointer_name: str) -> None: - add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) - errored_once = False - - def raise_if_above_10(input: int) -> int: - nonlocal errored_once - if input > 4: - if errored_once: - pass - else: - errored_once = True - raise ConnectionError("I will be retried") - if input > 10: - raise ValueError("Input is too large") - return input - - one = ( - Channel.subscribe_to(["input"]).join(["total"]) - | add_one - | Channel.write_to("output", "total") - | raise_if_above_10 - ) - - async with awith_checkpointer(checkpointer_name) as checkpointer: - app = Pregel( - nodes={"one": one}, - channels={ - "total": BinaryOperatorAggregate(int, operator.add), - "input": LastValue(int), - "output": LastValue(int), - }, - input_channels="input", - output_channels="output", - checkpointer=checkpointer, - retry_policy=RetryPolicy(), - ) - - # total starts out as 0, so output is 0+2=2 - assert await app.ainvoke(2, {"configurable": {"thread_id": "1"}}) == 2 - checkpoint = await checkpointer.aget({"configurable": {"thread_id": "1"}}) - assert checkpoint is not None - assert checkpoint["channel_values"].get("total") == 2 - # total is now 2, so output is 2+3=5 - assert await app.ainvoke(3, {"configurable": {"thread_id": "1"}}) == 5 - assert errored_once, "errored and retried" - checkpoint = await checkpointer.aget({"configurable": {"thread_id": "1"}}) - assert checkpoint is not None - assert checkpoint["channel_values"].get("total") == 7 - # total is now 2+5=7, so output would be 7+4=11, but raises ValueError - with pytest.raises(ValueError): - await app.ainvoke(4, {"configurable": {"thread_id": "1"}}) - # checkpoint is not updated - checkpoint = await checkpointer.aget({"configurable": {"thread_id": "1"}}) - assert checkpoint is not None - assert checkpoint["channel_values"].get("total") == 7 - # on a new thread, total starts out as 0, so output is 0+5=5 - assert await app.ainvoke(5, {"configurable": {"thread_id": "2"}}) == 5 - checkpoint = await checkpointer.aget({"configurable": {"thread_id": "1"}}) - assert checkpoint is not None - assert checkpoint["channel_values"].get("total") == 7 - checkpoint = await checkpointer.aget({"configurable": {"thread_id": "2"}}) - assert checkpoint is not None - assert checkpoint["channel_values"].get("total") == 5 - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_pending_writes_resume( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - class State(TypedDict): - value: Annotated[int, operator.add] - - class AwhileMaker: - def __init__(self, sleep: float, rtn: Union[Dict, Exception]) -> None: - self.sleep = sleep - self.rtn = rtn - self.reset() - - async def __call__(self, input: State) -> Any: - self.calls += 1 - await asyncio.sleep(self.sleep) - if isinstance(self.rtn, Exception): - raise self.rtn - else: - return self.rtn - - def reset(self): - self.calls = 0 - - one = AwhileMaker(0.1, {"value": 2}) - two = AwhileMaker(0.3, ConnectionError("I'm not good")) - builder = StateGraph(State) - builder.add_node("one", one) - builder.add_node("two", two, retry=RetryPolicy(max_attempts=2)) - builder.add_edge(START, "one") - builder.add_edge(START, "two") - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer) - - thread1: RunnableConfig = {"configurable": {"thread_id": "1"}} - with pytest.raises(ConnectionError, match="I'm not good"): - await graph.ainvoke({"value": 1}, thread1) - - # both nodes should have been called once - assert one.calls == 1 - assert two.calls == 2 - - # latest checkpoint should be before nodes "one", "two" - # but we should have applied pending writes from "one" - state = await graph.aget_state(thread1) - assert state is not None - assert state.values == {"value": 3} - assert state.next == ("two",) - assert state.tasks == ( - PregelTask(AnyStr(), "one", (PULL, "one"), result={"value": 2}), - PregelTask( - AnyStr(), - "two", - (PULL, "two"), - 'ConnectionError("I\'m not good")', - ), - ) - assert state.metadata == { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "thread_id": "1", - } - # get_state with checkpoint_id should not apply any pending writes - state = await graph.aget_state(state.config) - assert state is not None - assert state.values == {"value": 1} - assert state.next == ("one", "two") - # should contain pending write of "one" - checkpoint = await checkpointer.aget_tuple(thread1) - assert checkpoint is not None - # should contain error from "two" - expected_writes = [ - (AnyStr(), "one", "one"), - (AnyStr(), "value", 2), - (AnyStr(), ERROR, 'ConnectionError("I\'m not good")'), - ] - assert len(checkpoint.pending_writes) == 3 - assert all(w in expected_writes for w in checkpoint.pending_writes) - # both non-error pending writes come from same task - non_error_writes = [w for w in checkpoint.pending_writes if w[1] != ERROR] - assert non_error_writes[0][0] == non_error_writes[1][0] - # error write is from the other task - error_write = next(w for w in checkpoint.pending_writes if w[1] == ERROR) - assert error_write[0] != non_error_writes[0][0] - - # resume execution - with pytest.raises(ConnectionError, match="I'm not good"): - await graph.ainvoke(None, thread1) - - # node "one" succeeded previously, so shouldn't be called again - assert one.calls == 1 - # node "two" should have been called once again - assert two.calls == 4 - - # confirm no new checkpoints saved - state_two = await graph.aget_state(thread1) - assert state_two.metadata == state.metadata - - # resume execution, without exception - two.rtn = {"value": 3} - # both the pending write and the new write were applied, 1 + 2 + 3 = 6 - assert await graph.ainvoke(None, thread1) == {"value": 6} - - if "shallow" in checkpointer_name: - assert len([c async for c in checkpointer.alist(thread1)]) == 1 - return - - # check all final checkpoints - checkpoints = [c async for c in checkpointer.alist(thread1)] - # we should have 3 - assert len(checkpoints) == 3 - # the last one not too interesting for this test - assert checkpoints[0] == CheckpointTuple( - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - checkpoint={ - "v": 1, - "id": AnyStr(), - "ts": AnyStr(), - "pending_sends": [], - "versions_seen": { - "one": { - "start:one": AnyVersion(), - }, - "two": { - "start:two": AnyVersion(), - }, - "__input__": {}, - "__start__": { - "__start__": AnyVersion(), - }, - "__interrupt__": { - "value": AnyVersion(), - "__start__": AnyVersion(), - "start:one": AnyVersion(), - "start:two": AnyVersion(), - }, - }, - "channel_versions": { - "one": AnyVersion(), - "two": AnyVersion(), - "value": AnyVersion(), - "__start__": AnyVersion(), - "start:one": AnyVersion(), - "start:two": AnyVersion(), - }, - "channel_values": {"one": "one", "two": "two", "value": 6}, - }, - metadata={ - "parents": {}, - "step": 1, - "source": "loop", - "writes": {"one": {"value": 2}, "two": {"value": 3}}, - "thread_id": "1", - }, - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": checkpoints[1].config["configurable"][ - "checkpoint_id" - ], - } - }, - pending_writes=[], - ) - # the previous one we assert that pending writes contains both - # - original error - # - successful writes from resuming after preventing error - assert checkpoints[1] == CheckpointTuple( - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - checkpoint={ - "v": 1, - "id": AnyStr(), - "ts": AnyStr(), - "pending_sends": [], - "versions_seen": { - "__input__": {}, - "__start__": { - "__start__": AnyVersion(), - }, - }, - "channel_versions": { - "value": AnyVersion(), - "__start__": AnyVersion(), - "start:one": AnyVersion(), - "start:two": AnyVersion(), - }, - "channel_values": { - "value": 1, - "start:one": "__start__", - "start:two": "__start__", - }, - }, - metadata={ - "parents": {}, - "step": 0, - "source": "loop", - "writes": None, - "thread_id": "1", - }, - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": checkpoints[2].config["configurable"][ - "checkpoint_id" - ], - } - }, - pending_writes=UnsortedSequence( - (AnyStr(), "one", "one"), - (AnyStr(), "value", 2), - (AnyStr(), "__error__", 'ConnectionError("I\'m not good")'), - (AnyStr(), "two", "two"), - (AnyStr(), "value", 3), - ), - ) - assert checkpoints[2] == CheckpointTuple( - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - checkpoint={ - "v": 1, - "id": AnyStr(), - "ts": AnyStr(), - "pending_sends": [], - "versions_seen": {"__input__": {}}, - "channel_versions": { - "__start__": AnyVersion(), - }, - "channel_values": {"__start__": {"value": 1}}, - }, - metadata={ - "parents": {}, - "step": -1, - "source": "input", - "writes": {"__start__": {"value": 1}}, - "thread_id": "1", - }, - parent_config=None, - pending_writes=UnsortedSequence( - (AnyStr(), "value", 1), - (AnyStr(), "start:one", "__start__"), - (AnyStr(), "start:two", "__start__"), - ), - ) - - -@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC) -async def test_run_from_checkpoint_id_retains_previous_writes( - request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture -) -> None: - class MyState(TypedDict): - myval: Annotated[int, operator.add] - otherval: bool - - class Anode: - def __init__(self): - self.switch = False - - async def __call__(self, state: MyState): - self.switch = not self.switch - return {"myval": 2 if self.switch else 1, "otherval": self.switch} - - builder = StateGraph(MyState) - thenode = Anode() # Fun. - builder.add_node("node_one", thenode) - builder.add_node("node_two", thenode) - builder.add_edge(START, "node_one") - - def _getedge(src: str): - swap = "node_one" if src == "node_two" else "node_two" - - def _edge(st: MyState) -> Literal["__end__", "node_one", "node_two"]: - if st["myval"] > 3: - return END - if st["otherval"]: - return swap - return src - - return _edge - - builder.add_conditional_edges("node_one", _getedge("node_one")) - builder.add_conditional_edges("node_two", _getedge("node_two")) - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer) - - thread_id = uuid.uuid4() - thread1 = {"configurable": {"thread_id": str(thread_id)}} - - result = await graph.ainvoke({"myval": 1}, thread1) - assert result["myval"] == 4 - history = [c async for c in graph.aget_state_history(thread1)] - - assert len(history) == 4 - assert history[-1].values == {"myval": 0} - assert history[0].values == {"myval": 4, "otherval": False} - - second_run_config = { - **thread1, - "configurable": { - **thread1["configurable"], - "checkpoint_id": history[1].config["configurable"]["checkpoint_id"], - }, - } - second_result = await graph.ainvoke(None, second_run_config) - assert second_result == {"myval": 5, "otherval": True} - - new_history = [ - c - async for c in graph.aget_state_history( - {"configurable": {"thread_id": str(thread_id), "checkpoint_ns": ""}} - ) - ] - - assert len(new_history) == len(history) + 1 - for original, new in zip(history, new_history[1:]): - assert original.values == new.values - assert original.next == new.next - assert original.metadata["step"] == new.metadata["step"] - - def _get_tasks(hist: list, start: int): - return [h.tasks for h in hist[start:]] - - assert _get_tasks(new_history, 1) == _get_tasks(history, 0) - - -async def test_cond_edge_after_send() -> None: - class Node: - def __init__(self, name: str): - self.name = name - setattr(self, "__name__", name) - - async def __call__(self, state): - return [self.name] - - async def send_for_fun(state): - return [Send("2", state), Send("2", state)] - - async def route_to_three(state) -> Literal["3"]: - return "3" - - builder = StateGraph(Annotated[list, operator.add]) - builder.add_node(Node("1")) - builder.add_node(Node("2")) - builder.add_node(Node("3")) - builder.add_edge(START, "1") - builder.add_conditional_edges("1", send_for_fun) - builder.add_conditional_edges("2", route_to_three) - graph = builder.compile() - - assert await graph.ainvoke(["0"]) == ["0", "1", "2", "2", "3"] - - -async def test_concurrent_emit_sends() -> None: - class Node: - def __init__(self, name: str): - self.name = name - setattr(self, "__name__", name) - - async def __call__(self, state): - return ( - [self.name] - if isinstance(state, list) - else ["|".join((self.name, str(state)))] - ) - - async def send_for_fun(state): - return [Send("2", 1), Send("2", 2), "3.1"] - - async def send_for_profit(state): - return [Send("2", 3), Send("2", 4)] - - async def route_to_three(state) -> Literal["3"]: - return "3" - - builder = StateGraph(Annotated[list, operator.add]) - builder.add_node(Node("1")) - builder.add_node(Node("1.1")) - builder.add_node(Node("2")) - builder.add_node(Node("3")) - builder.add_node(Node("3.1")) - builder.add_edge(START, "1") - builder.add_edge(START, "1.1") - builder.add_conditional_edges("1", send_for_fun) - builder.add_conditional_edges("1.1", send_for_profit) - builder.add_conditional_edges("2", route_to_three) - graph = builder.compile() - assert await graph.ainvoke(["0"]) == ( - [ - "0", - "1", - "1.1", - "3.1", - "2|1", - "2|2", - "2|3", - "2|4", - "3", - ] - ) - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_send_sequences(checkpointer_name: str) -> None: - class Node: - def __init__(self, name: str): - self.name = name - setattr(self, "__name__", name) - - async def __call__(self, state): - update = ( - [self.name] - if isinstance(state, list) # or isinstance(state, Control) - else ["|".join((self.name, str(state)))] - ) - if isinstance(state, Command): - return replace(state, update=update) - else: - return update - - async def send_for_fun(state): - return [ - Send("2", Command(goto=Send("2", 3))), - Send("2", Command(goto=Send("2", 4))), - "3.1", - ] - - async def route_to_three(state) -> Literal["3"]: - return "3" - - builder = StateGraph(Annotated[list, operator.add]) - builder.add_node(Node("1")) - builder.add_node(Node("2")) - builder.add_node(Node("3")) - builder.add_node(Node("3.1")) - builder.add_edge(START, "1") - builder.add_conditional_edges("1", send_for_fun) - builder.add_conditional_edges("2", route_to_three) - graph = builder.compile() - assert await graph.ainvoke(["0"]) == [ - "0", - "1", - "3.1", - "2|Command(goto=Send(node='2', arg=3))", - "2|Command(goto=Send(node='2', arg=4))", - "3", - "2|3", - "2|4", - "3", - ] - - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer, interrupt_before=["3.1"]) - thread1 = {"configurable": {"thread_id": "1"}} - assert await graph.ainvoke(["0"], thread1) == [ - "0", - "1", - ] - assert await graph.ainvoke(None, thread1) == [ - "0", - "1", - "3.1", - "2|Command(goto=Send(node='2', arg=3))", - "2|Command(goto=Send(node='2', arg=4))", - "3", - "2|3", - "2|4", - "3", - ] - - -@NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_imp_task(checkpointer_name: str) -> None: - async with awith_checkpointer(checkpointer_name) as checkpointer: - mapper_calls = 0 - - @task() - async def mapper(input: int) -> str: - nonlocal mapper_calls - mapper_calls += 1 - await asyncio.sleep(0.1 * input) - return str(input) * 2 - - @entrypoint(checkpointer=checkpointer) - async def graph(input: list[int]) -> list[str]: - futures = [mapper(i) for i in input] - mapped = await asyncio.gather(*futures) - answer = interrupt("question") - return [m + answer for m in mapped] - - tracer = FakeTracer() - thread1 = {"configurable": {"thread_id": "1"}, "callbacks": [tracer]} - assert [c async for c in graph.astream([0, 1], thread1)] == [ - {"mapper": "00"}, - {"mapper": "11"}, - { - "__interrupt__": ( - Interrupt( - value="question", - resumable=True, - ns=[AnyStr("graph:")], - when="during", - ), - ) - }, - ] - assert mapper_calls == 2 - assert len(tracer.runs) == 1 - assert len(tracer.runs[0].child_runs) == 1 - entrypoint_run = tracer.runs[0].child_runs[0] - assert entrypoint_run.name == "graph" - mapper_runs = [r for r in entrypoint_run.child_runs if r.name == "mapper"] - assert len(mapper_runs) == 2 - assert any(r.inputs == {"input": 0} for r in mapper_runs) - assert any(r.inputs == {"input": 1} for r in mapper_runs) - - assert await graph.ainvoke(Command(resume="answer"), thread1) == [ - "00answer", - "11answer", - ] - assert mapper_calls == 2 - - -@NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_imp_nested(checkpointer_name: str) -> None: - async def mynode(input: list[str]) -> list[str]: - return [it + "a" for it in input] - - builder = StateGraph(list[str]) - builder.add_node(mynode) - builder.add_edge(START, "mynode") - add_a = builder.compile() - - @task - def submapper(input: int) -> str: - return str(input) - - @task - async def mapper(input: int) -> str: - await asyncio.sleep(input / 100) - return await submapper(input) * 2 - - async with awith_checkpointer(checkpointer_name) as checkpointer: - - @entrypoint(checkpointer=checkpointer) - async def graph(input: list[int]) -> list[str]: - futures = [mapper(i) for i in input] - mapped = await asyncio.gather(*futures) - answer = interrupt("question") - final = [m + answer for m in mapped] - return await add_a.ainvoke(final) - - assert graph.get_input_jsonschema() == { - "type": "array", - "items": {"type": "integer"}, - "title": "LangGraphInput", - } - assert graph.get_output_jsonschema() == { - "type": "array", - "items": {"type": "string"}, - "title": "LangGraphOutput", - } - - thread1 = {"configurable": {"thread_id": "1"}} - assert [c async for c in graph.astream([0, 1], thread1)] == [ - {"submapper": "0"}, - {"mapper": "00"}, - {"submapper": "1"}, - {"mapper": "11"}, - { - "__interrupt__": ( - Interrupt( - value="question", - resumable=True, - ns=[AnyStr("graph:")], - when="during", - ), - ) - }, - ] - - assert await graph.ainvoke(Command(resume="answer"), thread1) == [ - "00answera", - "11answera", - ] - - -@NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_imp_task_cancel(checkpointer_name: str) -> None: - async with awith_checkpointer(checkpointer_name) as checkpointer: - mapper_calls = 0 - mapper_cancels = 0 - - @task() - async def mapper(input: int) -> str: - nonlocal mapper_calls, mapper_cancels - mapper_calls += 1 - try: - await asyncio.sleep(1) - except asyncio.CancelledError: - mapper_cancels += 1 - raise - return str(input) * 2 - - @entrypoint(checkpointer=checkpointer) - async def graph(input: list[int]) -> list[str]: - futures = [mapper(i) for i in input] - await asyncio.sleep(0.1) - futures.pop().cancel() # cancel one - mapped = await asyncio.gather(*futures) - answer = interrupt("question") - return [m + answer for m in mapped] - - thread1 = {"configurable": {"thread_id": "1"}} - assert [c async for c in graph.astream([0, 1], thread1)] == [ - {"mapper": "00"}, - { - "__interrupt__": ( - Interrupt( - value="question", - resumable=True, - ns=[AnyStr("graph:")], - when="during", - ), - ) - }, - ] - assert mapper_calls == 2 - assert mapper_cancels == 1 - - assert await graph.ainvoke(Command(resume="answer"), thread1) == [ - "00answer", - ] - assert mapper_calls == 3 - assert mapper_cancels == 2 - - -@NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_imp_sync_from_async(checkpointer_name: str) -> None: - async with awith_checkpointer(checkpointer_name) as checkpointer: - - @task() - def foo(state: dict) -> dict: - return {"a": state["a"] + "foo", "b": "bar"} - - @task - def bar(a: str, b: str, c: Optional[str] = None) -> dict: - return {"a": a + b, "c": (c or "") + "bark"} - - @task() - def baz(state: dict) -> dict: - return {"a": state["a"] + "baz", "c": "something else"} - - @entrypoint(checkpointer=checkpointer) - def graph(state: dict) -> dict: - foo_result = foo(state).result() - fut_bar = bar(foo_result["a"], foo_result["b"]) - fut_baz = baz(fut_bar.result()) - return fut_baz.result() - - thread1 = {"configurable": {"thread_id": "1"}} - assert [c async for c in graph.astream({"a": "0"}, thread1)] == [ - {"foo": {"a": "0foo", "b": "bar"}}, - {"bar": {"a": "0foobar", "c": "bark"}}, - {"baz": {"a": "0foobarbaz", "c": "something else"}}, - {"graph": {"a": "0foobarbaz", "c": "something else"}}, - ] - - -@NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_imp_stream_order(checkpointer_name: str) -> None: - async with awith_checkpointer(checkpointer_name) as checkpointer: - - @task() - async def foo(state: dict) -> dict: - return {"a": state["a"] + "foo", "b": "bar"} - - @task - async def bar(a: str, b: str, c: Optional[str] = None) -> dict: - return {"a": a + b, "c": (c or "") + "bark"} - - @task() - async def baz(state: dict) -> dict: - return {"a": state["a"] + "baz", "c": "something else"} - - @entrypoint(checkpointer=checkpointer) - async def graph(state: dict) -> dict: - foo_res = await foo(state) - - fut_bar = bar(foo_res["a"], foo_res["b"]) - fut_baz = baz(await fut_bar) - return await fut_baz - - thread1 = {"configurable": {"thread_id": "1"}} - assert [c async for c in graph.astream({"a": "0"}, thread1)] == [ - {"foo": {"a": "0foo", "b": "bar"}}, - {"bar": {"a": "0foobar", "c": "bark"}}, - {"baz": {"a": "0foobarbaz", "c": "something else"}}, - {"graph": {"a": "0foobarbaz", "c": "something else"}}, - ] - - -@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC) -async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: - class InterruptOnce: - ticks: int = 0 - - def __call__(self, state): - self.ticks += 1 - if self.ticks == 1: - raise NodeInterrupt("Bahh") - return ["|".join(("flaky", str(state)))] - - class Node: - def __init__(self, name: str): - self.name = name - self.ticks = 0 - setattr(self, "__name__", name) - - def __call__(self, state): - self.ticks += 1 - update = ( - [self.name] - if isinstance(state, list) - else ["|".join((self.name, str(state)))] - ) - if isinstance(state, Command): - return replace(state, update=update) - else: - return update - - def send_for_fun(state): - return [ - Send("2", Command(goto=Send("2", 3))), - Send("2", Command(goto=Send("flaky", 4))), - "3.1", - ] - - def route_to_three(state) -> Literal["3"]: - return "3" - - builder = StateGraph(Annotated[list, operator.add]) - builder.add_node(Node("1")) - builder.add_node(Node("2")) - builder.add_node(Node("3")) - builder.add_node(Node("3.1")) - builder.add_node("flaky", InterruptOnce()) - builder.add_edge(START, "1") - builder.add_conditional_edges("1", send_for_fun) - builder.add_conditional_edges("2", route_to_three) - - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer) - thread1 = {"configurable": {"thread_id": "1"}} - assert await graph.ainvoke(["0"], thread1, debug=1) == [ - "0", - "1", - "3.1", - "2|Command(goto=Send(node='2', arg=3))", - "2|Command(goto=Send(node='flaky', arg=4))", - "3", - "2|3", - ] - assert builder.nodes["2"].runnable.func.ticks == 3 - assert builder.nodes["flaky"].runnable.func.ticks == 1 - # resume execution - assert await graph.ainvoke(None, thread1, debug=1) == [ - "0", - "1", - "3.1", - "2|Command(goto=Send(node='2', arg=3))", - "2|Command(goto=Send(node='flaky', arg=4))", - "3", - "2|3", - "flaky|4", - "3", - ] - # node "2" doesn't get called again, as we recover writes saved before - assert builder.nodes["2"].runnable.func.ticks == 3 - # node "flaky" gets called again, as it was interrupted - assert builder.nodes["flaky"].runnable.func.ticks == 2 - # check history - history = [c async for c in graph.aget_state_history(thread1)] - assert history == [ - StateSnapshot( - values=[ - "0", - "1", - "3.1", - "2|Command(goto=Send(node='2', arg=3))", - "2|Command(goto=Send(node='flaky', arg=4))", - "3", - "2|3", - "flaky|4", - "3", - ], - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"3": ["3"]}, - "thread_id": "1", - "step": 4, - "parents": {}, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - tasks=(), - ), - StateSnapshot( - values=[ - "0", - "1", - "3.1", - "2|Command(goto=Send(node='2', arg=3))", - "2|Command(goto=Send(node='flaky', arg=4))", - "3", - "2|3", - "flaky|4", - ], - next=("3",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"2": ["2|3"], "3": ["3"], "flaky": ["flaky|4"]}, - "thread_id": "1", - "step": 3, - "parents": {}, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - tasks=( - PregelTask( - id=AnyStr(), - name="3", - path=("__pregel_pull", "3"), - error=None, - interrupts=(), - state=None, - result=["3"], - ), - ), - ), - StateSnapshot( - values=[ - "0", - "1", - "3.1", - "2|Command(goto=Send(node='2', arg=3))", - "2|Command(goto=Send(node='flaky', arg=4))", - ], - next=("2", "flaky", "3"), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": { - "2": [ - ["2|Command(goto=Send(node='2', arg=3))"], - ["2|Command(goto=Send(node='flaky', arg=4))"], - ], - "3.1": ["3.1"], - }, - "thread_id": "1", - "step": 2, - "parents": {}, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - tasks=( - PregelTask( - id=AnyStr(), - name="2", - path=("__pregel_push", 0), - error=None, - interrupts=(), - state=None, - result=["2|3"], - ), - PregelTask( - id=AnyStr(), - name="flaky", - path=("__pregel_push", 1), - error=None, - interrupts=( - Interrupt( - value="Bahh", resumable=False, ns=None, when="during" - ), - ), - state=None, - result=["flaky|4"], - ), - PregelTask( - id=AnyStr(), - name="3", - path=("__pregel_pull", "3"), - error=None, - interrupts=(), - state=None, - result=["3"], - ), - ), - ), - StateSnapshot( - values=["0", "1"], - next=("2", "2", "3.1"), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"1": ["1"]}, - "thread_id": "1", - "step": 1, - "parents": {}, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - tasks=( - PregelTask( - id=AnyStr(), - name="2", - path=("__pregel_push", 0), - error=None, - interrupts=(), - state=None, - result=["2|Command(goto=Send(node='2', arg=3))"], - ), - PregelTask( - id=AnyStr(), - name="2", - path=("__pregel_push", 1), - error=None, - interrupts=(), - state=None, - result=["2|Command(goto=Send(node='flaky', arg=4))"], - ), - PregelTask( - id=AnyStr(), - name="3.1", - path=("__pregel_pull", "3.1"), - error=None, - interrupts=(), - state=None, - result=["3.1"], - ), - ), - ), - StateSnapshot( - values=["0"], - next=("1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": None, - "thread_id": "1", - "step": 0, - "parents": {}, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - tasks=( - PregelTask( - id=AnyStr(), - name="1", - path=("__pregel_pull", "1"), - error=None, - interrupts=(), - state=None, - result=["1"], - ), - ), - ), - StateSnapshot( - values=[], - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "input", - "writes": {"__start__": ["0"]}, - "thread_id": "1", - "step": -1, - "parents": {}, - }, - created_at=AnyStr(), - parent_config=None, - tasks=( - PregelTask( - id=AnyStr(), - name="__start__", - path=("__pregel_pull", "__start__"), - error=None, - interrupts=(), - state=None, - result=["0"], - ), - ), - ), - ] - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_send_react_interrupt(checkpointer_name: str) -> None: - from langchain_core.messages import AIMessage, HumanMessage, ToolCall, ToolMessage - - ai_message = AIMessage( - "", - id="ai1", - tool_calls=[ToolCall(name="foo", args={"hi": [1, 2, 3]}, id=AnyStr())], - ) - - async def agent(state): - return {"messages": ai_message} - - def route(state): - if isinstance(state["messages"][-1], AIMessage): - return [ - Send(call["name"], call) for call in state["messages"][-1].tool_calls - ] - - foo_called = 0 - - async def foo(call: ToolCall): - nonlocal foo_called - foo_called += 1 - return {"messages": ToolMessage(str(call["args"]), tool_call_id=call["id"])} - - builder = StateGraph(MessagesState) - builder.add_node(agent) - builder.add_node(foo) - builder.add_edge(START, "agent") - builder.add_conditional_edges("agent", route) - graph = builder.compile() - - assert await graph.ainvoke({"messages": [HumanMessage("hello")]}) == { - "messages": [ - _AnyIdHumanMessage(content="hello"), - _AnyIdAIMessage( - content="", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [1, 2, 3]}, - "id": "", - "type": "tool_call", - } - ], - ), - _AnyIdToolMessage( - content="{'hi': [1, 2, 3]}", - tool_call_id=AnyStr(), - ), - ] - } - assert foo_called == 1 - - async with awith_checkpointer(checkpointer_name) as checkpointer: - # simple interrupt-resume flow - foo_called = 0 - graph = builder.compile(checkpointer=checkpointer, interrupt_before=["foo"]) - thread1 = {"configurable": {"thread_id": "1"}} - assert await graph.ainvoke({"messages": [HumanMessage("hello")]}, thread1) == { - "messages": [ - _AnyIdHumanMessage(content="hello"), - _AnyIdAIMessage( - content="", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [1, 2, 3]}, - "id": "", - "type": "tool_call", - } - ], - ), - ] - } - assert foo_called == 0 - assert await graph.ainvoke(None, thread1) == { - "messages": [ - _AnyIdHumanMessage(content="hello"), - _AnyIdAIMessage( - content="", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [1, 2, 3]}, - "id": "", - "type": "tool_call", - } - ], - ), - _AnyIdToolMessage( - content="{'hi': [1, 2, 3]}", - tool_call_id=AnyStr(), - ), - ] - } - assert foo_called == 1 - - # interrupt-update-resume flow - foo_called = 0 - graph = builder.compile(checkpointer=checkpointer, interrupt_before=["foo"]) - thread1 = {"configurable": {"thread_id": "2"}} - assert await graph.ainvoke({"messages": [HumanMessage("hello")]}, thread1) == { - "messages": [ - _AnyIdHumanMessage(content="hello"), - _AnyIdAIMessage( - content="", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [1, 2, 3]}, - "id": "", - "type": "tool_call", - } - ], - ), - ] - } - assert foo_called == 0 - - # get state should show the pending task - state = await graph.aget_state(thread1) - assert state == StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="hello"), - _AnyIdAIMessage( - content="", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [1, 2, 3]}, - "id": "", - "type": "tool_call", - } - ], - ), - ] - }, - next=("foo",), - config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "step": 1, - "source": "loop", - "writes": { - "agent": { - "messages": AIMessage( - content="", - id="ai1", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [1, 2, 3]}, - "id": "", - "type": "tool_call", - } - ], - ) - } - }, - "parents": {}, - "thread_id": "2", - }, - created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), - tasks=( - PregelTask( - id=AnyStr(), - name="foo", - path=("__pregel_push", 0), - error=None, - interrupts=(), - state=None, - result=None, - ), - ), - ) - - # remove the tool call, clearing the pending task - await graph.aupdate_state( - thread1, {"messages": AIMessage("Bye now", id=ai_message.id, tool_calls=[])} - ) - - # tool call no longer in pending tasks - assert await graph.aget_state(thread1) == StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="hello"), - _AnyIdAIMessage( - content="Bye now", - tool_calls=[], - ), - ] - }, - next=(), - config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "step": 2, - "source": "update", - "writes": { - "agent": { - "messages": _AnyIdAIMessage( - content="Bye now", - tool_calls=[], - ) - } - }, - "parents": {}, - "thread_id": "2", - }, - created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), - tasks=(), - ) - - # tool call not executed - assert await graph.ainvoke(None, thread1) == { - "messages": [ - _AnyIdHumanMessage(content="hello"), - _AnyIdAIMessage(content="Bye now"), - ] - } - assert foo_called == 0 - - # interrupt-update-resume flow, creating new Send in update call - foo_called = 0 - graph = builder.compile(checkpointer=checkpointer, interrupt_before=["foo"]) - thread1 = {"configurable": {"thread_id": "3"}} - assert await graph.ainvoke({"messages": [HumanMessage("hello")]}, thread1) == { - "messages": [ - _AnyIdHumanMessage(content="hello"), - _AnyIdAIMessage( - content="", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [1, 2, 3]}, - "id": "", - "type": "tool_call", - } - ], - ), - ] - } - assert foo_called == 0 - - # get state should show the pending task - state = await graph.aget_state(thread1) - assert state == StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="hello"), - _AnyIdAIMessage( - content="", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [1, 2, 3]}, - "id": "", - "type": "tool_call", - } - ], - ), - ] - }, - next=("foo",), - config={ - "configurable": { - "thread_id": "3", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "step": 1, - "source": "loop", - "writes": { - "agent": { - "messages": AIMessage( - content="", - id="ai1", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [1, 2, 3]}, - "id": "", - "type": "tool_call", - } - ], - ) - } - }, - "parents": {}, - "thread_id": "3", - }, - created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { - "configurable": { - "thread_id": "3", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), - tasks=( - PregelTask( - id=AnyStr(), - name="foo", - path=("__pregel_push", 0), - error=None, - interrupts=(), - state=None, - result=None, - ), - ), - ) - - # replace the tool call, should clear previous send, create new one - await graph.aupdate_state( - thread1, - { - "messages": AIMessage( - "", - id=ai_message.id, - tool_calls=[ - { - "name": "foo", - "args": {"hi": [4, 5, 6]}, - "id": "tool1", - "type": "tool_call", - } - ], - ) - }, - ) - - # prev tool call no longer in pending tasks, new tool call is - assert await graph.aget_state(thread1) == StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="hello"), - _AnyIdAIMessage( - content="", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [4, 5, 6]}, - "id": "tool1", - "type": "tool_call", - } - ], - ), - ] - }, - next=("foo",), - config={ - "configurable": { - "thread_id": "3", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "step": 2, - "source": "update", - "writes": { - "agent": { - "messages": _AnyIdAIMessage( - content="", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [4, 5, 6]}, - "id": "tool1", - "type": "tool_call", - } - ], - ) - } - }, - "parents": {}, - "thread_id": "3", - }, - created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { - "configurable": { - "thread_id": "3", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), - tasks=( - PregelTask( - id=AnyStr(), - name="foo", - path=("__pregel_push", 0), - error=None, - interrupts=(), - state=None, - result=None, - ), - ), - ) - - # prev tool call not executed, new tool call is - assert await graph.ainvoke(None, thread1) == { - "messages": [ - _AnyIdHumanMessage(content="hello"), - AIMessage( - "", - id="ai1", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [4, 5, 6]}, - "id": "tool1", - "type": "tool_call", - } - ], - ), - _AnyIdToolMessage(content="{'hi': [4, 5, 6]}", tool_call_id="tool1"), - ] - } - assert foo_called == 1 - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_send_react_interrupt_control( - checkpointer_name: str, snapshot: SnapshotAssertion -) -> None: - from langchain_core.messages import AIMessage, HumanMessage, ToolCall, ToolMessage - - ai_message = AIMessage( - "", - id="ai1", - tool_calls=[ToolCall(name="foo", args={"hi": [1, 2, 3]}, id=AnyStr())], - ) - - async def agent(state) -> Command[Literal["foo"]]: - return Command( - update={"messages": ai_message}, - goto=[Send(call["name"], call) for call in ai_message.tool_calls], - ) - - foo_called = 0 - - async def foo(call: ToolCall): - nonlocal foo_called - foo_called += 1 - return {"messages": ToolMessage(str(call["args"]), tool_call_id=call["id"])} - - builder = StateGraph(MessagesState) - builder.add_node(agent) - builder.add_node(foo) - builder.add_edge(START, "agent") - graph = builder.compile() - assert graph.get_graph().draw_mermaid() == snapshot - - assert await graph.ainvoke({"messages": [HumanMessage("hello")]}) == { - "messages": [ - _AnyIdHumanMessage(content="hello"), - _AnyIdAIMessage( - content="", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [1, 2, 3]}, - "id": "", - "type": "tool_call", - } - ], - ), - _AnyIdToolMessage( - content="{'hi': [1, 2, 3]}", - tool_call_id=AnyStr(), - ), - ] - } - assert foo_called == 1 - - async with awith_checkpointer(checkpointer_name) as checkpointer: - # simple interrupt-resume flow - foo_called = 0 - graph = builder.compile(checkpointer=checkpointer, interrupt_before=["foo"]) - thread1 = {"configurable": {"thread_id": "1"}} - assert await graph.ainvoke({"messages": [HumanMessage("hello")]}, thread1) == { - "messages": [ - _AnyIdHumanMessage(content="hello"), - _AnyIdAIMessage( - content="", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [1, 2, 3]}, - "id": "", - "type": "tool_call", - } - ], - ), - ] - } - assert foo_called == 0 - assert await graph.ainvoke(None, thread1) == { - "messages": [ - _AnyIdHumanMessage(content="hello"), - _AnyIdAIMessage( - content="", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [1, 2, 3]}, - "id": "", - "type": "tool_call", - } - ], - ), - _AnyIdToolMessage( - content="{'hi': [1, 2, 3]}", - tool_call_id=AnyStr(), - ), - ] - } - assert foo_called == 1 - - # interrupt-update-resume flow - foo_called = 0 - graph = builder.compile(checkpointer=checkpointer, interrupt_before=["foo"]) - thread1 = {"configurable": {"thread_id": "2"}} - assert await graph.ainvoke({"messages": [HumanMessage("hello")]}, thread1) == { - "messages": [ - _AnyIdHumanMessage(content="hello"), - _AnyIdAIMessage( - content="", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [1, 2, 3]}, - "id": "", - "type": "tool_call", - } - ], - ), - ] - } - assert foo_called == 0 - - # get state should show the pending task - state = await graph.aget_state(thread1) - assert state == StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="hello"), - _AnyIdAIMessage( - content="", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [1, 2, 3]}, - "id": "", - "type": "tool_call", - } - ], - ), - ] - }, - next=("foo",), - config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "step": 1, - "source": "loop", - "writes": { - "agent": { - "messages": AIMessage( - content="", - id="ai1", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [1, 2, 3]}, - "id": "", - "type": "tool_call", - } - ], - ) - } - }, - "parents": {}, - "thread_id": "2", - }, - created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), - tasks=( - PregelTask( - id=AnyStr(), - name="foo", - path=("__pregel_push", 0), - error=None, - interrupts=(), - state=None, - result=None, - ), - ), - ) - - # remove the tool call, clearing the pending task - await graph.aupdate_state( - thread1, {"messages": AIMessage("Bye now", id=ai_message.id, tool_calls=[])} - ) - - # tool call no longer in pending tasks - assert await graph.aget_state(thread1) == StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="hello"), - _AnyIdAIMessage( - content="Bye now", - tool_calls=[], - ), - ] - }, - next=(), - config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "step": 2, - "source": "update", - "writes": { - "agent": { - "messages": _AnyIdAIMessage( - content="Bye now", - tool_calls=[], - ) - } - }, - "parents": {}, - "thread_id": "2", - }, - created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), - tasks=(), - ) - - # tool call not executed - assert await graph.ainvoke(None, thread1) == { - "messages": [ - _AnyIdHumanMessage(content="hello"), - _AnyIdAIMessage(content="Bye now"), - ] - } - assert foo_called == 0 - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_max_concurrency(checkpointer_name: str) -> None: - class Node: - def __init__(self, name: str): - self.name = name - setattr(self, "__name__", name) - self.currently = 0 - self.max_currently = 0 - - async def __call__(self, state): - self.currently += 1 - if self.currently > self.max_currently: - self.max_currently = self.currently - await asyncio.sleep(random.random() / 10) - self.currently -= 1 - return [state] - - def one(state): - return ["1"] - - def three(state): - return ["3"] - - async def send_to_many(state): - return [Send("2", idx) for idx in range(100)] - - async def route_to_three(state) -> Literal["3"]: - return "3" - - node2 = Node("2") - builder = StateGraph(Annotated[list, operator.add]) - builder.add_node("1", one) - builder.add_node(node2) - builder.add_node("3", three) - builder.add_edge(START, "1") - builder.add_conditional_edges("1", send_to_many) - builder.add_conditional_edges("2", route_to_three) - graph = builder.compile() - - assert await graph.ainvoke(["0"]) == ["0", "1", *range(100), "3"] - assert node2.max_currently == 100 - assert node2.currently == 0 - node2.max_currently = 0 - - assert await graph.ainvoke(["0"], {"max_concurrency": 10}) == [ - "0", - "1", - *range(100), - "3", - ] - assert node2.max_currently == 10 - assert node2.currently == 0 - - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer, interrupt_before=["2"]) - thread1 = {"max_concurrency": 10, "configurable": {"thread_id": "1"}} - - assert await graph.ainvoke(["0"], thread1, debug=True) == ["0", "1"] - state = await graph.aget_state(thread1) - assert state.values == ["0", "1"] - assert await graph.ainvoke(None, thread1) == ["0", "1", *range(100), "3"] - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_max_concurrency_control(checkpointer_name: str) -> None: - async def node1(state) -> Command[Literal["2"]]: - return Command(update=["1"], goto=[Send("2", idx) for idx in range(100)]) - - node2_currently = 0 - node2_max_currently = 0 - - async def node2(state) -> Command[Literal["3"]]: - nonlocal node2_currently, node2_max_currently - node2_currently += 1 - if node2_currently > node2_max_currently: - node2_max_currently = node2_currently - await asyncio.sleep(0.1) - node2_currently -= 1 - - return Command(update=[state], goto="3") - - async def node3(state) -> Literal["3"]: - return ["3"] - - builder = StateGraph(Annotated[list, operator.add]) - builder.add_node("1", node1) - builder.add_node("2", node2) - builder.add_node("3", node3) - builder.add_edge(START, "1") - graph = builder.compile() - - assert ( - graph.get_graph().draw_mermaid() - == """%%{init: {'flowchart': {'curve': 'linear'}}}%% -graph TD; - __start__([

__start__

]):::first - 1(1) - 2(2) - 3([3]):::last - __start__ --> 1; - 1 -.-> 2; - 2 -.-> 3; - classDef default fill:#f2f0ff,line-height:1.2 - classDef first fill-opacity:0 - classDef last fill:#bfb6fc -""" - ) - - assert await graph.ainvoke(["0"], debug=True) == ["0", "1", *range(100), "3"] - assert node2_max_currently == 100 - assert node2_currently == 0 - node2_max_currently = 0 - - assert await graph.ainvoke(["0"], {"max_concurrency": 10}) == [ - "0", - "1", - *range(100), - "3", - ] - assert node2_max_currently == 10 - assert node2_currently == 0 - - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer, interrupt_before=["2"]) - thread1 = {"max_concurrency": 10, "configurable": {"thread_id": "1"}} - - assert await graph.ainvoke(["0"], thread1) == ["0", "1"] - assert await graph.ainvoke(None, thread1) == ["0", "1", *range(100), "3"] - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_invoke_checkpoint_three( - mocker: MockerFixture, checkpointer_name: str -) -> None: - add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) - - def raise_if_above_10(input: int) -> int: - if input > 10: - raise ValueError("Input is too large") - return input - - one = ( - Channel.subscribe_to(["input"]).join(["total"]) - | add_one - | Channel.write_to("output", "total") - | raise_if_above_10 - ) - - async with awith_checkpointer(checkpointer_name) as checkpointer: - app = Pregel( - nodes={"one": one}, - channels={ - "total": BinaryOperatorAggregate(int, operator.add), - "input": LastValue(int), - "output": LastValue(int), - }, - input_channels="input", - output_channels="output", - checkpointer=checkpointer, - debug=True, - ) - - thread_1 = {"configurable": {"thread_id": "1"}} - # total starts out as 0, so output is 0+2=2 - assert await app.ainvoke(2, thread_1) == 2 - state = await app.aget_state(thread_1) - assert state is not None - assert state.values.get("total") == 2 - assert ( - state.config["configurable"]["checkpoint_id"] - == (await checkpointer.aget(thread_1))["id"] - ) - # total is now 2, so output is 2+3=5 - assert await app.ainvoke(3, thread_1) == 5 - state = await app.aget_state(thread_1) - assert state is not None - assert state.values.get("total") == 7 - assert ( - state.config["configurable"]["checkpoint_id"] - == (await checkpointer.aget(thread_1))["id"] - ) - # total is now 2+5=7, so output would be 7+4=11, but raises ValueError - with pytest.raises(ValueError): - await app.ainvoke(4, thread_1) - # checkpoint is not updated - state = await app.aget_state(thread_1) - assert state is not None - assert state.values.get("total") == 7 - assert state.next == ("one",) - """we checkpoint inputs and it failed on "one", so the next node is one""" - # we can recover from error by sending new inputs - assert await app.ainvoke(2, thread_1) == 9 - state = await app.aget_state(thread_1) - assert state is not None - assert state.values.get("total") == 16, "total is now 7+9=16" - assert state.next == () - - thread_2 = {"configurable": {"thread_id": "2"}} - # on a new thread, total starts out as 0, so output is 0+5=5 - assert await app.ainvoke(5, thread_2) == 5 - state = await app.aget_state({"configurable": {"thread_id": "1"}}) - assert state is not None - assert state.values.get("total") == 16 - assert state.next == () - state = await app.aget_state(thread_2) - assert state is not None - assert state.values.get("total") == 5 - assert state.next == () - - if "shallow" in checkpointer_name: - return - - assert len([c async for c in app.aget_state_history(thread_1, limit=1)]) == 1 - # list all checkpoints for thread 1 - thread_1_history = [c async for c in app.aget_state_history(thread_1)] - # there are 7 checkpoints - assert len(thread_1_history) == 7 - assert Counter(c.metadata["source"] for c in thread_1_history) == { - "input": 4, - "loop": 3, - } - # sorted descending - assert ( - thread_1_history[0].config["configurable"]["checkpoint_id"] - > thread_1_history[1].config["configurable"]["checkpoint_id"] - ) - # cursor pagination - cursored = [ - c - async for c in app.aget_state_history( - thread_1, limit=1, before=thread_1_history[0].config - ) - ] - assert len(cursored) == 1 - assert cursored[0].config == thread_1_history[1].config - # the last checkpoint - assert thread_1_history[0].values["total"] == 16 - # the first "loop" checkpoint - assert thread_1_history[-2].values["total"] == 2 - # can get each checkpoint using aget with config - assert (await checkpointer.aget(thread_1_history[0].config))[ - "id" - ] == thread_1_history[0].config["configurable"]["checkpoint_id"] - assert (await checkpointer.aget(thread_1_history[1].config))[ - "id" - ] == thread_1_history[1].config["configurable"]["checkpoint_id"] - - thread_1_next_config = await app.aupdate_state(thread_1_history[1].config, 10) - # update creates a new checkpoint - assert ( - thread_1_next_config["configurable"]["checkpoint_id"] - > thread_1_history[0].config["configurable"]["checkpoint_id"] - ) - # 1 more checkpoint in history - assert len([c async for c in app.aget_state_history(thread_1)]) == 8 - assert Counter( - [c.metadata["source"] async for c in app.aget_state_history(thread_1)] - ) == { - "update": 1, - "input": 4, - "loop": 3, - } - # the latest checkpoint is the updated one - assert await app.aget_state(thread_1) == await app.aget_state( - thread_1_next_config - ) - - -async def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None: - add_one = mocker.Mock(side_effect=lambda x: x + 1) - add_10_each = mocker.Mock(side_effect=lambda x: sorted(y + 10 for y in x)) - - one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") - chain_three = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") - chain_four = ( - Channel.subscribe_to("inbox") | add_10_each | Channel.write_to("output") - ) - - app = Pregel( - nodes={ - "one": one, - "chain_three": chain_three, - "chain_four": chain_four, - }, - channels={ - "inbox": Topic(int), - "output": LastValue(int), - "input": LastValue(int), - }, - input_channels="input", - output_channels="output", - ) - - # Then invoke app - # We get a single array result as chain_four waits for all publishers to finish - # before operating on all elements published to topic_two as an array - for _ in range(100): - assert await app.ainvoke(2) == [13, 13] - - assert await asyncio.gather(*(app.ainvoke(2) for _ in range(100))) == [ - [13, 13] for _ in range(100) - ] - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_invoke_join_then_call_other_pregel( - mocker: MockerFixture, checkpointer_name: str -) -> None: - add_one = mocker.Mock(side_effect=lambda x: x + 1) - add_10_each = mocker.Mock(side_effect=lambda x: [y + 10 for y in x]) - - inner_app = Pregel( - nodes={ - "one": Channel.subscribe_to("input") | add_one | Channel.write_to("output") - }, - channels={ - "output": LastValue(int), - "input": LastValue(int), - }, - input_channels="input", - output_channels="output", - ) - - one = ( - Channel.subscribe_to("input") - | add_10_each - | Channel.write_to("inbox_one").map() - ) - two = ( - Channel.subscribe_to("inbox_one") - | inner_app.map() - | sorted - | Channel.write_to("outbox_one") - ) - chain_three = Channel.subscribe_to("outbox_one") | sum | Channel.write_to("output") - - app = Pregel( - nodes={ - "one": one, - "two": two, - "chain_three": chain_three, - }, - channels={ - "inbox_one": Topic(int), - "outbox_one": LastValue(int), - "output": LastValue(int), - "input": LastValue(int), - }, - input_channels="input", - output_channels="output", - ) - - # Then invoke pubsub - for _ in range(10): - assert await app.ainvoke([2, 3]) == 27 - - assert await asyncio.gather(*(app.ainvoke([2, 3]) for _ in range(10))) == [ - 27 for _ in range(10) - ] - - async with awith_checkpointer(checkpointer_name) as checkpointer: - # add checkpointer - app.checkpointer = checkpointer - # subgraph is called twice, and that works - assert await app.ainvoke([2, 3], {"configurable": {"thread_id": "1"}}) == 27 - - # set inner graph checkpointer NeverCheckpoint - inner_app.checkpointer = False - # subgraph still called twice, but checkpointing for inner graph is disabled - assert await app.ainvoke([2, 3], {"configurable": {"thread_id": "1"}}) == 27 - - -async def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None: - add_one = mocker.Mock(side_effect=lambda x: x + 1) - - one = ( - Channel.subscribe_to("input") | add_one | Channel.write_to("output", "between") - ) - two = Channel.subscribe_to("between") | add_one | Channel.write_to("output") - - app = Pregel( - nodes={"one": one, "two": two}, - channels={ - "input": LastValue(int), - "between": LastValue(int), - "output": LastValue(int), - }, - stream_channels=["output", "between"], - input_channels="input", - output_channels="output", - ) - - # Then invoke pubsub - assert [c async for c in app.astream(2)] == [ - {"between": 3, "output": 3}, - {"between": 3, "output": 4}, - ] - - -async def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None: - add_one = mocker.Mock(side_effect=lambda x: x + 1) - one = Channel.subscribe_to("input") | add_one | Channel.write_to("between") - two = Channel.subscribe_to("between") | add_one - - app = Pregel( - nodes={"one": one, "two": two}, - channels={ - "input": LastValue(int), - "between": LastValue(int), - "output": LastValue(int), - }, - input_channels="input", - output_channels="output", - ) - - # It finishes executing (once no more messages being published) - # but returns nothing, as nothing was published to "output" topic - assert await app.ainvoke(2) is None - - -async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: - setup_sync = mocker.Mock() - cleanup_sync = mocker.Mock() - setup_async = mocker.Mock() - cleanup_async = mocker.Mock() - - @contextmanager - def an_int() -> Generator[int, None, None]: - setup_sync() - try: - yield 5 - finally: - cleanup_sync() - - @asynccontextmanager - async def an_int_async() -> AsyncGenerator[int, None]: - setup_async() - try: - yield 5 - finally: - cleanup_async() - - add_one = mocker.Mock(side_effect=lambda x: x + 1) - one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") - two = ( - Channel.subscribe_to("inbox") - | RunnableLambda(add_one).abatch - | Channel.write_to("output").abatch - ) - - app = Pregel( - nodes={"one": one, "two": two}, - channels={ - "input": LastValue(int), - "output": LastValue(int), - "inbox": Topic(int), - "ctx": Context(an_int, an_int_async), - }, - input_channels="input", - output_channels=["inbox", "output"], - stream_channels=["inbox", "output"], - ) - - async def aenumerate(aiter: AsyncIterator[Any]) -> AsyncIterator[tuple[int, Any]]: - i = 0 - async for chunk in aiter: - yield i, chunk - i += 1 - - assert setup_sync.call_count == 0 - assert cleanup_sync.call_count == 0 - assert setup_async.call_count == 0 - assert cleanup_async.call_count == 0 - async for i, chunk in aenumerate(app.astream(2)): - assert setup_sync.call_count == 0, "Sync context manager should not be used" - assert cleanup_sync.call_count == 0, "Sync context manager should not be used" - assert setup_async.call_count == 1, "Expected setup to be called once" - if i == 0: - assert chunk == {"inbox": [3]} - elif i == 1: - assert chunk == {"output": 4} - else: - assert False, "Expected only two chunks" - assert setup_sync.call_count == 0 - assert cleanup_sync.call_count == 0 - assert setup_async.call_count == 1, "Expected setup to be called once" - assert cleanup_async.call_count == 1, "Expected cleanup to be called once" - - -async def test_conditional_entrypoint_graph() -> None: - async def left(data: str) -> str: - return data + "->left" - - async def right(data: str) -> str: - return data + "->right" - - def should_start(data: str) -> str: - # Logic to decide where to start - if len(data) > 10: - return "go-right" - else: - return "go-left" - - # Define a new graph - workflow = Graph() - - workflow.add_node("left", left) - workflow.add_node("right", right) - - workflow.set_conditional_entry_point( - should_start, {"go-left": "left", "go-right": "right"} - ) - - workflow.add_conditional_edges("left", lambda data: END) - workflow.add_edge("right", END) - - app = workflow.compile() - - assert await app.ainvoke("what is weather in sf") == "what is weather in sf->right" - - assert [c async for c in app.astream("what is weather in sf")] == [ - {"right": "what is weather in sf->right"}, - ] - - -async def test_conditional_entrypoint_graph_state() -> None: - class AgentState(TypedDict, total=False): - input: str - output: str - steps: Annotated[list[str], operator.add] - - async def left(data: AgentState) -> AgentState: - return {"output": data["input"] + "->left"} - - async def right(data: AgentState) -> AgentState: - return {"output": data["input"] + "->right"} - - def should_start(data: AgentState) -> str: - assert data["steps"] == [], "Expected input to be read from the state" - # Logic to decide where to start - if len(data["input"]) > 10: - return "go-right" - else: - return "go-left" - - # Define a new graph - workflow = StateGraph(AgentState) - - workflow.add_node("left", left) - workflow.add_node("right", right) - - workflow.set_conditional_entry_point( - should_start, {"go-left": "left", "go-right": "right"} - ) - - workflow.add_conditional_edges("left", lambda data: END) - workflow.add_edge("right", END) - - app = workflow.compile() - - assert await app.ainvoke({"input": "what is weather in sf"}) == { - "input": "what is weather in sf", - "output": "what is weather in sf->right", - "steps": [], - } - - assert [c async for c in app.astream({"input": "what is weather in sf"})] == [ - {"right": {"output": "what is weather in sf->right"}}, - ] - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_in_one_fan_out_state_graph_waiting_edge(checkpointer_name: str) -> None: - def sorted_add( - x: list[str], y: Union[list[str], list[tuple[str, str]]] - ) -> list[str]: - if isinstance(y[0], tuple): - for rem, _ in y: - x.remove(rem) - y = [t[1] for t in y] - return sorted(operator.add(x, y)) - - class State(TypedDict, total=False): - query: str - answer: str - docs: Annotated[list[str], sorted_add] - - async def rewrite_query(data: State) -> State: - return {"query": f"query: {data['query']}"} - - async def analyzer_one(data: State) -> State: - return {"query": f"analyzed: {data['query']}"} - - async def retriever_one(data: State) -> State: - return {"docs": ["doc1", "doc2"]} - - async def retriever_two(data: State) -> State: - await asyncio.sleep(0.1) - return {"docs": ["doc3", "doc4"]} - - async def qa(data: State) -> State: - return {"answer": ",".join(data["docs"])} - - workflow = StateGraph(State) - - workflow.add_node("rewrite_query", rewrite_query) - workflow.add_node("analyzer_one", analyzer_one) - workflow.add_node("retriever_one", retriever_one) - workflow.add_node("retriever_two", retriever_two) - workflow.add_node("qa", qa) - - workflow.set_entry_point("rewrite_query") - workflow.add_edge("rewrite_query", "analyzer_one") - workflow.add_edge("analyzer_one", "retriever_one") - workflow.add_edge("rewrite_query", "retriever_two") - workflow.add_edge(["retriever_one", "retriever_two"], "qa") - workflow.set_finish_point("qa") - - app = workflow.compile() - - assert await app.ainvoke({"query": "what is weather in sf"}) == { - "query": "analyzed: query: what is weather in sf", - "docs": ["doc1", "doc2", "doc3", "doc4"], - "answer": "doc1,doc2,doc3,doc4", - } - - assert [c async for c in app.astream({"query": "what is weather in sf"})] == [ - {"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"]}}, - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] - - async with awith_checkpointer(checkpointer_name) as checkpointer: - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_after=["retriever_one"], - ) - config = {"configurable": {"thread_id": "1"}} - - assert [ - c - async for c in app_w_interrupt.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__": ()}, - ] - - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_in_one_fan_out_state_graph_waiting_edge_via_branch( - snapshot: SnapshotAssertion, checkpointer_name: str -) -> None: - def sorted_add( - x: list[str], y: Union[list[str], list[tuple[str, str]]] - ) -> list[str]: - if isinstance(y[0], tuple): - for rem, _ in y: - x.remove(rem) - y = [t[1] for t in y] - return sorted(operator.add(x, y)) - - class State(TypedDict, total=False): - query: str - answer: str - docs: Annotated[list[str], sorted_add] - - async def rewrite_query(data: State) -> State: - return {"query": f"query: {data['query']}"} - - async def analyzer_one(data: State) -> State: - return {"query": f"analyzed: {data['query']}"} - - async def retriever_one(data: State) -> State: - return {"docs": ["doc1", "doc2"]} - - async def retriever_two(data: State) -> State: - await asyncio.sleep(0.1) - return {"docs": ["doc3", "doc4"]} - - async def qa(data: State) -> State: - return {"answer": ",".join(data["docs"])} - - workflow = StateGraph(State) - - workflow.add_node("rewrite_query", rewrite_query) - workflow.add_node("analyzer_one", analyzer_one) - workflow.add_node("retriever_one", retriever_one) - workflow.add_node("retriever_two", retriever_two) - workflow.add_node("qa", qa) - - workflow.set_entry_point("rewrite_query") - workflow.add_edge("rewrite_query", "analyzer_one") - workflow.add_edge("analyzer_one", "retriever_one") - workflow.add_conditional_edges( - "rewrite_query", lambda _: "retriever_two", {"retriever_two": "retriever_two"} - ) - workflow.add_edge(["retriever_one", "retriever_two"], "qa") - workflow.set_finish_point("qa") - - app = workflow.compile() - - assert await app.ainvoke({"query": "what is weather in sf"}, debug=True) == { - "query": "analyzed: query: what is weather in sf", - "docs": ["doc1", "doc2", "doc3", "doc4"], - "answer": "doc1,doc2,doc3,doc4", - } - - assert [c async for c in app.astream({"query": "what is weather in sf"})] == [ - {"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"]}}, - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] - - async with awith_checkpointer(checkpointer_name) as checkpointer: - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_after=["retriever_one"], - ) - config = {"configurable": {"thread_id": "1"}} - - assert [ - c - async for c in app_w_interrupt.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__": ()}, - ] - - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( - snapshot: SnapshotAssertion, mocker: MockerFixture, checkpointer_name: str -) -> None: - from pydantic.v1 import BaseModel, ValidationError - - setup = mocker.Mock() - teardown = mocker.Mock() - - @asynccontextmanager - async def assert_ctx_once() -> AsyncIterator[None]: - assert setup.call_count == 0 - assert teardown.call_count == 0 - try: - yield - finally: - assert setup.call_count == 1 - assert teardown.call_count == 1 - setup.reset_mock() - teardown.reset_mock() - - @asynccontextmanager - async def make_httpx_client() -> AsyncIterator[httpx.AsyncClient]: - setup() - async with httpx.AsyncClient() as client: - try: - yield client - finally: - teardown() - - def sorted_add( - x: list[str], y: Union[list[str], list[tuple[str, str]]] - ) -> list[str]: - if isinstance(y[0], tuple): - for rem, _ in y: - x.remove(rem) - y = [t[1] for t in y] - return sorted(operator.add(x, y)) - - class State(BaseModel): - class Config: - arbitrary_types_allowed = True - - query: str - answer: Optional[str] = None - docs: Annotated[list[str], sorted_add] - client: Annotated[httpx.AsyncClient, Context(make_httpx_client)] - - class Input(BaseModel): - query: str - - class Output(BaseModel): - answer: str - docs: list[str] - - class StateUpdate(BaseModel): - query: Optional[str] = None - answer: Optional[str] = None - docs: Optional[list[str]] = None - - async def rewrite_query(data: State) -> State: - return {"query": f"query: {data.query}"} - - async def analyzer_one(data: State) -> State: - return StateUpdate(query=f"analyzed: {data.query}") - - async def retriever_one(data: State) -> State: - return {"docs": ["doc1", "doc2"]} - - async def retriever_two(data: State) -> State: - await asyncio.sleep(0.1) - return {"docs": ["doc3", "doc4"]} - - async def qa(data: State) -> State: - return {"answer": ",".join(data.docs)} - - async def decider(data: State) -> str: - assert isinstance(data, State) - return "retriever_two" - - workflow = StateGraph(State, input=Input, output=Output) - - workflow.add_node("rewrite_query", rewrite_query) - workflow.add_node("analyzer_one", analyzer_one) - workflow.add_node("retriever_one", retriever_one) - workflow.add_node("retriever_two", retriever_two) - workflow.add_node("qa", qa) - - workflow.set_entry_point("rewrite_query") - workflow.add_edge("rewrite_query", "analyzer_one") - workflow.add_edge("analyzer_one", "retriever_one") - workflow.add_conditional_edges( - "rewrite_query", decider, {"retriever_two": "retriever_two"} - ) - workflow.add_edge(["retriever_one", "retriever_two"], "qa") - workflow.set_finish_point("qa") - - app = workflow.compile() - - async with assert_ctx_once(): - with pytest.raises(ValidationError): - await app.ainvoke({"query": {}}) - - async with assert_ctx_once(): - assert await app.ainvoke({"query": "what is weather in sf"}) == { - "docs": ["doc1", "doc2", "doc3", "doc4"], - "answer": "doc1,doc2,doc3,doc4", - } - - async with assert_ctx_once(): - assert [c async for c in app.astream({"query": "what is weather in sf"})] == [ - {"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"]}}, - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] - - async with awith_checkpointer(checkpointer_name) as checkpointer: - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_after=["retriever_one"], - ) - config = {"configurable": {"thread_id": "1"}} - - async with assert_ctx_once(): - assert [ - c - async for c in app_w_interrupt.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__": ()}, - ] - - async with assert_ctx_once(): - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "query": "analyzed: query: what is weather in sf", - "answer": "doc1,doc2,doc3,doc4", - "docs": ["doc1", "doc2", "doc3", "doc4"], - }, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - "step": 4, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), - ) - - async with assert_ctx_once(): - assert await app_w_interrupt.aupdate_state( - config, {"docs": ["doc5"]}, as_node="rewrite_query" - ) == { - "configurable": { - "thread_id": "1", - "checkpoint_id": AnyStr(), - "checkpoint_ns": "", - } - } - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2( - snapshot: SnapshotAssertion, checkpointer_name: str -) -> None: - from pydantic import BaseModel, ValidationError - - def sorted_add( - x: list[str], y: Union[list[str], list[tuple[str, str]]] - ) -> list[str]: - if isinstance(y[0], tuple): - for rem, _ in y: - x.remove(rem) - y = [t[1] for t in y] - return sorted(operator.add(x, y)) - - class InnerObject(BaseModel): - yo: int - - class State(BaseModel): - query: str - inner: InnerObject - answer: Optional[str] = None - docs: Annotated[list[str], sorted_add] - - class StateUpdate(BaseModel): - query: Optional[str] = None - answer: Optional[str] = None - docs: Optional[list[str]] = None - - async def rewrite_query(data: State) -> State: - return {"query": f"query: {data.query}"} - - async def analyzer_one(data: State) -> State: - return StateUpdate(query=f"analyzed: {data.query}") - - async def retriever_one(data: State) -> State: - return {"docs": ["doc1", "doc2"]} - - async def retriever_two(data: State) -> State: - await asyncio.sleep(0.1) - return {"docs": ["doc3", "doc4"]} - - async def qa(data: State) -> State: - return {"answer": ",".join(data.docs)} - - async def decider(data: State) -> str: - assert isinstance(data, State) - return "retriever_two" - - workflow = StateGraph(State) - - workflow.add_node("rewrite_query", rewrite_query) - workflow.add_node("analyzer_one", analyzer_one) - workflow.add_node("retriever_one", retriever_one) - workflow.add_node("retriever_two", retriever_two) - workflow.add_node("qa", qa) - - workflow.set_entry_point("rewrite_query") - workflow.add_edge("rewrite_query", "analyzer_one") - workflow.add_edge("analyzer_one", "retriever_one") - workflow.add_conditional_edges( - "rewrite_query", decider, {"retriever_two": "retriever_two"} - ) - workflow.add_edge(["retriever_one", "retriever_two"], "qa") - workflow.set_finish_point("qa") - - app = workflow.compile() - - if SHOULD_CHECK_SNAPSHOTS: - assert app.get_graph().draw_mermaid(with_styles=False) == snapshot - assert app.get_input_schema().model_json_schema() == snapshot - assert app.get_output_schema().model_json_schema() == snapshot - - with pytest.raises(ValidationError): - await app.ainvoke({"query": {}}) - - assert await app.ainvoke( - {"query": "what is weather in sf", "inner": {"yo": 1}} - ) == { - "query": "analyzed: query: what is weather in sf", - "docs": ["doc1", "doc2", "doc3", "doc4"], - "answer": "doc1,doc2,doc3,doc4", - "inner": {"yo": 1}, - } - - assert [ - c - async for c in app.astream( - {"query": "what is weather in sf", "inner": {"yo": 1}} - ) - ] == [ - {"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"]}}, - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] - - async with awith_checkpointer(checkpointer_name) as checkpointer: - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_after=["retriever_one"], - ) - config = {"configurable": {"thread_id": "1"}} - - assert [ - c - async for c in app_w_interrupt.astream( - {"query": "what is weather in sf", "inner": {"yo": 1}}, 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__": ()}, - ] - - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] - - assert await app_w_interrupt.aupdate_state( - config, {"docs": ["doc5"]}, as_node="rewrite_query" - ) == { - "configurable": { - "thread_id": "1", - "checkpoint_id": AnyStr(), - "checkpoint_ns": "", - } - } - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular( - checkpointer_name: str, -) -> None: - def sorted_add( - x: list[str], y: Union[list[str], list[tuple[str, str]]] - ) -> list[str]: - if isinstance(y[0], tuple): - for rem, _ in y: - x.remove(rem) - y = [t[1] for t in y] - return sorted(operator.add(x, y)) - - class State(TypedDict, total=False): - query: str - answer: str - docs: Annotated[list[str], sorted_add] - - async def rewrite_query(data: State) -> State: - return {"query": f"query: {data['query']}"} - - async def analyzer_one(data: State) -> State: - await asyncio.sleep(0.1) - return {"query": f"analyzed: {data['query']}"} - - async def retriever_one(data: State) -> State: - return {"docs": ["doc1", "doc2"]} - - async def retriever_two(data: State) -> State: - await asyncio.sleep(0.2) - return {"docs": ["doc3", "doc4"]} - - async def qa(data: State) -> State: - return {"answer": ",".join(data["docs"])} - - workflow = StateGraph(State) - - workflow.add_node("rewrite_query", rewrite_query) - workflow.add_node("analyzer_one", analyzer_one) - workflow.add_node("retriever_one", retriever_one) - workflow.add_node("retriever_two", retriever_two) - workflow.add_node("qa", qa) - - workflow.set_entry_point("rewrite_query") - workflow.add_edge("rewrite_query", "analyzer_one") - workflow.add_edge("analyzer_one", "retriever_one") - workflow.add_edge("rewrite_query", "retriever_two") - workflow.add_edge(["retriever_one", "retriever_two"], "qa") - workflow.set_finish_point("qa") - - # silly edge, to make sure having been triggered before doesn't break - # semantics of named barrier (== waiting edges) - workflow.add_edge("rewrite_query", "qa") - - app = workflow.compile() - - assert await app.ainvoke({"query": "what is weather in sf"}) == { - "query": "analyzed: query: what is weather in sf", - "docs": ["doc1", "doc2", "doc3", "doc4"], - "answer": "doc1,doc2,doc3,doc4", - } - - assert [c async for c in app.astream({"query": "what is weather in sf"})] == [ - {"rewrite_query": {"query": "query: what is weather in sf"}}, - {"qa": {"answer": ""}}, - {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, - {"retriever_two": {"docs": ["doc3", "doc4"]}}, - {"retriever_one": {"docs": ["doc1", "doc2"]}}, - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] - - async with awith_checkpointer(checkpointer_name) as checkpointer: - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_after=["retriever_one"], - ) - config = {"configurable": {"thread_id": "1"}} - - assert [ - c - async for c in app_w_interrupt.astream( - {"query": "what is weather in sf"}, config - ) - ] == [ - {"rewrite_query": {"query": "query: what is weather in sf"}}, - {"qa": {"answer": ""}}, - {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, - {"retriever_two": {"docs": ["doc3", "doc4"]}}, - {"retriever_one": {"docs": ["doc1", "doc2"]}}, - {"__interrupt__": ()}, - ] - - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] - - -async def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None: - def sorted_add( - x: list[str], y: Union[list[str], list[tuple[str, str]]] - ) -> list[str]: - if isinstance(y[0], tuple): - for rem, _ in y: - x.remove(rem) - y = [t[1] for t in y] - return sorted(operator.add(x, y)) - - class State(TypedDict, total=False): - query: str - answer: str - docs: Annotated[list[str], sorted_add] - - async def rewrite_query(data: State) -> State: - return {"query": f"query: {data['query']}"} - - async def analyzer_one(data: State) -> State: - return {"query": f"analyzed: {data['query']}"} - - async def retriever_one(data: State) -> State: - return {"docs": ["doc1", "doc2"]} - - async def retriever_two(data: State) -> State: - await asyncio.sleep(0.1) - return {"docs": ["doc3", "doc4"]} - - async def qa(data: State) -> State: - return {"answer": ",".join(data["docs"])} - - async def decider(data: State) -> None: - return None - - def decider_cond(data: State) -> str: - if data["query"].count("analyzed") > 1: - return "qa" - else: - return "rewrite_query" - - workflow = StateGraph(State) - - workflow.add_node("rewrite_query", rewrite_query) - workflow.add_node("analyzer_one", analyzer_one) - workflow.add_node("retriever_one", retriever_one) - workflow.add_node("retriever_two", retriever_two) - workflow.add_node("decider", decider) - workflow.add_node("qa", qa) - - workflow.set_entry_point("rewrite_query") - workflow.add_edge("rewrite_query", "analyzer_one") - workflow.add_edge("analyzer_one", "retriever_one") - workflow.add_edge("rewrite_query", "retriever_two") - workflow.add_edge(["retriever_one", "retriever_two"], "decider") - workflow.add_conditional_edges("decider", decider_cond) - workflow.set_finish_point("qa") - - app = workflow.compile() - - assert await app.ainvoke({"query": "what is weather in sf"}) == { - "query": "analyzed: query: analyzed: query: what is weather in sf", - "answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4", - "docs": ["doc1", "doc1", "doc2", "doc2", "doc3", "doc3", "doc4", "doc4"], - } - - assert [c async for c in app.astream({"query": "what is weather in sf"})] == [ - {"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"]}}, - {"decider": None}, - {"rewrite_query": {"query": "query: analyzed: query: what is weather in sf"}}, - { - "analyzer_one": { - "query": "analyzed: query: analyzed: query: what is weather in sf" - } - }, - {"retriever_two": {"docs": ["doc3", "doc4"]}}, - {"retriever_one": {"docs": ["doc1", "doc2"]}}, - {"decider": None}, - {"qa": {"answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4"}}, - ] - - -async def test_in_one_fan_out_state_graph_waiting_edge_multiple_cond_edge() -> None: - def sorted_add( - x: list[str], y: Union[list[str], list[tuple[str, str]]] - ) -> list[str]: - if isinstance(y[0], tuple): - for rem, _ in y: - x.remove(rem) - y = [t[1] for t in y] - return sorted(operator.add(x, y)) - - class State(TypedDict, total=False): - query: str - answer: str - docs: Annotated[list[str], sorted_add] - - async def rewrite_query(data: State) -> State: - return {"query": f"query: {data['query']}"} - - async def retriever_picker(data: State) -> list[str]: - return ["analyzer_one", "retriever_two"] - - async def analyzer_one(data: State) -> State: - return {"query": f"analyzed: {data['query']}"} - - async def retriever_one(data: State) -> State: - return {"docs": ["doc1", "doc2"]} - - async def retriever_two(data: State) -> State: - await asyncio.sleep(0.1) - return {"docs": ["doc3", "doc4"]} - - async def qa(data: State) -> State: - return {"answer": ",".join(data["docs"])} - - async def decider(data: State) -> None: - return None - - def decider_cond(data: State) -> str: - if data["query"].count("analyzed") > 1: - return "qa" - else: - return "rewrite_query" - - workflow = StateGraph(State) - - workflow.add_node("rewrite_query", rewrite_query) - workflow.add_node("analyzer_one", analyzer_one) - workflow.add_node("retriever_one", retriever_one) - workflow.add_node("retriever_two", retriever_two) - workflow.add_node("decider", decider) - workflow.add_node("qa", qa) - - workflow.set_entry_point("rewrite_query") - workflow.add_conditional_edges("rewrite_query", retriever_picker) - workflow.add_edge("analyzer_one", "retriever_one") - workflow.add_edge(["retriever_one", "retriever_two"], "decider") - workflow.add_conditional_edges("decider", decider_cond) - workflow.set_finish_point("qa") - - app = workflow.compile() - - assert await app.ainvoke({"query": "what is weather in sf"}) == { - "query": "analyzed: query: analyzed: query: what is weather in sf", - "answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4", - "docs": ["doc1", "doc1", "doc2", "doc2", "doc3", "doc3", "doc4", "doc4"], - } - - assert [c async for c in app.astream({"query": "what is weather in sf"})] == [ - {"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"]}}, - {"decider": None}, - {"rewrite_query": {"query": "query: analyzed: query: what is weather in sf"}}, - { - "analyzer_one": { - "query": "analyzed: query: analyzed: query: what is weather in sf" - } - }, - {"retriever_two": {"docs": ["doc3", "doc4"]}}, - {"retriever_one": {"docs": ["doc1", "doc2"]}}, - {"decider": None}, - {"qa": {"answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4"}}, - ] - - -async def test_nested_graph(snapshot: SnapshotAssertion) -> None: - def never_called_fn(state: Any): - assert 0, "This function should never be called" - - never_called = RunnableLambda(never_called_fn) - - class InnerState(TypedDict): - my_key: str - my_other_key: str - - def up(state: InnerState): - return {"my_key": state["my_key"] + " there", "my_other_key": state["my_key"]} - - inner = StateGraph(InnerState) - inner.add_node("up", up) - inner.set_entry_point("up") - inner.set_finish_point("up") - - class State(TypedDict): - my_key: str - never_called: Any - - async def side(state: State): - return {"my_key": state["my_key"] + " and back again"} - - graph = StateGraph(State) - graph.add_node("inner", inner.compile()) - graph.add_node("side", side) - graph.set_entry_point("inner") - graph.add_edge("inner", "side") - graph.set_finish_point("side") - - app = graph.compile() - - assert await app.ainvoke({"my_key": "my value", "never_called": never_called}) == { - "my_key": "my value there and back again", - "never_called": never_called, - } - assert [ - chunk - async for chunk in app.astream( - {"my_key": "my value", "never_called": never_called} - ) - ] == [ - {"inner": {"my_key": "my value there"}}, - {"side": {"my_key": "my value there and back again"}}, - ] - assert [ - chunk - async for chunk in app.astream( - {"my_key": "my value", "never_called": never_called}, stream_mode="values" - ) - ] == [ - {"my_key": "my value", "never_called": never_called}, - {"my_key": "my value there", "never_called": never_called}, - {"my_key": "my value there and back again", "never_called": never_called}, - ] - times_called = 0 - async for event in app.astream_events( - {"my_key": "my value", "never_called": never_called}, - version="v2", - config={"run_id": UUID(int=0)}, - stream_mode="values", - ): - if event["event"] == "on_chain_end" and event["run_id"] == str(UUID(int=0)): - times_called += 1 - assert event["data"] == { - "output": { - "my_key": "my value there and back again", - "never_called": never_called, - } - } - assert times_called == 1 - times_called = 0 - async for event in app.astream_events( - {"my_key": "my value", "never_called": never_called}, - version="v2", - config={"run_id": UUID(int=0)}, - ): - if event["event"] == "on_chain_end" and event["run_id"] == str(UUID(int=0)): - times_called += 1 - assert event["data"] == { - "output": { - "my_key": "my value there and back again", - "never_called": never_called, - } - } - assert times_called == 1 - - chain = app | RunnablePassthrough() - - assert await chain.ainvoke( - {"my_key": "my value", "never_called": never_called} - ) == { - "my_key": "my value there and back again", - "never_called": never_called, - } - assert [ - chunk - async for chunk in chain.astream( - {"my_key": "my value", "never_called": never_called} - ) - ] == [ - {"inner": {"my_key": "my value there"}}, - {"side": {"my_key": "my value there and back again"}}, - ] - times_called = 0 - async for event in chain.astream_events( - {"my_key": "my value", "never_called": never_called}, - version="v2", - config={"run_id": UUID(int=0)}, - ): - if event["event"] == "on_chain_end" and event["run_id"] == str(UUID(int=0)): - times_called += 1 - assert event["data"] == { - "output": [ - {"inner": {"my_key": "my value there"}}, - {"side": {"my_key": "my value there and back again"}}, - ] - } - assert times_called == 1 - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_stream_subgraphs_during_execution(checkpointer_name: str) -> None: - class InnerState(TypedDict): - my_key: Annotated[str, operator.add] - my_other_key: str - - async def inner_1(state: InnerState): - return {"my_key": "got here", "my_other_key": state["my_key"]} - - async def inner_2(state: InnerState): - await asyncio.sleep(0.5) - return { - "my_key": " and there", - "my_other_key": state["my_key"], - } - - inner = StateGraph(InnerState) - inner.add_node("inner_1", inner_1) - inner.add_node("inner_2", inner_2) - inner.add_edge("inner_1", "inner_2") - inner.set_entry_point("inner_1") - inner.set_finish_point("inner_2") - - class State(TypedDict): - my_key: Annotated[str, operator.add] - - async def outer_1(state: State): - await asyncio.sleep(0.2) - return {"my_key": " and parallel"} - - async def outer_2(state: State): - return {"my_key": " and back again"} - - graph = StateGraph(State) - graph.add_node("inner", inner.compile()) - graph.add_node("outer_1", outer_1) - graph.add_node("outer_2", outer_2) - - graph.add_edge(START, "inner") - graph.add_edge(START, "outer_1") - graph.add_edge(["inner", "outer_1"], "outer_2") - graph.add_edge("outer_2", END) - - async with awith_checkpointer(checkpointer_name) as checkpointer: - app = graph.compile(checkpointer=checkpointer) - - start = perf_counter() - chunks: list[tuple[float, Any]] = [] - config = {"configurable": {"thread_id": "2"}} - async for c in app.astream({"my_key": ""}, config, subgraphs=True): - chunks.append((round(perf_counter() - start, 1), c)) - for idx in range(len(chunks)): - elapsed, c = chunks[idx] - chunks[idx] = (round(elapsed - chunks[0][0], 1), c) - - assert chunks == [ - # arrives before "inner" finishes - ( - FloatBetween(0.0, 0.1), - ( - (AnyStr("inner:"),), - {"inner_1": {"my_key": "got here", "my_other_key": ""}}, - ), - ), - (FloatBetween(0.2, 0.4), ((), {"outer_1": {"my_key": " and parallel"}})), - ( - FloatBetween(0.5, 0.8), - ( - (AnyStr("inner:"),), - {"inner_2": {"my_key": " and there", "my_other_key": "got here"}}, - ), - ), - (FloatBetween(0.5, 0.8), ((), {"inner": {"my_key": "got here and there"}})), - (FloatBetween(0.5, 0.8), ((), {"outer_2": {"my_key": " and back again"}})), - ] - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_stream_buffering_single_node(checkpointer_name: str) -> None: - class State(TypedDict): - my_key: Annotated[str, operator.add] - - async def node(state: State, writer: StreamWriter): - writer("Before sleep") - await asyncio.sleep(0.2) - writer("After sleep") - return {"my_key": "got here"} - - builder = StateGraph(State) - builder.add_node("node", node) - builder.add_edge(START, "node") - builder.add_edge("node", END) - - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer) - - start = perf_counter() - chunks: list[tuple[float, Any]] = [] - config = {"configurable": {"thread_id": "2"}} - async for c in graph.astream({"my_key": ""}, config, stream_mode="custom"): - chunks.append((round(perf_counter() - start, 1), c)) - - assert chunks == [ - (FloatBetween(0.0, 0.1), "Before sleep"), - (FloatBetween(0.2, 0.3), "After sleep"), - ] - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_nested_graph_interrupts_parallel(checkpointer_name: str) -> None: - class InnerState(TypedDict): - my_key: Annotated[str, operator.add] - my_other_key: str - - async def inner_1(state: InnerState): - await asyncio.sleep(0.1) - return {"my_key": "got here", "my_other_key": state["my_key"]} - - async def inner_2(state: InnerState): - return { - "my_key": " and there", - "my_other_key": state["my_key"], - } - - inner = StateGraph(InnerState) - inner.add_node("inner_1", inner_1) - inner.add_node("inner_2", inner_2) - inner.add_edge("inner_1", "inner_2") - inner.set_entry_point("inner_1") - inner.set_finish_point("inner_2") - - class State(TypedDict): - my_key: Annotated[str, operator.add] - - async def outer_1(state: State): - return {"my_key": " and parallel"} - - async def outer_2(state: State): - return {"my_key": " and back again"} - - graph = StateGraph(State) - graph.add_node( - "inner", - inner.compile(interrupt_before=["inner_2"]), - ) - graph.add_node("outer_1", outer_1) - graph.add_node("outer_2", outer_2) - - graph.add_edge(START, "inner") - graph.add_edge(START, "outer_1") - graph.add_edge(["inner", "outer_1"], "outer_2") - graph.set_finish_point("outer_2") - - async with awith_checkpointer(checkpointer_name) as checkpointer: - app = graph.compile(checkpointer=checkpointer) - - # test invoke w/ nested interrupt - config = {"configurable": {"thread_id": "1"}} - assert await app.ainvoke({"my_key": ""}, config, debug=True) == { - "my_key": " and parallel", - } - - assert await app.ainvoke(None, config, debug=True) == { - "my_key": "got here and there and parallel and back again", - } - - # below combo of assertions is asserting two things - # - outer_1 finishes before inner interrupts (because we see its output in stream, which only happens after node finishes) - # - the writes of outer are persisted in 1st call and used in 2nd call, ie outer isn't called again (because we dont see outer_1 output again in 2nd stream) - # test stream updates w/ nested interrupt - config = {"configurable": {"thread_id": "2"}} - assert [ - c async for c in app.astream({"my_key": ""}, config, subgraphs=True) - ] == [ - # we got to parallel node first - ((), {"outer_1": {"my_key": " and parallel"}}), - ( - (AnyStr("inner:"),), - {"inner_1": {"my_key": "got here", "my_other_key": ""}}, - ), - ((), {"__interrupt__": ()}), - ] - assert [c async for c in app.astream(None, config)] == [ - {"outer_1": {"my_key": " and parallel"}, "__metadata__": {"cached": True}}, - {"inner": {"my_key": "got here and there"}}, - {"outer_2": {"my_key": " and back again"}}, - ] - - # test stream values w/ nested interrupt - config = {"configurable": {"thread_id": "3"}} - assert [ - c async for c in app.astream({"my_key": ""}, config, stream_mode="values") - ] == [ - {"my_key": ""}, - {"my_key": " and parallel"}, - ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ - {"my_key": ""}, - {"my_key": "got here and there and parallel"}, - {"my_key": "got here and there and parallel and back again"}, - ] - - # # test interrupts BEFORE the parallel node - app = graph.compile(checkpointer=checkpointer, interrupt_before=["outer_1"]) - config = {"configurable": {"thread_id": "4"}} - assert [ - c async for c in app.astream({"my_key": ""}, config, stream_mode="values") - ] == [ - {"my_key": ""}, - ] - # while we're waiting for the node w/ interrupt inside to finish - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ - {"my_key": ""}, - {"my_key": " and parallel"}, - ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ - {"my_key": ""}, - {"my_key": "got here and there and parallel"}, - {"my_key": "got here and there and parallel and back again"}, - ] - - # test interrupts AFTER the parallel node - app = graph.compile(checkpointer=checkpointer, interrupt_after=["outer_1"]) - config = {"configurable": {"thread_id": "5"}} - assert [ - c async for c in app.astream({"my_key": ""}, config, stream_mode="values") - ] == [ - {"my_key": ""}, - {"my_key": " and parallel"}, - ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ - {"my_key": ""}, - {"my_key": "got here and there and parallel"}, - ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ - {"my_key": "got here and there and parallel"}, - {"my_key": "got here and there and parallel and back again"}, - ] - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_doubly_nested_graph_interrupts(checkpointer_name: str) -> None: - class State(TypedDict): - my_key: str - - class ChildState(TypedDict): - my_key: str - - class GrandChildState(TypedDict): - my_key: str - - async def grandchild_1(state: ChildState): - return {"my_key": state["my_key"] + " here"} - - async def grandchild_2(state: ChildState): - return { - "my_key": state["my_key"] + " and there", - } - - grandchild = StateGraph(GrandChildState) - grandchild.add_node("grandchild_1", grandchild_1) - grandchild.add_node("grandchild_2", grandchild_2) - grandchild.add_edge("grandchild_1", "grandchild_2") - grandchild.set_entry_point("grandchild_1") - grandchild.set_finish_point("grandchild_2") - - child = StateGraph(ChildState) - child.add_node( - "child_1", - grandchild.compile(interrupt_before=["grandchild_2"]), - ) - child.set_entry_point("child_1") - child.set_finish_point("child_1") - - async def parent_1(state: State): - return {"my_key": "hi " + state["my_key"]} - - async def parent_2(state: State): - return {"my_key": state["my_key"] + " and back again"} - - graph = StateGraph(State) - graph.add_node("parent_1", parent_1) - graph.add_node("child", child.compile()) - graph.add_node("parent_2", parent_2) - graph.set_entry_point("parent_1") - graph.add_edge("parent_1", "child") - graph.add_edge("child", "parent_2") - graph.set_finish_point("parent_2") - - async with awith_checkpointer(checkpointer_name) as checkpointer: - app = graph.compile(checkpointer=checkpointer) - - # test invoke w/ nested interrupt - config = {"configurable": {"thread_id": "1"}} - assert await app.ainvoke({"my_key": "my value"}, config, debug=True) == { - "my_key": "hi my value", - } - - assert await app.ainvoke(None, config, debug=True) == { - "my_key": "hi my value here and there and back again", - } - - # test stream updates w/ nested interrupt - nodes: list[str] = [] - config = { - "configurable": {"thread_id": "2", CONFIG_KEY_NODE_FINISHED: nodes.append} - } - assert [c async for c in app.astream({"my_key": "my value"}, config)] == [ - {"parent_1": {"my_key": "hi my value"}}, - {"__interrupt__": ()}, - ] - assert nodes == ["parent_1", "grandchild_1"] - assert [c async for c in app.astream(None, config)] == [ - {"child": {"my_key": "hi my value here and there"}}, - {"parent_2": {"my_key": "hi my value here and there and back again"}}, - ] - assert nodes == [ - "parent_1", - "grandchild_1", - "grandchild_2", - "child_1", - "child", - "parent_2", - ] - - # test stream values w/ nested interrupt - config = {"configurable": {"thread_id": "3"}} - assert [ - c - async for c in app.astream( - {"my_key": "my value"}, config, stream_mode="values" - ) - ] == [ - {"my_key": "my value"}, - {"my_key": "hi my value"}, - ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ - {"my_key": "hi my value"}, - {"my_key": "hi my value here and there"}, - {"my_key": "hi my value here and there and back again"}, - ] - - -async def test_checkpoint_metadata() -> None: - """This test verifies that a run's configurable fields are merged with the - previous checkpoint config for each step in the run. - """ - # set up test - from langchain_core.language_models.fake_chat_models import ( - FakeMessagesListChatModel, - ) - from langchain_core.messages import AIMessage, AnyMessage - from langchain_core.prompts import ChatPromptTemplate - from langchain_core.tools import tool - - # graph state - class BaseState(TypedDict): - messages: Annotated[list[AnyMessage], add_messages] - - # initialize graph nodes - @tool() - def search_api(query: str) -> str: - """Searches the API for the query.""" - return f"result for {query}" - - tools = [search_api] - - prompt = ChatPromptTemplate.from_messages( - [ - ("system", "You are a nice assistant."), - ("placeholder", "{messages}"), - ] - ) - - model = FakeMessagesListChatModel( - responses=[ - AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - }, - ], - ), - AIMessage(content="answer"), - ] - ) - - def agent(state: BaseState, config: RunnableConfig) -> BaseState: - formatted = prompt.invoke(state) - response = model.invoke(formatted) - return {"messages": response} - - def should_continue(data: BaseState) -> str: - # Logic to decide whether to continue in the loop or exit - if not data["messages"][-1].tool_calls: - return "exit" - else: - return "continue" - - # define graphs w/ and w/o interrupt - workflow = StateGraph(BaseState) - workflow.add_node("agent", agent) - workflow.add_node("tools", ToolNode(tools)) - workflow.set_entry_point("agent") - workflow.add_conditional_edges( - "agent", should_continue, {"continue": "tools", "exit": END} - ) - workflow.add_edge("tools", "agent") - - # graph w/o interrupt - checkpointer_1 = MemorySaverAssertCheckpointMetadata() - app = workflow.compile(checkpointer=checkpointer_1) - - # graph w/ interrupt - checkpointer_2 = MemorySaverAssertCheckpointMetadata() - app_w_interrupt = workflow.compile( - checkpointer=checkpointer_2, interrupt_before=["tools"] - ) - - # assertions - - # invoke graph w/o interrupt - await app.ainvoke( - {"messages": ["what is weather in sf"]}, - { - "configurable": { - "thread_id": "1", - "test_config_1": "foo", - "test_config_2": "bar", - }, - }, - ) - - config = {"configurable": {"thread_id": "1"}} - - # assert that checkpoint metadata contains the run's configurable fields - chkpnt_metadata_1 = (await checkpointer_1.aget_tuple(config)).metadata - assert chkpnt_metadata_1["thread_id"] == "1" - assert chkpnt_metadata_1["test_config_1"] == "foo" - assert chkpnt_metadata_1["test_config_2"] == "bar" - - # Verify that all checkpoint metadata have the expected keys. This check - # is needed because a run may have an arbitrary number of steps depending - # on how the graph is constructed. - chkpnt_tuples_1 = checkpointer_1.alist(config) - async for chkpnt_tuple in chkpnt_tuples_1: - assert chkpnt_tuple.metadata["thread_id"] == "1" - assert chkpnt_tuple.metadata["test_config_1"] == "foo" - assert chkpnt_tuple.metadata["test_config_2"] == "bar" - - # invoke graph, but interrupt before tool call - await app_w_interrupt.ainvoke( - {"messages": ["what is weather in sf"]}, - { - "configurable": { - "thread_id": "2", - "test_config_3": "foo", - "test_config_4": "bar", - }, - }, - ) - - config = {"configurable": {"thread_id": "2"}} - - # assert that checkpoint metadata contains the run's configurable fields - chkpnt_metadata_2 = (await checkpointer_2.aget_tuple(config)).metadata - assert chkpnt_metadata_2["thread_id"] == "2" - assert chkpnt_metadata_2["test_config_3"] == "foo" - assert chkpnt_metadata_2["test_config_4"] == "bar" - - # resume graph execution - await app_w_interrupt.ainvoke( - input=None, - config={ - "configurable": { - "thread_id": "2", - "test_config_3": "foo", - "test_config_4": "bar", - } - }, - ) - - # assert that checkpoint metadata contains the run's configurable fields - chkpnt_metadata_3 = (await checkpointer_2.aget_tuple(config)).metadata - assert chkpnt_metadata_3["thread_id"] == "2" - assert chkpnt_metadata_3["test_config_3"] == "foo" - assert chkpnt_metadata_3["test_config_4"] == "bar" - - # Verify that all checkpoint metadata have the expected keys. This check - # is needed because a run may have an arbitrary number of steps depending - # on how the graph is constructed. - chkpnt_tuples_2 = checkpointer_2.alist(config) - async for chkpnt_tuple in chkpnt_tuples_2: - assert chkpnt_tuple.metadata["thread_id"] == "2" - assert chkpnt_tuple.metadata["test_config_3"] == "foo" - assert chkpnt_tuple.metadata["test_config_4"] == "bar" - - -async def test_checkpointer_null_pending_writes() -> None: - class Node: - def __init__(self, name: str): - self.name = name - setattr(self, "__name__", name) - - def __call__(self, state): - return [self.name] - - builder = StateGraph(Annotated[list, operator.add]) - builder.add_node(Node("1")) - builder.add_edge(START, "1") - graph = builder.compile(checkpointer=MemorySaverNoPending()) - assert graph.invoke([], {"configurable": {"thread_id": "foo"}}) == ["1"] - assert graph.invoke([], {"configurable": {"thread_id": "foo"}}) == ["1"] * 2 - assert (await graph.ainvoke([], {"configurable": {"thread_id": "foo"}})) == [ - "1" - ] * 3 - assert (await graph.ainvoke([], {"configurable": {"thread_id": "foo"}})) == [ - "1" - ] * 4 - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -@pytest.mark.parametrize("store_name", ALL_STORES_ASYNC) -async def test_store_injected_async(checkpointer_name: str, store_name: str) -> None: - class State(TypedDict): - count: Annotated[int, operator.add] - - doc_id = str(uuid.uuid4()) - doc = {"some-key": "this-is-a-val"} - uid = uuid.uuid4().hex - namespace = (f"foo-{uid}", "bar") - thread_1 = str(uuid.uuid4()) - thread_2 = str(uuid.uuid4()) - - class Node: - def __init__(self, i: Optional[int] = None): - self.i = i - - async def __call__( - self, inputs: State, config: RunnableConfig, store: BaseStore - ): - assert isinstance(store, BaseStore) - await store.aput( - namespace - if self.i is not None - and config["configurable"]["thread_id"] in (thread_1, thread_2) - else (f"foo_{self.i}", "bar"), - doc_id, - { - **doc, - "from_thread": config["configurable"]["thread_id"], - "some_val": inputs["count"], - }, - ) - return {"count": 1} - - def other_node(inputs: State, config: RunnableConfig, store: BaseStore): - assert isinstance(store, BaseStore) - store.put(("not", "interesting"), "key", {"val": "val"}) - item = store.get(("not", "interesting"), "key") - assert item is not None - assert item.value == {"val": "val"} - return {"count": 0} - - builder = StateGraph(State) - builder.add_node("node", Node()) - builder.add_node("other_node", other_node) - builder.add_edge("__start__", "node") - builder.add_edge("node", "other_node") - - N = 500 - M = 1 - - for i in range(N): - builder.add_node(f"node_{i}", Node(i)) - builder.add_edge("__start__", f"node_{i}") - - async with ( - awith_checkpointer(checkpointer_name) as checkpointer, - awith_store(store_name) as the_store, - ): - graph = builder.compile(store=the_store, checkpointer=checkpointer) - - # Test batch operations with multiple threads - results = await graph.abatch( - [{"count": 0}] * M, - ([{"configurable": {"thread_id": str(uuid.uuid4())}}] * (M - 1)) - + [{"configurable": {"thread_id": thread_1}}], - ) - result = results[-1] - assert result == {"count": N + 1} - returned_doc = (await the_store.aget(namespace, doc_id)).value - assert returned_doc == {**doc, "from_thread": thread_1, "some_val": 0} - assert len((await the_store.asearch(namespace))) == 1 - - # Check results after another turn of the same thread - result = await graph.ainvoke( - {"count": 0}, {"configurable": {"thread_id": thread_1}} - ) - assert result == {"count": (N + 1) * 2} - returned_doc = (await the_store.aget(namespace, doc_id)).value - assert returned_doc == {**doc, "from_thread": thread_1, "some_val": N + 1} - assert len((await the_store.asearch(namespace))) == 1 - - # Test with a different thread - result = await graph.ainvoke( - {"count": 0}, {"configurable": {"thread_id": thread_2}} - ) - assert result == {"count": N + 1} - returned_doc = (await the_store.aget(namespace, doc_id)).value - assert returned_doc == { - **doc, - "from_thread": thread_2, - "some_val": 0, - } # Overwrites the whole doc - assert ( - len((await the_store.asearch(namespace))) == 1 - ) # still overwriting the same one - - -async def test_debug_retry(): - class State(TypedDict): - messages: Annotated[list[str], operator.add] - - def node(name): - async def _node(state: State): - return {"messages": [f"entered {name} node"]} - - return _node - - builder = StateGraph(State) - builder.add_node("one", node("one")) - builder.add_node("two", node("two")) - builder.add_edge(START, "one") - builder.add_edge("one", "two") - builder.add_edge("two", END) - - saver = InMemorySaver() - - graph = builder.compile(checkpointer=saver) - - config = {"configurable": {"thread_id": "1"}} - await graph.ainvoke({"messages": []}, config=config) - - # re-run step: 1 - async for c in saver.alist(config): - if c.metadata["step"] == 1: - target_config = c.parent_config - break - assert target_config is not None - - update_config = await graph.aupdate_state(target_config, values=None) - - events = [ - c async for c in graph.astream(None, config=update_config, stream_mode="debug") - ] - - checkpoint_events = list( - reversed([e["payload"] for e in events if e["type"] == "checkpoint"]) - ) - - checkpoint_history = { - c.config["configurable"]["checkpoint_id"]: c - async for c in graph.aget_state_history(config) - } - - def lax_normalize_config(config: Optional[dict]) -> Optional[dict]: - if config is None: - return None - return config["configurable"] - - for stream in checkpoint_events: - stream_conf = lax_normalize_config(stream["config"]) - stream_parent_conf = lax_normalize_config(stream["parent_config"]) - assert stream_conf != stream_parent_conf - - # ensure the streamed checkpoint == checkpoint from checkpointer.list() - history = checkpoint_history[stream["config"]["configurable"]["checkpoint_id"]] - history_conf = lax_normalize_config(history.config) - assert stream_conf == history_conf - - history_parent_conf = lax_normalize_config(history.parent_config) - assert stream_parent_conf == history_parent_conf - - -async def test_debug_subgraphs(): - class State(TypedDict): - messages: Annotated[list[str], operator.add] - - def node(name): - async def _node(state: State): - return {"messages": [f"entered {name} node"]} - - return _node - - parent = StateGraph(State) - child = StateGraph(State) - - child.add_node("c_one", node("c_one")) - child.add_node("c_two", node("c_two")) - child.add_edge(START, "c_one") - child.add_edge("c_one", "c_two") - child.add_edge("c_two", END) - - parent.add_node("p_one", node("p_one")) - parent.add_node("p_two", child.compile()) - parent.add_edge(START, "p_one") - parent.add_edge("p_one", "p_two") - parent.add_edge("p_two", END) - - graph = parent.compile(checkpointer=InMemorySaver()) - - config = {"configurable": {"thread_id": "1"}} - events = [ - c - async for c in graph.astream( - {"messages": []}, - config=config, - stream_mode="debug", - ) - ] - - checkpoint_events = list( - reversed([e["payload"] for e in events if e["type"] == "checkpoint"]) - ) - checkpoint_history = [c async for c in graph.aget_state_history(config)] - - assert len(checkpoint_events) == len(checkpoint_history) - - def normalize_config(config: Optional[dict]) -> Optional[dict]: - if config is None: - return None - return config["configurable"] - - for stream, history in zip(checkpoint_events, checkpoint_history): - assert stream["values"] == history.values - assert stream["next"] == list(history.next) - assert normalize_config(stream["config"]) == normalize_config(history.config) - assert normalize_config(stream["parent_config"]) == normalize_config( - history.parent_config - ) - - assert len(stream["tasks"]) == len(history.tasks) - for stream_task, history_task in zip(stream["tasks"], history.tasks): - assert stream_task["id"] == history_task.id - assert stream_task["name"] == history_task.name - assert stream_task["interrupts"] == history_task.interrupts - assert stream_task.get("error") == history_task.error - assert stream_task.get("state") == history_task.state - - -async def test_debug_nested_subgraphs(): - from collections import defaultdict - - class State(TypedDict): - messages: Annotated[list[str], operator.add] - - def node(name): - async def _node(state: State): - return {"messages": [f"entered {name} node"]} - - return _node - - grand_parent = StateGraph(State) - parent = StateGraph(State) - child = StateGraph(State) - - child.add_node("c_one", node("c_one")) - child.add_node("c_two", node("c_two")) - child.add_edge(START, "c_one") - child.add_edge("c_one", "c_two") - child.add_edge("c_two", END) - - parent.add_node("p_one", node("p_one")) - parent.add_node("p_two", child.compile()) - parent.add_edge(START, "p_one") - parent.add_edge("p_one", "p_two") - parent.add_edge("p_two", END) - - grand_parent.add_node("gp_one", node("gp_one")) - grand_parent.add_node("gp_two", parent.compile()) - grand_parent.add_edge(START, "gp_one") - grand_parent.add_edge("gp_one", "gp_two") - grand_parent.add_edge("gp_two", END) - - graph = grand_parent.compile(checkpointer=InMemorySaver()) - - config = {"configurable": {"thread_id": "1"}} - events = [ - c - async for c in graph.astream( - {"messages": []}, - config=config, - stream_mode="debug", - subgraphs=True, - ) - ] - - stream_ns: dict[tuple, dict] = defaultdict(list) - for ns, e in events: - if e["type"] == "checkpoint": - stream_ns[ns].append(e["payload"]) - - assert list(stream_ns.keys()) == [ - (), - (AnyStr("gp_two:"),), - (AnyStr("gp_two:"), AnyStr("p_two:")), - ] - - history_ns = {} - for ns in stream_ns.keys(): - - async def get_history(): - history = [ - c - async for c in graph.aget_state_history( - {"configurable": {"thread_id": "1", "checkpoint_ns": "|".join(ns)}} - ) - ] - return history[::-1] - - history_ns[ns] = await get_history() - - def normalize_config(config: Optional[dict]) -> Optional[dict]: - if config is None: - return None - - clean_config = {} - clean_config["thread_id"] = config["configurable"]["thread_id"] - clean_config["checkpoint_id"] = config["configurable"]["checkpoint_id"] - clean_config["checkpoint_ns"] = config["configurable"]["checkpoint_ns"] - if "checkpoint_map" in config["configurable"]: - clean_config["checkpoint_map"] = config["configurable"]["checkpoint_map"] - - return clean_config - - for checkpoint_events, checkpoint_history in zip( - stream_ns.values(), history_ns.values() - ): - for stream, history in zip(checkpoint_events, checkpoint_history): - assert stream["values"] == history.values - assert stream["next"] == list(history.next) - assert normalize_config(stream["config"]) == normalize_config( - history.config - ) - assert normalize_config(stream["parent_config"]) == normalize_config( - history.parent_config - ) - - assert len(stream["tasks"]) == len(history.tasks) - for stream_task, history_task in zip(stream["tasks"], history.tasks): - assert stream_task["id"] == history_task.id - assert stream_task["name"] == history_task.name - assert stream_task["interrupts"] == history_task.interrupts - assert stream_task.get("error") == history_task.error - assert stream_task.get("state") == history_task.state - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_parent_command(checkpointer_name: str) -> None: - from langchain_core.messages import BaseMessage - from langchain_core.tools import tool - - @tool(return_direct=True) - def get_user_name() -> Command: - """Retrieve user name""" - return Command(update={"user_name": "Meow"}, graph=Command.PARENT) - - subgraph_builder = StateGraph(MessagesState) - subgraph_builder.add_node("tool", get_user_name) - subgraph_builder.add_edge(START, "tool") - subgraph = subgraph_builder.compile() - - class CustomParentState(TypedDict): - messages: Annotated[list[BaseMessage], add_messages] - # this key is not available to the child graph - user_name: str - - builder = StateGraph(CustomParentState) - builder.add_node("alice", subgraph) - builder.add_edge(START, "alice") - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer) - - config = {"configurable": {"thread_id": "1"}} - - assert await graph.ainvoke( - {"messages": [("user", "get user name")]}, config - ) == { - "messages": [ - _AnyIdHumanMessage( - content="get user name", additional_kwargs={}, response_metadata={} - ), - ], - "user_name": "Meow", - } - assert await graph.aget_state(config) == StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage( - content="get user name", - additional_kwargs={}, - response_metadata={}, - ), - ], - "user_name": "Meow", - }, - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": { - "alice": { - "messages": [ - _AnyIdHumanMessage( - content="get user name", - additional_kwargs={}, - response_metadata={}, - ), - ], - "user_name": "Meow", - } - }, - "thread_id": "1", - "step": 1, - "parents": {}, - }, - created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), - tasks=(), - ) - - -@NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_interrupt_subgraph(checkpointer_name: str): - class State(TypedDict): - baz: str - - def foo(state): - return {"baz": "foo"} - - def bar(state): - value = interrupt("Please provide baz value:") - return {"baz": value} - - child_builder = StateGraph(State) - child_builder.add_node(bar) - child_builder.add_edge(START, "bar") - - builder = StateGraph(State) - builder.add_node(foo) - builder.add_node("bar", child_builder.compile()) - builder.add_edge(START, "foo") - builder.add_edge("foo", "bar") - - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer) - - thread1 = {"configurable": {"thread_id": "1"}} - # First run, interrupted at bar - assert await graph.ainvoke({"baz": ""}, thread1) - # Resume with answer - assert await graph.ainvoke(Command(resume="bar"), thread1) - - -@NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_interrupt_multiple(checkpointer_name: str): - class State(TypedDict): - my_key: Annotated[str, operator.add] - - async def node(s: State) -> State: - answer = interrupt({"value": 1}) - answer2 = interrupt({"value": 2}) - return {"my_key": answer + " " + answer2} - - builder = StateGraph(State) - builder.add_node("node", node) - builder.add_edge(START, "node") - - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer) - thread1 = {"configurable": {"thread_id": "1"}} - - assert [ - e async for e in graph.astream({"my_key": "DE", "market": "DE"}, thread1) - ] == [ - { - "__interrupt__": ( - Interrupt( - value={"value": 1}, - resumable=True, - ns=[AnyStr("node:")], - when="during", - ), - ) - } - ] - - assert [ - event - async for event in graph.astream( - Command(resume="answer 1", update={"my_key": "foofoo"}), - thread1, - stream_mode="updates", - ) - ] == [ - { - "__interrupt__": ( - Interrupt( - value={"value": 2}, - resumable=True, - ns=[AnyStr("node:")], - when="during", - ), - ) - } - ] - - assert [ - event - async for event in graph.astream( - Command(resume="answer 2"), thread1, stream_mode="updates" - ) - ] == [ - {"node": {"my_key": "answer 1 answer 2"}}, - ] - - -@NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_interrupt_loop(checkpointer_name: str): - class State(TypedDict): - age: int - other: str - - async def ask_age(s: State): - """Ask an expert for help.""" - question = "How old are you?" - value = None - for _ in range(10): - value: str = interrupt(question) - if not value.isdigit() or int(value) < 18: - question = "invalid response" - value = None - else: - break - - return {"age": int(value)} - - builder = StateGraph(State) - builder.add_node("node", ask_age) - builder.add_edge(START, "node") - - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer) - thread1 = {"configurable": {"thread_id": "1"}} - - assert [e async for e in graph.astream({"other": ""}, thread1)] == [ - { - "__interrupt__": ( - Interrupt( - value="How old are you?", - resumable=True, - ns=[AnyStr("node:")], - when="during", - ), - ) - } - ] - - assert [ - event - async for event in graph.astream( - Command(resume="13"), - thread1, - ) - ] == [ - { - "__interrupt__": ( - Interrupt( - value="invalid response", - resumable=True, - ns=[AnyStr("node:")], - when="during", - ), - ) - } - ] - - assert [ - event - async for event in graph.astream( - Command(resume="15"), - thread1, - ) - ] == [ - { - "__interrupt__": ( - Interrupt( - value="invalid response", - resumable=True, - ns=[AnyStr("node:")], - when="during", - ), - ) - } - ] - - assert [ - event async for event in graph.astream(Command(resume="19"), thread1) - ] == [ - {"node": {"age": 19}}, - ] - - -@NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_interrupt_functional(checkpointer_name: str) -> None: - @task - async def foo(state: dict) -> dict: - return {"a": state["a"] + "foo"} - - @task - async def bar(state: dict) -> dict: - return {"a": state["a"] + "bar", "b": state["b"]} - - async with awith_checkpointer(checkpointer_name) as checkpointer: - - @entrypoint(checkpointer=checkpointer) - async def graph(inputs: dict) -> dict: - foo_result = await foo(inputs) - value = interrupt("Provide value for bar:") - bar_input = {**foo_result, "b": value} - bar_result = await bar(bar_input) - return bar_result - - config = {"configurable": {"thread_id": "1"}} - # First run, interrupted at bar - await graph.ainvoke({"a": ""}, config) - # Resume with an answer - res = await graph.ainvoke(Command(resume="bar"), config) - assert res == {"a": "foobar", "b": "bar"} - - -@NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_interrupt_task_functional(checkpointer_name: str) -> None: - @task - async def foo(state: dict) -> dict: - return {"a": state["a"] + "foo"} - - @task - async def bar(state: dict) -> dict: - value = interrupt("Provide value for bar:") - return {"a": state["a"] + value} - - async with awith_checkpointer(checkpointer_name) as checkpointer: - - @entrypoint(checkpointer=checkpointer) - async def graph(inputs: dict) -> dict: - foo_result = await foo(inputs) - bar_result = await bar(foo_result) - return bar_result - - config = {"configurable": {"thread_id": "1"}} - # First run, interrupted at bar - await graph.ainvoke({"a": ""}, config) - # Resume with an answer - res = await graph.ainvoke(Command(resume="bar"), config) - assert res == {"a": "foobar"} - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_command_with_static_breakpoints(checkpointer_name: str) -> None: - """Test that we can use Command to resume and update with static breakpoints.""" - - class State(TypedDict): - """The graph state.""" - - foo: str - - def node1(state: State): - return { - "foo": state["foo"] + "|node-1", - } - - def node2(state: State): - return { - "foo": state["foo"] + "|node-2", - } - - builder = StateGraph(State) - builder.add_node("node1", node1) - builder.add_node("node2", node2) - builder.add_edge(START, "node1") - builder.add_edge("node1", "node2") - - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer, interrupt_before=["node1"]) - config = {"configurable": {"thread_id": str(uuid.uuid4())}} - - # Start the graph and interrupt at the first node - await graph.ainvoke({"foo": "abc"}, config) - result = await graph.ainvoke(Command(update={"foo": "def"}), config) - assert result == {"foo": "def|node-1|node-2"} - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_multistep_plan(checkpointer_name: str): - from langchain_core.messages import AnyMessage - - class State(TypedDict, total=False): - plan: list[Union[str, list[str]]] - messages: Annotated[list[AnyMessage], add_messages] - - def planner(state: State): - if state.get("plan") is None: - # create plan somehow - plan = ["step1", ["step2", "step3"], "step4"] - # pick the first step to execute next - first_step, *plan = plan - # put the rest of plan in state - return Command(goto=first_step, update={"plan": plan}) - elif state["plan"]: - # go to the next step of the plan - next_step, *next_plan = state["plan"] - return Command(goto=next_step, update={"plan": next_plan}) - else: - # the end of the plan - pass - - def step1(state: State): - return Command(goto="planner", update={"messages": [("human", "step1")]}) - - def step2(state: State): - return Command(goto="planner", update={"messages": [("human", "step2")]}) - - def step3(state: State): - return Command(goto="planner", update={"messages": [("human", "step3")]}) - - def step4(state: State): - return Command(goto="planner", update={"messages": [("human", "step4")]}) - - builder = StateGraph(State) - builder.add_node(planner) - builder.add_node(step1) - builder.add_node(step2) - builder.add_node(step3) - builder.add_node(step4) - builder.add_edge(START, "planner") - - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer) - - config = {"configurable": {"thread_id": "1"}} - - assert await graph.ainvoke({"messages": [("human", "start")]}, config) == { - "messages": [ - _AnyIdHumanMessage(content="start"), - _AnyIdHumanMessage(content="step1"), - _AnyIdHumanMessage(content="step2"), - _AnyIdHumanMessage(content="step3"), - _AnyIdHumanMessage(content="step4"), - ], - "plan": [], - } - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_command_goto_with_static_breakpoints(checkpointer_name: str) -> None: - """Use Command goto with static breakpoints.""" - - class State(TypedDict): - """The graph state.""" - - foo: Annotated[str, operator.add] - - def node1(state: State): - return { - "foo": "|node-1", - } - - def node2(state: State): - return { - "foo": "|node-2", - } - - builder = StateGraph(State) - builder.add_node("node1", node1) - builder.add_node("node2", node2) - builder.add_edge(START, "node1") - builder.add_edge("node1", "node2") - - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer, interrupt_before=["node1"]) - - config = {"configurable": {"thread_id": str(uuid.uuid4())}} - - # Start the graph and interrupt at the first node - await graph.ainvoke({"foo": "abc"}, config) - result = await graph.ainvoke(Command(goto=["node2"]), config) - assert result == {"foo": "abc|node-1|node-2|node-2"} - - -async def test_nested_graph_state_error_handling(): - """Test error handling when updating state in nested graphs.""" - - class State(TypedDict): - count: int - - def child_node(state: State): - return {"count": state["count"] + 1} - - child = StateGraph(State) - child.add_node("child", child_node) - child.add_edge(START, "child") - - parent = StateGraph(State) - parent.add_node("child_graph", child.compile()) - parent.add_edge(START, "child_graph") - - app = parent.compile(checkpointer=MemorySaver()) - - # Test invalid state update on parent - with pytest.raises(InvalidUpdateError): - await app.aupdate_state( - {"configurable": {"thread_id": "1"}}, {"invalid_key": "value"} - ) - - # Test invalid state update on child - with pytest.raises(InvalidUpdateError): - await app.aupdate_state( - {"configurable": {"thread_id": "1", "checkpoint_ns": "child_graph"}}, - {"invalid_key": "value"}, - ) - - -async def test_parallel_node_execution(): - """Test that parallel nodes execute concurrently.""" - - class State(TypedDict): - results: Annotated[list[str], operator.add] - - async def slow_node(state: State): - await asyncio.sleep(1) - return {"results": ["slow"]} - - async def fast_node(state: State): - await asyncio.sleep(2) - return {"results": ["fast"]} - - builder = StateGraph(State) - builder.add_node("slow", slow_node) - builder.add_node("fast", fast_node) - builder.add_edge(START, "slow") - builder.add_edge(START, "fast") - - graph = builder.compile() - - start = perf_counter() - result = await graph.ainvoke({"results": []}) - duration = perf_counter() - start - - # Fast node result should be available first - assert "fast" in result["results"][0] - - # Total duration should be less than sum of both nodes - assert duration < 3.0 - - -@NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_multiple_interrupt_state_persistence(checkpointer_name: str) -> None: - """Test that state is preserved correctly across multiple interrupts.""" - - class State(TypedDict): - steps: Annotated[list[str], operator.add] - - def interruptible_node(state: State): - first = interrupt("First interrupt") - second = interrupt("Second interrupt") - return {"steps": [first, second]} - - builder = StateGraph(State) - builder.add_node("node", interruptible_node) - builder.add_edge(START, "node") - - async with awith_checkpointer(checkpointer_name) as checkpointer: - app = builder.compile(checkpointer=checkpointer) - config = {"configurable": {"thread_id": "1"}} - - # First execution - should hit first interrupt - await app.ainvoke({"steps": []}, config) - - # State should still be empty since node hasn't returned - state = await app.aget_state(config) - assert state.values == {"steps": []} - - # Resume after first interrupt - should hit second interrupt - await app.ainvoke(Command(resume="step1"), config) - - # State should still be empty since node hasn't returned - state = await app.aget_state(config) - assert state.values == {"steps": []} - - # Resume after second interrupt - node should complete - result = await app.ainvoke(Command(resume="step2"), config) - - # Now state should contain both steps since node returned - assert result["steps"] == ["step1", "step2"] - state = await app.aget_state(config) - assert state.values["steps"] == ["step1", "step2"] - - -async def test_concurrent_execution(): - """Test concurrent execution with async nodes.""" - - class State(TypedDict): - counter: Annotated[int, operator.add] - - results = deque() - - async def slow_node(state: State): - await asyncio.sleep(0.1) - return {"counter": 1} - - builder = StateGraph(State) - builder.add_node("node", slow_node) - builder.add_edge(START, "node") - graph = builder.compile() - - async def run_graph(): - result = await graph.ainvoke({"counter": 0}) - results.append(result) - - # Create and gather tasks - tasks = [run_graph() for _ in range(10)] - await asyncio.gather(*tasks) - - # Verify results are independent - assert len(results) == 10 - for result in results: - assert result["counter"] == 1 - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_checkpoint_recovery_async(checkpointer_name: str): - """Test recovery from checkpoints after failures with async nodes.""" - - class State(TypedDict): - steps: Annotated[list[str], operator.add] - attempt: int # Track number of attempts - - async def failing_node(state: State): - # Fail on first attempt, succeed on retry - if state["attempt"] == 1: - raise RuntimeError("Simulated failure") - await asyncio.sleep(0.1) # Simulate async work - return {"steps": ["node1"]} - - async def second_node(state: State): - await asyncio.sleep(0.1) # Simulate async work - return {"steps": ["node2"]} - - builder = StateGraph(State) - builder.add_node("node1", failing_node) - builder.add_node("node2", second_node) - builder.add_edge(START, "node1") - builder.add_edge("node1", "node2") - - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer) - config = {"configurable": {"thread_id": "1"}} - - # First attempt should fail - with pytest.raises(RuntimeError): - await graph.ainvoke({"steps": ["start"], "attempt": 1}, config) - - # Verify checkpoint state - state = await graph.aget_state(config) - assert state is not None - assert state.values == {"steps": ["start"], "attempt": 1} # input state saved - assert state.next == ("node1",) # Should retry failed node - - # Retry with updated attempt count - result = await graph.ainvoke({"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 = [c async for c in graph.aget_state_history(config)] - assert len(history) == 6 # Initial + failed attempt + successful attempt - - # Verify the error was recorded in checkpoint - failed_checkpoint = next(c for c in history if c.tasks and c.tasks[0].error) - assert "RuntimeError('Simulated failure')" in failed_checkpoint.tasks[0].error - - -async def test_multiple_updates_root() -> None: - def node_a(state): - return [Command(update="a1"), Command(update="a2")] - - def node_b(state): - return "b" - - graph = ( - StateGraph(Annotated[str, operator.add]) - .add_sequence([node_a, node_b]) - .add_edge(START, "node_a") - .compile() - ) - - assert await graph.ainvoke("") == "a1a2b" - - # only streams the last update from node_a - assert [c async for c in graph.astream("", stream_mode="updates")] == [ - {"node_a": ["a1", "a2"]}, - {"node_b": "b"}, - ] - - -async def test_multiple_updates() -> None: - class State(TypedDict): - foo: Annotated[str, operator.add] - - def node_a(state): - return [Command(update={"foo": "a1"}), Command(update={"foo": "a2"})] - - def node_b(state): - return {"foo": "b"} - - graph = ( - StateGraph(State) - .add_sequence([node_a, node_b]) - .add_edge(START, "node_a") - .compile() - ) - - assert await graph.ainvoke({"foo": ""}) == { - "foo": "a1a2b", - } - - # only streams the last update from node_a - assert [c async for c in graph.astream({"foo": ""}, stream_mode="updates")] == [ - {"node_a": [{"foo": "a1"}, {"foo": "a2"}]}, - {"node_b": {"foo": "b"}}, - ] - - -@NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_falsy_return_from_task(checkpointer_name: str) -> None: - """Test with a falsy return from a task.""" - - @task - async def falsy_task() -> bool: - return False - - async with awith_checkpointer(checkpointer_name) as checkpointer: - - @entrypoint(checkpointer=checkpointer) - async def graph(state: dict) -> dict: - """React tool.""" - await falsy_task() - interrupt("test") - - configurable = {"configurable": {"thread_id": str(uuid.uuid4())}} - await graph.ainvoke({"a": 5}, configurable) - await graph.ainvoke(Command(resume="123"), configurable) - - -@NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_multiple_interrupts_functional(checkpointer_name: str) -> None: - """Test multiple interrupts with functional API.""" - from langgraph.func import entrypoint, task - - counter = 0 - - @task - async def double(x: int) -> int: - """Increment the counter.""" - nonlocal counter - counter += 1 - return 2 * x - - async with awith_checkpointer(checkpointer_name) as checkpointer: - - @entrypoint(checkpointer=checkpointer) - async def graph(state: dict) -> dict: - """React tool.""" - - values = [] - - for idx in [1, 2, 3]: - values.extend([await double(idx), interrupt({"a": "boo"})]) - - return {"values": values} - - configurable = {"configurable": {"thread_id": str(uuid.uuid4())}} - await graph.ainvoke({}, configurable) - await graph.ainvoke(Command(resume="a"), configurable) - await graph.ainvoke(Command(resume="b"), configurable) - result = await graph.ainvoke(Command(resume="c"), configurable) - # `double` value should be cached appropriately when used w/ `interrupt` - assert result == { - "values": [2, "a", 4, "b", 6, "c"], - } - assert counter == 3 - - -@NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_double_interrupt_subgraph(checkpointer_name: str) -> None: - class AgentState(TypedDict): - input: str - - def node_1(state: AgentState): - result = interrupt("interrupt node 1") - return {"input": result} - - def node_2(state: AgentState): - result = interrupt("interrupt node 2") - return {"input": result} - - subgraph_builder = ( - StateGraph(AgentState) - .add_node("node_1", node_1) - .add_node("node_2", node_2) - .add_edge(START, "node_1") - .add_edge("node_1", "node_2") - .add_edge("node_2", END) - ) - - async with awith_checkpointer(checkpointer_name) as checkpointer: - # invoke the sub graph - subgraph = subgraph_builder.compile(checkpointer=checkpointer) - thread = {"configurable": {"thread_id": str(uuid.uuid4())}} - assert [c async for c in subgraph.astream({"input": "test"}, thread)] == [ - { - "__interrupt__": ( - Interrupt( - value="interrupt node 1", - resumable=True, - ns=[AnyStr("node_1:")], - when="during", - ), - ) - }, - ] - # resume from the first interrupt - assert [c async for c in subgraph.astream(Command(resume="123"), thread)] == [ - { - "node_1": {"input": "123"}, - }, - { - "__interrupt__": ( - Interrupt( - value="interrupt node 2", - resumable=True, - ns=[AnyStr("node_2:")], - when="during", - ), - ) - }, - ] - # resume from the second interrupt - assert [c async for c in subgraph.astream(Command(resume="123"), thread)] == [ - { - "node_2": {"input": "123"}, - }, - ] - - subgraph = subgraph_builder.compile() - - def invoke_sub_agent(state: AgentState): - return subgraph.invoke(state) - - parent_agent = ( - StateGraph(AgentState) - .add_node("invoke_sub_agent", invoke_sub_agent) - .add_edge(START, "invoke_sub_agent") - .add_edge("invoke_sub_agent", END) - .compile(checkpointer=checkpointer) - ) - - assert [c async for c in parent_agent.astream({"input": "test"}, thread)] == [ - { - "__interrupt__": ( - Interrupt( - value="interrupt node 1", - resumable=True, - ns=[AnyStr("invoke_sub_agent:"), AnyStr("node_1:")], - when="during", - ), - ) - }, - ] - - # resume from the first interrupt - assert [ - c async for c in parent_agent.astream(Command(resume=True), thread) - ] == [ - { - "__interrupt__": ( - Interrupt( - value="interrupt node 2", - resumable=True, - ns=[AnyStr("invoke_sub_agent:"), AnyStr("node_2:")], - when="during", - ), - ) - } - ] - - # resume from 2nd interrupt - assert [ - c async for c in parent_agent.astream(Command(resume=True), thread) - ] == [ - { - "invoke_sub_agent": {"input": True}, - }, - ] - - -@NEEDS_CONTEXTVARS -async def test_async_streaming_with_functional_api() -> None: - """Test streaming with functional API. - - This test verifies that we're able to stream results as they're being generated - rather than have all the results arrive at once after the graph has completed. - - The time of arrival between the two updates corresponding to the two `slow` tasks - should be greater than the time delay between the two tasks. - """ - - time_delay = 0.01 - - @task() - async def slow() -> dict: - await asyncio.sleep(time_delay) # Simulate a delay of 10 ms - return {"tic": asyncio.get_running_loop().time()} - - @entrypoint() - async def graph(inputs: dict) -> list: - first = await slow() - second = await slow() - return [first, second] - - arrival_times = [] - - async for chunk in graph.astream({}): - if "slow" not in chunk: # We'll just look at the updates from `slow` - continue - arrival_times.append(asyncio.get_running_loop().time()) - - assert len(arrival_times) == 2 - delta = arrival_times[1] - arrival_times[0] - # Delta cannot be less than 10 ms if it is streaming as results are generated. - assert delta > time_delay - - -@NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_multiple_subgraphs(checkpointer_name: str) -> None: - class State(TypedDict): - a: int - b: int - - class Output(TypedDict): - result: int - - async with awith_checkpointer(checkpointer_name) as checkpointer: - # Define the subgraphs - async def add(state): - return {"result": state["a"] + state["b"]} - - add_subgraph = ( - StateGraph(State, output=Output) - .add_node(add) - .add_edge(START, "add") - .compile() - ) - - async def multiply(state): - return {"result": state["a"] * state["b"]} - - multiply_subgraph = ( - StateGraph(State, output=Output) - .add_node(multiply) - .add_edge(START, "multiply") - .compile() - ) - - # Test calling the same subgraph multiple times - async def call_same_subgraph(state): - result = await add_subgraph.ainvoke(state) - another_result = await add_subgraph.ainvoke( - {"a": result["result"], "b": 10} - ) - return another_result - - parent_call_same_subgraph = ( - StateGraph(State, output=Output) - .add_node(call_same_subgraph) - .add_edge(START, "call_same_subgraph") - .compile(checkpointer=checkpointer) - ) - config = {"configurable": {"thread_id": "1"}} - assert await parent_call_same_subgraph.ainvoke({"a": 2, "b": 3}, config) == { - "result": 15 - } - - # Test calling multiple subgraphs - class Output(TypedDict): - add_result: int - multiply_result: int - - async def call_multiple_subgraphs(state): - add_result = await add_subgraph.ainvoke(state) - multiply_result = await multiply_subgraph.ainvoke(state) - return { - "add_result": add_result["result"], - "multiply_result": multiply_result["result"], - } - - parent_call_multiple_subgraphs = ( - StateGraph(State, output=Output) - .add_node(call_multiple_subgraphs) - .add_edge(START, "call_multiple_subgraphs") - .compile(checkpointer=checkpointer) - ) - config = {"configurable": {"thread_id": "2"}} - assert await parent_call_multiple_subgraphs.ainvoke( - {"a": 2, "b": 3}, config - ) == { - "add_result": 5, - "multiply_result": 6, - } - - -@NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_multiple_subgraphs_functional(checkpointer_name: str) -> None: - async with awith_checkpointer(checkpointer_name) as checkpointer: - # Define addition subgraph - @entrypoint() - async def add(inputs): - a, b = inputs - return a + b - - # Define multiplication subgraph using tasks - @task - async def multiply_task(a, b): - return a * b - - @entrypoint() - async def multiply(inputs): - return await multiply_task(*inputs) - - # Test calling the same subgraph multiple times - @task - async def call_same_subgraph(a, b): - result = await add.ainvoke([a, b]) - another_result = await add.ainvoke([result, 10]) - return another_result - - @entrypoint(checkpointer=checkpointer) - async def parent_call_same_subgraph(inputs): - return await call_same_subgraph(*inputs) - - config = {"configurable": {"thread_id": "1"}} - assert await parent_call_same_subgraph.ainvoke([2, 3], config) == 15 - - # Test calling multiple subgraphs - @task - async def call_multiple_subgraphs(a, b): - add_result = await add.ainvoke([a, b]) - multiply_result = await multiply.ainvoke([a, b]) - return [add_result, multiply_result] - - @entrypoint(checkpointer=checkpointer) - async def parent_call_multiple_subgraphs(inputs): - return await call_multiple_subgraphs(*inputs) - - config = {"configurable": {"thread_id": "2"}} - assert await parent_call_multiple_subgraphs.ainvoke([2, 3], config) == [5, 6] - - -@NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_multiple_subgraphs_mixed_entrypoint(checkpointer_name: str) -> None: - """Test calling multiple StateGraph subgraphs from an entrypoint.""" - - class State(TypedDict): - a: int - b: int - - class Output(TypedDict): - result: int - - async with awith_checkpointer(checkpointer_name) as checkpointer: - # Define the subgraphs - async def add(state): - return {"result": state["a"] + state["b"]} - - add_subgraph = ( - StateGraph(State, output=Output) - .add_node(add) - .add_edge(START, "add") - .compile() - ) - - async def multiply(state): - return {"result": state["a"] * state["b"]} - - multiply_subgraph = ( - StateGraph(State, output=Output) - .add_node(multiply) - .add_edge(START, "multiply") - .compile() - ) - - # Test calling the same subgraph multiple times - @task - async def call_same_subgraph(a, b): - result = (await add_subgraph.ainvoke({"a": a, "b": b}))["result"] - another_result = (await add_subgraph.ainvoke({"a": result, "b": 10}))[ - "result" - ] - return another_result - - @entrypoint(checkpointer=checkpointer) - async def parent_call_same_subgraph(inputs): - return await call_same_subgraph(*inputs) - - config = {"configurable": {"thread_id": "1"}} - assert await parent_call_same_subgraph.ainvoke([2, 3], config) == 15 - - # Test calling multiple subgraphs - @task - async def call_multiple_subgraphs(a, b): - add_result = (await add_subgraph.ainvoke({"a": a, "b": b}))["result"] - multiply_result = (await multiply_subgraph.ainvoke({"a": a, "b": b}))[ - "result" - ] - return [add_result, multiply_result] - - @entrypoint(checkpointer=checkpointer) - async def parent_call_multiple_subgraphs(inputs): - return await call_multiple_subgraphs(*inputs) - - config = {"configurable": {"thread_id": "2"}} - assert await parent_call_multiple_subgraphs.ainvoke([2, 3], config) == [5, 6] - - -@NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_multiple_subgraphs_mixed_state_graph( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - """Test calling multiple entrypoint "subgraphs" from a StateGraph.""" - async with awith_checkpointer(checkpointer_name) as checkpointer: - - class State(TypedDict): - a: int - b: int - - class Output(TypedDict): - result: int - - # Define addition subgraph - @entrypoint() - async def add(inputs): - a, b = inputs - return a + b - - # Define multiplication subgraph using tasks - @task - async def multiply_task(a, b): - return a * b - - @entrypoint() - async def multiply(inputs): - return await multiply_task(*inputs) - - # Test calling the same subgraph multiple times - async def call_same_subgraph(state): - result = await add.ainvoke([state["a"], state["b"]]) - another_result = await add.ainvoke([result, 10]) - return {"result": another_result} - - parent_call_same_subgraph = ( - StateGraph(State, output=Output) - .add_node(call_same_subgraph) - .add_edge(START, "call_same_subgraph") - .compile(checkpointer=checkpointer) - ) - config = {"configurable": {"thread_id": "1"}} - assert await parent_call_same_subgraph.ainvoke({"a": 2, "b": 3}, config) == { - "result": 15 - } - - # Test calling multiple subgraphs - class Output(TypedDict): - add_result: int - multiply_result: int - - async def call_multiple_subgraphs(state): - add_result = await add.ainvoke([state["a"], state["b"]]) - multiply_result = await multiply.ainvoke([state["a"], state["b"]]) - return { - "add_result": add_result, - "multiply_result": multiply_result, - } - - parent_call_multiple_subgraphs = ( - StateGraph(State, output=Output) - .add_node(call_multiple_subgraphs) - .add_edge(START, "call_multiple_subgraphs") - .compile(checkpointer=checkpointer) - ) - config = {"configurable": {"thread_id": "2"}} - assert await parent_call_multiple_subgraphs.ainvoke( - {"a": 2, "b": 3}, config - ) == { - "add_result": 5, - "multiply_result": 6, - } - - -@NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_multiple_subgraphs_checkpointer(checkpointer_name: str) -> None: - async with awith_checkpointer(checkpointer_name) as checkpointer: - - class SubgraphState(TypedDict): - sub_counter: Annotated[int, operator.add] - - async def subgraph_node(state): - return {"sub_counter": 2} - - sub_graph_1 = ( - StateGraph(SubgraphState) - .add_node(subgraph_node) - .add_edge(START, "subgraph_node") - .compile(checkpointer=True) - ) - - class OtherSubgraphState(TypedDict): - other_sub_counter: Annotated[int, operator.add] - - async def other_subgraph_node(state): - return {"other_sub_counter": 3} - - sub_graph_2 = ( - StateGraph(OtherSubgraphState) - .add_node(other_subgraph_node) - .add_edge(START, "other_subgraph_node") - .compile() - ) - - class ParentState(TypedDict): - parent_counter: int - - async def parent_node(state): - result = await sub_graph_1.ainvoke({"sub_counter": state["parent_counter"]}) - other_result = await sub_graph_2.ainvoke( - {"other_sub_counter": result["sub_counter"]} - ) - return {"parent_counter": other_result["other_sub_counter"]} - - parent_graph = ( - StateGraph(ParentState) - .add_node(parent_node) - .add_edge(START, "parent_node") - .compile(checkpointer=checkpointer) - ) - - config = {"configurable": {"thread_id": "1"}} - assert await parent_graph.ainvoke({"parent_counter": 0}, config) == { - "parent_counter": 5 - } - assert await parent_graph.ainvoke({"parent_counter": 0}, config) == { - "parent_counter": 7 - } - config = {"configurable": {"thread_id": "2"}} - assert [ - c - async for c in parent_graph.astream( - {"parent_counter": 0}, config, subgraphs=True, stream_mode="updates" - ) - ] == [ - (("parent_node",), {"subgraph_node": {"sub_counter": 2}}), - ( - (AnyStr("parent_node:"), "1"), - {"other_subgraph_node": {"other_sub_counter": 3}}, - ), - ((), {"parent_node": {"parent_counter": 5}}), - ] - assert [ - c - async for c in parent_graph.astream( - {"parent_counter": 0}, config, subgraphs=True, stream_mode="updates" - ) - ] == [ - (("parent_node",), {"subgraph_node": {"sub_counter": 2}}), - ( - (AnyStr("parent_node:"), "1"), - {"other_subgraph_node": {"other_sub_counter": 3}}, - ), - ((), {"parent_node": {"parent_counter": 7}}), - ] - - -@NEEDS_CONTEXTVARS -async def test_async_entrypoint_without_checkpointer() -> None: - """Test no checkpointer.""" - states = [] - config = {"configurable": {"thread_id": "1"}} - - # Test without previous - @entrypoint() - async def foo(inputs: Any) -> Any: - states.append(inputs) - return inputs - - assert (await foo.ainvoke({"a": "1"}, config)) == {"a": "1"} - - @entrypoint() - async def foo(inputs: Any, *, previous: Any) -> Any: - states.append(previous) - return {"previous": previous, "current": inputs} - - assert (await foo.ainvoke({"a": "1"}, config)) == { - "current": {"a": "1"}, - "previous": None, - } - assert (await foo.ainvoke({"a": "1"}, config)) == { - "current": {"a": "1"}, - "previous": None, - } - - -async def test_entrypoint_from_async_generator() -> None: - """@entrypoint does not support sync generators.""" - with pytest.raises(NotImplementedError): - - @entrypoint(checkpointer=MemorySaver()) - async def foo(inputs) -> Any: - yield "a" - yield "b" - - -@NEEDS_CONTEXTVARS -async def test_named_tasks_functional() -> None: - class Foo: - async def foo(self, value: str) -> dict: - return value + "foo" - - f = Foo() - - # class method task - foo = task(f.foo, name="custom_foo") - other_foo = task(f.foo, name="other_foo") - - # regular function task - @task(name="custom_bar") - async def bar(value: str) -> dict: - return value + "|bar" - - async def baz(update: str, value: str) -> dict: - return value + f"|{update}" - - # partial function task (unnamed) - baz_task = task(functools.partial(baz, "baz")) - # partial function task (named_) - custom_baz_task = task(functools.partial(baz, "custom_baz"), name="custom_baz") - - class Qux: - def __call__(self, value: str) -> dict: - return value + "|qux" - - qux_task = task(Qux(), name="qux") - - @entrypoint() - async def workflow(inputs: dict) -> dict: - foo_result = await foo(inputs) - await other_foo(inputs) - bar_result = await bar(foo_result) - baz_result = await baz_task(bar_result) - custom_baz_result = await custom_baz_task(baz_result) - qux_result = await qux_task(custom_baz_result) - return qux_result - - assert [c async for c in workflow.astream("", stream_mode="updates")] == [ - {"custom_foo": "foo"}, - {"other_foo": "foo"}, - {"custom_bar": "foo|bar"}, - {"baz": "foo|bar|baz"}, - {"custom_baz": "foo|bar|baz|custom_baz"}, - {"qux": "foo|bar|baz|custom_baz|qux"}, - {"workflow": "foo|bar|baz|custom_baz|qux"}, - ] - - -@NEEDS_CONTEXTVARS -async def test_overriding_injectable_args_with_async_task() -> None: - """Test overriding injectable args in tasks.""" - from langgraph.store.memory import InMemoryStore - - @task - async def foo(store: BaseStore, writer: StreamWriter, value: Any) -> None: - assert store is value - assert writer is value - - @entrypoint(store=InMemoryStore()) - async def main(inputs, store: BaseStore) -> str: - assert store is not None - await foo(store=None, writer=None, value=None) - await foo(store="hello", writer="hello", value="hello") - return "OK" - - assert await main.ainvoke({}) == "OK" - - -async def test_tags_stream_mode_messages() -> None: - model = GenericFakeChatModel(messages=iter(["foo"]), tags=["meow"]) - - async def call_model(state, config): - return {"messages": await model.ainvoke(state["messages"], config)} - - graph = ( - StateGraph(MessagesState) - .add_node(call_model) - .add_edge(START, "call_model") - .compile() - ) - assert [ - c - async for c in graph.astream( - { - "messages": "hi", - }, - stream_mode="messages", - ) - ] == [ - ( - _AnyIdAIMessageChunk(content="foo"), - { - "langgraph_step": 1, - "langgraph_node": "call_model", - "langgraph_triggers": ["start:call_model"], - "langgraph_path": ("__pregel_pull", "call_model"), - "langgraph_checkpoint_ns": AnyStr("call_model:"), - "checkpoint_ns": AnyStr("call_model:"), - "ls_provider": "genericfakechatmodel", - "ls_model_type": "chat", - "tags": ["meow"], - }, - ) - ] - - -async def test_stream_messages_dedupe_inputs() -> None: - from langchain_core.messages import AIMessage - - async def call_model(state): - return {"messages": AIMessage("hi", id="1")} - - async def route(state): - return Command(goto="node_2", graph=Command.PARENT) - - subgraph = ( - StateGraph(MessagesState) - .add_node(call_model) - .add_node(route) - .add_edge(START, "call_model") - .add_edge("call_model", "route") - .compile() - ) - - graph = ( - StateGraph(MessagesState) - .add_node("node_1", subgraph) - .add_node("node_2", lambda state: state) - .add_edge(START, "node_1") - .compile() - ) - - chunks = [ - chunk - async for ns, chunk in graph.astream( - {"messages": "hi"}, stream_mode="messages", subgraphs=True - ) - ] - - assert len(chunks) == 1 - assert chunks[0][0] == AIMessage("hi", id="1") - assert chunks[0][1]["langgraph_node"] == "call_model" - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_stream_messages_dedupe_state(checkpointer_name: str) -> None: - async with awith_checkpointer(checkpointer_name) as checkpointer: - from langchain_core.messages import AIMessage - - to_emit = [AIMessage("bye", id="1"), AIMessage("bye again", id="2")] - - async def call_model(state): - return {"messages": to_emit.pop(0)} - - async def route(state): - return Command(goto="node_2", graph=Command.PARENT) - - subgraph = ( - StateGraph(MessagesState) - .add_node(call_model) - .add_node(route) - .add_edge(START, "call_model") - .add_edge("call_model", "route") - .compile() - ) - - graph = ( - StateGraph(MessagesState) - .add_node("node_1", subgraph) - .add_node("node_2", lambda state: state) - .add_edge(START, "node_1") - .compile(checkpointer=checkpointer) - ) - - thread1 = {"configurable": {"thread_id": "1"}} - - chunks = [ - chunk - async for ns, chunk in graph.astream( - {"messages": "hi"}, thread1, stream_mode="messages", subgraphs=True - ) - ] - - assert len(chunks) == 1 - assert chunks[0][0] == AIMessage("bye", id="1") - assert chunks[0][1]["langgraph_node"] == "call_model" - - chunks = [ - chunk - async for ns, chunk in graph.astream( - {"messages": "hi again"}, - thread1, - stream_mode="messages", - subgraphs=True, - ) - ] - - assert len(chunks) == 1 - assert chunks[0][0] == AIMessage("bye again", id="2") - assert chunks[0][1]["langgraph_node"] == "call_model" - - -@NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_interrupt_subgraph_reenter_checkpointer_true( - checkpointer_name: str, -) -> None: - async with awith_checkpointer(checkpointer_name) as checkpointer: - - class SubgraphState(TypedDict): - foo: str - bar: str - - class ParentState(TypedDict): - foo: str - counter: int - - called = [] - bar_values = [] - - async def subnode_1(state: SubgraphState): - called.append("subnode_1") - bar_values.append(state.get("bar")) - return {"foo": "subgraph_1"} - - async def subnode_2(state: SubgraphState): - called.append("subnode_2") - value = interrupt("Provide value") - value += "baz" - return {"foo": "subgraph_2", "bar": value} - - subgraph = ( - StateGraph(SubgraphState) - .add_node(subnode_1) - .add_node(subnode_2) - .add_edge(START, "subnode_1") - .add_edge("subnode_1", "subnode_2") - .compile(checkpointer=True) - ) - - async def call_subgraph(state: ParentState): - called.append("call_subgraph") - return await subgraph.ainvoke(state) - - async def node(state: ParentState): - called.append("parent") - if state["counter"] < 1: - return Command( - goto="call_subgraph", update={"counter": state["counter"] + 1} - ) - - return {"foo": state["foo"] + "|" + "parent"} - - parent = ( - StateGraph(ParentState) - .add_node(call_subgraph) - .add_node(node) - .add_edge(START, "call_subgraph") - .add_edge("call_subgraph", "node") - .compile(checkpointer=checkpointer) - ) - - config = {"configurable": {"thread_id": "1"}} - assert await parent.ainvoke({"foo": "", "counter": 0}, config) == { - "foo": "", - "counter": 0, - } - assert await parent.ainvoke(Command(resume="bar"), config) == { - "foo": "subgraph_2", - "counter": 1, - } - assert await parent.ainvoke(Command(resume="qux"), config) == { - "foo": "subgraph_2|parent", - "counter": 1, - } - assert called == [ - "call_subgraph", - "subnode_1", - "subnode_2", - "call_subgraph", - "subnode_2", - "parent", - "call_subgraph", - "subnode_1", - "subnode_2", - "call_subgraph", - "subnode_2", - "parent", - ] - - # invoke parent again (new turn) - assert await parent.ainvoke({"foo": "meow", "counter": 0}, config) == { - "foo": "meow", - "counter": 0, - } - # confirm that we preserve the state values from the previous invocation - assert bar_values == [None, "barbaz", "quxbaz"] diff --git a/libs/langgraph/tests/test_remote_graph.py b/libs/langgraph/tests/test_remote_graph.py deleted file mode 100644 index 70857ed61..000000000 --- a/libs/langgraph/tests/test_remote_graph.py +++ /dev/null @@ -1,798 +0,0 @@ -from unittest.mock import AsyncMock, MagicMock - -import pytest -from langchain_core.runnables.graph import ( - Edge as DrawableEdge, -) -from langchain_core.runnables.graph import ( - Node as DrawableNode, -) -from langgraph_sdk.schema import StreamPart - -from langgraph.errors import GraphInterrupt -from langgraph.pregel.remote import RemoteGraph -from langgraph.pregel.types import StateSnapshot - - -def test_with_config(): - # set up test - remote_pregel = RemoteGraph( - "test_graph_id", - config={ - "configurable": { - "foo": "bar", - "thread_id": "thread_id_1", - } - }, - ) - - # call method / assertions - config = {"configurable": {"hello": "world"}} - remote_pregel_copy = remote_pregel.with_config(config) - - # assert that a copy was returned - assert remote_pregel_copy != remote_pregel - # assert that configs were merged - assert remote_pregel_copy.config == { - "configurable": { - "foo": "bar", - "thread_id": "thread_id_1", - "hello": "world", - } - } - - -def test_get_graph(): - # set up test - mock_sync_client = MagicMock() - mock_sync_client.assistants.get_graph.return_value = { - "nodes": [ - {"id": "__start__", "type": "schema", "data": "__start__"}, - {"id": "__end__", "type": "schema", "data": "__end__"}, - { - "id": "agent", - "type": "runnable", - "data": { - "id": ["langgraph", "utils", "RunnableCallable"], - "name": "agent_1", - }, - }, - ], - "edges": [ - {"source": "__start__", "target": "agent"}, - {"source": "agent", "target": "__end__"}, - ], - } - - remote_pregel = RemoteGraph("test_graph_id", sync_client=mock_sync_client) - - # call method / assertions - drawable_graph = remote_pregel.get_graph() - - assert drawable_graph.nodes == { - "__start__": DrawableNode( - id="__start__", name="__start__", data="__start__", metadata=None - ), - "__end__": DrawableNode( - id="__end__", name="__end__", data="__end__", metadata=None - ), - "agent": DrawableNode( - id="agent", - name="agent_1", - data={"id": ["langgraph", "utils", "RunnableCallable"], "name": "agent_1"}, - metadata=None, - ), - } - - assert drawable_graph.edges == [ - DrawableEdge(source="__start__", target="agent"), - DrawableEdge(source="agent", target="__end__"), - ] - - -@pytest.mark.anyio -async def test_aget_graph(): - # set up test - mock_async_client = AsyncMock() - mock_async_client.assistants.get_graph.return_value = { - "nodes": [ - {"id": "__start__", "type": "schema", "data": "__start__"}, - {"id": "__end__", "type": "schema", "data": "__end__"}, - { - "id": "agent", - "type": "runnable", - "data": { - "id": ["langgraph", "utils", "RunnableCallable"], - "name": "agent_1", - }, - }, - ], - "edges": [ - {"source": "__start__", "target": "agent"}, - {"source": "agent", "target": "__end__"}, - ], - } - - remote_pregel = RemoteGraph("test_graph_id", client=mock_async_client) - - # call method / assertions - drawable_graph = await remote_pregel.aget_graph() - - assert drawable_graph.nodes == { - "__start__": DrawableNode( - id="__start__", name="__start__", data="__start__", metadata=None - ), - "__end__": DrawableNode( - id="__end__", name="__end__", data="__end__", metadata=None - ), - "agent": DrawableNode( - id="agent", - name="agent_1", - data={"id": ["langgraph", "utils", "RunnableCallable"], "name": "agent_1"}, - metadata=None, - ), - } - - assert drawable_graph.edges == [ - DrawableEdge(source="__start__", target="agent"), - DrawableEdge(source="agent", target="__end__"), - ] - - -def test_get_state(): - # set up test - mock_sync_client = MagicMock() - mock_sync_client.threads.get_state.return_value = { - "values": {"messages": [{"type": "human", "content": "hello"}]}, - "next": None, - "checkpoint": { - "thread_id": "thread_1", - "checkpoint_ns": "ns", - "checkpoint_id": "checkpoint_1", - "checkpoint_map": {}, - }, - "metadata": {}, - "created_at": "timestamp", - "parent_checkpoint": None, - "tasks": [], - } - - # call method / assertions - remote_pregel = RemoteGraph( - "test_graph_id", - sync_client=mock_sync_client, - ) - - config = {"configurable": {"thread_id": "thread1"}} - state_snapshot = remote_pregel.get_state(config) - - assert state_snapshot == StateSnapshot( - values={"messages": [{"type": "human", "content": "hello"}]}, - next=(), - config={ - "configurable": { - "thread_id": "thread_1", - "checkpoint_ns": "ns", - "checkpoint_id": "checkpoint_1", - "checkpoint_map": {}, - } - }, - metadata={}, - created_at="timestamp", - parent_config=None, - tasks=(), - ) - - -@pytest.mark.anyio -async def test_aget_state(): - mock_async_client = AsyncMock() - mock_async_client.threads.get_state.return_value = { - "values": {"messages": [{"type": "human", "content": "hello"}]}, - "next": None, - "checkpoint": { - "thread_id": "thread_1", - "checkpoint_ns": "ns", - "checkpoint_id": "checkpoint_2", - "checkpoint_map": {}, - }, - "metadata": {}, - "created_at": "timestamp", - "parent_checkpoint": { - "thread_id": "thread_1", - "checkpoint_ns": "ns", - "checkpoint_id": "checkpoint_1", - "checkpoint_map": {}, - }, - "tasks": [], - } - - # call method / assertions - remote_pregel = RemoteGraph( - "test_graph_id", - client=mock_async_client, - ) - - config = {"configurable": {"thread_id": "thread1"}} - state_snapshot = await remote_pregel.aget_state(config) - - assert state_snapshot == StateSnapshot( - values={"messages": [{"type": "human", "content": "hello"}]}, - next=(), - config={ - "configurable": { - "thread_id": "thread_1", - "checkpoint_ns": "ns", - "checkpoint_id": "checkpoint_2", - "checkpoint_map": {}, - } - }, - metadata={}, - created_at="timestamp", - parent_config={ - "configurable": { - "thread_id": "thread_1", - "checkpoint_ns": "ns", - "checkpoint_id": "checkpoint_1", - "checkpoint_map": {}, - } - }, - tasks=(), - ) - - -def test_get_state_history(): - # set up test - mock_sync_client = MagicMock() - mock_sync_client.threads.get_history.return_value = [ - { - "values": {"messages": [{"type": "human", "content": "hello"}]}, - "next": None, - "checkpoint": { - "thread_id": "thread_1", - "checkpoint_ns": "ns", - "checkpoint_id": "checkpoint_1", - "checkpoint_map": {}, - }, - "metadata": {}, - "created_at": "timestamp", - "parent_checkpoint": None, - "tasks": [], - } - ] - - # call method / assertions - remote_pregel = RemoteGraph( - "test_graph_id", - sync_client=mock_sync_client, - ) - - config = {"configurable": {"thread_id": "thread1"}} - state_history_snapshot = list( - remote_pregel.get_state_history(config, filter=None, before=None, limit=None) - ) - - assert len(state_history_snapshot) == 1 - assert state_history_snapshot[0] == StateSnapshot( - values={"messages": [{"type": "human", "content": "hello"}]}, - next=(), - config={ - "configurable": { - "thread_id": "thread_1", - "checkpoint_ns": "ns", - "checkpoint_id": "checkpoint_1", - "checkpoint_map": {}, - } - }, - metadata={}, - created_at="timestamp", - parent_config=None, - tasks=(), - ) - - -@pytest.mark.anyio -async def test_aget_state_history(): - # set up test - mock_async_client = AsyncMock() - mock_async_client.threads.get_history.return_value = [ - { - "values": {"messages": [{"type": "human", "content": "hello"}]}, - "next": None, - "checkpoint": { - "thread_id": "thread_1", - "checkpoint_ns": "ns", - "checkpoint_id": "checkpoint_1", - "checkpoint_map": {}, - }, - "metadata": {}, - "created_at": "timestamp", - "parent_checkpoint": None, - "tasks": [], - } - ] - - # call method / assertions - remote_pregel = RemoteGraph( - "test_graph_id", - client=mock_async_client, - ) - - config = {"configurable": {"thread_id": "thread1"}} - state_history_snapshot = [] - async for state_snapshot in remote_pregel.aget_state_history( - config, filter=None, before=None, limit=None - ): - state_history_snapshot.append(state_snapshot) - - assert len(state_history_snapshot) == 1 - assert state_history_snapshot[0] == StateSnapshot( - values={"messages": [{"type": "human", "content": "hello"}]}, - next=(), - config={ - "configurable": { - "thread_id": "thread_1", - "checkpoint_ns": "ns", - "checkpoint_id": "checkpoint_1", - "checkpoint_map": {}, - } - }, - metadata={}, - created_at="timestamp", - parent_config=None, - tasks=(), - ) - - -def test_update_state(): - # set up test - mock_sync_client = MagicMock() - mock_sync_client.threads.update_state.return_value = { - "checkpoint": { - "thread_id": "thread_1", - "checkpoint_ns": "ns", - "checkpoint_id": "checkpoint_1", - "checkpoint_map": {}, - } - } - - # call method / assertions - remote_pregel = RemoteGraph( - "test_graph_id", - sync_client=mock_sync_client, - ) - - config = {"configurable": {"thread_id": "thread1"}} - response = remote_pregel.update_state(config, {"key": "value"}) - - assert response == { - "configurable": { - "thread_id": "thread_1", - "checkpoint_ns": "ns", - "checkpoint_id": "checkpoint_1", - "checkpoint_map": {}, - } - } - - -@pytest.mark.anyio -async def test_aupdate_state(): - # set up test - mock_async_client = AsyncMock() - mock_async_client.threads.update_state.return_value = { - "checkpoint": { - "thread_id": "thread_1", - "checkpoint_ns": "ns", - "checkpoint_id": "checkpoint_1", - "checkpoint_map": {}, - } - } - - # call method / assertions - remote_pregel = RemoteGraph( - "test_graph_id", - client=mock_async_client, - ) - - config = {"configurable": {"thread_id": "thread1"}} - response = await remote_pregel.aupdate_state(config, {"key": "value"}) - - assert response == { - "configurable": { - "thread_id": "thread_1", - "checkpoint_ns": "ns", - "checkpoint_id": "checkpoint_1", - "checkpoint_map": {}, - } - } - - -def test_stream(): - # set up test - mock_sync_client = MagicMock() - mock_sync_client.runs.stream.return_value = [ - StreamPart(event="values", data={"chunk": "data1"}), - StreamPart(event="values", data={"chunk": "data2"}), - StreamPart(event="values", data={"chunk": "data3"}), - StreamPart(event="updates", data={"chunk": "data4"}), - StreamPart(event="updates", data={"__interrupt__": ()}), - ] - - # call method / assertions - remote_pregel = RemoteGraph( - "test_graph_id", - sync_client=mock_sync_client, - ) - - # stream modes doesn't include 'updates' - stream_parts = [] - with pytest.raises(GraphInterrupt): - for stream_part in remote_pregel.stream( - {"input": "data"}, - config={"configurable": {"thread_id": "thread_1"}}, - stream_mode="values", - ): - stream_parts.append(stream_part) - - assert stream_parts == [ - {"chunk": "data1"}, - {"chunk": "data2"}, - {"chunk": "data3"}, - ] - - mock_sync_client.runs.stream.return_value = [ - StreamPart(event="updates", data={"chunk": "data3"}), - StreamPart(event="updates", data={"chunk": "data4"}), - StreamPart(event="updates", data={"__interrupt__": ()}), - ] - - # default stream_mode is updates - stream_parts = [] - with pytest.raises(GraphInterrupt): - for stream_part in remote_pregel.stream( - {"input": "data"}, - config={"configurable": {"thread_id": "thread_1"}}, - ): - stream_parts.append(stream_part) - - assert stream_parts == [ - {"chunk": "data3"}, - {"chunk": "data4"}, - ] - - # list stream_mode includes mode names - stream_parts = [] - with pytest.raises(GraphInterrupt): - for stream_part in remote_pregel.stream( - {"input": "data"}, - config={"configurable": {"thread_id": "thread_1"}}, - stream_mode=["updates"], - ): - stream_parts.append(stream_part) - - assert stream_parts == [ - ("updates", {"chunk": "data3"}), - ("updates", {"chunk": "data4"}), - ] - - # subgraphs + list modes - stream_parts = [] - with pytest.raises(GraphInterrupt): - for stream_part in remote_pregel.stream( - {"input": "data"}, - config={"configurable": {"thread_id": "thread_1"}}, - stream_mode=["updates"], - subgraphs=True, - ): - stream_parts.append(stream_part) - - assert stream_parts == [ - ((), "updates", {"chunk": "data3"}), - ((), "updates", {"chunk": "data4"}), - ] - - # subgraphs + single mode - stream_parts = [] - with pytest.raises(GraphInterrupt): - for stream_part in remote_pregel.stream( - {"input": "data"}, - config={"configurable": {"thread_id": "thread_1"}}, - subgraphs=True, - ): - stream_parts.append(stream_part) - - assert stream_parts == [ - ((), {"chunk": "data3"}), - ((), {"chunk": "data4"}), - ] - - -@pytest.mark.anyio -async def test_astream(): - # set up test - mock_async_client = MagicMock() - async_iter = MagicMock() - async_iter.__aiter__.return_value = [ - StreamPart(event="values", data={"chunk": "data1"}), - StreamPart(event="values", data={"chunk": "data2"}), - StreamPart(event="values", data={"chunk": "data3"}), - StreamPart(event="updates", data={"chunk": "data4"}), - StreamPart(event="updates", data={"__interrupt__": ()}), - ] - mock_async_client.runs.stream.return_value = async_iter - - # call method / assertions - remote_pregel = RemoteGraph( - "test_graph_id", - client=mock_async_client, - ) - - # stream modes doesn't include 'updates' - stream_parts = [] - with pytest.raises(GraphInterrupt): - async for stream_part in remote_pregel.astream( - {"input": "data"}, - config={"configurable": {"thread_id": "thread_1"}}, - stream_mode="values", - ): - stream_parts.append(stream_part) - - assert stream_parts == [ - {"chunk": "data1"}, - {"chunk": "data2"}, - {"chunk": "data3"}, - ] - - async_iter = MagicMock() - async_iter.__aiter__.return_value = [ - StreamPart(event="updates", data={"chunk": "data3"}), - StreamPart(event="updates", data={"chunk": "data4"}), - StreamPart(event="updates", data={"__interrupt__": ()}), - ] - mock_async_client.runs.stream.return_value = async_iter - - # default stream_mode is updates - stream_parts = [] - with pytest.raises(GraphInterrupt): - async for stream_part in remote_pregel.astream( - {"input": "data"}, - config={"configurable": {"thread_id": "thread_1"}}, - ): - stream_parts.append(stream_part) - - assert stream_parts == [ - {"chunk": "data3"}, - {"chunk": "data4"}, - ] - - # list stream_mode includes mode names - stream_parts = [] - with pytest.raises(GraphInterrupt): - async for stream_part in remote_pregel.astream( - {"input": "data"}, - config={"configurable": {"thread_id": "thread_1"}}, - stream_mode=["updates"], - ): - stream_parts.append(stream_part) - - assert stream_parts == [ - ("updates", {"chunk": "data3"}), - ("updates", {"chunk": "data4"}), - ] - - # subgraphs + list modes - stream_parts = [] - with pytest.raises(GraphInterrupt): - async for stream_part in remote_pregel.astream( - {"input": "data"}, - config={"configurable": {"thread_id": "thread_1"}}, - stream_mode=["updates"], - subgraphs=True, - ): - stream_parts.append(stream_part) - - assert stream_parts == [ - ((), "updates", {"chunk": "data3"}), - ((), "updates", {"chunk": "data4"}), - ] - - # subgraphs + single mode - stream_parts = [] - with pytest.raises(GraphInterrupt): - async for stream_part in remote_pregel.astream( - {"input": "data"}, - config={"configurable": {"thread_id": "thread_1"}}, - subgraphs=True, - ): - stream_parts.append(stream_part) - - assert stream_parts == [ - ((), {"chunk": "data3"}), - ((), {"chunk": "data4"}), - ] - - async_iter = MagicMock() - async_iter.__aiter__.return_value = [ - StreamPart(event="updates|my|subgraph", data={"chunk": "data3"}), - StreamPart(event="updates|hello|subgraph", data={"chunk": "data4"}), - StreamPart(event="updates|bye|subgraph", data={"__interrupt__": ()}), - ] - mock_async_client.runs.stream.return_value = async_iter - - # subgraphs + list modes - stream_parts = [] - with pytest.raises(GraphInterrupt): - async for stream_part in remote_pregel.astream( - {"input": "data"}, - config={"configurable": {"thread_id": "thread_1"}}, - stream_mode=["updates"], - subgraphs=True, - ): - stream_parts.append(stream_part) - - assert stream_parts == [ - (("my", "subgraph"), "updates", {"chunk": "data3"}), - (("hello", "subgraph"), "updates", {"chunk": "data4"}), - ] - - # subgraphs + single mode - stream_parts = [] - with pytest.raises(GraphInterrupt): - async for stream_part in remote_pregel.astream( - {"input": "data"}, - config={"configurable": {"thread_id": "thread_1"}}, - subgraphs=True, - ): - stream_parts.append(stream_part) - - assert stream_parts == [ - (("my", "subgraph"), {"chunk": "data3"}), - (("hello", "subgraph"), {"chunk": "data4"}), - ] - - -def test_invoke(): - # set up test - mock_sync_client = MagicMock() - mock_sync_client.runs.stream.return_value = [ - StreamPart(event="values", data={"chunk": "data1"}), - StreamPart(event="values", data={"chunk": "data2"}), - StreamPart( - event="values", data={"messages": [{"type": "human", "content": "world"}]} - ), - ] - - # call method / assertions - remote_pregel = RemoteGraph( - "test_graph_id", - sync_client=mock_sync_client, - ) - - config = {"configurable": {"thread_id": "thread_1"}} - result = remote_pregel.invoke( - {"input": {"messages": [{"type": "human", "content": "hello"}]}}, config - ) - - assert result == {"messages": [{"type": "human", "content": "world"}]} - - -@pytest.mark.anyio -async def test_ainvoke(): - # set up test - mock_async_client = MagicMock() - async_iter = MagicMock() - async_iter.__aiter__.return_value = [ - StreamPart(event="values", data={"chunk": "data1"}), - StreamPart(event="values", data={"chunk": "data2"}), - StreamPart( - event="values", data={"messages": [{"type": "human", "content": "world"}]} - ), - ] - mock_async_client.runs.stream.return_value = async_iter - - # call method / assertions - remote_pregel = RemoteGraph( - "test_graph_id", - client=mock_async_client, - ) - - config = {"configurable": {"thread_id": "thread_1"}} - result = await remote_pregel.ainvoke( - {"input": {"messages": [{"type": "human", "content": "hello"}]}}, config - ) - - assert result == {"messages": [{"type": "human", "content": "world"}]} - - -@pytest.mark.skip("Unskip this test to manually test the LangGraph Cloud integration") -@pytest.mark.anyio -async def test_langgraph_cloud_integration(): - from langgraph_sdk.client import get_client, get_sync_client - - from langgraph.checkpoint.memory import MemorySaver - from langgraph.graph import END, START, MessagesState, StateGraph - - # create RemotePregel instance - client = get_client() - sync_client = get_sync_client() - remote_pregel = RemoteGraph( - "agent", - client=client, - sync_client=sync_client, - ) - - # define graph - workflow = StateGraph(MessagesState) - workflow.add_node("agent", remote_pregel) - workflow.add_edge(START, "agent") - workflow.add_edge("agent", END) - app = workflow.compile(checkpointer=MemorySaver()) - - # test invocation - input = { - "messages": [ - { - "role": "human", - "content": "What's the weather in SF?", - } - ] - } - - # test invoke - response = app.invoke( - input, - config={"configurable": {"thread_id": "39a6104a-34e7-4f83-929c-d9eb163003c9"}}, - interrupt_before=["agent"], - ) - print("response:", response["messages"][-1].content) - - # test stream - async for chunk in app.astream( - input, - config={"configurable": {"thread_id": "2dc3e3e7-39ac-4597-aa57-4404b944e82a"}}, - subgraphs=True, - stream_mode=["debug", "messages"], - ): - print("chunk:", chunk) - - # test stream events - async for chunk in remote_pregel.astream_events( - input, - config={"configurable": {"thread_id": "2dc3e3e7-39ac-4597-aa57-4404b944e82a"}}, - version="v2", - subgraphs=True, - stream_mode=[], - ): - print("chunk:", chunk) - - # test get state - state_snapshot = await remote_pregel.aget_state( - config={"configurable": {"thread_id": "2dc3e3e7-39ac-4597-aa57-4404b944e82a"}}, - subgraphs=True, - ) - print("state snapshot:", state_snapshot) - - # test update state - response = await remote_pregel.aupdate_state( - config={"configurable": {"thread_id": "6645e002-ed50-4022-92a3-d0d186fdf812"}}, - values={ - "messages": [ - { - "role": "ai", - "content": "Hello world again!", - } - ] - }, - ) - print("response:", response) - - # test get history - async for state in remote_pregel.aget_state_history( - config={"configurable": {"thread_id": "2dc3e3e7-39ac-4597-aa57-4404b944e82a"}}, - ): - print("state snapshot:", state) - - # test get graph - remote_pregel.graph_id = "fe096781-5601-53d2-b2f6-0d3403f7e9ca" # must be UUID - graph = await remote_pregel.aget_graph(xray=True) - print("graph:", graph) diff --git a/libs/langgraph/tests/test_runnable.py b/libs/langgraph/tests/test_runnable.py deleted file mode 100644 index 0a81be368..000000000 --- a/libs/langgraph/tests/test_runnable.py +++ /dev/null @@ -1,261 +0,0 @@ -from __future__ import annotations - -from typing import Any, Optional - -import pytest - -from langgraph.store.base import BaseStore -from langgraph.types import StreamWriter -from langgraph.utils.runnable import RunnableCallable - -pytestmark = pytest.mark.anyio - - -def test_runnable_callable_func_accepts(): - def sync_func(x: Any) -> str: - return f"{x}" - - async def async_func(x: Any) -> str: - return f"{x}" - - def func_with_store(x: Any, store: BaseStore) -> str: - return f"{x}" - - def func_with_writer(x: Any, writer: StreamWriter) -> str: - return f"{x}" - - async def afunc_with_store(x: Any, store: BaseStore) -> str: - return f"{x}" - - async def afunc_with_writer(x: Any, writer: StreamWriter) -> str: - return f"{x}" - - runnables = { - "sync": RunnableCallable(sync_func), - "async": RunnableCallable(func=None, afunc=async_func), - "with_store": RunnableCallable(func_with_store), - "with_writer": RunnableCallable(func_with_writer), - "awith_store": RunnableCallable(afunc_with_store), - "awith_writer": RunnableCallable(afunc_with_writer), - } - - expected_store = {"with_store": True, "awith_store": True} - expected_writer = {"with_writer": True, "awith_writer": True} - - for name, runnable in runnables.items(): - if expected_writer.get(name, False): - assert "writer" in runnable.func_accepts - else: - assert "writer" not in runnable.func_accepts - - if expected_store.get(name, False): - assert "store" in runnable.func_accepts - else: - assert "store" not in runnable.func_accepts - - -async def test_runnable_callable_basic(): - def sync_func(x: Any) -> str: - return f"{x}" - - async def async_func(x: Any) -> str: - return f"{x}" - - runnable_sync = RunnableCallable(sync_func) - runnable_async = RunnableCallable(func=None, afunc=async_func) - - result_sync = runnable_sync.invoke("test") - assert result_sync == "test" - - # Test asynchronous ainvoke - result_async = await runnable_async.ainvoke("test") - assert result_async == "test" - - -def test_runnable_callable_injectable_arguments() -> None: - """Test injectable arguments for RunnableCallable. - - This test verifies that injectable arguments like BaseStore work correctly. - It tests: - - Optional store injection - - Required store injection - - Store injection via config - - Store injection override behavior - - Store value injection and validation - """ - - # Test Optional[BaseStore] annotation. - def func_optional_store(inputs: Any, store: Optional[BaseStore]) -> str: - """Test function that accepts an optional store parameter.""" - assert store is None - return "success" - - assert RunnableCallable(func_optional_store).invoke({"x": "1"}) == "success" - - # Test BaseStore annotation - def func_required_store(inputs: Any, store: BaseStore) -> str: - """Test function that requires a store parameter.""" - assert store is None - return "success" - - with pytest.raises(ValueError): - # Should fail b/c store is not Optional and config is not populated with store. - assert RunnableCallable(func_required_store).invoke({}) == "success" - - # Manually provide store - assert RunnableCallable(func_required_store).invoke({}, store=None) == "success" - - # Specify a value for store in the config - assert ( - RunnableCallable(func_required_store).invoke( - {}, config={"configurable": {"__pregel_store": None}} - ) - == "success" - ) - - # Specify a value for store in config, but override with None - assert ( - RunnableCallable(func_optional_store).invoke( - {"x": "1"}, - store=None, - config={"configurable": {"__pregel_store": "foobar"}}, - ) - == "success" - ) - - # Set of tests where we verify that 'foobar' is injected as the store value. - def func_required_store_v2(inputs: Any, store: BaseStore) -> str: - """Test function that requires a store parameter and validates its value. - - The store value is expected to be 'foobar' when injected. - """ - assert store == "foobar" - return "success" - - assert ( - RunnableCallable(func_required_store_v2).invoke( - {}, config={"configurable": {"__pregel_store": "foobar"}} - ) - == "success" - ) - - assert RunnableCallable(func_required_store_v2).invoke( - # And manual override takes precedence. - {}, - store="foobar", - config={"configurable": {"__pregel_store": "barbar"}}, - ) - - -async def test_runnable_callable_injectable_arguments_async() -> None: - """Test injectable arguments for async RunnableCallable. - - This test verifies that injectable arguments like BaseStore work correctly - in the async context. It tests: - - Optional store injection - - Required store injection - - Store injection via config - - Store injection override behavior - """ - - # Test Optional[BaseStore] annotation. - def func_optional_store(inputs: Any, store: Optional[BaseStore]) -> str: - """Test function that accepts an optional store parameter.""" - assert store is None - return "success" - - async def afunc_optional_store(inputs: Any, store: Optional[BaseStore]) -> str: - """Async version of func_optional_store.""" - assert store is None - return "success" - - assert ( - await RunnableCallable( - func=func_optional_store, afunc=afunc_optional_store - ).ainvoke({"x": "1"}) - == "success" - ) - - # Test BaseStore annotation - def func_required_store(inputs: Any, store: BaseStore) -> str: - """Test function that requires a store parameter.""" - assert store is None - return "success" - - async def afunc_required_store(inputs: Any, store: BaseStore) -> str: - """Async version of func_required_store.""" - assert store is None - return "success" - - with pytest.raises(ValueError): - # Should fail b/c store is not Optional and config is not populated with store. - assert ( - await RunnableCallable( - func=func_required_store, afunc=afunc_required_store - ).ainvoke({}) - == "success" - ) - - # Manually provide store - assert ( - await RunnableCallable( - func=func_required_store, afunc=afunc_required_store - ).ainvoke({}, store=None) - == "success" - ) - - # Specify a value for store in the config - assert ( - await RunnableCallable( - func=func_required_store, afunc=afunc_required_store - ).ainvoke({}, config={"configurable": {"__pregel_store": None}}) - == "success" - ) - - # Specify a value for store in config, but override with None - assert ( - await RunnableCallable( - func=func_optional_store, afunc=afunc_optional_store - ).ainvoke( - {"x": "1"}, - store=None, - config={"configurable": {"__pregel_store": "foobar"}}, - ) - == "success" - ) - - # Set of tests where we verify that 'foobar' is injected as the store value. - def func_required_store_v2(inputs: Any, store: BaseStore) -> str: - """Test function that requires a store parameter with specific value. - - The store parameter is expected to be 'foobar' when injected. - """ - assert store == "foobar" - return "success" - - async def afunc_required_store_v2(inputs: Any, store: BaseStore) -> str: - """Async version of func_required_store_v2. - - The store parameter is expected to be 'foobar' when injected. - """ - assert store == "foobar" - return "success" - - assert ( - await RunnableCallable( - func=func_required_store_v2, afunc=afunc_required_store_v2 - ).ainvoke({}, config={"configurable": {"__pregel_store": "foobar"}}) - == "success" - ) - - assert ( - await RunnableCallable( - func=func_required_store_v2, afunc=afunc_required_store_v2 - ).ainvoke( - # And manual override takes precedence. - {}, - store="foobar", - config={"configurable": {"__pregel_store": "barbar"}}, - ) - == "success" - ) diff --git a/libs/langgraph/tests/test_state.py b/libs/langgraph/tests/test_state.py deleted file mode 100644 index 5af406669..000000000 --- a/libs/langgraph/tests/test_state.py +++ /dev/null @@ -1,330 +0,0 @@ -import inspect -import warnings -from dataclasses import dataclass, field -from typing import Annotated as Annotated2 -from typing import Any, Optional - -import pytest -from langchain_core.runnables import RunnableConfig, RunnableLambda -from pydantic.v1 import BaseModel -from typing_extensions import Annotated, NotRequired, Required, TypedDict - -from langgraph.graph.state import StateGraph, _get_node_name, _warn_invalid_state_schema -from langgraph.managed.shared_value import SharedValue - - -class State(BaseModel): - foo: str - bar: int - - -class State2(TypedDict): - foo: str - bar: int - - -@pytest.mark.parametrize( - "schema", - [ - {"foo": "bar"}, - ["hi", lambda x, y: x + y], - State(foo="bar", bar=1), - State2(foo="bar", bar=1), - ], -) -def test_warns_invalid_schema(schema: Any): - with pytest.warns(UserWarning): - _warn_invalid_state_schema(schema) - - -@pytest.mark.parametrize( - "schema", - [ - Annotated[dict, lambda x, y: y], - Annotated2[list, lambda x, y: y], - dict, - State, - State2, - ], -) -def test_doesnt_warn_valid_schema(schema: Any): - # Assert the function does not raise a warning - with warnings.catch_warnings(): - warnings.simplefilter("error") - _warn_invalid_state_schema(schema) - - -def test_state_schema_with_type_hint(): - class InputState(TypedDict): - question: str - - class OutputState(TypedDict): - input_state: InputState - - class FooState(InputState): - foo: str - - def complete_hint(state: InputState) -> OutputState: - return {"input_state": state} - - def miss_first_hint(state, config: RunnableConfig) -> OutputState: - return {"input_state": state} - - def only_return_hint(state, config) -> OutputState: - return {"input_state": state} - - def miss_all_hint(state, config): - return {"input_state": state} - - def pre_foo(_) -> FooState: - return {"foo": "bar"} - - def pre_bar(_) -> FooState: - return {"foo": "bar"} - - class Foo: - def __call__(self, state: FooState) -> OutputState: - assert state.pop("foo") == "bar" - return {"input_state": state} - - class Bar: - def my_node(self, state: FooState) -> OutputState: - assert state.pop("foo") == "bar" - return {"input_state": state} - - graph = StateGraph(InputState, output=OutputState) - actions = [ - complete_hint, - miss_first_hint, - only_return_hint, - miss_all_hint, - pre_foo, - Foo(), - pre_bar, - Bar().my_node, - ] - - for action in actions: - graph.add_node(action) - - def get_name(action) -> str: - return getattr(action, "__name__", action.__class__.__name__) - - graph.set_entry_point(get_name(actions[0])) - for i in range(len(actions) - 1): - graph.add_edge(get_name(actions[i]), get_name(actions[i + 1])) - graph.set_finish_point(get_name(actions[-1])) - - graph = graph.compile() - - input_state = InputState(question="Hello World!") - output_state = OutputState(input_state=input_state) - foo_state = FooState(foo="bar") - for i, c in enumerate(graph.stream(input_state, stream_mode="updates")): - node_name = get_name(actions[i]) - if node_name in {"pre_foo", "pre_bar"}: - assert c[node_name] == foo_state - else: - assert c[node_name] == output_state - - -@pytest.mark.parametrize("total_", [True, False]) -def test_state_schema_optional_values(total_: bool): - class SomeParentState(TypedDict): - val0a: str - val0b: Optional[str] - - class InputState(SomeParentState, total=total_): # type: ignore - val1: str - val2: Optional[str] - val3: Required[str] - val4: NotRequired[dict] - val5: Annotated[Required[str], "foo"] - val6: Annotated[NotRequired[str], "bar"] - - class OutputState(SomeParentState, total=total_): # type: ignore - out_val1: str - out_val2: Optional[str] - out_val3: Required[str] - out_val4: NotRequired[dict] - out_val5: Annotated[Required[str], "foo"] - out_val6: Annotated[NotRequired[str], "bar"] - - class State(InputState): # this would be ignored - val4: dict - some_shared_channel: Annotated[str, SharedValue.on("assistant_id")] = field( - default="foo" - ) - - builder = StateGraph(State, input=InputState, output=OutputState) - builder.add_node("n", lambda x: x) - builder.add_edge("__start__", "n") - graph = builder.compile() - json_schema = graph.get_input_jsonschema() - - if total_ is False: - expected_required = set() - expected_optional = {"val2", "val1"} - else: - expected_required = {"val1"} - - expected_optional = {"val2"} - - # The others should always have precedence based on the required annotation - expected_required |= {"val0a", "val3", "val5"} - expected_optional |= {"val0b", "val4", "val6"} - - assert set(json_schema.get("required", set())) == expected_required - assert ( - set(json_schema["properties"].keys()) == expected_required | expected_optional - ) - - # Check output schema. Should be the same process - output_schema = graph.get_output_jsonschema() - if total_ is False: - expected_required = set() - expected_optional = {"out_val2", "out_val1"} - else: - expected_required = {"out_val1"} - expected_optional = {"out_val2"} - - expected_required |= {"val0a", "out_val3", "out_val5"} - expected_optional |= {"val0b", "out_val4", "out_val6"} - - assert set(output_schema.get("required", set())) == expected_required - assert ( - set(output_schema["properties"].keys()) == expected_required | expected_optional - ) - - -@pytest.mark.parametrize("kw_only_", [False, True]) -def test_state_schema_default_values(kw_only_: bool): - kwargs = {} - if "kw_only" in inspect.signature(dataclass).parameters: - kwargs = {"kw_only": kw_only_} - - @dataclass(**kwargs) - class InputState: - val1: str - val2: Optional[int] - val3: Annotated[Optional[float], "optional annotated"] - val4: Optional[str] = None - val5: list[int] = field(default_factory=lambda: [1, 2, 3]) - val6: dict[str, int] = field(default_factory=lambda: {"a": 1}) - val7: str = field(default=...) - val8: Annotated[int, "some metadata"] = 42 - val9: Annotated[str, "more metadata"] = field(default="some foo") - val10: str = "default" - val11: Annotated[list[str], "annotated list"] = field( - default_factory=lambda: ["a", "b"] - ) - some_shared_channel: Annotated[str, SharedValue.on("assistant_id")] = field( - default="foo" - ) - - builder = StateGraph(InputState) - builder.add_node("n", lambda x: x) - builder.add_edge("__start__", "n") - graph = builder.compile() - for json_schema in [graph.get_input_jsonschema(), graph.get_output_jsonschema()]: - expected_required = {"val1", "val7"} - expected_optional = { - "val2", - "val3", - "val4", - "val5", - "val6", - "val8", - "val9", - "val10", - "val11", - } - - assert set(json_schema.get("required", set())) == expected_required - assert ( - set(json_schema["properties"].keys()) == expected_required | expected_optional - ) - - -def test_raises_invalid_managed(): - class BadInputState(TypedDict): - some_thing: str - some_input_channel: Annotated[str, SharedValue.on("assistant_id")] - - class InputState(TypedDict): - some_thing: str - some_input_channel: str - - class BadOutputState(TypedDict): - some_thing: str - some_output_channel: Annotated[str, SharedValue.on("assistant_id")] - - class OutputState(TypedDict): - some_thing: str - some_output_channel: str - - class State(TypedDict): - some_thing: str - some_channel: Annotated[str, SharedValue.on("assistant_id")] - - # All OK - StateGraph(State, input=InputState, output=OutputState) - StateGraph(State) - StateGraph(State, input=State, output=State) - StateGraph(State, input=InputState) - StateGraph(State, input=InputState) - - bad_input_examples = [ - (State, BadInputState, OutputState), - (State, BadInputState, BadOutputState), - (State, BadInputState, State), - (State, BadInputState, None), - ] - for _state, _inp, _outp in bad_input_examples: - with pytest.raises( - ValueError, - match="Invalid managed channels detected in BadInputState: some_input_channel. Managed channels are not permitted in Input/Output schema.", - ): - StateGraph(_state, input=_inp, output=_outp) - bad_output_examples = [ - (State, InputState, BadOutputState), - (State, None, BadOutputState), - ] - for _state, _inp, _outp in bad_output_examples: - with pytest.raises( - ValueError, - match="Invalid managed channels detected in BadOutputState: some_output_channel. Managed channels are not permitted in Input/Output schema.", - ): - StateGraph(_state, input=_inp, output=_outp) - - -def test__get_node_name() -> None: - # default runnable name - assert _get_node_name(RunnableLambda(func=lambda x: x)) == "RunnableLambda" - # custom runnable name - assert ( - _get_node_name(RunnableLambda(name="my_runnable", func=lambda x: x)) - == "my_runnable" - ) - - # lambda - assert _get_node_name(lambda x: x) == "" - - # regular function - def func(state): - return - - assert _get_node_name(func) == "func" - - class MyClass: - def __call__(self, state): - return - - def class_method(self, state): - return - - # callable class - assert _get_node_name(MyClass()) == "MyClass" - - # class method - assert _get_node_name(MyClass().class_method) == "class_method" diff --git a/libs/langgraph/tests/test_utils.py b/libs/langgraph/tests/test_utils.py index 3549bd574..982f6bd5e 100644 --- a/libs/langgraph/tests/test_utils.py +++ b/libs/langgraph/tests/test_utils.py @@ -19,7 +19,7 @@ import pytest from typing_extensions import Annotated, NotRequired, Required, TypedDict from langgraph.graph import END, StateGraph -from langgraph.graph.graph import CompiledGraph +from langgraph.graph.state import CompiledStateGraph from langgraph.utils.config import _is_not_empty from langgraph.utils.fields import ( _is_optional_type, @@ -104,7 +104,7 @@ def test_is_generator() -> None: @pytest.fixture -def rt_graph() -> CompiledGraph: +def rt_graph() -> CompiledStateGraph: class State(TypedDict): foo: int node_run_id: int @@ -121,7 +121,7 @@ def rt_graph() -> CompiledGraph: return graph.compile() -def test_runnable_callable_tracing_nested(rt_graph: CompiledGraph) -> None: +def test_runnable_callable_tracing_nested(rt_graph: CompiledStateGraph) -> None: with patch("langsmith.client.Client", spec=langsmith.Client) as mock_client: with patch("langchain_core.tracers.langchain.get_client") as mock_get_client: mock_get_client.return_value = mock_client @@ -134,7 +134,9 @@ def test_runnable_callable_tracing_nested(rt_graph: CompiledGraph) -> None: sys.version_info < (3, 11), reason="Python 3.11+ is required for async contextvars support", ) -async def test_runnable_callable_tracing_nested_async(rt_graph: CompiledGraph) -> None: +async def test_runnable_callable_tracing_nested_async( + rt_graph: CompiledStateGraph, +) -> None: with patch("langsmith.client.Client", spec=langsmith.Client) as mock_client: with patch("langchain_core.tracers.langchain.get_client") as mock_get_client: mock_get_client.return_value = mock_client