mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-19 22:25:44 +02:00
Compare commits
32
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
3a860ad537 | ||
|
|
7051bccc30 | ||
|
|
199e41b228 | ||
|
|
229a9e19a8 | ||
|
|
3c3a1a1f35 | ||
|
|
f11127648e | ||
|
|
29f833b1a7 | ||
|
|
7a3ea42743 | ||
|
|
a94902db8a | ||
|
|
9fd152ef3a | ||
|
|
03bc9ba6e6 | ||
|
|
7fe6f88876 | ||
|
|
1d88affd29 | ||
|
|
6906e12edb | ||
|
|
16bfa80b58 | ||
|
|
00964b18f6 | ||
|
|
0d5c6201d3 | ||
|
|
86d2847dab | ||
|
|
b3a4eaa967 | ||
|
|
87fc519ce7 | ||
|
|
ef3a1ee997 | ||
|
|
311e16dffd | ||
|
|
c83b8f6d04 | ||
|
|
62d3a85b07 | ||
|
|
810ae0ef51 | ||
|
|
2ff49d2200 | ||
|
|
d0567dc7be | ||
|
|
ea64ac5c07 | ||
|
|
090b53ccc1 | ||
|
|
0e872e7482 | ||
|
|
3ad966e057 | ||
|
|
89a0859928 |
@@ -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,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"
|
||||
|
||||
@@ -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."""
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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]))
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
)
|
||||
|
||||
@@ -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)}"
|
||||
|
||||
@@ -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]],
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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 = [
|
||||
|
||||
@@ -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),
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
@@ -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")
|
||||
|
||||
Generated
+3
-3
@@ -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"
|
||||
|
||||
@@ -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'}}}%%
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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,
|
||||
|
||||
+748
-241
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -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):
|
||||
|
||||
@@ -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)
|
||||
@@ -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"]
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,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"
|
||||
|
||||
Reference in New Issue
Block a user