diff --git a/.github/workflows/_integration_test.yml b/.github/workflows/_integration_test.yml index c2d7cb895..1cfdb5fde 100644 --- a/.github/workflows/_integration_test.yml +++ b/.github/workflows/_integration_test.yml @@ -4,7 +4,7 @@ on: workflow_call: env: - POETRY_VERSION: "1.7.1" + POETRY_VERSION: "2.1.2" jobs: build: @@ -71,4 +71,3 @@ jobs: working-directory: libs/cli/js-examples run: | langgraph build -t langgraph-test-e - \ No newline at end of file diff --git a/.github/workflows/_lint.yml b/.github/workflows/_lint.yml index 6d63e11c6..97afe5ed6 100644 --- a/.github/workflows/_lint.yml +++ b/.github/workflows/_lint.yml @@ -9,7 +9,7 @@ on: description: "From which folder this pipeline executes" env: - POETRY_VERSION: "1.7.1" + POETRY_VERSION: "2.1.2" # This env var allows us to get inline annotations when ruff has complaints. RUFF_OUTPUT_FORMAT: github @@ -50,12 +50,6 @@ jobs: working-directory: ${{ inputs.working-directory }} run: poetry check - - name: Check lock file - if: steps.changed-files.outputs.all - shell: bash - working-directory: ${{ inputs.working-directory }} - run: poetry check --lock - - name: Install dependencies if: steps.changed-files.outputs.all # Also installs dev/lint/test/typing dependencies, to ensure we have diff --git a/.github/workflows/_test.yml b/.github/workflows/_test.yml index 64dfacd51..85b6ac353 100644 --- a/.github/workflows/_test.yml +++ b/.github/workflows/_test.yml @@ -9,7 +9,7 @@ on: description: "From which folder this pipeline executes" env: - POETRY_VERSION: "1.7.1" + POETRY_VERSION: "2.1.2" jobs: build: @@ -39,12 +39,6 @@ jobs: username: ${{ secrets.DOCKERHUB_USERNAME }} password: ${{ secrets.DOCKERHUB_RO_TOKEN }} - - name: Check Lock - shell: bash - working-directory: ${{ inputs.working-directory }} - run: | - poetry check --lock - - name: Install dependencies shell: bash working-directory: ${{ inputs.working-directory }} diff --git a/.github/workflows/_test_langgraph.yml b/.github/workflows/_test_langgraph.yml index ba1cbd08e..6ff952e79 100644 --- a/.github/workflows/_test_langgraph.yml +++ b/.github/workflows/_test_langgraph.yml @@ -4,7 +4,7 @@ on: workflow_call: env: - POETRY_VERSION: "1.7.1" + POETRY_VERSION: "2.1.2" jobs: build: diff --git a/.github/workflows/_test_release.yml b/.github/workflows/_test_release.yml index 46e065d33..21a8c9951 100644 --- a/.github/workflows/_test_release.yml +++ b/.github/workflows/_test_release.yml @@ -9,7 +9,7 @@ on: description: "From which folder this pipeline executes" env: - POETRY_VERSION: "1.7.1" + POETRY_VERSION: "2.1.2" PYTHON_VERSION: "3.10" jobs: diff --git a/.github/workflows/_test_scheduler_kafka.yml b/.github/workflows/_test_scheduler_kafka.yml index 1f0edf420..8d0008871 100644 --- a/.github/workflows/_test_scheduler_kafka.yml +++ b/.github/workflows/_test_scheduler_kafka.yml @@ -4,7 +4,7 @@ on: workflow_call: env: - POETRY_VERSION: "1.7.1" + POETRY_VERSION: "2.1.2" jobs: build: diff --git a/.github/workflows/baseline.yml b/.github/workflows/baseline.yml index 26a4f2ef0..4f3a33d73 100644 --- a/.github/workflows/baseline.yml +++ b/.github/workflows/baseline.yml @@ -8,7 +8,7 @@ on: - "libs/**" env: - POETRY_VERSION: "1.7.1" + POETRY_VERSION: "2.1.2" jobs: benchmark: diff --git a/.github/workflows/bench.yml b/.github/workflows/bench.yml index 0d0fc84b6..1222c4d4e 100644 --- a/.github/workflows/bench.yml +++ b/.github/workflows/bench.yml @@ -6,7 +6,7 @@ on: - "libs/**" env: - POETRY_VERSION: "1.7.1" + POETRY_VERSION: "2.1.2" jobs: benchmark: diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 9bed24785..6f359ad67 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -17,7 +17,7 @@ concurrency: cancel-in-progress: true env: - POETRY_VERSION: "1.7.1" + POETRY_VERSION: "2.1.2" jobs: changes: diff --git a/.github/workflows/deploy_docs.yml b/.github/workflows/deploy_docs.yml index 490169946..ec17db77c 100644 --- a/.github/workflows/deploy_docs.yml +++ b/.github/workflows/deploy_docs.yml @@ -10,7 +10,7 @@ on: workflow_dispatch: env: - POETRY_VERSION: "1.7.1" + POETRY_VERSION: "2.1.2" permissions: contents: read diff --git a/.github/workflows/link_check.yml b/.github/workflows/link_check.yml index a42b5fb73..f0d0c35b1 100644 --- a/.github/workflows/link_check.yml +++ b/.github/workflows/link_check.yml @@ -12,7 +12,7 @@ on: workflow_dispatch: env: - POETRY_VERSION: "1.7.1" + POETRY_VERSION: "2.1.2" jobs: markdown-link-check: @@ -42,8 +42,8 @@ jobs: - name: Check README.md is in sync run: | - if ! diff -q README.md libs/langgraph/README.md >/dev/null; then - echo "README.md is out of sync with libs/langgraph/README.md" - diff -C 3 README.md libs/langgraph/README.md - exit 1 - fi + if ! diff -q README.md libs/langgraph/README.md >/dev/null; then + echo "README.md is out of sync with libs/langgraph/README.md" + diff -C 3 README.md libs/langgraph/README.md + exit 1 + fi diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index c93db70a3..8c1a07085 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -10,7 +10,7 @@ on: env: PYTHON_VERSION: "3.11" - POETRY_VERSION: "1.7.1" + POETRY_VERSION: "2.1.2" jobs: build: diff --git a/.github/workflows/run_notebooks.yml b/.github/workflows/run_notebooks.yml index 16abf71d7..07d52e93e 100644 --- a/.github/workflows/run_notebooks.yml +++ b/.github/workflows/run_notebooks.yml @@ -9,7 +9,7 @@ on: type: string description: "JSON string of changed files" schedule: - - cron: '0 13 * * *' + - cron: "0 13 * * *" defaults: run: @@ -30,12 +30,12 @@ jobs: uses: "./.github/actions/poetry_setup" with: python-version: 3.11 - poetry-version: 1.7.1 + poetry-version: 2.1.2 cache-key: test-langgraph-notebooks - name: Install dependencies run: | - poetry install --with test + poetry install --with test --no-root poetry run pip install jupyter - name: Start services diff --git a/docs/docs/how-tos/ttl/configure_ttl.md b/docs/docs/how-tos/ttl/configure_ttl.md index 9fdbf8eb0..de431946c 100644 --- a/docs/docs/how-tos/ttl/configure_ttl.md +++ b/docs/docs/how-tos/ttl/configure_ttl.md @@ -2,7 +2,7 @@ !!! tip "Prerequisites" - This guide assumes familiarity with the [LangGraph Platform](../../concepts/index.md#langgraph-platform), [Persistence](../../concepts/persistence.md), and [Cross-thread persistence](../../concepts/store.md) concepts. + This guide assumes familiarity with the [LangGraph Platform](../../concepts/index.md#langgraph-platform), [Persistence](../../concepts/persistence.md), and [Cross-thread persistence](../../concepts/persistence.md#memory-store) concepts. ???+ note "LangGraph platform only" diff --git a/libs/cli/langgraph_cli/config.py b/libs/cli/langgraph_cli/config.py index 5c512b58c..29807dae0 100644 --- a/libs/cli/langgraph_cli/config.py +++ b/libs/cli/langgraph_cli/config.py @@ -406,6 +406,13 @@ class Config(TypedDict, total=False): """ +PIP_CLEANUP_LINES = """# -- Removing pip from the final image ~<:===~~~ -- +RUN pip uninstall -y pip setuptools wheel && \ + rm -rf /usr/local/lib/python*/site-packages/pip* /usr/local/lib/python*/site-packages/setuptools* /usr/local/lib/python*/site-packages/wheel* && \ + find /usr/local/bin -name "pip*" -delete +# -- End of pip removal --""" + + def _parse_version(version_str: str) -> tuple[int, int]: """Parse a version string into a tuple of (major, minor).""" try: @@ -1141,6 +1148,8 @@ ADD {relpath} /deps/{name} "", ui_inst_str, "", + PIP_CLEANUP_LINES, # Add pip cleanup after all installations are complete + "", f"WORKDIR {local_deps.working_dir}" if local_deps.working_dir else "", ] diff --git a/libs/cli/pyproject.toml b/libs/cli/pyproject.toml index e350076f9..2040ad217 100644 --- a/libs/cli/pyproject.toml +++ b/libs/cli/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langgraph-cli" -version = "0.1.89" +version = "0.1.90" description = "CLI for interacting with LangGraph API" authors = [] license = "MIT" diff --git a/libs/cli/tests/unit_tests/cli/test_cli.py b/libs/cli/tests/unit_tests/cli/test_cli.py index 53954da5e..4140ee9a5 100644 --- a/libs/cli/tests/unit_tests/cli/test_cli.py +++ b/libs/cli/tests/unit_tests/cli/test_cli.py @@ -2,13 +2,14 @@ import json import pathlib import shutil import tempfile +import textwrap from contextlib import contextmanager from pathlib import Path from click.testing import CliRunner from langgraph_cli.cli import cli, prepare_args_and_stdin -from langgraph_cli.config import Config, validate_config +from langgraph_cli.config import PIP_CLEANUP_LINES, Config, validate_config from langgraph_cli.docker import DEFAULT_POSTGRES_URI, DockerCapabilities, Version from langgraph_cli.util import clean_empty_lines @@ -143,6 +144,7 @@ services: RUN PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/* # -- End of local dependencies install -- ENV LANGSERVE_GRAPHS='{{"agent": "agent.py:graph"}}' +{textwrap.indent(textwrap.dedent(PIP_CLEANUP_LINES), " ")} WORKDIR /deps/cli develop: diff --git a/libs/cli/tests/unit_tests/test_config.py b/libs/cli/tests/unit_tests/test_config.py index dfa8cbbef..cb27b0d0c 100644 --- a/libs/cli/tests/unit_tests/test_config.py +++ b/libs/cli/tests/unit_tests/test_config.py @@ -2,11 +2,13 @@ import json import os import pathlib import tempfile +import textwrap import click import pytest from langgraph_cli.config import ( + PIP_CLEANUP_LINES, config_to_compose, config_to_docker, validate_config, @@ -208,7 +210,7 @@ def test_config_to_docker_simple(): ), "langchain/langgraph-api", ) - expected_docker_stdin = """\ + expected_docker_stdin = f"""\ FROM langchain/langgraph-api:3.11 # -- Installing local requirements -- COPY --from=__outer_requirements.txt requirements.txt /deps/__outer_graphs_reqs_a/graphs_reqs_a/requirements.txt @@ -242,8 +244,9 @@ RUN set -ex && \\ # -- Installing all local dependencies -- RUN PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/* # -- End of local dependencies install -- -ENV LANGGRAPH_HTTP='{"app": "/deps/examples/my_app.py:app"}' -ENV LANGSERVE_GRAPHS='{"agent": "/deps/__outer_unit_tests/unit_tests/agent.py:graph"}' +ENV LANGGRAPH_HTTP='{{"app": "/deps/examples/my_app.py:app"}}' +ENV LANGSERVE_GRAPHS='{{"agent": "/deps/__outer_unit_tests/unit_tests/agent.py:graph"}}' +{PIP_CLEANUP_LINES} WORKDIR /deps/__outer_unit_tests/unit_tests\ """ assert clean_empty_lines(actual_docker_stdin) == expected_docker_stdin @@ -263,7 +266,8 @@ def test_config_to_docker_outside_path(): validate_config({"dependencies": [".", ".."], "graphs": graphs}), "langchain/langgraph-api", ) - expected_docker_stdin = """\ + expected_docker_stdin = ( + """\ FROM langchain/langgraph-api:3.11 # -- Adding non-package dependency unit_tests -- ADD . /deps/__outer_unit_tests/unit_tests @@ -291,8 +295,12 @@ RUN set -ex && \\ RUN PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/* # -- End of local dependencies install -- ENV LANGSERVE_GRAPHS='{"agent": "/deps/__outer_unit_tests/unit_tests/agent.py:graph"}' +""" + + PIP_CLEANUP_LINES + + """ WORKDIR /deps/__outer_unit_tests/unit_tests\ """ + ) assert clean_empty_lines(actual_docker_stdin) == expected_docker_stdin assert additional_contexts == { "__outer_tests": str(pathlib.Path(__file__).parent.parent.absolute()), @@ -312,7 +320,8 @@ def test_config_to_docker_pipconfig(): ), "langchain/langgraph-api", ) - expected_docker_stdin = """\ + expected_docker_stdin = ( + """\ FROM langchain/langgraph-api:3.11 ADD pipconfig.txt /pipconfig.txt # -- Adding non-package dependency unit_tests -- @@ -330,8 +339,12 @@ RUN set -ex && \\ RUN PIP_CONFIG_FILE=/pipconfig.txt PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/* # -- End of local dependencies install -- ENV LANGSERVE_GRAPHS='{"agent": "/deps/__outer_unit_tests/unit_tests/agent.py:graph"}' +""" + + PIP_CLEANUP_LINES + + """ WORKDIR /deps/__outer_unit_tests/unit_tests\ """ + ) assert clean_empty_lines(actual_docker_stdin) == expected_docker_stdin assert additional_contexts == {} @@ -368,7 +381,7 @@ def test_config_to_docker_local_deps(): ), "langchain/langgraph-api-custom", ) - expected_docker_stdin = """\ + expected_docker_stdin = f"""\ FROM langchain/langgraph-api-custom:3.11 # -- Adding non-package dependency graphs -- ADD ./graphs /deps/__outer_graphs/src @@ -384,7 +397,8 @@ RUN set -ex && \\ # -- Installing all local dependencies -- RUN PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/* # -- End of local dependencies install -- -ENV LANGSERVE_GRAPHS='{"agent": "/deps/__outer_graphs/src/agent.py:graph"}'\ +ENV LANGSERVE_GRAPHS='{{"agent": "/deps/__outer_graphs/src/agent.py:graph"}}' +{PIP_CLEANUP_LINES}\ """ assert clean_empty_lines(actual_docker_stdin) == expected_docker_stdin assert additional_contexts == {} @@ -411,7 +425,8 @@ dependencies = ["langchain"]""" "langchain/langgraph-api", ) os.remove(pyproject_path) - expected_docker_stdin = """FROM langchain/langgraph-api:3.11 + expected_docker_stdin = ( + """FROM langchain/langgraph-api:3.11 # -- Adding local package . -- ADD . /deps/unit_tests # -- End of local package . -- @@ -419,7 +434,12 @@ ADD . /deps/unit_tests RUN PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/* # -- End of local dependencies install -- ENV LANGSERVE_GRAPHS='{"agent": "/deps/unit_tests/graphs/agent.py:graph"}' -WORKDIR /deps/unit_tests""" +""" + + PIP_CLEANUP_LINES + + "\n" + + "WORKDIR /deps/unit_tests" + "" + ) assert clean_empty_lines(actual_docker_stdin) == expected_docker_stdin assert additional_contexts == {} @@ -439,7 +459,7 @@ def test_config_to_docker_end_to_end(): ), "langchain/langgraph-api", ) - expected_docker_stdin = """FROM langchain/langgraph-api:3.12 + expected_docker_stdin = f"""FROM langchain/langgraph-api:3.12 ARG meow ARG foo ADD pipconfig.txt /pipconfig.txt @@ -458,7 +478,8 @@ RUN set -ex && \\ # -- Installing all local dependencies -- RUN PIP_CONFIG_FILE=/pipconfig.txt PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/* # -- End of local dependencies install -- -ENV LANGSERVE_GRAPHS='{"agent": "/deps/__outer_graphs/src/agent.py:graph"}'""" +ENV LANGSERVE_GRAPHS='{{"agent": "/deps/__outer_graphs/src/agent.py:graph"}}' +{PIP_CLEANUP_LINES}""" assert clean_empty_lines(actual_docker_stdin) == expected_docker_stdin assert additional_contexts == {} @@ -509,7 +530,7 @@ def test_config_to_docker_gen_ui_python(): "langchain/langgraph-api", ) - expected_docker_stdin = """FROM langchain/langgraph-api:3.11 + expected_docker_stdin = f"""FROM langchain/langgraph-api:3.11 RUN /storage/install-node.sh # -- Adding non-package dependency unit_tests -- ADD . /deps/__outer_unit_tests/unit_tests @@ -525,12 +546,13 @@ RUN set -ex && \\ # -- Installing all local dependencies -- RUN PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/* # -- End of local dependencies install -- -ENV LANGSERVE_GRAPHS='{"agent": "/deps/__outer_unit_tests/unit_tests/agent.py:graph"}' +ENV LANGSERVE_GRAPHS='{{"agent": "/deps/__outer_unit_tests/unit_tests/agent.py:graph"}}' # -- Installing UI dependencies -- -ENV LANGGRAPH_UI='{"agent": "./graphs/agent.ui.jsx"}' -ENV LANGGRAPH_UI_CONFIG='{"shared": ["nuqs"]}' +ENV LANGGRAPH_UI='{{"agent": "./graphs/agent.ui.jsx"}}' +ENV LANGGRAPH_UI_CONFIG='{{"shared": ["nuqs"]}}' RUN cd /deps/__outer_unit_tests/unit_tests && npm i && tsx /api/langgraph_api/js/build.mts # -- End of UI dependencies install -- +{PIP_CLEANUP_LINES} WORKDIR /deps/__outer_unit_tests/unit_tests""" assert clean_empty_lines(actual_docker_stdin) == expected_docker_stdin @@ -540,8 +562,8 @@ WORKDIR /deps/__outer_unit_tests/unit_tests""" # config_to_compose def test_config_to_compose_simple_config(): graphs = {"agent": "./agent.py:graph"} - expected_compose_stdin = """\ - + # Create a properly indented version of PIP_CLEANUP_LINES for compose files + expected_compose_stdin = f""" pull_policy: build build: context: . @@ -561,7 +583,8 @@ def test_config_to_compose_simple_config(): # -- Installing all local dependencies -- RUN PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/* # -- End of local dependencies install -- - ENV LANGSERVE_GRAPHS='{"agent": "/deps/__outer_unit_tests/unit_tests/agent.py:graph"}' + ENV LANGSERVE_GRAPHS='{{"agent": "/deps/__outer_unit_tests/unit_tests/agent.py:graph"}}' +{textwrap.indent(textwrap.dedent(PIP_CLEANUP_LINES), " ")} WORKDIR /deps/__outer_unit_tests/unit_tests """ actual_compose_stdin = config_to_compose( @@ -569,12 +592,15 @@ def test_config_to_compose_simple_config(): validate_config({"dependencies": ["."], "graphs": graphs}), "langchain/langgraph-api", ) - assert clean_empty_lines(actual_compose_stdin) == expected_compose_stdin + assert ( + clean_empty_lines(actual_compose_stdin).strip() + == expected_compose_stdin.strip() + ) def test_config_to_compose_env_vars(): graphs = {"agent": "./agent.py:graph"} - expected_compose_stdin = """ OPENAI_API_KEY: "key" + expected_compose_stdin = f""" OPENAI_API_KEY: "key" pull_policy: build build: @@ -595,7 +621,8 @@ def test_config_to_compose_env_vars(): # -- Installing all local dependencies -- RUN PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/* # -- End of local dependencies install -- - ENV LANGSERVE_GRAPHS='{"agent": "/deps/__outer_unit_tests/unit_tests/agent.py:graph"}' + ENV LANGSERVE_GRAPHS='{{"agent": "/deps/__outer_unit_tests/unit_tests/agent.py:graph"}}' +{textwrap.indent(textwrap.dedent(PIP_CLEANUP_LINES), " ")} WORKDIR /deps/__outer_unit_tests/unit_tests """ openai_api_key = "key" @@ -615,7 +642,7 @@ def test_config_to_compose_env_vars(): def test_config_to_compose_env_file(): graphs = {"agent": "./agent.py:graph"} - expected_compose_stdin = """\ + expected_compose_stdin = f"""\ env_file: .env pull_policy: build build: @@ -636,7 +663,8 @@ def test_config_to_compose_env_file(): # -- Installing all local dependencies -- RUN PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/* # -- End of local dependencies install -- - ENV LANGSERVE_GRAPHS='{"agent": "/deps/__outer_unit_tests/unit_tests/agent.py:graph"}' + ENV LANGSERVE_GRAPHS='{{"agent": "/deps/__outer_unit_tests/unit_tests/agent.py:graph"}}' +{textwrap.indent(textwrap.dedent(PIP_CLEANUP_LINES), " ")} WORKDIR /deps/__outer_unit_tests/unit_tests """ actual_compose_stdin = config_to_compose( @@ -649,7 +677,7 @@ def test_config_to_compose_env_file(): def test_config_to_compose_watch(): graphs = {"agent": "./agent.py:graph"} - expected_compose_stdin = """\ + expected_compose_stdin = f"""\ pull_policy: build build: @@ -670,7 +698,8 @@ def test_config_to_compose_watch(): # -- Installing all local dependencies -- RUN PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/* # -- End of local dependencies install -- - ENV LANGSERVE_GRAPHS='{"agent": "/deps/__outer_unit_tests/unit_tests/agent.py:graph"}' + ENV LANGSERVE_GRAPHS='{{"agent": "/deps/__outer_unit_tests/unit_tests/agent.py:graph"}}' +{textwrap.indent(textwrap.dedent(PIP_CLEANUP_LINES), " ")} WORKDIR /deps/__outer_unit_tests/unit_tests develop: @@ -692,7 +721,7 @@ def test_config_to_compose_watch(): def test_config_to_compose_end_to_end(): # test all of the above + langgraph API path graphs = {"agent": "./agent.py:graph"} - expected_compose_stdin = """\ + expected_compose_stdin = f"""\ env_file: .env pull_policy: build build: @@ -713,7 +742,8 @@ def test_config_to_compose_end_to_end(): # -- Installing all local dependencies -- RUN PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/* # -- End of local dependencies install -- - ENV LANGSERVE_GRAPHS='{"agent": "/deps/__outer_unit_tests/unit_tests/agent.py:graph"}' + ENV LANGSERVE_GRAPHS='{{"agent": "/deps/__outer_unit_tests/unit_tests/agent.py:graph"}}' +{textwrap.indent(textwrap.dedent(PIP_CLEANUP_LINES), " ")} WORKDIR /deps/__outer_unit_tests/unit_tests develop: diff --git a/libs/langgraph/bench/__main__.py b/libs/langgraph/bench/__main__.py index 966e25d1d..f75f2122b 100644 --- a/libs/langgraph/bench/__main__.py +++ b/libs/langgraph/bench/__main__.py @@ -26,6 +26,7 @@ async def arun(graph: Pregel, input: dict): "configurable": {"thread_id": str(uuid4())}, "recursion_limit": 1000000000, }, + checkpoint_during=False, ) ] ) @@ -42,6 +43,7 @@ async def arun_first_event_latency(graph: Pregel, input: dict) -> None: "configurable": {"thread_id": str(uuid4())}, "recursion_limit": 1000000000, }, + checkpoint_during=False, ) try: @@ -61,6 +63,7 @@ def run(graph: Pregel, input: dict): "configurable": {"thread_id": str(uuid4())}, "recursion_limit": 1000000000, }, + checkpoint_during=False, ) ] ) @@ -77,6 +80,7 @@ def run_first_event_latency(graph: Pregel, input: dict) -> None: "configurable": {"thread_id": str(uuid4())}, "recursion_limit": 1000000000, }, + checkpoint_during=False, ) try: diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index 4fde5e7dd..8bcbb2b94 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -83,6 +83,8 @@ CONFIG_KEY_PREVIOUS = sys.intern("__pregel_previous") # holds the previous return value from a stateful Pregel graph. CONFIG_KEY_RUNNER_SUBMIT = sys.intern("__pregel_runner_submit") # holds a function that receives tasks from runner, executes them and returns results +CONFIG_KEY_CHECKPOINT_DURING = sys.intern("__pregel_checkpoint_during") +# holds a boolean indicating whether to checkpoint during the run (or only at the end) # --- Other constants --- PUSH = sys.intern("__pregel_push") diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 545c8faf1..c9821248e 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -39,7 +39,6 @@ from langchain_core.runnables.utils import ( ConfigurableFieldSpec, get_unique_config_specs, ) -from langchain_core.tracers._streaming import _StreamingCallbackHandler from pydantic import BaseModel from typing_extensions import Self @@ -54,6 +53,7 @@ from langgraph.checkpoint.base import ( ) from langgraph.constants import ( CONF, + CONFIG_KEY_CHECKPOINT_DURING, CONFIG_KEY_CHECKPOINT_ID, CONFIG_KEY_CHECKPOINT_NS, CONFIG_KEY_CHECKPOINTER, @@ -125,6 +125,11 @@ from langgraph.utils.fields import get_enhanced_type_hints from langgraph.utils.pydantic import create_model, is_supported_by_pydantic from langgraph.utils.queue import AsyncQueue, SyncQueue # type: ignore[attr-defined] +try: + from langchain_core.tracers._streaming import _StreamingCallbackHandler +except ImportError: + _StreamingCallbackHandler = None # type: ignore + WriteValue = Union[Callable[[Input], Output], Any] @@ -2094,6 +2099,7 @@ class Pregel(PregelProtocol): output_keys: Optional[Union[str, Sequence[str]]] = None, interrupt_before: Optional[Union[All, Sequence[str]]] = None, interrupt_after: Optional[Union[All, Sequence[str]]] = None, + checkpoint_during: Optional[bool] = None, debug: Optional[bool] = None, subgraphs: bool = False, ) -> Iterator[Union[dict[str, Any], Any]]: @@ -2115,6 +2121,7 @@ class Pregel(PregelProtocol): output_keys: The keys to stream, defaults to all non-context channels. interrupt_before: Nodes to interrupt before, defaults to all nodes in the graph. interrupt_after: Nodes to interrupt after, defaults to all nodes in the graph. + checkpoint_during: Whether to checkpoint intermediate steps, defaults to True. If False, only the final checkpoint is saved. debug: Whether to print debug information during execution, defaults to False. subgraphs: Whether to stream subgraphs, defaults to False. @@ -2276,6 +2283,9 @@ class Pregel(PregelProtocol): config[CONF][CONFIG_KEY_STREAM_WRITER] = lambda c: stream.put( ((), "custom", c) ) + # set checkpointing mode for subgraphs + if checkpoint_during is not None: + config[CONF][CONFIG_KEY_CHECKPOINT_DURING] = checkpoint_during with SyncPregelLoop( input, input_model=self.input_model, @@ -2291,6 +2301,9 @@ class Pregel(PregelProtocol): interrupt_after=interrupt_after_, manager=run_manager, debug=debug, + checkpoint_during=checkpoint_during + if checkpoint_during is not None + else config[CONF].get(CONFIG_KEY_CHECKPOINT_DURING, True), trigger_to_nodes=self.trigger_to_nodes, migrate_checkpoint=self._migrate_checkpoint, ) as loop: @@ -2373,6 +2386,7 @@ class Pregel(PregelProtocol): output_keys: Optional[Union[str, Sequence[str]]] = None, interrupt_before: Optional[Union[All, Sequence[str]]] = None, interrupt_after: Optional[Union[All, Sequence[str]]] = None, + checkpoint_during: Optional[bool] = None, debug: Optional[bool] = None, subgraphs: bool = False, ) -> AsyncIterator[Union[dict[str, Any], Any]]: @@ -2394,6 +2408,7 @@ class Pregel(PregelProtocol): output_keys: The keys to stream, defaults to all non-context channels. interrupt_before: Nodes to interrupt before, defaults to all nodes in the graph. interrupt_after: Nodes to interrupt after, defaults to all nodes in the graph. + checkpoint_during: Whether to checkpoint intermediate steps, defaults to True. If False, only the final checkpoint is saved. debug: Whether to print debug information during execution, defaults to False. subgraphs: Whether to stream subgraphs, defaults to False. @@ -2529,13 +2544,17 @@ class Pregel(PregelProtocol): run_id=config.get("run_id"), ) # if running from astream_log() run each proc with streaming - do_stream = next( - ( - cast(_StreamingCallbackHandler, h) - for h in run_manager.handlers - if isinstance(h, _StreamingCallbackHandler) - ), - None, + do_stream = ( + next( + ( + cast(_StreamingCallbackHandler, h) + for h in run_manager.handlers + if isinstance(h, _StreamingCallbackHandler) + ), + None, + ) + if _StreamingCallbackHandler is not None + else False ) try: # assign defaults @@ -2571,6 +2590,9 @@ class Pregel(PregelProtocol): stream.put_nowait, ((), "custom", c) ) ) + # set checkpointing mode for subgraphs + if checkpoint_during is not None: + config[CONF][CONFIG_KEY_CHECKPOINT_DURING] = checkpoint_during async with AsyncPregelLoop( input, input_model=self.input_model, @@ -2586,6 +2608,9 @@ class Pregel(PregelProtocol): interrupt_after=interrupt_after_, manager=run_manager, debug=debug, + checkpoint_during=checkpoint_during + if checkpoint_during is not None + else config[CONF].get(CONFIG_KEY_CHECKPOINT_DURING, True), trigger_to_nodes=self.trigger_to_nodes, migrate_checkpoint=self._migrate_checkpoint, ) as loop: @@ -2661,6 +2686,7 @@ class Pregel(PregelProtocol): output_keys: Optional[Union[str, Sequence[str]]] = None, interrupt_before: Optional[Union[All, Sequence[str]]] = None, interrupt_after: Optional[Union[All, Sequence[str]]] = None, + checkpoint_during: Optional[bool] = None, debug: Optional[bool] = None, **kwargs: Any, ) -> Union[dict[str, Any], Any]: @@ -2692,6 +2718,7 @@ class Pregel(PregelProtocol): output_keys=output_keys, interrupt_before=interrupt_before, interrupt_after=interrupt_after, + checkpoint_during=checkpoint_during, debug=debug, **kwargs, ): @@ -2713,6 +2740,7 @@ class Pregel(PregelProtocol): output_keys: Optional[Union[str, Sequence[str]]] = None, interrupt_before: Optional[Union[All, Sequence[str]]] = None, interrupt_after: Optional[Union[All, Sequence[str]]] = None, + checkpoint_during: Optional[bool] = None, debug: Optional[bool] = None, **kwargs: Any, ) -> Union[dict[str, Any], Any]: @@ -2745,6 +2773,7 @@ class Pregel(PregelProtocol): output_keys=output_keys, interrupt_before=interrupt_before, interrupt_after=interrupt_after, + checkpoint_during=checkpoint_during, debug=debug, **kwargs, ): diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 4c4ed6dde..e1de05602 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -63,6 +63,7 @@ from langgraph.constants import ( RESUME, SCHEDULED, TAG_HIDDEN, + TASKS, ) from langgraph.errors import ( CheckpointNotLatest, @@ -155,7 +156,7 @@ class PregelLoop(LoopProtocol): manager: Union[None, AsyncParentRunManager, ParentRunManager] interrupt_after: Union[All, Sequence[str]] interrupt_before: Union[All, Sequence[str]] - checkpoint_every_step: bool + checkpoint_during: bool debug: bool checkpointer_get_next_version: GetNextVersion @@ -180,6 +181,7 @@ class PregelLoop(LoopProtocol): channels: Mapping[str, BaseChannel] managed: ManagedValueMapping checkpoint: Checkpoint + checkpoint_id_saved: str checkpoint_ns: tuple[str, ...] checkpoint_config: RunnableConfig checkpoint_metadata: CheckpointMetadata @@ -215,7 +217,7 @@ class PregelLoop(LoopProtocol): debug: bool = False, migrate_checkpoint: Optional[Callable[[Checkpoint], None]] = None, trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, - checkpoint_every_step: bool = True, + checkpoint_during: bool = True, ) -> None: super().__init__( step=0, @@ -241,7 +243,7 @@ class PregelLoop(LoopProtocol): ) self._migrate_checkpoint = migrate_checkpoint self.trigger_to_nodes = trigger_to_nodes - self.checkpoint_every_step = checkpoint_every_step + self.checkpoint_during = checkpoint_during self.debug = debug if self.stream is not None and CONFIG_KEY_STREAM in config[CONF]: self.stream = DuplexStream(self.stream, config[CONF][CONFIG_KEY_STREAM]) @@ -294,29 +296,19 @@ class PregelLoop(LoopProtocol): """Put writes for a task, to be read by the next tick.""" if not writes: return + # always checkpoint writes containing Send, as they are fetched from the + # parent checkpoint, not the current one + checkpoint_during = self.checkpoint_during or any(w[0] == TASKS for w in writes) # deduplicate writes to special channels, last write wins if all(w[0] in WRITES_IDX_MAP for w in writes): writes = list({w[0]: w for w in writes}.values()) + # remove existing writes for this task + self.checkpoint_pending_writes = [ + w for w in self.checkpoint_pending_writes if w[0] != task_id + ] # save writes - for c, v in writes: - if ( - c in WRITES_IDX_MAP - and ( - idx := next( - ( - i - for i, w in enumerate(self.checkpoint_pending_writes) - if w[0] == task_id and w[1] == c - ), - None, - ) - ) - is not None - ): - self.checkpoint_pending_writes[idx] = (task_id, c, v) - else: - self.checkpoint_pending_writes.append((task_id, c, v)) - if self.checkpointer_put_writes is not None: + self.checkpoint_pending_writes.extend((task_id, c, v) for c, v in writes) + if checkpoint_during and self.checkpointer_put_writes is not None: config = patch_configurable( self.checkpoint_config, { @@ -349,6 +341,46 @@ class PregelLoop(LoopProtocol): if hasattr(self, "tasks"): self._output_writes(task_id, writes) + def _put_pending_writes(self) -> None: + if self.checkpointer_put_writes is None: + return + if not self.checkpoint_pending_writes: + return + # patch config + config = patch_configurable( + self.checkpoint_config, + { + CONFIG_KEY_CHECKPOINT_NS: self.config[CONF].get( + CONFIG_KEY_CHECKPOINT_NS, "" + ), + CONFIG_KEY_CHECKPOINT_ID: self.checkpoint["id"], + }, + ) + # group by task id + by_task = defaultdict(list) + for task_id, channel, value in self.checkpoint_pending_writes: + by_task[task_id].append((channel, value)) + # submit writes to checkpointer + for task_id, writes in by_task.items(): + if self.checkpointer_put_writes_accepts_task_path and hasattr( + self, "tasks" + ): + task = self.tasks.get(task_id) + self.submit( + self.checkpointer_put_writes, + config, + writes, + task_id, + task_path_str(task.path) if task else "", + ) + else: + self.submit( + self.checkpointer_put_writes, + config, + writes, + task_id, + ) + def accept_push( self, task: PregelExecutableTask, write_idx: int, call: Optional[Call] = None ) -> Optional[PregelExecutableTask]: @@ -711,32 +743,44 @@ class PregelLoop(LoopProtocol): def _put_checkpoint(self, metadata: CheckpointMetadata) -> None: # assign step and parents - metadata["step"] = self.step - metadata["parents"] = self.config[CONF].get(CONFIG_KEY_CHECKPOINT_MAP, {}) - # debug flag - if self.debug: - print_step_checkpoint( - metadata, - self.channels, - ( - [self.stream_keys] - if isinstance(self.stream_keys, str) - else self.stream_keys - ), - ) + exiting = metadata is self.checkpoint_metadata + if exiting and self.checkpoint["id"] == self.checkpoint_id_saved: + # checkpoint already saved + return + if not exiting: + metadata["step"] = self.step + metadata["parents"] = self.config[CONF].get(CONFIG_KEY_CHECKPOINT_MAP, {}) + self.checkpoint_metadata = metadata + # debug flag + if self.debug: + print_step_checkpoint( + metadata, + self.channels, + ( + [self.stream_keys] + if isinstance(self.stream_keys, str) + else self.stream_keys + ), + ) + self.checkpoint_id_prev = self.checkpoint["id"] if self.step > -1 else None + # do checkpoint? + do_checkpoint = self._checkpointer_put_after_previous is not None and ( + exiting or self.checkpoint_during + ) + # create new checkpoint + self.checkpoint = create_checkpoint( + self.checkpoint, + self.channels if do_checkpoint else None, + self.step, + id=self.checkpoint["id"] if exiting else None, + ) # bail if no checkpointer - if self._checkpointer_put_after_previous is not None: + if do_checkpoint and self._checkpointer_put_after_previous is not None: for k, v in self.config["metadata"].items(): if k in EXCLUDED_METADATA_KEYS: continue metadata.setdefault(k, v) # type: ignore - # create new checkpoint - self.checkpoint = create_checkpoint( - self.checkpoint, self.channels, self.step - ) - self.checkpoint_metadata = metadata - self.prev_checkpoint_config = ( self.checkpoint_config if CONFIG_KEY_CHECKPOINT_ID in self.checkpoint_config[CONF] @@ -747,6 +791,8 @@ class PregelLoop(LoopProtocol): **self.checkpoint_config, CONF: { **self.checkpoint_config[CONF], + # this is guaranteed to be set by code above + CONFIG_KEY_CHECKPOINT_ID: self.checkpoint_id_prev, CONFIG_KEY_CHECKPOINT_NS: self.config[CONF].get( CONFIG_KEY_CHECKPOINT_NS, "" ), @@ -777,8 +823,9 @@ class PregelLoop(LoopProtocol): CONFIG_KEY_CHECKPOINT_ID: self.checkpoint["id"], }, } - # increment step - self.step += 1 + if not exiting: + # increment step + self.step += 1 def _update_mv(self, key: str, values: Sequence[Any]) -> None: raise NotImplementedError @@ -789,6 +836,10 @@ class PregelLoop(LoopProtocol): exc_value: Optional[BaseException], traceback: Optional[TracebackType], ) -> Optional[bool]: + # persist current checkpoint and writes + if not self.checkpoint_during: + self._put_checkpoint(self.checkpoint_metadata) + self._put_pending_writes() # suppress interrupt suppress = isinstance(exc_value, GraphInterrupt) and not self.is_nested if suppress: @@ -907,6 +958,7 @@ class SyncPregelLoop(PregelLoop, ContextManager): debug: bool = False, migrate_checkpoint: Optional[Callable[[Checkpoint], None]] = None, trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, + checkpoint_during: bool = True, ) -> None: super().__init__( input, @@ -925,6 +977,7 @@ class SyncPregelLoop(PregelLoop, ContextManager): debug=debug, migrate_checkpoint=migrate_checkpoint, trigger_to_nodes=trigger_to_nodes, + checkpoint_during=checkpoint_during, ) self.stack = ExitStack() if checkpointer: @@ -1004,6 +1057,7 @@ class SyncPregelLoop(PregelLoop, ContextManager): }, } self.prev_checkpoint_config = saved.parent_config + self.checkpoint_id_saved = saved.checkpoint["id"] self.checkpoint = saved.checkpoint self.checkpoint_metadata = saved.metadata self.checkpoint_pending_writes = ( @@ -1054,6 +1108,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): debug: bool = False, migrate_checkpoint: Optional[Callable[[Checkpoint], None]] = None, trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, + checkpoint_during: bool = True, ) -> None: super().__init__( input, @@ -1072,6 +1127,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): debug=debug, migrate_checkpoint=migrate_checkpoint, trigger_to_nodes=trigger_to_nodes, + checkpoint_during=checkpoint_during, ) self.stack = AsyncExitStack() if checkpointer: @@ -1151,6 +1207,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): }, } self.prev_checkpoint_config = saved.parent_config + self.checkpoint_id_saved = saved.checkpoint["id"] self.checkpoint = saved.checkpoint self.checkpoint_metadata = saved.metadata self.checkpoint_pending_writes = ( diff --git a/libs/langgraph/langgraph/pregel/messages.py b/libs/langgraph/langgraph/pregel/messages.py index 867012fa6..5766c2b63 100644 --- a/libs/langgraph/langgraph/pregel/messages.py +++ b/libs/langgraph/langgraph/pregel/messages.py @@ -7,6 +7,7 @@ from typing import ( List, Optional, Sequence, + TypeVar, Union, cast, ) @@ -15,11 +16,16 @@ from uuid import UUID, uuid4 from langchain_core.callbacks import BaseCallbackHandler from langchain_core.messages import BaseMessage from langchain_core.outputs import ChatGenerationChunk, LLMResult -from langchain_core.tracers._streaming import T, _StreamingCallbackHandler from langgraph.constants import NS_SEP, TAG_HIDDEN, TAG_NOSTREAM from langgraph.types import StreamChunk +try: + from langchain_core.tracers._streaming import _StreamingCallbackHandler +except ImportError: + _StreamingCallbackHandler = object # type: ignore + +T = TypeVar("T") Meta = tuple[tuple[str, ...], dict[str, Any]] diff --git a/libs/langgraph/langgraph/pregel/remote.py b/libs/langgraph/langgraph/pregel/remote.py index 07d44cb66..13cc7af4f 100644 --- a/libs/langgraph/langgraph/pregel/remote.py +++ b/libs/langgraph/langgraph/pregel/remote.py @@ -10,7 +10,6 @@ from typing import ( cast, ) -import orjson from langchain_core.runnables import RunnableConfig from langchain_core.runnables.graph import ( Edge as DrawableEdge, @@ -35,6 +34,8 @@ from typing_extensions import Self from langgraph.checkpoint.base import CheckpointMetadata from langgraph.constants import ( CONF, + CONFIG_KEY_CHECKPOINT_ID, + CONFIG_KEY_CHECKPOINT_MAP, CONFIG_KEY_CHECKPOINT_NS, CONFIG_KEY_STREAM, INTERRUPT, @@ -46,6 +47,14 @@ from langgraph.pregel.types import All, PregelTask, StateSnapshot, StreamMode from langgraph.types import Command, Interrupt, StreamProtocol from langgraph.utils.config import merge_configs +CONF_DROPLIST = frozenset( + ( + CONFIG_KEY_CHECKPOINT_MAP, + CONFIG_KEY_CHECKPOINT_ID, + CONFIG_KEY_CHECKPOINT_NS, + ), +) + class RemoteException(Exception): """Exception raised when an error occurs in the remote graph.""" @@ -290,47 +299,26 @@ class RemoteGraph(PregelProtocol): } def _sanitize_config(self, config: RunnableConfig) -> RunnableConfig: - reserved_configurable_keys = frozenset( - [ - "callbacks", - "checkpoint_map", - "checkpoint_id", - "checkpoint_ns", - ] - ) - - def _sanitize_obj(obj: Any) -> Any: - """Remove non-JSON serializable fields from the given object.""" - if isinstance(obj, dict): - return {k: _sanitize_obj(v) for k, v in obj.items()} - elif isinstance(obj, list): - return [_sanitize_obj(v) for v in obj] - else: - try: - orjson.dumps(obj) - return obj - except orjson.JSONEncodeError: - return None - - # Remove non-JSON serializable fields from the config. - config = _sanitize_obj(config) - - # Only include configurable keys that are not reserved and - # not starting with "__pregel_" prefix. - new_configurable = { - k: v - for k, v in config["configurable"].items() - if k not in reserved_configurable_keys and not k.startswith("__pregel_") - } - - sanitized: RunnableConfig = { - "tags": config.get("tags") or [], - "metadata": config.get("metadata") or {}, - "configurable": new_configurable, - } + """Sanitize the config to remove non-serializable fields.""" + sanitized: RunnableConfig = {} if "recursion_limit" in config: sanitized["recursion_limit"] = config["recursion_limit"] - + if "tags" in config: + sanitized["tags"] = [tag for tag in config["tags"] if isinstance(tag, str)] + if "metadata" in config: + sanitized["metadata"] = {} + for k, v in config["metadata"].items(): + if isinstance(k, str) and isinstance(v, (str, int, float, bool)): + sanitized["metadata"][k] = v + if "configurable" in config: + sanitized["configurable"] = {} + for k, v in config["configurable"].items(): + if ( + isinstance(k, str) + and k not in CONF_DROPLIST + and isinstance(v, (str, int, float, bool)) + ): + sanitized["configurable"][k] = v return sanitized def get_state( diff --git a/libs/langgraph/langgraph/utils/runnable.py b/libs/langgraph/langgraph/utils/runnable.py index f21420ad0..4dab44023 100644 --- a/libs/langgraph/langgraph/utils/runnable.py +++ b/libs/langgraph/langgraph/utils/runnable.py @@ -36,7 +36,6 @@ from langchain_core.runnables.config import ( var_child_runnable_config, ) from langchain_core.runnables.utils import Input, Output -from langchain_core.tracers._streaming import _StreamingCallbackHandler from typing_extensions import TypeGuard from langgraph.constants import ( @@ -54,6 +53,11 @@ from langgraph.utils.config import ( patch_config, ) +try: + from langchain_core.tracers._streaming import _StreamingCallbackHandler +except ImportError: + _StreamingCallbackHandler = None # type: ignore + def _set_config_context( config: RunnableConfig, @@ -683,13 +687,15 @@ class RunnableSeq(Runnable): iterator = step.stream(input, config, **kwargs) else: iterator = step.transform(iterator, config) - if stream_handler := next( - ( - cast(_StreamingCallbackHandler, h) - for h in run_manager.handlers - if isinstance(h, _StreamingCallbackHandler) - ), - None, + if _StreamingCallbackHandler is not None and ( + stream_handler := next( + ( + cast(_StreamingCallbackHandler, h) + for h in run_manager.handlers + if isinstance(h, _StreamingCallbackHandler) + ), + None, + ) ): # populates streamed_output in astream_log() output if needed iterator = stream_handler.tap_output_iter(run_manager.run_id, iterator) @@ -749,13 +755,15 @@ class RunnableSeq(Runnable): aiterator = step.atransform(aiterator, config) if hasattr(aiterator, "aclose"): stack.push_async_callback(aiterator.aclose) - if stream_handler := next( - ( - cast(_StreamingCallbackHandler, h) - for h in run_manager.handlers - if isinstance(h, _StreamingCallbackHandler) - ), - None, + if _StreamingCallbackHandler is not None and ( + stream_handler := next( + ( + cast(_StreamingCallbackHandler, h) + for h in run_manager.handlers + if isinstance(h, _StreamingCallbackHandler) + ), + None, + ) ): # populates streamed_output in astream_log() output if needed aiterator = stream_handler.tap_output_aiter( diff --git a/libs/langgraph/pyproject.toml b/libs/langgraph/pyproject.toml index cece38afb..722b2e018 100644 --- a/libs/langgraph/pyproject.toml +++ b/libs/langgraph/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langgraph" -version = "0.3.25" +version = "0.3.27" description = "Building stateful, multi-actor applications with LLMs" authors = [] license = "MIT" diff --git a/libs/langgraph/tests/test_interruption.py b/libs/langgraph/tests/test_interruption.py index aa543cf00..f305d94d6 100644 --- a/libs/langgraph/tests/test_interruption.py +++ b/libs/langgraph/tests/test_interruption.py @@ -1,20 +1,20 @@ import pytest -from pytest_mock import MockerFixture from typing_extensions import TypedDict from langgraph.graph import END, START, StateGraph from tests.conftest import ( - ALL_CHECKPOINTERS_ASYNC, - ALL_CHECKPOINTERS_SYNC, + REGULAR_CHECKPOINTERS_ASYNC, + REGULAR_CHECKPOINTERS_SYNC, awith_checkpointer, ) pytestmark = pytest.mark.anyio -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +@pytest.mark.parametrize("checkpoint_during", [True, False]) +@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC) def test_interruption_without_state_updates( - request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture + request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool ) -> None: """Test interruption without state updates. This test confirms that interrupting doesn't require a state key having been updated in the prev step""" @@ -40,20 +40,27 @@ def test_interruption_without_state_updates( initial_input = {"input": "hello world"} thread = {"configurable": {"thread_id": "1"}} - graph.invoke(initial_input, thread, debug=True) + graph.invoke(initial_input, thread, checkpoint_during=checkpoint_during) assert graph.get_state(thread).next == ("step_2",) + n_checkpoints = len([c for c in graph.get_state_history(thread)]) + assert n_checkpoints == (3 if checkpoint_during else 1) - graph.invoke(None, thread, debug=True) + graph.invoke(None, thread, checkpoint_during=checkpoint_during) assert graph.get_state(thread).next == ("step_3",) + n_checkpoints = len([c for c in graph.get_state_history(thread)]) + assert n_checkpoints == (4 if checkpoint_during else 2) - graph.invoke(None, thread, debug=True) + graph.invoke(None, thread, checkpoint_during=checkpoint_during) assert graph.get_state(thread).next == () + n_checkpoints = len([c for c in graph.get_state_history(thread)]) + assert n_checkpoints == (5 if checkpoint_during else 3) -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +@pytest.mark.parametrize("checkpoint_during", [True, False]) +@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC) async def test_interruption_without_state_updates_async( - checkpointer_name: str, mocker: MockerFixture -): + checkpointer_name: str, checkpoint_during: bool +) -> None: """Test interruption without state updates. This test confirms that interrupting doesn't require a state key having been updated in the prev step""" @@ -78,11 +85,17 @@ async def test_interruption_without_state_updates_async( initial_input = {"input": "hello world"} thread = {"configurable": {"thread_id": "1"}} - await graph.ainvoke(initial_input, thread, debug=True) + await graph.ainvoke(initial_input, thread, checkpoint_during=checkpoint_during) assert (await graph.aget_state(thread)).next == ("step_2",) + n_checkpoints = len([c async for c in graph.aget_state_history(thread)]) + assert n_checkpoints == (3 if checkpoint_during else 1) - await graph.ainvoke(None, thread, debug=True) + await graph.ainvoke(None, thread, checkpoint_during=checkpoint_during) assert (await graph.aget_state(thread)).next == ("step_3",) + n_checkpoints = len([c async for c in graph.aget_state_history(thread)]) + assert n_checkpoints == (4 if checkpoint_during else 2) - await graph.ainvoke(None, thread, debug=True) + await graph.ainvoke(None, thread, checkpoint_during=checkpoint_during) assert (await graph.aget_state(thread)).next == () + n_checkpoints = len([c async for c in graph.aget_state_history(thread)]) + assert n_checkpoints == (5 if checkpoint_during else 3) diff --git a/libs/langgraph/tests/test_large_cases.py b/libs/langgraph/tests/test_large_cases.py index f15cd5c23..f5be91cee 100644 --- a/libs/langgraph/tests/test_large_cases.py +++ b/libs/langgraph/tests/test_large_cases.py @@ -7258,9 +7258,10 @@ def test_branch_then( ) -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +@pytest.mark.parametrize("checkpoint_during", [True, False]) +@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC) def test_send_dedupe_on_resume( - request: pytest.FixtureRequest, checkpointer_name: str + request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool ) -> None: checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") @@ -7316,7 +7317,7 @@ def test_send_dedupe_on_resume( graph = builder.compile(checkpointer=checkpointer) thread1 = {"configurable": {"thread_id": "1"}} - assert graph.invoke(["0"], thread1, debug=1) == [ + assert graph.invoke(["0"], thread1, checkpoint_during=checkpoint_during) == [ "0", "1", "3.1", @@ -7333,12 +7334,11 @@ def test_send_dedupe_on_resume( pytest.xfail("TODO: shallow checkpointer reports wrong next set") assert state.next == ("flaky",) # check history - if "shallow" not in checkpointer_name: - history = [c for c in graph.get_state_history(thread1)] - assert len(history) == 4 + history = [c for c in graph.get_state_history(thread1)] + assert len(history) == (4 if checkpoint_during else 1) # resume execution - assert graph.invoke(None, thread1, debug=1) == [ + assert graph.invoke(None, thread1, checkpoint_during=checkpoint_during) == [ "0", "1", "3.1", @@ -7358,6 +7358,7 @@ def test_send_dedupe_on_resume( assert state.next == () # check history history = [c for c in graph.get_state_history(thread1)] + assert len(history) == (6 if checkpoint_during else 2) expected_history = [ StateSnapshot( values=[ @@ -7494,13 +7495,9 @@ def test_send_dedupe_on_resume( name="flaky", path=("__pregel_push", 1), error=None, - interrupts=( - Interrupt( - value="Bahh", resumable=False, ns=None, when="during" - ), - ), + interrupts=(Interrupt(value="Bahh", resumable=False, ns=None),), state=None, - result=["flaky|4"], + result=["flaky|4"] if checkpoint_during else None, ), PregelTask( id=AnyStr(), @@ -7637,10 +7634,11 @@ def test_send_dedupe_on_resume( ), ), ] - if "shallow" in checkpointer_name: - expected_history = expected_history[:1] - - assert history == expected_history + if checkpoint_during: + assert history == expected_history + else: + assert history[0] == expected_history[0] + assert history[1] == expected_history[2] @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 2d1542477..be3f94ec7 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -1121,10 +1121,14 @@ def test_invoke_checkpoint_two( assert checkpoint["channel_values"].get("total") == 5 +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_pending_writes_resume( - request: pytest.FixtureRequest, checkpointer_name: str + request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool ) -> None: + if not checkpoint_during and "shallow" in checkpointer_name: + pytest.skip("Checkpointing during execution not supported") + checkpointer: BaseCheckpointSaver = request.getfixturevalue( f"checkpointer_{checkpointer_name}" ) @@ -1150,17 +1154,19 @@ def test_pending_writes_resume( self.calls = 0 one = AwhileMaker(0.1, {"value": 2}) - two = AwhileMaker(0.3, ConnectionError("I'm not good")) + two = AwhileMaker(0.2, ConnectionError("I'm not good")) builder = StateGraph(State) builder.add_node("one", one) - builder.add_node("two", two, retry=RetryPolicy(max_attempts=2)) + builder.add_node( + "two", two, retry=RetryPolicy(max_attempts=2, initial_interval=0, jitter=False) + ) builder.add_edge(START, "one") builder.add_edge(START, "two") graph = builder.compile(checkpointer=checkpointer) thread1: RunnableConfig = {"configurable": {"thread_id": "1"}} with pytest.raises(ConnectionError, match="I'm not good"): - graph.invoke({"value": 1}, thread1) + graph.invoke({"value": 1}, thread1, checkpoint_during=checkpoint_during) # both nodes should have been called once assert one.calls == 1 @@ -1206,7 +1212,7 @@ def test_pending_writes_resume( # resume execution with pytest.raises(ConnectionError, match="I'm not good"): - graph.invoke(None, thread1) + graph.invoke(None, thread1, checkpoint_during=checkpoint_during) # node "one" succeeded previously, so shouldn't be called again assert one.calls == 1 @@ -1220,7 +1226,9 @@ def test_pending_writes_resume( # resume execution, without exception two.rtn = {"value": 3} # both the pending write and the new write were applied, 1 + 2 + 3 = 6 - assert graph.invoke(None, thread1) == {"value": 6} + assert graph.invoke(None, thread1, checkpoint_during=checkpoint_during) == { + "value": 6 + } if "shallow" in checkpointer_name: assert len(list(checkpointer.list(thread1))) == 1 @@ -1229,7 +1237,7 @@ def test_pending_writes_resume( # check all final checkpoints checkpoints = [c for c in checkpointer.list(thread1)] # we should have 3 - assert len(checkpoints) == 3 + assert len(checkpoints) == (3 if checkpoint_during else 2) # the last one not too interesting for this test assert checkpoints[0] == CheckpointTuple( config={ @@ -1331,15 +1339,26 @@ def test_pending_writes_resume( "configurable": { "thread_id": "1", "checkpoint_ns": "", - "checkpoint_id": checkpoints[2].config["configurable"]["checkpoint_id"], + "checkpoint_id": checkpoints[2].config["configurable"]["checkpoint_id"] + if checkpoint_during + else AnyStr(), } }, pending_writes=UnsortedSequence( (AnyStr(), "value", 2), (AnyStr(), "__error__", 'ConnectionError("I\'m not good")'), (AnyStr(), "value", 3), + ) + if checkpoint_during + else UnsortedSequence( + (AnyStr(), "value", 2), + (AnyStr(), "__error__", 'ConnectionError("I\'m not good")'), + # the write against the previous checkpoint is not saved, as it is + # produced in a run where only the next checkpoint (the last) is saved ), ) + if not checkpoint_during: + return assert checkpoints[2] == CheckpointTuple( config={ "configurable": { @@ -1497,8 +1516,14 @@ def test_send_sequences() -> None: ] +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_imp_task(request: pytest.FixtureRequest, checkpointer_name: str) -> None: +def test_imp_task( + request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool +) -> None: + if not checkpoint_during and "shallow" in checkpointer_name: + pytest.skip("Checkpointing during execution not supported") + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") mapper_calls = 0 @@ -1564,7 +1589,7 @@ def test_imp_task(request: pytest.FixtureRequest, checkpointer_name: str) -> Non } thread1 = {"configurable": {"thread_id": "1"}} - assert [*graph.stream([0, 1], thread1)] == [ + assert [*graph.stream([0, 1], thread1, checkpoint_during=checkpoint_during)] == [ {"mapper": "00"}, {"mapper": "11"}, { @@ -1580,17 +1605,23 @@ def test_imp_task(request: pytest.FixtureRequest, checkpointer_name: str) -> Non ] assert mapper_calls == 2 - assert graph.invoke(Command(resume="answer"), thread1) == [ + assert graph.invoke( + Command(resume="answer"), thread1, checkpoint_during=checkpoint_during + ) == [ "00answer", "11answer", ] assert mapper_calls == 2 +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_imp_nested( - request: pytest.FixtureRequest, checkpointer_name: str, snapshot: SnapshotAssertion + request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool ) -> None: + if not checkpoint_during and "shallow" in checkpointer_name: + pytest.skip("Checkpointing during execution not supported") + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") def mynode(input: list[str]) -> list[str]: @@ -1632,7 +1663,7 @@ def test_imp_nested( } thread1 = {"configurable": {"thread_id": "1"}} - assert [*graph.stream([0, 1], thread1)] == [ + assert [*graph.stream([0, 1], thread1, checkpoint_during=checkpoint_during)] == [ {"submapper": "0"}, {"mapper": "00"}, {"submapper": "1"}, @@ -1649,16 +1680,22 @@ def test_imp_nested( }, ] - assert graph.invoke(Command(resume="answer"), thread1) == [ + assert graph.invoke( + Command(resume="answer"), thread1, checkpoint_during=checkpoint_during + ) == [ "00answera", "11answera", ] +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_imp_stream_order( - request: pytest.FixtureRequest, checkpointer_name: str, snapshot: SnapshotAssertion + request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool ) -> None: + if not checkpoint_during and "shallow" in checkpointer_name: + pytest.skip("Checkpointing during execution not supported") + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") @task() @@ -1681,7 +1718,10 @@ def test_imp_stream_order( return fut_baz.result() thread1 = {"configurable": {"thread_id": "1"}} - assert [c for c in graph.stream({"a": "0"}, thread1)] == [ + assert [ + c + for c in graph.stream({"a": "0"}, thread1, checkpoint_during=checkpoint_during) + ] == [ { "foo": ( "0foo", @@ -3807,10 +3847,14 @@ def test_nested_graph(snapshot: SnapshotAssertion) -> None: ] +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_subgraph_checkpoint_true( - request: pytest.FixtureRequest, checkpointer_name: str + request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool ) -> None: + if not checkpoint_during and "shallow" in checkpointer_name: + pytest.skip("Unsupported combo") + checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) class InnerState(TypedDict): @@ -3842,7 +3886,12 @@ def test_subgraph_checkpoint_true( app = graph.compile(checkpointer=checkpointer) config = {"configurable": {"thread_id": "2"}} - assert [c for c in app.stream({"my_key": ""}, config, subgraphs=True)] == [ + assert [ + c + for c in app.stream( + {"my_key": ""}, config, subgraphs=True, checkpoint_during=checkpoint_during + ) + ] == [ (("inner",), {"inner_1": {"my_key": " got here", "my_other_key": ""}}), (("inner",), {"inner_2": {"my_key": " and there"}}), ((), {"inner": {"my_key": " got here and there"}}), @@ -3867,10 +3916,14 @@ def test_subgraph_checkpoint_true( ] +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_subgraph_checkpoint_true_interrupt( - request: pytest.FixtureRequest, checkpointer_name: str + request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool ) -> None: + if not checkpoint_during and "shallow" in checkpointer_name: + pytest.skip("Unsupported combo") + checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) # Define subgraph @@ -3909,15 +3962,18 @@ def test_subgraph_checkpoint_true_interrupt( builder.add_edge(START, "node_1") builder.add_edge("node_1", "node_2") - checkpointer = MemorySaver() graph = builder.compile(checkpointer=checkpointer) config = {"configurable": {"thread_id": "1"}} - assert graph.invoke({"foo": "foo"}, config) == {"foo": "hi! foo"} + assert graph.invoke( + {"foo": "foo"}, config, checkpoint_during=checkpoint_during + ) == {"foo": "hi! foo"} assert graph.get_state(config, subgraphs=True).tasks[0].state.values == { "bar": "hi! foo" } - assert graph.invoke(Command(resume="baz"), config) == {"foo": "hi! foobaz"} + assert graph.invoke( + Command(resume="baz"), config, checkpoint_during=checkpoint_during + ) == {"foo": "hi! foobaz"} @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) @@ -4033,10 +4089,14 @@ def test_stream_buffering_single_node( ] +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_nested_graph_interrupts_parallel( - request: pytest.FixtureRequest, checkpointer_name: str + request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool ) -> None: + if not checkpoint_during and "shallow" in checkpointer_name: + pytest.skip("Unsupported combo") + checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) class InnerState(TypedDict): @@ -4083,11 +4143,11 @@ def test_nested_graph_interrupts_parallel( # test invoke w/ nested interrupt config = {"configurable": {"thread_id": "1"}} - assert app.invoke({"my_key": ""}, config, debug=True) == { + assert app.invoke({"my_key": ""}, config, checkpoint_during=checkpoint_during) == { "my_key": " and parallel", } - assert app.invoke(None, config, debug=True) == { + assert app.invoke(None, config, checkpoint_during=checkpoint_during) == { "my_key": "got here and there and parallel and back again", } @@ -4096,13 +4156,17 @@ def test_nested_graph_interrupts_parallel( # - the writes of outer are persisted in 1st call and used in 2nd call, ie outer isn't called again (because we dont see outer_1 output again in 2nd stream) # test stream updates w/ nested interrupt config = {"configurable": {"thread_id": "2"}} - assert [*app.stream({"my_key": ""}, config, subgraphs=True)] == [ + assert [ + *app.stream( + {"my_key": ""}, config, subgraphs=True, checkpoint_during=checkpoint_during + ) + ] == [ # we got to parallel node first ((), {"outer_1": {"my_key": " and parallel"}}), ((AnyStr("inner:"),), {"inner_1": {"my_key": "got here", "my_other_key": ""}}), ((), {"__interrupt__": ()}), ] - assert [*app.stream(None, config)] == [ + assert [*app.stream(None, config, checkpoint_during=checkpoint_during)] == [ {"outer_1": {"my_key": " and parallel"}, "__metadata__": {"cached": True}}, {"inner": {"my_key": "got here and there"}}, {"outer_2": {"my_key": " and back again"}}, @@ -4110,11 +4174,22 @@ def test_nested_graph_interrupts_parallel( # test stream values w/ nested interrupt config = {"configurable": {"thread_id": "3"}} - assert [*app.stream({"my_key": ""}, config, stream_mode="values")] == [ + assert [ + *app.stream( + {"my_key": ""}, + config, + stream_mode="values", + checkpoint_during=checkpoint_during, + ) + ] == [ {"my_key": ""}, {"my_key": " and parallel"}, ] - assert [*app.stream(None, config, stream_mode="values")] == [ + assert [ + *app.stream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ {"my_key": ""}, {"my_key": "got here and there and parallel"}, {"my_key": "got here and there and parallel and back again"}, @@ -4123,15 +4198,28 @@ def test_nested_graph_interrupts_parallel( # test interrupts BEFORE the parallel node app = graph.compile(checkpointer=checkpointer, interrupt_before=["outer_1"]) config = {"configurable": {"thread_id": "4"}} - assert [*app.stream({"my_key": ""}, config, stream_mode="values")] == [ - {"my_key": ""} - ] + assert [ + *app.stream( + {"my_key": ""}, + config, + stream_mode="values", + checkpoint_during=checkpoint_during, + ) + ] == [{"my_key": ""}] # while we're waiting for the node w/ interrupt inside to finish - assert [*app.stream(None, config, stream_mode="values")] == [ + assert [ + *app.stream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ {"my_key": ""}, {"my_key": " and parallel"}, ] - assert [*app.stream(None, config, stream_mode="values")] == [ + assert [ + *app.stream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ {"my_key": ""}, {"my_key": "got here and there and parallel"}, {"my_key": "got here and there and parallel and back again"}, @@ -4140,24 +4228,43 @@ def test_nested_graph_interrupts_parallel( # test interrupts AFTER the parallel node app = graph.compile(checkpointer=checkpointer, interrupt_after=["outer_1"]) config = {"configurable": {"thread_id": "5"}} - assert [*app.stream({"my_key": ""}, config, stream_mode="values")] == [ + assert [ + *app.stream( + {"my_key": ""}, + config, + stream_mode="values", + checkpoint_during=checkpoint_during, + ) + ] == [ {"my_key": ""}, {"my_key": " and parallel"}, ] - assert [*app.stream(None, config, stream_mode="values")] == [ + assert [ + *app.stream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ {"my_key": ""}, {"my_key": "got here and there and parallel"}, ] - assert [*app.stream(None, config, stream_mode="values")] == [ + assert [ + *app.stream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ {"my_key": "got here and there and parallel"}, {"my_key": "got here and there and parallel and back again"}, ] +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_doubly_nested_graph_interrupts( - request: pytest.FixtureRequest, checkpointer_name: str + request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool ) -> None: + if not checkpoint_during and "shallow" in checkpointer_name: + pytest.skip("Unsupported combo") + checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) class State(TypedDict): @@ -4211,11 +4318,13 @@ def test_doubly_nested_graph_interrupts( # test invoke w/ nested interrupt config = {"configurable": {"thread_id": "1"}} - assert app.invoke({"my_key": "my value"}, config, debug=True) == { + assert app.invoke( + {"my_key": "my value"}, config, checkpoint_during=checkpoint_during + ) == { "my_key": "hi my value", } - assert app.invoke(None, config, debug=True) == { + assert app.invoke(None, config, checkpoint_during=checkpoint_during) == { "my_key": "hi my value here and there and back again", } @@ -4224,12 +4333,14 @@ def test_doubly_nested_graph_interrupts( config = { "configurable": {"thread_id": "2", CONFIG_KEY_NODE_FINISHED: nodes.append} } - assert [*app.stream({"my_key": "my value"}, config)] == [ + assert [ + *app.stream({"my_key": "my value"}, config, checkpoint_during=checkpoint_during) + ] == [ {"parent_1": {"my_key": "hi my value"}}, {"__interrupt__": ()}, ] assert nodes == ["parent_1", "grandchild_1"] - assert [*app.stream(None, config)] == [ + assert [*app.stream(None, config, checkpoint_during=checkpoint_during)] == [ {"child": {"my_key": "hi my value here and there"}}, {"parent_2": {"my_key": "hi my value here and there and back again"}}, ] @@ -4244,11 +4355,22 @@ def test_doubly_nested_graph_interrupts( # test stream values w/ nested interrupt config = {"configurable": {"thread_id": "3"}} - assert [*app.stream({"my_key": "my value"}, config, stream_mode="values")] == [ + assert [ + *app.stream( + {"my_key": "my value"}, + config, + stream_mode="values", + checkpoint_during=checkpoint_during, + ) + ] == [ {"my_key": "my value"}, {"my_key": "hi my value"}, ] - assert [*app.stream(None, config, stream_mode="values")] == [ + assert [ + *app.stream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ {"my_key": "hi my value"}, {"my_key": "hi my value here and there"}, {"my_key": "hi my value here and there and back again"}, diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 375beb261..0e74f0b74 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -1947,10 +1947,14 @@ async def test_invoke_checkpoint(mocker: MockerFixture, checkpointer_name: str) assert checkpoint["channel_values"].get("total") == 5 +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_pending_writes_resume( - request: pytest.FixtureRequest, checkpointer_name: str + checkpointer_name: str, checkpoint_during: bool ) -> None: + if not checkpoint_during and "shallow" in checkpointer_name: + pytest.skip("Checkpointing during execution not supported") + class State(TypedDict): value: Annotated[int, operator.add] @@ -1972,10 +1976,12 @@ async def test_pending_writes_resume( self.calls = 0 one = AwhileMaker(0.1, {"value": 2}) - two = AwhileMaker(0.3, ConnectionError("I'm not good")) + two = AwhileMaker(0.2, ConnectionError("I'm not good")) builder = StateGraph(State) builder.add_node("one", one) - builder.add_node("two", two, retry=RetryPolicy(max_attempts=2)) + builder.add_node( + "two", two, retry=RetryPolicy(max_attempts=2, initial_interval=0, jitter=False) + ) builder.add_edge(START, "one") builder.add_edge(START, "two") async with awith_checkpointer(checkpointer_name) as checkpointer: @@ -1983,7 +1989,9 @@ async def test_pending_writes_resume( thread1: RunnableConfig = {"configurable": {"thread_id": "1"}} with pytest.raises(ConnectionError, match="I'm not good"): - await graph.ainvoke({"value": 1}, thread1) + await graph.ainvoke( + {"value": 1}, thread1, checkpoint_during=checkpoint_during + ) # both nodes should have been called once assert one.calls == 1 @@ -2034,7 +2042,7 @@ async def test_pending_writes_resume( # resume execution with pytest.raises(ConnectionError, match="I'm not good"): - await graph.ainvoke(None, thread1) + await graph.ainvoke(None, thread1, checkpoint_during=checkpoint_during) # node "one" succeeded previously, so shouldn't be called again assert one.calls == 1 @@ -2048,7 +2056,9 @@ async def test_pending_writes_resume( # resume execution, without exception two.rtn = {"value": 3} # both the pending write and the new write were applied, 1 + 2 + 3 = 6 - assert await graph.ainvoke(None, thread1) == {"value": 6} + assert await graph.ainvoke( + None, thread1, checkpoint_during=checkpoint_during + ) == {"value": 6} if "shallow" in checkpointer_name: assert len([c async for c in checkpointer.alist(thread1)]) == 1 @@ -2057,7 +2067,7 @@ async def test_pending_writes_resume( # check all final checkpoints checkpoints = [c async for c in checkpointer.alist(thread1)] # we should have 3 - assert len(checkpoints) == 3 + assert len(checkpoints) == (3 if checkpoint_during else 2) # the last one not too interesting for this test assert checkpoints[0] == CheckpointTuple( config={ @@ -2163,15 +2173,26 @@ async def test_pending_writes_resume( "checkpoint_ns": "", "checkpoint_id": checkpoints[2].config["configurable"][ "checkpoint_id" - ], + ] + if checkpoint_during + else AnyStr(), } }, pending_writes=UnsortedSequence( (AnyStr(), "value", 2), (AnyStr(), "__error__", 'ConnectionError("I\'m not good")'), (AnyStr(), "value", 3), + ) + if checkpoint_during + else UnsortedSequence( + (AnyStr(), "value", 2), + (AnyStr(), "__error__", 'ConnectionError("I\'m not good")'), + # the write against the previous checkpoint is not saved, as it is + # produced in a run where only the next checkpoint (the last) is saved ), ) + if not checkpoint_during: + return assert checkpoints[2] == CheckpointTuple( config={ "configurable": { @@ -2209,7 +2230,7 @@ async def test_pending_writes_resume( @pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC) async def test_run_from_checkpoint_id_retains_previous_writes( - request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture + checkpointer_name: str, ) -> None: class MyState(TypedDict): myval: Annotated[int, operator.add] @@ -2254,8 +2275,8 @@ async def test_run_from_checkpoint_id_retains_previous_writes( history = [c async for c in graph.aget_state_history(thread1)] assert len(history) == 4 - assert history[-1].values == {"myval": 0} assert history[0].values == {"myval": 4, "otherval": False} + assert history[-1].values == {"myval": 0} second_run_config = { **thread1, @@ -2432,8 +2453,12 @@ async def test_send_sequences(checkpointer_name: str) -> None: @NEEDS_CONTEXTVARS +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_imp_task(checkpointer_name: str) -> None: +async def test_imp_task(checkpointer_name: str, checkpoint_during: bool) -> None: + if not checkpoint_during and "shallow" in checkpointer_name: + pytest.skip("Checkpointing during execution not supported") + async with awith_checkpointer(checkpointer_name) as checkpointer: mapper_calls = 0 @@ -2453,7 +2478,12 @@ async def test_imp_task(checkpointer_name: str) -> None: tracer = FakeTracer() thread1 = {"configurable": {"thread_id": "1"}, "callbacks": [tracer]} - assert [c async for c in graph.astream([0, 1], thread1)] == [ + assert [ + c + async for c in graph.astream( + [0, 1], thread1, checkpoint_during=checkpoint_during + ) + ] == [ {"mapper": "00"}, {"mapper": "11"}, { @@ -2477,7 +2507,9 @@ async def test_imp_task(checkpointer_name: str) -> None: assert any(r.inputs == {"input": 0} for r in mapper_runs) assert any(r.inputs == {"input": 1} for r in mapper_runs) - assert await graph.ainvoke(Command(resume="answer"), thread1) == [ + assert await graph.ainvoke( + Command(resume="answer"), thread1, checkpoint_during=checkpoint_during + ) == [ "00answer", "11answer", ] @@ -2485,8 +2517,12 @@ async def test_imp_task(checkpointer_name: str) -> None: @NEEDS_CONTEXTVARS +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_imp_nested(checkpointer_name: str) -> None: +async def test_imp_nested(checkpointer_name: str, checkpoint_during: bool) -> None: + if not checkpoint_during and "shallow" in checkpointer_name: + pytest.skip("Checkpointing during execution not supported") + async def mynode(input: list[str]) -> list[str]: return [it + "a" for it in input] @@ -2526,7 +2562,12 @@ async def test_imp_nested(checkpointer_name: str) -> None: } thread1 = {"configurable": {"thread_id": "1"}} - assert [c async for c in graph.astream([0, 1], thread1)] == [ + assert [ + c + async for c in graph.astream( + [0, 1], thread1, checkpoint_during=checkpoint_during + ) + ] == [ {"submapper": "0"}, {"mapper": "00"}, {"submapper": "1"}, @@ -2543,15 +2584,21 @@ async def test_imp_nested(checkpointer_name: str) -> None: }, ] - assert await graph.ainvoke(Command(resume="answer"), thread1) == [ + assert await graph.ainvoke( + Command(resume="answer"), thread1, checkpoint_during=checkpoint_during + ) == [ "00answera", "11answera", ] @NEEDS_CONTEXTVARS +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_imp_task_cancel(checkpointer_name: str) -> None: +async def test_imp_task_cancel(checkpointer_name: str, checkpoint_during: bool) -> None: + if not checkpoint_during and "shallow" in checkpointer_name: + pytest.skip("Checkpointing during execution not supported") + async with awith_checkpointer(checkpointer_name) as checkpointer: mapper_calls = 0 mapper_cancels = 0 @@ -2577,7 +2624,12 @@ async def test_imp_task_cancel(checkpointer_name: str) -> None: return [m + answer for m in mapped] thread1 = {"configurable": {"thread_id": "1"}} - assert [c async for c in graph.astream([0, 1], thread1)] == [ + assert [ + c + async for c in graph.astream( + [0, 1], thread1, checkpoint_during=checkpoint_during + ) + ] == [ {"mapper": "00"}, { "__interrupt__": ( @@ -2593,7 +2645,9 @@ async def test_imp_task_cancel(checkpointer_name: str) -> None: assert mapper_calls == 2 assert mapper_cancels == 1 - assert await graph.ainvoke(Command(resume="answer"), thread1) == [ + assert await graph.ainvoke( + Command(resume="answer"), thread1, checkpoint_during=checkpoint_during + ) == [ "00answer", ] assert mapper_calls == 3 @@ -2601,8 +2655,14 @@ async def test_imp_task_cancel(checkpointer_name: str) -> None: @NEEDS_CONTEXTVARS +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_imp_sync_from_async(checkpointer_name: str) -> None: +async def test_imp_sync_from_async( + checkpointer_name: str, checkpoint_during: bool +) -> None: + if not checkpoint_during and "shallow" in checkpointer_name: + pytest.skip("Checkpointing during execution not supported") + async with awith_checkpointer(checkpointer_name) as checkpointer: @task() @@ -2625,7 +2685,12 @@ async def test_imp_sync_from_async(checkpointer_name: str) -> None: return fut_baz.result() thread1 = {"configurable": {"thread_id": "1"}} - assert [c async for c in graph.astream({"a": "0"}, thread1)] == [ + assert [ + c + async for c in graph.astream( + {"a": "0"}, thread1, checkpoint_during=checkpoint_during + ) + ] == [ {"foo": {"a": "0foo", "b": "bar"}}, {"bar": {"a": "0foobar", "c": "bark"}}, {"baz": {"a": "0foobarbaz", "c": "something else"}}, @@ -2634,8 +2699,14 @@ async def test_imp_sync_from_async(checkpointer_name: str) -> None: @NEEDS_CONTEXTVARS +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_imp_stream_order(checkpointer_name: str) -> None: +async def test_imp_stream_order( + checkpointer_name: str, checkpoint_during: bool +) -> None: + if not checkpoint_during and "shallow" in checkpointer_name: + pytest.skip("Checkpointing during execution not supported") + async with awith_checkpointer(checkpointer_name) as checkpointer: @task() @@ -2659,7 +2730,12 @@ async def test_imp_stream_order(checkpointer_name: str) -> None: return await fut_baz thread1 = {"configurable": {"thread_id": "1"}} - assert [c async for c in graph.astream({"a": "0"}, thread1)] == [ + assert [ + c + async for c in graph.astream( + {"a": "0"}, thread1, checkpoint_during=checkpoint_during + ) + ] == [ {"foo": {"a": "0foo", "b": "bar"}}, {"bar": {"a": "0foobar", "c": "bark"}}, {"baz": {"a": "0foobarbaz", "c": "something else"}}, @@ -2667,8 +2743,11 @@ async def test_imp_stream_order(checkpointer_name: str) -> None: ] +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC) -async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: +async def test_send_dedupe_on_resume( + checkpointer_name: str, checkpoint_during: bool +) -> None: class InterruptOnce: ticks: int = 0 @@ -2719,7 +2798,9 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: async with awith_checkpointer(checkpointer_name) as checkpointer: graph = builder.compile(checkpointer=checkpointer) thread1 = {"configurable": {"thread_id": "1"}} - assert await graph.ainvoke(["0"], thread1, debug=1) == [ + assert await graph.ainvoke( + ["0"], thread1, checkpoint_during=checkpoint_during + ) == [ "0", "1", "3.1", @@ -2731,7 +2812,9 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: assert builder.nodes["2"].runnable.func.ticks == 3 assert builder.nodes["flaky"].runnable.func.ticks == 1 # resume execution - assert await graph.ainvoke(None, thread1, debug=1) == [ + assert await graph.ainvoke( + None, thread1, checkpoint_during=checkpoint_during + ) == [ "0", "1", "3.1", @@ -2748,7 +2831,8 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: assert builder.nodes["flaky"].runnable.func.ticks == 2 # check history history = [c async for c in graph.aget_state_history(thread1)] - assert history == [ + assert len(history) == (6 if checkpoint_during else 2) + expected_history = [ StateSnapshot( values=[ "0", @@ -2884,13 +2968,9 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: name="flaky", path=("__pregel_push", 1), error=None, - interrupts=( - Interrupt( - value="Bahh", resumable=False, ns=None, when="during" - ), - ), + interrupts=(Interrupt(value="Bahh", resumable=False, ns=None),), state=None, - result=["flaky|4"], + result=["flaky|4"] if checkpoint_during else None, ), PregelTask( id=AnyStr(), @@ -3027,6 +3107,11 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: ), ), ] + if checkpoint_during: + assert history == expected_history + else: + assert history[0] == expected_history[0] + assert history[1] == expected_history[2] @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) @@ -5350,6 +5435,132 @@ async def test_nested_graph(snapshot: SnapshotAssertion) -> None: assert times_called == 1 +@pytest.mark.parametrize("checkpoint_during", [True, False]) +@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC) +async def test_subgraph_checkpoint_true( + checkpointer_name: str, checkpoint_during: bool +) -> None: + class InnerState(TypedDict): + my_key: Annotated[str, operator.add] + my_other_key: str + + def inner_1(state: InnerState): + return {"my_key": " got here", "my_other_key": state["my_key"]} + + def inner_2(state: InnerState): + return {"my_key": " and there"} + + inner = StateGraph(InnerState) + inner.add_node("inner_1", inner_1) + inner.add_node("inner_2", inner_2) + inner.add_edge("inner_1", "inner_2") + inner.set_entry_point("inner_1") + inner.set_finish_point("inner_2") + + class State(TypedDict): + my_key: str + + graph = StateGraph(State) + graph.add_node("inner", inner.compile(checkpointer=True)) + graph.add_edge(START, "inner") + graph.add_conditional_edges( + "inner", lambda s: "inner" if s["my_key"].count("there") < 2 else END + ) + + async with awith_checkpointer(checkpointer_name) as checkpointer: + app = graph.compile(checkpointer=checkpointer) + + config = {"configurable": {"thread_id": "2"}} + assert [ + c + async for c in app.astream( + {"my_key": ""}, + config, + subgraphs=True, + checkpoint_during=checkpoint_during, + ) + ] == [ + (("inner",), {"inner_1": {"my_key": " got here", "my_other_key": ""}}), + (("inner",), {"inner_2": {"my_key": " and there"}}), + ((), {"inner": {"my_key": " got here and there"}}), + ( + ("inner",), + { + "inner_1": { + "my_key": " got here", + "my_other_key": " got here and there got here and there", + } + }, + ), + (("inner",), {"inner_2": {"my_key": " and there"}}), + ( + (), + { + "inner": { + "my_key": " got here and there got here and there got here and there" + } + }, + ), + ] + + +@NEEDS_CONTEXTVARS +@pytest.mark.parametrize("checkpoint_during", [True, False]) +@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC) +async def test_subgraph_checkpoint_true_interrupt( + checkpointer_name: str, checkpoint_during: bool +) -> None: + # Define subgraph + class SubgraphState(TypedDict): + # note that none of these keys are shared with the parent graph state + bar: str + baz: str + + def subgraph_node_1(state: SubgraphState): + baz_value = interrupt("Provide baz value") + return {"baz": baz_value} + + def subgraph_node_2(state: SubgraphState): + return {"bar": state["bar"] + state["baz"]} + + subgraph_builder = StateGraph(SubgraphState) + subgraph_builder.add_node(subgraph_node_1) + subgraph_builder.add_node(subgraph_node_2) + subgraph_builder.add_edge(START, "subgraph_node_1") + subgraph_builder.add_edge("subgraph_node_1", "subgraph_node_2") + subgraph = subgraph_builder.compile(checkpointer=True) + + class ParentState(TypedDict): + foo: str + + def node_1(state: ParentState): + return {"foo": "hi! " + state["foo"]} + + async def node_2(state: ParentState, config: RunnableConfig): + response = await subgraph.ainvoke({"bar": state["foo"]}) + return {"foo": response["bar"]} + + builder = StateGraph(ParentState) + builder.add_node("node_1", node_1) + builder.add_node("node_2", node_2) + builder.add_edge(START, "node_1") + builder.add_edge("node_1", "node_2") + + async with awith_checkpointer(checkpointer_name) as checkpointer: + graph = builder.compile(checkpointer=checkpointer) + config = {"configurable": {"thread_id": "1"}} + + assert await graph.ainvoke( + {"foo": "foo"}, config, checkpoint_during=checkpoint_during + ) == {"foo": "hi! foo"} + assert (await graph.aget_state(config, subgraphs=True)).tasks[ + 0 + ].state.values == {"bar": "hi! foo"} + assert await graph.ainvoke( + Command(resume="baz"), config, checkpoint_during=checkpoint_during + ) == {"foo": "hi! foobaz"} + + @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_stream_subgraphs_during_execution(checkpointer_name: str) -> None: class InnerState(TypedDict): @@ -5458,8 +5669,11 @@ async def test_stream_buffering_single_node(checkpointer_name: str) -> None: ] +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_nested_graph_interrupts_parallel(checkpointer_name: str) -> None: +async def test_nested_graph_interrupts_parallel( + checkpointer_name: str, checkpoint_during: bool +) -> None: class InnerState(TypedDict): my_key: Annotated[str, operator.add] my_other_key: str @@ -5508,11 +5722,13 @@ async def test_nested_graph_interrupts_parallel(checkpointer_name: str) -> None: # test invoke w/ nested interrupt config = {"configurable": {"thread_id": "1"}} - assert await app.ainvoke({"my_key": ""}, config, debug=True) == { + assert await app.ainvoke( + {"my_key": ""}, config, checkpoint_during=checkpoint_during + ) == { "my_key": " and parallel", } - assert await app.ainvoke(None, config, debug=True) == { + assert await app.ainvoke(None, config, checkpoint_during=checkpoint_during) == { "my_key": "got here and there and parallel and back again", } @@ -5522,7 +5738,13 @@ async def test_nested_graph_interrupts_parallel(checkpointer_name: str) -> None: # test stream updates w/ nested interrupt config = {"configurable": {"thread_id": "2"}} assert [ - c async for c in app.astream({"my_key": ""}, config, subgraphs=True) + c + async for c in app.astream( + {"my_key": ""}, + config, + subgraphs=True, + checkpoint_during=checkpoint_during, + ) ] == [ # we got to parallel node first ((), {"outer_1": {"my_key": " and parallel"}}), @@ -5532,7 +5754,12 @@ async def test_nested_graph_interrupts_parallel(checkpointer_name: str) -> None: ), ((), {"__interrupt__": ()}), ] - assert [c async for c in app.astream(None, config)] == [ + assert [ + c + async for c in app.astream( + None, config, checkpoint_during=checkpoint_during + ) + ] == [ {"outer_1": {"my_key": " and parallel"}, "__metadata__": {"cached": True}}, {"inner": {"my_key": "got here and there"}}, {"outer_2": {"my_key": " and back again"}}, @@ -5541,12 +5768,23 @@ async def test_nested_graph_interrupts_parallel(checkpointer_name: str) -> None: # test stream values w/ nested interrupt config = {"configurable": {"thread_id": "3"}} assert [ - c async for c in app.astream({"my_key": ""}, config, stream_mode="values") + c + async for c in app.astream( + {"my_key": ""}, + config, + stream_mode="values", + checkpoint_during=checkpoint_during, + ) ] == [ {"my_key": ""}, {"my_key": " and parallel"}, ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + assert [ + c + async for c in app.astream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ {"my_key": ""}, {"my_key": "got here and there and parallel"}, {"my_key": "got here and there and parallel and back again"}, @@ -5556,16 +5794,32 @@ async def test_nested_graph_interrupts_parallel(checkpointer_name: str) -> None: app = graph.compile(checkpointer=checkpointer, interrupt_before=["outer_1"]) config = {"configurable": {"thread_id": "4"}} assert [ - c async for c in app.astream({"my_key": ""}, config, stream_mode="values") + c + async for c in app.astream( + {"my_key": ""}, + config, + stream_mode="values", + checkpoint_during=checkpoint_during, + ) ] == [ {"my_key": ""}, ] # while we're waiting for the node w/ interrupt inside to finish - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + assert [ + c + async for c in app.astream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ {"my_key": ""}, {"my_key": " and parallel"}, ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + assert [ + c + async for c in app.astream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ {"my_key": ""}, {"my_key": "got here and there and parallel"}, {"my_key": "got here and there and parallel and back again"}, @@ -5575,23 +5829,42 @@ async def test_nested_graph_interrupts_parallel(checkpointer_name: str) -> None: app = graph.compile(checkpointer=checkpointer, interrupt_after=["outer_1"]) config = {"configurable": {"thread_id": "5"}} assert [ - c async for c in app.astream({"my_key": ""}, config, stream_mode="values") + c + async for c in app.astream( + {"my_key": ""}, + config, + stream_mode="values", + checkpoint_during=checkpoint_during, + ) ] == [ {"my_key": ""}, {"my_key": " and parallel"}, ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + assert [ + c + async for c in app.astream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ {"my_key": ""}, {"my_key": "got here and there and parallel"}, ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + assert [ + c + async for c in app.astream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ {"my_key": "got here and there and parallel"}, {"my_key": "got here and there and parallel and back again"}, ] +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_doubly_nested_graph_interrupts(checkpointer_name: str) -> None: +async def test_doubly_nested_graph_interrupts( + checkpointer_name: str, checkpoint_during: bool +) -> None: class State(TypedDict): my_key: str @@ -5644,11 +5917,13 @@ async def test_doubly_nested_graph_interrupts(checkpointer_name: str) -> None: # test invoke w/ nested interrupt config = {"configurable": {"thread_id": "1"}} - assert await app.ainvoke({"my_key": "my value"}, config, debug=True) == { + assert await app.ainvoke( + {"my_key": "my value"}, config, checkpoint_during=checkpoint_during + ) == { "my_key": "hi my value", } - assert await app.ainvoke(None, config, debug=True) == { + assert await app.ainvoke(None, config, checkpoint_during=checkpoint_during) == { "my_key": "hi my value here and there and back again", } @@ -5657,12 +5932,22 @@ async def test_doubly_nested_graph_interrupts(checkpointer_name: str) -> None: config = { "configurable": {"thread_id": "2", CONFIG_KEY_NODE_FINISHED: nodes.append} } - assert [c async for c in app.astream({"my_key": "my value"}, config)] == [ + assert [ + c + async for c in app.astream( + {"my_key": "my value"}, config, checkpoint_during=checkpoint_during + ) + ] == [ {"parent_1": {"my_key": "hi my value"}}, {"__interrupt__": ()}, ] assert nodes == ["parent_1", "grandchild_1"] - assert [c async for c in app.astream(None, config)] == [ + assert [ + c + async for c in app.astream( + None, config, checkpoint_during=checkpoint_during + ) + ] == [ {"child": {"my_key": "hi my value here and there"}}, {"parent_2": {"my_key": "hi my value here and there and back again"}}, ] @@ -5680,13 +5965,21 @@ async def test_doubly_nested_graph_interrupts(checkpointer_name: str) -> None: assert [ c async for c in app.astream( - {"my_key": "my value"}, config, stream_mode="values" + {"my_key": "my value"}, + config, + stream_mode="values", + checkpoint_during=checkpoint_during, ) ] == [ {"my_key": "my value"}, {"my_key": "hi my value"}, ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + assert [ + c + async for c in app.astream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ {"my_key": "hi my value"}, {"my_key": "hi my value here and there"}, {"my_key": "hi my value here and there and back again"}, diff --git a/libs/sdk-js/.gitignore b/libs/sdk-js/.gitignore index 128c6dc98..db1a582d0 100644 --- a/libs/sdk-js/.gitignore +++ b/libs/sdk-js/.gitignore @@ -6,6 +6,10 @@ client.cjs client.js client.d.ts client.d.cts +auth.cjs +auth.js +auth.d.ts +auth.d.cts react.cjs react.js react.d.ts diff --git a/libs/sdk-js/langchain.config.js b/libs/sdk-js/langchain.config.js index 21d796c87..0c8a691cd 100644 --- a/libs/sdk-js/langchain.config.js +++ b/libs/sdk-js/langchain.config.js @@ -14,6 +14,7 @@ export const config = { entrypoints: { index: "index", client: "client", + auth: "auth/index", react: "react/index", "react-ui": "react-ui/index", "react-ui/server": "react-ui/server/index", diff --git a/libs/sdk-js/package.json b/libs/sdk-js/package.json index 64ddcf047..29510712b 100644 --- a/libs/sdk-js/package.json +++ b/libs/sdk-js/package.json @@ -1,6 +1,6 @@ { "name": "@langchain/langgraph-sdk", - "version": "0.0.63", + "version": "0.0.65", "description": "Client library for interacting with the LangGraph API", "type": "module", "packageManager": "yarn@1.22.19", @@ -72,6 +72,15 @@ "import": "./client.js", "require": "./client.cjs" }, + "./auth": { + "types": { + "import": "./auth.d.ts", + "require": "./auth.d.cts", + "default": "./auth.d.ts" + }, + "import": "./auth.js", + "require": "./auth.cjs" + }, "./react": { "types": { "import": "./react.d.ts", @@ -111,6 +120,10 @@ "client.js", "client.d.ts", "client.d.cts", + "auth.cjs", + "auth.js", + "auth.d.ts", + "auth.d.cts", "react.cjs", "react.js", "react.d.ts", diff --git a/libs/sdk-js/src/auth/error.ts b/libs/sdk-js/src/auth/error.ts new file mode 100644 index 000000000..ace6df54c --- /dev/null +++ b/libs/sdk-js/src/auth/error.ts @@ -0,0 +1,80 @@ +const HTTP_STATUS_MAPPING: { [key: number]: string } = { + 100: "Continue", + 101: "Switching Protocols", + 102: "Processing", + 103: "Early Hints", + 200: "OK", + 201: "Created", + 202: "Accepted", + 203: "Non-Authoritative Information", + 204: "No Content", + 205: "Reset Content", + 206: "Partial Content", + 207: "Multi-Status", + 208: "Already Reported", + 226: "IM Used", + 300: "Multiple Choices", + 301: "Moved Permanently", + 302: "Found", + 303: "See Other", + 304: "Not Modified", + 305: "Use Proxy", + 307: "Temporary Redirect", + 308: "Permanent Redirect", + 400: "Bad Request", + 401: "Unauthorized", + 402: "Payment Required", + 403: "Forbidden", + 404: "Not Found", + 405: "Method Not Allowed", + 406: "Not Acceptable", + 407: "Proxy Authentication Required", + 408: "Request Timeout", + 409: "Conflict", + 410: "Gone", + 411: "Length Required", + 412: "Precondition Failed", + 413: "Request Entity Too Large", + 414: "Request-URI Too Long", + 415: "Unsupported Media Type", + 416: "Requested Range Not Satisfiable", + 417: "Expectation Failed", + 418: "I'm a Teapot", + 421: "Misdirected Request", + 422: "Unprocessable Entity", + 423: "Locked", + 424: "Failed Dependency", + 425: "Too Early", + 426: "Upgrade Required", + 428: "Precondition Required", + 429: "Too Many Requests", + 431: "Request Header Fields Too Large", + 451: "Unavailable For Legal Reasons", + 500: "Internal Server Error", + 501: "Not Implemented", + 502: "Bad Gateway", + 503: "Service Unavailable", + 504: "Gateway Timeout", + 505: "HTTP Version Not Supported", + 506: "Variant Also Negotiates", + 507: "Insufficient Storage", + 508: "Loop Detected", + 510: "Not Extended", + 511: "Network Authentication Required", +}; + +export class HTTPException extends Error { + status: number; + headers: HeadersInit; + + constructor( + status: number, + options?: { message?: string; headers?: HeadersInit; cause?: unknown }, + ) { + super(options?.message ?? HTTP_STATUS_MAPPING[status] ?? "Unknown error", { + cause: options?.cause, + }); + this.status = status; + this.headers = options?.headers ?? {}; + } +} diff --git a/libs/sdk-js/src/auth/index.ts b/libs/sdk-js/src/auth/index.ts new file mode 100644 index 000000000..cf995438d --- /dev/null +++ b/libs/sdk-js/src/auth/index.ts @@ -0,0 +1,39 @@ +import type { + AuthenticateCallback, + AnyCallback, + CallbackEvent, + OnCallback, + BaseAuthReturn, + ToUserLike, + BaseUser, +} from "./types.js"; + +export class Auth< + TExtra = {}, + TAuthReturn extends BaseAuthReturn = BaseAuthReturn, + TUser extends BaseUser = ToUserLike, +> { + "~handlerCache": { + authenticate?: AuthenticateCallback; + callbacks?: Record; + } = {}; + + authenticate( + cb: AuthenticateCallback, + ): Auth { + this["~handlerCache"].authenticate = cb; + return this as unknown as Auth; + } + + on(event: T, callback: OnCallback): this { + this["~handlerCache"].callbacks ??= {}; + const events: string[] = Array.isArray(event) ? event : [event]; + for (const event of events) { + this["~handlerCache"].callbacks[event] = callback as AnyCallback; + } + return this; + } +} + +export type { Filters, ResourceActionType } from "./types.js"; +export { HTTPException } from "./error.js"; diff --git a/libs/sdk-js/src/auth/types.ts b/libs/sdk-js/src/auth/types.ts new file mode 100644 index 000000000..d2a39e05b --- /dev/null +++ b/libs/sdk-js/src/auth/types.ts @@ -0,0 +1,345 @@ +type Maybe = T | null | undefined; +type PromiseMaybe = Promise | T; + +interface AssistantConfig { + tags?: Maybe; + recursion_limit?: Maybe; + configurable?: Maybe<{ + thread_id?: Maybe; + thread_ts?: Maybe; + [key: string]: unknown; + }>; +} + +interface AssistantCreate { + assistant_id?: Maybe; + metadata?: Maybe>; + config?: Maybe; + if_exists?: Maybe<"raise" | "do_nothing">; + name?: Maybe; + graph_id: string; +} + +interface AssistantRead { + assistant_id: string; + metadata?: Maybe>; +} + +interface AssistantUpdate { + assistant_id: string; + metadata?: Maybe>; + config?: Maybe; + graph_id?: Maybe; + name?: Maybe; + version?: Maybe; +} + +interface AssistantDelete { + assistant_id: string; +} + +interface AssistantSearch { + graph_id?: Maybe; + metadata?: Maybe>; + limit?: Maybe; + offset?: Maybe; +} + +interface ThreadCreate { + thread_id?: Maybe; + metadata?: Maybe>; + if_exists?: Maybe<"raise" | "do_nothing">; +} + +interface ThreadRead { + thread_id?: Maybe; +} + +interface ThreadUpdate { + thread_id?: Maybe; + metadata?: Maybe>; + action?: Maybe<"interrupt" | "rollback">; +} + +interface ThreadDelete { + thread_id?: Maybe; + run_id?: Maybe; +} + +interface ThreadSearch { + thread_id?: Maybe; + status?: Maybe<"idle" | "busy" | "interrupted" | "error" | (string & {})>; + metadata?: Maybe>; + values?: Maybe>; + limit?: Maybe; + offset?: Maybe; +} + +interface CronCreate { + payload?: Maybe>; + schedule: string; + cron_id?: Maybe; + thread_id?: Maybe; + user_id?: Maybe; + end_time?: Maybe; +} + +interface CronRead { + cron_id: string; +} + +interface CronUpdate { + cron_id: string; + payload?: Maybe>; + schedule?: Maybe; +} + +interface CronDelete { + cron_id: string; +} + +interface CronSearch { + assistant_id?: Maybe; + thread_id?: Maybe; + limit?: Maybe; + offset?: Maybe; +} + +interface StorePut { + namespace: string[]; + key: string; + value: Record; +} + +interface StoreGet { + namespace: Maybe; + key: string; +} + +interface StoreSearch { + namespace?: Maybe; + filter?: Maybe>; + limit?: Maybe; + offset?: Maybe; + query?: Maybe; +} + +interface StoreListNamespaces { + namespace?: Maybe; + suffix?: Maybe; + max_depth?: Maybe; + limit?: Maybe; + offset?: Maybe; +} + +interface StoreDelete { + namespace?: Maybe; + key: string; +} + +interface RunsCreate { + thread_id?: Maybe; + assistant_id: string; + run_id: string; + status: Maybe< + "pending" | "running" | "error" | "success" | "timeout" | "interrupted" + >; + metadata?: Maybe>; + prevent_insert_if_inflight?: Maybe; + multitask_strategy?: Maybe<"interrupt" | "rollback" | "reject" | "enqueue">; + if_not_exists?: Maybe<"reject" | "create">; + after_seconds?: Maybe; + kwargs: Record; +} + +export interface ResourceActionType { + ["threads:create"]: ThreadCreate; + ["threads:read"]: ThreadRead; + ["threads:update"]: ThreadUpdate; + ["threads:delete"]: ThreadDelete; + ["threads:search"]: ThreadSearch; + ["threads:create_run"]: RunsCreate; + + ["assistants:create"]: AssistantCreate; + ["assistants:read"]: AssistantRead; + ["assistants:update"]: AssistantUpdate; + ["assistants:delete"]: AssistantDelete; + ["assistants:search"]: AssistantSearch; + + ["crons:create"]: CronCreate; + ["crons:read"]: CronRead; + ["crons:update"]: CronUpdate; + ["crons:delete"]: CronDelete; + ["crons:search"]: CronSearch; + + ["store:put"]: StorePut; + ["store:get"]: StoreGet; + ["store:search"]: StoreSearch; + ["store:list_namespaces"]: StoreListNamespaces; + ["store:delete"]: StoreDelete; +} +interface ResourceType { + threads: + | "threads:create" + | "threads:read" + | "threads:update" + | "threads:delete" + | "threads:search" + | "threads:create_run"; + + assistants: + | "assistants:create" + | "assistants:read" + | "assistants:update" + | "assistants:delete" + | "assistants:search"; + crons: + | "crons:create" + | "crons:read" + | "crons:update" + | "crons:delete" + | "crons:search"; + + store: + | "store:put" + | "store:get" + | "store:search" + | "store:list_namespaces" + | "store:delete"; +} +interface ActionType { + "*:create": "threads:create" | "assistants:create" | "crons:create"; + + "*:read": "threads:read" | "assistants:read" | "crons:read"; + + "*:update": "threads:update" | "assistants:update" | "crons:update"; + + "*:delete": + | "threads:delete" + | "assistants:delete" + | "crons:delete" + | "store:delete"; + + "*:search": + | "threads:search" + | "assistants:search" + | "crons:search" + | "store:search"; + + "*:create_run": "threads:create_run"; + + "*:put": "store:put"; + + "*:get": "store:get"; + + "*:list_namespaces": "store:list_namespaces"; +} + +export type BaseAuthReturn = + | { + is_authenticated?: boolean; + display_name?: string; + identity: string; + permissions: string[]; + } + | string; + +export interface BaseUser { + is_authenticated: boolean; + display_name: string; + identity: string; + permissions: string[]; +} + +export type ToUserLike = T extends string + ? { + is_authenticated: boolean; + display_name: string; + identity: string; + permissions: string[]; + } + : Omit & { + is_authenticated: boolean; + display_name: string; + }; + +type CallbackParameter< + Resource extends string = string, + Action extends string = string, + Value extends unknown = unknown, + TUser extends BaseUser = BaseUser, +> = { + resource: Resource; + action: Action; + value: Value; + user: TUser; + permissions: string[]; +}; + +type ContextMap = { + [ActionType in keyof ResourceActionType]: CallbackParameter< + ActionType extends `${infer Resource}:${string}` ? Resource : never, + ActionType, + ResourceActionType[ActionType], + BaseUser + >; +}; + +type ActionCallbackParameter< + T extends keyof ActionType, + TUser extends BaseUser = BaseUser, +> = ContextMap[ActionType[T]] & { user: TUser }; +type AuthCallbackParameter< + T extends keyof ResourceActionType, + TUser extends BaseUser = BaseUser, +> = ContextMap[T] & { user: TUser }; +type ResourceCallbackParameter< + T extends keyof ResourceType, + TUser extends BaseUser = BaseUser, +> = ContextMap[ResourceType[T]] & { user: TUser }; + +export type Filters = { + [key in TKey]: string | { [op in "$contains" | "$eq"]?: string }; +}; + +export interface AuthenticateCallback { + (request: Request): PromiseMaybe; +} + +type OnKey = keyof ResourceType | keyof ActionType | keyof ResourceActionType; + +type OnSingleParameter< + T extends OnKey, + TUser extends BaseUser = BaseUser, +> = T extends keyof ResourceType + ? ResourceCallbackParameter + : T extends keyof ActionType + ? ActionCallbackParameter + : T extends keyof ResourceActionType + ? AuthCallbackParameter + : never; + +type OnParameter< + T extends "*" | OnKey | OnKey[], + TUser extends BaseUser = BaseUser, +> = T extends OnKey[] + ? OnSingleParameter + : T extends "*" + ? AuthCallbackParameter + : T extends OnKey + ? OnSingleParameter + : never; + +export type AnyCallback = ( + request: CallbackParameter, +) => void | boolean | Filters; + +export type CallbackEvent = "*" | OnKey | OnKey[]; + +export type OnCallback< + T extends CallbackEvent, + TUser extends BaseUser = BaseUser, + TMetadata extends Record = Record, +> = ( + request: OnParameter, +) => void | boolean | Filters; diff --git a/libs/sdk-js/tsconfig.json b/libs/sdk-js/tsconfig.json index 9c6561a09..d5c1626ab 100644 --- a/libs/sdk-js/tsconfig.json +++ b/libs/sdk-js/tsconfig.json @@ -2,11 +2,7 @@ "extends": "@tsconfig/recommended", "compilerOptions": { "target": "ES2021", - "lib": [ - "ES2021", - "ES2022.Object", - "DOM" - ], + "lib": ["ES2021", "ES2022.Object", "ES2022.Error", "DOM"], "module": "NodeNext", "moduleResolution": "nodenext", "esModuleInterop": true, @@ -22,24 +18,14 @@ "jsx": "react-jsx", "outDir": "dist" }, - "include": [ - "src/**/*" - ], - "exclude": [ - "node_modules", - "dist", - "coverage" - ], + "include": ["src/**/*"], + "exclude": ["node_modules", "dist", "coverage"], "includeVersion": true, "typedocOptions": { - "entryPoints": [ - "src/client.ts" - ], + "entryPoints": ["src/client.ts"], "readme": "none", "out": "docs", - "plugin": [ - "typedoc-plugin-markdown" - ], + "plugin": ["typedoc-plugin-markdown"], "excludePrivate": true, "excludeProtected": true, "excludeExternals": false