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

Remove features and dependencies

Browse Source 
- 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
This commit is contained in:
Nuno Campos
2025-03-01 13:53:27 -08:00
parent 9284b57ba0
commit eb57c06896
67 changed files with 1301 additions and 33581 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py
+14 -15
View File
@@ -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"]
libs/checkpoint-postgres/langgraph/checkpoint/postgres/_ainternal.py
-24
View File
@@ -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)}")
libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py
-485
View File
@@ -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"]
libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py
+3 -3
View File
@@ -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.
libs/checkpoint-postgres/langgraph/checkpoint/postgres/shallow.py
-928
View File
@@ -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()
libs/checkpoint-postgres/langgraph/store/postgres/__init__.py
+1 -2
View File
@@ -1,4 +1,3 @@
from langgraph.store.postgres.aio import AsyncPostgresStore
from langgraph.store.postgres.base import PostgresStore
__all__ = ["AsyncPostgresStore", "PostgresStore"]
__all__ = ["PostgresStore"]
libs/checkpoint-postgres/langgraph/store/postgres/aio.py
-434
View File
@@ -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
libs/checkpoint-postgres/langgraph/store/postgres/base.py
+2 -6
View File
@@ -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):
libs/checkpoint/langgraph/checkpoint/base/__init__.py
+37 -152
View File
@@ -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()
libs/checkpoint/langgraph/checkpoint/memory/__init__.py
+18 -111
View File
@@ -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")
libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py
+3 -9
View File
@@ -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(
libs/checkpoint/langgraph/store/base/__init__.py
+2 -6
View File
@@ -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"
libs/checkpoint/langgraph/store/base/embed.py
+7 -73
View File
@@ -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",
]
libs/checkpoint/langgraph/store/memory/__init__.py
+1 -2
View File
@@ -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,
libs/checkpoint/tests/embed_test_utils.py
+1 -1
View File
@@ -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):
libs/langgraph/Makefile
+1 -1
View File
@@ -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
libs/langgraph/langgraph/channels/__init__.py
-2
View File
@@ -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",
libs/langgraph/langgraph/channels/context.py
-5
View File
@@ -1,5 +0,0 @@
from langgraph.managed.context import Context as ContextManagedValue
Context = ContextManagedValue.of
__all__ = ["Context"]
libs/langgraph/langgraph/config.py
+1 -56
View File
@@ -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!'}
```
libs/langgraph/langgraph/constants.py
+1 -1
View File
@@ -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")
libs/langgraph/langgraph/func/__init__.py
-443
View File
@@ -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,
)
libs/langgraph/langgraph/func/py.typed
View File
libs/langgraph/langgraph/graph/__init__.py
+2 -5
View File
@@ -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",
]
libs/langgraph/langgraph/graph/graph.py
-642
View File
@@ -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(),
}
libs/langgraph/langgraph/graph/message.py
+36 -128
View File
@@ -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]
libs/langgraph/langgraph/graph/state.py
+327 -204
View File
@@ -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<State>.
@@ -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)
},
)
libs/langgraph/langgraph/managed/__init__.py
-3
View File
@@ -1,3 +0,0 @@
from langgraph.managed.is_last_step import IsLastStep, RemainingSteps
__all__ = ["IsLastStep", "RemainingSteps"]
libs/langgraph/langgraph/managed/base.py
-104
View File
@@ -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]
libs/langgraph/langgraph/managed/context.py
-108
View File
@@ -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
libs/langgraph/langgraph/managed/is_last_step.py
-19
View File
@@ -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]
libs/langgraph/langgraph/managed/py.typed
View File
libs/langgraph/langgraph/managed/shared_value.py
-123
View File
@@ -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))
libs/langgraph/langgraph/pregel/__init__.py
+248 -1220
View File
File diff suppressed because it is too large Load Diff
libs/langgraph/langgraph/pregel/algo.py
+11 -172
View File
@@ -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):
libs/langgraph/langgraph/pregel/call.py
-233
View File
@@ -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 == "<locals>":
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
libs/langgraph/langgraph/pregel/debug.py
+9 -54
View File
@@ -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]],
libs/langgraph/langgraph/pregel/executor.py
+63 -100
View File
@@ -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)
libs/langgraph/langgraph/pregel/io.py
+2 -20
View File
@@ -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")
libs/langgraph/langgraph/pregel/loop.py
+13 -285
View File
@@ -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
libs/langgraph/langgraph/pregel/manager.py
+9 -92
View File
@@ -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()
}
libs/langgraph/langgraph/pregel/messages.py
-185
View File
@@ -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)
libs/langgraph/langgraph/pregel/protocol.py
+11 -75
View File
@@ -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,
libs/langgraph/langgraph/pregel/read.py
+15 -121
View File
@@ -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
libs/langgraph/langgraph/pregel/remote.py
-841
View File
@@ -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
libs/langgraph/langgraph/pregel/retry.py
+1 -93
View File
@@ -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})
libs/langgraph/langgraph/pregel/runner.py
+17 -378
View File
@@ -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
libs/langgraph/langgraph/pregel/utils.py
+2 -18
View File
@@ -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
libs/langgraph/langgraph/pregel/write.py
+5 -66
View File
@@ -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
libs/langgraph/langgraph/types.py
+11 -14
View File
@@ -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
libs/langgraph/langgraph/utils/config.py
+152 -116
View File
@@ -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(),
libs/langgraph/langgraph/utils/pydantic.py
-37
View File
@@ -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 {}))
libs/langgraph/langgraph/utils/queue.py
+1 -37
View File
@@ -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"]
libs/langgraph/langgraph/utils/runnable.py
+73 -543
View File
@@ -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)
libs/langgraph/tests/__snapshots__/test_pregel.ambr
-1852
View File
File diff suppressed because it is too large Load Diff
libs/langgraph/tests/conftest.py
+5 -283
View File
@@ -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",
]
libs/langgraph/tests/messages.py
+2 -1
View File
@@ -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
libs/langgraph/tests/test_algo.py
+1 -3
View File
@@ -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,
libs/langgraph/tests/test_channels.py
-2
View File
@@ -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)
libs/langgraph/tests/test_interruption.py
-42
View File
@@ -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 == ()
libs/langgraph/tests/test_large_cases.py
+4 -3396
View File
File diff suppressed because it is too large Load Diff
libs/langgraph/tests/test_large_cases_async.py
-7395
View File
File diff suppressed because it is too large Load Diff
libs/langgraph/tests/test_pregel.py
+183 -2713
View File
File diff suppressed because it is too large Load Diff
libs/langgraph/tests/test_pregel_async.py
-7706
View File
File diff suppressed because it is too large Load Diff
libs/langgraph/tests/test_remote_graph.py
-798
View File
@@ -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)
libs/langgraph/tests/test_runnable.py
-261
View File
@@ -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"
)
libs/langgraph/tests/test_state.py
-330
View File
@@ -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) == "<lambda>"
# 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"
libs/langgraph/tests/test_utils.py
+6 -4
View File
@@ -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