Compare commits

..
Author SHA1 Message Date
Nuno Campos 3a860ad537 0.2.48 2024-11-13 17:45:04 -08:00
Nuno Campos 7051bccc30 checkpoint 2.0.4 2024-11-13 17:37:01 -08:00
Nuno Campos 199e41b228 sdk py 0.1.36 2024-11-13 17:07:07 -08:00
Nuno Campos 229a9e19a8 sdk-py: Add command arg for creating runs 2024-11-13 17:06:56 -08:00
Nuno Campos 3c3a1a1f35 0.2.47 2024-11-13 14:02:25 -08:00
Nuno CamposandGitHub f11127648e Merge pull request #2346 from langchain-ai/nc/4nov/send-eager
lib: Execute Sends in the superstep that originated them (feature-flagged)
2024-11-13 13:39:50 -08:00
Nuno CamposandGitHub 29f833b1a7 Merge pull request #2393 from langchain-ai/nc/11nov/command-resume
lib: Add interrupt() function
2024-11-13 13:34:19 -08:00
Nuno CamposandGitHub 7a3ea42743 Merge pull request #2410 from langchain-ai/nc/13nov/command-dataclass
Nc/13nov/command dataclass
2024-11-13 13:28:13 -08:00
Nuno Campos a94902db8a Lint 2024-11-13 13:17:02 -08:00
Nuno Campos 9fd152ef3a format 2024-11-13 13:14:03 -08:00
Nuno Campos 03bc9ba6e6 Make Command a dataclass 2024-11-13 13:11:28 -08:00
Nuno Campos 7fe6f88876 Add resumeable/ns properties to Interrupt 2024-11-13 12:51:55 -08:00
Nuno CamposandGitHub 1d88affd29 Merge pull request #2400 from langchain-ai/nc/12nov/command
Make Command accept generic arg for destinations
2024-11-13 10:03:27 -08:00
Eugene YurtsevandGitHub 6906e12edb cli: add ability to output docker compose file (#2379)
* Add ability to output docker compose file `langgraph dockerfile
Dockerfile --add-docker-compose`
* Add emoji in places
2024-11-13 12:42:58 -05:00
Nuno Campos 16bfa80b58 Make Command accept generic arg for destinations 2024-11-12 16:18:23 -08:00
Nuno Campos 00964b18f6 Undo 2024-11-11 18:12:25 -08:00
Nuno Campos 0d5c6201d3 Disable in py 3.10 or below for async 2024-11-11 18:09:58 -08:00
Nuno Campos 86d2847dab Use neg idx 2024-11-11 17:57:45 -08:00
Nuno Campos b3a4eaa967 Remove print 2024-11-11 17:56:10 -08:00
Nuno Campos 87fc519ce7 Remove print 2024-11-11 17:54:56 -08:00
Nuno Campos ef3a1ee997 Undo 2024-11-11 17:54:04 -08:00
Nuno Campos 311e16dffd Remove prints 2024-11-11 17:52:18 -08:00
Nuno Campos c83b8f6d04 Update 2024-11-11 17:49:58 -08:00
Nuno Campos 62d3a85b07 Add sync test 2024-11-11 17:46:54 -08:00
Nuno Campos 810ae0ef51 lib: Add interrupt() function
- This works similarly to the input() function from stdlib
- calling it in a node interrupts execution
- invoking the graph with Command(resume=...) will set ... as the return value of interrupt() so that the node can access the "answer" to the "question"
- This PR also starts the work to control the graph on invoke/stream with Command() input, to be continued in a future PR
2024-11-11 17:44:03 -08:00
Nuno Campos 2ff49d2200 Update 2024-11-11 15:43:20 -08:00
Nuno Campos d0567dc7be Add feature flag (default off) so we can merge this before releasing
- Add additional ci job to test with FF on
2024-11-11 15:38:48 -08:00
Nuno Campos ea64ac5c07 Update 2024-11-11 14:18:14 -08:00
Nuno Campos 090b53ccc1 Lint 2024-11-11 14:14:23 -08:00
Nuno Campos 0e872e7482 Lin t 2024-11-11 14:12:33 -08:00
Nuno Campos 3ad966e057 Update 2024-11-11 14:08:59 -08:00
Nuno Campos 89a0859928 Execute Sends in same super step that triggered them
- Keep old code path for compatibility with existing checkpoints
- Keep a similar order of application of updates, in some cases there will be no visible change
- Update task path for Sends to contain the path of all the parent tasks (multiple parents when a Send task creates another Send)
- That lineage path is used to ensure order of application of updates respects their logical lineage (ie updates from parents always applied before their child tasks)
- Move Interrupt writes to use negative indexes, which allow replacing/shadowing (when task is re-run it may interrupt again, or succeed)
- Runner will now attempt to schedule new Send tasks as soon as the write is received (ie while the originating node is still running)
- Update kafka scheduler to support new Send behavior
2024-11-11 14:01:16 -08:00
41 changed files with 3357 additions and 816 deletions
+8 -1
View File
@@ -19,14 +19,19 @@ jobs:
- "3.13"
core-version:
- "latest"
ff-send-v2:
- "false"
include:
- python-version: "3.11"
core-version: ">=0.2.42,<0.3.0"
- python-version: "3.11"
core-version: "latest"
ff-send-v2: "true"
defaults:
run:
working-directory: libs/langgraph
name: "test #${{ matrix.python-version }} (langchain-core: ${{ matrix.core-version }})"
name: "test #${{ matrix.python-version }} (langchain-core: ${{ matrix.core-version }}, ff-send-v2: ${{ matrix.ff-send-v2 }})"
steps:
- uses: actions/checkout@v4
- name: Set up Python ${{ matrix.python-version }} + Poetry ${{ env.POETRY_VERSION }}
@@ -52,6 +57,8 @@ jobs:
- name: Run tests
shell: bash
env:
LANGGRAPH_FF_SEND_V2: ${{ matrix.ff-send-v2 }}
run: |
make test
@@ -24,6 +24,8 @@ from langgraph.checkpoint.serde.base import SerializerProtocol, maybe_add_typed_
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
from langgraph.checkpoint.serde.types import (
ERROR,
INTERRUPT,
RESUME,
SCHEDULED,
ChannelProtocol,
SendProtocol,
@@ -449,4 +451,4 @@ Special writes (e.g. errors) map to negative indices, to avoid those writes from
conflicting with regular writes.
Each Checkpointer implementation should use this mapping in put_writes.
"""
WRITES_IDX_MAP = {ERROR: -1, SCHEDULED: -2}
WRITES_IDX_MAP = {ERROR: -1, SCHEDULED: -2, INTERRUPT: -3, RESUME: -4}
@@ -25,7 +25,7 @@ from langchain_core.load.serializable import Serializable
from zoneinfo import ZoneInfo
from langgraph.checkpoint.serde.base import SerializerProtocol
from langgraph.checkpoint.serde.types import CommandProtocol, SendProtocol
from langgraph.checkpoint.serde.types import SendProtocol
from langgraph.store.base import Item
LC_REVIVER = Reviver()
@@ -122,11 +122,6 @@ class JsonPlusSerializer(SerializerProtocol):
return self._encode_constructor_args(
obj.__class__, kwargs={"node": obj.node, "arg": obj.arg}
)
elif isinstance(obj, CommandProtocol):
return self._encode_constructor_args(
obj.__class__,
kwargs={k: getattr(obj, k) for k in obj.__all_slots__},
)
elif isinstance(obj, (bytes, bytearray)):
return self._encode_constructor_args(
obj.__class__, method="fromhex", args=(obj.hex(),)
@@ -407,17 +402,6 @@ def _msgpack_default(obj: Any) -> Union[str, msgpack.ExtType]:
(obj.__class__.__module__, obj.__class__.__name__, (obj.node, obj.arg)),
),
)
elif isinstance(obj, CommandProtocol):
return msgpack.ExtType(
EXT_CONSTRUCTOR_KW_ARGS,
_msgpack_enc(
(
obj.__class__.__module__,
obj.__class__.__name__,
{k: getattr(obj, k) for k in obj.__all_slots__},
),
),
)
elif dataclasses.is_dataclass(obj):
# doesn't use dataclasses.asdict to avoid deepcopy and recursion
return msgpack.ExtType(
@@ -4,7 +4,6 @@ from typing import (
Protocol,
Sequence,
TypeVar,
Union,
runtime_checkable,
)
@@ -12,6 +11,8 @@ from typing_extensions import Self
ERROR = "__error__"
SCHEDULED = "__scheduled__"
INTERRUPT = "__interrupt__"
RESUME = "__resume__"
TASKS = "__pregel_tasks"
Value = TypeVar("Value", covariant=True)
@@ -49,11 +50,3 @@ class SendProtocol(Protocol):
def __repr__(self) -> str: ...
def __eq__(self, value: object) -> bool: ...
@runtime_checkable
class CommandProtocol(Protocol):
# Mirrors langgraph.types.Command
update: Optional[dict[str, Any]]
send: Union[Any, Sequence[Any]]
__all_slots__: set[str]
+1 -1
View File
@@ -1,6 +1,6 @@
[tool.poetry]
name = "langgraph-checkpoint"
version = "2.0.3"
version = "2.0.4"
description = "Library with base interfaces for LangGraph checkpoint savers."
authors = []
license = "MIT"
+140 -9
View File
@@ -6,6 +6,7 @@ from typing import Callable, Optional, Sequence
import click
import click.exceptions
from click import secho
import langgraph_cli.config
import langgraph_cli.docker
@@ -170,9 +171,7 @@ def cli():
is_flag=True,
help="Wait for services to start before returning. Implies --detach",
)
@cli.command(
help="Start langgraph API server. For local testing, requires a LangSmith API key with access to LangGraph Cloud closed beta. Requires a license key for production use."
)
@cli.command(help="🚀 Launch LangGraph API server.")
@log_command
def up(
config: pathlib.Path,
@@ -337,7 +336,7 @@ def _build(
)
@click.argument("docker_build_args", nargs=-1, type=click.UNPROCESSED)
@cli.command(
help="Build langgraph API server docker image",
help="📦 Build LangGraph API server Docker image.",
context_settings=dict(
ignore_unknown_options=True,
),
@@ -360,14 +359,86 @@ def build(
)
def _get_docker_ignore_content() -> str:
"""Return the content of a .dockerignore file.
This file is used to exclude files and directories from the Docker build context.
It may be overly broad, but it's better to be safe than sorry.
The main goal is to exclude .env files by default.
"""
return """\
# Ignore node_modules and other dependency directories
node_modules
bower_components
vendor
# Ignore logs and temporary files
*.log
*.tmp
*.swp
# Ignore .env files and other environment files
.env
.env.*
*.local
# Ignore git-related files
.git
.gitignore
# Ignore Docker-related files and configs
.dockerignore
docker-compose.yml
# Ignore build and cache directories
dist
build
.cache
__pycache__
# Ignore IDE and editor configurations
.vscode
.idea
*.sublime-project
*.sublime-workspace
.DS_Store # macOS-specific
# Ignore test and coverage files
coverage
*.coverage
*.test.js
*.spec.js
tests
"""
@OPT_CONFIG
@click.argument("save_path", type=click.Path(resolve_path=True))
@cli.command(help="Generate a Dockerfile for langgraph API server")
@cli.command(
help="🐳 Generate a Dockerfile for the LangGraph API server, with Docker Compose options."
)
@click.option(
# Add a flag for adding a docker-compose.yml file as part of the output
"--add-docker-compose",
help=(
"Add additional files for running the LangGraph API server with "
"docker-compose. These files include a docker-compose.yml, .env file, "
"and a .dockerignore file."
),
is_flag=True,
)
@log_command
def dockerfile(save_path: pathlib.Path, config: pathlib.Path):
with open(config) as f:
def dockerfile(save_path: str, config: pathlib.Path, add_docker_compose: bool) -> None:
save_path = pathlib.Path(save_path).absolute()
secho(f"🔍 Validating configuration at path: {config}", fg="yellow")
with open(config, encoding="utf-8") as f:
config_json = langgraph_cli.config.validate_config(json.load(f))
with open(save_path, "w") as f:
secho("✅ Configuration validated!", fg="green")
secho(f"📝 Generating Dockerfile at {save_path}", fg="yellow")
with open(str(save_path), "w", encoding="utf-8") as f:
f.write(
langgraph_cli.config.config_to_docker(
config,
@@ -377,6 +448,66 @@ def dockerfile(save_path: pathlib.Path, config: pathlib.Path):
else "langchain/langgraph-api",
)
)
secho("✅ Created: Dockerfile", fg="green")
if add_docker_compose:
# Add docker compose and related files
# Add .dockerignore file in the same directory as the Dockerfile
with open(str(save_path.parent / ".dockerignore"), "w", encoding="utf-8") as f:
f.write(_get_docker_ignore_content())
secho("✅ Created: .dockerignore", fg="green")
# Generate a docker-compose.yml file
path = str(save_path.parent / "docker-compose.yml")
with open(path, "w", encoding="utf-8") as f:
with Runner() as runner:
capabilities = langgraph_cli.docker.check_capabilities(runner)
compose_dict = langgraph_cli.docker.compose_as_dict(
capabilities,
port=8123,
)
# Add .env file to the docker-compose.yml for the langgraph-api service
compose_dict["services"]["langgraph-api"]["env_file"] = [".env"]
# Add the Dockerfile to the build context
compose_dict["services"]["langgraph-api"]["build"] = {
"context": ".",
"dockerfile": save_path.name,
}
f.write(langgraph_cli.docker.dict_to_yaml(compose_dict))
secho("✅ Created: docker-compose.yml", fg="green")
# Check if the .env file exists in the same directory as the Dockerfile
if not (save_path.parent / ".env").exists():
# Also add an empty .env file
with open(str(save_path.parent / ".env"), "w", encoding="utf-8") as f:
f.writelines(
[
"# Uncomment the following line to add your LangSmith API key",
"\n",
"# LANGSMITH_API_KEY=your-api-key",
"\n",
"# Or if you have a LangGraph Cloud license key, "
"then uncomment the following line: ",
"\n",
"# LANGGRAPH_CLOUD_LICENSE_KEY=your-license-key",
"\n",
"# Add any other environment variables go below...",
]
)
secho("✅ Created: .env", fg="green")
else:
# Do nothing since the .env file already exists. Not a great
# idea to overwrite in case the user has added custom env vars set
# in the .env file already.
secho(" Skipped: .env. It already exists!", fg="yellow")
secho(
f"🎉 Files generated successfully at path {save_path.parent}!",
fg="cyan",
bold=True,
)
@click.argument("path", required=False)
@@ -385,7 +516,7 @@ def dockerfile(save_path: pathlib.Path, config: pathlib.Path):
type=str,
help=TEMPLATE_HELP_STRING,
)
@cli.command("new", help="Create a new LangGraph project from a template.")
@cli.command("new", help="🌱 Create a new LangGraph project from a template.")
@log_command
def new(path: Optional[str], template: Optional[str]) -> None:
"""Create a new LangGraph project from a template."""
+145 -93
View File
@@ -12,34 +12,6 @@ DEFAULT_POSTGRES_URI = (
"postgres://postgres:postgres@langgraph-postgres:5432/postgres?sslmode=disable"
)
REDIS = """
langgraph-redis:
image: redis:6
healthcheck:
test: redis-cli ping
interval: 5s
timeout: 1s
retries: 5
"""
DB = """
langgraph-postgres:
image: postgres:16
ports:
- "5433:5432"
environment:
POSTGRES_DB: postgres
POSTGRES_USER: postgres
POSTGRES_PASSWORD: postgres
volumes:
- langgraph-data:/var/lib/postgresql/data
healthcheck:
test: pg_isready -U postgres
start_period: 10s
timeout: 1s
retries: 5
"""
class Version(NamedTuple):
major: int
@@ -116,28 +88,149 @@ def check_capabilities(runner) -> DockerCapabilities:
def debugger_compose(
*, port: Optional[int] = None, base_url: Optional[str] = None
) -> str:
) -> dict:
if port is None:
return ""
compose_str = """
langgraph-debugger:
image: langchain/langgraph-debugger
restart: on-failure
depends_on:
langgraph-postgres:
condition: service_healthy
ports:
- "{port}:3968"
"""
config = {
"langgraph-debugger": {
"image": "langchain/langgraph-debugger",
"restart": "on-failure",
"depends_on": {
"langgraph-postgres": {"condition": "service_healthy"},
},
"ports": [f'"{port}:3968"'],
}
}
if base_url:
compose_str += """
environment:
VITE_STUDIO_LOCAL_GRAPH_URL: {base_url}
"""
config["langgraph-debugger"]["environment"] = {
"VITE_STUDIO_LOCAL_GRAPH_URL": base_url
}
return compose_str.format(port=port, base_url=base_url)
return config
# Function to convert dictionary to YAML
def dict_to_yaml(d: dict, *, indent: int = 0) -> str:
"""Convert a dictionary to a YAML string."""
yaml_str = ""
for idx, (key, value) in enumerate(d.items()):
# Format things in a visually appealing way
# Use an extra newline for top-level keys only
if idx >= 1 and indent < 2:
yaml_str += "\n"
space = " " * indent
if isinstance(value, dict):
yaml_str += f"{space}{key}:\n" + dict_to_yaml(value, indent=indent + 1)
elif isinstance(value, list):
yaml_str += f"{space}{key}:\n"
for item in value:
yaml_str += f"{space} - {item}\n"
else:
yaml_str += f"{space}{key}: {value}\n"
return yaml_str
def compose_as_dict(
capabilities: DockerCapabilities,
*,
port: int,
debugger_port: Optional[int] = None,
debugger_base_url: Optional[str] = None,
# postgres://user:password@host:port/database?option=value
postgres_uri: Optional[str] = None,
) -> dict:
"""Create a docker compose file as a dictionary in YML style."""
if postgres_uri is None:
include_db = True
postgres_uri = DEFAULT_POSTGRES_URI
else:
include_db = False
# The services below are defined in a non-intuitive order to match
# the existing unit tests for this function.
# It's fine to re-order just requires updating the unit tests, so it should
# be done with caution.
# Define the Redis service first as per the test order
services = {
"langgraph-redis": {
"image": "redis:6",
"healthcheck": {
"test": "redis-cli ping",
"interval": "5s",
"timeout": "1s",
"retries": 5,
},
}
}
# Add Postgres service before langgraph-api if it is needed
if include_db:
services["langgraph-postgres"] = {
"image": "postgres:16",
"ports": ['"5433:5432"'],
"environment": {
"POSTGRES_DB": "postgres",
"POSTGRES_USER": "postgres",
"POSTGRES_PASSWORD": "postgres",
},
"volumes": ["langgraph-data:/var/lib/postgresql/data"],
"healthcheck": {
"test": "pg_isready -U postgres",
"start_period": "10s",
"timeout": "1s",
"retries": 5,
},
}
if capabilities.healthcheck_start_interval:
services["langgraph-postgres"]["healthcheck"]["interval"] = "60s"
services["langgraph-postgres"]["healthcheck"]["start_interval"] = "1s"
else:
services["langgraph-postgres"]["healthcheck"]["interval"] = "5s"
# Add optional debugger service if debugger_port is specified
if debugger_port:
services["langgraph-debugger"] = debugger_compose(
port=debugger_port, base_url=debugger_base_url
)["langgraph-debugger"]
# Add langgraph-api service
services["langgraph-api"] = {
"ports": [f'"{port}:8000"'],
"depends_on": {
"langgraph-redis": {"condition": "service_healthy"},
},
"environment": {
"REDIS_URI": "redis://langgraph-redis:6379",
"POSTGRES_URI": postgres_uri,
},
}
# If Postgres is included, add it to the dependencies of langgraph-api
if include_db:
services["langgraph-api"]["depends_on"]["langgraph-postgres"] = {
"condition": "service_healthy"
}
# Additional healthcheck for langgraph-api if required
if capabilities.healthcheck_start_interval:
services["langgraph-api"]["healthcheck"] = {
"test": "python /api/healthcheck.py",
"interval": "60s",
"start_interval": "1s",
"start_period": "10s",
}
# Final compose dictionary with volumes included if needed
compose_dict = {}
if include_db:
compose_dict["volumes"] = {"langgraph-data": {"driver": "local"}}
compose_dict["services"] = services
return compose_dict
def compose(
@@ -149,54 +242,13 @@ def compose(
# postgres://user:password@host:port/database?option=value
postgres_uri: Optional[str] = None,
) -> str:
if postgres_uri is None:
include_db = True
postgres_uri = DEFAULT_POSTGRES_URI
else:
include_db = False
db = DB.format() if include_db else ""
volumes = (
"""volumes:
langgraph-data:
driver: local
"""
if include_db
else ""
"""Create a docker compose file as a string."""
compose_content = compose_as_dict(
capabilities,
port=port,
debugger_port=debugger_port,
debugger_base_url=debugger_base_url,
postgres_uri=postgres_uri,
)
if db:
if capabilities.healthcheck_start_interval:
db += """
interval: 60s
start_interval: 1s"""
else:
db += """
interval: 5s"""
compose_str = f"""{volumes}services:
{REDIS}
{db}
{debugger_compose(port=debugger_port, base_url=debugger_base_url)}
langgraph-api:
ports:
- "{port}:8000\"
depends_on:
langgraph-redis:
condition: service_healthy"""
if include_db:
compose_str += """
langgraph-postgres:
condition: service_healthy"""
compose_str += f"""
environment:
REDIS_URI: redis://langgraph-redis:6379
POSTGRES_URI: {postgres_uri}
"""
if capabilities.healthcheck_start_interval:
compose_str += """ healthcheck:
test: python /api/healthcheck.py
interval: 60s
start_interval: 1s
start_period: 10s"""
compose_str = dict_to_yaml(compose_content)
return compose_str
+113 -2
View File
@@ -1,4 +1,9 @@
import json
import pathlib
import shutil
import tempfile
from contextlib import contextmanager
from pathlib import Path
from click.testing import CliRunner
@@ -14,7 +19,26 @@ DEFAULT_DOCKER_CAPABILITIES = DockerCapabilities(
)
def test_prepare_args_and_stdin():
@contextmanager
def temporary_config_folder(config_content: dict):
# Create a temporary directory
temp_dir = tempfile.mkdtemp()
try:
# Define the path for the config.json file
config_path = Path(temp_dir) / "config.json"
# Write the provided dictionary content to config.json
with open(config_path, "w", encoding="utf-8") as config_file:
json.dump(config_content, config_file)
# Yield the temporary directory path for use within the context
yield config_path.parent
finally:
# Cleanup the temporary directory and its contents
shutil.rmtree(temp_dir)
def test_prepare_args_and_stdin() -> None:
# this basically serves as an end-to-end test for using config and docker helpers
config_path = pathlib.Path("./langgraph.json")
config = validate_config(
@@ -28,7 +52,7 @@ def test_prepare_args_and_stdin():
capabilities=DEFAULT_DOCKER_CAPABILITIES,
config_path=config_path,
config=config,
docker_compose="custom-docker-compose.yml",
docker_compose=pathlib.Path("custom-docker-compose.yml"),
port=port,
debugger_port=debugger_port,
debugger_base_url=debugger_graph_url,
@@ -131,3 +155,90 @@ def test_version_option() -> None:
assert (
"LangGraph CLI, version" in result.output
), "Expected version information in output"
def test_dockerfile_command_basic() -> None:
"""Test the 'dockerfile' command with basic configuration."""
runner = CliRunner()
config_content = {
"node_version": "20", # Add any other necessary configuration fields
"graphs": {"agent": "agent.py:graph"},
}
with temporary_config_folder(config_content) as temp_dir:
save_path = temp_dir / "Dockerfile"
result = runner.invoke(
cli,
["dockerfile", str(save_path), "--config", str(temp_dir / "config.json")],
)
# Assert command was successful
assert result.exit_code == 0, result.output
assert "✅ Created: Dockerfile" in result.output
# Check if Dockerfile was created
assert save_path.exists()
def test_dockerfile_command_with_docker_compose() -> None:
"""Test the 'dockerfile' command with Docker Compose configuration."""
runner = CliRunner()
config_content = {
"dependencies": ["./my_agent"],
"graphs": {"agent": "./my_agent/agent.py:graph"},
"env": ".env",
}
with temporary_config_folder(config_content) as temp_dir:
save_path = temp_dir / "Dockerfile"
# Add agent.py file
agent_path = temp_dir / "my_agent" / "agent.py"
agent_path.parent.mkdir(parents=True, exist_ok=True)
agent_path.touch()
result = runner.invoke(
cli,
[
"dockerfile",
str(save_path),
"--config",
str(temp_dir / "config.json"),
"--add-docker-compose",
],
)
# Assert command was successful
assert result.exit_code == 0
assert "✅ Created: Dockerfile" in result.output
assert "✅ Created: .dockerignore" in result.output
assert "✅ Created: docker-compose.yml" in result.output
assert (
"✅ Created: .env" in result.output or " Skipped: .env" in result.output
)
assert "🎉 Files generated successfully" in result.output
# Check if Dockerfile, .dockerignore, docker-compose.yml, and .env were created
assert save_path.exists()
assert (temp_dir / ".dockerignore").exists()
assert (temp_dir / "docker-compose.yml").exists()
assert (temp_dir / ".env").exists() or " Skipped: .env" in result.output
def test_dockerfile_command_with_bad_config() -> None:
"""Test the 'dockerfile' command with basic configuration."""
runner = CliRunner()
config_content = {
"node_version": "20" # Add any other necessary configuration fields
}
with temporary_config_folder(config_content) as temp_dir:
save_path = temp_dir / "Dockerfile"
result = runner.invoke(
cli,
["dockerfile", str(save_path), "--config", str(temp_dir / "conf.json")],
)
# Assert command was successful
assert result.exit_code == 2
assert "conf.json' does not exist" in result.output
+1 -1
View File
@@ -49,7 +49,7 @@ test:
exit $$EXIT_CODE
test_watch:
make start-postgres && poetry run ptw . -- --ff -v -x -n auto --dist worksteal --snapshot-update --tb short $(TEST); \
make start-postgres && poetry run ptw . -- --ff -vv -x -n auto --dist worksteal --snapshot-update --tb short $(TEST); \
EXIT_CODE=$$?; \
make stop-postgres; \
exit $$EXIT_CODE
+11
View File
@@ -1,4 +1,5 @@
import sys
from os import getenv
from types import MappingProxyType
from typing import Any, Literal, Mapping, cast
@@ -10,6 +11,7 @@ from langgraph.types import Interrupt, Send # noqa: F401
# --- Empty read-only containers ---
EMPTY_MAP: Mapping[str, Any] = MappingProxyType({})
EMPTY_SEQ: tuple[str, ...] = tuple()
MISSING = object()
# --- Public constants ---
TAG_NOSTREAM = sys.intern("langsmith:nostream")
@@ -28,6 +30,8 @@ INPUT = sys.intern("__input__")
# for values passed as input to the graph
INTERRUPT = sys.intern("__interrupt__")
# for dynamic interrupts raised by nodes
RESUME = sys.intern("__resume__")
# for values passed to resume a node after an interrupt
ERROR = sys.intern("__error__")
# for errors raised by nodes
NO_WRITES = sys.intern("__no_writes__")
@@ -69,6 +73,8 @@ CONFIG_KEY_CHECKPOINT_NS = sys.intern("checkpoint_ns")
# holds the current checkpoint_ns, "" for root graph
CONFIG_KEY_NODE_FINISHED = sys.intern("__pregel_node_finished")
# callback to be called when a node is finished
CONFIG_KEY_RESUME_VALUE = sys.intern("__pregel_resume_value")
# holds the value that "answers" an interrupt() call
# --- Other constants ---
PUSH = sys.intern("__pregel_push")
@@ -81,12 +87,17 @@ NS_END = sys.intern(":")
# for checkpoint_ns, for each level, separates the namespace from the task_id
CONF = cast(Literal["configurable"], sys.intern("configurable"))
# key for the configurable dict in RunnableConfig
FF_SEND_V2 = getenv("LANGGRAPH_FF_SEND_V2", "false").lower() == "true"
# temporary flag to enable new Send semantics
NULL_TASK_ID = sys.intern("00000000-0000-0000-0000-000000000000")
# the task_id to use for writes that are not associated with a task
RESERVED = {
TAG_HIDDEN,
# reserved write keys
INPUT,
INTERRUPT,
RESUME,
ERROR,
NO_WRITES,
SCHEDULED,
+1 -1
View File
@@ -70,7 +70,7 @@ class NodeInterrupt(GraphInterrupt):
"""Raised by a node to interrupt execution."""
def __init__(self, value: Any) -> None:
super().__init__([Interrupt(value)])
super().__init__([Interrupt(value=value)])
class GraphDelegate(Exception):
+14 -16
View File
@@ -1,3 +1,4 @@
import dataclasses
import inspect
import logging
import typing
@@ -14,7 +15,6 @@ from typing import (
Optional,
Sequence,
Type,
TypeVar,
Union,
cast,
get_args,
@@ -50,15 +50,13 @@ from langgraph.managed.base import (
from langgraph.pregel.read import ChannelRead, PregelNode
from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry
from langgraph.store.base import BaseStore
from langgraph.types import All, Checkpointer, Command, RetryPolicy
from langgraph.types import _DC_KWARGS, All, Checkpointer, Command, N, RetryPolicy
from langgraph.utils.fields import get_field_default
from langgraph.utils.pydantic import create_model
from langgraph.utils.runnable import RunnableCallable, coerce_to_runnable
logger = logging.getLogger(__name__)
N = TypeVar("N")
def _warn_invalid_state_schema(schema: Union[Type[Any], Any]) -> None:
if isinstance(schema, type):
@@ -81,20 +79,20 @@ def _get_node_name(node: RunnableLike) -> str:
raise TypeError(f"Unsupported node type: {type(node)}")
class GraphCommand(Command, Generic[N]):
@dataclasses.dataclass(**_DC_KWARGS)
class GraphCommand(Generic[N], Command[N]):
"""One or more commands to update a StateGraph's state and go to, or send messages to nodes."""
__slots__ = ("goto",)
goto: Union[str, Sequence[str]] = ()
def __init__(
self,
*,
update: Optional[dict[str, Any]] = None,
goto: Union[str, Sequence[str]] = (),
send: Union[Send, Sequence[Send]] = (),
) -> None:
super().__init__(update=update, send=send)
self.goto = goto
def __repr__(self) -> str:
# get all non-None values
contents = ", ".join(
f"{key}={value!r}"
for key, value in dataclasses.asdict(self).items()
if value
)
return f"Command({contents})"
class StateNodeSpec(NamedTuple):
@@ -389,7 +387,7 @@ class StateGraph(Graph):
input = input_hint
if (
(rtn := hints.get("return"))
and get_origin(rtn) is GraphCommand
and get_origin(rtn) in (Command, GraphCommand)
and (rargs := get_args(rtn))
and get_origin(rargs[0]) is Literal
and (vals := get_args(rargs[0]))
+112 -17
View File
@@ -66,9 +66,12 @@ from langgraph.constants import (
CONFIG_KEY_STREAM_WRITER,
CONFIG_KEY_TASK_ID,
ERROR,
INPUT,
INTERRUPT,
NS_END,
NS_SEP,
NULL_TASK_ID,
PUSH,
SCHEDULED,
)
from langgraph.errors import (
@@ -98,7 +101,13 @@ from langgraph.pregel.utils import find_subgraph_pregel, get_new_channel_version
from langgraph.pregel.validate import validate_graph, validate_keys
from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
from langgraph.store.base import BaseStore
from langgraph.types import All, Checkpointer, LoopProtocol, StateSnapshot, StreamMode
from langgraph.types import (
All,
Checkpointer,
LoopProtocol,
StateSnapshot,
StreamMode,
)
from langgraph.utils.config import (
ensure_config,
merge_configs,
@@ -468,6 +477,7 @@ class Pregel(PregelProtocol):
# tasks for this checkpoint
next_tasks = prepare_next_tasks(
saved.checkpoint,
saved.pending_writes or [],
self.nodes,
channels,
managed,
@@ -511,6 +521,15 @@ class Pregel(PregelProtocol):
config, subgraphs=True
)
# apply pending writes
if null_writes := [
w[1:] for w in saved.pending_writes or [] if w[0] == NULL_TASK_ID
]:
apply_writes(
saved.checkpoint,
channels,
[PregelTaskWrites((), INPUT, null_writes, [])],
None,
)
if apply_pending_writes and saved.pending_writes:
for tid, k, v in saved.pending_writes:
if k in (ERROR, INTERRUPT, SCHEDULED):
@@ -570,6 +589,7 @@ class Pregel(PregelProtocol):
# tasks for this checkpoint
next_tasks = prepare_next_tasks(
saved.checkpoint,
saved.pending_writes or [],
self.nodes,
channels,
managed,
@@ -613,6 +633,15 @@ class Pregel(PregelProtocol):
config, subgraphs=True
)
# apply pending writes
if null_writes := [
w[1:] for w in saved.pending_writes or [] if w[0] == NULL_TASK_ID
]:
apply_writes(
saved.checkpoint,
channels,
[PregelTaskWrites((), INPUT, null_writes, [])],
None,
)
if apply_pending_writes and saved.pending_writes:
for tid, k, v in saved.pending_writes:
if k in (ERROR, INTERRUPT, SCHEDULED):
@@ -878,6 +907,7 @@ class Pregel(PregelProtocol):
# tasks for this checkpoint
next_tasks = prepare_next_tasks(
checkpoint,
saved.pending_writes or [],
self.nodes,
channels,
managed,
@@ -888,6 +918,18 @@ class Pregel(PregelProtocol):
checkpointer=self.checkpointer or None,
manager=None,
)
# apply null writes
if null_writes := [
w[1:]
for w in saved.pending_writes or []
if w[0] == NULL_TASK_ID
]:
apply_writes(
saved.checkpoint,
channels,
[PregelTaskWrites((), INPUT, null_writes, [])],
None,
)
# apply writes from tasks that already ran
for tid, k, v in saved.pending_writes or []:
if k in (ERROR, INTERRUPT, SCHEDULED):
@@ -922,6 +964,7 @@ class Pregel(PregelProtocol):
# tasks for this checkpoint
next_tasks = prepare_next_tasks(
checkpoint,
saved.pending_writes,
self.nodes,
channels,
managed,
@@ -932,6 +975,16 @@ class Pregel(PregelProtocol):
checkpointer=self.checkpointer or None,
manager=None,
)
# apply null writes
if null_writes := [
w[1:] for w in saved.pending_writes or [] if w[0] == NULL_TASK_ID
]:
apply_writes(
saved.checkpoint,
channels,
[PregelTaskWrites((), INPUT, null_writes, [])],
None,
)
# apply writes
for tid, k, v in saved.pending_writes:
if k in (ERROR, INTERRUPT, SCHEDULED):
@@ -1001,8 +1054,14 @@ class Pregel(PregelProtocol):
),
)
# save task writes
if saved:
checkpointer.put_writes(checkpoint_config, task.writes, task_id)
# channel writes are saved to current checkpoint
# push writes are saved to next checkpoint
channel_writes, push_writes = (
[w for w in task.writes if w[0] != PUSH],
[w for w in task.writes if w[0] == PUSH],
)
if saved and channel_writes:
checkpointer.put_writes(checkpoint_config, channel_writes, task_id)
# apply to checkpoint and save
mv_writes = apply_writes(
checkpoint, channels, [task], checkpointer.get_next_version
@@ -1023,6 +1082,8 @@ class Pregel(PregelProtocol):
checkpoint_previous_versions, checkpoint["channel_versions"]
),
)
if push_writes:
checkpointer.put_writes(next_config, push_writes, task_id)
return patch_checkpoint_map(next_config, saved.metadata if saved else None)
async def aupdate_state(
@@ -1088,6 +1149,7 @@ class Pregel(PregelProtocol):
# tasks for this checkpoint
next_tasks = prepare_next_tasks(
checkpoint,
saved.pending_writes or [],
self.nodes,
channels,
managed,
@@ -1098,6 +1160,18 @@ class Pregel(PregelProtocol):
checkpointer=self.checkpointer or None,
manager=None,
)
# apply null writes
if null_writes := [
w[1:]
for w in saved.pending_writes or []
if w[0] == NULL_TASK_ID
]:
apply_writes(
saved.checkpoint,
channels,
[PregelTaskWrites((), INPUT, null_writes, [])],
None,
)
# apply writes from tasks that already ran
for tid, k, v in saved.pending_writes or []:
if k in (ERROR, INTERRUPT, SCHEDULED):
@@ -1132,6 +1206,7 @@ class Pregel(PregelProtocol):
# tasks for this checkpoint
next_tasks = prepare_next_tasks(
checkpoint,
saved.pending_writes,
self.nodes,
channels,
managed,
@@ -1142,6 +1217,16 @@ class Pregel(PregelProtocol):
checkpointer=self.checkpointer or None,
manager=None,
)
# apply null writes
if null_writes := [
w[1:] for w in saved.pending_writes or [] if w[0] == NULL_TASK_ID
]:
apply_writes(
saved.checkpoint,
channels,
[PregelTaskWrites((), INPUT, null_writes, [])],
None,
)
for tid, k, v in saved.pending_writes:
if k in (ERROR, INTERRUPT, SCHEDULED):
continue
@@ -1208,14 +1293,23 @@ class Pregel(PregelProtocol):
),
)
# save task writes
if saved:
await checkpointer.aput_writes(checkpoint_config, writes, task_id)
# channel writes are saved to current checkpoint
# push writes are saved to next checkpoint
channel_writes, push_writes = (
[w for w in task.writes if w[0] != PUSH],
[w for w in task.writes if w[0] == PUSH],
)
if saved and channel_writes:
await checkpointer.aput_writes(
checkpoint_config, channel_writes, task_id
)
# apply to checkpoint and save
mv_writes = apply_writes(
checkpoint, channels, [task], checkpointer.get_next_version
)
assert not mv_writes, "Can't write to SharedValues from update_state"
checkpoint = create_checkpoint(checkpoint, channels, step + 1)
# save checkpoint, after applying writes
next_config = await checkpointer.aput(
checkpoint_config,
checkpoint,
@@ -1230,6 +1324,9 @@ class Pregel(PregelProtocol):
checkpoint_previous_versions, checkpoint["channel_versions"]
),
)
# save push writes
if push_writes:
await checkpointer.aput_writes(next_config, push_writes, task_id)
return patch_checkpoint_map(next_config, saved.metadata if saved else None)
def _defaults(
@@ -1432,12 +1529,16 @@ class Pregel(PregelProtocol):
specs=self.channels,
output_keys=output_keys,
stream_keys=self.stream_channels_asis,
interrupt_before=interrupt_before_,
interrupt_after=interrupt_after_,
manager=run_manager,
debug=debug,
) as loop:
# create runner
runner = PregelRunner(
submit=loop.submit,
put_writes=loop.put_writes,
schedule_task=loop.accept_push,
node_finished=config[CONF].get(CONFIG_KEY_NODE_FINISHED),
)
# enable subgraph streaming
@@ -1468,12 +1569,7 @@ class Pregel(PregelProtocol):
# channel updates from step N are only visible in step N+1
# channels are guaranteed to be immutable for the duration of the step,
# with channel updates applied only at the transition between steps
while loop.tick(
input_keys=self.input_channels,
interrupt_before=interrupt_before_,
interrupt_after=interrupt_after_,
manager=run_manager,
):
while loop.tick(input_keys=self.input_channels):
for _ in runner.tick(
loop.tasks.values(),
timeout=self.step_timeout,
@@ -1654,11 +1750,15 @@ class Pregel(PregelProtocol):
specs=self.channels,
output_keys=output_keys,
stream_keys=self.stream_channels_asis,
interrupt_before=interrupt_before_,
interrupt_after=interrupt_after_,
manager=run_manager,
) as loop:
# create runner
runner = PregelRunner(
submit=loop.submit,
put_writes=loop.put_writes,
schedule_task=loop.accept_push,
use_astream=do_stream is not None,
node_finished=config[CONF].get(CONFIG_KEY_NODE_FINISHED),
)
@@ -1678,12 +1778,7 @@ class Pregel(PregelProtocol):
# channel updates from step N are only visible in step N+1
# channels are guaranteed to be immutable for the duration of the step,
# with channel updates applied only at the transition between steps
while loop.tick(
input_keys=self.input_channels,
interrupt_before=interrupt_before_,
interrupt_after=interrupt_after_,
manager=run_manager,
):
while loop.tick(input_keys=self.input_channels):
async for _ in runner.atick(
loop.tasks.values(),
timeout=self.step_timeout,
+191 -46
View File
@@ -25,6 +25,7 @@ from langgraph.channels.base import BaseChannel
from langgraph.checkpoint.base import (
BaseCheckpointSaver,
Checkpoint,
PendingWrite,
V,
copy_checkpoint,
)
@@ -35,17 +36,21 @@ from langgraph.constants import (
CONFIG_KEY_CHECKPOINT_NS,
CONFIG_KEY_CHECKPOINTER,
CONFIG_KEY_READ,
CONFIG_KEY_RESUME_VALUE,
CONFIG_KEY_SEND,
CONFIG_KEY_STORE,
CONFIG_KEY_TASK_ID,
EMPTY_SEQ,
INTERRUPT,
MISSING,
NO_WRITES,
NS_END,
NS_SEP,
NULL_TASK_ID,
PULL,
PUSH,
RESERVED,
RESUME,
TAG_HIDDEN,
TASKS,
Send,
@@ -68,7 +73,7 @@ class WritesProtocol(Protocol):
Implemented by PregelTaskWrites and PregelExecutableTask."""
@property
def path(self) -> tuple[Union[str, int], ...]: ...
def path(self) -> tuple[Union[str, int, tuple], ...]: ...
@property
def name(self) -> str: ...
@@ -84,7 +89,7 @@ class PregelTaskWrites(NamedTuple):
"""Simplest implementation of WritesProtocol, for usage with writes that
don't originate from a runnable task, eg. graph input, update_state, etc."""
path: tuple[Union[str, int], ...]
path: tuple[Union[str, int, tuple], ...]
name: str
writes: Sequence[tuple[str, Any]]
triggers: Sequence[str]
@@ -172,7 +177,7 @@ def local_write(
"""Function injected under CONFIG_KEY_SEND in task config, to write to channels.
Validates writes and forwards them to `commit` function."""
for chan, value in writes:
if chan == TASKS:
if chan in (PUSH, TASKS):
if not isinstance(value, Send):
raise InvalidUpdateError(f"Expected Send, got {value}")
if value.node not in process_keys:
@@ -194,8 +199,13 @@ def apply_writes(
"""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."""
# sort tasks on path
tasks = sorted(tasks, key=lambda t: t.path)
# 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)
tasks = sorted(tasks, key=lambda t: t.path[:3])
# if no task has triggers this is applying writes from the null task only
# so we don't do anything other than update the channels written to
bump_step = any(t.triggers for t in tasks)
# update seen versions
for task in tasks:
@@ -227,7 +237,7 @@ def apply_writes(
)
# clear pending sends
if checkpoint["pending_sends"]:
if checkpoint["pending_sends"] and bump_step:
checkpoint["pending_sends"].clear()
# Group writes by channel
@@ -235,9 +245,9 @@ def apply_writes(
pending_writes_by_managed: dict[str, list[Any]] = defaultdict(list)
for task in tasks:
for chan, val in task.writes:
if chan == NO_WRITES:
if chan in (NO_WRITES, PUSH, RESUME, INTERRUPT):
pass
elif chan == TASKS:
elif chan == TASKS: # TODO: remove branch in 1.0
checkpoint["pending_sends"].append(val)
elif chan in channels:
pending_writes_by_channel[chan].append(val)
@@ -262,13 +272,14 @@ def apply_writes(
updated_channels.add(chan)
# Channels that weren't updated in this step are notified of a new step
for chan in channels:
if chan not in updated_channels:
if channels[chan].update([]) and get_next_version is not None:
checkpoint["channel_versions"][chan] = get_next_version(
max_version,
channels[chan],
)
if bump_step:
for chan in channels:
if chan not in updated_channels:
if channels[chan].update([]) and get_next_version is not None:
checkpoint["channel_versions"][chan] = get_next_version(
max_version,
channels[chan],
)
# Return managed values writes to be applied externally
return pending_writes_by_managed
@@ -277,6 +288,7 @@ def apply_writes(
@overload
def prepare_next_tasks(
checkpoint: Checkpoint,
pending_writes: Sequence[PendingWrite],
processes: Mapping[str, PregelNode],
channels: Mapping[str, BaseChannel],
managed: ManagedValueMapping,
@@ -293,6 +305,7 @@ def prepare_next_tasks(
@overload
def prepare_next_tasks(
checkpoint: Checkpoint,
pending_writes: Sequence[PendingWrite],
processes: Mapping[str, PregelNode],
channels: Mapping[str, BaseChannel],
managed: ManagedValueMapping,
@@ -308,6 +321,7 @@ def prepare_next_tasks(
def prepare_next_tasks(
checkpoint: Checkpoint,
pending_writes: Sequence[PendingWrite],
processes: Mapping[str, PregelNode],
channels: Mapping[str, BaseChannel],
managed: ManagedValueMapping,
@@ -322,13 +336,14 @@ def prepare_next_tasks(
"""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
by edges)."""
tasks: dict[str, Union[PregelTask, PregelExecutableTask]] = {}
# Consume pending packets
for idx, _ in enumerate(checkpoint["pending_sends"]):
tasks: list[Union[PregelTask, PregelExecutableTask]] = []
# Consume pending_sends from previous step (legacy version of Send)
for idx, _ in enumerate(checkpoint["pending_sends"]): # TODO: remove branch in 1.0
if task := prepare_single_task(
(PUSH, idx),
None,
checkpoint=checkpoint,
pending_writes=pending_writes,
processes=processes,
channels=channels,
managed=managed,
@@ -339,7 +354,7 @@ def prepare_next_tasks(
checkpointer=checkpointer,
manager=manager,
):
tasks[task.id] = task
tasks.append(task)
# Check if any processes should be run in next step
# If so, prepare the values to be passed to them
for name in processes:
@@ -347,6 +362,7 @@ def prepare_next_tasks(
(PULL, name),
None,
checkpoint=checkpoint,
pending_writes=pending_writes,
processes=processes,
channels=channels,
managed=managed,
@@ -357,15 +373,74 @@ def prepare_next_tasks(
checkpointer=checkpointer,
manager=manager,
):
tasks[task.id] = task
return tasks
tasks.append(task)
# Consume pending Sends from this step (new version of Send)
if any(c == PUSH for _, c, _ in pending_writes):
# group writes by task id
grouped_by_task = defaultdict(list)
for tid, c, _ in pending_writes:
grouped_by_task[tid].append(c)
# prepare send tasks from grouped writes
# 1. start from sends originating from existing tasks
tidx = 0
while tidx < len(tasks):
task = tasks[tidx]
if twrites := grouped_by_task.pop(task.id, None):
for idx, c in enumerate(twrites):
if c != PUSH:
continue
if next_task := prepare_single_task(
(PUSH, task.path, idx, task.id),
None,
checkpoint=checkpoint,
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(next_task)
tidx += 1
# key tasks by id
task_map = {t.id: t for t in tasks}
# 2. create new tasks for remaining sends (eg. from update_state)
for tid, writes in grouped_by_task.items():
task = task_map.get(tid)
for idx, c in enumerate(writes):
if c != PUSH:
continue
if next_task := prepare_single_task(
(PUSH, task.path if task else (), idx, tid),
None,
checkpoint=checkpoint,
pending_writes=pending_writes,
processes=processes,
channels=channels,
managed=managed,
config=config,
step=step,
for_execution=for_execution,
store=store,
checkpointer=checkpointer,
manager=manager,
):
task_map[next_task.id] = next_task
else:
task_map = {t.id: t for t in tasks}
return task_map
def prepare_single_task(
task_path: tuple[str, Union[int, str]],
task_path: tuple[Union[str, int, tuple], ...],
task_id_checksum: Optional[str],
*,
checkpoint: Checkpoint,
pending_writes: Sequence[PendingWrite],
processes: Mapping[str, PregelNode],
channels: Mapping[str, BaseChannel],
managed: ManagedValueMapping,
@@ -383,31 +458,74 @@ def prepare_single_task(
parent_ns = configurable.get(CONFIG_KEY_CHECKPOINT_NS, "")
if task_path[0] == PUSH:
idx = int(task_path[1])
if idx >= len(checkpoint["pending_sends"]):
return
packet = checkpoint["pending_sends"][idx]
if not isinstance(packet, Send):
logger.warning(
f"Ignoring invalid packet type {type(packet)} in pending sends"
if len(task_path) == 2: # TODO: remove branch in 1.0
# legacy SEND tasks, executed in superstep n+1
# (PUSH, idx of pending send)
idx = cast(int, task_path[1])
if idx >= len(checkpoint["pending_sends"]):
return
packet = checkpoint["pending_sends"][idx]
if not isinstance(packet, Send):
logger.warning(
f"Ignoring invalid packet type {type(packet)} in pending sends"
)
return
if packet.node not in processes:
logger.warning(
f"Ignoring unknown node name {packet.node} in pending sends"
)
return
# create task id
triggers = [PUSH]
checkpoint_ns = (
f"{parent_ns}{NS_SEP}{packet.node}" if parent_ns else packet.node
)
task_id = _uuid5_str(
checkpoint_id,
checkpoint_ns,
str(step),
packet.node,
PUSH,
str(idx),
)
elif len(task_path) == 4:
# new PUSH tasks, executed in superstep n
# (PUSH, parent task path, idx of PUSH write, id of parent task)
task_path_t = cast(tuple[str, tuple, int, str], task_path)
writes_for_path = [w for w in pending_writes if w[0] == task_path_t[3]]
if task_path_t[2] >= len(writes_for_path):
logger.warning(
f"Ignoring invalid write index {task_path[2]} in pending writes"
)
return
packet = writes_for_path[task_path_t[2]][2]
if not isinstance(packet, Send):
logger.warning(
f"Ignoring invalid packet type {type(packet)} in pending writes"
)
return
if packet.node not in processes:
logger.warning(
f"Ignoring unknown node name {packet.node} in pending writes"
)
return
# create task id
triggers = [PUSH]
checkpoint_ns = (
f"{parent_ns}{NS_SEP}{packet.node}" if parent_ns else packet.node
)
task_id = _uuid5_str(
checkpoint_id,
checkpoint_ns,
str(step),
packet.node,
PUSH,
_tuple_str(task_path[1]),
str(task_path[2]),
)
else:
logger.warning(f"Ignoring invalid PUSH task path {task_path}")
return
if packet.node not in processes:
logger.warning(f"Ignoring unknown node name {packet.node} in pending sends")
return
# create task id
triggers = [PUSH]
checkpoint_ns = (
f"{parent_ns}{NS_SEP}{packet.node}" if parent_ns else packet.node
)
task_id = _uuid5_str(
checkpoint_id,
checkpoint_ns,
str(step),
packet.node,
PUSH,
str(idx),
)
task_checkpoint_ns = f"{checkpoint_ns}:{task_id}"
metadata = {
"langgraph_step": step,
@@ -417,7 +535,7 @@ def prepare_single_task(
"langgraph_checkpoint_ns": task_checkpoint_ns,
}
if task_id_checksum is not None:
assert task_id == task_id_checksum
assert task_id == task_id_checksum, f"{task_id} != {task_id_checksum}"
if for_execution:
proc = processes[packet.node]
if node := proc.node:
@@ -469,6 +587,14 @@ def prepare_single_task(
},
CONFIG_KEY_CHECKPOINT_ID: None,
CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns,
CONFIG_KEY_RESUME_VALUE: next(
(
v
for tid, c, v in pending_writes
if tid in (NULL_TASK_ID, task_id) and c == RESUME
),
MISSING,
),
},
),
triggers,
@@ -481,6 +607,7 @@ def prepare_single_task(
else:
return PregelTask(task_id, packet.node, task_path)
elif task_path[0] == PULL:
# (PULL, node name)
name = cast(str, task_path[1])
if name not in processes:
return
@@ -577,6 +704,15 @@ def prepare_single_task(
},
CONFIG_KEY_CHECKPOINT_ID: None,
CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns,
CONFIG_KEY_RESUME_VALUE: next(
(
v
for tid, c, v in pending_writes
if tid in (NULL_TASK_ID, task_id)
and c == RESUME
),
MISSING,
),
},
),
triggers,
@@ -642,3 +778,12 @@ def _uuid5_str(namespace: bytes, *parts: str) -> str:
sha.update(b"".join(p.encode() for p in parts))
hex = sha.hexdigest()
return f"{hex[:8]}-{hex[8:12]}-{hex[12:16]}-{hex[16:20]}-{hex[20:32]}"
def _tuple_str(tup: Union[str, int, tuple]) -> str:
"""Generate a string representation of a tuple."""
return (
f"({', '.join(_tuple_str(x) for x in tup)})"
if isinstance(tup, (tuple, list))
else str(tup)
)
+1 -1
View File
@@ -208,7 +208,7 @@ def print_step_tasks(step: int, next_tasks: list[PregelExecutableTask]) -> None:
print(
f"{get_colored_text(f'[{step}:tasks]', color='blue')} "
+ get_bolded_text(
f"Starting step {step} with {n_tasks} task{'s' if n_tasks != 1 else ''}:\n"
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)}"
+52 -2
View File
@@ -1,11 +1,31 @@
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.constants import EMPTY_SEQ, ERROR, INTERRUPT, TAG_HIDDEN
from langgraph.constants import (
EMPTY_SEQ,
ERROR,
FF_SEND_V2,
INTERRUPT,
NULL_TASK_ID,
PUSH,
RESUME,
TAG_HIDDEN,
TASKS,
)
from langgraph.pregel.log import logger
from langgraph.types import PregelExecutableTask
from langgraph.types import Command, PregelExecutableTask, Send
def is_task_id(task_id: str) -> bool:
"""Check if a string is a valid task id."""
try:
UUID(task_id)
except ValueError:
return False
return True
def read_channel(
@@ -44,6 +64,36 @@ def read_channels(
return values
def map_command(
cmd: Command,
) -> Iterator[tuple[str, str, Any]]:
"""Map input chunk to a sequence of pending writes in the form (channel, value)."""
if cmd.send:
if isinstance(cmd.send, (tuple, list)):
sends = cmd.send
else:
sends = [cmd.send]
for send in sends:
if not isinstance(send, Send):
raise TypeError(
f"In Command.send, expected Send, got {type(send).__name__}"
)
yield (NULL_TASK_ID, PUSH if FF_SEND_V2 else TASKS, send)
if cmd.resume:
if isinstance(cmd.resume, dict) and all(is_task_id(k) for k in cmd.resume):
for tid, resume in cmd.resume.items():
yield (tid, RESUME, resume)
else:
yield (NULL_TASK_ID, RESUME, cmd.resume)
if cmd.update:
if not isinstance(cmd.update, dict):
raise TypeError(
f"Expected cmd.update to be a dict mapping channel names to update values, got {type(cmd.update).__name__}"
)
for k, v in cmd.update.items():
yield (NULL_TASK_ID, k, v)
def map_input(
input_channels: Union[str, Sequence[str]],
chunk: Optional[Union[dict[str, Any], Any]],
+147 -17
View File
@@ -1,6 +1,6 @@
import asyncio
import concurrent.futures
from collections import deque
from collections import defaultdict, deque
from contextlib import AsyncExitStack, ExitStack
from types import TracebackType
from typing import (
@@ -52,6 +52,9 @@ from langgraph.constants import (
INPUT,
INTERRUPT,
NS_SEP,
NULL_TASK_ID,
PUSH,
RESUME,
SCHEDULED,
TAG_HIDDEN,
)
@@ -74,6 +77,7 @@ from langgraph.pregel.algo import (
apply_writes,
increment,
prepare_next_tasks,
prepare_single_task,
should_interrupt,
)
from langgraph.pregel.debug import (
@@ -90,6 +94,7 @@ from langgraph.pregel.executor import (
Submit,
)
from langgraph.pregel.io import (
map_command,
map_input,
map_output_updates,
map_output_values,
@@ -100,7 +105,13 @@ from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager
from langgraph.pregel.read import PregelNode
from langgraph.pregel.utils import get_new_channel_versions
from langgraph.store.base import BaseStore
from langgraph.types import All, LoopProtocol, PregelExecutableTask, StreamProtocol
from langgraph.types import (
All,
Command,
LoopProtocol,
PregelExecutableTask,
StreamProtocol,
)
from langgraph.utils.config import patch_configurable
V = TypeVar("V")
@@ -130,6 +141,9 @@ class PregelLoop(LoopProtocol):
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]]
checkpointer_get_next_version: GetNextVersion
checkpointer_put_writes: Optional[
@@ -162,6 +176,7 @@ class PregelLoop(LoopProtocol):
"pending", "done", "interrupt_before", "interrupt_after", "out_of_steps"
]
tasks: dict[str, PregelExecutableTask]
to_interrupt: list[PregelExecutableTask]
output: Union[None, dict[str, Any], Any] = None
# public
@@ -178,6 +193,9 @@ class PregelLoop(LoopProtocol):
specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
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,
check_subgraphs: bool = True,
debug: bool = False,
) -> None:
@@ -194,6 +212,9 @@ class PregelLoop(LoopProtocol):
self.specs = specs
self.output_keys = output_keys
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]
@@ -261,15 +282,60 @@ class PregelLoop(LoopProtocol):
task_id,
)
# output writes
self._output_writes(task_id, writes)
if hasattr(self, "tasks"):
self._output_writes(task_id, writes)
def accept_push(
self, task: PregelExecutableTask, write_idx: int
) -> Optional[PregelExecutableTask]:
"""Accept a PUSH from a task, potentially returning a new task to start."""
# don't start if an earlier PUSH has already triggered an interrupt
if self.to_interrupt:
return
# don't start if we should interrupt *after* the original task
if 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),
None,
checkpoint=self.checkpoint,
pending_writes=[(task.id, *w) for w in task.writes],
processes=self.nodes,
channels=self.channels,
managed=self.managed,
config=self.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 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
if not pushed.writes:
return pushed
def tick(
self,
*,
input_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,
) -> bool:
"""Execute a single iteration of the Pregel loop.
Returns True if more iterations are needed."""
@@ -278,6 +344,10 @@ class PregelLoop(LoopProtocol):
if self.input not in (INPUT_DONE, INPUT_RESUMING):
self._first(input_keys=input_keys)
elif self.to_interrupt:
# if we need to interrupt, do so
self.status = "interrupt_before"
raise GraphInterrupt()
elif all(task.writes for task in self.tasks.values()):
writes = [w for t in self.tasks.values() for w in t.writes]
# debug flag
@@ -322,7 +392,9 @@ class PregelLoop(LoopProtocol):
}
)
# after execution, check if we should interrupt
if should_interrupt(self.checkpoint, interrupt_after, self.tasks.values()):
if should_interrupt(
self.checkpoint, self.interrupt_after, self.tasks.values()
):
self.status = "interrupt_after"
raise GraphInterrupt()
else:
@@ -333,19 +405,33 @@ class PregelLoop(LoopProtocol):
self.status = "out_of_steps"
return False
# apply NULL writes
if null_writes := [
w[1:] for w in self.checkpoint_pending_writes if w[0] == NULL_TASK_ID
]:
mv_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)
# prepare next tasks
self.tasks = prepare_next_tasks(
self.checkpoint,
self.checkpoint_pending_writes,
self.nodes,
self.channels,
self.managed,
self.config,
self.step,
for_execution=True,
manager=manager,
manager=self.manager,
store=self.store,
checkpointer=self.checkpointer,
)
self.to_interrupt = []
# produce debug output
if self._checkpointer_put_after_previous is not None:
@@ -387,15 +473,12 @@ class PregelLoop(LoopProtocol):
# if all tasks have finished, re-tick
if all(task.writes for task in self.tasks.values()):
return self.tick(
input_keys=input_keys,
interrupt_after=interrupt_after,
interrupt_before=interrupt_before,
manager=manager,
)
return self.tick(input_keys=input_keys)
# before execution, check if we should interrupt
if should_interrupt(self.checkpoint, interrupt_before, self.tasks.values()):
if should_interrupt(
self.checkpoint, self.interrupt_before, self.tasks.values()
):
self.status = "interrupt_before"
raise GraphInterrupt()
@@ -417,7 +500,7 @@ class PregelLoop(LoopProtocol):
def _match_writes(self, tasks: Mapping[str, PregelExecutableTask]) -> None:
for tid, k, v in self.checkpoint_pending_writes:
if k in (ERROR, INTERRUPT):
if k in (ERROR, INTERRUPT, RESUME):
continue
if task := tasks.get(tid):
if k == SCHEDULED:
@@ -449,8 +532,20 @@ class PregelLoop(LoopProtocol):
self._emit(
"values", map_output_values, self.output_keys, True, self.channels
)
# map command to writes
elif isinstance(self.input, Command):
writes: defaultdict[str, list[tuple[str, Any]]] = defaultdict(list)
# group writes by task ID
for tid, c, v in map_command(self.input):
writes[tid].append((c, v))
if not writes:
raise EmptyInputError("Received empty Command input")
# save writes
for tid, ws in writes.items():
self.put_writes(tid, ws)
# map inputs to channel updates
elif input_writes := deque(map_input(input_keys, self.input)):
# TODO shouldn't these writes be passed to put_writes too?
# check if we should delegate (used by subgraphs in distributed mode)
if self.config[CONF].get(CONFIG_KEY_DELEGATE):
raise GraphDelegate(
@@ -464,6 +559,7 @@ class PregelLoop(LoopProtocol):
# discard any unfinished tasks from previous checkpoint
discard_tasks = prepare_next_tasks(
self.checkpoint,
self.checkpoint_pending_writes,
self.nodes,
self.channels,
self.managed,
@@ -577,11 +673,33 @@ class PregelLoop(LoopProtocol):
# save final output
self.output = read_channels(self.channels, self.output_keys)
if suppress:
# suppress interrupt
# emit one last "values" event, with pending writes applied
if (
hasattr(self, "tasks")
and self.checkpoint_pending_writes
and any(task.writes for task in self.tasks.values())
):
mv_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,
self.output_keys,
[w for t in self.tasks.values() for w in t.writes],
self.channels,
)
# emit INTERRUPT event
self._emit(
"updates",
lambda: iter([{INTERRUPT: cast(GraphInterrupt, exc_value).args[0]}]),
)
# suppress interrupt
return True
def _emit(
@@ -635,6 +753,9 @@ class SyncPregelLoop(PregelLoop, ContextManager):
checkpointer: Optional[BaseCheckpointSaver],
nodes: Mapping[str, PregelNode],
specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
manager: Union[None, AsyncParentRunManager, ParentRunManager] = None,
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,
check_subgraphs: bool = True,
@@ -650,7 +771,10 @@ class SyncPregelLoop(PregelLoop, ContextManager):
specs=specs,
output_keys=output_keys,
stream_keys=stream_keys,
interrupt_after=interrupt_after,
interrupt_before=interrupt_before,
check_subgraphs=check_subgraphs,
manager=manager,
debug=debug,
)
self.stack = ExitStack()
@@ -761,6 +885,9 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
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,
check_subgraphs: bool = True,
@@ -776,7 +903,10 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
specs=specs,
output_keys=output_keys,
stream_keys=stream_keys,
interrupt_after=interrupt_after,
interrupt_before=interrupt_before,
check_subgraphs=check_subgraphs,
manager=manager,
debug=debug,
)
self.stack = AsyncExitStack()
+18 -2
View File
@@ -2,9 +2,15 @@ import asyncio
import logging
import random
import time
from typing import Optional, Sequence
from functools import partial
from typing import Any, Callable, Optional, Sequence
from langgraph.constants import CONF, CONFIG_KEY_CHECKPOINT_NS, CONFIG_KEY_RESUMING
from langgraph.constants import (
CONF,
CONFIG_KEY_CHECKPOINT_NS,
CONFIG_KEY_RESUMING,
CONFIG_KEY_SEND,
)
from langgraph.errors import _SEEN_CHECKPOINT_NS, GraphInterrupt
from langgraph.types import PregelExecutableTask, RetryPolicy
from langgraph.utils.config import patch_configurable
@@ -15,12 +21,17 @@ logger = logging.getLogger(__name__)
def run_with_retry(
task: PregelExecutableTask,
retry_policy: Optional[RetryPolicy],
writer: Optional[
Callable[[PregelExecutableTask, Sequence[tuple[str, Any]]], None]
] = 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 writer is not None:
config = patch_configurable(config, {CONFIG_KEY_SEND: partial(writer, task)})
while True:
try:
# clear any writes from previous attempts
@@ -84,12 +95,17 @@ async def arun_with_retry(
task: PregelExecutableTask,
retry_policy: Optional[RetryPolicy],
stream: bool = False,
writer: Optional[
Callable[[PregelExecutableTask, Sequence[tuple[str, Any]]], None]
] = 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 writer is not None:
config = patch_configurable(config, {CONFIG_KEY_SEND: partial(writer, task)})
while True:
try:
# clear any writes from previous attempts
+115 -27
View File
@@ -14,7 +14,15 @@ from typing import (
cast,
)
from langgraph.constants import ERROR, INTERRUPT, NO_WRITES, TAG_HIDDEN
from langgraph.constants import (
CONF,
CONFIG_KEY_SEND,
ERROR,
INTERRUPT,
NO_WRITES,
PUSH,
TAG_HIDDEN,
)
from langgraph.errors import GraphDelegate, GraphInterrupt
from langgraph.pregel.executor import Submit
from langgraph.pregel.retry import arun_with_retry, run_with_retry
@@ -31,6 +39,9 @@ class PregelRunner:
*,
submit: Submit,
put_writes: Callable[[str, Sequence[tuple[str, Any]]], None],
schedule_task: Callable[
[PregelExecutableTask, int], Optional[PregelExecutableTask]
],
use_astream: bool = False,
node_finished: Optional[Callable[[str], None]] = None,
) -> None:
@@ -38,6 +49,7 @@ 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,
@@ -48,27 +60,58 @@ 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]]
) -> None:
prev_length = len(task.writes)
# delegate to the underlying writer
task.config[CONF][CONFIG_KEY_SEND](writes)
for idx, w in enumerate(task.writes):
# find the index for the newly inserted writes
if idx < prev_length:
continue
assert writes[idx - prev_length] is w
# bail if not a PUSH write
if w[0] != PUSH:
continue
# schedule the next task, if the callback returns one
if next_task := self.schedule_task(task, idx):
# if the parent task was retried,
# the next task might already be running
if any(
t == next_task.id for t in futures.values() if t is not None
):
continue
# schedule the next task
futures[
self.submit(
run_with_retry,
next_task,
retry_policy,
writer=writer,
__reraise_on_exit__=reraise,
)
] = next_task
tasks = tuple(tasks)
futures: dict[concurrent.futures.Future, Optional[PregelExecutableTask]] = {}
# give control back to the caller
yield
# fast path if single task with no timeout and no waiter
if len(tasks) == 1 and timeout is None and get_waiter is None:
t = tasks[0]
try:
run_with_retry(t, retry_policy)
run_with_retry(t, retry_policy, writer=writer)
self.commit(t, None)
except Exception as exc:
self.commit(t, exc)
if reraise:
raise
return
if not futures: # maybe `t` schuduled another task
return
# add waiter task if requested
if get_waiter is not None:
futures: dict[concurrent.futures.Future, Optional[PregelExecutableTask]] = {
get_waiter(): None
}
else:
futures = {}
futures[get_waiter()] = None
# 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
@@ -79,10 +122,11 @@ class PregelRunner:
run_with_retry,
t,
retry_policy,
writer=writer,
__reraise_on_exit__=reraise,
)
] = t
all_futures = futures.copy()
done_futures: set[concurrent.futures.Future] = set()
end_time = timeout + time.monotonic() if timeout else None
while len(futures) > (1 if get_waiter is not None else 0):
done, inflight = concurrent.futures.wait(
@@ -99,6 +143,8 @@ class PregelRunner:
if inflight and get_waiter is not None:
futures[get_waiter()] = None
else:
# store for panic check
done_futures.add(fut)
# task finished, commit writes
self.commit(task, _exception(fut))
else:
@@ -110,7 +156,10 @@ class PregelRunner:
# give control back to the caller
yield
# panic on failure or timeout
_panic_or_proceed(all_futures, panic=reraise)
_panic_or_proceed(
done_futures.union(f for f, t in futures.items() if t is not None),
panic=reraise,
)
async def atick(
self,
@@ -121,28 +170,67 @@ class PregelRunner:
retry_policy: Optional[RetryPolicy] = None,
get_waiter: Optional[Callable[[], asyncio.Future[None]]] = None,
) -> AsyncIterator[None]:
def writer(
task: PregelExecutableTask, writes: Sequence[tuple[str, Any]]
) -> None:
prev_length = len(task.writes)
# delegate to the underlying writer
task.config[CONF][CONFIG_KEY_SEND](writes)
for idx, w in enumerate(task.writes):
# find the index for the newly inserted writes
if idx < prev_length:
continue
assert writes[idx - prev_length] is w
# bail if not a PUSH write
if w[0] != PUSH:
continue
# schedule the next task, if the callback returns one
if next_task := self.schedule_task(task, idx):
# if the parent task was retried,
# the next task might already be running
if any(
t == next_task.id for t in futures.values() if t is not None
):
continue
# schedule the next task
futures[
cast(
asyncio.Future,
self.submit(
arun_with_retry,
next_task,
retry_policy,
stream=self.use_astream,
writer=writer,
__name__=t.name,
__cancel_on_exit__=True,
__reraise_on_exit__=reraise,
),
)
] = next_task
loop = asyncio.get_event_loop()
tasks = tuple(tasks)
futures: dict[asyncio.Future, Optional[PregelExecutableTask]] = {}
# 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)
await arun_with_retry(
t, retry_policy, stream=self.use_astream, writer=writer
)
self.commit(t, None)
except Exception as exc:
self.commit(t, exc)
if reraise:
raise
return
if not futures: # maybe `t` schuduled another task
return
# add waiter task if requested
if get_waiter is not None:
futures: dict[asyncio.Future, Optional[PregelExecutableTask]] = {
get_waiter(): None
}
else:
futures = {}
futures[get_waiter()] = None
# 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
@@ -156,13 +244,14 @@ class PregelRunner:
t,
retry_policy,
stream=self.use_astream,
writer=writer,
__name__=t.name,
__cancel_on_exit__=True,
__reraise_on_exit__=reraise,
),
)
] = t
all_futures = futures.copy()
done_futures: set[asyncio.Future] = set()
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(
@@ -179,6 +268,8 @@ class PregelRunner:
if inflight and get_waiter is not None:
futures[get_waiter()] = None
else:
# store for panic check
done_futures.add(fut)
# task finished, commit writes
self.commit(task, _exception(fut))
else:
@@ -194,7 +285,9 @@ class PregelRunner:
fut.cancel()
# panic on failure or timeout
_panic_or_proceed(
all_futures, timeout_exc_cls=asyncio.TimeoutError, panic=reraise
done_futures.union(f for f, t in futures.items() if t is not None),
timeout_exc_cls=asyncio.TimeoutError,
panic=reraise,
)
def commit(
@@ -250,10 +343,7 @@ def _exception(
def _panic_or_proceed(
futs: Union[
dict[concurrent.futures.Future, Optional[PregelExecutableTask]],
dict[asyncio.Future, Optional[PregelExecutableTask]],
],
futs: Union[set[concurrent.futures.Future], set[asyncio.Future]],
*,
timeout_exc_cls: Type[Exception] = TimeoutError,
panic: bool = True,
@@ -261,10 +351,8 @@ def _panic_or_proceed(
"""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()
for fut, val in futs.items():
if val is None:
continue
elif fut.done():
for fut in futs:
if fut.done():
done.add(fut)
else:
inflight.add(fut)
+7 -3
View File
@@ -14,7 +14,7 @@ from typing import (
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.constants import CONF, CONFIG_KEY_SEND, FF_SEND_V2, PUSH, TASKS, Send
from langgraph.errors import InvalidUpdateError
from langgraph.utils.runnable import RunnableCallable
@@ -112,14 +112,18 @@ class ChannelWrite(RunnableCallable):
# validate
for w in writes:
if isinstance(w, ChannelWriteEntry):
if w.channel == TASKS:
if w.channel in (TASKS, PUSH):
raise InvalidUpdateError(
"Cannot write to the reserved channel TASKS"
)
if w.value is PASSTHROUGH:
raise InvalidUpdateError("PASSTHROUGH value must be replaced")
# split packets and entries
sends = [(TASKS, packet) for packet in writes if isinstance(packet, Send)]
sends = [
(PUSH if FF_SEND_V2 else TASKS, packet)
for packet in writes
if isinstance(packet, Send)
]
entries = [write for write in writes if isinstance(write, ChannelWriteEntry)]
# process entries into values
values = [
+50 -39
View File
@@ -1,14 +1,18 @@
import dataclasses
import sys
from collections import deque
from dataclasses import dataclass
from typing import (
TYPE_CHECKING,
Any,
Callable,
Generic,
Hashable,
Literal,
NamedTuple,
Optional,
Sequence,
Type,
TypeVar,
Union,
cast,
)
@@ -44,6 +48,11 @@ StreamWriter = Callable[[Any], None]
Always injected into nodes if requested as a keyword argument, but it's a no-op
when not using stream_mode="custom"."""
if sys.version_info >= (3, 10):
_DC_KWARGS = {"kw_only": True, "slots": True, "frozen": True}
else:
_DC_KWARGS = {"frozen": True}
def default_retry_on(exc: Exception) -> bool:
import httpx
@@ -101,16 +110,18 @@ class CachePolicy(NamedTuple):
pass
@dataclass
@dataclasses.dataclass(**_DC_KWARGS)
class Interrupt:
value: Any
resumable: bool = False
ns: Optional[Sequence[str]] = None
when: Literal["during"] = "during"
class PregelTask(NamedTuple):
id: str
name: str
path: tuple[Union[str, int], ...]
path: tuple[Union[str, int, tuple], ...]
error: Optional[Exception] = None
interrupts: tuple[Interrupt, ...] = ()
state: Union[None, RunnableConfig, "StateSnapshot"] = None
@@ -127,7 +138,7 @@ class PregelExecutableTask(NamedTuple):
retry_policy: Optional[RetryPolicy]
cache_policy: Optional[CachePolicy]
id: str
path: tuple[Union[str, int], ...]
path: tuple[Union[str, int, tuple], ...]
scheduled: bool = False
@@ -221,51 +232,26 @@ class Send:
)
class Command:
N = TypeVar("N", bound=Hashable)
@dataclasses.dataclass(**_DC_KWARGS)
class Command(Generic[N]):
"""One or more commands to update the graph's state and send messages to nodes."""
__slots__ = ("update", "send")
def __init__(
self,
*,
update: Optional[dict[str, Any]] = None,
send: Union[Send, Sequence[Send]] = (),
) -> None:
self.update = update
self.send = send
@property
def __all_slots__(self) -> set[str]:
# get all slots from mro
slots = set()
for cls in type(self).__mro__:
if ss := getattr(cls, "__slots__", ()):
if isinstance(ss, str):
slots.add(ss)
else:
slots.update(ss)
return slots
update: Optional[dict[str, Any]] = None
send: Union[Send, Sequence[Send]] = ()
resume: Optional[Union[Any, dict[str, Any]]] = None
def __repr__(self) -> str:
# get all non-None values
contents = ", ".join(
f"{key}={value!r}"
for key in self.__all_slots__
if (value := getattr(self, key))
for key, value in dataclasses.asdict(self).items()
if value
)
return f"Command({contents})"
def __eq__(self, value: Any) -> bool:
return type(value) is type(self) and all(
getattr(self, key) == getattr(value, key) for key in self.__all_slots__
)
def copy(self, **kwargs: Any) -> Self:
for slot in self.__all_slots__:
kwargs.setdefault(slot, getattr(self, slot))
return self.__class__(**kwargs)
StreamChunk = tuple[tuple[str, ...], str, Any]
@@ -307,3 +293,28 @@ class LoopProtocol:
self.store = store
self.step = step
self.stop = stop
def interrupt(value: Any) -> Any:
from langgraph.constants import (
CONFIG_KEY_CHECKPOINT_NS,
CONFIG_KEY_RESUME_VALUE,
MISSING,
NS_SEP,
)
from langgraph.errors import GraphInterrupt
from langgraph.utils.config import get_configurable
conf = get_configurable()
if (resume := conf.get(CONFIG_KEY_RESUME_VALUE, MISSING)) and resume is not MISSING:
return resume
else:
raise GraphInterrupt(
(
Interrupt(
value=value,
resumable=True,
ns=cast(str, conf[CONFIG_KEY_CHECKPOINT_NS]).split(NS_SEP),
),
)
)
+17
View File
@@ -1,3 +1,5 @@
import asyncio
import sys
from collections import ChainMap
from typing import Any, Optional, Sequence
@@ -290,3 +292,18 @@ def ensure_config(*configs: Optional[RunnableConfig]) -> RunnableConfig:
):
empty["metadata"][key] = value
return empty
def get_configurable() -> dict[str, Any]:
if sys.version_info < (3, 11):
try:
if asyncio.current_task():
raise RuntimeError(
"Python 3.11 or later required to use this in an async context"
)
except RuntimeError:
pass
if var_config := var_child_runnable_config.get():
return var_config[CONF]
else:
raise RuntimeError("Called get_configurable outside of a runnable context")
+3 -3
View File
@@ -1359,7 +1359,7 @@ typing-extensions = ">=4.7"
[[package]]
name = "langgraph-checkpoint"
version = "2.0.2"
version = "2.0.4"
description = "Library with base interfaces for LangGraph checkpoint savers."
optional = false
python-versions = "^3.9.0,<4.0"
@@ -1429,7 +1429,7 @@ url = "../checkpoint-sqlite"
[[package]]
name = "langgraph-sdk"
version = "0.1.35"
version = "0.1.36"
description = "SDK for interacting with LangGraph API"
optional = false
python-versions = "^3.9.0,<4.0"
@@ -3425,4 +3425,4 @@ type = ["pytest-mypy"]
[metadata]
lock-version = "2.0"
python-versions = ">=3.9.0,<4.0"
content-hash = "f0b7800f90041227ba4d7b72482c70fabdbacc95e8a3a784fdebc88721728a91"
content-hash = "9bf5668d3f70f3b77457906732404a6401583a5966f70a72ef10a68f2a5b27ad"
+2 -2
View File
@@ -1,6 +1,6 @@
[tool.poetry]
name = "langgraph"
version = "0.2.46"
version = "0.2.48"
description = "Building stateful, multi-actor applications with LLMs"
authors = []
license = "MIT"
@@ -10,7 +10,7 @@ repository = "https://www.github.com/langchain-ai/langgraph"
[tool.poetry.dependencies]
python = ">=3.9.0,<4.0"
langchain-core = ">=0.2.43,<0.4.0,!=0.3.0,!=0.3.1,!=0.3.2,!=0.3.3,!=0.3.4,!=0.3.5,!=0.3.6,!=0.3.7,!=0.3.8,!=0.3.9,!=0.3.10,!=0.3.11,!=0.3.12,!=0.3.13,!=0.3.14"
langgraph-checkpoint = "^2.0.0"
langgraph-checkpoint = "^2.0.4"
langgraph-sdk = "^0.1.32"
[tool.poetry.group.dev.dependencies]
@@ -5108,6 +5108,81 @@
'''
# ---
# name: test_send_react_interrupt_control[memory]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
agent(agent)
foo([foo]):::last
__start__ --> agent;
agent -.-> foo;
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_send_react_interrupt_control[postgres]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
agent(agent)
foo([foo]):::last
__start__ --> agent;
agent -.-> foo;
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_send_react_interrupt_control[postgres_pipe]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
agent(agent)
foo([foo]):::last
__start__ --> agent;
agent -.-> foo;
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_send_react_interrupt_control[postgres_pool]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
agent(agent)
foo([foo]):::last
__start__ --> agent;
agent -.-> foo;
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_send_react_interrupt_control[sqlite]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
agent(agent)
foo([foo]):::last
__start__ --> agent;
agent -.-> foo;
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_simple_multi_edge
'''
graph TD;
@@ -1302,6 +1302,81 @@
+---------+
'''
# ---
# name: test_send_react_interrupt_control[memory]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
agent(agent)
foo([foo]):::last
__start__ --> agent;
agent -.-> foo;
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_send_react_interrupt_control[postgres_aio]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
agent(agent)
foo([foo]):::last
__start__ --> agent;
agent -.-> foo;
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_send_react_interrupt_control[postgres_aio_pipe]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
agent(agent)
foo([foo]):::last
__start__ --> agent;
agent -.-> foo;
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_send_react_interrupt_control[postgres_aio_pool]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
agent(agent)
foo([foo]):::last
__start__ --> agent;
agent -.-> foo;
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_send_react_interrupt_control[sqlite_aio]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
agent(agent)
foo([foo]):::last
__start__ --> agent;
agent -.-> foo;
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_weather_subgraph[duckdb_aio]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
-2
View File
@@ -327,7 +327,6 @@ async def awith_store(store_name: Optional[str]) -> AsyncIterator[BaseStore]:
ALL_CHECKPOINTERS_SYNC = [
"memory",
"sqlite",
"duckdb",
"postgres",
"postgres_pipe",
"postgres_pool",
@@ -335,7 +334,6 @@ ALL_CHECKPOINTERS_SYNC = [
ALL_CHECKPOINTERS_ASYNC = [
"memory",
"sqlite_aio",
"duckdb_aio",
"postgres_aio",
"postgres_aio_pipe",
"postgres_aio_pool",
+9 -1
View File
@@ -11,13 +11,21 @@ def test_prepare_next_tasks() -> None:
with ChannelsManager({}, checkpoint, config) as (channels, managed):
assert (
prepare_next_tasks(
checkpoint, processes, channels, managed, config, 0, for_execution=False
checkpoint,
{},
processes,
channels,
managed,
config,
0,
for_execution=False,
)
== {}
)
assert (
prepare_next_tasks(
checkpoint,
{},
processes,
channels,
managed,
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
+1 -1
View File
@@ -19,7 +19,7 @@ test:
exit $$EXIT_CODE
test_watch:
make start-services && poetry run ptw . -- $(TEST_PATH); \
make start-services && poetry run ptw . -- -x $(TEST_PATH); \
EXIT_CODE=$$?; \
make stop-services; \
exit $$EXIT_CODE
@@ -38,7 +38,7 @@ from langgraph.scheduler.kafka.types import (
Sendable,
Topics,
)
from langgraph.types import LoopProtocol, RetryPolicy
from langgraph.types import LoopProtocol, PregelExecutableTask, RetryPolicy
from langgraph.utils.config import patch_configurable
@@ -198,6 +198,7 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager):
msg["task"]["path"],
msg["task"]["id"],
checkpoint=saved.checkpoint,
pending_writes=saved.pending_writes or [],
processes=graph.nodes,
channels=channels,
managed=managed,
@@ -211,6 +212,7 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager):
runner = PregelRunner(
submit=submit,
put_writes=partial(self._put_writes, submit, msg["config"]),
schedule_task=self._schedule_task,
)
async for _ in runner.atick([task], reraise=False):
pass
@@ -239,6 +241,14 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager):
)
await fut
def _schedule_task(
self,
task: PregelExecutableTask,
idx: int,
) -> None:
# will be scheduled by orchestrator when executor finishes
pass
def _put_writes(
self,
submit: Submit,
@@ -400,6 +410,7 @@ class KafkaExecutor(AbstractContextManager):
msg["task"]["path"],
msg["task"]["id"],
checkpoint=saved.checkpoint,
pending_writes=saved.pending_writes or [],
processes=graph.nodes,
channels=channels,
managed=managed,
@@ -412,6 +423,7 @@ class KafkaExecutor(AbstractContextManager):
runner = PregelRunner(
submit=submit,
put_writes=partial(self._put_writes, submit, msg["config"]),
schedule_task=self._schedule_task,
)
for _ in runner.tick([task], reraise=False):
pass
@@ -440,6 +452,14 @@ class KafkaExecutor(AbstractContextManager):
)
fut.result()
def _schedule_task(
self,
task: PregelExecutableTask,
idx: int,
) -> None:
# will be scheduled by orchestrator when executor finishes
pass
def _put_writes(
self,
submit: Submit,
@@ -161,18 +161,18 @@ class AsyncKafkaOrchestrator(AbstractAsyncContextManager):
specs=graph.channels,
output_keys=graph.output_channels,
stream_keys=graph.stream_channels,
interrupt_after=graph.interrupt_after_nodes,
interrupt_before=graph.interrupt_before_nodes,
check_subgraphs=False,
) as loop:
if loop.tick(
input_keys=graph.input_channels,
interrupt_after=graph.interrupt_after_nodes,
interrupt_before=graph.interrupt_before_nodes,
):
if loop.tick(input_keys=graph.input_channels):
# wait for checkpoint to be saved
if hasattr(loop, "_put_checkpoint_fut"):
await loop._put_checkpoint_fut
# schedule any new tasks
if new_tasks := [t for t in loop.tasks.values() if not t.scheduled]:
if new_tasks := [
t for t in loop.tasks.values() if not t.scheduled and not t.writes
]:
# send messages to executor
futures = await asyncio.gather(
*(
@@ -351,18 +351,18 @@ class KafkaOrchestrator(AbstractContextManager):
specs=graph.channels,
output_keys=graph.output_channels,
stream_keys=graph.stream_channels,
interrupt_after=graph.interrupt_after_nodes,
interrupt_before=graph.interrupt_before_nodes,
check_subgraphs=False,
) as loop:
if loop.tick(
input_keys=graph.input_channels,
interrupt_after=graph.interrupt_after_nodes,
interrupt_before=graph.interrupt_before_nodes,
):
if loop.tick(input_keys=graph.input_channels):
# wait for checkpoint to be saved
if hasattr(loop, "_put_checkpoint_fut"):
loop._put_checkpoint_fut.result()
# schedule any new tasks
if new_tasks := [t for t in loop.tasks.values() if not t.scheduled]:
if new_tasks := [
t for t in loop.tasks.values() if not t.scheduled and not t.writes
]:
# send messages to executor
futures = [
self.producer.send(
@@ -24,8 +24,8 @@ class MessageToOrchestrator(TypedDict):
class ExecutorTask(TypedDict):
id: str
path: tuple[str, ...]
id: Optional[str]
path: tuple[Union[str, int], ...]
class MessageToExecutor(TypedDict):
+208
View File
@@ -0,0 +1,208 @@
import operator
from typing import (
Annotated,
Literal,
Union,
)
import pytest
from aiokafka import AIOKafkaProducer
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.constants import FF_SEND_V2, START
from langgraph.errors import NodeInterrupt
from langgraph.graph.state import CompiledStateGraph, GraphCommand, StateGraph
from langgraph.scheduler.kafka import serde
from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics
from langgraph.types import Send
from tests.any import AnyDict
from tests.drain import drain_topics_async
pytestmark = pytest.mark.anyio
def mk_push_graph(
checkpointer: BaseCheckpointSaver,
) -> CompiledStateGraph:
# copied from test_send_dedupe_on_resume
class InterruptOnce:
ticks: int = 0
def __call__(self, state):
self.ticks += 1
if self.ticks == 1:
raise NodeInterrupt("Bahh")
return ["|".join(("flaky", str(state)))]
class Node:
def __init__(self, name: str):
self.name = name
self.ticks = 0
self.__name__ = name
def __call__(self, state):
self.ticks += 1
update = (
[self.name]
if isinstance(state, list)
else ["|".join((self.name, str(state)))]
)
if isinstance(state, GraphCommand):
return state.copy(update=update)
else:
return update
def send_for_fun(state):
return [
Send("2", GraphCommand(send=Send("2", 3))),
Send("2", GraphCommand(send=Send("flaky", 4))),
"3.1",
]
def route_to_three(state) -> Literal["3"]:
return "3"
builder = StateGraph(Annotated[list, operator.add])
builder.add_node(Node("1"))
builder.add_node(Node("2"))
builder.add_node(Node("3"))
builder.add_node(Node("3.1"))
builder.add_node("flaky", InterruptOnce())
builder.add_edge(START, "1")
builder.add_conditional_edges("1", send_for_fun)
builder.add_conditional_edges("2", route_to_three)
return builder.compile(checkpointer=checkpointer)
async def test_push_graph(topics: Topics, acheckpointer: BaseCheckpointSaver) -> None:
if not FF_SEND_V2:
pytest.skip("Test requires FF_SEND_V2")
input = ["0"]
config = {"configurable": {"thread_id": "1"}}
graph = mk_push_graph(acheckpointer)
graph_compare = mk_push_graph(acheckpointer)
# start a new run
async with AIOKafkaProducer(value_serializer=serde.dumps) as producer:
await producer.send_and_wait(
topics.orchestrator,
MessageToOrchestrator(input=input, config=config),
)
# drain topics
orch_msgs, exec_msgs = await drain_topics_async(topics, graph)
# check state
state = await graph.aget_state(config)
assert all(not t.error for t in state.tasks)
assert state.next == ("flaky",)
assert (
state.values
== await graph_compare.ainvoke(input, {"configurable": {"thread_id": "2"}})
== [
"0",
"1",
"2|Control(send=Send(node='2', arg=3))",
"2|Control(send=Send(node='flaky', arg=4))",
"2|3",
]
)
# check history
history = [c async for c in graph.aget_state_history(config)]
assert len(history) == 2
# check messages
assert orch_msgs == [MessageToOrchestrator(input=input, config=config)] + [
{
"config": {
"callbacks": None,
"configurable": {
"__pregel_ensure_latest": True,
"__pregel_dedupe_tasks": True,
"__pregel_resuming": False,
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_ns": "",
"thread_id": "1",
},
"metadata": AnyDict(),
"recursion_limit": 25,
"tags": [],
},
"input": None,
"finally_send": None,
}
for c in reversed(history)
for _ in c.tasks
]
assert exec_msgs == [
{
"config": {
"callbacks": None,
"configurable": {
"__pregel_ensure_latest": True,
"__pregel_dedupe_tasks": True,
"__pregel_resuming": False,
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_ns": "",
"thread_id": "1",
},
"metadata": AnyDict(),
"recursion_limit": 25,
"tags": [],
},
"task": {
"id": t.id,
"path": _convert_path(t.path),
},
"finally_send": None,
}
for c in reversed(history)
for t in c.tasks
]
# resume the thread
async with AIOKafkaProducer(value_serializer=serde.dumps) as producer:
await producer.send_and_wait(
topics.orchestrator,
MessageToOrchestrator(input=None, config=config),
)
orch_msgs, exec_msgs = await drain_topics_async(topics, graph)
# check final state
state = await graph.aget_state(config)
assert state.next == ()
assert (
state.values
== await graph_compare.ainvoke(None, {"configurable": {"thread_id": "2"}})
== [
"0",
"1",
"2|Control(send=Send(node='2', arg=3))",
"2|Control(send=Send(node='flaky', arg=4))",
"2|3",
"flaky|4",
"3",
"3.1",
]
)
# check history
history = [c async for c in graph.aget_state_history(config)]
assert len(history) == 4
# check executions
# node "2" doesn't get called again, as we recover writes saved before
assert graph.builder.nodes["2"].runnable.func.ticks == 3
# node "flaky" gets called again, as it was interrupted
assert graph.builder.nodes["flaky"].runnable.func.ticks == 2
def _convert_path(
path: tuple[Union[str, int, tuple], ...],
) -> list[Union[str, int, list]]:
return list(_convert_path(p) if isinstance(p, tuple) else p for p in path)
@@ -0,0 +1,210 @@
import operator
from typing import (
Annotated,
Literal,
Union,
)
import pytest
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.constants import FF_SEND_V2, START
from langgraph.errors import NodeInterrupt
from langgraph.graph.state import CompiledStateGraph, GraphCommand, StateGraph
from langgraph.scheduler.kafka import serde
from langgraph.scheduler.kafka.default_sync import DefaultProducer
from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics
from langgraph.types import Send
from tests.any import AnyDict
from tests.drain import drain_topics
pytestmark = pytest.mark.anyio
def mk_push_graph(
checkpointer: BaseCheckpointSaver,
) -> CompiledStateGraph:
# copied from test_send_dedupe_on_resume
class InterruptOnce:
ticks: int = 0
def __call__(self, state):
self.ticks += 1
if self.ticks == 1:
raise NodeInterrupt("Bahh")
return ["|".join(("flaky", str(state)))]
class Node:
def __init__(self, name: str):
self.name = name
self.ticks = 0
self.__name__ = name
def __call__(self, state):
self.ticks += 1
update = (
[self.name]
if isinstance(state, list)
else ["|".join((self.name, str(state)))]
)
if isinstance(state, GraphCommand):
return state.copy(update=update)
else:
return update
def send_for_fun(state):
return [
Send("2", GraphCommand(send=Send("2", 3))),
Send("2", GraphCommand(send=Send("flaky", 4))),
"3.1",
]
def route_to_three(state) -> Literal["3"]:
return "3"
builder = StateGraph(Annotated[list, operator.add])
builder.add_node(Node("1"))
builder.add_node(Node("2"))
builder.add_node(Node("3"))
builder.add_node(Node("3.1"))
builder.add_node("flaky", InterruptOnce())
builder.add_edge(START, "1")
builder.add_conditional_edges("1", send_for_fun)
builder.add_conditional_edges("2", route_to_three)
return builder.compile(checkpointer=checkpointer)
def test_push_graph(topics: Topics, acheckpointer: BaseCheckpointSaver) -> None:
if not FF_SEND_V2:
pytest.skip("Test requires FF_SEND_V2")
input = ["0"]
config = {"configurable": {"thread_id": "1"}}
graph = mk_push_graph(acheckpointer)
graph_compare = mk_push_graph(acheckpointer)
# start a new run
with DefaultProducer() as producer:
producer.send(
topics.orchestrator,
value=serde.dumps(MessageToOrchestrator(input=input, config=config)),
)
producer.flush()
# drain topics
orch_msgs, exec_msgs = drain_topics(topics, graph)
# check state
state = graph.get_state(config)
assert all(not t.error for t in state.tasks)
assert state.next == ("flaky",)
assert (
state.values
== graph_compare.invoke(input, {"configurable": {"thread_id": "2"}})
== [
"0",
"1",
"2|Control(send=Send(node='2', arg=3))",
"2|Control(send=Send(node='flaky', arg=4))",
"2|3",
]
)
# check history
history = [c for c in graph.get_state_history(config)]
assert len(history) == 2
# check messages
assert orch_msgs == [MessageToOrchestrator(input=input, config=config)] + [
{
"config": {
"callbacks": None,
"configurable": {
"__pregel_ensure_latest": True,
"__pregel_dedupe_tasks": True,
"__pregel_resuming": False,
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_ns": "",
"thread_id": "1",
},
"metadata": AnyDict(),
"recursion_limit": 25,
"tags": [],
},
"input": None,
"finally_send": None,
}
for c in reversed(history)
for _ in c.tasks
]
assert exec_msgs == [
{
"config": {
"callbacks": None,
"configurable": {
"__pregel_ensure_latest": True,
"__pregel_dedupe_tasks": True,
"__pregel_resuming": False,
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_ns": "",
"thread_id": "1",
},
"metadata": AnyDict(),
"recursion_limit": 25,
"tags": [],
},
"task": {
"id": t.id,
"path": _convert_path(t.path),
},
"finally_send": None,
}
for c in reversed(history)
for t in c.tasks
]
# resume the thread
with DefaultProducer() as producer:
producer.send(
topics.orchestrator,
value=serde.dumps(MessageToOrchestrator(input=None, config=config)),
)
producer.flush()
orch_msgs, exec_msgs = drain_topics(topics, graph)
# check final state
state = graph.get_state(config)
assert state.next == ()
assert (
state.values
== graph_compare.invoke(None, {"configurable": {"thread_id": "2"}})
== [
"0",
"1",
"2|Control(send=Send(node='2', arg=3))",
"2|Control(send=Send(node='flaky', arg=4))",
"2|3",
"flaky|4",
"3",
"3.1",
]
)
# check history
history = [c for c in graph.get_state_history(config)]
assert len(history) == 4
# check executions
# node "2" doesn't get called again, as we recover writes saved before
assert graph.builder.nodes["2"].runnable.func.ticks == 3
# node "flaky" gets called again, as it was interrupted
assert graph.builder.nodes["flaky"].runnable.func.ticks == 2
def _convert_path(
path: tuple[Union[str, int, tuple], ...],
) -> list[Union[str, int, list]]:
return list(_convert_path(p) if isinstance(p, tuple) else p for p in path)
+12 -6
View File
@@ -194,8 +194,9 @@ async def test_subgraph_w_interrupt(
"__pregel_ensure_latest": True,
"__pregel_dedupe_tasks": True,
"__pregel_resuming": False,
'__pregel_store': None,
"__pregel_store": None,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_resume_value": None,
"checkpoint_id": None,
"checkpoint_map": {
"": history[0].config["configurable"]["checkpoint_id"]
@@ -258,8 +259,9 @@ async def test_subgraph_w_interrupt(
"__pregel_ensure_latest": True,
"__pregel_dedupe_tasks": True,
"__pregel_resuming": False,
'__pregel_store': None,
"__pregel_store": None,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_resume_value": None,
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_map": {
"": history[0].config["configurable"]["checkpoint_id"]
@@ -352,8 +354,9 @@ async def test_subgraph_w_interrupt(
"__pregel_ensure_latest": True,
"__pregel_dedupe_tasks": True,
"__pregel_resuming": False,
'__pregel_store': None,
"__pregel_store": None,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_resume_value": None,
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_map": {
"": history[0].config["configurable"]["checkpoint_id"]
@@ -456,8 +459,9 @@ async def test_subgraph_w_interrupt(
"__pregel_ensure_latest": True,
"__pregel_dedupe_tasks": True,
"__pregel_resuming": True,
'__pregel_store': None,
"__pregel_store": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_resume_value": None,
"checkpoint_id": None,
"checkpoint_map": {
"": history[1].config["configurable"]["checkpoint_id"]
@@ -515,8 +519,9 @@ async def test_subgraph_w_interrupt(
"__pregel_ensure_latest": True,
"__pregel_dedupe_tasks": True,
"__pregel_resuming": True,
'__pregel_store': None,
"__pregel_store": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_resume_value": None,
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_map": {
"": history[1].config["configurable"]["checkpoint_id"]
@@ -630,8 +635,9 @@ async def test_subgraph_w_interrupt(
"__pregel_ensure_latest": True,
"__pregel_dedupe_tasks": True,
"__pregel_resuming": True,
'__pregel_store': None,
"__pregel_store": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_resume_value": None,
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_map": {
"": history[1].config["configurable"]["checkpoint_id"]
@@ -193,8 +193,9 @@ def test_subgraph_w_interrupt(
"__pregel_ensure_latest": True,
"__pregel_dedupe_tasks": True,
"__pregel_resuming": False,
'__pregel_store': None,
"__pregel_store": None,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_resume_value": None,
"checkpoint_id": None,
"checkpoint_map": {
"": history[0].config["configurable"]["checkpoint_id"]
@@ -255,10 +256,11 @@ def test_subgraph_w_interrupt(
"__pregel_read": None,
"__pregel_send": None,
"__pregel_ensure_latest": True,
'__pregel_store': None,
"__pregel_store": None,
"__pregel_dedupe_tasks": True,
"__pregel_resuming": False,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_resume_value": None,
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_map": {
"": history[0].config["configurable"]["checkpoint_id"]
@@ -350,9 +352,10 @@ def test_subgraph_w_interrupt(
"__pregel_send": None,
"__pregel_ensure_latest": True,
"__pregel_dedupe_tasks": True,
'__pregel_store': None,
"__pregel_store": None,
"__pregel_resuming": False,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_resume_value": None,
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_map": {
"": history[0].config["configurable"]["checkpoint_id"]
@@ -453,9 +456,10 @@ def test_subgraph_w_interrupt(
"__pregel_send": None,
"__pregel_ensure_latest": True,
"__pregel_dedupe_tasks": True,
'__pregel_store': None,
"__pregel_store": None,
"__pregel_resuming": True,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_resume_value": None,
"checkpoint_id": None,
"checkpoint_map": {
"": history[1].config["configurable"]["checkpoint_id"]
@@ -512,9 +516,10 @@ def test_subgraph_w_interrupt(
"__pregel_send": None,
"__pregel_ensure_latest": True,
"__pregel_dedupe_tasks": True,
'__pregel_store': None,
"__pregel_store": None,
"__pregel_resuming": True,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_resume_value": None,
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_map": {
"": history[1].config["configurable"]["checkpoint_id"]
@@ -628,8 +633,9 @@ def test_subgraph_w_interrupt(
"__pregel_ensure_latest": True,
"__pregel_dedupe_tasks": True,
"__pregel_resuming": True,
'__pregel_store': None,
"__pregel_store": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_resume_value": None,
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_map": {
"": history[1].config["configurable"]["checkpoint_id"]
+16
View File
@@ -36,6 +36,7 @@ from langgraph_sdk.schema import (
AssistantVersion,
CancelAction,
Checkpoint,
Command,
Config,
Cron,
DisconnectMode,
@@ -1174,6 +1175,7 @@ class RunsClient:
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,
@@ -1197,6 +1199,7 @@ class RunsClient:
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,
@@ -1217,6 +1220,7 @@ class RunsClient:
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,
@@ -1241,6 +1245,7 @@ class RunsClient:
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: A command to execute. Cannot be combined with input.
stream_mode: The stream mode(s) to use.
stream_subgraphs: Whether to stream output from subgraphs.
metadata: Metadata to assign to the run.
@@ -1291,6 +1296,7 @@ class RunsClient:
""" # noqa: E501
payload = {
"input": input,
"command": command,
"config": config,
"metadata": metadata,
"stream_mode": stream_mode,
@@ -1324,6 +1330,7 @@ class RunsClient:
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,
@@ -1343,6 +1350,7 @@ class RunsClient:
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,
@@ -1363,6 +1371,7 @@ class RunsClient:
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,
@@ -1385,6 +1394,7 @@ class RunsClient:
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: A command to execute. Cannot be combined with input.
stream_mode: The stream mode(s) to use.
stream_subgraphs: Whether to stream output from subgraphs.
metadata: Metadata to assign to the run.
@@ -1471,6 +1481,7 @@ class RunsClient:
""" # noqa: E501
payload = {
"input": input,
"command": command,
"stream_mode": stream_mode,
"stream_subgraphs": stream_subgraphs,
"config": config,
@@ -1508,6 +1519,7 @@ class RunsClient:
assistant_id: str,
*,
input: Optional[dict] = None,
command: Optional[Command] = None,
metadata: Optional[dict] = None,
config: Optional[Config] = None,
checkpoint: Optional[Checkpoint] = None,
@@ -1529,6 +1541,7 @@ class RunsClient:
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,
@@ -1547,6 +1560,7 @@ class RunsClient:
assistant_id: str,
*,
input: Optional[dict] = None,
command: Optional[Command] = None,
metadata: Optional[dict] = None,
config: Optional[Config] = None,
checkpoint: Optional[Checkpoint] = None,
@@ -1569,6 +1583,7 @@ class RunsClient:
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: A command to execute. Cannot be combined with input.
metadata: Metadata to assign to the run.
config: The configuration for the assistant.
checkpoint: The checkpoint to resume from.
@@ -1635,6 +1650,7 @@ class RunsClient:
""" # noqa: E501
payload = {
"input": input,
"command": command,
"config": config,
"metadata": metadata,
"assistant_id": assistant_id,
+11
View File
@@ -339,3 +339,14 @@ class StreamPart(NamedTuple):
"""The type of event for this stream part."""
data: dict
"""The data payload associated with the event."""
class Send(TypedDict):
node: str
input: Optional[dict[str, Any]]
class Command(TypedDict, total=False):
send: Union[Send, Sequence[Send]]
update: dict[str, Any]
resume: Any
+1 -1
View File
@@ -1,6 +1,6 @@
[tool.poetry]
name = "langgraph-sdk"
version = "0.1.35"
version = "0.1.36"
description = "SDK for interacting with LangGraph API"
authors = []
license = "MIT"