mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-18 05:35:43 +02:00
Compare commits
29
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
0361554fcf | ||
|
|
4332a9515d | ||
|
|
64b99c187a | ||
|
|
afa37d2059 | ||
|
|
d70b659adb | ||
|
|
15f0765d60 | ||
|
|
fe538d4bcb | ||
|
|
4e26a5cf2e | ||
|
|
20f091a277 | ||
|
|
0071bd1e1c | ||
|
|
d36e6ceaaf | ||
|
|
a3feaef2eb | ||
|
|
c6fe26510e | ||
|
|
efbd02a27d | ||
|
|
a91bf116cb | ||
|
|
6a6c3ed84c | ||
|
|
988dd237d2 | ||
|
|
63f5f15c04 | ||
|
|
6fc1c602ab | ||
|
|
2b65308508 | ||
|
|
46dd424a7e | ||
|
|
784821705b | ||
|
|
65172c2a43 | ||
|
|
f4a9d17d24 | ||
|
|
f416480e9d | ||
|
|
61e47cb137 | ||
|
|
d4bbb66963 | ||
|
|
45b5f386e5 | ||
|
|
a1ec55abc5 |
@@ -35,10 +35,10 @@ Create a new app from the `react-agent` template. This template is a simple agen
|
||||
|
||||
## Install Dependencies
|
||||
|
||||
In the root of your new LangGraph app, install the dependencies:
|
||||
In the root of your new LangGraph app, install the dependencies in `edit` mode so your local changes are used by the server:
|
||||
|
||||
```shell
|
||||
pip install .
|
||||
pip install -e .
|
||||
```
|
||||
|
||||
## Create a `.env` file
|
||||
|
||||
@@ -5,7 +5,6 @@ from typing import Any, Optional
|
||||
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from psycopg import Capabilities, Connection, Cursor, Pipeline
|
||||
from psycopg.errors import UndefinedTable
|
||||
from psycopg.rows import DictRow, dict_row
|
||||
from psycopg.types.json import Jsonb
|
||||
from psycopg_pool import ConnectionPool
|
||||
@@ -76,16 +75,15 @@ class PostgresSaver(BasePostgresSaver):
|
||||
the first time checkpointer is used.
|
||||
"""
|
||||
with self._cursor() as cur:
|
||||
try:
|
||||
row = cur.execute(
|
||||
"SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
|
||||
).fetchone()
|
||||
if row is None:
|
||||
version = -1
|
||||
else:
|
||||
version = row["v"]
|
||||
except UndefinedTable:
|
||||
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 :],
|
||||
|
||||
@@ -5,7 +5,6 @@ from typing import Any, Optional
|
||||
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline, Capabilities
|
||||
from psycopg.errors import UndefinedTable
|
||||
from psycopg.rows import DictRow, dict_row
|
||||
from psycopg.types.json import Jsonb
|
||||
from psycopg_pool import AsyncConnectionPool
|
||||
@@ -81,17 +80,15 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
the first time checkpointer is used.
|
||||
"""
|
||||
async with self._cursor() as cur:
|
||||
try:
|
||||
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"]
|
||||
except UndefinedTable:
|
||||
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 :],
|
||||
|
||||
@@ -56,6 +56,7 @@ class Migration(NamedTuple):
|
||||
|
||||
sql: str
|
||||
params: Optional[dict[str, Any]] = None
|
||||
condition: Optional[Callable[["BasePostgresStore"], bool]] = None
|
||||
|
||||
|
||||
MIGRATIONS: Sequence[str] = [
|
||||
@@ -104,11 +105,29 @@ CREATE TABLE IF NOT EXISTS store_vectors (
|
||||
),
|
||||
},
|
||||
),
|
||||
# TODO: Add an HNSW or IVFFlat index depending on config
|
||||
# First must improve the search query when filtering by
|
||||
# namespace
|
||||
Migration(
|
||||
"""
|
||||
CREATE INDEX IF NOT EXISTS store_vectors_embedding_idx ON store_vectors
|
||||
USING %(index_type)s (embedding %(ops)s)%(index_params)s;
|
||||
""",
|
||||
condition=lambda store: bool(
|
||||
store.index_config and _get_index_params(store)[0] != "flat"
|
||||
),
|
||||
params={
|
||||
"index_type": lambda store: _get_index_params(store)[0],
|
||||
"ops": lambda store: _get_vector_type_ops(store),
|
||||
"index_params": lambda store: (
|
||||
" WITH ("
|
||||
+ ", ".join(f"{k}={v}" for k, v in _get_index_params(store)[1].items())
|
||||
+ ")"
|
||||
if _get_index_params(store)[1]
|
||||
else ""
|
||||
),
|
||||
},
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
C = TypeVar("C", bound=Union[_pg_internal.Conn, _ainternal.Conn])
|
||||
|
||||
|
||||
@@ -140,6 +159,8 @@ class PoolConfig(TypedDict, total=False):
|
||||
class ANNIndexConfig(TypedDict, total=False):
|
||||
"""Configuration for vector index in PostgreSQL store."""
|
||||
|
||||
kind: Literal["hnsw", "ivfflat", "flat"]
|
||||
"""Type of index to use: 'hnsw' for Hierarchical Navigable Small World, or 'ivfflat' for Inverted File Flat."""
|
||||
vector_type: Literal["vector", "halfvec"]
|
||||
"""Type of vector storage to use.
|
||||
Options:
|
||||
@@ -148,6 +169,35 @@ class ANNIndexConfig(TypedDict, total=False):
|
||||
"""
|
||||
|
||||
|
||||
class HNSWConfig(ANNIndexConfig, total=False):
|
||||
"""Configuration for HNSW (Hierarchical Navigable Small World) index."""
|
||||
|
||||
kind: Literal["hnsw"] # type: ignore[misc]
|
||||
m: int
|
||||
"""Maximum number of connections per layer. Default is 16."""
|
||||
ef_construction: int
|
||||
"""Size of dynamic candidate list for index construction. Default is 64."""
|
||||
|
||||
|
||||
class IVFFlatConfig(ANNIndexConfig, total=False):
|
||||
"""IVFFlat index divides vectors into lists, and then searches a subset of those lists that are closest to the query vector. It has faster build times and uses less memory than HNSW, but has lower query performance (in terms of speed-recall tradeoff).
|
||||
|
||||
Three keys to achieving good recall are:
|
||||
1. Create the index after the table has some data
|
||||
2. Choose an appropriate number of lists - a good place to start is rows / 1000 for up to 1M rows and sqrt(rows) for over 1M rows
|
||||
3. When querying, specify an appropriate number of probes (higher is better for recall, lower is better for speed) - a good place to start is sqrt(lists)
|
||||
"""
|
||||
|
||||
kind: Literal["ivfflat"] # type: ignore[misc]
|
||||
nlist: int
|
||||
"""Number of inverted lists (clusters) for IVF index.
|
||||
|
||||
Determines the number of clusters used in the index structure.
|
||||
Higher values can improve search speed but increase index size and build time.
|
||||
Typically set to the square root of the number of vectors in the index.
|
||||
"""
|
||||
|
||||
|
||||
class PostgresIndexConfig(IndexConfig, total=False):
|
||||
"""Configuration for vector embeddings in PostgreSQL store with pgvector-specific options.
|
||||
|
||||
@@ -321,7 +371,7 @@ class BasePostgresStore(Generic[C]):
|
||||
if op.query and self.index_config:
|
||||
embedding_requests.append((idx, op.query))
|
||||
|
||||
score_operator = _get_distance_operator(self)
|
||||
score_operator, post_operator = _get_distance_operator(self)
|
||||
vector_type = (
|
||||
cast(PostgresIndexConfig, self.index_config)
|
||||
.get("ann_index_config", {})
|
||||
@@ -351,18 +401,28 @@ class BasePostgresStore(Generic[C]):
|
||||
if not filter_conditions
|
||||
else " AND " + " AND ".join(filter_conditions)
|
||||
)
|
||||
if op.namespace_prefix:
|
||||
prefix_filter_str = f"WHERE s.prefix LIKE %s {filter_str} "
|
||||
ns_args: Sequence = (f"{_namespace_to_text(op.namespace_prefix)}%",)
|
||||
else:
|
||||
ns_args = ()
|
||||
if filter_str:
|
||||
prefix_filter_str = f"WHERE {filter_str} "
|
||||
else:
|
||||
prefix_filter_str = ""
|
||||
|
||||
base_query = f"""
|
||||
WITH scored AS (
|
||||
SELECT s.prefix, s.key, s.value, s.created_at, s.updated_at, {score_operator} AS score
|
||||
SELECT s.prefix, s.key, s.value, s.created_at, s.updated_at, {score_operator} AS neg_score
|
||||
FROM store s
|
||||
JOIN store_vectors sv ON s.prefix = sv.prefix AND s.key = sv.key
|
||||
WHERE s.prefix LIKE %s {filter_str}
|
||||
ORDER BY {score_operator} DESC
|
||||
{prefix_filter_str}
|
||||
ORDER BY {score_operator} ASC
|
||||
LIMIT %s
|
||||
)
|
||||
SELECT * FROM (
|
||||
SELECT DISTINCT ON (prefix, key)
|
||||
prefix, key, value, created_at, updated_at, score
|
||||
prefix, key, value, created_at, updated_at, {post_operator} as score
|
||||
FROM scored
|
||||
ORDER BY prefix, key, score DESC
|
||||
) AS unique_docs
|
||||
@@ -372,7 +432,7 @@ class BasePostgresStore(Generic[C]):
|
||||
"""
|
||||
params = [
|
||||
_PLACEHOLDER, # Vector placeholder
|
||||
f"{_namespace_to_text(op.namespace_prefix)}%",
|
||||
*ns_args,
|
||||
*filter_params,
|
||||
_PLACEHOLDER,
|
||||
expanded_limit,
|
||||
@@ -702,7 +762,6 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
|
||||
_paramslist[i] = embedding
|
||||
|
||||
for (idx, _), (query, params) in zip(search_ops, queries):
|
||||
# Execute the actual query
|
||||
cur.execute(query, params)
|
||||
rows = cast(list[Row], cur.fetchall())
|
||||
results[idx] = [
|
||||
@@ -765,6 +824,8 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
|
||||
for v, migration in enumerate(
|
||||
self.VECTOR_MIGRATIONS[version + 1 :], start=version + 1
|
||||
):
|
||||
if migration.condition and not migration.condition(self):
|
||||
continue
|
||||
sql = migration.sql
|
||||
if migration.params:
|
||||
params = {
|
||||
@@ -823,6 +884,18 @@ def _get_vector_type_ops(store: BasePostgresStore) -> str:
|
||||
return f"{type_prefix}_{distance_suffix}"
|
||||
|
||||
|
||||
def _get_index_params(store: Any) -> tuple[str, dict[str, Any]]:
|
||||
"""Get the index type and configuration based on config."""
|
||||
if not store.index_config:
|
||||
return "hnsw", {}
|
||||
|
||||
config = cast(PostgresIndexConfig, store.index_config)
|
||||
index_config = config.get("ann_index_config", _DEFAULT_ANN_CONFIG).copy()
|
||||
kind = index_config.pop("kind", "hnsw")
|
||||
index_config.pop("vector_type", None)
|
||||
return kind, index_config
|
||||
|
||||
|
||||
def _namespace_to_text(
|
||||
namespace: tuple[str, ...], handle_wildcards: bool = False
|
||||
) -> str:
|
||||
@@ -915,7 +988,7 @@ def _decode_ns_bytes(namespace: Union[str, bytes, list]) -> tuple[str, ...]:
|
||||
return tuple(namespace.split("."))
|
||||
|
||||
|
||||
def _get_distance_operator(store: Any) -> str:
|
||||
def _get_distance_operator(store: Any) -> tuple[str, str]:
|
||||
"""Get the distance operator and score expression based on config."""
|
||||
# Note: Today, we are not using ANN indices due to restrictions
|
||||
# on PGVector's support for mixing vector and non-vector filters
|
||||
@@ -936,12 +1009,22 @@ def _get_distance_operator(store: Any) -> str:
|
||||
config = cast(PostgresIndexConfig, store.index_config)
|
||||
distance_type = config.get("distance_type", "cosine")
|
||||
|
||||
# Return the operator and the score expression
|
||||
# The operator is used in the CTE and will be compatible with an ASCENDING ORDER
|
||||
# sort clause.
|
||||
# The score expression is used in the final query and will be compatible with
|
||||
# a DESCENDING ORDER sort clause and the user's expectations of what the similarity score
|
||||
# should be.
|
||||
if distance_type == "l2":
|
||||
return "1 - (sv.embedding <-> %s::%s)"
|
||||
# Final: "-(sv.embedding <-> %s::%s)"
|
||||
# We return the "l2 similarity" so that the sorting order is the same
|
||||
return "sv.embedding <-> %s::%s", "-scored.neg_score"
|
||||
elif distance_type == "inner_product":
|
||||
return "-(sv.embedding <#> %s::%s)"
|
||||
else: # cosine
|
||||
return "1 - (sv.embedding <=> %s::%s)"
|
||||
# Final: "-(sv.embedding <#> %s::%s)"
|
||||
return "sv.embedding <#> %s::%s", "-(scored.neg_score)"
|
||||
else: # cosine similarity
|
||||
# Final: "1 - (sv.embedding <=> %s::%s)"
|
||||
return "sv.embedding <=> %s::%s", "1 - scored.neg_score"
|
||||
|
||||
|
||||
def _ensure_index_config(
|
||||
|
||||
Generated
+1
-1
@@ -1219,4 +1219,4 @@ watchmedo = ["PyYAML (>=3.10)"]
|
||||
[metadata]
|
||||
lock-version = "2.0"
|
||||
python-versions = "^3.9.0,<4.0"
|
||||
content-hash = "35bd5ff50127337dbbdd1b9937bae6190d8f5c9a6808e334c7f45e5a471773fe"
|
||||
content-hash = "d64fe96797a79103d952c13d4dc0a0296c468b6615b4872aa01cb34479ca4104"
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[tool.poetry]
|
||||
name = "langgraph-checkpoint-postgres"
|
||||
version = "2.0.5"
|
||||
version = "2.0.7"
|
||||
description = "Library with a Postgres implementation of LangGraph checkpoint saver."
|
||||
authors = []
|
||||
license = "MIT"
|
||||
@@ -12,8 +12,8 @@ packages = [{ include = "langgraph" }]
|
||||
python = "^3.9.0,<4.0"
|
||||
langgraph-checkpoint = "^2.0.7"
|
||||
orjson = ">=3.10.1"
|
||||
psycopg = "^3.0.0"
|
||||
psycopg-pool = "^3.0.0"
|
||||
psycopg = "^3.2.0"
|
||||
psycopg-pool = "^3.2.0"
|
||||
|
||||
[tool.poetry.group.dev.dependencies]
|
||||
ruff = "^0.6.2"
|
||||
|
||||
@@ -7,6 +7,7 @@ from psycopg.rows import DictRow, dict_row
|
||||
|
||||
from tests.embed_test_utils import CharacterEmbeddings
|
||||
|
||||
DEFAULT_POSTGRES_URI = "postgres://postgres:postgres@localhost:5441/"
|
||||
DEFAULT_URI = "postgres://postgres:postgres@localhost:5441/postgres?sslmode=disable"
|
||||
|
||||
|
||||
|
||||
@@ -1,7 +1,14 @@
|
||||
# type: ignore
|
||||
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from psycopg import AsyncConnection
|
||||
from psycopg.rows import dict_row
|
||||
from psycopg_pool import AsyncConnectionPool
|
||||
|
||||
from langgraph.checkpoint.base import (
|
||||
Checkpoint,
|
||||
@@ -10,104 +17,212 @@ from langgraph.checkpoint.base import (
|
||||
empty_checkpoint,
|
||||
)
|
||||
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
|
||||
from tests.conftest import DEFAULT_URI
|
||||
from tests.conftest import DEFAULT_POSTGRES_URI
|
||||
|
||||
|
||||
class TestAsyncPostgresSaver:
|
||||
@pytest.fixture(autouse=True)
|
||||
async def setup(self) -> None:
|
||||
# objects for test setup
|
||||
self.config_1: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": "thread-1",
|
||||
# for backwards compatibility testing
|
||||
"thread_ts": "1",
|
||||
"checkpoint_ns": "",
|
||||
}
|
||||
}
|
||||
self.config_2: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": "thread-2",
|
||||
"checkpoint_id": "2",
|
||||
"checkpoint_ns": "",
|
||||
}
|
||||
}
|
||||
self.config_3: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": "thread-2",
|
||||
"checkpoint_id": "2-inner",
|
||||
"checkpoint_ns": "inner",
|
||||
}
|
||||
}
|
||||
@asynccontextmanager
|
||||
async def _pool_saver():
|
||||
"""Fixture for pool mode testing."""
|
||||
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, "row_factory": dict_row},
|
||||
) 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}")
|
||||
|
||||
self.chkpnt_1: Checkpoint = empty_checkpoint()
|
||||
self.chkpnt_2: Checkpoint = create_checkpoint(self.chkpnt_1, {}, 1)
|
||||
self.chkpnt_3: Checkpoint = empty_checkpoint()
|
||||
|
||||
self.metadata_1: CheckpointMetadata = {
|
||||
"source": "input",
|
||||
"step": 2,
|
||||
"writes": {},
|
||||
"score": 1,
|
||||
@asynccontextmanager
|
||||
async def _pipe_saver():
|
||||
"""Fixture for pipeline mode testing."""
|
||||
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:
|
||||
async with await AsyncConnection.connect(
|
||||
DEFAULT_POSTGRES_URI + database,
|
||||
autocommit=True,
|
||||
prepare_threshold=0,
|
||||
row_factory=dict_row,
|
||||
) as conn:
|
||||
async with conn.pipeline() as pipe:
|
||||
checkpointer = AsyncPostgresSaver(conn, pipe=pipe)
|
||||
await checkpointer.setup()
|
||||
async with conn.pipeline() as pipe:
|
||||
checkpointer = AsyncPostgresSaver(conn, 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 _base_saver():
|
||||
"""Fixture for regular connection mode testing."""
|
||||
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:
|
||||
async with await AsyncConnection.connect(
|
||||
DEFAULT_POSTGRES_URI + database,
|
||||
autocommit=True,
|
||||
prepare_threshold=0,
|
||||
row_factory=dict_row,
|
||||
) as conn:
|
||||
checkpointer = AsyncPostgresSaver(conn)
|
||||
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 _saver(name: str):
|
||||
if name == "base":
|
||||
async with _base_saver() as saver:
|
||||
yield saver
|
||||
elif name == "pool":
|
||||
async with _pool_saver() as saver:
|
||||
yield saver
|
||||
elif name == "pipe":
|
||||
async with _pipe_saver() as saver:
|
||||
yield saver
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def test_data():
|
||||
"""Fixture providing test data for checkpoint tests."""
|
||||
config_1: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": "thread-1",
|
||||
# for backwards compatibility testing
|
||||
"thread_ts": "1",
|
||||
"checkpoint_ns": "",
|
||||
}
|
||||
self.metadata_2: CheckpointMetadata = {
|
||||
"source": "loop",
|
||||
}
|
||||
config_2: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": "thread-2",
|
||||
"checkpoint_id": "2",
|
||||
"checkpoint_ns": "",
|
||||
}
|
||||
}
|
||||
config_3: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": "thread-2",
|
||||
"checkpoint_id": "2-inner",
|
||||
"checkpoint_ns": "inner",
|
||||
}
|
||||
}
|
||||
|
||||
chkpnt_1: Checkpoint = empty_checkpoint()
|
||||
chkpnt_2: Checkpoint = create_checkpoint(chkpnt_1, {}, 1)
|
||||
chkpnt_3: Checkpoint = empty_checkpoint()
|
||||
|
||||
metadata_1: CheckpointMetadata = {
|
||||
"source": "input",
|
||||
"step": 2,
|
||||
"writes": {},
|
||||
"score": 1,
|
||||
}
|
||||
metadata_2: CheckpointMetadata = {
|
||||
"source": "loop",
|
||||
"step": 1,
|
||||
"writes": {"foo": "bar"},
|
||||
"score": None,
|
||||
}
|
||||
metadata_3: CheckpointMetadata = {}
|
||||
|
||||
return {
|
||||
"configs": [config_1, config_2, config_3],
|
||||
"checkpoints": [chkpnt_1, chkpnt_2, chkpnt_3],
|
||||
"metadata": [metadata_1, metadata_2, metadata_3],
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
|
||||
async def test_asearch(request, saver_name: str, test_data) -> None:
|
||||
async with _saver(saver_name) as saver:
|
||||
configs = test_data["configs"]
|
||||
checkpoints = test_data["checkpoints"]
|
||||
metadata = test_data["metadata"]
|
||||
|
||||
await saver.aput(configs[0], checkpoints[0], metadata[0], {})
|
||||
await saver.aput(configs[1], checkpoints[1], metadata[1], {})
|
||||
await saver.aput(configs[2], checkpoints[2], metadata[2], {})
|
||||
|
||||
# call method / assertions
|
||||
query_1 = {"source": "input"} # search by 1 key
|
||||
query_2 = {
|
||||
"step": 1,
|
||||
"writes": {"foo": "bar"},
|
||||
"score": None,
|
||||
}
|
||||
self.metadata_3: CheckpointMetadata = {}
|
||||
async with AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver:
|
||||
await saver.setup()
|
||||
} # search by multiple keys
|
||||
query_3: dict[str, Any] = {} # search by no keys, return all checkpoints
|
||||
query_4 = {"source": "update", "step": 1} # no match
|
||||
|
||||
async def test_asearch(self) -> None:
|
||||
async with AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver:
|
||||
await saver.aput(self.config_1, self.chkpnt_1, self.metadata_1, {})
|
||||
await saver.aput(self.config_2, self.chkpnt_2, self.metadata_2, {})
|
||||
await saver.aput(self.config_3, self.chkpnt_3, self.metadata_3, {})
|
||||
search_results_1 = [c async for c in saver.alist(None, filter=query_1)]
|
||||
assert len(search_results_1) == 1
|
||||
assert search_results_1[0].metadata == metadata[0]
|
||||
|
||||
# call method / assertions
|
||||
query_1 = {"source": "input"} # search by 1 key
|
||||
query_2 = {
|
||||
"step": 1,
|
||||
"writes": {"foo": "bar"},
|
||||
} # search by multiple keys
|
||||
query_3: dict[str, Any] = {} # search by no keys, return all checkpoints
|
||||
query_4 = {"source": "update", "step": 1} # no match
|
||||
search_results_2 = [c async for c in saver.alist(None, filter=query_2)]
|
||||
assert len(search_results_2) == 1
|
||||
assert search_results_2[0].metadata == metadata[1]
|
||||
|
||||
search_results_1 = [c async for c in saver.alist(None, filter=query_1)]
|
||||
assert len(search_results_1) == 1
|
||||
assert search_results_1[0].metadata == self.metadata_1
|
||||
search_results_3 = [c async for c in saver.alist(None, filter=query_3)]
|
||||
assert len(search_results_3) == 3
|
||||
|
||||
search_results_2 = [c async for c in saver.alist(None, filter=query_2)]
|
||||
assert len(search_results_2) == 1
|
||||
assert search_results_2[0].metadata == self.metadata_2
|
||||
search_results_4 = [c async for c in saver.alist(None, filter=query_4)]
|
||||
assert len(search_results_4) == 0
|
||||
|
||||
search_results_3 = [c async for c in saver.alist(None, filter=query_3)]
|
||||
assert len(search_results_3) == 3
|
||||
# search by config (defaults to checkpoints across all namespaces)
|
||||
search_results_5 = [
|
||||
c async for c in saver.alist({"configurable": {"thread_id": "thread-2"}})
|
||||
]
|
||||
assert len(search_results_5) == 2
|
||||
assert {
|
||||
search_results_5[0].config["configurable"]["checkpoint_ns"],
|
||||
search_results_5[1].config["configurable"]["checkpoint_ns"],
|
||||
} == {"", "inner"}
|
||||
|
||||
search_results_4 = [c async for c in saver.alist(None, filter=query_4)]
|
||||
assert len(search_results_4) == 0
|
||||
|
||||
# search by config (defaults to checkpoints across all namespaces)
|
||||
search_results_5 = [
|
||||
c
|
||||
async for c in saver.alist({"configurable": {"thread_id": "thread-2"}})
|
||||
]
|
||||
assert len(search_results_5) == 2
|
||||
assert {
|
||||
search_results_5[0].config["configurable"]["checkpoint_ns"],
|
||||
search_results_5[1].config["configurable"]["checkpoint_ns"],
|
||||
} == {"", "inner"}
|
||||
|
||||
# TODO: test before and limit params
|
||||
|
||||
async def test_null_chars(self) -> None:
|
||||
async with AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver:
|
||||
config = await saver.aput(
|
||||
self.config_1, self.chkpnt_1, {"my_key": "\x00abc"}, {}
|
||||
)
|
||||
assert (await saver.aget_tuple(config)).metadata["my_key"] == "abc" # type: ignore
|
||||
assert [c async for c in saver.alist(None, filter={"my_key": "abc"})][
|
||||
0
|
||||
].metadata["my_key"] == "abc"
|
||||
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
|
||||
async def test_null_chars(request, saver_name: str, test_data) -> None:
|
||||
async with _saver(saver_name) as saver:
|
||||
config = await saver.aput(
|
||||
test_data["configs"][0],
|
||||
test_data["checkpoints"][0],
|
||||
{"my_key": "\x00abc"},
|
||||
{},
|
||||
)
|
||||
assert (await saver.aget_tuple(config)).metadata["my_key"] == "abc" # type: ignore
|
||||
assert [c async for c in saver.alist(None, filter={"my_key": "abc"})][
|
||||
0
|
||||
].metadata["my_key"] == "abc"
|
||||
|
||||
@@ -634,6 +634,7 @@ def test_embed_with_path_operation_config(
|
||||
distance_type: str,
|
||||
) -> None:
|
||||
"""Test operation-level field configuration for vector search."""
|
||||
|
||||
with _create_vector_store(
|
||||
vector_type,
|
||||
distance_type,
|
||||
@@ -695,3 +696,89 @@ def test_embed_with_path_operation_config(
|
||||
# assert len(results) == 3
|
||||
# doc5_result = next(r for r in results if r.key == "doc5")
|
||||
# assert doc5_result.score is None
|
||||
|
||||
|
||||
def _cosine_similarity(X: list[float], Y: list[list[float]]) -> list[float]:
|
||||
"""
|
||||
Compute cosine similarity between a vector X and a matrix Y.
|
||||
Lazy import numpy for efficiency.
|
||||
"""
|
||||
|
||||
similarities = []
|
||||
for y in Y:
|
||||
dot_product = sum(a * b for a, b in zip(X, y))
|
||||
norm1 = sum(a * a for a in X) ** 0.5
|
||||
norm2 = sum(a * a for a in y) ** 0.5
|
||||
similarity = dot_product / (norm1 * norm2) if norm1 > 0 and norm2 > 0 else 0.0
|
||||
similarities.append(similarity)
|
||||
|
||||
return similarities
|
||||
|
||||
|
||||
def _inner_product(X: list[float], Y: list[list[float]]) -> list[float]:
|
||||
"""
|
||||
Compute inner product between a vector X and a matrix Y.
|
||||
Lazy import numpy for efficiency.
|
||||
"""
|
||||
|
||||
similarities = []
|
||||
for y in Y:
|
||||
similarity = sum(a * b for a, b in zip(X, y))
|
||||
similarities.append(similarity)
|
||||
|
||||
return similarities
|
||||
|
||||
|
||||
def _neg_l2_distance(X: list[float], Y: list[list[float]]) -> list[float]:
|
||||
"""
|
||||
Compute l2 distance between a vector X and a matrix Y.
|
||||
Lazy import numpy for efficiency.
|
||||
"""
|
||||
|
||||
similarities = []
|
||||
for y in Y:
|
||||
similarity = sum((a - b) ** 2 for a, b in zip(X, y)) ** 0.5
|
||||
similarities.append(-similarity)
|
||||
|
||||
return similarities
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"vector_type,distance_type",
|
||||
[
|
||||
("vector", "cosine"),
|
||||
("vector", "inner_product"),
|
||||
("halfvec", "l2"),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("query", ["aaa", "bbb", "ccc", "abcd", "poisson"])
|
||||
def test_scores(
|
||||
fake_embeddings: CharacterEmbeddings,
|
||||
vector_type: str,
|
||||
distance_type: str,
|
||||
query: str,
|
||||
) -> None:
|
||||
"""Test operation-level field configuration for vector search."""
|
||||
with _create_vector_store(
|
||||
vector_type,
|
||||
distance_type,
|
||||
fake_embeddings,
|
||||
text_fields=["key0"],
|
||||
) as store:
|
||||
doc = {
|
||||
"key0": "aaa",
|
||||
}
|
||||
store.put(("test",), "doc", doc, index=["key0", "key1"])
|
||||
|
||||
results = store.search((), query=query)
|
||||
vec0 = fake_embeddings.embed_query(doc["key0"])
|
||||
vec1 = fake_embeddings.embed_query(query)
|
||||
if distance_type == "cosine":
|
||||
similarities = _cosine_similarity(vec1, [vec0])
|
||||
elif distance_type == "inner_product":
|
||||
similarities = _inner_product(vec1, [vec0])
|
||||
elif distance_type == "l2":
|
||||
similarities = _neg_l2_distance(vec1, [vec0])
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0].score == pytest.approx(similarities[0], abs=1e-3)
|
||||
|
||||
@@ -1,7 +1,14 @@
|
||||
# type: ignore
|
||||
|
||||
from contextlib import contextmanager
|
||||
from typing import Any
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from psycopg import Connection
|
||||
from psycopg.rows import dict_row
|
||||
from psycopg_pool import ConnectionPool
|
||||
|
||||
from langgraph.checkpoint.base import (
|
||||
Checkpoint,
|
||||
@@ -10,103 +17,199 @@ from langgraph.checkpoint.base import (
|
||||
empty_checkpoint,
|
||||
)
|
||||
from langgraph.checkpoint.postgres import PostgresSaver
|
||||
from tests.conftest import DEFAULT_URI
|
||||
from tests.conftest import DEFAULT_POSTGRES_URI
|
||||
|
||||
|
||||
class TestPostgresSaver:
|
||||
@pytest.fixture(autouse=True)
|
||||
def setup(self) -> None:
|
||||
# objects for test setup
|
||||
self.config_1: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": "thread-1",
|
||||
# for backwards compatibility testing
|
||||
"thread_ts": "1",
|
||||
"checkpoint_ns": "",
|
||||
}
|
||||
}
|
||||
self.config_2: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": "thread-2",
|
||||
"checkpoint_id": "2",
|
||||
"checkpoint_ns": "",
|
||||
}
|
||||
}
|
||||
self.config_3: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": "thread-2",
|
||||
"checkpoint_id": "2-inner",
|
||||
"checkpoint_ns": "inner",
|
||||
}
|
||||
}
|
||||
@contextmanager
|
||||
def _pool_saver():
|
||||
"""Fixture for pool mode testing."""
|
||||
database = f"test_{uuid4().hex[:16]}"
|
||||
# create unique db
|
||||
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
|
||||
conn.execute(f"CREATE DATABASE {database}")
|
||||
try:
|
||||
# yield checkpointer
|
||||
with ConnectionPool(
|
||||
DEFAULT_POSTGRES_URI + database,
|
||||
max_size=10,
|
||||
kwargs={"autocommit": True, "row_factory": dict_row},
|
||||
) as pool:
|
||||
checkpointer = PostgresSaver(pool)
|
||||
checkpointer.setup()
|
||||
yield checkpointer
|
||||
finally:
|
||||
# drop unique db
|
||||
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
|
||||
conn.execute(f"DROP DATABASE {database}")
|
||||
|
||||
self.chkpnt_1: Checkpoint = empty_checkpoint()
|
||||
self.chkpnt_2: Checkpoint = create_checkpoint(self.chkpnt_1, {}, 1)
|
||||
self.chkpnt_3: Checkpoint = empty_checkpoint()
|
||||
|
||||
self.metadata_1: CheckpointMetadata = {
|
||||
"source": "input",
|
||||
"step": 2,
|
||||
"writes": {},
|
||||
"score": 1,
|
||||
@contextmanager
|
||||
def _pipe_saver():
|
||||
"""Fixture for pipeline mode testing."""
|
||||
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:
|
||||
with Connection.connect(
|
||||
DEFAULT_POSTGRES_URI + database,
|
||||
autocommit=True,
|
||||
prepare_threshold=0,
|
||||
row_factory=dict_row,
|
||||
) as conn:
|
||||
with conn.pipeline() as pipe:
|
||||
checkpointer = PostgresSaver(conn, pipe=pipe)
|
||||
checkpointer.setup()
|
||||
with conn.pipeline() as pipe:
|
||||
checkpointer = PostgresSaver(conn, pipe=pipe)
|
||||
yield checkpointer
|
||||
finally:
|
||||
# drop unique db
|
||||
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
|
||||
conn.execute(f"DROP DATABASE {database}")
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _base_saver():
|
||||
"""Fixture for regular connection mode testing."""
|
||||
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:
|
||||
with Connection.connect(
|
||||
DEFAULT_POSTGRES_URI + database,
|
||||
autocommit=True,
|
||||
prepare_threshold=0,
|
||||
row_factory=dict_row,
|
||||
) as conn:
|
||||
checkpointer = PostgresSaver(conn)
|
||||
checkpointer.setup()
|
||||
yield checkpointer
|
||||
finally:
|
||||
# drop unique db
|
||||
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
|
||||
conn.execute(f"DROP DATABASE {database}")
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _saver(name: str):
|
||||
if name == "base":
|
||||
with _base_saver() as saver:
|
||||
yield saver
|
||||
elif name == "pool":
|
||||
with _pool_saver() as saver:
|
||||
yield saver
|
||||
elif name == "pipe":
|
||||
with _pipe_saver() as saver:
|
||||
yield saver
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def test_data():
|
||||
"""Fixture providing test data for checkpoint tests."""
|
||||
config_1: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": "thread-1",
|
||||
# for backwards compatibility testing
|
||||
"thread_ts": "1",
|
||||
"checkpoint_ns": "",
|
||||
}
|
||||
self.metadata_2: CheckpointMetadata = {
|
||||
"source": "loop",
|
||||
}
|
||||
config_2: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": "thread-2",
|
||||
"checkpoint_id": "2",
|
||||
"checkpoint_ns": "",
|
||||
}
|
||||
}
|
||||
config_3: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": "thread-2",
|
||||
"checkpoint_id": "2-inner",
|
||||
"checkpoint_ns": "inner",
|
||||
}
|
||||
}
|
||||
|
||||
chkpnt_1: Checkpoint = empty_checkpoint()
|
||||
chkpnt_2: Checkpoint = create_checkpoint(chkpnt_1, {}, 1)
|
||||
chkpnt_3: Checkpoint = empty_checkpoint()
|
||||
|
||||
metadata_1: CheckpointMetadata = {
|
||||
"source": "input",
|
||||
"step": 2,
|
||||
"writes": {},
|
||||
"score": 1,
|
||||
}
|
||||
metadata_2: CheckpointMetadata = {
|
||||
"source": "loop",
|
||||
"step": 1,
|
||||
"writes": {"foo": "bar"},
|
||||
"score": None,
|
||||
}
|
||||
metadata_3: CheckpointMetadata = {}
|
||||
|
||||
return {
|
||||
"configs": [config_1, config_2, config_3],
|
||||
"checkpoints": [chkpnt_1, chkpnt_2, chkpnt_3],
|
||||
"metadata": [metadata_1, metadata_2, metadata_3],
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
|
||||
def test_search(saver_name: str, test_data) -> None:
|
||||
with _saver(saver_name) as saver:
|
||||
configs = test_data["configs"]
|
||||
checkpoints = test_data["checkpoints"]
|
||||
metadata = test_data["metadata"]
|
||||
|
||||
saver.put(configs[0], checkpoints[0], metadata[0], {})
|
||||
saver.put(configs[1], checkpoints[1], metadata[1], {})
|
||||
saver.put(configs[2], checkpoints[2], metadata[2], {})
|
||||
|
||||
# call method / assertions
|
||||
query_1 = {"source": "input"} # search by 1 key
|
||||
query_2 = {
|
||||
"step": 1,
|
||||
"writes": {"foo": "bar"},
|
||||
"score": None,
|
||||
}
|
||||
self.metadata_3: CheckpointMetadata = {}
|
||||
with PostgresSaver.from_conn_string(DEFAULT_URI) as saver:
|
||||
saver.setup()
|
||||
} # search by multiple keys
|
||||
query_3: dict[str, Any] = {} # search by no keys, return all checkpoints
|
||||
query_4 = {"source": "update", "step": 1} # no match
|
||||
|
||||
def test_search(self) -> None:
|
||||
with PostgresSaver.from_conn_string(DEFAULT_URI) as saver:
|
||||
# save checkpoints
|
||||
saver.put(self.config_1, self.chkpnt_1, self.metadata_1, {})
|
||||
saver.put(self.config_2, self.chkpnt_2, self.metadata_2, {})
|
||||
saver.put(self.config_3, self.chkpnt_3, self.metadata_3, {})
|
||||
search_results_1 = list(saver.list(None, filter=query_1))
|
||||
assert len(search_results_1) == 1
|
||||
assert search_results_1[0].metadata == metadata[0]
|
||||
|
||||
# call method / assertions
|
||||
query_1 = {"source": "input"} # search by 1 key
|
||||
query_2 = {
|
||||
"step": 1,
|
||||
"writes": {"foo": "bar"},
|
||||
} # search by multiple keys
|
||||
query_3: dict[str, Any] = {} # search by no keys, return all checkpoints
|
||||
query_4 = {"source": "update", "step": 1} # no match
|
||||
search_results_2 = list(saver.list(None, filter=query_2))
|
||||
assert len(search_results_2) == 1
|
||||
assert search_results_2[0].metadata == metadata[1]
|
||||
|
||||
search_results_1 = list(saver.list(None, filter=query_1))
|
||||
assert len(search_results_1) == 1
|
||||
assert search_results_1[0].metadata == self.metadata_1
|
||||
search_results_3 = list(saver.list(None, filter=query_3))
|
||||
assert len(search_results_3) == 3
|
||||
|
||||
search_results_2 = list(saver.list(None, filter=query_2))
|
||||
assert len(search_results_2) == 1
|
||||
assert search_results_2[0].metadata == self.metadata_2
|
||||
search_results_4 = list(saver.list(None, filter=query_4))
|
||||
assert len(search_results_4) == 0
|
||||
|
||||
search_results_3 = list(saver.list(None, filter=query_3))
|
||||
assert len(search_results_3) == 3
|
||||
# search by config (defaults to checkpoints across all namespaces)
|
||||
search_results_5 = list(saver.list({"configurable": {"thread_id": "thread-2"}}))
|
||||
assert len(search_results_5) == 2
|
||||
assert {
|
||||
search_results_5[0].config["configurable"]["checkpoint_ns"],
|
||||
search_results_5[1].config["configurable"]["checkpoint_ns"],
|
||||
} == {"", "inner"}
|
||||
|
||||
search_results_4 = list(saver.list(None, filter=query_4))
|
||||
assert len(search_results_4) == 0
|
||||
|
||||
# search by config (defaults to checkpoints across all namespaces)
|
||||
search_results_5 = list(
|
||||
saver.list({"configurable": {"thread_id": "thread-2"}})
|
||||
)
|
||||
assert len(search_results_5) == 2
|
||||
assert {
|
||||
search_results_5[0].config["configurable"]["checkpoint_ns"],
|
||||
search_results_5[1].config["configurable"]["checkpoint_ns"],
|
||||
} == {"", "inner"}
|
||||
|
||||
# TODO: test before and limit params
|
||||
|
||||
def test_null_chars(self) -> None:
|
||||
with PostgresSaver.from_conn_string(DEFAULT_URI) as saver:
|
||||
config = saver.put(self.config_1, self.chkpnt_1, {"my_key": "\x00abc"}, {})
|
||||
assert saver.get_tuple(config).metadata["my_key"] == "abc" # type: ignore
|
||||
assert (
|
||||
list(saver.list(None, filter={"my_key": "abc"}))[0].metadata["my_key"] # type: ignore
|
||||
== "abc"
|
||||
)
|
||||
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
|
||||
def test_null_chars(saver_name: str, test_data) -> None:
|
||||
with _saver(saver_name) as saver:
|
||||
config = saver.put(
|
||||
test_data["configs"][0],
|
||||
test_data["checkpoints"][0],
|
||||
{"my_key": "\x00abc"},
|
||||
{},
|
||||
)
|
||||
assert saver.get_tuple(config).metadata["my_key"] == "abc" # type: ignore
|
||||
assert (
|
||||
list(saver.list(None, filter={"my_key": "abc"}))[0].metadata["my_key"]
|
||||
== "abc"
|
||||
)
|
||||
|
||||
@@ -550,6 +550,12 @@ def dockerfile(save_path: str, config: pathlib.Path, add_docker_compose: bool) -
|
||||
type=int,
|
||||
help="Enable remote debugging by listening on specified port. Requires debugpy to be installed",
|
||||
)
|
||||
@click.option(
|
||||
"--wait-for-client",
|
||||
is_flag=True,
|
||||
help="Wait for a debugger client to connect to the debug port before starting the server",
|
||||
default=False,
|
||||
)
|
||||
@cli.command(
|
||||
"dev",
|
||||
help="🏃♀️➡️ Run LangGraph API server in development mode with hot reloading and debugging support",
|
||||
@@ -563,6 +569,7 @@ def dev(
|
||||
n_jobs_per_worker: Optional[int],
|
||||
no_browser: bool,
|
||||
debug_port: Optional[int],
|
||||
wait_for_client: bool,
|
||||
):
|
||||
"""CLI entrypoint for running the LangGraph API server."""
|
||||
try:
|
||||
@@ -595,9 +602,6 @@ def dev(
|
||||
sys.path.append(str(dep_path))
|
||||
|
||||
graphs = config_json.get("graphs", {})
|
||||
additional_config = {}
|
||||
if config_json.get("store"):
|
||||
additional_config["store"] = config_json["store"]
|
||||
|
||||
run_server(
|
||||
host,
|
||||
@@ -607,8 +611,9 @@ def dev(
|
||||
n_jobs_per_worker=n_jobs_per_worker,
|
||||
open_browser=not no_browser,
|
||||
debug_port=debug_port,
|
||||
env=config_json.get("env", None),
|
||||
config=additional_config,
|
||||
env=config_json.get("env"),
|
||||
store=config_json.get("store"),
|
||||
wait_for_client=wait_for_client,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -392,14 +392,12 @@ RUN set -ex && \\
|
||||
],
|
||||
)
|
||||
)
|
||||
additional_config = {}
|
||||
if config.get("store"):
|
||||
additional_config["store"] = config["store"]
|
||||
store_config = config.get("store")
|
||||
env_additional_config = (
|
||||
""
|
||||
if not additional_config
|
||||
if not store_config
|
||||
else f"""
|
||||
ENV LANGGRAPH_CONFIG='{json.dumps(additional_config)}'
|
||||
ENV LANGGRAPH_STORE='{json.dumps(store_config)}'
|
||||
"""
|
||||
)
|
||||
return f"""FROM {base_image}:{config['python_version']}
|
||||
@@ -439,14 +437,12 @@ def node_config_to_docker(config_path: pathlib.Path, config: Config, base_image:
|
||||
install_cmd = "npm ci"
|
||||
else:
|
||||
install_cmd = "npm i"
|
||||
additional_config = {}
|
||||
if config.get("store"):
|
||||
additional_config["store"] = config["store"]
|
||||
store_config = config.get("store")
|
||||
env_additional_config = (
|
||||
""
|
||||
if not additional_config
|
||||
if not store_config
|
||||
else f"""
|
||||
ENV LANGGRAPH_CONFIG='{json.dumps(additional_config)}'
|
||||
ENV LANGGRAPH_STORE='{json.dumps(store_config)}'
|
||||
"""
|
||||
)
|
||||
return f"""FROM {base_image}:{config['node_version']}
|
||||
|
||||
Generated
+4
-4
@@ -565,13 +565,13 @@ langgraph-sdk = ">=0.1.32,<0.2.0"
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-api"
|
||||
version = "0.0.5"
|
||||
version = "0.0.6"
|
||||
description = ""
|
||||
optional = true
|
||||
python-versions = "<4.0,>=3.11.0"
|
||||
files = [
|
||||
{file = "langgraph_api-0.0.5-py3-none-any.whl", hash = "sha256:9c981c489924f5d7e67ce7a3a9908ede15ac266e699deff6f5de691af5b49931"},
|
||||
{file = "langgraph_api-0.0.5.tar.gz", hash = "sha256:f7ff041f1706152a2587916f0373f513e456cde067c94f8266ff29f6af908d20"},
|
||||
{file = "langgraph_api-0.0.6-py3-none-any.whl", hash = "sha256:f64b13959d721143f6a023af5b9ffc9aa054064af98d21d5d8090cda7e7bffd2"},
|
||||
{file = "langgraph_api-0.0.6.tar.gz", hash = "sha256:badac44fa1ec979509e56fc0da57eeb5f278ee5871f27803f73ea6d8822c21b9"},
|
||||
]
|
||||
|
||||
[package.dependencies]
|
||||
@@ -1561,4 +1561,4 @@ inmem = ["langgraph-api", "python-dotenv"]
|
||||
[metadata]
|
||||
lock-version = "2.0"
|
||||
python-versions = "^3.9.0,<4.0"
|
||||
content-hash = "7388b141c48dd6cfa33b504822e5348ac514ac2e2711d37368d7cb7103c8340a"
|
||||
content-hash = "8eaaa66d9e6e447699e3bcee336dfe779b58c956f8c2ad6678008a07be935838"
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[tool.poetry]
|
||||
name = "langgraph-cli"
|
||||
version = "0.1.60"
|
||||
version = "0.1.61"
|
||||
description = "CLI for interacting with LangGraph API"
|
||||
authors = []
|
||||
license = "MIT"
|
||||
@@ -14,7 +14,7 @@ langgraph = "langgraph_cli.cli:cli"
|
||||
[tool.poetry.dependencies]
|
||||
python = "^3.9.0,<4.0"
|
||||
click = "^8.1.7"
|
||||
langgraph-api = { version = ">=0.0.5,<0.1.0", optional = true, python = ">=3.11,<4.0" }
|
||||
langgraph-api = { version = ">=0.0.6,<0.1.0", optional = true, python = ">=3.11,<4.0" }
|
||||
python-dotenv = { version = ">=0.8.0", optional = true }
|
||||
|
||||
[tool.poetry.group.dev.dependencies]
|
||||
|
||||
@@ -374,6 +374,11 @@ class Graph:
|
||||
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():
|
||||
@@ -395,10 +400,6 @@ class Graph:
|
||||
for name, spec in self.nodes.items():
|
||||
if spec.ends:
|
||||
all_targets.update(spec.ends)
|
||||
# validate targets
|
||||
for node in self.nodes:
|
||||
if node not in all_targets:
|
||||
raise ValueError(f"Node `{node}` is not reachable")
|
||||
for target in all_targets:
|
||||
if target not in self.nodes and target != END:
|
||||
raise ValueError(f"Found edge ending at unknown node `{target}`")
|
||||
|
||||
@@ -595,7 +595,7 @@ def prepare_single_task(
|
||||
for tid, c, v in pending_writes
|
||||
if tid in (NULL_TASK_ID, task_id) and c == RESUME
|
||||
),
|
||||
MISSING,
|
||||
configurable.get(CONFIG_KEY_RESUME_VALUE, MISSING),
|
||||
),
|
||||
},
|
||||
),
|
||||
@@ -720,7 +720,7 @@ def prepare_single_task(
|
||||
if tid in (NULL_TASK_ID, task_id)
|
||||
and c == RESUME
|
||||
),
|
||||
MISSING,
|
||||
configurable.get(CONFIG_KEY_RESUME_VALUE, MISSING),
|
||||
),
|
||||
},
|
||||
),
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
from dataclasses import asdict
|
||||
from typing import (
|
||||
Any,
|
||||
AsyncIterator,
|
||||
@@ -27,6 +28,7 @@ from langgraph_sdk.client import (
|
||||
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
|
||||
|
||||
@@ -41,7 +43,7 @@ from langgraph.constants import (
|
||||
from langgraph.errors import GraphInterrupt
|
||||
from langgraph.pregel.protocol import PregelProtocol
|
||||
from langgraph.pregel.types import All, PregelTask, StateSnapshot, StreamMode
|
||||
from langgraph.types import Interrupt, StreamProtocol
|
||||
from langgraph.types import Command, Interrupt, StreamProtocol
|
||||
from langgraph.utils.config import merge_configs
|
||||
|
||||
|
||||
@@ -597,11 +599,17 @@ class RemoteGraph(PregelProtocol):
|
||||
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,
|
||||
@@ -680,11 +688,17 @@ class RemoteGraph(PregelProtocol):
|
||||
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,
|
||||
|
||||
@@ -160,7 +160,7 @@ def test_graph_validation() -> None:
|
||||
workflow = Graph()
|
||||
workflow.add_node("agent", logic)
|
||||
workflow.set_finish_point("agent")
|
||||
with pytest.raises(ValueError, match="not reachable"):
|
||||
with pytest.raises(ValueError, match="must have an entrypoint"):
|
||||
workflow.compile()
|
||||
|
||||
workflow = Graph()
|
||||
@@ -207,18 +207,6 @@ def test_graph_validation() -> None:
|
||||
with pytest.raises(ValueError, match="unknown"): # extra is not defined
|
||||
workflow.compile()
|
||||
|
||||
workflow = Graph()
|
||||
workflow.add_node("agent", logic)
|
||||
workflow.add_node("tools", logic)
|
||||
workflow.add_node("extra", logic)
|
||||
workflow.set_entry_point("agent")
|
||||
workflow.add_conditional_edges("agent", logic, {"continue": "tools", "exit": END})
|
||||
workflow.add_edge("tools", "agent")
|
||||
with pytest.raises(
|
||||
ValueError, match="Node `extra` is not reachable"
|
||||
): # extra is not reachable
|
||||
workflow.compile()
|
||||
|
||||
workflow = Graph()
|
||||
workflow.add_node("agent", logic)
|
||||
workflow.add_node("tools", logic)
|
||||
@@ -276,6 +264,25 @@ def test_graph_validation() -> None:
|
||||
graph.invoke({"hello": "there"})
|
||||
|
||||
|
||||
def test_graph_validation_with_command() -> None:
|
||||
class State(TypedDict):
|
||||
foo: str
|
||||
bar: str
|
||||
|
||||
def node_a(state: State):
|
||||
return GraphCommand(goto="b", update={"foo": "bar"})
|
||||
|
||||
def node_b(state: State):
|
||||
return GraphCommand(goto=END, update={"bar": "baz"})
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("a", node_a)
|
||||
builder.add_node("b", node_b)
|
||||
builder.add_edge(START, "a")
|
||||
graph = builder.compile()
|
||||
assert graph.invoke({"foo": ""}) == {"foo": "bar", "bar": "baz"}
|
||||
|
||||
|
||||
def test_checkpoint_errors() -> None:
|
||||
class FaultyGetCheckpointer(MemorySaver):
|
||||
def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
|
||||
@@ -8728,6 +8735,176 @@ def test_copy_checkpoint(
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
|
||||
def test_dynamic_interrupt_subgraph(
|
||||
request: pytest.FixtureRequest, checkpointer_name: str
|
||||
) -> None:
|
||||
checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
|
||||
|
||||
class SubgraphState(TypedDict):
|
||||
my_key: str
|
||||
market: str
|
||||
|
||||
tool_two_node_count = 0
|
||||
|
||||
def tool_two_node(s: SubgraphState) -> SubgraphState:
|
||||
nonlocal tool_two_node_count
|
||||
tool_two_node_count += 1
|
||||
if s["market"] == "DE":
|
||||
answer = interrupt("Just because...")
|
||||
else:
|
||||
answer = " all good"
|
||||
return {"my_key": answer}
|
||||
|
||||
subgraph = StateGraph(SubgraphState)
|
||||
subgraph.add_node("do", tool_two_node, retry=RetryPolicy())
|
||||
subgraph.add_edge(START, "do")
|
||||
|
||||
class State(TypedDict):
|
||||
my_key: Annotated[str, operator.add]
|
||||
market: str
|
||||
|
||||
tool_two_graph = StateGraph(State)
|
||||
tool_two_graph.add_node("tool_two", subgraph.compile())
|
||||
tool_two_graph.add_edge(START, "tool_two")
|
||||
tool_two = tool_two_graph.compile()
|
||||
|
||||
tracer = FakeTracer()
|
||||
assert tool_two.invoke(
|
||||
{"my_key": "value", "market": "DE"}, {"callbacks": [tracer]}
|
||||
) == {
|
||||
"my_key": "value",
|
||||
"market": "DE",
|
||||
}
|
||||
assert tool_two_node_count == 1, "interrupts aren't retried"
|
||||
assert len(tracer.runs) == 1
|
||||
run = tracer.runs[0]
|
||||
assert run.end_time is not None
|
||||
assert run.error is None
|
||||
assert run.outputs == {"market": "DE", "my_key": "value"}
|
||||
|
||||
assert tool_two.invoke({"my_key": "value", "market": "US"}) == {
|
||||
"my_key": "value all good",
|
||||
"market": "US",
|
||||
}
|
||||
|
||||
tool_two = tool_two_graph.compile(checkpointer=checkpointer)
|
||||
|
||||
# missing thread_id
|
||||
with pytest.raises(ValueError, match="thread_id"):
|
||||
tool_two.invoke({"my_key": "value", "market": "DE"})
|
||||
|
||||
# flow: interrupt -> resume with answer
|
||||
thread2 = {"configurable": {"thread_id": "2"}}
|
||||
# stop when about to enter node
|
||||
assert [
|
||||
c for c in tool_two.stream({"my_key": "value ⛰️", "market": "DE"}, thread2)
|
||||
] == [
|
||||
{
|
||||
"__interrupt__": (
|
||||
Interrupt(
|
||||
value="Just because...",
|
||||
resumable=True,
|
||||
ns=[AnyStr("tool_two:"), AnyStr("do:")],
|
||||
),
|
||||
)
|
||||
},
|
||||
]
|
||||
# resume with answer
|
||||
assert [c for c in tool_two.stream(Command(resume=" my answer"), thread2)] == [
|
||||
{"tool_two": {"my_key": " my answer", "market": "DE"}},
|
||||
]
|
||||
|
||||
# flow: interrupt -> clear tasks
|
||||
thread1 = {"configurable": {"thread_id": "1"}}
|
||||
# stop when about to enter node
|
||||
assert tool_two.invoke({"my_key": "value ⛰️", "market": "DE"}, thread1) == {
|
||||
"my_key": "value ⛰️",
|
||||
"market": "DE",
|
||||
}
|
||||
assert [
|
||||
c.metadata
|
||||
for c in tool_two.checkpointer.list(
|
||||
{"configurable": {"thread_id": "1", "checkpoint_ns": ""}}
|
||||
)
|
||||
] == [
|
||||
{
|
||||
"parents": {},
|
||||
"source": "loop",
|
||||
"step": 0,
|
||||
"writes": None,
|
||||
"thread_id": "1",
|
||||
},
|
||||
{
|
||||
"parents": {},
|
||||
"source": "input",
|
||||
"step": -1,
|
||||
"writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}},
|
||||
"thread_id": "1",
|
||||
},
|
||||
]
|
||||
assert tool_two.get_state(thread1) == StateSnapshot(
|
||||
values={"my_key": "value ⛰️", "market": "DE"},
|
||||
next=("tool_two",),
|
||||
tasks=(
|
||||
PregelTask(
|
||||
AnyStr(),
|
||||
"tool_two",
|
||||
(PULL, "tool_two"),
|
||||
interrupts=(
|
||||
Interrupt(
|
||||
value="Just because...",
|
||||
resumable=True,
|
||||
ns=[AnyStr("tool_two:"), AnyStr("do:")],
|
||||
),
|
||||
),
|
||||
state={
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": AnyStr("tool_two:"),
|
||||
}
|
||||
},
|
||||
),
|
||||
),
|
||||
config=tool_two.checkpointer.get_tuple(thread1).config,
|
||||
created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"],
|
||||
metadata={
|
||||
"parents": {},
|
||||
"source": "loop",
|
||||
"step": 0,
|
||||
"writes": None,
|
||||
"thread_id": "1",
|
||||
},
|
||||
parent_config=[
|
||||
*tool_two.checkpointer.list(
|
||||
{"configurable": {"thread_id": "1", "checkpoint_ns": ""}}, limit=2
|
||||
)
|
||||
][-1].config,
|
||||
)
|
||||
# clear the interrupt and next tasks
|
||||
tool_two.update_state(thread1, None, as_node=END)
|
||||
# interrupt and next tasks are cleared
|
||||
assert tool_two.get_state(thread1) == StateSnapshot(
|
||||
values={"my_key": "value ⛰️", "market": "DE"},
|
||||
next=(),
|
||||
tasks=(),
|
||||
config=tool_two.checkpointer.get_tuple(thread1).config,
|
||||
created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"],
|
||||
metadata={
|
||||
"parents": {},
|
||||
"source": "update",
|
||||
"step": 1,
|
||||
"writes": {},
|
||||
"thread_id": "1",
|
||||
},
|
||||
parent_config=[
|
||||
*tool_two.checkpointer.list(
|
||||
{"configurable": {"thread_id": "1", "checkpoint_ns": ""}}, limit=2
|
||||
)
|
||||
][-1].config,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
|
||||
def test_start_branch_then(
|
||||
snapshot: SnapshotAssertion, request: pytest.FixtureRequest, checkpointer_name: str
|
||||
@@ -14471,3 +14648,35 @@ def test_parent_command(request: pytest.FixtureRequest, checkpointer_name: str)
|
||||
},
|
||||
tasks=(),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
|
||||
def test_interrupt_subgraph(request: pytest.FixtureRequest, checkpointer_name: str):
|
||||
checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
|
||||
|
||||
class State(TypedDict):
|
||||
baz: str
|
||||
|
||||
def foo(state):
|
||||
return {"baz": "foo"}
|
||||
|
||||
def bar(state):
|
||||
value = interrupt("Please provide baz value:")
|
||||
return {"baz": value}
|
||||
|
||||
child_builder = StateGraph(State)
|
||||
child_builder.add_node(bar)
|
||||
child_builder.add_edge(START, "bar")
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node(foo)
|
||||
builder.add_node("bar", child_builder.compile())
|
||||
builder.add_edge(START, "foo")
|
||||
builder.add_edge("foo", "bar")
|
||||
graph = builder.compile(checkpointer=checkpointer)
|
||||
|
||||
thread1 = {"configurable": {"thread_id": "1"}}
|
||||
# First run, interrupted at bar
|
||||
assert graph.invoke({"baz": ""}, thread1)
|
||||
# Resume with answer
|
||||
assert graph.invoke(Command(resume="bar"), thread1)
|
||||
|
||||
@@ -429,6 +429,189 @@ async def test_dynamic_interrupt(checkpointer_name: str) -> None:
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
sys.version_info < (3, 11),
|
||||
reason="Python 3.11+ is required for async contextvars support",
|
||||
)
|
||||
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
|
||||
async def test_dynamic_interrupt_subgraph(checkpointer_name: str) -> None:
|
||||
class SubgraphState(TypedDict):
|
||||
my_key: str
|
||||
market: str
|
||||
|
||||
tool_two_node_count = 0
|
||||
|
||||
def tool_two_node(s: SubgraphState) -> SubgraphState:
|
||||
nonlocal tool_two_node_count
|
||||
tool_two_node_count += 1
|
||||
if s["market"] == "DE":
|
||||
answer = interrupt("Just because...")
|
||||
else:
|
||||
answer = " all good"
|
||||
return {"my_key": answer}
|
||||
|
||||
subgraph = StateGraph(SubgraphState)
|
||||
subgraph.add_node("do", tool_two_node, retry=RetryPolicy())
|
||||
subgraph.add_edge(START, "do")
|
||||
|
||||
class State(TypedDict):
|
||||
my_key: Annotated[str, operator.add]
|
||||
market: str
|
||||
|
||||
tool_two_graph = StateGraph(State)
|
||||
tool_two_graph.add_node("tool_two", subgraph.compile())
|
||||
tool_two_graph.add_edge(START, "tool_two")
|
||||
tool_two = tool_two_graph.compile()
|
||||
|
||||
tracer = FakeTracer()
|
||||
assert await tool_two.ainvoke(
|
||||
{"my_key": "value", "market": "DE"}, {"callbacks": [tracer]}
|
||||
) == {
|
||||
"my_key": "value",
|
||||
"market": "DE",
|
||||
}
|
||||
assert tool_two_node_count == 1, "interrupts aren't retried"
|
||||
assert len(tracer.runs) == 1
|
||||
run = tracer.runs[0]
|
||||
assert run.end_time is not None
|
||||
assert run.error is None
|
||||
assert run.outputs == {"market": "DE", "my_key": "value"}
|
||||
|
||||
assert await tool_two.ainvoke({"my_key": "value", "market": "US"}) == {
|
||||
"my_key": "value all good",
|
||||
"market": "US",
|
||||
}
|
||||
|
||||
async with awith_checkpointer(checkpointer_name) as checkpointer:
|
||||
tool_two = tool_two_graph.compile(checkpointer=checkpointer)
|
||||
|
||||
# missing thread_id
|
||||
with pytest.raises(ValueError, match="thread_id"):
|
||||
await tool_two.ainvoke({"my_key": "value", "market": "DE"})
|
||||
|
||||
# flow: interrupt -> resume with answer
|
||||
thread2 = {"configurable": {"thread_id": "2"}}
|
||||
# stop when about to enter node
|
||||
assert [
|
||||
c
|
||||
async for c in tool_two.astream(
|
||||
{"my_key": "value ⛰️", "market": "DE"}, thread2
|
||||
)
|
||||
] == [
|
||||
{
|
||||
"__interrupt__": (
|
||||
Interrupt(
|
||||
value="Just because...",
|
||||
resumable=True,
|
||||
ns=[AnyStr("tool_two:"), AnyStr("do:")],
|
||||
),
|
||||
)
|
||||
},
|
||||
]
|
||||
# resume with answer
|
||||
assert [
|
||||
c async for c in tool_two.astream(Command(resume=" my answer"), thread2)
|
||||
] == [
|
||||
{"tool_two": {"my_key": " my answer", "market": "DE"}},
|
||||
]
|
||||
|
||||
# flow: interrupt -> clear
|
||||
thread1 = {"configurable": {"thread_id": "1"}}
|
||||
thread1root = {"configurable": {"thread_id": "1", "checkpoint_ns": ""}}
|
||||
# stop when about to enter node
|
||||
assert [
|
||||
c
|
||||
async for c in tool_two.astream(
|
||||
{"my_key": "value ⛰️", "market": "DE"}, thread1
|
||||
)
|
||||
] == [
|
||||
{
|
||||
"__interrupt__": (
|
||||
Interrupt(
|
||||
value="Just because...",
|
||||
resumable=True,
|
||||
ns=[AnyStr("tool_two:"), AnyStr("do:")],
|
||||
),
|
||||
)
|
||||
},
|
||||
]
|
||||
assert [c.metadata async for c in tool_two.checkpointer.alist(thread1root)] == [
|
||||
{
|
||||
"parents": {},
|
||||
"source": "loop",
|
||||
"step": 0,
|
||||
"writes": None,
|
||||
"thread_id": "1",
|
||||
},
|
||||
{
|
||||
"parents": {},
|
||||
"source": "input",
|
||||
"step": -1,
|
||||
"writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}},
|
||||
"thread_id": "1",
|
||||
},
|
||||
]
|
||||
tup = await tool_two.checkpointer.aget_tuple(thread1)
|
||||
assert await tool_two.aget_state(thread1) == StateSnapshot(
|
||||
values={"my_key": "value ⛰️", "market": "DE"},
|
||||
next=("tool_two",),
|
||||
tasks=(
|
||||
PregelTask(
|
||||
AnyStr(),
|
||||
"tool_two",
|
||||
(PULL, "tool_two"),
|
||||
interrupts=(
|
||||
Interrupt(
|
||||
value="Just because...",
|
||||
resumable=True,
|
||||
ns=[AnyStr("tool_two:"), AnyStr("do:")],
|
||||
),
|
||||
),
|
||||
state={
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": AnyStr("tool_two:"),
|
||||
}
|
||||
},
|
||||
),
|
||||
),
|
||||
config=tup.config,
|
||||
created_at=tup.checkpoint["ts"],
|
||||
metadata={
|
||||
"parents": {},
|
||||
"source": "loop",
|
||||
"step": 0,
|
||||
"writes": None,
|
||||
"thread_id": "1",
|
||||
},
|
||||
parent_config=[
|
||||
c async for c in tool_two.checkpointer.alist(thread1root, limit=2)
|
||||
][-1].config,
|
||||
)
|
||||
|
||||
# clear the interrupt and next tasks
|
||||
await tool_two.aupdate_state(thread1, None, as_node=END)
|
||||
# interrupt is cleared, as well as the next tasks
|
||||
tup = await tool_two.checkpointer.aget_tuple(thread1)
|
||||
assert await tool_two.aget_state(thread1) == StateSnapshot(
|
||||
values={"my_key": "value ⛰️", "market": "DE"},
|
||||
next=(),
|
||||
tasks=(),
|
||||
config=tup.config,
|
||||
created_at=tup.checkpoint["ts"],
|
||||
metadata={
|
||||
"parents": {},
|
||||
"source": "update",
|
||||
"step": 1,
|
||||
"writes": {},
|
||||
"thread_id": "1",
|
||||
},
|
||||
parent_config=[
|
||||
c async for c in tool_two.checkpointer.alist(thread1root, limit=2)
|
||||
][-1].config,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not FF_SEND_V2, reason="send v2 is not enabled")
|
||||
@pytest.mark.skipif(
|
||||
sys.version_info < (3, 11),
|
||||
@@ -12677,3 +12860,39 @@ async def test_parent_command(checkpointer_name: str) -> None:
|
||||
},
|
||||
tasks=(),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
sys.version_info < (3, 11),
|
||||
reason="Python 3.11+ is required for async contextvars support",
|
||||
)
|
||||
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
|
||||
async def test_interrupt_subgraph(checkpointer_name: str):
|
||||
class State(TypedDict):
|
||||
baz: str
|
||||
|
||||
def foo(state):
|
||||
return {"baz": "foo"}
|
||||
|
||||
def bar(state):
|
||||
value = interrupt("Please provide baz value:")
|
||||
return {"baz": value}
|
||||
|
||||
child_builder = StateGraph(State)
|
||||
child_builder.add_node(bar)
|
||||
child_builder.add_edge(START, "bar")
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node(foo)
|
||||
builder.add_node("bar", child_builder.compile())
|
||||
builder.add_edge(START, "foo")
|
||||
builder.add_edge("foo", "bar")
|
||||
|
||||
async with awith_checkpointer(checkpointer_name) as checkpointer:
|
||||
graph = builder.compile(checkpointer=checkpointer)
|
||||
|
||||
thread1 = {"configurable": {"thread_id": "1"}}
|
||||
# First run, interrupted at bar
|
||||
assert await graph.ainvoke({"baz": ""}, thread1)
|
||||
# Resume with answer
|
||||
assert await graph.ainvoke(Command(resume="bar"), thread1)
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@langchain/langgraph-sdk",
|
||||
"version": "0.0.29",
|
||||
"version": "0.0.30",
|
||||
"description": "Client library for interacting with the LangGraph API",
|
||||
"type": "module",
|
||||
"packageManager": "yarn@1.22.19",
|
||||
|
||||
@@ -7,6 +7,7 @@ import {
|
||||
GraphSchema,
|
||||
Metadata,
|
||||
Run,
|
||||
RunStatus,
|
||||
Thread,
|
||||
ThreadState,
|
||||
Cron,
|
||||
@@ -944,12 +945,18 @@ export class RunsClient extends BaseClient {
|
||||
* Defaults to 0.
|
||||
*/
|
||||
offset?: number;
|
||||
|
||||
/**
|
||||
* Status of the run to filter by.
|
||||
*/
|
||||
status?: RunStatus;
|
||||
},
|
||||
): Promise<Run[]> {
|
||||
return this.fetch<Run[]>(`/threads/${threadId}/runs`, {
|
||||
params: {
|
||||
limit: options?.limit ?? 10,
|
||||
offset: options?.offset ?? 0,
|
||||
status: options?.status ?? undefined,
|
||||
},
|
||||
});
|
||||
}
|
||||
@@ -1014,19 +1021,28 @@ export class RunsClient extends BaseClient {
|
||||
*
|
||||
* @param threadId The ID of the thread.
|
||||
* @param runId The ID of the run.
|
||||
* @param signal An optional abort signal.
|
||||
* @returns An async generator yielding stream parts.
|
||||
*/
|
||||
async *joinStream(
|
||||
threadId: string,
|
||||
runId: string,
|
||||
signal?: AbortSignal,
|
||||
options?:
|
||||
| { signal?: AbortSignal; cancelOnDisconnect?: boolean }
|
||||
| AbortSignal,
|
||||
): AsyncGenerator<{ event: StreamEvent; data: any }> {
|
||||
const opts =
|
||||
typeof options === "object" &&
|
||||
options != null &&
|
||||
options instanceof AbortSignal
|
||||
? { signal: options }
|
||||
: options;
|
||||
|
||||
const response = await this.asyncCaller.fetch(
|
||||
...this.prepareFetchOptions(`/threads/${threadId}/runs/${runId}/stream`, {
|
||||
method: "GET",
|
||||
timeoutMs: null,
|
||||
signal,
|
||||
signal: opts?.signal,
|
||||
params: { cancel_on_disconnect: opts?.cancelOnDisconnect ? "1" : "0" },
|
||||
}),
|
||||
);
|
||||
|
||||
@@ -1041,7 +1057,7 @@ export class RunsClient extends BaseClient {
|
||||
async start(ctrl) {
|
||||
parser = createParser((event) => {
|
||||
if (
|
||||
(signal && signal.aborted) ||
|
||||
(opts?.signal && opts.signal.aborted) ||
|
||||
(event.type === "event" && event.data === "[DONE]")
|
||||
) {
|
||||
ctrl.terminate();
|
||||
|
||||
@@ -2,7 +2,7 @@ import type { JSONSchema7 } from "json-schema";
|
||||
|
||||
type Optional<T> = T | null | undefined;
|
||||
|
||||
type RunStatus =
|
||||
export type RunStatus =
|
||||
| "pending"
|
||||
| "running"
|
||||
| "error"
|
||||
|
||||
@@ -51,6 +51,7 @@ from langgraph_sdk.schema import (
|
||||
OnConflictBehavior,
|
||||
Run,
|
||||
RunCreate,
|
||||
RunStatus,
|
||||
SearchItemsResponse,
|
||||
StreamMode,
|
||||
StreamPart,
|
||||
@@ -1684,7 +1685,12 @@ class RunsClient:
|
||||
return response
|
||||
|
||||
async def list(
|
||||
self, thread_id: str, *, limit: int = 10, offset: int = 0
|
||||
self,
|
||||
thread_id: str,
|
||||
*,
|
||||
limit: int = 10,
|
||||
offset: int = 0,
|
||||
status: Optional[RunStatus] = None,
|
||||
) -> List[Run]:
|
||||
"""List runs.
|
||||
|
||||
@@ -1692,6 +1698,7 @@ class RunsClient:
|
||||
thread_id: The thread ID to list runs for.
|
||||
limit: The maximum number of results to return.
|
||||
offset: The number of results to skip.
|
||||
status: The status of the run to filter by.
|
||||
|
||||
Returns:
|
||||
List[Run]: The runs for the thread.
|
||||
@@ -1705,9 +1712,13 @@ class RunsClient:
|
||||
)
|
||||
|
||||
""" # noqa: E501
|
||||
return await self.http.get(
|
||||
f"/threads/{thread_id}/runs?limit={limit}&offset={offset}"
|
||||
)
|
||||
params = {
|
||||
"limit": limit,
|
||||
"offset": offset,
|
||||
}
|
||||
if status is not None:
|
||||
params["status"] = status
|
||||
return await self.http.get(f"/threads/{thread_id}/runs", params=params)
|
||||
|
||||
async def get(self, thread_id: str, run_id: str) -> Run:
|
||||
"""Get a run.
|
||||
@@ -1785,7 +1796,9 @@ class RunsClient:
|
||||
""" # noqa: E501
|
||||
return await self.http.get(f"/threads/{thread_id}/runs/{run_id}/join")
|
||||
|
||||
def join_stream(self, thread_id: str, run_id: str) -> AsyncIterator[StreamPart]:
|
||||
def join_stream(
|
||||
self, thread_id: str, run_id: str, *, cancel_on_disconnect: bool = False
|
||||
) -> AsyncIterator[StreamPart]:
|
||||
"""Stream output from a run in real-time, until the run is done.
|
||||
Output is not buffered, so any output produced before this call will
|
||||
not be received here.
|
||||
@@ -1793,6 +1806,7 @@ class RunsClient:
|
||||
Args:
|
||||
thread_id: The thread ID to join.
|
||||
run_id: The run ID to join.
|
||||
cancel_on_disconnect: Whether to cancel the run when the stream is disconnected.
|
||||
|
||||
Returns:
|
||||
None
|
||||
@@ -1805,7 +1819,11 @@ class RunsClient:
|
||||
)
|
||||
|
||||
""" # noqa: E501
|
||||
return self.http.stream(f"/threads/{thread_id}/runs/{run_id}/stream", "GET")
|
||||
return self.http.stream(
|
||||
f"/threads/{thread_id}/runs/{run_id}/stream",
|
||||
"GET",
|
||||
params={"cancel_on_disconnect": cancel_on_disconnect},
|
||||
)
|
||||
|
||||
async def delete(self, thread_id: str, run_id: str) -> None:
|
||||
"""Delete a run.
|
||||
@@ -3303,6 +3321,7 @@ class SyncRunsClient:
|
||||
assistant_id: str,
|
||||
*,
|
||||
input: Optional[dict] = None,
|
||||
command: Optional[Command] = None,
|
||||
stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
|
||||
stream_subgraphs: bool = False,
|
||||
metadata: Optional[dict] = None,
|
||||
@@ -3326,6 +3345,7 @@ class SyncRunsClient:
|
||||
assistant_id: str,
|
||||
*,
|
||||
input: Optional[dict] = None,
|
||||
command: Optional[Command] = None,
|
||||
stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
|
||||
stream_subgraphs: bool = False,
|
||||
metadata: Optional[dict] = None,
|
||||
@@ -3346,6 +3366,7 @@ class SyncRunsClient:
|
||||
assistant_id: str,
|
||||
*,
|
||||
input: Optional[dict] = None,
|
||||
command: Optional[Command] = None,
|
||||
stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
|
||||
stream_subgraphs: bool = False,
|
||||
metadata: Optional[dict] = None,
|
||||
@@ -3370,6 +3391,7 @@ class SyncRunsClient:
|
||||
assistant_id: The assistant ID or graph name to stream from.
|
||||
If using graph name, will default to first assistant created from that graph.
|
||||
input: The input to the graph.
|
||||
command: The command to execute.
|
||||
stream_mode: The stream mode(s) to use.
|
||||
stream_subgraphs: Whether to stream output from subgraphs.
|
||||
metadata: Metadata to assign to the run.
|
||||
@@ -3420,6 +3442,7 @@ class SyncRunsClient:
|
||||
""" # noqa: E501
|
||||
payload = {
|
||||
"input": input,
|
||||
"command": command,
|
||||
"config": config,
|
||||
"metadata": metadata,
|
||||
"stream_mode": stream_mode,
|
||||
@@ -3453,6 +3476,7 @@ class SyncRunsClient:
|
||||
assistant_id: str,
|
||||
*,
|
||||
input: Optional[dict] = None,
|
||||
command: Optional[Command] = None,
|
||||
stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
|
||||
stream_subgraphs: bool = False,
|
||||
metadata: Optional[dict] = None,
|
||||
@@ -3472,6 +3496,7 @@ class SyncRunsClient:
|
||||
assistant_id: str,
|
||||
*,
|
||||
input: Optional[dict] = None,
|
||||
command: Optional[Command] = None,
|
||||
stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
|
||||
stream_subgraphs: bool = False,
|
||||
metadata: Optional[dict] = None,
|
||||
@@ -3492,6 +3517,7 @@ class SyncRunsClient:
|
||||
assistant_id: str,
|
||||
*,
|
||||
input: Optional[dict] = None,
|
||||
command: Optional[Command] = None,
|
||||
stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
|
||||
stream_subgraphs: bool = False,
|
||||
metadata: Optional[dict] = None,
|
||||
@@ -3514,6 +3540,7 @@ class SyncRunsClient:
|
||||
assistant_id: The assistant ID or graph name to stream from.
|
||||
If using graph name, will default to first assistant created from that graph.
|
||||
input: The input to the graph.
|
||||
command: The command to execute.
|
||||
stream_mode: The stream mode(s) to use.
|
||||
stream_subgraphs: Whether to stream output from subgraphs.
|
||||
metadata: Metadata to assign to the run.
|
||||
@@ -3600,6 +3627,7 @@ class SyncRunsClient:
|
||||
""" # noqa: E501
|
||||
payload = {
|
||||
"input": input,
|
||||
"command": command,
|
||||
"stream_mode": stream_mode,
|
||||
"stream_subgraphs": stream_subgraphs,
|
||||
"config": config,
|
||||
@@ -3637,6 +3665,7 @@ class SyncRunsClient:
|
||||
assistant_id: str,
|
||||
*,
|
||||
input: Optional[dict] = None,
|
||||
command: Optional[Command] = None,
|
||||
metadata: Optional[dict] = None,
|
||||
config: Optional[Config] = None,
|
||||
checkpoint: Optional[Checkpoint] = None,
|
||||
@@ -3657,6 +3686,7 @@ class SyncRunsClient:
|
||||
assistant_id: str,
|
||||
*,
|
||||
input: Optional[dict] = None,
|
||||
command: Optional[Command] = None,
|
||||
metadata: Optional[dict] = None,
|
||||
config: Optional[Config] = None,
|
||||
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
|
||||
@@ -3674,6 +3704,7 @@ class SyncRunsClient:
|
||||
assistant_id: str,
|
||||
*,
|
||||
input: Optional[dict] = None,
|
||||
command: Optional[Command] = None,
|
||||
metadata: Optional[dict] = None,
|
||||
config: Optional[Config] = None,
|
||||
checkpoint: Optional[Checkpoint] = None,
|
||||
@@ -3695,6 +3726,7 @@ class SyncRunsClient:
|
||||
assistant_id: The assistant ID or graph name to run.
|
||||
If using graph name, will default to first assistant created from that graph.
|
||||
input: The input to the graph.
|
||||
command: The command to execute.
|
||||
metadata: Metadata to assign to the run.
|
||||
config: The configuration for the assistant.
|
||||
checkpoint: The checkpoint to resume from.
|
||||
@@ -3761,6 +3793,7 @@ class SyncRunsClient:
|
||||
""" # noqa: E501
|
||||
payload = {
|
||||
"input": input,
|
||||
"command": command,
|
||||
"config": config,
|
||||
"metadata": metadata,
|
||||
"assistant_id": assistant_id,
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
"""Data models for interacting with the LangGraph API."""
|
||||
|
||||
from datetime import datetime
|
||||
from typing import Any, Literal, NamedTuple, Optional, Sequence, TypedDict, Union
|
||||
from typing import Any, Dict, Literal, NamedTuple, Optional, Sequence, TypedDict, Union
|
||||
|
||||
Json = Optional[dict[str, Any]]
|
||||
"""Represents a JSON-like structure, which can be None or a dictionary with string keys and any values."""
|
||||
@@ -176,6 +176,19 @@ class Assistant(AssistantBase):
|
||||
"""The name of the assistant"""
|
||||
|
||||
|
||||
class Interrupt(TypedDict, total=False):
|
||||
"""Represents an interruption in the execution flow."""
|
||||
|
||||
value: Any
|
||||
"""The value associated with the interrupt."""
|
||||
when: Literal["during"]
|
||||
"""When the interrupt occurred."""
|
||||
resumable: bool
|
||||
"""Whether the interrupt can be resumed."""
|
||||
ns: Optional[list[str]]
|
||||
"""Optional namespace for the interrupt."""
|
||||
|
||||
|
||||
class Thread(TypedDict):
|
||||
"""Represents a conversation thread."""
|
||||
|
||||
@@ -191,6 +204,8 @@ class Thread(TypedDict):
|
||||
"""The status of the thread, one of 'idle', 'busy', 'interrupted'."""
|
||||
values: Json
|
||||
"""The current state of the thread."""
|
||||
interrupts: Dict[str, list[Interrupt]]
|
||||
"""Interrupts which were thrown in this thread"""
|
||||
|
||||
|
||||
class ThreadTask(TypedDict):
|
||||
@@ -199,7 +214,7 @@ class ThreadTask(TypedDict):
|
||||
id: str
|
||||
name: str
|
||||
error: Optional[str]
|
||||
interrupts: list[dict]
|
||||
interrupts: list[Interrupt]
|
||||
checkpoint: Optional[Checkpoint]
|
||||
state: Optional["ThreadState"]
|
||||
result: Optional[dict[str, Any]]
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[tool.poetry]
|
||||
name = "langgraph-sdk"
|
||||
version = "0.1.40"
|
||||
version = "0.1.41"
|
||||
description = "SDK for interacting with LangGraph API"
|
||||
authors = []
|
||||
license = "MIT"
|
||||
|
||||
Reference in New Issue
Block a user