cli0.1.50

cli: Fix crash when subprocess has a very long stdout/stderr line

.github/ISSUE_TEMPLATE/bug-report.yml

@@ -1,5 +1,5 @@
 name: "\U0001F41B Bug Report"
-description: Report a bug in LangChain. To report a security issue, please instead use the security option below. For questions, please use the GitHub Discussions.
+description: Report a bug in LangGraph. To report a security issue, please instead use the security option below. For questions, please use the GitHub Discussions.
 labels: ["02 Bug Report"]
 body:
   - type: markdown
@@ -15,11 +15,11 @@ body:
         Relevant links to check before filing a bug report to see if your issue has already been reported, fixed or
         if there's another way to solve your problem:
         
+        [LangGraph documentation](https://langchain-ai.github.io/langgraph/).
         [LangChain documentation with the integrated search](https://python.langchain.com/docs/get_started/introduction),
-        [API Reference](https://api.python.langchain.com/en/stable/),
-        [GitHub search](https://github.com/langchain-ai/langchain),
-        [LangChain Github Discussions](https://github.com/langchain-ai/langchain/discussions),
-        [LangChain Github Issues](https://github.com/langchain-ai/langchain/issues?q=is%3Aissue),
+        [GitHub search](https://github.com/langchain-ai/langgraph),
+        [LangChain Github Discussions](https://github.com/langchain-ai/langgraph/discussions),
+        [LangChain Github Issues](https://github.com/langchain-ai/langgraph/issues),
         [LangChain ChatBot](https://chat.langchain.com/)
   - type: checkboxes
     id: checks
@@ -54,13 +54,13 @@ body:
         * Avoid screenshots when possible, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.
 
       placeholder: |
-        from langchain_core.runnables import RunnableLambda
+        from langgraph.graph import StateGraph
 
         def bad_code(inputs) -> int:
           raise NotImplementedError('For demo purpose')
           
-          chain = RunnableLambda(bad_code)
-          chain.invoke('Hello!')
+        chain = StateGraph(list)
+        chain.invoke('Hello!')
       render: python
   - type: textarea
     id: error
@@ -82,7 +82,7 @@ body:
 
         Write a short description telling what you are doing, what you expect to happen, and what is currently happening.
       placeholder: |
-        * I'm trying to use the `langchain` library to do X.
+        * I'm trying to use the `langgraph` library to do X.
         * I expect to see Y.
         * Instead, it does Z.
     validations:
@@ -102,7 +102,7 @@ body:
         
         python -m langchain_core.sys_info
       placeholder: |
-        "pip freeze | grep langchain"
+        "pip freeze | grep langgraph"
         platform
         python version
         

.github/ISSUE_TEMPLATE/config.yml

@@ -3,13 +3,13 @@ version: 2.1
 contact_links:
   - name: 🤔 Question or Problem
     about: Ask a question or ask about a problem in GitHub Discussions.
-    url: https://www.github.com/langchain-ai/langchain/discussions/categories/q-a
-  - name: Discord
-    url: https://discord.gg/6adMQxSpJS
-    about: General community discussions
+    url: https://github.com/langchain-ai/langgraph/discussions/categories/q-a
   - name: Feature Request
-    url: https://www.github.com/langchain-ai/langchain/discussions/categories/ideas
+    url: https://github.com/langchain-ai/langgraph/discussions/categories/ideas
     about: Suggest a feature or an idea
   - name: Show and tell
     about: Show what you built with LangChain
-    url: https://www.github.com/langchain-ai/langchain/discussions/categories/show-and-tell
+    url: https://github.com/langchain-ai/langgraph/discussions/categories/show-and-tell
+  - name: Slack
+    url:  https://www.langchain.com/join-community
+    about: General community discussions

.github/ISSUE_TEMPLATE/documentation.yml

@@ -1,5 +1,5 @@
 name: Documentation
-description: Report an issue related to the LangChain documentation.
+description: Report an issue related to the LangGraph documentation.
 title: "DOC: <Please write a comprehensive title after the 'DOC: ' prefix>"
 labels: [03 - Documentation]
 

docs/.gitignore

@@ -1,2 +1,3 @@
 *.ipynb
 site/
+docs/tutorials/**/*.png

docs/_scripts/copy_notebooks.py

@@ -21,17 +21,26 @@ _MANUAL = {
         "stream-updates.ipynb",
         "stream-multiple.ipynb",
         "streaming-tokens.ipynb",
+        "streaming-tokens-without-langchain.ipynb",
         "streaming-content.ipynb",
         "streaming-events-from-within-tools.ipynb",
+        "streaming-events-from-within-tools-without-langchain.ipynb",
         "streaming-from-final-node.ipynb",
         "persistence.ipynb",
-        "managing-conversation-history.ipynb",
+        "memory/manage-conversation-history.ipynb",
+        "memory/delete-messages.ipynb",
+        "memory/add-summary-conversation-history.ipynb",
         "persistence_postgres.ipynb",
+        "persistence_mongodb.ipynb",
+        "persistence_redis.ipynb",
         "visualization.ipynb",
         "state-model.ipynb",
         "subgraph.ipynb",
         "force-calling-a-tool-first.ipynb",
         "pass-run-time-values-to-tools.ipynb",
+        "tool-calling.ipynb",
+        "tool-calling-errors.ipynb",
+        "pass-config-to-tools.ipynb",
         "dynamic-returning-direct.ipynb",
         "managing-agent-steps.ipynb",
         "respond-in-format.ipynb",
@@ -47,6 +56,7 @@ _MANUAL = {
         "human_in_the_loop/time-travel.ipynb",
         "human_in_the_loop/edit-graph-state.ipynb",
         "human_in_the_loop/wait-user-input.ipynb",
+        "node-retries.ipynb",
     ],
     "tutorials": [
         "introduction.ipynb",
@@ -82,7 +92,6 @@ _HIDE = set(
         "dynamically-returning-directly.ipynb",
         "force-calling-a-tool-first.ipynb",
         "managing-agent-steps.ipynb",
-        "pass-run-time-values-to-tools.ipynb",
         "respond-in-format.ipynb",
         "quickstart.ipynb",
         "human-in-the-loop.ipynb",

docs/docs/cloud/concepts/api.md

@@ -63,7 +63,7 @@ There are many occasions where the graph cannot run completely autonomously. For
 
 ### Double Texting
 
-Many times users might interact with your graph in unintended ways. For instance, a user may send one message and before the graph has finished running send a second message. To solve this issue of "double-texting" (i.e. prompting the graph a second time before the first run has finished), Langgraph has provided four different solutions, all of which are covered in the [Double Texting how-tos](../how-tos/index.md#double-texting). These options are:
+Many times users might interact with your graph in unintended ways. For instance, a user may send one message and before the graph has finished running send a second message. To solve this issue of "double-texting" (i.e. prompting the graph a second time before the first run has finished), LangGraph has provided four different solutions, all of which are covered in the [Double Texting how-tos](../how-tos/index.md#double-texting). These options are:
 
 - `reject`: This is the simplest option, this just rejects any follow up runs and does not allow double texting. See the [how-to guide](../how-tos/reject_concurrent.md) for configuring the reject double text option.
 - `enqueue`: This is a relatively simple option which continues the first run until it completes the whole run, then sends the new input as a separate run. See the [how-to guide](../how-tos/enqueue_concurrent.md) for configuring the enqueue double text option.
@@ -101,3 +101,7 @@ The LangGraph Cloud offers several features to support secure and robost deploym
 ### Authentication
 
 LangGraph applications deployed to LangGraph Cloud are automatically configured with LangSmith authentication. In order to call the API, a valid <a href="https://docs.smith.langchain.com/how_to_guides/setup/create_account_api_key#api-keys" target="_blank">LangSmith API key</a> is required.
+
+### Local Testing
+
+Before deploying your app in production to LangGraph Cloud, you may wish to test out your graph locally in order to ensure that everything is running as expected. Luckily, LangGraph makes this easy for you through use of the LangGraph CLI. Read more in this [how-to guide](../deployment/test_locally.md) or look at the [CLI reference](../reference/cli.md) to learn more.

docs/docs/cloud/concepts/cloud.md

@@ -0,0 +1,28 @@
+# Cloud Concepts
+
+This page describes the high-level concepts of the LangGraph Cloud deployment.
+
+## Deployment
+
+A deployment is an instance of a LangGraph API. A single deployment can have many [revisions](#revision). When a deployment is created, all of the necessary infrastructure (e.g. database, containers, secrets store) are automatically provisioned. See the [architecture diagram](#architecture) below for more details.
+
+See the [how-to guide](../deployment/cloud.md#create-new-deployment) for creating a new deployment.
+
+## Revision
+
+A revision is an iteration of a [deployment](#deployment). When a new deployment is created, an initial revision is automatically created. To deploy new code changes or update environment variable configurations for a deployment, a new revision must be created. When a revision is created, a new container image is built automatically.
+
+See the [how-to guide](../deployment/cloud.md#create-new-revision) for creating a new revision.
+
+## Asynchronous Deployment
+
+Infrastructure for [deployments](#deployment) and [revisions](#revision) are provisioned and deployed asynchronously. They are not deployed immediately after submission. Currently, deployment can take up to several minutes.
+
+## Architecture
+
+!!! warning "Subject to Change"
+    The LangGraph Cloud deployment architecture may change in the future.
+
+A high-level diagram of a LangGraph Cloud deployment.
+
+![diagram](langgraph_cloud_architecture.png)

docs/docs/cloud/deployment/cloud.md

@@ -2,9 +2,10 @@
 
 LangGraph Cloud is available within <a href="https://www.langchain.com/langsmith" target="_blank">LangSmith</a>. To deploy a LangGraph Cloud API, navigate to the <a href="https://smith.langchain.com/" target="_blank">LangSmith UI</a>.
 
-## Setup GitHub Repository
+## Prerequisites
 
-LangGraph Cloud applications are deployed from GitHub repositories. Configure and upload a LangGraph Cloud application to a GitHub repository in order to deploy it to LangGraph Cloud.
+1. LangGraph Cloud applications are deployed from GitHub repositories. Configure and upload a LangGraph Cloud application to a GitHub repository in order to deploy it to LangGraph Cloud.
+1. [Verify that the LangGraph API runs locally](test_locally.md). If the API does not build and run successfully (i.e. `langgraph up`), deploying to LangGraph Cloud will fail as well.
 
 ## Create New Deployment
 
@@ -45,8 +46,33 @@ Starting from the <a href="https://smith.langchain.com/" target="_blank">LangSmi
         1. Update the value of existing secrets or environment variables.
 1. Select `Submit`. After a few seconds, the `New Revision` modal will close and the new revision will be queued for deployment.
 
-## Asynchronous Deployment
+## View Build and Deployment Logs
 
-New [deployments](#create-new-deployment) and [revisions](#create-new-revision) are provisioned and deployed asynchronously. They are not deployed immediately after submission. Currently, deployment can take up to several minutes.
+Build and deployment logs are available for each revision.
 
-The `Deployment` view continually updates the status of pending revisions.
+Starting from the `Deployment` view...
+
+1. Select the desired revision from the `Revisions` table. A panel slides open from the right-hand side and the `Build` tab is selected by default, which displays build logs for the revision.
+1. In the panel, select the `Deploy` tab to view deployment logs for the revision.
+1. Within the `Deploy` tab, adjust the date/time range picker as needed. By default, the date/time range picker is set to the `Last 15 minutes`.
+
+## Interrupt Revision
+
+Interrupting a revision will stop deployment of the revision.
+
+!!! warning "Undefined Behavior"
+    Interrupted revisions have undefined behavior. This is only useful if you need to deploy a new revision and you already have a revision "stuck" in progress. In the future, this feature may be removed.
+
+Starting from the `Deployment` view...
+
+1. Select the menu icon (three dots) on the right-hand side of the row for the desired revision from the `Revisions` table.
+1. Select `Interrupt` from the menu.
+1. A modal will appear. Review the confirmation message. Select `Interrupt revision`.
+
+## Delete Deployment
+
+Starting from the <a href="https://smith.langchain.com/" target="_blank">LangSmith UI</a>...
+
+1. In the left-hand navigation panel, select `Deployments`. The `Deployments` view contains a list of existing LangGraph Cloud deployments.
+1. Select the menu icon (three dots) on the right-hand side of the row for the desired deployment and select `Delete`.
+1. A `Confirmation` modal will appear. Select `Delete`.

docs/docs/cloud/deployment/setup.md

@@ -101,6 +101,9 @@ Example `langgraph.json` file:
 
 Note that the variable name of the `CompiledGraph` appears at the end of the value of each subkey in the top-level `graphs` key (i.e. `:<variable_name>`).
 
+!!! warning "Configuration Location"
+    The LangGraph API configuration file must be placed in a directory that is at the same level or higher than the Python files that contain compiled graphs and associated dependencies.
+
 Example file directory:
 
 ```bash

docs/docs/cloud/deployment/setup_pyproject.md

@@ -102,7 +102,7 @@ agent = graph_workflow.compile()
 ```
 
 !!! warning "Assign `CompiledGraph` to Variable"
-The build process for LangGraph Cloud requires that the `CompiledGraph` object be assigned to a variable at the top-level of a Python module.
+    The build process for LangGraph Cloud requires that the `CompiledGraph` object be assigned to a variable at the top-level of a Python module.
 
 Example file directory:
 
@@ -133,6 +133,9 @@ Example `langgraph.json` file:
 
 Note that the variable name of the `CompiledGraph` appears at the end of the value of each subkey in the top-level `graphs` key (i.e. `:<variable_name>`).
 
+!!! warning "Configuration Location"
+    The LangGraph API configuration file must be placed in a directory that is at the same level or higher than the Python files that contain compiled graphs and associated dependencies.
+
 Example file directory:
 
 ```bash

docs/docs/cloud/deployment/test_locally.md

@@ -0,0 +1,99 @@
+# How to test a LangGraph app locally
+
+This guide assumes you have a LangGraph app correctly set up with a proper configuration file and a corresponding compiled graph, and that you have a proper LangChain API key.
+
+Testing locally ensures that there are no errors or conflicts with Python dependencies and confirms that the configuration file is specified correctly.
+
+## Setup
+
+Install the proper packages:
+
+```shell
+pip install langgraph-cli
+```
+
+Ensure you have an API key, which you can create from the LangSmith UI (Settings > API Keys). This is required to authenticate that you have LangGraph Cloud access. After you have saved the key to a safe place, place the following line in your `.env` file:
+
+```python
+LANGCHAIN_API_KEY = *********
+```
+
+## Start the API server
+
+Once you have downloaded the CLI, you can run the following command to start the API server for local testing:
+
+```shell
+langgraph up
+```
+
+This will start up the LangGraph API server locally. If this runs successfully, you should see something like:
+
+```shell
+Ready!
+- API: http://localhost:8123
+2024-06-26 19:20:41,056:INFO:uvicorn.access 127.0.0.1:44138 - "GET /ok HTTP/1.1" 200
+```
+
+### Interact with the server
+
+We can now interact with the API server using the LangGraph SDK. First, we need to start our client, select our assistant (in this case a graph we called "agent", make sure to select the proper assistant you wish to test).
+
+=== "Python"
+
+    ```python
+    from langgraph_sdk import get_client
+
+    # only pass the url argument to get_client() if you changed the default port when calling langgraph up
+    client = get_client()
+    assistant_id = "agent"
+    thread = await client.threads.create()
+    ```
+
+=== "Javascript"
+
+    ```js
+    import { Client } from "@langchain/langgraph-sdk";
+
+    // only set the apiUrl if you changed the default port when calling langgraph up
+    const client = new Client();
+    const assistantId = "agent"
+    const thread = await client.threads.create();
+    ```
+
+Now we can invoke our graph to ensure it is working. Make sure to change the input to match the proper schema for your graph. 
+
+=== "Python"
+
+    ```python
+    input = {"messages": [{"role": "human", "content": "what's the weather in sf"}]}
+    async for chunk in client.runs.stream(
+        thread["thread_id"],
+        assistant_id,
+        input=input,
+        stream_mode="updates",
+    ):
+        print(f"Receiving new event of type: {chunk.event}...")
+        print(chunk.data)
+        print("\n\n")
+    ```
+=== "Javascript"
+
+    ```js
+    const input = { "messages": [{ "role": "human", "content": "what's the weather in sf"}] }
+
+    const streamResponse = client.runs.stream(
+      thread["thread_id"],
+      assistantId,
+      {
+        input: input,
+        streamMode: "updates",
+      }
+    );
+    for await (const chunk of streamResponse) {
+      console.log(`Receiving new event of type: ${chunk.event}...`);
+      console.log(chunk.data);
+      console.log("\n\n");
+    }
+    ```
+
+If your graph works correctly, you should see your graph output displayed in the console. Of course, there are many more ways you might need to test your graph, for a full list of commands you can send with the SDK, see the [Python](https://langchain-ai.github.io/langgraph/cloud/reference/sdk/python_sdk_ref/) and [JS/TS](https://langchain-ai.github.io/langgraph/cloud/reference/sdk/js_ts_sdk_ref/) references.

docs/docs/cloud/how-tos/enqueue_concurrent.md

@@ -1,6 +1,6 @@
 ## Enqueue
 
-This guide assumes knowledge of what double-texting is, which you can learn about in the [double-texting conceptual guide](../concepts/index.md#double-texting).
+This guide assumes knowledge of what double-texting is, which you can learn about in the [double-texting conceptual guide](../concepts/api.md#double-texting).
 
 The guide covers the `enqueue` option for double texting, which adds the interruptions to a queue and executes them in the order they are received by the client. Below is a quick example of using the `enqueue` option.
 

docs/docs/cloud/how-tos/human_in_the_loop_breakpoint.md

@@ -43,7 +43,7 @@ We can do this by adding `interrupt_before=["action"]`, which tells us to interr
 We can do this either when compiling the graph or when kicking off a run.
 Here we will do it when kicking of a run, if you would like to to do it at compile time you need to edit the python file where your graph is defined and add the `interrupt_before` parameter when you call `.compile`.
 
-First let's access our hosted Langgraph instance through the SDK:
+First let's access our hosted LangGraph instance through the SDK:
 
 And, now let's compile it with a breakpoint before the tool node:
 

docs/docs/cloud/how-tos/human_in_the_loop_time_travel.md

@@ -1,6 +1,6 @@
 # How to Replay and Branch from Prior States
 
-With Langgraph Cloud you have the ability to return to any of your prior states and either re-run the graph to reproduce issues noticed during testing, or branch out in a different way from what was originally done in the prior states. In this guide we will show a quick example of how to rerun past states and how to branch off from previous states as well.
+With LangGraph Cloud you have the ability to return to any of your prior states and either re-run the graph to reproduce issues noticed during testing, or branch out in a different way from what was originally done in the prior states. In this guide we will show a quick example of how to rerun past states and how to branch off from previous states as well.
 
 ## Setup
 

docs/docs/cloud/how-tos/index.md

@@ -11,10 +11,13 @@ Welcome to the LangGraph Cloud how-to guides! These guides provide practical, st
 
 LangGraph Cloud gives you best in class observability, testing, and hosting services. Read more about them in these how to guides:
 
-- [How to set up app for deployment](../deployment/setup.md)
+- [How to set up app for deployment (requirements.txt)](../deployment/setup.md)
+- [How to set up app for deployment (pyproject.toml)](../deployment/setup_pyproject.md)
+- [How to test locally](../deployment/test_locally.md)
 - [How to deploy to LangGraph cloud](../deployment/cloud.md)
 - [How to self-host](../deployment/self_hosted.md)
 
+
 ## Streaming
 
 Streaming the results of your LLM application is vital for ensuring a good user experience, especially when your graph may call multiple models and take a long time to fully complete a run. Read about how to stream values from your graph in these how to guides:

docs/docs/cloud/how-tos/interrupt_concurrent.md

@@ -1,6 +1,6 @@
 ## Interrupt
 
-This guide assumes knowledge of what double-texting is, which you can learn about in the [double-texting conceptual guide](../concepts/index.md#double-texting).
+This guide assumes knowledge of what double-texting is, which you can learn about in the [double-texting conceptual guide](../concepts/api.md#double-texting).
 
 The guide covers the `interrupt` option for double texting, which interrupts the prior run of the graph and starts a new one with the double-text. This option does not delete the first run, but rather keeps it in the database but sets its status to `interrupted`. Below is a quick example of using the `interrupt` option.
 

docs/docs/cloud/index.md

@@ -1,5 +1,11 @@
 # LangGraph Cloud (beta)
 
+!!! tip
+    - LangGraph is an MIT-licensed open-source library, which we are committed to maintaining and growing for the community.
+    - LangGraph Cloud is an optional managed hosting service for LangGraph, which provides additional features geared towards production deployments.
+    - We are actively contributing improvements back to LangGraph informed by our work on LangGraph Cloud.
+    - You can always deploy LangGraph applications on your own infrastructure using the open-source LangGraph project.
+
 !!! danger "Important"
     LangGraph Cloud is a closed source, paid product in an invite-only stage. We are currently focused on providing high bandwidth support to make our select early customers successful. If you are interested in applying for access, please fill out [this form](https://www.langchain.com/langgraph-cloud-beta).
 
@@ -12,7 +18,7 @@
 
 LangGraph Cloud is a managed service for deploying and hosting LangGraph applications. Deploying applications with LangGraph Cloud shortens the time-to-market for developers. With one click, deploy a production-ready API with built-in persistence for your LangGraph application. LangGraph Cloud APIs are horizontally scalable and deployed with durable storage.
 
-The LangGraph Cloud API exposes functionality of your LangGraph application through [Assistants](./concepts/index.md#assistants). An assistant abstracts the cognitive architecture of your graph. Invoke an assistant by calling the pre-built [API endpoints](./reference/api/api_ref.md).
+The LangGraph Cloud API exposes functionality of your LangGraph application through [Assistants](./concepts/api.md#assistants). An assistant abstracts the cognitive architecture of your graph. Invoke an assistant by calling the pre-built [API endpoints](./reference/api/api_ref.md).
 
 LangGraph Cloud is seamlessly integrated with [LangSmith](https://www.langchain.com/langsmith) and is accessible from within the LangSmith UI.
 
@@ -31,5 +37,5 @@ The LangGraph Cloud API supports key LangGraph features in addition to new funct
 
 - [Tutorials](./quick_start.md): Learn to build and deploy applications for LangGraph Cloud.
 - [How-to Guides](./how-tos/index.md): Learn how to set up a LangGraph application for deployment and implement features of the LangGraph Cloud API such as streaming tokens, configuring double texting, and creating cron jobs. Go here if you want to copy and run a specific code snippet.
-- [Conceptual Guides](./concepts/index.md): In-depth explanations of the core data models (e.g. assistants) and key features (e.g. double texting) of the LangGraph Cloud API.
+- [Conceptual Guides](./concepts/api.md): In-depth explanations of the core data models (e.g. assistants), key features of the LangGraph Cloud API (e.g. double texting), and the architecture of a LangGraph Cloud deployment.
 - [Reference](./reference/api/api_ref.md): References for the LangGraph Cloud API, the corresponding Python and JS/TS SDKs, the LangGraph CLI, and deployment environment variables.

docs/docs/cloud/quick_start.md

@@ -74,13 +74,13 @@ In order to do this we can first install the LangGraph CLI
 pip install langgraph-cli
 ```
 
-We can then stand up a simple test server. The server this stands up is INCREDIBLY simple - it is just a single endpoint and has no persistence. **This should not be used for hosting your application, only for testing the build and basic functionality.**
+We can then test our API server locally. This requires access to LangGraph closed beta. In order to run the server locally, you will need to add your `LANGSMITH_API_KEY` to the .env file so we can validate you have access to LangGraph closed beta.
 
 ```shell
-langgraph test
+langgraph up
 ```
 
-This will test building of the agent server. If this runs successfully, you should see something like:
+This will start up the LangGraph API server locally. If this runs successfully, you should see something like:
 
 ```shell
 Ready!
@@ -88,7 +88,7 @@ Ready!
 2024-06-26 19:20:41,056:INFO:uvicorn.access 127.0.0.1:44138 - "GET /ok HTTP/1.1" 200
 ```
 
-You can now test this out! Again, we only expose a single simple endpoint (for streaming stateless runs). This is intended to allow you to test that the agent is properly set up, but should **NOT** but used for production purposes. To test it out, you can go to another terminal window and run:
+You can now test this out! **Note: this local server is intended SOLELY for local testing purposes and is not performant enough for production applications, so please do not use it as such.** To test it out, you can go to another terminal window and run:
 
 ```shell
 curl --request POST \
@@ -186,14 +186,12 @@ Once you have tested that your hosted graph works as expected using LangGraph St
 
 First, make sure you have the SDK installed by calling `pip install langgraph_sdk`.
 
-Before using, you need to get the URL of your LangGraph deployment. You can find this on the auto generated documentation page here:
-
-![Base URL of LangGraph deployment](./deployment/img/base_url.png)
+Before using, you need to get the URL of your LangGraph deployment. You can find this in the `Deployment` view. Click the URL to copy it to the clipboard.
 
 You also need to make sure you have set up your API key properly so you can authenticate with LangGraph Cloud.
 
 ```shell
-export LANGCHAIN_API_KEY=...
+export LANGSMITH_API_KEY=...
 ```
 
 The first thing to do when using the SDK is to setup our client, access our assistant, and create a thread to execute a run on:

docs/docs/cloud/reference/cli.md

@@ -83,6 +83,31 @@ langgraph build [OPTIONS]
 | `-c, --config FILE` | `langgraph.json` | Path to configuration file declaring dependencies, graphs and environment variables. |
 | `--help` | | Display command documentation. |
 
+### `up`
+Start langgraph API server. For local testing, requires a LangSmith API key with access to LangGraph Cloud closed beta. Requires a license key for production use.
+
+**Usage**
+```
+langgraph up [OPTIONS]
+```
+
+**Options**
+
+| Option | Default | Description |
+| ------ | ------- | ----------- |
+| `--wait` | | Wait for services to start before returning. Implies --detach |
+| `--postgres-uri TEXT` | Local database | Postgres URI to use for the database. |
+| `--watch` | | Restart on file changes |
+| `--debugger-base-url TEXT` | `http://127.0.0.1:[PORT]` | URL used by the debugger to access LangGraph API. |
+| `--debugger-port INTEGER` | | Pull the debugger image locally and serve the UI on specified port |
+| `--verbose` | | Show more output from the server logs. |
+| `-c, --config FILE` | `langgraph.json` | Path to configuration file declaring dependencies, graphs and environment variables. |
+| `-d, --docker-compose FILE` | | Path to docker-compose.yml file with additional services to launch. |
+| `-p, --port INTEGER` | `8123` | Port to expose. Example: `langgraph test --port 8000` |
+| `--pull / --no-pull` | `pull` | Pull latest images. Use --no-pull for running the server with locally-built images. Example: `langgraph up --no-pull` |
+| `--recreate / --no-recreate` | `no-recreate` | Recreate containers even if their configuration and image haven't changed |
+| `--help` | | Display command documentation. | 
+
 ### `test`
 Test your LangGraph in the cloud. The only function you can call from the SDK after testing your graph is `client.runs.stream(thread_id=None, ...)`
 

docs/docs/concepts/low_level.md

@@ -14,7 +14,9 @@ By composing `Nodes` and `Edges`, you can create complex, looping workflows that
 
 In short: _nodes do the work. edges tell what to do next_.
 
-LangGraph's underlying graph algorithm uses [message passing](https://en.wikipedia.org/wiki/Message_passing) to define a general program. When a `Node` completes, it sends a message along one or more edges to other node(s). These nodes run their functions, pass the resulting messages to the next set of nodes, and on and on it goes. Inspired by [Pregel](https://research.google/pubs/pregel-a-system-for-large-scale-graph-processing/), the program proceeds in discrete "super-steps" that are all executed conceptually in parallel. Whenever the graph is run, all the nodes start in an `inactive` state. Whenever an incoming edge (or "channel") receives a new message (state), the node becomes `active`, runs the function, and responds with updates. At the end of each superstep, each node votes to `halt` by marking itself as `inactive` if it has no more incoming messages. The graph terminates when all nodes are `inactive` and when no messages are in transit.
+LangGraph's underlying graph algorithm uses [message passing](https://en.wikipedia.org/wiki/Message_passing) to define a general program. When a Node completes its operation, it sends messages along one or more edges to other node(s). These recipient nodes then execute their functions, pass the resulting messages to the next set of nodes, and the process continues. Inspired by Google's [Pregel](https://research.google/pubs/pregel-a-system-for-large-scale-graph-processing/) system, the program proceeds in discrete "super-steps."
+
+A super-step can be considered a single iteration over the graph nodes. Nodes that run in parallel are part of the same super-step, while nodes that run sequentially belong to separate super-steps. At the start of graph execution, all nodes begin in an `inactive` state. A node becomes `active` when it receives a new message (state) on any of its incoming edges (or "channels"). The active node then runs its function and responds with updates. At the end of each super-step, nodes with no incoming messages vote to `halt` by marking themselves as `inactive`. The graph execution terminates when all nodes are `inactive` and no messages are in transit.
 
 ### StateGraph
 
@@ -178,7 +180,7 @@ graph.add_edge("node_a", "node_b")
 If you want to **optionally** route to 1 or more edges (or optionally terminate), you can use the [add_conditional_edges][langgraph.graph.StateGraph.add_conditional_edges] method. This method accepts the name of a node and a "routing function" to call after that node is executed:
 
 ```python
-graph.add_edge("node_a", routing_function)
+graph.add_conditional_edges("node_a", routing_function)
 ```
 
 Similar to nodes, the `routing_function` accept the current `state` of the graph and return a value.
@@ -188,7 +190,7 @@ By default, the return value `routing_function` is used as the name of the node
 You can optionally provide a dictionary that maps the `routing_function`'s output to the name of the next node.
 
 ```python
-graph.add_edge("node_a", routing_function, {True: "node_b", False: "node_c"})
+graph.add_conditional_edges("node_a", routing_function, {True: "node_b", False: "node_c"})
 ```
 
 ### Entry Point
@@ -232,11 +234,9 @@ graph.add_conditional_edges("node_a", continue_to_jokes)
 
 ## Checkpointer
 
-One of the main benefits of LangGraph is that it comes backed by a persistence layer. This is accomplished via [checkpointers][basecheckpointsaver].
+LangGraph has a built-in persistence layer, implemented through [checkpointers][basecheckpointsaver]. When you use a checkpointer with a graph, you can interact with the state of that graph. When you use a checkpointer with a graph, you can interact with and manage the graph's state. The checkpointer saves a _checkpoint_ of the graph state at every super-step, enabling several powerful capabilities:
 
-Checkpointers can be used to save a _checkpoint_ of the state of a graph after all steps of the graph. This allows for several things.
-
-First, it allows for [human-in-the-loop workflows](agentic_concepts.md#human-in-the-loop), as it allows humans to inspect, interrupt, and approve steps. Checkpointers are needed for these workflows as the human has to be able to view the state of a graph at any point in time, and the graph has to be to resume execution after the human has made any updates to the state.
+First, checkpointers facilitate [human-in-the-loop workflows](agentic_concepts.md#human-in-the-loop) workflows by allowing humans to inspect, interrupt, and approve steps.Checkpointers are needed for these workflows as the human has to be able to view the state of a graph at any point in time, and the graph has to be to resume execution after the human has made any updates to the state.
 
 Second, it allows for ["memory"](agentic_concepts.md#memory) between interactions. You can use checkpointers to create threads and save the state of a thread after a graph executes. In the case of repeated human interactions (like conversations) any follow up messages can be sent to that checkpoint, which will retain its memory of previous ones.
 
@@ -244,8 +244,7 @@ See [this guide](../how-tos/persistence.ipynb) for how to add a checkpointer to
 
 ## Threads
 
-When using a checkpointer, you must specify a `thread_id` or `thread_ts` when running the graph.
-Threads are used to checkpoint multiple different runs. This can be used to enable a multi-tenant chat applications.
+Threads enable the checkpointing of multiple different runs, making them essential for multi-tenant chat applications and other scenarios where maintaining separate states is necessary. A thread is a unique ID assigned to a series of checkpoints saved by a checkpointer. When using a checkpointer, you must specify a `thread_id` or `thread_ts` when running the graph.
 
 `thread_id` is simply the ID of a thread. This is always required
 
@@ -262,14 +261,7 @@ See [this guide](../how-tos/persistence.ipynb) for how to use threads.
 
 ## Checkpointer state
 
-When you use a checkpointer with a graph, you can interact with the state of that graph.
-This usually done when enabling different human-in-the-loop interaction patterns.
-Each time you run the graph, the checkpointer creates several checkpoints every time a
-node or set of nodes finishes running.
-The most recent checkpoint is the current state of the thread.
-When interacting with the checkpointer state, you must specify a [thread identifier](#threads).
-
-Each checkpoint has two properties:
+ When interacting with the checkpointer state, you must specify a [thread identifier](#threads).Each checkpoint saved by the checkpointer has two properties:
 
 - **values**: This is the value of the state at this point in time.
 - **next**: This is a tuple of the nodes to execute next in the graph.

docs/docs/how-tos/index.md

@@ -22,8 +22,12 @@ These how-to guides show how to achieve that controllability.
 LangGraph makes it easy to persist state across graph runs. The guide below shows how to add persistence to your graph.
 
 - [How to add persistence ("memory") to your graph](persistence.ipynb)
-- [How to manage conversation history](managing-conversation-history.ipynb)
+- [How to manage conversation history](memory/manage-conversation-history.ipynb)
+- [How to delete messages](memory/delete-messages.ipynb)
+- [How to add summary conversation memory](memory/add-summary-conversation-history.ipynb)
 - [How to create a custom checkpointer using Postgres](persistence_postgres.ipynb)
+- [How to create a custom checkpointer using MongoDB](persistence_mongodb.ipynb)
+- [How to create a custom checkpointer using Redis](persistence_redis.ipynb)
 
 ## Human in the Loop
 
@@ -43,17 +47,28 @@ These guides show how to use different streaming modes.
 - [How to stream full state of your graph](stream-values.ipynb)
 - [How to stream state updates of your graph](stream-updates.ipynb)
 - [How to stream LLM tokens](streaming-tokens.ipynb)
+- [How to stream LLM tokens without LangChain models](streaming-tokens-without-langchain.ipynb)
 - [How to stream arbitrarily nested content](streaming-content.ipynb)
 - [How to configure multiple streaming modes at the same time](stream-multiple.ipynb)
 - [How to stream events from within a tool](streaming-events-from-within-tools.ipynb)
+- [How to stream events from within a tool without LangChain models](streaming-events-from-within-tools-without-langchain.ipynb)
 - [How to stream events from the final node](streaming-from-final-node.ipynb)
 
+## Tool calling
+
+- [How to call tools using ToolNode](tool-calling.ipynb)
+- [How to handle tool calling errors](tool-calling-errors.ipynb)
+- [How to pass graph state to tools](pass-run-time-values-to-tools.ipynb)
+- [How to pass config to tools](pass-config-to-tools.ipynb)
+
 ## Other
+
 - [How to run graph asynchronously](async.ipynb)
 - [How to visualize your graph](visualization.ipynb)
 - [How to add runtime configuration to your graph](configuration.ipynb)
 - [How to use a Pydantic model as your state](state-model.ipynb)
 - [How to use a context object in state](state-context-key.ipynb)
+- [How to add node retries](node-retries.ipynb)
 
 ## Prebuilt ReAct Agent
 

docs/docs/reference/checkpoints.md

@@ -9,17 +9,19 @@ You can [compile][langgraph.graph.MessageGraph.compile] any LangGraph workflow w
 
 ### Checkpoint
 
-::: langgraph.checkpoint.Checkpoint
+::: langgraph.checkpoint.base.Checkpoint
+
+### CheckpointMetadata
+
+::: langgraph.checkpoint.base.CheckpointMetadata
 
 ### BaseCheckpointSaver
 
 ::: langgraph.checkpoint.base.BaseCheckpointSaver
-handler: python
 
 ### SerializerProtocol
 
 ::: langgraph.checkpoint.SerializerProtocol
-handler: python
 
 ## Implementations
 
@@ -28,12 +30,10 @@ LangGraph also natively provides the following checkpoint implementations.
 ### MemorySaver
 
 ::: langgraph.checkpoint.memory.MemorySaver
-handler: python
 
 ### AsyncSqliteSaver
 
 ::: langgraph.checkpoint.aiosqlite.AsyncSqliteSaver
-handler: python
 
 ### SqliteSaver
 

docs/docs/reference/graphs.md

@@ -65,4 +65,8 @@ builder.add_conditional_edges("my_node", my_condition)
 
 ## Send
 
-::: langgraph.constants.Send
+::: langgraph.constants.Send
+
+## RetryPolicy
+
+::: langgraph.pregel.types.RetryPolicy

docs/docs/reference/prebuilt.md

@@ -55,4 +55,13 @@ from langgraph.prebuilt import tools_condition
 from langgraph.prebuilt import ValidationNode
 ```
 
-::: langgraph.prebuilt.ValidationNode
+::: langgraph.prebuilt.ValidationNode
+
+## InjectedState
+
+```python
+from langgraph.prebuilt import InjectedState
+```
+
+::: langgraph.prebuilt.InjectedState
+    handler: python

docs/mkdocs.yml

@@ -12,8 +12,6 @@ theme:
     repo: fontawesome/brands/git-alt
   features:
     - announce.dismiss
-    - content.action.edit
-    - content.action.view
     - content.code.annotate
     - content.code.copy
     - content.code.select
@@ -133,8 +131,12 @@ nav:
           - Create map-reduce branches for parallel execution: how-tos/map-reduce.ipynb
       - Persistence:
           - Add persistence ("memory"): how-tos/persistence.ipynb
-          - Manage conversation history: how-tos/managing-conversation-history.ipynb
+          - Manage conversation history: how-tos/memory/manage-conversation-history.ipynb
+          - Delete messages: how-tos/memory/delete-messages.ipynb
+          - Add summary of the conversation history: how-tos/memory/add-summary-conversation-history.ipynb
           - Create custom checkpointer using Postgres: how-tos/persistence_postgres.ipynb
+          - Create custom checkpointer using MongoDB: how-tos/persistence_mongodb.ipynb
+          - Create custom checkpointer using Redis: how-tos/persistence_redis.ipynb
       - Human-in-the-loop:
           - Add breakpoints: how-tos/human_in_the_loop/breakpoints.ipynb
           - Wait for user input: how-tos/human_in_the_loop/wait-user-input.ipynb
@@ -144,16 +146,24 @@ nav:
           - Stream full state: how-tos/stream-values.ipynb
           - Stream state updates: how-tos/stream-updates.ipynb
           - Stream LLM tokens: how-tos/streaming-tokens.ipynb
+          - Stream LLM tokens without LangChain models: how-tos/streaming-tokens-without-langchain.ipynb
           - Stream arbitrarily nested content: how-tos/streaming-content.ipynb
           - Configure multiple streaming modes: how-tos/stream-multiple.ipynb
           - Stream events from within tools: how-tos/streaming-events-from-within-tools.ipynb
+          - Stream events from within tools without LangChain models: how-tos/streaming-events-from-within-tools-without-langchain.ipynb
           - Stream events from the final node: how-tos/streaming-from-final-node.ipynb
+      - Tool calling:
+          - Call tools using ToolNode: how-tos/tool-calling.ipynb
+          - Handle tool calling errors: how-tos/tool-calling-errors.ipynb
+          - Pass graph state to tools: how-tos/pass-run-time-values-to-tools.ipynb
+          - Pass config to tools: how-tos/pass-config-to-tools.ipynb
       - Other:
           - Run graph asynchronously: how-tos/async.ipynb
           - Visualize your graph: how-tos/visualization.ipynb
           - Add runtime configuration: how-tos/configuration.ipynb
           - Use Pydantic model as state: how-tos/state-model.ipynb
           - Use a context object in state: how-tos/state-context-key.ipynb
+          - Add node retries: how-tos/node-retries.ipynb
       - Prebuilt ReAct Agent:
           - Create a ReAct agent: how-tos/create-react-agent.ipynb
           - Add memory to a ReAct agent: how-tos/create-react-agent-memory.ipynb
@@ -179,6 +189,7 @@ nav:
           - Deployment:
               - Setup App: "cloud/deployment/setup.md"
               - Setup App (pyproject.toml): "cloud/deployment/setup_pyproject.md"
+              - Test App Locally: "cloud/deployment/test_locally.md"
               - Deploy to Cloud: "cloud/deployment/cloud.md"
               - Self-Host: "cloud/deployment/self_hosted.md"
           - Streaming:
@@ -211,7 +222,9 @@ nav:
               - Configure Agents: "cloud/how-tos/cloud_examples/configuration_cloud.ipynb"
               - Convert LangGraph calls to LangGraph Cloud calls: "cloud/how-tos/cloud_examples/langgraph_to_langgraph_cloud.ipynb"
               - Integrate Webhooks: 'cloud/how-tos/cloud_examples/webhooks.ipynb'
-      - Conceptual Guides: "cloud/concepts/index.md"
+      - Conceptual Guides:
+          - API Concepts: "cloud/concepts/api.md"
+          - Cloud Concepts: "cloud/concepts/cloud.md"
       - Reference:
           - API: "cloud/reference/api/api_ref.md"
           - SDK:

docs/overrides/main.html

@@ -146,6 +146,13 @@
       display: none !important;
     }
 
+    .jupyter-wrapper .jp-CodeCell .jp-Cell-inputWrapper .jp-InputPrompt.jp-InputArea-prompt {
+      display: none !important;
+    }
+
+    .jupyter-wrapper .jp-Notebook .jp-Cell .jp-OutputPrompt {
+      display: none !important;
+    }
   </style>
 {% endblock %}
 

examples/chatbots/customer-support.ipynb

@@ -2,11 +2,502 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "id": "4c37bb65-6e2c-42e4-bfa7-9df10e2652a0",
+   "id": "d9d1a28b-c2a1-4246-b1c2-c58d6938f798",
    "metadata": {},
    "source": [
-    "This example has moved! Check out the [Customer Support Tutorial](../customer-support/customer-support.ipynb) for more information."
+    "# Customer Support\n",
+    "\n",
+    "Here, we show an example of building a customer support chatbot.\n",
+    "\n",
+    "This customer support chatbot interacts with SQL database to answer questions.\n",
+    "We will use a mock SQL database to get started: the [Chinook](https://www.sqlitetutorial.net/sqlite-sample-database/) database.\n",
+    "This database is about sales from a music store: what songs and album exists, customer orders, things like that.\n",
+    "\n",
+    "This chatbot has two different states: \n",
+    "1. Music: the user can inquire about different songs and albums present in the store\n",
+    "2. Account: the user can ask questions about their account\n",
+    "\n",
+    "Under the hood, this is handled by two separate agents. \n",
+    "Each has a specific prompt and tools related to their objective. \n",
+    "There is also a generic agent who is responsible for routing between these two agents as needed.\n",
+    "\n",
+    "Note: This is a very simple example! For a more complete tutorial on building a customer support bot, check out the [Customer Support Tutorial](../customer-support/customer-support.ipynb) for more information."
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "35abc013-2613-4a49-a806-939dcf13ccf3",
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [],
+   "source": ["%%capture --no-stderr\n%pip install -U langgraph langchain-community langchain-openai scikit-learn"]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9431e7f1-07fa-49d9-ac45-29613703dcc1",
+   "metadata": {},
+   "source": [
+    "## Load the data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "3d1ef253-6b0c-4481-868c-e1fe84f2c8ff",
+   "metadata": {},
+   "outputs": [],
+   "source": ["import requests\n\nurl = \"https://storage.googleapis.com/benchmarks-artifacts/chinook/Chinook.db\"\nresponse = requests.get(url)\n\nwith open(\"Chinook.db\", \"wb\") as file:\n    file.write(response.content)"]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "61f7ef9c",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "['Album',\n",
+       " 'Artist',\n",
+       " 'Customer',\n",
+       " 'Employee',\n",
+       " 'Genre',\n",
+       " 'Invoice',\n",
+       " 'InvoiceLine',\n",
+       " 'MediaType',\n",
+       " 'Playlist',\n",
+       " 'PlaylistTrack',\n",
+       " 'Track']"
+      ]
+     },
+     "execution_count": 2,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": ["from langchain_community.utilities import SQLDatabase\n\ndb = SQLDatabase.from_uri(\"sqlite:///Chinook.db\")\ndb.get_usable_table_names()"]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1cf668e4-8cb4-4de1-bc5e-c90284bf74bc",
+   "metadata": {},
+   "source": [
+    "## Load an LLM\n",
+    "\n",
+    "We will load a language model to use.\n",
+    "For this demo we will use OpenAI."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "d9ea4e80-30e6-4d46-b480-35f0be2fb055",
+   "metadata": {},
+   "outputs": [],
+   "source": ["from langchain_openai import ChatOpenAI\n\nmodel = ChatOpenAI(temperature=0, streaming=True, model=\"gpt-4o\")"]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "73907422-7e05-431e-b06d-256c9ec1f6f6",
+   "metadata": {},
+   "source": [
+    "## Load Other Modules\n",
+    "\n",
+    "Load other modules we will use.\n",
+    "\n",
+    "All of the tools our agents will use will be custom tools. As such, we will use the `@tool` decorator to create custom tools.\n",
+    "\n",
+    "We will pass in messages to the agent, so we load `HumanMessage` and `SystemMessage`"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "ea958e9f-ab1f-49b5-bd85-16332055297c",
+   "metadata": {},
+   "outputs": [],
+   "source": ["from langchain_core.messages import HumanMessage, SystemMessage"]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "35271d4d-2a1c-41be-9359-a3a7c3fed3d9",
+   "metadata": {},
+   "source": [
+    "## Define the Customer Agent\n",
+    "\n",
+    "This agent is responsible for looking up customer information.\n",
+    "It will have a specific prompt as well a specific tool to look up information about that customer (after asking for their user id)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "975b039a",
+   "metadata": {},
+   "outputs": [],
+   "source": ["# This tool is given to the agent to look up information about a customer\ndef get_customer_info(customer_id: int):\n    \"\"\"Look up customer info given their ID. ALWAYS make sure you have the customer ID before invoking this.\"\"\"\n    return db.run(f\"SELECT * FROM Customer WHERE CustomerID = {customer_id};\")"]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "1d5fa446",
+   "metadata": {},
+   "outputs": [],
+   "source": ["customer_prompt = \"\"\"Your job is to help a user update their profile.\n\nYou only have certain tools you can use. These tools require specific input. If you don't know the required input, then ask the user for it.\n\nIf you are unable to help the user, you can \"\"\"\n\n\ndef get_customer_messages(messages):\n    return [SystemMessage(content=customer_prompt)] + messages\n\n\ncustomer_chain = get_customer_messages | model.bind_tools([get_customer_info])"]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "904a9485-3857-458e-8b9d-33bc33842bc9",
+   "metadata": {},
+   "source": [
+    "## Define the Music Agent\n",
+    "\n",
+    "This agent is responsible for figuring out information about music. To do that, we will create a prompt and various tools for looking up information about music\n",
+    "\n",
+    "First, we will create indexes for looking up artists and track names.\n",
+    "This will allow us to look up artists and tracks without having to spell their names exactly right."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "a8604a3b-b484-4b2b-a914-4236cb98c524",
+   "metadata": {},
+   "outputs": [],
+   "source": ["from langchain_community.vectorstores import SKLearnVectorStore\nfrom langchain_openai import OpenAIEmbeddings\n\nartists = db._execute(\"select * from Artist\")\nsongs = db._execute(\"select * from Track\")\nartist_retriever = SKLearnVectorStore.from_texts(\n    [a[\"Name\"] for a in artists], OpenAIEmbeddings(), metadatas=artists\n).as_retriever()\nsong_retriever = SKLearnVectorStore.from_texts(\n    [a[\"Name\"] for a in songs], OpenAIEmbeddings(), metadatas=songs\n).as_retriever()"]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ac7eb264-c572-4925-ad55-a1d52a18b1c0",
+   "metadata": {},
+   "source": [
+    "First, let's create a tool for getting albums by artist."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "0a2a2b74",
+   "metadata": {},
+   "outputs": [],
+   "source": ["def get_albums_by_artist(artist):\n    \"\"\"Get albums by an artist (or similar artists).\"\"\"\n    docs = artist_retriever.get_relevant_documents(artist)\n    artist_ids = \", \".join([str(d.metadata[\"ArtistId\"]) for d in docs])\n    return db.run(\n        f\"SELECT Title, Name FROM Album LEFT JOIN Artist ON Album.ArtistId = Artist.ArtistId WHERE Album.ArtistId in ({artist_ids});\",\n        include_columns=True,\n    )"]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "45e85066-f2fc-490e-992d-cd66c9cd6486",
+   "metadata": {},
+   "source": [
+    "Next, lets create a tool for getting tracks by an artist"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "da533f50",
+   "metadata": {},
+   "outputs": [],
+   "source": ["def get_tracks_by_artist(artist):\n    \"\"\"Get songs by an artist (or similar artists).\"\"\"\n    docs = artist_retriever.invoke(artist)\n    artist_ids = \", \".join([str(d.metadata[\"ArtistId\"]) for d in docs])\n    return db.run(\n        f\"SELECT Track.Name as SongName, Artist.Name as ArtistName FROM Album LEFT JOIN Artist ON Album.ArtistId = Artist.ArtistId LEFT JOIN Track ON Track.AlbumId = Album.AlbumId WHERE Album.ArtistId in ({artist_ids});\",\n        include_columns=True,\n    )"]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "bb0e50ab-b059-427c-924b-f8072d8db23c",
+   "metadata": {},
+   "source": [
+    "Finally, let's create a tool for looking up songs by their name."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "b3c07010",
+   "metadata": {},
+   "outputs": [],
+   "source": ["def check_for_songs(song_title):\n    \"\"\"Check if a song exists by its name.\"\"\"\n    return song_retriever.invoke(song_title)"]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "88388ff8-38b5-4e4e-a24d-de8c3670bd2b",
+   "metadata": {},
+   "source": [
+    "Create the chain to call the relevant tools"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "72a14d5c",
+   "metadata": {},
+   "outputs": [],
+   "source": ["song_system_message = \"\"\"Your job is to help a customer find any songs they are looking for. \n\nYou only have certain tools you can use. If a customer asks you to look something up that you don't know how, politely tell them what you can help with.\n\nWhen looking up artists and songs, sometimes the artist/song will not be found. In that case, the tools will return information \\\non similar songs and artists. This is intentional, it is not the tool messing up.\"\"\"\n\n\ndef get_song_messages(messages):\n    return [SystemMessage(content=song_system_message)] + messages\n\n\nsong_recc_chain = get_song_messages | model.bind_tools(\n    [get_albums_by_artist, get_tracks_by_artist, check_for_songs]\n)"]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "cff15eb0-62c7-451d-a5f9-4576b24c879e",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "AIMessage(content='', additional_kwargs={'tool_calls': [{'index': 0, 'id': 'call_aXa9rSRXvTCJabrMY6AqkSV8', 'function': {'arguments': '{\"artist\":\"Amy Winehouse\"}', 'name': 'get_tracks_by_artist'}, 'type': 'function'}]}, response_metadata={'finish_reason': 'tool_calls', 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_ce0793330f'}, id='run-60633269-fd02-43e8-b434-28e3b8b69fb1-0', tool_calls=[{'name': 'get_tracks_by_artist', 'args': {'artist': 'Amy Winehouse'}, 'id': 'call_aXa9rSRXvTCJabrMY6AqkSV8'}])"
+      ]
+     },
+     "execution_count": 12,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": ["msgs = [HumanMessage(content=\"hi! can you help me find songs by amy whinehouse?\")]\nsong_recc_chain.invoke(msgs)"]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0a42c293-0816-4f3c-b4a3-5b9f3a0665d1",
+   "metadata": {},
+   "source": [
+    "## Define the Generic Agent\n",
+    "\n",
+    "We now define a generic agent that is responsible for handling initial inquiries and routing to the right sub agent."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "73e74268",
+   "metadata": {},
+   "outputs": [],
+   "source": ["from langchain_core.messages import AIMessage, HumanMessage, SystemMessage\nfrom langchain_core.pydantic_v1 import BaseModel, Field\n\n\nclass Router(BaseModel):\n    \"\"\"Call this if you are able to route the user to the appropriate representative.\"\"\"\n\n    choice: str = Field(description=\"should be one of: music, customer\")\n\n\nsystem_message = \"\"\"Your job is to help as a customer service representative for a music store.\n\nYou should interact politely with customers to try to figure out how you can help. You can help in a few ways:\n\n- Updating user information: if a customer wants to update the information in the user database. Call the router with `customer`\n- Recommending music: if a customer wants to find some music or information about music. Call the router with `music`\n\nIf the user is asking or wants to ask about updating or accessing their information, send them to that route.\nIf the user is asking or wants to ask about music, send them to that route.\nOtherwise, respond.\"\"\"\n\n\ndef get_messages(messages):\n    return [SystemMessage(content=system_message)] + messages"]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "ddf27314",
+   "metadata": {},
+   "outputs": [],
+   "source": ["chain = get_messages | model.bind_tools([Router])"]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "id": "3c896f34",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "AIMessage(content='', additional_kwargs={'tool_calls': [{'index': 0, 'id': 'call_0aaFPPCWDiAoPyXQX2PS8TcJ', 'function': {'arguments': '{\"choice\":\"music\"}', 'name': 'Router'}, 'type': 'function'}]}, response_metadata={'finish_reason': 'tool_calls', 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_ce0793330f'}, id='run-73d51d75-b7c5-49fe-b558-105ede7c75d1-0', tool_calls=[{'name': 'Router', 'args': {'choice': 'music'}, 'id': 'call_0aaFPPCWDiAoPyXQX2PS8TcJ'}])"
+      ]
+     },
+     "execution_count": 15,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": ["msgs = [HumanMessage(content=\"hi! can you help me find a good song?\")]\nchain.invoke(msgs)"]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "id": "40d86f59",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "AIMessage(content='', additional_kwargs={'tool_calls': [{'index': 0, 'id': 'call_Okla9DfMHIPs5TslS6KPaoBA', 'function': {'arguments': '{\"choice\":\"customer\"}', 'name': 'Router'}, 'type': 'function'}]}, response_metadata={'finish_reason': 'tool_calls', 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_4008e3b719'}, id='run-0d82d4f8-f4eb-4b16-add8-3e7fdffd6332-0', tool_calls=[{'name': 'Router', 'args': {'choice': 'customer'}, 'id': 'call_Okla9DfMHIPs5TslS6KPaoBA'}])"
+      ]
+     },
+     "execution_count": 16,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": ["msgs = [HumanMessage(content=\"hi! what's the email you have for me?\")]\nchain.invoke(msgs)"]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "id": "bd6ddd8b-7500-46a7-811d-3bcb937bda51",
+   "metadata": {},
+   "outputs": [],
+   "source": ["from langchain_core.messages import AIMessage\n\n\ndef add_name(message, name):\n    _dict = message.dict()\n    _dict[\"name\"] = name\n    return AIMessage(**_dict)"]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "id": "27494de5-8345-4c23-bc0e-81e0dd5d47d8",
+   "metadata": {},
+   "outputs": [],
+   "source": ["import json\n\nfrom langgraph.graph import END, START\n\n\ndef _get_last_ai_message(messages):\n    for m in messages[::-1]:\n        if isinstance(m, AIMessage):\n            return m\n    return None\n\n\ndef _is_tool_call(msg):\n    return hasattr(msg, \"additional_kwargs\") and \"tool_calls\" in msg.additional_kwargs\n\n\ndef _route(messages):\n    last_message = messages[-1]\n    if isinstance(last_message, AIMessage):\n        if not last_message.tool_calls:\n            return END\n        else:\n            if last_message.name == \"general\":\n                if len(last_message.tool_calls) > 1:\n                    raise ValueError(\"Too many tools\")\n                return last_message.tool_calls[0][\"args\"][\"choice\"]\n            else:\n                return \"tools\"\n    last_m = _get_last_ai_message(messages)\n    if last_m is None:\n        return \"general\"\n    if last_m.name == \"music\":\n        return \"music\"\n    elif last_m.name == \"customer\":\n        return \"customer\"\n    else:\n        return \"general\""]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "id": "8aec704a-46fe-4fb3-bdee-11c3bbffc370",
+   "metadata": {},
+   "outputs": [],
+   "source": ["from langgraph.prebuilt import ToolNode\n\ntools = [get_albums_by_artist, get_tracks_by_artist, check_for_songs, get_customer_info]\ntool_node = ToolNode(tools)"]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "id": "4d5b75c6-73e0-4922-a765-a15be63f869e",
+   "metadata": {},
+   "outputs": [],
+   "source": ["def _filter_out_routes(messages):\n    ms = []\n    for m in messages:\n        if _is_tool_call(m):\n            if m.name == \"general\":\n                continue\n        ms.append(m)\n    return ms"]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "id": "fd4dbf98-dbb3-411a-bad6-2bb334072aaf",
+   "metadata": {},
+   "outputs": [],
+   "source": ["from functools import partial\n\ngeneral_node = _filter_out_routes | chain | partial(add_name, name=\"general\")\nmusic_node = _filter_out_routes | song_recc_chain | partial(add_name, name=\"music\")\ncustomer_node = _filter_out_routes | customer_chain | partial(add_name, name=\"customer\")"]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "id": "dcade924",
+   "metadata": {},
+   "outputs": [],
+   "source": ["from langgraph.checkpoint.sqlite import SqliteSaver\n\nfrom langgraph.graph import MessageGraph\n\nmemory = SqliteSaver.from_conn_string(\":memory:\")\ngraph = MessageGraph()\nnodes = {\n    \"general\": \"general\",\n    \"music\": \"music\",\n    END: END,\n    \"tools\": \"tools\",\n    \"customer\": \"customer\",\n}\n# Define a new graph\nworkflow = MessageGraph()\nworkflow.add_node(\"general\", general_node)\nworkflow.add_node(\"music\", music_node)\nworkflow.add_node(\"customer\", customer_node)\nworkflow.add_node(\"tools\", tool_node)\nworkflow.add_conditional_edges(\"general\", _route, nodes)\nworkflow.add_conditional_edges(\"tools\", _route, nodes)\nworkflow.add_conditional_edges(\"music\", _route, nodes)\nworkflow.add_conditional_edges(\"customer\", _route, nodes)\nworkflow.add_conditional_edges(START, _route, nodes)\ngraph = workflow.compile()"]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 23,
+   "id": "ac65d6d2",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdin",
+     "output_type": "stream",
+     "text": [
+      "User (q/Q to quit):  what music do you have?\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Output from node 'general':\n",
+      "---\n",
+      "content='' additional_kwargs={'tool_calls': [{'index': 0, 'id': 'call_iste6NuKvZou8O9QudOectOU', 'function': {'arguments': '{\"choice\":\"music\"}', 'name': 'Router'}, 'type': 'function'}]} response_metadata={'finish_reason': 'tool_calls', 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_4008e3b719'} name='general' id='run-9eb940ff-6592-43ae-aa34-22c3d630ac65-0' tool_calls=[{'name': 'Router', 'args': {'choice': 'music'}, 'id': 'call_iste6NuKvZou8O9QudOectOU'}]\n",
+      "\n",
+      "---\n",
+      "\n",
+      "Output from node 'music':\n",
+      "---\n",
+      "content=\"I can help you find songs and albums by specific artists, or check if a particular song exists. Just let me know the name of the artist or song you're interested in!\" response_metadata={'finish_reason': 'stop', 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_4008e3b719'} name='music' id='run-91f560e7-ffa5-437f-afda-27490cbd1efe-0'\n",
+      "\n",
+      "---\n",
+      "\n"
+     ]
+    },
+    {
+     "name": "stdin",
+     "output_type": "stream",
+     "text": [
+      "User (q/Q to quit):  how about shakira?\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Output from node 'general':\n",
+      "---\n",
+      "content='' additional_kwargs={'tool_calls': [{'index': 0, 'id': 'call_fH4oKyA3U9aQy3p31MYXv2VP', 'function': {'arguments': '{\"choice\":\"music\"}', 'name': 'Router'}, 'type': 'function'}]} response_metadata={'finish_reason': 'tool_calls', 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_d576307f90'} name='general' id='run-6f2eee09-9e3e-4011-a045-196cc5baa1ee-0' tool_calls=[{'name': 'Router', 'args': {'choice': 'music'}, 'id': 'call_fH4oKyA3U9aQy3p31MYXv2VP'}]\n",
+      "\n",
+      "---\n",
+      "\n",
+      "Output from node 'music':\n",
+      "---\n",
+      "content='' additional_kwargs={'tool_calls': [{'index': 0, 'id': 'call_qivZqsI8zQAqSDP2jsvHyR7T', 'function': {'arguments': '{\"artist\": \"Shakira\"}', 'name': 'get_albums_by_artist'}, 'type': 'function'}, {'index': 1, 'id': 'call_GER0B3vlAjxcvYOYq1NGlV4r', 'function': {'arguments': '{\"artist\": \"Shakira\"}', 'name': 'get_tracks_by_artist'}, 'type': 'function'}]} response_metadata={'finish_reason': 'tool_calls', 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_ce0793330f'} name='music' id='run-38e2eca8-771c-48a2-83be-023fd17fc6f6-0' tool_calls=[{'name': 'get_albums_by_artist', 'args': {'artist': 'Shakira'}, 'id': 'call_qivZqsI8zQAqSDP2jsvHyR7T'}, {'name': 'get_tracks_by_artist', 'args': {'artist': 'Shakira'}, 'id': 'call_GER0B3vlAjxcvYOYq1NGlV4r'}]\n",
+      "\n",
+      "---\n",
+      "\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/Users/wfh/code/lc/langchain/libs/core/langchain_core/_api/deprecation.py:139: LangChainDeprecationWarning: The method `BaseRetriever.get_relevant_documents` was deprecated in langchain-core 0.1.46 and will be removed in 0.3.0. Use invoke instead.\n",
+      "  warn_deprecated(\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Output from node 'tools':\n",
+      "---\n",
+      "[ToolMessage(content=\"[{'Title': 'Supernatural', 'Name': 'Santana'}, {'Title': 'Santana - As Years Go By', 'Name': 'Santana'}, {'Title': 'Santana Live', 'Name': 'Santana'}, {'Title': 'Lulu Santos - RCA 100 Anos De Música - Álbum 01', 'Name': 'Lulu Santos'}, {'Title': 'Lulu Santos - RCA 100 Anos De Música - Álbum 02', 'Name': 'Lulu Santos'}]\", name='get_albums_by_artist', id='14ad12f3-afa1-4375-a89d-e878babf2d95', tool_call_id='call_qivZqsI8zQAqSDP2jsvHyR7T'), ToolMessage(content='[{\\'SongName\\': \\'(Da Le) Yaleo\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'Love Of My Life\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'Put Your Lights On\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'Africa Bamba\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'Smooth\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'Do You Like The Way\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'Maria Maria\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'Migra\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'Corazon Espinado\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'Wishing It Was\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'El Farol\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'Primavera\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'The Calling\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'Jingo\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'El Corazon Manda\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'La Puesta Del Sol\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'Persuasion\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'As The Years Go by\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'Soul Sacrifice\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'Fried Neckbones And Home Fries\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'Santana Jam\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'Evil Ways\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \"We\\'ve Got To Get Together/Jingo\", \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'Rock Me\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \"Just Ain\\'t Good Enough\", \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'Funky Piano\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'The Way You Do To Mer\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'Assim Caminha A Humanidade\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Um Pro Outro\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Casa\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Condição\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Satisfação\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Brumário\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Sábado À Noite\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'A Cura\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Atrás Do Trio Elétrico\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Tudo Bem\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Toda Forma De Amor\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Sereia\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Se Você Pensa\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Lá Vem O Sol (Here Comes The Sun)\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Honolulu\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Dancin´Days\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Aviso Aos Navegantes\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Hyperconectividade\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'O Descobridor Dos Sete Mares\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Um Certo Alguém\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Fullgás\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Aquilo\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Senta A Pua\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Ro-Que-Se-Da-Ne\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Tudo Igual\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Fogo De Palha\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Assaltaram A Gramática\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'O Último Romântico (Ao Vivo)\\', \\'ArtistName\\': \\'Lulu Santos\\'}]', name='get_tracks_by_artist', id='4ddcebc3-8e4e-42d4-8ae4-3ce3a62b548c', tool_call_id='call_GER0B3vlAjxcvYOYq1NGlV4r')]\n",
+      "\n",
+      "---\n",
+      "\n",
+      "Output from node 'music':\n",
+      "---\n",
+      "content=\"It seems I couldn't find specific albums or songs by Shakira, but I did find some related artists and their works. Here are some albums and songs by Santana and Lulu Santos:\\n\\n### Albums:\\n- **Santana:**\\n  - Supernatural\\n  - Santana - As Years Go By\\n  - Santana Live\\n\\n- **Lulu Santos:**\\n  - Lulu Santos - RCA 100 Anos De Música - Álbum 01\\n  - Lulu Santos - RCA 100 Anos De Música - Álbum 02\\n\\n### Songs:\\n- **Santana:**\\n  - (Da Le) Yaleo\\n  - Love Of My Life\\n  - Put Your Lights On\\n  - Africa Bamba\\n  - Smooth\\n  - Maria Maria\\n  - Corazon Espinado\\n  - Jingo\\n  - Evil Ways\\n\\n- **Lulu Santos:**\\n  - Assim Caminha A Humanidade\\n  - Um Pro Outro\\n  - Casa\\n  - Condição\\n  - Satisfação\\n  - A Cura\\n  - Atrás Do Trio Elétrico\\n  - Toda Forma De Amor\\n  - Sereia\\n  - Se Você Pensa\\n\\nIf you have any other specific artists or songs in mind, feel free to let me know!\" response_metadata={'finish_reason': 'stop', 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_d576307f90'} name='music' id='run-3d90d7cc-e7f7-48dc-bffd-00765c3f5d13-0'\n",
+      "\n",
+      "---\n",
+      "\n"
+     ]
+    },
+    {
+     "name": "stdin",
+     "output_type": "stream",
+     "text": [
+      "User (q/Q to quit):  hm cool\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Output from node 'general':\n",
+      "---\n",
+      "content='' additional_kwargs={'tool_calls': [{'index': 0, 'id': 'call_Nc5d0TWNbpnVJFeYJdQFuFGd', 'function': {'arguments': '{\"choice\":\"music\"}', 'name': 'Router'}, 'type': 'function'}]} response_metadata={'finish_reason': 'tool_calls', 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_4008e3b719'} name='general' id='run-19a3588e-f04c-48db-a54f-95d6b930ee3e-0' tool_calls=[{'name': 'Router', 'args': {'choice': 'music'}, 'id': 'call_Nc5d0TWNbpnVJFeYJdQFuFGd'}]\n",
+      "\n",
+      "---\n",
+      "\n",
+      "Output from node 'music':\n",
+      "---\n",
+      "content='' additional_kwargs={'tool_calls': [{'index': 0, 'id': 'call_weMHkgu3GYaMwXZM6hnCZd0z', 'function': {'arguments': '{\"artist\": \"Shakira\"}', 'name': 'get_albums_by_artist'}, 'type': 'function'}, {'index': 1, 'id': 'call_rXlVbPiEHbt10CNDUJ5GA2ZQ', 'function': {'arguments': '{\"artist\": \"Shakira\"}', 'name': 'get_tracks_by_artist'}, 'type': 'function'}]} response_metadata={'finish_reason': 'tool_calls', 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_ce0793330f'} name='music' id='run-71bf8550-b5a8-423f-b463-83d9ff51391b-0' tool_calls=[{'name': 'get_albums_by_artist', 'args': {'artist': 'Shakira'}, 'id': 'call_weMHkgu3GYaMwXZM6hnCZd0z'}, {'name': 'get_tracks_by_artist', 'args': {'artist': 'Shakira'}, 'id': 'call_rXlVbPiEHbt10CNDUJ5GA2ZQ'}]\n",
+      "\n",
+      "---\n",
+      "\n",
+      "Output from node 'tools':\n",
+      "---\n",
+      "[ToolMessage(content=\"[{'Title': 'Supernatural', 'Name': 'Santana'}, {'Title': 'Santana - As Years Go By', 'Name': 'Santana'}, {'Title': 'Santana Live', 'Name': 'Santana'}, {'Title': 'Lulu Santos - RCA 100 Anos De Música - Álbum 01', 'Name': 'Lulu Santos'}, {'Title': 'Lulu Santos - RCA 100 Anos De Música - Álbum 02', 'Name': 'Lulu Santos'}]\", name='get_albums_by_artist', id='52ada997-83f6-4500-b0a7-1104d4d38eb9', tool_call_id='call_weMHkgu3GYaMwXZM6hnCZd0z'), ToolMessage(content='[{\\'SongName\\': \\'(Da Le) Yaleo\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'Love Of My Life\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'Put Your Lights On\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'Africa Bamba\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'Smooth\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'Do You Like The Way\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'Maria Maria\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'Migra\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'Corazon Espinado\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'Wishing It Was\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'El Farol\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'Primavera\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'The Calling\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'Jingo\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'El Corazon Manda\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'La Puesta Del Sol\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'Persuasion\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'As The Years Go by\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'Soul Sacrifice\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'Fried Neckbones And Home Fries\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'Santana Jam\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'Evil Ways\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \"We\\'ve Got To Get Together/Jingo\", \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'Rock Me\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \"Just Ain\\'t Good Enough\", \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'Funky Piano\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'The Way You Do To Mer\\', \\'ArtistName\\': \\'Santana\\'}, {\\'SongName\\': \\'Assim Caminha A Humanidade\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Um Pro Outro\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Casa\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Condição\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Satisfação\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Brumário\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Sábado À Noite\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'A Cura\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Atrás Do Trio Elétrico\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Tudo Bem\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Toda Forma De Amor\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Sereia\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Se Você Pensa\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Lá Vem O Sol (Here Comes The Sun)\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Honolulu\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Dancin´Days\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Aviso Aos Navegantes\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Hyperconectividade\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'O Descobridor Dos Sete Mares\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Um Certo Alguém\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Fullgás\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Aquilo\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Senta A Pua\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Ro-Que-Se-Da-Ne\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Tudo Igual\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Fogo De Palha\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'Assaltaram A Gramática\\', \\'ArtistName\\': \\'Lulu Santos\\'}, {\\'SongName\\': \\'O Último Romântico (Ao Vivo)\\', \\'ArtistName\\': \\'Lulu Santos\\'}]', name='get_tracks_by_artist', id='7583e077-2fb4-44bf-bb00-b89d599ae84d', tool_call_id='call_rXlVbPiEHbt10CNDUJ5GA2ZQ')]\n",
+      "\n",
+      "---\n",
+      "\n",
+      "Output from node 'music':\n",
+      "---\n",
+      "content=\"It seems like I couldn't find specific albums or songs by Shakira. However, I did find some related artists and their works. Here are some albums and songs by Santana and Lulu Santos:\\n\\n### Albums:\\n1. **Santana**\\n   - Supernatural\\n   - Santana - As Years Go By\\n   - Santana Live\\n\\n2. **Lulu Santos**\\n   - RCA 100 Anos De Música - Álbum 01\\n   - RCA 100 Anos De Música - Álbum 02\\n\\n### Songs:\\n1. **Santana**\\n   - (Da Le) Yaleo\\n   - Love Of My Life\\n   - Put Your Lights On\\n   - Africa Bamba\\n   - Smooth\\n   - Maria Maria\\n   - Corazon Espinado\\n   - Jingo\\n   - Evil Ways\\n\\n2. **Lulu Santos**\\n   - Assim Caminha A Humanidade\\n   - Um Pro Outro\\n   - Casa\\n   - Condição\\n   - Satisfação\\n   - A Cura\\n   - Atrás Do Trio Elétrico\\n   - Toda Forma De Amor\\n   - Sereia\\n\\nIf you have any other artists or songs in mind, feel free to let me know!\" response_metadata={'finish_reason': 'stop', 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_d576307f90'} name='music' id='run-2418b45b-2762-49ca-b166-a0fee309ed9e-0'\n",
+      "\n",
+      "---\n",
+      "\n"
+     ]
+    },
+    {
+     "name": "stdin",
+     "output_type": "stream",
+     "text": [
+      "User (q/Q to quit):  q\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "AI: Byebye\n"
+     ]
+    }
+   ],
+   "source": ["import uuid\n\nfrom langchain_core.messages import HumanMessage\n\nfrom langgraph.graph.graph import START\n\nhistory = []\nwhile True:\n    user = input(\"User (q/Q to quit): \")\n    if user in {\"q\", \"Q\"}:\n        print(\"AI: Byebye\")\n        break\n    history.append(HumanMessage(content=user))\n    async for output in graph.astream(history):\n        for key, value in output.items():\n            print(f\"Output from node '{key}':\")\n            print(\"---\")\n            print(value)\n        print(\"\\n---\\n\")"]
   }
  ],
  "metadata": {

examples/cloud_examples/webhooks.ipynb

@@ -6,7 +6,7 @@
    "source": [
     "# Use Webhooks\n",
     "\n",
-    "You may wish to use webhooks in your client, especially when using async streams in case you want to update something in your service once the API call to Langgraph Cloud has finished running. To do so, you will need to expose an endpoint that can accept POST requests, and then pass it to your API request in the \"webhook\" parameter.\n",
+    "You may wish to use webhooks in your client, especially when using async streams in case you want to update something in your service once the API call to LangGraph Cloud has finished running. To do so, you will need to expose an endpoint that can accept POST requests, and then pass it to your API request in the \"webhook\" parameter.\n",
     "\n",
     "Currently, the SDK has not exposed this endpoint but you can access it through curl commands as follows.\n",
     "\n",

examples/create-react-agent-system-prompt.ipynb

@@ -9,7 +9,7 @@
     "\n",
     "This tutorial will show how to add a custom system prompt to the prebuilt ReAct agent. Please see [this tutorial](./create-react-agent.ipynb) for how to get started with the prebuilt ReAct agent\n",
     "\n",
-    "You can add a custom system prompt by passing a string to the `messages_modifier` param."
+    "You can add a custom system prompt by passing a string to the `state_modifier` param."
    ]
   },
   {
@@ -112,7 +112,7 @@
     "\n",
     "from langgraph.prebuilt import create_react_agent\n",
     "\n",
-    "graph = create_react_agent(model, tools=tools, messages_modifier=prompt)"
+    "graph = create_react_agent(model, tools=tools, state_modifier=prompt)"
    ]
   },
   {

examples/human_in_the_loop/breakpoints.ipynb

@@ -133,21 +133,26 @@
     "from langgraph.checkpoint.memory import MemorySaver\n",
     "from IPython.display import Image, display\n",
     "\n",
+    "\n",
     "class State(TypedDict):\n",
     "    input: str\n",
     "\n",
+    "\n",
     "def step_1(state):\n",
     "    print(\"---Step 1---\")\n",
     "    pass\n",
     "\n",
+    "\n",
     "def step_2(state):\n",
     "    print(\"---Step 2---\")\n",
     "    pass\n",
     "\n",
+    "\n",
     "def step_3(state):\n",
     "    print(\"---Step 3---\")\n",
     "    pass\n",
     "\n",
+    "\n",
     "builder = StateGraph(State)\n",
     "builder.add_node(\"step_1\", step_1)\n",
     "builder.add_node(\"step_2\", step_2)\n",
@@ -160,7 +165,7 @@
     "# Set up memory\n",
     "memory = MemorySaver()\n",
     "\n",
-    "# Add \n",
+    "# Add\n",
     "graph = builder.compile(checkpointer=memory, interrupt_before=[\"step_3\"])\n",
     "\n",
     "# View\n",
@@ -222,8 +227,7 @@
     "\n",
     "user_approval = input(\"Do you want to go to Step 3? (yes/no): \")\n",
     "\n",
-    "if user_approval.lower() == 'yes':\n",
-    "    \n",
+    "if user_approval.lower() == \"yes\":\n",
     "    # If approved, continue the graph execution\n",
     "    for event in graph.stream(None, thread, stream_mode=\"values\"):\n",
     "        print(event)\n",

examples/introduction.ipynb

@@ -277,21 +277,21 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Assistant: Langgraph is a new open-source deep learning framework that focuses on enabling efficient training and deployment of large language models. Some key things to know about Langgraph:\n",
+      "Assistant: LangGraph is a new open-source deep learning framework that focuses on enabling efficient training and deployment of large language models. Some key things to know about LangGraph:\n",
       "\n",
-      "1. Efficient Training: Langgraph is designed to accelerate the training of large language models by leveraging advanced optimization techniques and parallelization strategies.\n",
+      "1. Efficient Training: LangGraph is designed to accelerate the training of large language models by leveraging advanced optimization techniques and parallelization strategies.\n",
       "\n",
-      "2. Modular Architecture: Langgraph has a modular architecture that allows for easy customization and extension of language models, making it flexible for a variety of NLP tasks.\n",
+      "2. Modular Architecture: LangGraph has a modular architecture that allows for easy customization and extension of language models, making it flexible for a variety of NLP tasks.\n",
       "\n",
       "3. Hardware Acceleration: The framework is optimized for both CPU and GPU hardware, allowing for efficient model deployment on a wide range of devices.\n",
       "\n",
-      "4. Scalability: Langgraph is designed to handle large-scale language models with billions of parameters, enabling the development of state-of-the-art NLP applications.\n",
+      "4. Scalability: LangGraph is designed to handle large-scale language models with billions of parameters, enabling the development of state-of-the-art NLP applications.\n",
       "\n",
-      "5. Open-Source: Langgraph is an open-source project, allowing developers and researchers to collaborate, contribute, and build upon the framework.\n",
+      "5. Open-Source: LangGraph is an open-source project, allowing developers and researchers to collaborate, contribute, and build upon the framework.\n",
       "\n",
-      "6. Performance: The goal of Langgraph is to provide superior performance and efficiency compared to existing deep learning frameworks, particularly for training and deploying large language models.\n",
+      "6. Performance: The goal of LangGraph is to provide superior performance and efficiency compared to existing deep learning frameworks, particularly for training and deploying large language models.\n",
       "\n",
-      "Overall, Langgraph is a promising new deep learning framework that aims to address the challenges of building and deploying advanced natural language processing models at scale. It is an active area of research and development, with the potential to drive further advancements in the field of language AI.\n"
+      "Overall, LangGraph is a promising new deep learning framework that aims to address the challenges of building and deploying advanced natural language processing models at scale. It is an active area of research and development, with the potential to drive further advancements in the field of language AI.\n"
      ]
     },
     {
@@ -1224,7 +1224,8 @@
     "graph = graph_builder.compile(checkpointer=memory)\n",
     "```\n",
     "</pre>\n",
-    "</pre>details"
+    "</pre>\n",
+    "</details>"
    ]
   },
   {

examples/memory/add-summary-conversation-history.ipynb

@@ -0,0 +1,542 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "51466c8d-8ce4-4b3d-be4e-18fdbeda5f53",
+   "metadata": {},
+   "source": [
+    "# How to add summary of the conversation history\n",
+    "\n",
+    "One of the most common use cases for persistence is to use it to keep track of conversation history. This is great - it makes it easy to continue conversations. As conversations get longer and longer, however, this conversation history can build up and take up more and more of the context window. This can often be undesirable as it leads to more expensive and longer calls to the LLM, and potentially ones that error. One way to work around that is to create a summary of the conversation to date, and use that with the past N messages. This guide will go through an example of how to do that.\n",
+    "\n",
+    "This will involve a few steps:\n",
+    "- Check if the conversation is too long (can be done by checking number of messages or length of messages)\n",
+    "- If yes, the create summary (will need a prompt for this)\n",
+    "- Then remove all except the last N messages\n",
+    "\n",
+    "A big part of this is deleting old messages. For an in depth guide on how to do that, see [this guide](./delete-messages.ipynb)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7cbd446a-808f-4394-be92-d45ab818953c",
+   "metadata": {},
+   "source": [
+    "## Setup\n",
+    "\n",
+    "First, let's set up the packages we're going to want to use"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "af4ce0ba-7596-4e5f-8bf8-0b0bd6e62833",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%capture --no-stderr\n",
+    "%pip install --quiet -U langgraph langchain_anthropic"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0abe11f4-62ed-4dc4-8875-3db21e260d1d",
+   "metadata": {},
+   "source": [
+    "Next, we need to set API keys for Anthropic (the LLM we will use)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "c903a1cf-2977-4e2d-ad7d-8b3946821d89",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import getpass\n",
+    "import os\n",
+    "\n",
+    "\n",
+    "def _set_env(var: str):\n",
+    "    if not os.environ.get(var):\n",
+    "        os.environ[var] = getpass.getpass(f\"{var}: \")\n",
+    "\n",
+    "\n",
+    "_set_env(\"ANTHROPIC_API_KEY\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f0ed46a8-effe-4596-b0e1-a6a29ee16f5c",
+   "metadata": {},
+   "source": [
+    "Optionally, we can set API key for [LangSmith tracing](https://smith.langchain.com/), which will give us best-in-class observability."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "95e25aec-7c9f-4a63-b143-225d0e9a79c3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\n",
+    "_set_env(\"LANGCHAIN_API_KEY\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "84835fdb-a5f3-4c90-85f3-0e6257650aba",
+   "metadata": {},
+   "source": [
+    "## Build the chatbot\n",
+    "\n",
+    "Let's now build the chatbot."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 25,
+   "id": "378899a9-3b9a-4748-95b6-eb00e0828677",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from typing import Literal\n",
+    "\n",
+    "from langchain_anthropic import ChatAnthropic\n",
+    "from langchain_core.messages import SystemMessage, RemoveMessage\n",
+    "from langgraph.checkpoint.sqlite import SqliteSaver\n",
+    "from langgraph.graph import MessagesState, StateGraph, START, END\n",
+    "\n",
+    "memory = SqliteSaver.from_conn_string(\":memory:\")\n",
+    "\n",
+    "\n",
+    "# We will add a `summary` attribute (in addition to `messages` key,\n",
+    "# which MessagesState already has)\n",
+    "class State(MessagesState):\n",
+    "    summary: str\n",
+    "\n",
+    "\n",
+    "# We will use this model for both the conversation and the summarization\n",
+    "model = ChatAnthropic(model_name=\"claude-3-haiku-20240307\")\n",
+    "\n",
+    "\n",
+    "# Define the logic to call the model\n",
+    "def call_model(state: State):\n",
+    "    # If a summary exists, we add this in as a system message\n",
+    "    summary = state.get(\"summary\", \"\")\n",
+    "    if summary:\n",
+    "        system_message = f\"Summary of conversation earlier: {summary}\"\n",
+    "        messages = [SystemMessage(content=system_message)] + state[\"messages\"]\n",
+    "    else:\n",
+    "        messages = state[\"messages\"]\n",
+    "    response = model.invoke(messages)\n",
+    "    # We return a list, because this will get added to the existing list\n",
+    "    return {\"messages\": [response]}\n",
+    "\n",
+    "\n",
+    "# We now define the logic for determining whether to end or summarize the conversation\n",
+    "def should_continue(state: State) -> Literal[\"summarize_conversation\", END]:\n",
+    "    \"\"\"Return the next node to execute.\"\"\"\n",
+    "    messages = state[\"messages\"]\n",
+    "    # If there are more than six messages, then we summarize the conversation\n",
+    "    if len(messages) > 6:\n",
+    "        return \"summarize_conversation\"\n",
+    "    # Otherwise we can just end\n",
+    "    return END\n",
+    "\n",
+    "\n",
+    "def summarize_conversation(state: State):\n",
+    "    # First, we summarize the conversation\n",
+    "    summary = state.get(\"summary\", \"\")\n",
+    "    if summary:\n",
+    "        # If a summary already exists, we use a different system prompt\n",
+    "        # to summarize it than if one didn't\n",
+    "        summary_message = (\n",
+    "            f\"This is summary of the conversation to date: {summary}\\n\\n\"\n",
+    "            \"Extend the summary by taking into account the new messages above:\"\n",
+    "        )\n",
+    "    else:\n",
+    "        summary_message = \"Create a summary of the conversation above:\"\n",
+    "\n",
+    "    messages = state[\"messages\"] + [HumanMessage(content=summary_message)]\n",
+    "    response = model.invoke(messages)\n",
+    "    # We now need to delete messages that we no longer want to show up\n",
+    "    # I will delete all but the last two messages, but you can change this\n",
+    "    delete_messages = [RemoveMessage(id=m.id) for m in state[\"messages\"][:-2]]\n",
+    "    return {\"summary\": response.content, \"messages\": delete_messages}\n",
+    "\n",
+    "\n",
+    "# Define a new graph\n",
+    "workflow = StateGraph(State)\n",
+    "\n",
+    "# Define the conversation node and the summarize node\n",
+    "workflow.add_node(\"conversation\", call_model)\n",
+    "workflow.add_node(summarize_conversation)\n",
+    "\n",
+    "# Set the entrypoint as conversation\n",
+    "workflow.add_edge(START, \"conversation\")\n",
+    "\n",
+    "# We now add a conditional edge\n",
+    "workflow.add_conditional_edges(\n",
+    "    # First, we define the start node. We use `conversation`.\n",
+    "    # This means these are the edges taken after the `conversation` node is called.\n",
+    "    \"conversation\",\n",
+    "    # Next, we pass in the function that will determine which node is called next.\n",
+    "    should_continue,\n",
+    ")\n",
+    "\n",
+    "# We now add a normal edge from `summarize_conversation` to END.\n",
+    "# This means that after `summarize_conversation` is called, we end.\n",
+    "workflow.add_edge(\"summarize_conversation\", END)\n",
+    "\n",
+    "# Finally, we compile it!\n",
+    "app = workflow.compile(checkpointer=memory)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "41c2872e-04b3-4c44-9e03-9e84a5230adf",
+   "metadata": {},
+   "source": [
+    "## Using the graph"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 26,
+   "id": "dc697132-8fa1-4bf5-9722-56a9859331ab",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def print_update(update):\n",
+    "    for k, v in update.items():\n",
+    "        for m in v[\"messages\"]:\n",
+    "            m.pretty_print()\n",
+    "        if \"summary\" in v:\n",
+    "            print(v[\"summary\"])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 27,
+   "id": "57b27553-21be-43e5-ac48-d1d0a3aa0dca",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "================================\u001b[1m Human Message \u001b[0m=================================\n",
+      "\n",
+      "hi! I'm bob\n",
+      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
+      "\n",
+      "It's nice to meet you, Bob! I'm an AI assistant created by Anthropic. How can I help you today?\n",
+      "================================\u001b[1m Human Message \u001b[0m=================================\n",
+      "\n",
+      "what's my name?\n",
+      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
+      "\n",
+      "Your name is Bob, as you told me at the beginning of our conversation.\n",
+      "================================\u001b[1m Human Message \u001b[0m=================================\n",
+      "\n",
+      "i like the celtics!\n",
+      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
+      "\n",
+      "That's great, the Celtics are a fun team to follow! Basketball is an exciting sport. Do you have a favorite Celtics player or a favorite moment from a Celtics game you've watched? I'd be happy to discuss the team and the sport with you.\n"
+     ]
+    }
+   ],
+   "source": [
+    "from langchain_core.messages import HumanMessage\n",
+    "\n",
+    "config = {\"configurable\": {\"thread_id\": \"4\"}}\n",
+    "input_message = HumanMessage(content=\"hi! I'm bob\")\n",
+    "input_message.pretty_print()\n",
+    "for event in app.stream({\"messages\": [input_message]}, config, stream_mode=\"updates\"):\n",
+    "    print_update(event)\n",
+    "\n",
+    "input_message = HumanMessage(content=\"what's my name?\")\n",
+    "input_message.pretty_print()\n",
+    "for event in app.stream({\"messages\": [input_message]}, config, stream_mode=\"updates\"):\n",
+    "    print_update(event)\n",
+    "\n",
+    "input_message = HumanMessage(content=\"i like the celtics!\")\n",
+    "input_message.pretty_print()\n",
+    "for event in app.stream({\"messages\": [input_message]}, config, stream_mode=\"updates\"):\n",
+    "    print_update(event)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9760e219-a7fc-4d81-b4e8-1334c5afc510",
+   "metadata": {},
+   "source": [
+    "We can see that so far no summarization has happened - this is because there are only six messages in the list."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 28,
+   "id": "935265a0-d511-475a-8a0d-b3c3cc5e42a0",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'messages': [HumanMessage(content=\"hi! I'm bob\", id='6534853d-b8a7-44b9-837b-eb7abaf7ebf7'),\n",
+       "  AIMessage(content=\"It's nice to meet you, Bob! I'm an AI assistant created by Anthropic. How can I help you today?\", response_metadata={'id': 'msg_015wCFew2vwMQJcpUh2VZ5ah', 'model': 'claude-3-haiku-20240307', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 12, 'output_tokens': 30}}, id='run-0d33008b-1094-4f5e-94ce-293283fc3024-0'),\n",
+       "  HumanMessage(content=\"what's my name?\", id='0a4f203a-b95a-42a9-b1c5-bb20f68b3251'),\n",
+       "  AIMessage(content='Your name is Bob, as you told me at the beginning of our conversation.', response_metadata={'id': 'msg_01PLp8wg2xDsJbNR9uCtxcGz', 'model': 'claude-3-haiku-20240307', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 50, 'output_tokens': 19}}, id='run-3815dd4d-ee0c-4fc2-9889-f6dd40325961-0'),\n",
+       "  HumanMessage(content='i like the celtics!', id='ac128172-42d1-4390-b7cc-7bcb2d22ee48'),\n",
+       "  AIMessage(content=\"That's great, the Celtics are a fun team to follow! Basketball is an exciting sport. Do you have a favorite Celtics player or a favorite moment from a Celtics game you've watched? I'd be happy to discuss the team and the sport with you.\", response_metadata={'id': 'msg_01CSg5avZEx6CKcZsSvSVXpr', 'model': 'claude-3-haiku-20240307', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 78, 'output_tokens': 61}}, id='run-698faa28-0f72-495f-8ebe-e948664d2200-0')]}"
+      ]
+     },
+     "execution_count": 28,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "values = app.get_state(config).values\n",
+    "values"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "bb40eddb-9a31-4410-a4c0-9762e2d89e56",
+   "metadata": {},
+   "source": [
+    "Now let's send another message in"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 29,
+   "id": "048805a4-3d97-4e76-ac45-8d80d4364c46",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "================================\u001b[1m Human Message \u001b[0m=================================\n",
+      "\n",
+      "i like how much they win\n",
+      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
+      "\n",
+      "That's understandable, the Celtics have been one of the more successful NBA franchises over the years. Their history of winning championships is very impressive. It's always fun to follow a team that regularly competes for titles. What do you think has been the key to the Celtics' sustained success? Is there a particular era or team that stands out as your favorite?\n",
+      "================================\u001b[1m Remove Message \u001b[0m================================\n",
+      "\n",
+      "\n",
+      "================================\u001b[1m Remove Message \u001b[0m================================\n",
+      "\n",
+      "\n",
+      "================================\u001b[1m Remove Message \u001b[0m================================\n",
+      "\n",
+      "\n",
+      "================================\u001b[1m Remove Message \u001b[0m================================\n",
+      "\n",
+      "\n",
+      "================================\u001b[1m Remove Message \u001b[0m================================\n",
+      "\n",
+      "\n",
+      "================================\u001b[1m Remove Message \u001b[0m================================\n",
+      "\n",
+      "\n",
+      "Here is a summary of our conversation so far:\n",
+      "\n",
+      "- You introduced yourself as Bob and said you like the Boston Celtics basketball team.\n",
+      "- I acknowledged that it's nice to meet you, Bob, and noted that you had shared your name earlier in the conversation.\n",
+      "- You expressed that you like how much the Celtics win, and I agreed that their history of sustained success and championship pedigree is impressive.\n",
+      "- I asked if you have a favorite Celtics player or moment that stands out to you, and invited further discussion about the team and the sport of basketball.\n",
+      "- The overall tone has been friendly and conversational, with me trying to engage with your interest in the Celtics by asking follow-up questions.\n"
+     ]
+    }
+   ],
+   "source": [
+    "input_message = HumanMessage(content=\"i like how much they win\")\n",
+    "input_message.pretty_print()\n",
+    "for event in app.stream({\"messages\": [input_message]}, config, stream_mode=\"updates\"):\n",
+    "    print_update(event)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6b196367-6151-4982-9430-3db7373de06e",
+   "metadata": {},
+   "source": [
+    "If we check the state now, we can see that we have a summary of the conversation, as well as the last two messages"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 30,
+   "id": "09ebb693-4738-4474-a095-6491def5c5f9",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'messages': [HumanMessage(content='i like how much they win', id='bb916ce7-534c-4d48-9f92-e269f9dc4859'),\n",
+       "  AIMessage(content=\"That's understandable, the Celtics have been one of the more successful NBA franchises over the years. Their history of winning championships is very impressive. It's always fun to follow a team that regularly competes for titles. What do you think has been the key to the Celtics' sustained success? Is there a particular era or team that stands out as your favorite?\", response_metadata={'id': 'msg_01B7TMagaM8xBnYXLSMwUDAG', 'model': 'claude-3-haiku-20240307', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 148, 'output_tokens': 82}}, id='run-c5aa9a8f-7983-4a7f-9c1e-0c0055334ac1-0')],\n",
+       " 'summary': \"Here is a summary of our conversation so far:\\n\\n- You introduced yourself as Bob and said you like the Boston Celtics basketball team.\\n- I acknowledged that it's nice to meet you, Bob, and noted that you had shared your name earlier in the conversation.\\n- You expressed that you like how much the Celtics win, and I agreed that their history of sustained success and championship pedigree is impressive.\\n- I asked if you have a favorite Celtics player or moment that stands out to you, and invited further discussion about the team and the sport of basketball.\\n- The overall tone has been friendly and conversational, with me trying to engage with your interest in the Celtics by asking follow-up questions.\"}"
+      ]
+     },
+     "execution_count": 30,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "values = app.get_state(config).values\n",
+    "values"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "966e4177-c0fc-4fd0-a494-dd03f7f2fddb",
+   "metadata": {},
+   "source": [
+    "We can now resume having a conversation! Note that even though we only have the last two messages, we can still ask it questions about things mentioned earlier in the conversation (because we summarized those)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 31,
+   "id": "7094c5ab-66f8-42ff-b1c3-90c8a9468e62",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "================================\u001b[1m Human Message \u001b[0m=================================\n",
+      "\n",
+      "what's my name?\n",
+      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
+      "\n",
+      "In our conversation so far, you introduced yourself as Bob. I acknowledged that earlier when you had shared your name.\n"
+     ]
+    }
+   ],
+   "source": [
+    "input_message = HumanMessage(content=\"what's my name?\")\n",
+    "input_message.pretty_print()\n",
+    "for event in app.stream({\"messages\": [input_message]}, config, stream_mode=\"updates\"):\n",
+    "    print_update(event)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 32,
+   "id": "40e5db8e-9db9-4ac7-9d76-a99fd4034bf3",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "================================\u001b[1m Human Message \u001b[0m=================================\n",
+      "\n",
+      "what NFL team do you think I like?\n",
+      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
+      "\n",
+      "I don't actually have any information about what NFL team you might like. In our conversation so far, you've only mentioned that you're a fan of the Boston Celtics basketball team. I don't have any prior knowledge about your preferences for NFL teams. Unless you provide me with that information, I don't have a basis to guess which NFL team you might be a fan of.\n"
+     ]
+    }
+   ],
+   "source": [
+    "input_message = HumanMessage(content=\"what NFL team do you think I like?\")\n",
+    "input_message.pretty_print()\n",
+    "for event in app.stream({\"messages\": [input_message]}, config, stream_mode=\"updates\"):\n",
+    "    print_update(event)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 33,
+   "id": "0a1a0fda-5309-45f0-9465-9f3dff604d74",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "================================\u001b[1m Human Message \u001b[0m=================================\n",
+      "\n",
+      "i like the patriots!\n",
+      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
+      "\n",
+      "Okay, got it! Thanks for sharing that you're also a fan of the New England Patriots in the NFL. That makes sense, given your interest in other Boston sports teams like the Celtics. The Patriots have also had a very successful run over the past couple of decades, winning multiple Super Bowls. It's fun to follow winning franchises like the Celtics and Patriots. Do you have a favorite Patriots player or moment that stands out to you?\n",
+      "================================\u001b[1m Remove Message \u001b[0m================================\n",
+      "\n",
+      "\n",
+      "================================\u001b[1m Remove Message \u001b[0m================================\n",
+      "\n",
+      "\n",
+      "================================\u001b[1m Remove Message \u001b[0m================================\n",
+      "\n",
+      "\n",
+      "================================\u001b[1m Remove Message \u001b[0m================================\n",
+      "\n",
+      "\n",
+      "================================\u001b[1m Remove Message \u001b[0m================================\n",
+      "\n",
+      "\n",
+      "================================\u001b[1m Remove Message \u001b[0m================================\n",
+      "\n",
+      "\n",
+      "Okay, extending the summary with the new information:\n",
+      "\n",
+      "- You initially introduced yourself as Bob and said you like the Boston Celtics basketball team. \n",
+      "- I acknowledged that and we discussed your appreciation for the Celtics' history of winning.\n",
+      "- You then asked what your name was, and I reminded you that you had introduced yourself as Bob earlier in the conversation.\n",
+      "- You followed up by asking what NFL team I thought you might like, and I explained that I didn't have any prior information about your NFL team preferences.\n",
+      "- You then revealed that you are also a fan of the New England Patriots, which made sense given your Celtics fandom.\n",
+      "- I responded positively to this new information, noting the Patriots' own impressive success and dynasty over the past couple of decades.\n",
+      "- I then asked if you have a particular favorite Patriots player or moment that stands out to you, continuing the friendly, conversational tone.\n",
+      "\n",
+      "Overall, the discussion has focused on your sports team preferences, with you sharing that you are a fan of both the Celtics and the Patriots. I've tried to engage with your interests and ask follow-up questions to keep the dialogue flowing.\n"
+     ]
+    }
+   ],
+   "source": [
+    "input_message = HumanMessage(content=\"i like the patriots!\")\n",
+    "input_message.pretty_print()\n",
+    "for event in app.stream({\"messages\": [input_message]}, config, stream_mode=\"updates\"):\n",
+    "    print_update(event)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "67d26013-1362-4cee-b135-ab5c3c4eb3d0",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.1"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

examples/memory/delete-messages.ipynb

@@ -0,0 +1,500 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "51466c8d-8ce4-4b3d-be4e-18fdbeda5f53",
+   "metadata": {},
+   "source": [
+    "# How to delete messages\n",
+    "\n",
+    "One of the common states for a graph is a list of messages. Usually you only add messages to that state. However, sometimes you may want to remove messages (either by directly modifying the state or as part of the graph). To do that, you can use the `RemoveMessage` modifier. In this guide, we will cover how to do that.\n",
+    "\n",
+    "The key idea is that each state key has a `reducer` key. This key specifies how to combine updates to the state. The default `MessagesState` has a messages key, and the reducer for that key accepts these `RemoveMessage` modifiers. That reducer then uses these `RemoveMessage` to delete messages from the key.\n",
+    "\n",
+    "So note that just because your graph state has a key that is a list of messages, it doesn't mean that that this `RemoveMessage` modifier will work. You also have to have a `reducer` defined that knows how to work with this.\n",
+    "\n",
+    "**NOTE**: Many models expect certain rules around lists of messages. For example, some expect them to start with a `user` message, others expect all messages with tool calls to be followed by a tool message. **When deleting messages, you will want to make sure you don't violate these rules.**"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7cbd446a-808f-4394-be92-d45ab818953c",
+   "metadata": {},
+   "source": [
+    "## Setup\n",
+    "\n",
+    "First, let's build a simple graph that uses messages. Note that it's using the `MessagesState` which has the required `reducer`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "af4ce0ba-7596-4e5f-8bf8-0b0bd6e62833",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%capture --no-stderr\n",
+    "%pip install --quiet -U langgraph langchain_anthropic"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0abe11f4-62ed-4dc4-8875-3db21e260d1d",
+   "metadata": {},
+   "source": [
+    "Next, we need to set API keys for Anthropic (the LLM we will use)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "c903a1cf-2977-4e2d-ad7d-8b3946821d89",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdin",
+     "output_type": "stream",
+     "text": [
+      "ANTHROPIC_API_KEY:  ········\n"
+     ]
+    }
+   ],
+   "source": [
+    "import getpass\n",
+    "import os\n",
+    "\n",
+    "\n",
+    "def _set_env(var: str):\n",
+    "    if not os.environ.get(var):\n",
+    "        os.environ[var] = getpass.getpass(f\"{var}: \")\n",
+    "\n",
+    "\n",
+    "_set_env(\"ANTHROPIC_API_KEY\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f0ed46a8-effe-4596-b0e1-a6a29ee16f5c",
+   "metadata": {},
+   "source": [
+    "Optionally, we can set API key for [LangSmith tracing](https://smith.langchain.com/), which will give us best-in-class observability."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "95e25aec-7c9f-4a63-b143-225d0e9a79c3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\n",
+    "_set_env(\"LANGCHAIN_API_KEY\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4767ef1c-a7cf-41f8-a301-558988cb7ac5",
+   "metadata": {},
+   "source": [
+    "## Build the agent\n",
+    "Let's now build a simple ReAct style agent."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "378899a9-3b9a-4748-95b6-eb00e0828677",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from typing import Literal\n",
+    "\n",
+    "from langchain_anthropic import ChatAnthropic\n",
+    "from langchain_core.tools import tool\n",
+    "\n",
+    "from langgraph.checkpoint.sqlite import SqliteSaver\n",
+    "from langgraph.graph import MessagesState, StateGraph, START\n",
+    "from langgraph.prebuilt import ToolNode\n",
+    "\n",
+    "memory = SqliteSaver.from_conn_string(\":memory:\")\n",
+    "\n",
+    "\n",
+    "@tool\n",
+    "def search(query: str):\n",
+    "    \"\"\"Call to surf the web.\"\"\"\n",
+    "    # This is a placeholder for the actual implementation\n",
+    "    # Don't let the LLM know this though 😊\n",
+    "    return [\n",
+    "        \"It's sunny in San Francisco, but you better look out if you're a Gemini 😈.\"\n",
+    "    ]\n",
+    "\n",
+    "\n",
+    "tools = [search]\n",
+    "tool_node = ToolNode(tools)\n",
+    "model = ChatAnthropic(model_name=\"claude-3-haiku-20240307\")\n",
+    "bound_model = model.bind_tools(tools)\n",
+    "\n",
+    "\n",
+    "def should_continue(state: MessagesState) -> Literal[\"action\", \"__end__\"]:\n",
+    "    \"\"\"Return the next node to execute.\"\"\"\n",
+    "    last_message = state[\"messages\"][-1]\n",
+    "    # If there is no function call, then we finish\n",
+    "    if not last_message.tool_calls:\n",
+    "        return \"__end__\"\n",
+    "    # Otherwise if there is, we continue\n",
+    "    return \"action\"\n",
+    "\n",
+    "\n",
+    "# Define the function that calls the model\n",
+    "def call_model(state: MessagesState):\n",
+    "    response = model.invoke(state[\"messages\"])\n",
+    "    # We return a list, because this will get added to the existing list\n",
+    "    return {\"messages\": response}\n",
+    "\n",
+    "\n",
+    "# Define a new graph\n",
+    "workflow = StateGraph(MessagesState)\n",
+    "\n",
+    "# Define the two nodes we will cycle between\n",
+    "workflow.add_node(\"agent\", call_model)\n",
+    "workflow.add_node(\"action\", tool_node)\n",
+    "\n",
+    "# Set the entrypoint as `agent`\n",
+    "# This means that this node is the first one called\n",
+    "workflow.add_edge(START, \"agent\")\n",
+    "\n",
+    "# We now add a conditional edge\n",
+    "workflow.add_conditional_edges(\n",
+    "    # First, we define the start node. We use `agent`.\n",
+    "    # This means these are the edges taken after the `agent` node is called.\n",
+    "    \"agent\",\n",
+    "    # Next, we pass in the function that will determine which node is called next.\n",
+    "    should_continue,\n",
+    ")\n",
+    "\n",
+    "# We now add a normal edge from `tools` to `agent`.\n",
+    "# This means that after `tools` is called, `agent` node is called next.\n",
+    "workflow.add_edge(\"action\", \"agent\")\n",
+    "\n",
+    "# Finally, we compile it!\n",
+    "# This compiles it into a LangChain Runnable,\n",
+    "# meaning you can use it as you would any other runnable\n",
+    "app = workflow.compile(checkpointer=memory)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "57b27553-21be-43e5-ac48-d1d0a3aa0dca",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "================================\u001b[1m Human Message \u001b[0m=================================\n",
+      "\n",
+      "hi! I'm bob\n",
+      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
+      "\n",
+      "It's nice to meet you, Bob! I'm Claude, an AI assistant created by Anthropic. How can I assist you today?\n",
+      "================================\u001b[1m Human Message \u001b[0m=================================\n",
+      "\n",
+      "what's my name?\n",
+      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
+      "\n",
+      "Your name is Bob, as you introduced yourself at the beginning of our conversation.\n"
+     ]
+    }
+   ],
+   "source": [
+    "from langchain_core.messages import HumanMessage\n",
+    "\n",
+    "config = {\"configurable\": {\"thread_id\": \"2\"}}\n",
+    "input_message = HumanMessage(content=\"hi! I'm bob\")\n",
+    "for event in app.stream({\"messages\": [input_message]}, config, stream_mode=\"values\"):\n",
+    "    event[\"messages\"][-1].pretty_print()\n",
+    "\n",
+    "\n",
+    "input_message = HumanMessage(content=\"what's my name?\")\n",
+    "for event in app.stream({\"messages\": [input_message]}, config, stream_mode=\"values\"):\n",
+    "    event[\"messages\"][-1].pretty_print()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2fb0de5b-30ec-42d4-813a-7ad63fe1c367",
+   "metadata": {},
+   "source": [
+    "## Manually deleting messages\n",
+    "\n",
+    "First, we will cover how to manually delete messages. Let's take a look at the current state of the thread:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "8a850529-d038-48f7-b5a2-8d4d2923f83a",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[HumanMessage(content=\"hi! I'm bob\", id='bc1c6dd2-3bb9-4aa9-b7af-3c6af7e173ea'),\n",
+       " AIMessage(content=\"It's nice to meet you, Bob! I'm Claude, an AI assistant created by Anthropic. How can I assist you today?\", response_metadata={'id': 'msg_01XPSAenmSqK8rX2WgPZHfz7', 'model': 'claude-3-haiku-20240307', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 12, 'output_tokens': 32}}, id='run-1c69af09-adb1-412d-9010-2456e5a555fb-0', usage_metadata={'input_tokens': 12, 'output_tokens': 32, 'total_tokens': 44}),\n",
+       " HumanMessage(content=\"what's my name?\", id='f3c71afe-8ce2-4ed0-991e-65021f03b0a5'),\n",
+       " AIMessage(content='Your name is Bob, as you introduced yourself at the beginning of our conversation.', response_metadata={'id': 'msg_01BPZdwsjuMAbC1YAkqawXaF', 'model': 'claude-3-haiku-20240307', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 52, 'output_tokens': 19}}, id='run-b2eb9137-2f4e-446f-95f5-3d5f621a2cf8-0', usage_metadata={'input_tokens': 52, 'output_tokens': 19, 'total_tokens': 71})]"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "messages = app.get_state(config).values[\"messages\"]\n",
+    "messages"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "81be8a0a-1e94-4302-bd84-d1b72e3c501c",
+   "metadata": {},
+   "source": [
+    "We can call `update_state` and pass in the id of the first message. This will delete that message."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "df1a0970-7e64-4170-beef-2855d10eef42",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/Users/vadymbarda/.virtualenvs/langgraph/lib/python3.11/site-packages/langchain_core/_api/beta_decorator.py:87: LangChainBetaWarning: The class `RemoveMessage` is in beta. It is actively being worked on, so the API may change.\n",
+      "  warn_beta(\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "{'configurable': {'thread_id': '2',\n",
+       "  'thread_ts': '1ef42d00-d9ad-6f24-8005-feb089654def'}}"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from langchain_core.messages import RemoveMessage\n",
+    "\n",
+    "app.update_state(config, {\"messages\": RemoveMessage(id=messages[0].id)})"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9c9127ae-0d42-42b8-957f-ea69a5da555f",
+   "metadata": {},
+   "source": [
+    "If we now look at the messages, we can verify that the first one was deleted."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "8bfe4ffa-e170-43bc-aec4-6e36ac620931",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[AIMessage(content=\"It's nice to meet you, Bob! I'm Claude, an AI assistant created by Anthropic. How can I assist you today?\", response_metadata={'id': 'msg_01XPSAenmSqK8rX2WgPZHfz7', 'model': 'claude-3-haiku-20240307', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 12, 'output_tokens': 32}}, id='run-1c69af09-adb1-412d-9010-2456e5a555fb-0', usage_metadata={'input_tokens': 12, 'output_tokens': 32, 'total_tokens': 44}),\n",
+       " HumanMessage(content=\"what's my name?\", id='f3c71afe-8ce2-4ed0-991e-65021f03b0a5'),\n",
+       " AIMessage(content='Your name is Bob, as you introduced yourself at the beginning of our conversation.', response_metadata={'id': 'msg_01BPZdwsjuMAbC1YAkqawXaF', 'model': 'claude-3-haiku-20240307', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 52, 'output_tokens': 19}}, id='run-b2eb9137-2f4e-446f-95f5-3d5f621a2cf8-0', usage_metadata={'input_tokens': 52, 'output_tokens': 19, 'total_tokens': 71})]"
+      ]
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "messages = app.get_state(config).values[\"messages\"]\n",
+    "messages"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ef129a75-4cad-44d7-b532-eb37b0553c0c",
+   "metadata": {},
+   "source": [
+    "## Programmatically deleting messages\n",
+    "\n",
+    "We can also delete messages programmatically from inside the graph. Here we'll modify the graph to delete any old messages (longer than 3 messages ago) at the end of a graph run."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "bb22ede0-e153-4fd0-a4c0-f9af2f7663b1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from langchain_core.messages import RemoveMessage\n",
+    "from langgraph.graph import END\n",
+    "\n",
+    "\n",
+    "def delete_messages(state):\n",
+    "    messages = state[\"messages\"]\n",
+    "    if len(messages) > 3:\n",
+    "        return {\"messages\": [RemoveMessage(id=m.id) for m in messages[:-3]]}\n",
+    "\n",
+    "\n",
+    "# We need to modify the logic to call delete_messages rather than end right away\n",
+    "def should_continue(state: MessagesState) -> Literal[\"action\", \"delete_messages\"]:\n",
+    "    \"\"\"Return the next node to execute.\"\"\"\n",
+    "    last_message = state[\"messages\"][-1]\n",
+    "    # If there is no function call, then we call our delete_messages function\n",
+    "    if not last_message.tool_calls:\n",
+    "        return \"delete_messages\"\n",
+    "    # Otherwise if there is, we continue\n",
+    "    return \"action\"\n",
+    "\n",
+    "\n",
+    "# Define a new graph\n",
+    "workflow = StateGraph(MessagesState)\n",
+    "workflow.add_node(\"agent\", call_model)\n",
+    "workflow.add_node(\"action\", tool_node)\n",
+    "\n",
+    "# This is our new node we're defining\n",
+    "workflow.add_node(delete_messages)\n",
+    "\n",
+    "\n",
+    "workflow.add_edge(START, \"agent\")\n",
+    "workflow.add_conditional_edges(\n",
+    "    \"agent\",\n",
+    "    should_continue,\n",
+    ")\n",
+    "workflow.add_edge(\"action\", \"agent\")\n",
+    "\n",
+    "# This is the new edge we're adding: after we delete messages, we finish\n",
+    "workflow.add_edge(\"delete_messages\", END)\n",
+    "app = workflow.compile(checkpointer=memory)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "52cbdef6-7db7-45a2-8194-de4f8929bd1f",
+   "metadata": {},
+   "source": [
+    "We can now try this out. We can call the graph twice and then check the state"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "3975f34c-c243-40ea-b9d2-424d50a48dc9",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[('human', \"hi! I'm bob\")]\n",
+      "[('human', \"hi! I'm bob\"), ('ai', \"Hello Bob! It's nice to meet you. I'm an AI assistant created by Anthropic. I'm here to help with any questions or tasks you might have. Please let me know how I can assist you.\")]\n",
+      "[('human', \"hi! I'm bob\"), ('ai', \"Hello Bob! It's nice to meet you. I'm an AI assistant created by Anthropic. I'm here to help with any questions or tasks you might have. Please let me know how I can assist you.\"), ('human', \"what's my name?\")]\n",
+      "[('human', \"hi! I'm bob\"), ('ai', \"Hello Bob! It's nice to meet you. I'm an AI assistant created by Anthropic. I'm here to help with any questions or tasks you might have. Please let me know how I can assist you.\"), ('human', \"what's my name?\"), ('ai', 'You said your name is Bob, so that is the name I have for you.')]\n",
+      "[('ai', \"Hello Bob! It's nice to meet you. I'm an AI assistant created by Anthropic. I'm here to help with any questions or tasks you might have. Please let me know how I can assist you.\"), ('human', \"what's my name?\"), ('ai', 'You said your name is Bob, so that is the name I have for you.')]\n"
+     ]
+    }
+   ],
+   "source": [
+    "from langchain_core.messages import HumanMessage\n",
+    "\n",
+    "config = {\"configurable\": {\"thread_id\": \"3\"}}\n",
+    "input_message = HumanMessage(content=\"hi! I'm bob\")\n",
+    "for event in app.stream({\"messages\": [input_message]}, config, stream_mode=\"values\"):\n",
+    "    print([(message.type, message.content) for message in event[\"messages\"]])\n",
+    "\n",
+    "\n",
+    "input_message = HumanMessage(content=\"what's my name?\")\n",
+    "for event in app.stream({\"messages\": [input_message]}, config, stream_mode=\"values\"):\n",
+    "    print([(message.type, message.content) for message in event[\"messages\"]])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "67b2fd2a-14a1-4c47-8632-f8cbb0ba1d35",
+   "metadata": {},
+   "source": [
+    "If we now check the state, we should see that it is only three messages long. This is because we just deleted the earlier messages - otherwise it would be four!"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "a3e15abb-81d8-4072-9f10-61ae0fd61dac",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[AIMessage(content=\"Hello Bob! It's nice to meet you. I'm an AI assistant created by Anthropic. I'm here to help with any questions or tasks you might have. Please let me know how I can assist you.\", response_metadata={'id': 'msg_01XPEgPPbcnz5BbGWUDWTmzG', 'model': 'claude-3-haiku-20240307', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 12, 'output_tokens': 48}}, id='run-eded3820-b6a9-4d66-9210-03ca41787ce6-0', usage_metadata={'input_tokens': 12, 'output_tokens': 48, 'total_tokens': 60}),\n",
+       " HumanMessage(content=\"what's my name?\", id='a0ea2097-3280-402b-92e1-67177b807ae8'),\n",
+       " AIMessage(content='You said your name is Bob, so that is the name I have for you.', response_metadata={'id': 'msg_01JGT62pxhrhN4SykZ57CSjW', 'model': 'claude-3-haiku-20240307', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 68, 'output_tokens': 20}}, id='run-ace3519c-81f8-45fe-a777-91f42d48b3a3-0', usage_metadata={'input_tokens': 68, 'output_tokens': 20, 'total_tokens': 88})]"
+      ]
+     },
+     "execution_count": 11,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "messages = app.get_state(config).values[\"messages\"]\n",
+    "messages"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "359cfeae-d43a-46ee-9069-a1cab9a5720a",
+   "metadata": {},
+   "source": [
+    "Remember, when deleting messages you will want to make sure that the remaining message list is still valid. This message list **may actually not be** - this is because it currently starts with an AI message, which some models do not allow."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4d7222cd-5767-42f0-bc69-10615127eba5",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "langgraph",
+   "language": "python",
+   "name": "langgraph"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

examples/memory/manage-conversation-history.ipynb

@@ -7,7 +7,12 @@
    "source": [
     "# How to manage conversation history\n",
     "\n",
-    "One of the most common use cases for persistence is to use it to keep track of conversation history. This is great - it makes it easy to continue conversations. As conversations get longer and longer, however, this conversation history can build up and take up more and more of the context window. This can often be undesirable as it leads to more expensive and longer calls to the LLM, and potentially ones that error. In this notebook we will discuss a few strategies for how to deal with this."
+    "One of the most common use cases for persistence is to use it to keep track of conversation history. This is great - it makes it easy to continue conversations. As conversations get longer and longer, however, this conversation history can build up and take up more and more of the context window. This can often be undesirable as it leads to more expensive and longer calls to the LLM, and potentially ones that error. In order to prevent this from happening, you need to probably manage the conversation history.\n",
+    "\n",
+    "Note: this guide focuses on how to do this in LangGraph, where you can fully customize how this is done. If you want a more off-the-shelf solution, you can look into functionality provided in LangChain:\n",
+    "\n",
+    "- [How to filter messages](https://python.langchain.com/v0.2/docs/how_to/filter_messages/)\n",
+    "- [How to trim messages](https://python.langchain.com/v0.2/docs/how_to/trim_messages/)"
    ]
   },
   {
@@ -82,6 +87,7 @@
    "id": "4767ef1c-a7cf-41f8-a301-558988cb7ac5",
    "metadata": {},
    "source": [
+    "## Build the agent\n",
     "Let's now build a simple ReAct style agent."
    ]
   },
@@ -343,6 +349,25 @@
     "for event in app.stream({\"messages\": [input_message]}, config, stream_mode=\"values\"):\n",
     "    event[\"messages\"][-1].pretty_print()"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "454102b6-7112-4710-aa08-ba675e8be14c",
+   "metadata": {},
+   "source": [
+    "In the above example we defined the `filter_messages` function ourselves. We also provide off-the-shelf ways to trim and filter messages in LangChain. \n",
+    "\n",
+    "- [How to filter messages](https://python.langchain.com/v0.2/docs/how_to/filter_messages/)\n",
+    "- [How to trim messages](https://python.langchain.com/v0.2/docs/how_to/trim_messages/)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "686861bb-ec32-46f3-b7b3-fdac106f22f6",
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {

examples/node-retries.ipynb

@@ -0,0 +1,120 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# How to add node retry policies\n",
+    "\n",
+    "There are many use cases where you may wish for your node to have a custom retry policy, for example if you are calling an API, querying a database, or calling an LLM, etc. \n",
+    "\n",
+    "In order to configure the retry policty, you have to pass the `retry` parameter to the `add_node` function. The `retry` parameter takes in a `RetryPolicy` named tuple object. Below we instantiate a `RetryPolicy` object with the default parameters:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "RetryPolicy(initial_interval=0.5, backoff_factor=2.0, max_interval=128.0, max_attempts=3, jitter=True, retry_on=<function default_retry_on at 0x1157419e0>)"
+      ]
+     },
+     "execution_count": 15,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from langgraph.pregel import RetryPolicy\n",
+    "\n",
+    "RetryPolicy()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "If you want more information on what each of the parameters does, be sure to read the [reference](https://langchain-ai.github.io/langgraph/reference/graphs/#retrypolicy).\n",
+    "\n",
+    "## Passing a retry policy to a node\n",
+    "\n",
+    "Lastly, we can pass `RetryPolicy` objects when we call the `add_node` function. In the example below we pass two different retry policies to each of our nodes:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import operator\n",
+    "import sqlite3\n",
+    "from typing import Annotated, Sequence, TypedDict\n",
+    "\n",
+    "from langchain_anthropic import ChatAnthropic\n",
+    "from langchain_core.messages import BaseMessage\n",
+    "\n",
+    "from langgraph.graph import END, StateGraph, START\n",
+    "from langchain_community.utilities import SQLDatabase\n",
+    "from langchain_core.messages import AIMessage\n",
+    "\n",
+    "db = SQLDatabase.from_uri(\"sqlite:///:memory:\")\n",
+    "\n",
+    "model = ChatAnthropic(model_name=\"claude-2.1\")\n",
+    "\n",
+    "\n",
+    "class AgentState(TypedDict):\n",
+    "    messages: Annotated[Sequence[BaseMessage], operator.add]\n",
+    "\n",
+    "\n",
+    "def query_database(state):\n",
+    "    query_result = db.run(\"SELECT * FROM Artist LIMIT 10;\")\n",
+    "    return {\"messages\": [AIMessage(content=query_result)]}\n",
+    "\n",
+    "\n",
+    "def call_model(state):\n",
+    "    response = model.invoke(state[\"messages\"])\n",
+    "    return {\"messages\": [response]}\n",
+    "\n",
+    "\n",
+    "# Define a new graph\n",
+    "workflow = StateGraph(AgentState)\n",
+    "workflow.add_node(\n",
+    "    \"query_database\",\n",
+    "    query_database,\n",
+    "    retry=RetryPolicy(retry_on=sqlite3.OperationalError),\n",
+    ")\n",
+    "workflow.add_node(\"model\", call_model, retry=RetryPolicy(max_attempts=5))\n",
+    "workflow.add_edge(START, \"model\")\n",
+    "workflow.add_edge(\"model\", \"query_database\")\n",
+    "workflow.add_edge(\"query_database\", END)\n",
+    "\n",
+    "app = workflow.compile()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "env",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

examples/persistence.ipynb

@@ -38,14 +38,6 @@
     "</div>    "
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "0c3fde0a",
-   "metadata": {},
-   "outputs": [],
-   "source": [""]
-  },
   {
    "cell_type": "markdown",
    "id": "7cbd446a-808f-4394-be92-d45ab818953c",
@@ -62,7 +54,10 @@
    "id": "af4ce0ba-7596-4e5f-8bf8-0b0bd6e62833",
    "metadata": {},
    "outputs": [],
-   "source": ["%%capture --no-stderr\n%pip install --quiet -U langgraph langchain_anthropic"]
+   "source": [
+    "%%capture --no-stderr\n",
+    "%pip install --quiet -U langgraph langchain_anthropic"
+   ]
   },
   {
    "cell_type": "markdown",
@@ -78,7 +73,18 @@
    "id": "c903a1cf-2977-4e2d-ad7d-8b3946821d89",
    "metadata": {},
    "outputs": [],
-   "source": ["import getpass\nimport os\n\n\ndef _set_env(var: str):\n    if not os.environ.get(var):\n        os.environ[var] = getpass.getpass(f\"{var}: \")\n\n\n_set_env(\"ANTHROPIC_API_KEY\")"]
+   "source": [
+    "import getpass\n",
+    "import os\n",
+    "\n",
+    "\n",
+    "def _set_env(var: str):\n",
+    "    if not os.environ.get(var):\n",
+    "        os.environ[var] = getpass.getpass(f\"{var}: \")\n",
+    "\n",
+    "\n",
+    "_set_env(\"ANTHROPIC_API_KEY\")"
+   ]
   },
   {
    "cell_type": "markdown",
@@ -94,7 +100,10 @@
    "id": "95e25aec-7c9f-4a63-b143-225d0e9a79c3",
    "metadata": {},
    "outputs": [],
-   "source": ["os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\n_set_env(\"LANGCHAIN_API_KEY\")"]
+   "source": [
+    "os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\n",
+    "_set_env(\"LANGCHAIN_API_KEY\")"
+   ]
   },
   {
    "cell_type": "markdown",
@@ -112,7 +121,22 @@
    "id": "14619607",
    "metadata": {},
    "outputs": [],
-   "source": ["from typing import Annotated\n\nfrom typing_extensions import TypedDict\n\nfrom langgraph.graph.message import add_messages\n\n# Add messages essentially does this with more\n# robust handling\n# def add_messages(left: list, right: list):\n#     return left + right\n\n\nclass State(TypedDict):\n    messages: Annotated[list, add_messages]"]
+   "source": [
+    "from typing import Annotated\n",
+    "\n",
+    "from typing_extensions import TypedDict\n",
+    "\n",
+    "from langgraph.graph.message import add_messages\n",
+    "\n",
+    "# Add messages essentially does this with more\n",
+    "# robust handling\n",
+    "# def add_messages(left: list, right: list):\n",
+    "#     return left + right\n",
+    "\n",
+    "\n",
+    "class State(TypedDict):\n",
+    "    messages: Annotated[list, add_messages]"
+   ]
   },
   {
    "cell_type": "markdown",
@@ -132,7 +156,19 @@
    "id": "d7ef57dd-5d6e-4ad3-9377-a92201c1310e",
    "metadata": {},
    "outputs": [],
-   "source": ["from langchain_core.tools import tool\n\n\n@tool\ndef search(query: str):\n    \"\"\"Call to surf the web.\"\"\"\n    # This is a placeholder for the actual implementation\n    return [\"The answer to your question lies within.\"]\n\n\ntools = [search]"]
+   "source": [
+    "from langchain_core.tools import tool\n",
+    "\n",
+    "\n",
+    "@tool\n",
+    "def search(query: str):\n",
+    "    \"\"\"Call to surf the web.\"\"\"\n",
+    "    # This is a placeholder for the actual implementation\n",
+    "    return [\"The answer to your question lies within.\"]\n",
+    "\n",
+    "\n",
+    "tools = [search]"
+   ]
   },
   {
    "cell_type": "markdown",
@@ -149,7 +185,11 @@
    "id": "5cf3331e-ccb3-41c8-aeb9-a840a94d41e7",
    "metadata": {},
    "outputs": [],
-   "source": ["from langgraph.prebuilt import ToolNode\n\ntool_node = ToolNode(tools)"]
+   "source": [
+    "from langgraph.prebuilt import ToolNode\n",
+    "\n",
+    "tool_node = ToolNode(tools)"
+   ]
   },
   {
    "cell_type": "markdown",
@@ -179,7 +219,13 @@
    "id": "892b54b9-75f0-4804-9ed0-88b5e5532989",
    "metadata": {},
    "outputs": [],
-   "source": ["from langchain_openai import ChatOpenAI\n\n# We will set streaming=True so that we can stream tokens\n# See the streaming section for more information on this.\nmodel = ChatOpenAI(temperature=0, streaming=True)"]
+   "source": [
+    "from langchain_openai import ChatOpenAI\n",
+    "\n",
+    "# We will set streaming=True so that we can stream tokens\n",
+    "# See the streaming section for more information on this.\n",
+    "model = ChatOpenAI(temperature=0, streaming=True)"
+   ]
   },
   {
    "cell_type": "markdown",
@@ -197,7 +243,9 @@
    "id": "cd3cbae5-d92c-4559-a4aa-44721b80d107",
    "metadata": {},
    "outputs": [],
-   "source": ["bound_model = model.bind_tools(tools)"]
+   "source": [
+    "bound_model = model.bind_tools(tools)"
+   ]
   },
   {
    "cell_type": "markdown",
@@ -232,7 +280,27 @@
    "id": "3b541bb9-900c-40d0-964d-7b5dfee30667",
    "metadata": {},
    "outputs": [],
-   "source": ["# Define the function that determines whether to continue or not\nfrom typing import Literal\n\n\ndef should_continue(state: State) -> Literal[\"action\", \"__end__\"]:\n    \"\"\"Return the next node to execute.\"\"\"\n    last_message = state[\"messages\"][-1]\n    # If there is no function call, then we finish\n    if not last_message.tool_calls:\n        return \"__end__\"\n    # Otherwise if there is, we continue\n    return \"action\"\n\n\n# Define the function that calls the model\ndef call_model(state: State):\n    response = model.invoke(state[\"messages\"])\n    # We return a list, because this will get added to the existing list\n    return {\"messages\": response}"]
+   "source": [
+    "# Define the function that determines whether to continue or not\n",
+    "from typing import Literal\n",
+    "\n",
+    "\n",
+    "def should_continue(state: State) -> Literal[\"action\", \"__end__\"]:\n",
+    "    \"\"\"Return the next node to execute.\"\"\"\n",
+    "    last_message = state[\"messages\"][-1]\n",
+    "    # If there is no function call, then we finish\n",
+    "    if not last_message.tool_calls:\n",
+    "        return \"__end__\"\n",
+    "    # Otherwise if there is, we continue\n",
+    "    return \"action\"\n",
+    "\n",
+    "\n",
+    "# Define the function that calls the model\n",
+    "def call_model(state: State):\n",
+    "    response = model.invoke(state[\"messages\"])\n",
+    "    # We return a list, because this will get added to the existing list\n",
+    "    return {\"messages\": response}"
+   ]
   },
   {
    "cell_type": "markdown",
@@ -248,7 +316,33 @@
    "id": "812b4e70-4956-4415-8880-db48b3dcbad2",
    "metadata": {},
    "outputs": [],
-   "source": ["from langgraph.graph import StateGraph, START\n\n# Define a new graph\nworkflow = StateGraph(State)\n\n# Define the two nodes we will cycle between\nworkflow.add_node(\"agent\", call_model)\nworkflow.add_node(\"action\", tool_node)\n\n# Set the entrypoint as `agent`\n# This means that this node is the first one called\nworkflow.add_edge(START, \"agent\")\n\n# We now add a conditional edge\nworkflow.add_conditional_edges(\n    # First, we define the start node. We use `agent`.\n    # This means these are the edges taken after the `agent` node is called.\n    \"agent\",\n    # Next, we pass in the function that will determine which node is called next.\n    should_continue,\n)\n\n# We now add a normal edge from `tools` to `agent`.\n# This means that after `tools` is called, `agent` node is called next.\nworkflow.add_edge(\"action\", \"agent\")"]
+   "source": [
+    "from langgraph.graph import StateGraph, START\n",
+    "\n",
+    "# Define a new graph\n",
+    "workflow = StateGraph(State)\n",
+    "\n",
+    "# Define the two nodes we will cycle between\n",
+    "workflow.add_node(\"agent\", call_model)\n",
+    "workflow.add_node(\"action\", tool_node)\n",
+    "\n",
+    "# Set the entrypoint as `agent`\n",
+    "# This means that this node is the first one called\n",
+    "workflow.add_edge(START, \"agent\")\n",
+    "\n",
+    "# We now add a conditional edge\n",
+    "workflow.add_conditional_edges(\n",
+    "    # First, we define the start node. We use `agent`.\n",
+    "    # This means these are the edges taken after the `agent` node is called.\n",
+    "    \"agent\",\n",
+    "    # Next, we pass in the function that will determine which node is called next.\n",
+    "    should_continue,\n",
+    ")\n",
+    "\n",
+    "# We now add a normal edge from `tools` to `agent`.\n",
+    "# This means that after `tools` is called, `agent` node is called next.\n",
+    "workflow.add_edge(\"action\", \"agent\")"
+   ]
   },
   {
    "cell_type": "markdown",
@@ -266,7 +360,11 @@
    "id": "6845ed6a-d155-4105-9160-28849877248b",
    "metadata": {},
    "outputs": [],
-   "source": ["from langgraph.checkpoint.sqlite import SqliteSaver\n\nmemory = SqliteSaver.from_conn_string(\":memory:\")"]
+   "source": [
+    "from langgraph.checkpoint.sqlite import SqliteSaver\n",
+    "\n",
+    "memory = SqliteSaver.from_conn_string(\":memory:\")"
+   ]
   },
   {
    "cell_type": "code",
@@ -274,7 +372,25 @@
    "id": "79d29875-8aa8-434c-9f20-1c58346a6249",
    "metadata": {},
    "outputs": [],
-   "source": ["# Finally, we compile it!\n# This compiles it into a LangChain Runnable,\n# meaning you can use it as you would any other runnable\napp = workflow.compile(checkpointer=memory)"]
+   "source": [
+    "# Finally, we compile it!\n",
+    "# This compiles it into a LangChain Runnable,\n",
+    "# meaning you can use it as you would any other runnable\n",
+    "app = workflow.compile(checkpointer=memory)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7654ebcc-2179-41b4-92d1-6666f6f8634f",
+   "metadata": {},
+   "source": [
+    "<div class=\"admonition tip\">\n",
+    "    <p class=\"admonition-title\">Note</p>\n",
+    "    <p>\n",
+    "        If you're using LangGraph Cloud, you <strong>don't need</strong> to pass checkpointer when compiling the graph, since it's done automatically.\n",
+    "    </p>\n",
+    "</div>"
+   ]
   },
   {
    "cell_type": "code",
@@ -293,7 +409,15 @@
      "output_type": "display_data"
     }
    ],
-   "source": ["from IPython.display import Image, display\n\ntry:\n    display(Image(app.get_graph().draw_mermaid_png()))\nexcept Exception:\n    # This requires some extra dependencies and is optional\n    pass"]
+   "source": [
+    "from IPython.display import Image, display\n",
+    "\n",
+    "try:\n",
+    "    display(Image(app.get_graph().draw_mermaid_png()))\n",
+    "except Exception:\n",
+    "    # This requires some extra dependencies and is optional\n",
+    "    pass"
+   ]
   },
   {
    "cell_type": "markdown",
@@ -324,7 +448,14 @@
      ]
     }
    ],
-   "source": ["from langchain_core.messages import HumanMessage\n\nconfig = {\"configurable\": {\"thread_id\": \"2\"}}\ninput_message = HumanMessage(content=\"hi! I'm bob\")\nfor event in app.stream({\"messages\": [input_message]}, config, stream_mode=\"values\"):\n    event[\"messages\"][-1].pretty_print()"]
+   "source": [
+    "from langchain_core.messages import HumanMessage\n",
+    "\n",
+    "config = {\"configurable\": {\"thread_id\": \"2\"}}\n",
+    "input_message = HumanMessage(content=\"hi! I'm bob\")\n",
+    "for event in app.stream({\"messages\": [input_message]}, config, stream_mode=\"values\"):\n",
+    "    event[\"messages\"][-1].pretty_print()"
+   ]
   },
   {
    "cell_type": "code",
@@ -345,7 +476,11 @@
      ]
     }
    ],
-   "source": ["input_message = HumanMessage(content=\"what is my name?\")\nfor event in app.stream({\"messages\": [input_message]}, config, stream_mode=\"values\"):\n    event[\"messages\"][-1].pretty_print()"]
+   "source": [
+    "input_message = HumanMessage(content=\"what is my name?\")\n",
+    "for event in app.stream({\"messages\": [input_message]}, config, stream_mode=\"values\"):\n",
+    "    event[\"messages\"][-1].pretty_print()"
+   ]
   },
   {
    "cell_type": "markdown",
@@ -374,7 +509,15 @@
      ]
     }
    ],
-   "source": ["input_message = HumanMessage(content=\"what is my name?\")\nfor event in app.stream(\n    {\"messages\": [input_message]},\n    {\"configurable\": {\"thread_id\": \"3\"}},\n    stream_mode=\"values\",\n):\n    event[\"messages\"][-1].pretty_print()"]
+   "source": [
+    "input_message = HumanMessage(content=\"what is my name?\")\n",
+    "for event in app.stream(\n",
+    "    {\"messages\": [input_message]},\n",
+    "    {\"configurable\": {\"thread_id\": \"3\"}},\n",
+    "    stream_mode=\"values\",\n",
+    "):\n",
+    "    event[\"messages\"][-1].pretty_print()"
+   ]
   },
   {
    "cell_type": "markdown",
@@ -403,7 +546,15 @@
      ]
     }
    ],
-   "source": ["input_message = HumanMessage(content=\"You forgot??\")\nfor event in app.stream(\n    {\"messages\": [input_message]},\n    {\"configurable\": {\"thread_id\": \"2\"}},\n    stream_mode=\"values\",\n):\n    event[\"messages\"][-1].pretty_print()"]
+   "source": [
+    "input_message = HumanMessage(content=\"You forgot??\")\n",
+    "for event in app.stream(\n",
+    "    {\"messages\": [input_message]},\n",
+    "    {\"configurable\": {\"thread_id\": \"2\"}},\n",
+    "    stream_mode=\"values\",\n",
+    "):\n",
+    "    event[\"messages\"][-1].pretty_print()"
+   ]
   },
   {
    "cell_type": "code",
@@ -411,7 +562,7 @@
    "id": "eb20430f",
    "metadata": {},
    "outputs": [],
-   "source": [""]
+   "source": []
   }
  ],
  "metadata": {
@@ -430,7 +581,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.2"
+   "version": "3.11.9"
   }
  },
  "nbformat": 4,

examples/persistence_postgres.ipynb

@@ -51,6 +51,7 @@
     "    Union,\n",
     "    Tuple,\n",
     "    List,\n",
+    "    Sequence,\n",
     ")\n",
     "\n",
     "import psycopg\n",
@@ -174,6 +175,15 @@
     "        metadata BYTEA NOT NULL,\n",
     "        PRIMARY KEY (thread_id, thread_ts)\n",
     "    );\n",
+    "    CREATE TABLE IF NOT EXISTS writes (\n",
+    "        thread_id TEXT NOT NULL,\n",
+    "        thread_ts TEXT NOT NULL,\n",
+    "        task_id TEXT NOT NULL,\n",
+    "        idx INTEGER NOT NULL,\n",
+    "        channel TEXT NOT NULL,\n",
+    "        value BYTEA,\n",
+    "        PRIMARY KEY (thread_id, thread_ts, task_id, idx)\n",
+    "    );\n",
     "    \"\"\"\n",
     "\n",
     "    @staticmethod\n",
@@ -196,13 +206,13 @@
     "    def drop_tables(connection: psycopg.Connection, /) -> None:\n",
     "        \"\"\"Drop the table for the checkpoint saver.\"\"\"\n",
     "        with connection.cursor() as cur:\n",
-    "            cur.execute(\"DROP TABLE IF EXISTS checkpoints;\")\n",
+    "            cur.execute(\"DROP TABLE IF EXISTS checkpoints, writes;\")\n",
     "\n",
     "    @staticmethod\n",
     "    async def adrop_tables(connection: psycopg.AsyncConnection, /) -> None:\n",
     "        \"\"\"Drop the table for the checkpoint saver.\"\"\"\n",
     "        async with connection.cursor() as cur:\n",
-    "            await cur.execute(\"DROP TABLE IF EXISTS checkpoints;\")\n",
+    "            await cur.execute(\"DROP TABLE IF EXISTS checkpoints, writes;\")\n",
     "\n",
     "    UPSERT_CHECKPOINT_QUERY = \"\"\"\n",
     "    INSERT INTO checkpoints \n",
@@ -239,7 +249,7 @@
     "                    self.UPSERT_CHECKPOINT_QUERY,\n",
     "                    (\n",
     "                        thread_id,\n",
-    "                        checkpoint[\"ts\"],\n",
+    "                        checkpoint[\"id\"],\n",
     "                        parent_ts if parent_ts else None,\n",
     "                        self.serde.dumps(checkpoint),\n",
     "                        self.serde.dumps(metadata),\n",
@@ -249,7 +259,7 @@
     "        return {\n",
     "            \"configurable\": {\n",
     "                \"thread_id\": thread_id,\n",
-    "                \"thread_ts\": checkpoint[\"ts\"],\n",
+    "                \"thread_ts\": checkpoint[\"id\"],\n",
     "            },\n",
     "        }\n",
     "\n",
@@ -278,7 +288,7 @@
     "                    self.UPSERT_CHECKPOINT_QUERY,\n",
     "                    (\n",
     "                        thread_id,\n",
-    "                        checkpoint[\"ts\"],\n",
+    "                        checkpoint[\"id\"],\n",
     "                        parent_ts if parent_ts else None,\n",
     "                        self.serde.dumps(checkpoint),\n",
     "                        self.serde.dumps(metadata),\n",
@@ -288,10 +298,67 @@
     "        return {\n",
     "            \"configurable\": {\n",
     "                \"thread_id\": thread_id,\n",
-    "                \"thread_ts\": checkpoint[\"ts\"],\n",
+    "                \"thread_ts\": checkpoint[\"id\"],\n",
     "            },\n",
     "        }\n",
     "\n",
+    "    UPSERT_WRITES_QUERY = \"\"\"\n",
+    "    INSERT INTO writes\n",
+    "        (thread_id, thread_ts, task_id, idx, channel, value)\n",
+    "    VALUES\n",
+    "        (%s, %s, %s, %s, %s, %s)\n",
+    "    ON CONFLICT (thread_id, thread_ts, task_id, idx)\n",
+    "    DO UPDATE SET value = EXCLUDED.value;\n",
+    "    \"\"\"\n",
+    "\n",
+    "    def put_writes(\n",
+    "        self,\n",
+    "        config: RunnableConfig,\n",
+    "        writes: Sequence[Tuple[str, Any]],\n",
+    "        task_id: str,\n",
+    "    ) -> None:\n",
+    "        with self._get_sync_connection() as conn:\n",
+    "            with conn.cursor() as cur:\n",
+    "                cur.executemany(\n",
+    "                    self.UPSERT_WRITES_QUERY,\n",
+    "                    [\n",
+    "                        (\n",
+    "                            str(config[\"configurable\"][\"thread_id\"]),\n",
+    "                            str(config[\"configurable\"][\"thread_ts\"]),\n",
+    "                            task_id,\n",
+    "                            idx,\n",
+    "                            channel,\n",
+    "                            self.serde.dumps(value),\n",
+    "                        )\n",
+    "                        for idx, (channel, value) in enumerate(writes)\n",
+    "                    ],\n",
+    "                )\n",
+    "            conn.commit()\n",
+    "\n",
+    "    async def aput_writes(\n",
+    "        self,\n",
+    "        config: RunnableConfig,\n",
+    "        writes: Sequence[Tuple[str, Any]],\n",
+    "        task_id: str,\n",
+    "    ) -> None:\n",
+    "        async with self._get_async_connection() as conn:\n",
+    "            async with conn.cursor() as cur:\n",
+    "                await cur.executemany(\n",
+    "                    self.UPSERT_WRITES_QUERY,\n",
+    "                    [\n",
+    "                        (\n",
+    "                            str(config[\"configurable\"][\"thread_id\"]),\n",
+    "                            str(config[\"configurable\"][\"thread_ts\"]),\n",
+    "                            task_id,\n",
+    "                            idx,\n",
+    "                            channel,\n",
+    "                            self.serde.dumps(value),\n",
+    "                        )\n",
+    "                        for idx, (channel, value) in enumerate(writes)\n",
+    "                    ],\n",
+    "                )\n",
+    "            await conn.commit()\n",
+    "\n",
     "    LIST_CHECKPOINTS_QUERY_STR = \"\"\"\n",
     "    SELECT checkpoint, metadata, thread_ts, parent_ts\n",
     "    FROM checkpoints\n",
@@ -404,6 +471,7 @@
     "        thread_ts = config[\"configurable\"].get(\"thread_ts\")\n",
     "        with self._get_sync_connection() as conn:\n",
     "            with conn.cursor() as cur:\n",
+    "                # find the latest checkpoint for the thread_id\n",
     "                if thread_ts:\n",
     "                    cur.execute(\n",
     "                        self.GET_CHECKPOINT_BY_TS_QUERY,\n",
@@ -412,22 +480,6 @@
     "                            \"thread_ts\": thread_ts,\n",
     "                        },\n",
     "                    )\n",
-    "                    value = cur.fetchone()\n",
-    "                    if value:\n",
-    "                        checkpoint, metadata, thread_ts, parent_ts = value\n",
-    "                    return CheckpointTuple(\n",
-    "                        config=config,\n",
-    "                        checkpoint=self.serde.loads(checkpoint),\n",
-    "                        metadata=self.serde.loads(metadata),\n",
-    "                        parent_config={\n",
-    "                            \"configurable\": {\n",
-    "                                \"thread_id\": thread_id,\n",
-    "                                \"thread_ts\": thread_ts,\n",
-    "                            }\n",
-    "                        }\n",
-    "                        if thread_ts\n",
-    "                        else None,\n",
-    "                    )\n",
     "                else:\n",
     "                    cur.execute(\n",
     "                        self.GET_CHECKPOINT_QUERY,\n",
@@ -435,28 +487,44 @@
     "                            \"thread_id\": thread_id,\n",
     "                        },\n",
     "                    )\n",
-    "                    value = cur.fetchone()\n",
-    "                    if value:\n",
-    "                        checkpoint, metadata, thread_ts, parent_ts = value\n",
-    "                        return CheckpointTuple(\n",
-    "                            config={\n",
-    "                                \"configurable\": {\n",
-    "                                    \"thread_id\": thread_id,\n",
-    "                                    \"thread_ts\": thread_ts,\n",
-    "                                }\n",
-    "                            },\n",
-    "                            checkpoint=self.serde.loads(checkpoint),\n",
-    "                            metadata=self.serde.loads(metadata),\n",
-    "                            parent_config={\n",
-    "                                \"configurable\": {\n",
-    "                                    \"thread_id\": thread_id,\n",
-    "                                    \"thread_ts\": parent_ts,\n",
-    "                                }\n",
+    "\n",
+    "                # if a checkpoint is found, return it\n",
+    "                if value := cur.fetchone():\n",
+    "                    checkpoint, metadata, thread_ts, parent_ts = value\n",
+    "                    if not config[\"configurable\"].get(\"thread_ts\"):\n",
+    "                        config = {\n",
+    "                            \"configurable\": {\n",
+    "                                \"thread_id\": thread_id,\n",
+    "                                \"thread_ts\": thread_ts,\n",
     "                            }\n",
-    "                            if parent_ts\n",
-    "                            else None,\n",
-    "                        )\n",
-    "        return None\n",
+    "                        }\n",
+    "\n",
+    "                    # find any pending writes\n",
+    "                    cur.execute(\n",
+    "                        \"SELECT task_id, channel, value FROM writes WHERE thread_id = %(thread_id)s AND thread_ts = %(thread_ts)s\",\n",
+    "                        {\n",
+    "                            \"thread_id\": thread_id,\n",
+    "                            \"thread_ts\": thread_ts,\n",
+    "                        },\n",
+    "                    )\n",
+    "                    # deserialize the checkpoint and metadata\n",
+    "                    return CheckpointTuple(\n",
+    "                        config=config,\n",
+    "                        checkpoint=self.serde.loads(checkpoint),\n",
+    "                        metadata=self.serde.loads(metadata),\n",
+    "                        parent_config={\n",
+    "                            \"configurable\": {\n",
+    "                                \"thread_id\": thread_id,\n",
+    "                                \"thread_ts\": parent_ts,\n",
+    "                            }\n",
+    "                        }\n",
+    "                        if parent_ts\n",
+    "                        else None,\n",
+    "                        pending_writes=[\n",
+    "                            (task_id, channel, self.serde.loads(value))\n",
+    "                            for task_id, channel, value in cur\n",
+    "                        ],\n",
+    "                    )\n",
     "\n",
     "    async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:\n",
     "        \"\"\"Get the checkpoint tuple for the given configuration.\n",
@@ -474,6 +542,7 @@
     "        thread_ts = config[\"configurable\"].get(\"thread_ts\")\n",
     "        async with self._get_async_connection() as conn:\n",
     "            async with conn.cursor() as cur:\n",
+    "                # find the latest checkpoint for the thread_id\n",
     "                if thread_ts:\n",
     "                    await cur.execute(\n",
     "                        self.GET_CHECKPOINT_BY_TS_QUERY,\n",
@@ -482,22 +551,6 @@
     "                            \"thread_ts\": thread_ts,\n",
     "                        },\n",
     "                    )\n",
-    "                    value = await cur.fetchone()\n",
-    "                    if value:\n",
-    "                        checkpoint, metadata, thread_ts, parent_ts = value\n",
-    "                    return CheckpointTuple(\n",
-    "                        config=config,\n",
-    "                        checkpoint=self.serde.loads(checkpoint),\n",
-    "                        metadata=self.serde.loads(metadata),\n",
-    "                        parent_config={\n",
-    "                            \"configurable\": {\n",
-    "                                \"thread_id\": thread_id,\n",
-    "                                \"thread_ts\": thread_ts,\n",
-    "                            }\n",
-    "                        }\n",
-    "                        if thread_ts\n",
-    "                        else None,\n",
-    "                    )\n",
     "                else:\n",
     "                    await cur.execute(\n",
     "                        self.GET_CHECKPOINT_QUERY,\n",
@@ -505,29 +558,43 @@
     "                            \"thread_id\": thread_id,\n",
     "                        },\n",
     "                    )\n",
-    "                    value = await cur.fetchone()\n",
-    "                    if value:\n",
-    "                        checkpoint, metadata, thread_ts, parent_ts = value\n",
-    "                        return CheckpointTuple(\n",
-    "                            config={\n",
-    "                                \"configurable\": {\n",
-    "                                    \"thread_id\": thread_id,\n",
-    "                                    \"thread_ts\": thread_ts,\n",
-    "                                }\n",
-    "                            },\n",
-    "                            checkpoint=self.serde.loads(checkpoint),\n",
-    "                            metadata=self.serde.loads(metadata),\n",
-    "                            parent_config={\n",
-    "                                \"configurable\": {\n",
-    "                                    \"thread_id\": thread_id,\n",
-    "                                    \"thread_ts\": parent_ts,\n",
-    "                                }\n",
+    "                # if a checkpoint is found, return it\n",
+    "                if value := await cur.fetchone():\n",
+    "                    checkpoint, metadata, thread_ts, parent_ts = value\n",
+    "                    if not config[\"configurable\"].get(\"thread_ts\"):\n",
+    "                        config = {\n",
+    "                            \"configurable\": {\n",
+    "                                \"thread_id\": thread_id,\n",
+    "                                \"thread_ts\": thread_ts,\n",
     "                            }\n",
-    "                            if parent_ts\n",
-    "                            else None,\n",
-    "                        )\n",
+    "                        }\n",
     "\n",
-    "        return None\n",
+    "                    # find any pending writes\n",
+    "                    await cur.execute(\n",
+    "                        \"SELECT task_id, channel, value FROM writes WHERE thread_id = %(thread_id)s AND thread_ts = %(thread_ts)s\",\n",
+    "                        {\n",
+    "                            \"thread_id\": thread_id,\n",
+    "                            \"thread_ts\": thread_ts,\n",
+    "                        },\n",
+    "                    )\n",
+    "                    # deserialize the checkpoint and metadata\n",
+    "                    return CheckpointTuple(\n",
+    "                        config=config,\n",
+    "                        checkpoint=self.serde.loads(checkpoint),\n",
+    "                        metadata=self.serde.loads(metadata),\n",
+    "                        parent_config={\n",
+    "                            \"configurable\": {\n",
+    "                                \"thread_id\": thread_id,\n",
+    "                                \"thread_ts\": parent_ts,\n",
+    "                            }\n",
+    "                        }\n",
+    "                        if parent_ts\n",
+    "                        else None,\n",
+    "                        pending_writes=[\n",
+    "                            (task_id, channel, self.serde.loads(value))\n",
+    "                            async for task_id, channel, value in cur\n",
+    "                        ],\n",
+    "                    )\n",
     "\n",
     "    def _search_where(\n",
     "        self,\n",
@@ -701,10 +768,10 @@
     {
      "data": {
       "text/plain": [
-       "{'messages': [HumanMessage(content=\"what's the weather in sf\", id='8481d24b-d505-40fc-8c55-347277d50f2a'),\n",
-       "  AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_cGApaEshW0YNp5WsirAcFTBv', 'function': {'arguments': '{\"city\":\"sf\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 14, 'prompt_tokens': 57, 'total_tokens': 71}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_ce0793330f', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-bf85107a-02a4-48d1-83ba-5448498f5320-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'sf'}, 'id': 'call_cGApaEshW0YNp5WsirAcFTBv'}], usage_metadata={'input_tokens': 57, 'output_tokens': 14, 'total_tokens': 71}),\n",
-       "  ToolMessage(content=\"It's always sunny in sf\", name='get_weather', id='f685da3d-b196-44a1-9fe5-c7ce4d3d7d1d', tool_call_id='call_cGApaEshW0YNp5WsirAcFTBv'),\n",
-       "  AIMessage(content='The weather in San Francisco is currently sunny.', response_metadata={'token_usage': {'completion_tokens': 10, 'prompt_tokens': 84, 'total_tokens': 94}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_4008e3b719', 'finish_reason': 'stop', 'logprobs': None}, id='run-0d681a28-8116-4cb2-a25f-e4f7949481eb-0', usage_metadata={'input_tokens': 84, 'output_tokens': 10, 'total_tokens': 94})]}"
+       "{'messages': [HumanMessage(content=\"what's the weather in sf\", id='bc87fac7-1da1-4818-a43b-6ba7c9b9b3e4'),\n",
+       "  AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_MjkmibJlXeuNchL6B8qpIjOW', 'function': {'arguments': '{\"city\":\"sf\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 14, 'prompt_tokens': 57, 'total_tokens': 71}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_d33f7b429e', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-b9de0cab-f310-4f74-897e-97014072c001-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'sf'}, 'id': 'call_MjkmibJlXeuNchL6B8qpIjOW', 'type': 'tool_call'}], usage_metadata={'input_tokens': 57, 'output_tokens': 14, 'total_tokens': 71}),\n",
+       "  ToolMessage(content=\"It's always sunny in sf\", name='get_weather', id='8d8f9596-a683-4644-a898-1e303b5a01ea', tool_call_id='call_MjkmibJlXeuNchL6B8qpIjOW'),\n",
+       "  AIMessage(content='The weather in San Francisco is currently sunny.', response_metadata={'token_usage': {'completion_tokens': 10, 'prompt_tokens': 84, 'total_tokens': 94}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_dd932ca5d1', 'finish_reason': 'stop', 'logprobs': None}, id='run-4b5282a3-e7a6-42ee-ad0f-e6013a745a88-0', usage_metadata={'input_tokens': 84, 'output_tokens': 10, 'total_tokens': 94})]}"
       ]
      },
      "execution_count": 8,
@@ -726,12 +793,12 @@
      "data": {
       "text/plain": [
        "{'v': 1,\n",
-       " 'ts': '2024-06-27T01:06:18.970784+00:00',\n",
-       " 'id': '1ef34217-5fac-63f8-8003-9f1ea087dd14',\n",
-       " 'channel_values': {'messages': [HumanMessage(content=\"what's the weather in sf\", id='8481d24b-d505-40fc-8c55-347277d50f2a'),\n",
-       "   AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_cGApaEshW0YNp5WsirAcFTBv', 'function': {'arguments': '{\"city\":\"sf\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 14, 'prompt_tokens': 57, 'total_tokens': 71}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_ce0793330f', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-bf85107a-02a4-48d1-83ba-5448498f5320-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'sf'}, 'id': 'call_cGApaEshW0YNp5WsirAcFTBv'}], usage_metadata={'input_tokens': 57, 'output_tokens': 14, 'total_tokens': 71}),\n",
-       "   ToolMessage(content=\"It's always sunny in sf\", name='get_weather', id='f685da3d-b196-44a1-9fe5-c7ce4d3d7d1d', tool_call_id='call_cGApaEshW0YNp5WsirAcFTBv'),\n",
-       "   AIMessage(content='The weather in San Francisco is currently sunny.', response_metadata={'token_usage': {'completion_tokens': 10, 'prompt_tokens': 84, 'total_tokens': 94}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_4008e3b719', 'finish_reason': 'stop', 'logprobs': None}, id='run-0d681a28-8116-4cb2-a25f-e4f7949481eb-0', usage_metadata={'input_tokens': 84, 'output_tokens': 10, 'total_tokens': 94})],\n",
+       " 'ts': '2024-07-12T15:21:51.891852+00:00',\n",
+       " 'id': '1ef40627-6fb2-6962-8003-b74d816658c5',\n",
+       " 'channel_values': {'messages': [HumanMessage(content=\"what's the weather in sf\", id='bc87fac7-1da1-4818-a43b-6ba7c9b9b3e4'),\n",
+       "   AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_MjkmibJlXeuNchL6B8qpIjOW', 'function': {'arguments': '{\"city\":\"sf\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 14, 'prompt_tokens': 57, 'total_tokens': 71}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_d33f7b429e', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-b9de0cab-f310-4f74-897e-97014072c001-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'sf'}, 'id': 'call_MjkmibJlXeuNchL6B8qpIjOW', 'type': 'tool_call'}], usage_metadata={'input_tokens': 57, 'output_tokens': 14, 'total_tokens': 71}),\n",
+       "   ToolMessage(content=\"It's always sunny in sf\", name='get_weather', id='8d8f9596-a683-4644-a898-1e303b5a01ea', tool_call_id='call_MjkmibJlXeuNchL6B8qpIjOW'),\n",
+       "   AIMessage(content='The weather in San Francisco is currently sunny.', response_metadata={'token_usage': {'completion_tokens': 10, 'prompt_tokens': 84, 'total_tokens': 94}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_dd932ca5d1', 'finish_reason': 'stop', 'logprobs': None}, id='run-4b5282a3-e7a6-42ee-ad0f-e6013a745a88-0', usage_metadata={'input_tokens': 84, 'output_tokens': 10, 'total_tokens': 94})],\n",
        "  'agent': 'agent'},\n",
        " 'channel_versions': {'__start__': 2,\n",
        "  'messages': 5,\n",
@@ -790,7 +857,7 @@
     {
      "data": {
       "text/plain": [
-       "CheckpointTuple(config={'configurable': {'thread_id': '2', 'thread_ts': '2024-06-27T01:06:20.104263+00:00'}}, checkpoint={'v': 1, 'ts': '2024-06-27T01:06:20.104263+00:00', 'id': '1ef34217-6a7b-687e-8003-ab84b18fa9dc', 'channel_values': {'messages': [HumanMessage(content=\"what's the weather in sf\", id='3ef79bb9-4249-419d-b142-8d18de63c2eb'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_zqfYbNrD30Qa9YHFSI7nJVJN', 'function': {'arguments': '{\"city\":\"sf\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 14, 'prompt_tokens': 57, 'total_tokens': 71}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_ce0793330f', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-5dd2d9d1-e291-43e4-86cd-96a50a546f48-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'sf'}, 'id': 'call_zqfYbNrD30Qa9YHFSI7nJVJN'}], usage_metadata={'input_tokens': 57, 'output_tokens': 14, 'total_tokens': 71}), ToolMessage(content=\"It's always sunny in sf\", name='get_weather', id='f3eadc9a-9340-41d3-a0fb-521a78e6c438', tool_call_id='call_zqfYbNrD30Qa9YHFSI7nJVJN'), AIMessage(content='The weather in San Francisco is sunny! 🌞', response_metadata={'token_usage': {'completion_tokens': 11, 'prompt_tokens': 84, 'total_tokens': 95}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_ce0793330f', 'finish_reason': 'stop', 'logprobs': None}, id='run-f897e17e-8f0b-4140-be91-35f351609d82-0', usage_metadata={'input_tokens': 84, 'output_tokens': 11, 'total_tokens': 95})], 'agent': 'agent'}, 'channel_versions': {'__start__': 2, 'messages': 5, 'start:agent': 3, 'agent': 5, 'branch:agent:should_continue:tools': 4, 'tools': 5}, 'versions_seen': {'__start__': {'__start__': 1}, 'agent': {'start:agent': 3, 'tools': 4}, 'tools': {'branch:agent:should_continue:tools': 3}}, 'pending_sends': []}, metadata={'source': 'loop', 'step': 3, 'writes': {'agent': {'messages': [AIMessage(content='The weather in San Francisco is sunny! 🌞', response_metadata={'token_usage': {'completion_tokens': 11, 'prompt_tokens': 84, 'total_tokens': 95}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_ce0793330f', 'finish_reason': 'stop', 'logprobs': None}, id='run-f897e17e-8f0b-4140-be91-35f351609d82-0', usage_metadata={'input_tokens': 84, 'output_tokens': 11, 'total_tokens': 95})]}}}, parent_config={'configurable': {'thread_id': '2', 'thread_ts': '1ef34217-659f-62ec-8002-7f9b6b4c2a2f'}})"
+       "CheckpointTuple(config={'configurable': {'thread_id': '2', 'thread_ts': '1ef40627-7d58-6422-8003-de6e83a8c293'}}, checkpoint={'v': 1, 'ts': '2024-07-12T15:21:53.322868+00:00', 'id': '1ef40627-7d58-6422-8003-de6e83a8c293', 'channel_values': {'messages': [HumanMessage(content=\"what's the weather in sf\", id='8d0209ed-a8c2-42ae-8e77-cc71a9cca29d'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_BO4zHHp0JkEWtrtaEqFHkDjK', 'function': {'arguments': '{\"city\":\"sf\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 14, 'prompt_tokens': 57, 'total_tokens': 71}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_d33f7b429e', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-1f76b350-6a33-4de7-9276-59725b1ac101-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'sf'}, 'id': 'call_BO4zHHp0JkEWtrtaEqFHkDjK', 'type': 'tool_call'}], usage_metadata={'input_tokens': 57, 'output_tokens': 14, 'total_tokens': 71}), ToolMessage(content=\"It's always sunny in sf\", name='get_weather', id='c1bb1a24-62a8-4b43-b90e-b00899c112a8', tool_call_id='call_BO4zHHp0JkEWtrtaEqFHkDjK'), AIMessage(content='The weather in San Francisco is currently sunny.', response_metadata={'token_usage': {'completion_tokens': 10, 'prompt_tokens': 84, 'total_tokens': 94}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_dd932ca5d1', 'finish_reason': 'stop', 'logprobs': None}, id='run-7576d437-4938-48b9-b2cf-e4809d92742d-0', usage_metadata={'input_tokens': 84, 'output_tokens': 10, 'total_tokens': 94})], 'agent': 'agent'}, 'channel_versions': {'__start__': 2, 'messages': 5, 'start:agent': 3, 'agent': 5, 'branch:agent:should_continue:tools': 4, 'tools': 5}, 'versions_seen': {'__start__': {'__start__': 1}, 'agent': {'start:agent': 3, 'tools': 4}, 'tools': {'branch:agent:should_continue:tools': 3}}, 'pending_sends': []}, metadata={'source': 'loop', 'step': 3, 'writes': {'agent': {'messages': [AIMessage(content='The weather in San Francisco is currently sunny.', response_metadata={'token_usage': {'completion_tokens': 10, 'prompt_tokens': 84, 'total_tokens': 94}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_dd932ca5d1', 'finish_reason': 'stop', 'logprobs': None}, id='run-7576d437-4938-48b9-b2cf-e4809d92742d-0', usage_metadata={'input_tokens': 84, 'output_tokens': 10, 'total_tokens': 94})]}}}, parent_config={'configurable': {'thread_id': '2', 'thread_ts': '1ef40627-775a-6746-8002-a3967bf0eae6'}}, pending_writes=[])"
       ]
      },
      "execution_count": 11,
@@ -879,7 +946,7 @@
     {
      "data": {
       "text/plain": [
-       "CheckpointTuple(config={'configurable': {'thread_id': '3', 'thread_ts': '2024-06-27T01:06:21.430879+00:00'}}, checkpoint={'v': 1, 'ts': '2024-06-27T01:06:21.430879+00:00', 'id': '1ef34217-7722-6438-8003-0a35c0f65f25', 'channel_values': {'messages': [HumanMessage(content=\"what's the weather in nyc\", id='cb5a1b8c-f329-4bee-857a-5d2c3d020147'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_gPUiz98ETlDNXttk98w6iThc', 'function': {'arguments': '{\"city\":\"nyc\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 15, 'prompt_tokens': 58, 'total_tokens': 73}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_ce0793330f', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-2d35ac1c-0a99-4538-a8fc-4c182d9df69a-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'nyc'}, 'id': 'call_gPUiz98ETlDNXttk98w6iThc'}], usage_metadata={'input_tokens': 58, 'output_tokens': 15, 'total_tokens': 73}), ToolMessage(content='It might be cloudy in nyc', name='get_weather', id='f5f80d01-92e8-4961-b2b0-3b4d190c4d26', tool_call_id='call_gPUiz98ETlDNXttk98w6iThc'), AIMessage(content='The weather in NYC might be cloudy.', response_metadata={'token_usage': {'completion_tokens': 9, 'prompt_tokens': 88, 'total_tokens': 97}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_d576307f90', 'finish_reason': 'stop', 'logprobs': None}, id='run-aea6306e-f6cc-4f87-aaf4-b0dcfbdf9684-0', usage_metadata={'input_tokens': 88, 'output_tokens': 9, 'total_tokens': 97})], 'agent': 'agent'}, 'channel_versions': {'__start__': 2, 'messages': 5, 'start:agent': 3, 'agent': 5, 'branch:agent:should_continue:tools': 4, 'tools': 5}, 'versions_seen': {'__start__': {'__start__': 1}, 'agent': {'start:agent': 3, 'tools': 4}, 'tools': {'branch:agent:should_continue:tools': 3}}, 'pending_sends': []}, metadata={'source': 'loop', 'step': 3, 'writes': {'agent': {'messages': [AIMessage(content='The weather in NYC might be cloudy.', response_metadata={'token_usage': {'completion_tokens': 9, 'prompt_tokens': 88, 'total_tokens': 97}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_d576307f90', 'finish_reason': 'stop', 'logprobs': None}, id='run-aea6306e-f6cc-4f87-aaf4-b0dcfbdf9684-0', usage_metadata={'input_tokens': 88, 'output_tokens': 9, 'total_tokens': 97})]}}}, parent_config={'configurable': {'thread_id': '3', 'thread_ts': '1ef34217-7099-68be-8002-1a2b75cf9196'}})"
+       "CheckpointTuple(config={'configurable': {'thread_id': '3', 'thread_ts': '1ef40627-8b0e-6b02-8003-68a7a04ea6a5'}}, checkpoint={'v': 1, 'ts': '2024-07-12T15:21:54.760751+00:00', 'id': '1ef40627-8b0e-6b02-8003-68a7a04ea6a5', 'channel_values': {'messages': [HumanMessage(content=\"what's the weather in nyc\", id='108ac72d-f658-4ae0-af57-af481adc8aa5'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_XY5TyZEwF5nbdNTWjjiqGtdS', 'function': {'arguments': '{\"city\":\"nyc\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 15, 'prompt_tokens': 58, 'total_tokens': 73}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_d33f7b429e', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-454e2142-6f18-4676-ac4b-91f89ea7a6d4-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'nyc'}, 'id': 'call_XY5TyZEwF5nbdNTWjjiqGtdS', 'type': 'tool_call'}], usage_metadata={'input_tokens': 58, 'output_tokens': 15, 'total_tokens': 73}), ToolMessage(content='It might be cloudy in nyc', name='get_weather', id='2d86514f-b8f0-439b-ab94-68c731309c63', tool_call_id='call_XY5TyZEwF5nbdNTWjjiqGtdS'), AIMessage(content='The weather in NYC might be cloudy.', response_metadata={'token_usage': {'completion_tokens': 9, 'prompt_tokens': 88, 'total_tokens': 97}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_d33f7b429e', 'finish_reason': 'stop', 'logprobs': None}, id='run-90ae3622-b480-4964-b689-9c1a572112f1-0', usage_metadata={'input_tokens': 88, 'output_tokens': 9, 'total_tokens': 97})], 'agent': 'agent'}, 'channel_versions': {'__start__': 2, 'messages': 5, 'start:agent': 3, 'agent': 5, 'branch:agent:should_continue:tools': 4, 'tools': 5}, 'versions_seen': {'__start__': {'__start__': 1}, 'agent': {'start:agent': 3, 'tools': 4}, 'tools': {'branch:agent:should_continue:tools': 3}}, 'pending_sends': []}, metadata={'source': 'loop', 'step': 3, 'writes': {'agent': {'messages': [AIMessage(content='The weather in NYC might be cloudy.', response_metadata={'token_usage': {'completion_tokens': 9, 'prompt_tokens': 88, 'total_tokens': 97}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_d33f7b429e', 'finish_reason': 'stop', 'logprobs': None}, id='run-90ae3622-b480-4964-b689-9c1a572112f1-0', usage_metadata={'input_tokens': 88, 'output_tokens': 9, 'total_tokens': 97})]}}}, parent_config={'configurable': {'thread_id': '3', 'thread_ts': '1ef40627-860c-63d4-8002-49a92ae87052'}}, pending_writes=[])"
       ]
      },
      "execution_count": 15,
@@ -927,11 +994,11 @@
     {
      "data": {
       "text/plain": [
-       "[CheckpointTuple(config={'configurable': {'thread_id': '4', 'thread_ts': '2024-06-27T01:06:22.556635+00:00'}}, checkpoint={'v': 1, 'ts': '2024-06-27T01:06:22.556635+00:00', 'id': '1ef34217-81de-6c5a-8003-04d6ed02b672', 'channel_values': {'messages': [HumanMessage(content=\"what's the weather in nyc\", id='9df9dda1-5e80-4938-bf24-9207738d9f5f'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_TJWuPJovXruJVafVcI540OfS', 'function': {'arguments': '{\"city\":\"nyc\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 15, 'prompt_tokens': 58, 'total_tokens': 73}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_4008e3b719', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-5ba0db1c-45e7-4389-975f-489135ff50b4-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'nyc'}, 'id': 'call_TJWuPJovXruJVafVcI540OfS'}], usage_metadata={'input_tokens': 58, 'output_tokens': 15, 'total_tokens': 73}), ToolMessage(content='It might be cloudy in nyc', name='get_weather', id='04dd3f8a-c063-4885-bfcd-901aaeb63797', tool_call_id='call_TJWuPJovXruJVafVcI540OfS'), AIMessage(content='The weather in NYC might be cloudy.', response_metadata={'token_usage': {'completion_tokens': 9, 'prompt_tokens': 88, 'total_tokens': 97}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_d576307f90', 'finish_reason': 'stop', 'logprobs': None}, id='run-de373cba-dc06-4603-bed1-ba4534a2a45a-0', usage_metadata={'input_tokens': 88, 'output_tokens': 9, 'total_tokens': 97})], 'agent': 'agent'}, 'channel_versions': {'__start__': 2, 'messages': 5, 'start:agent': 3, 'agent': 5, 'branch:agent:should_continue:tools': 4, 'tools': 5}, 'versions_seen': {'__start__': {'__start__': 1}, 'agent': {'start:agent': 3, 'tools': 4}, 'tools': {'branch:agent:should_continue:tools': 3}}, 'pending_sends': []}, metadata={'source': 'loop', 'step': 3, 'writes': {'agent': {'messages': [AIMessage(content='The weather in NYC might be cloudy.', response_metadata={'token_usage': {'completion_tokens': 9, 'prompt_tokens': 88, 'total_tokens': 97}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_d576307f90', 'finish_reason': 'stop', 'logprobs': None}, id='run-de373cba-dc06-4603-bed1-ba4534a2a45a-0', usage_metadata={'input_tokens': 88, 'output_tokens': 9, 'total_tokens': 97})]}}}, parent_config={'configurable': {'thread_id': '4', 'thread_ts': '2024-06-27T01:06:22.556635+00:00'}}),\n",
-       " CheckpointTuple(config={'configurable': {'thread_id': '4', 'thread_ts': '2024-06-27T01:06:22.010058+00:00'}}, checkpoint={'v': 1, 'ts': '2024-06-27T01:06:22.010058+00:00', 'id': '1ef34217-7ca8-64ca-8002-c710e3b07a9c', 'channel_values': {'messages': [HumanMessage(content=\"what's the weather in nyc\", id='9df9dda1-5e80-4938-bf24-9207738d9f5f'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_TJWuPJovXruJVafVcI540OfS', 'function': {'arguments': '{\"city\":\"nyc\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 15, 'prompt_tokens': 58, 'total_tokens': 73}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_4008e3b719', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-5ba0db1c-45e7-4389-975f-489135ff50b4-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'nyc'}, 'id': 'call_TJWuPJovXruJVafVcI540OfS'}], usage_metadata={'input_tokens': 58, 'output_tokens': 15, 'total_tokens': 73}), ToolMessage(content='It might be cloudy in nyc', name='get_weather', id='04dd3f8a-c063-4885-bfcd-901aaeb63797', tool_call_id='call_TJWuPJovXruJVafVcI540OfS')], 'tools': 'tools'}, 'channel_versions': {'__start__': 2, 'messages': 4, 'start:agent': 3, 'agent': 4, 'branch:agent:should_continue:tools': 4, 'tools': 4}, 'versions_seen': {'__start__': {'__start__': 1}, 'agent': {'start:agent': 2}, 'tools': {'branch:agent:should_continue:tools': 3}}, 'pending_sends': []}, metadata={'source': 'loop', 'step': 2, 'writes': {'tools': {'messages': [ToolMessage(content='It might be cloudy in nyc', name='get_weather', id='04dd3f8a-c063-4885-bfcd-901aaeb63797', tool_call_id='call_TJWuPJovXruJVafVcI540OfS')]}}}, parent_config={'configurable': {'thread_id': '4', 'thread_ts': '2024-06-27T01:06:22.010058+00:00'}}),\n",
-       " CheckpointTuple(config={'configurable': {'thread_id': '4', 'thread_ts': '2024-06-27T01:06:22.004456+00:00'}}, checkpoint={'v': 1, 'ts': '2024-06-27T01:06:22.004456+00:00', 'id': '1ef34217-7c9a-6a64-8001-65cb5b71ff36', 'channel_values': {'messages': [HumanMessage(content=\"what's the weather in nyc\", id='9df9dda1-5e80-4938-bf24-9207738d9f5f'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_TJWuPJovXruJVafVcI540OfS', 'function': {'arguments': '{\"city\":\"nyc\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 15, 'prompt_tokens': 58, 'total_tokens': 73}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_4008e3b719', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-5ba0db1c-45e7-4389-975f-489135ff50b4-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'nyc'}, 'id': 'call_TJWuPJovXruJVafVcI540OfS'}], usage_metadata={'input_tokens': 58, 'output_tokens': 15, 'total_tokens': 73})], 'agent': 'agent', 'branch:agent:should_continue:tools': 'agent'}, 'channel_versions': {'__start__': 2, 'messages': 3, 'start:agent': 3, 'agent': 3, 'branch:agent:should_continue:tools': 3}, 'versions_seen': {'__start__': {'__start__': 1}, 'agent': {'start:agent': 2}, 'tools': {}}, 'pending_sends': []}, metadata={'source': 'loop', 'step': 1, 'writes': {'agent': {'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_TJWuPJovXruJVafVcI540OfS', 'function': {'arguments': '{\"city\":\"nyc\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 15, 'prompt_tokens': 58, 'total_tokens': 73}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_4008e3b719', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-5ba0db1c-45e7-4389-975f-489135ff50b4-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'nyc'}, 'id': 'call_TJWuPJovXruJVafVcI540OfS'}], usage_metadata={'input_tokens': 58, 'output_tokens': 15, 'total_tokens': 73})]}}}, parent_config={'configurable': {'thread_id': '4', 'thread_ts': '2024-06-27T01:06:22.004456+00:00'}}),\n",
-       " CheckpointTuple(config={'configurable': {'thread_id': '4', 'thread_ts': '2024-06-27T01:06:21.494039+00:00'}}, checkpoint={'v': 1, 'ts': '2024-06-27T01:06:21.494039+00:00', 'id': '1ef34217-77bc-67e0-8000-0dcc15e3e392', 'channel_values': {'messages': [HumanMessage(content=\"what's the weather in nyc\", id='9df9dda1-5e80-4938-bf24-9207738d9f5f')], 'start:agent': '__start__'}, 'channel_versions': {'__start__': 2, 'messages': 2, 'start:agent': 2}, 'versions_seen': {'__start__': {'__start__': 1}, 'agent': {}, 'tools': {}}, 'pending_sends': []}, metadata={'source': 'loop', 'step': 0, 'writes': None}, parent_config={'configurable': {'thread_id': '4', 'thread_ts': '2024-06-27T01:06:21.494039+00:00'}}),\n",
-       " CheckpointTuple(config={'configurable': {'thread_id': '4', 'thread_ts': '2024-06-27T01:06:21.491815+00:00'}}, checkpoint={'v': 1, 'ts': '2024-06-27T01:06:21.491815+00:00', 'id': '1ef34217-77b7-6114-bfff-cceb4fa6453c', 'channel_values': {'messages': [], '__start__': {'messages': [['human', \"what's the weather in nyc\"]]}}, 'channel_versions': {'__start__': 1}, 'versions_seen': {}, 'pending_sends': []}, metadata={'source': 'input', 'step': -1, 'writes': {'messages': [['human', \"what's the weather in nyc\"]]}}, parent_config=None)]"
+       "[CheckpointTuple(config={'configurable': {'thread_id': '4', 'thread_ts': '1ef40627-96b9-6682-8003-134aebfec1e9'}}, checkpoint={'v': 1, 'ts': '2024-07-12T15:21:55.984109+00:00', 'id': '1ef40627-96b9-6682-8003-134aebfec1e9', 'channel_values': {'messages': [HumanMessage(content=\"what's the weather in nyc\", id='1c1e48ba-fa25-4190-a847-459828f44579'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_pS4ybOXkIDOmS93jZ8wOYGfU', 'function': {'arguments': '{\"city\":\"nyc\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 15, 'prompt_tokens': 58, 'total_tokens': 73}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_d33f7b429e', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-07c7ee03-64f7-462a-9249-615572156216-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'nyc'}, 'id': 'call_pS4ybOXkIDOmS93jZ8wOYGfU', 'type': 'tool_call'}], usage_metadata={'input_tokens': 58, 'output_tokens': 15, 'total_tokens': 73}), ToolMessage(content='It might be cloudy in nyc', name='get_weather', id='50e612d7-c770-44dd-b128-4bfdbd7d5b7d', tool_call_id='call_pS4ybOXkIDOmS93jZ8wOYGfU'), AIMessage(content='The weather in NYC might be cloudy.', response_metadata={'token_usage': {'completion_tokens': 9, 'prompt_tokens': 88, 'total_tokens': 97}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_d33f7b429e', 'finish_reason': 'stop', 'logprobs': None}, id='run-924e10c9-7005-4cbf-a92e-3ce63b54092f-0', usage_metadata={'input_tokens': 88, 'output_tokens': 9, 'total_tokens': 97})], 'agent': 'agent'}, 'channel_versions': {'__start__': 2, 'messages': 5, 'start:agent': 3, 'agent': 5, 'branch:agent:should_continue:tools': 4, 'tools': 5}, 'versions_seen': {'__start__': {'__start__': 1}, 'agent': {'start:agent': 3, 'tools': 4}, 'tools': {'branch:agent:should_continue:tools': 3}}, 'pending_sends': []}, metadata={'source': 'loop', 'step': 3, 'writes': {'agent': {'messages': [AIMessage(content='The weather in NYC might be cloudy.', response_metadata={'token_usage': {'completion_tokens': 9, 'prompt_tokens': 88, 'total_tokens': 97}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_d33f7b429e', 'finish_reason': 'stop', 'logprobs': None}, id='run-924e10c9-7005-4cbf-a92e-3ce63b54092f-0', usage_metadata={'input_tokens': 88, 'output_tokens': 9, 'total_tokens': 97})]}}}, parent_config={'configurable': {'thread_id': '4', 'thread_ts': '1ef40627-96b9-6682-8003-134aebfec1e9'}}, pending_writes=None),\n",
+       " CheckpointTuple(config={'configurable': {'thread_id': '4', 'thread_ts': '1ef40627-91a0-6100-8002-e404dda477d4'}}, checkpoint={'v': 1, 'ts': '2024-07-12T15:21:55.449447+00:00', 'id': '1ef40627-91a0-6100-8002-e404dda477d4', 'channel_values': {'messages': [HumanMessage(content=\"what's the weather in nyc\", id='1c1e48ba-fa25-4190-a847-459828f44579'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_pS4ybOXkIDOmS93jZ8wOYGfU', 'function': {'arguments': '{\"city\":\"nyc\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 15, 'prompt_tokens': 58, 'total_tokens': 73}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_d33f7b429e', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-07c7ee03-64f7-462a-9249-615572156216-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'nyc'}, 'id': 'call_pS4ybOXkIDOmS93jZ8wOYGfU', 'type': 'tool_call'}], usage_metadata={'input_tokens': 58, 'output_tokens': 15, 'total_tokens': 73}), ToolMessage(content='It might be cloudy in nyc', name='get_weather', id='50e612d7-c770-44dd-b128-4bfdbd7d5b7d', tool_call_id='call_pS4ybOXkIDOmS93jZ8wOYGfU')], 'tools': 'tools'}, 'channel_versions': {'__start__': 2, 'messages': 4, 'start:agent': 3, 'agent': 4, 'branch:agent:should_continue:tools': 4, 'tools': 4}, 'versions_seen': {'__start__': {'__start__': 1}, 'agent': {'start:agent': 2}, 'tools': {'branch:agent:should_continue:tools': 3}}, 'pending_sends': []}, metadata={'source': 'loop', 'step': 2, 'writes': {'tools': {'messages': [ToolMessage(content='It might be cloudy in nyc', name='get_weather', id='50e612d7-c770-44dd-b128-4bfdbd7d5b7d', tool_call_id='call_pS4ybOXkIDOmS93jZ8wOYGfU')]}}}, parent_config={'configurable': {'thread_id': '4', 'thread_ts': '1ef40627-91a0-6100-8002-e404dda477d4'}}, pending_writes=None),\n",
+       " CheckpointTuple(config={'configurable': {'thread_id': '4', 'thread_ts': '1ef40627-9194-66de-8001-86c8d77c2d7c'}}, checkpoint={'v': 1, 'ts': '2024-07-12T15:21:55.444687+00:00', 'id': '1ef40627-9194-66de-8001-86c8d77c2d7c', 'channel_values': {'messages': [HumanMessage(content=\"what's the weather in nyc\", id='1c1e48ba-fa25-4190-a847-459828f44579'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_pS4ybOXkIDOmS93jZ8wOYGfU', 'function': {'arguments': '{\"city\":\"nyc\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 15, 'prompt_tokens': 58, 'total_tokens': 73}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_d33f7b429e', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-07c7ee03-64f7-462a-9249-615572156216-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'nyc'}, 'id': 'call_pS4ybOXkIDOmS93jZ8wOYGfU', 'type': 'tool_call'}], usage_metadata={'input_tokens': 58, 'output_tokens': 15, 'total_tokens': 73})], 'agent': 'agent', 'branch:agent:should_continue:tools': 'agent'}, 'channel_versions': {'__start__': 2, 'messages': 3, 'start:agent': 3, 'agent': 3, 'branch:agent:should_continue:tools': 3}, 'versions_seen': {'__start__': {'__start__': 1}, 'agent': {'start:agent': 2}, 'tools': {}}, 'pending_sends': []}, metadata={'source': 'loop', 'step': 1, 'writes': {'agent': {'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_pS4ybOXkIDOmS93jZ8wOYGfU', 'function': {'arguments': '{\"city\":\"nyc\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 15, 'prompt_tokens': 58, 'total_tokens': 73}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_d33f7b429e', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-07c7ee03-64f7-462a-9249-615572156216-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'nyc'}, 'id': 'call_pS4ybOXkIDOmS93jZ8wOYGfU', 'type': 'tool_call'}], usage_metadata={'input_tokens': 58, 'output_tokens': 15, 'total_tokens': 73})]}}}, parent_config={'configurable': {'thread_id': '4', 'thread_ts': '1ef40627-9194-66de-8001-86c8d77c2d7c'}}, pending_writes=None),\n",
+       " CheckpointTuple(config={'configurable': {'thread_id': '4', 'thread_ts': '1ef40627-8b8a-6b1c-8000-55b423aa733b'}}, checkpoint={'v': 1, 'ts': '2024-07-12T15:21:54.811566+00:00', 'id': '1ef40627-8b8a-6b1c-8000-55b423aa733b', 'channel_values': {'messages': [HumanMessage(content=\"what's the weather in nyc\", id='1c1e48ba-fa25-4190-a847-459828f44579')], 'start:agent': '__start__'}, 'channel_versions': {'__start__': 2, 'messages': 2, 'start:agent': 2}, 'versions_seen': {'__start__': {'__start__': 1}, 'agent': {}, 'tools': {}}, 'pending_sends': []}, metadata={'source': 'loop', 'step': 0, 'writes': None}, parent_config={'configurable': {'thread_id': '4', 'thread_ts': '1ef40627-8b8a-6b1c-8000-55b423aa733b'}}, pending_writes=None),\n",
+       " CheckpointTuple(config={'configurable': {'thread_id': '4', 'thread_ts': '1ef40627-8b88-62b8-bfff-9922bbf9342b'}}, checkpoint={'v': 1, 'ts': '2024-07-12T15:21:54.810527+00:00', 'id': '1ef40627-8b88-62b8-bfff-9922bbf9342b', 'channel_values': {'messages': [], '__start__': {'messages': [['human', \"what's the weather in nyc\"]]}}, 'channel_versions': {'__start__': 1}, 'versions_seen': {}, 'pending_sends': []}, metadata={'source': 'input', 'step': -1, 'writes': {'messages': [['human', \"what's the weather in nyc\"]]}}, parent_config=None, pending_writes=None)]"
       ]
      },
      "execution_count": 17,

examples/persistence_redis.ipynb

@@ -0,0 +1,842 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "51466c8d-8ce4-4b3d-be4e-18fdbeda5f53",
+   "metadata": {},
+   "source": [
+    "# How to create a custom checkpointer using Redis\n",
+    "\n",
+    "When creating LangGraph agents, you can also set them up so that they persist their state. This allows you to do things like interact with an agent multiple times and have it remember previous interactions. Make sure that you have Redis running on port `6379` for going through this tutorial\n",
+    "\n",
+    "This example shows how to use `Redis` as the backend for persisting checkpoint state.\n",
+    "\n",
+    "NOTE: this is just an example implementation. You can implement your own checkpointer using a different database or modify this one as long as it conforms to the `BaseCheckpointSaver` interface."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0aac2830",
+   "metadata": {},
+   "source": [
+    "## Install the necessary libraries for Redis on Python"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "faadfb1b-cebe-4dcf-82fd-34044c380bc4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%capture --no-stderr\n",
+    "%pip install -U redis langgraph"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a6a4e417",
+   "metadata": {},
+   "source": [
+    "## Checkpointer implementation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "a35dba8e-5562-4803-ad80-160f53592dd7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\"\"\"Implementation of a langgraph checkpoint saver using Redis.\"\"\"\n",
+    "from contextlib import asynccontextmanager, contextmanager\n",
+    "from typing import Any, AsyncGenerator, Generator, Union, Tuple, Optional\n",
+    "\n",
+    "import redis\n",
+    "from redis.asyncio import Redis as AsyncRedis, ConnectionPool as AsyncConnectionPool\n",
+    "from langchain_core.runnables import RunnableConfig\n",
+    "from langgraph.checkpoint import BaseCheckpointSaver\n",
+    "from langgraph.serde.jsonplus import JsonPlusSerializer\n",
+    "from langgraph.checkpoint.base import Checkpoint, CheckpointMetadata, CheckpointTuple\n",
+    "import logging\n",
+    "\n",
+    "logging.basicConfig(level=logging.INFO)\n",
+    "logger = logging.getLogger(__name__)\n",
+    "\n",
+    "\n",
+    "class JsonAndBinarySerializer(JsonPlusSerializer):\n",
+    "    def _default(self, obj: Any) -> Any:\n",
+    "        if isinstance(obj, (bytes, bytearray)):\n",
+    "            return self._encode_constructor_args(\n",
+    "                obj.__class__, method=\"fromhex\", args=[obj.hex()]\n",
+    "            )\n",
+    "        return super()._default(obj)\n",
+    "\n",
+    "    def dumps(self, obj: Any) -> str:\n",
+    "        try:\n",
+    "            if isinstance(obj, (bytes, bytearray)):\n",
+    "                return obj.hex()\n",
+    "            return super().dumps(obj)\n",
+    "        except Exception as e:\n",
+    "            logger.error(f\"Serialization error: {e}\")\n",
+    "            raise\n",
+    "\n",
+    "    def loads(self, s: str, is_binary: bool = False) -> Any:\n",
+    "        try:\n",
+    "            if is_binary:\n",
+    "                return bytes.fromhex(s)\n",
+    "            return super().loads(s)\n",
+    "        except Exception as e:\n",
+    "            logger.error(f\"Deserialization error: {e}\")\n",
+    "            raise\n",
+    "\n",
+    "\n",
+    "def initialize_sync_pool(\n",
+    "    host: str = \"localhost\", port: int = 6379, db: int = 0, **kwargs\n",
+    ") -> redis.ConnectionPool:\n",
+    "    \"\"\"Initialize a synchronous Redis connection pool.\"\"\"\n",
+    "    try:\n",
+    "        pool = redis.ConnectionPool(host=host, port=port, db=db, **kwargs)\n",
+    "        logger.info(\n",
+    "            f\"Synchronous Redis pool initialized with host={host}, port={port}, db={db}\"\n",
+    "        )\n",
+    "        return pool\n",
+    "    except Exception as e:\n",
+    "        logger.error(f\"Error initializing sync pool: {e}\")\n",
+    "        raise\n",
+    "\n",
+    "\n",
+    "def initialize_async_pool(\n",
+    "    url: str = \"redis://localhost\", **kwargs\n",
+    ") -> AsyncConnectionPool:\n",
+    "    \"\"\"Initialize an asynchronous Redis connection pool.\"\"\"\n",
+    "    try:\n",
+    "        pool = AsyncConnectionPool.from_url(url, **kwargs)\n",
+    "        logger.info(f\"Asynchronous Redis pool initialized with url={url}\")\n",
+    "        return pool\n",
+    "    except Exception as e:\n",
+    "        logger.error(f\"Error initializing async pool: {e}\")\n",
+    "        raise\n",
+    "\n",
+    "\n",
+    "@contextmanager\n",
+    "def _get_sync_connection(\n",
+    "    connection: Union[redis.Redis, redis.ConnectionPool, None]\n",
+    ") -> Generator[redis.Redis, None, None]:\n",
+    "    conn = None\n",
+    "    try:\n",
+    "        if isinstance(connection, redis.Redis):\n",
+    "            yield connection\n",
+    "        elif isinstance(connection, redis.ConnectionPool):\n",
+    "            conn = redis.Redis(connection_pool=connection)\n",
+    "            yield conn\n",
+    "        else:\n",
+    "            raise ValueError(\"Invalid sync connection object.\")\n",
+    "    except redis.ConnectionError as e:\n",
+    "        logger.error(f\"Sync connection error: {e}\")\n",
+    "        raise\n",
+    "    finally:\n",
+    "        if conn:\n",
+    "            conn.close()\n",
+    "\n",
+    "\n",
+    "@asynccontextmanager\n",
+    "async def _get_async_connection(\n",
+    "    connection: Union[AsyncRedis, AsyncConnectionPool, None]\n",
+    ") -> AsyncGenerator[AsyncRedis, None]:\n",
+    "    conn = None\n",
+    "    try:\n",
+    "        if isinstance(connection, AsyncRedis):\n",
+    "            yield connection\n",
+    "        elif isinstance(connection, AsyncConnectionPool):\n",
+    "            conn = AsyncRedis(connection_pool=connection)\n",
+    "            yield conn\n",
+    "        else:\n",
+    "            raise ValueError(\"Invalid async connection object.\")\n",
+    "    except redis.ConnectionError as e:\n",
+    "        logger.error(f\"Async connection error: {e}\")\n",
+    "        raise\n",
+    "    finally:\n",
+    "        if conn:\n",
+    "            await conn.aclose()\n",
+    "\n",
+    "\n",
+    "class RedisSaver(BaseCheckpointSaver):\n",
+    "    sync_connection: Optional[Union[redis.Redis, redis.ConnectionPool]] = None\n",
+    "    async_connection: Optional[Union[AsyncRedis, AsyncConnectionPool]] = None\n",
+    "\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        sync_connection: Optional[Union[redis.Redis, redis.ConnectionPool]] = None,\n",
+    "        async_connection: Optional[Union[AsyncRedis, AsyncConnectionPool]] = None,\n",
+    "    ):\n",
+    "        super().__init__(serde=JsonAndBinarySerializer())\n",
+    "        self.sync_connection = sync_connection\n",
+    "        self.async_connection = async_connection\n",
+    "\n",
+    "    def put(\n",
+    "        self,\n",
+    "        config: RunnableConfig,\n",
+    "        checkpoint: Checkpoint,\n",
+    "        metadata: CheckpointMetadata,\n",
+    "    ) -> RunnableConfig:\n",
+    "        thread_id = config[\"configurable\"][\"thread_id\"]\n",
+    "        parent_ts = config[\"configurable\"].get(\"thread_ts\")\n",
+    "        key = f\"checkpoint:{thread_id}:{checkpoint['ts']}\"\n",
+    "        try:\n",
+    "            with _get_sync_connection(self.sync_connection) as conn:\n",
+    "                conn.hset(\n",
+    "                    key,\n",
+    "                    mapping={\n",
+    "                        \"checkpoint\": self.serde.dumps(checkpoint),\n",
+    "                        \"metadata\": self.serde.dumps(metadata),\n",
+    "                        \"parent_ts\": parent_ts if parent_ts else \"\",\n",
+    "                    },\n",
+    "                )\n",
+    "                logger.info(\n",
+    "                    f\"Checkpoint stored successfully for thread_id: {thread_id}, ts: {checkpoint['ts']}\"\n",
+    "                )\n",
+    "        except Exception as e:\n",
+    "            logger.error(f\"Failed to put checkpoint: {e}\")\n",
+    "            raise\n",
+    "        return {\n",
+    "            \"configurable\": {\n",
+    "                \"thread_id\": thread_id,\n",
+    "                \"thread_ts\": checkpoint[\"ts\"],\n",
+    "            },\n",
+    "        }\n",
+    "\n",
+    "    async def aput(\n",
+    "        self,\n",
+    "        config: RunnableConfig,\n",
+    "        checkpoint: Checkpoint,\n",
+    "        metadata: CheckpointMetadata,\n",
+    "    ) -> RunnableConfig:\n",
+    "        thread_id = config[\"configurable\"][\"thread_id\"]\n",
+    "        parent_ts = config[\"configurable\"].get(\"thread_ts\")\n",
+    "        key = f\"checkpoint:{thread_id}:{checkpoint['ts']}\"\n",
+    "        try:\n",
+    "            async with _get_async_connection(self.async_connection) as conn:\n",
+    "                await conn.hset(\n",
+    "                    key,\n",
+    "                    mapping={\n",
+    "                        \"checkpoint\": self.serde.dumps(checkpoint),\n",
+    "                        \"metadata\": self.serde.dumps(metadata),\n",
+    "                        \"parent_ts\": parent_ts if parent_ts else \"\",\n",
+    "                    },\n",
+    "                )\n",
+    "                logger.info(\n",
+    "                    f\"Checkpoint stored successfully for thread_id: {thread_id}, ts: {checkpoint['ts']}\"\n",
+    "                )\n",
+    "        except Exception as e:\n",
+    "            logger.error(f\"Failed to aput checkpoint: {e}\")\n",
+    "            raise\n",
+    "        return {\n",
+    "            \"configurable\": {\n",
+    "                \"thread_id\": thread_id,\n",
+    "                \"thread_ts\": checkpoint[\"ts\"],\n",
+    "            },\n",
+    "        }\n",
+    "\n",
+    "    def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:\n",
+    "        thread_id = config[\"configurable\"][\"thread_id\"]\n",
+    "        thread_ts = config[\"configurable\"].get(\"thread_ts\", None)\n",
+    "        try:\n",
+    "            with _get_sync_connection(self.sync_connection) as conn:\n",
+    "                if thread_ts:\n",
+    "                    key = f\"checkpoint:{thread_id}:{thread_ts}\"\n",
+    "                else:\n",
+    "                    all_keys = conn.keys(f\"checkpoint:{thread_id}:*\")\n",
+    "                    if not all_keys:\n",
+    "                        logger.info(f\"No checkpoints found for thread_id: {thread_id}\")\n",
+    "                        return None\n",
+    "                    latest_key = max(all_keys, key=lambda k: k.decode().split(\":\")[-1])\n",
+    "                    key = latest_key.decode()\n",
+    "                checkpoint_data = conn.hgetall(key)\n",
+    "                if not checkpoint_data:\n",
+    "                    logger.info(f\"No valid checkpoint data found for key: {key}\")\n",
+    "                    return None\n",
+    "                checkpoint = self.serde.loads(checkpoint_data[b\"checkpoint\"].decode())\n",
+    "                metadata = self.serde.loads(checkpoint_data[b\"metadata\"].decode())\n",
+    "                parent_ts = checkpoint_data.get(b\"parent_ts\", b\"\").decode()\n",
+    "                parent_config = (\n",
+    "                    {\"configurable\": {\"thread_id\": thread_id, \"thread_ts\": parent_ts}}\n",
+    "                    if parent_ts\n",
+    "                    else None\n",
+    "                )\n",
+    "                logger.info(\n",
+    "                    f\"Checkpoint retrieved successfully for thread_id: {thread_id}, ts: {thread_ts}\"\n",
+    "                )\n",
+    "                return CheckpointTuple(\n",
+    "                    config=config,\n",
+    "                    checkpoint=checkpoint,\n",
+    "                    metadata=metadata,\n",
+    "                    parent_config=parent_config,\n",
+    "                )\n",
+    "        except Exception as e:\n",
+    "            logger.error(f\"Failed to get checkpoint tuple: {e}\")\n",
+    "            raise\n",
+    "\n",
+    "    async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:\n",
+    "        thread_id = config[\"configurable\"][\"thread_id\"]\n",
+    "        thread_ts = config[\"configurable\"].get(\"thread_ts\", None)\n",
+    "        try:\n",
+    "            async with _get_async_connection(self.async_connection) as conn:\n",
+    "                if thread_ts:\n",
+    "                    key = f\"checkpoint:{thread_id}:{thread_ts}\"\n",
+    "                else:\n",
+    "                    all_keys = await conn.keys(f\"checkpoint:{thread_id}:*\")\n",
+    "                    if not all_keys:\n",
+    "                        logger.info(f\"No checkpoints found for thread_id: {thread_id}\")\n",
+    "                        return None\n",
+    "                    latest_key = max(all_keys, key=lambda k: k.decode().split(\":\")[-1])\n",
+    "                    key = latest_key.decode()\n",
+    "                checkpoint_data = await conn.hgetall(key)\n",
+    "                if not checkpoint_data:\n",
+    "                    logger.info(f\"No valid checkpoint data found for key: {key}\")\n",
+    "                    return None\n",
+    "                checkpoint = self.serde.loads(checkpoint_data[b\"checkpoint\"].decode())\n",
+    "                metadata = self.serde.loads(checkpoint_data[b\"metadata\"].decode())\n",
+    "                parent_ts = checkpoint_data.get(b\"parent_ts\", b\"\").decode()\n",
+    "                parent_config = (\n",
+    "                    {\"configurable\": {\"thread_id\": thread_id, \"thread_ts\": parent_ts}}\n",
+    "                    if parent_ts\n",
+    "                    else None\n",
+    "                )\n",
+    "                logger.info(\n",
+    "                    f\"Checkpoint retrieved successfully for thread_id: {thread_id}, ts: {thread_ts}\"\n",
+    "                )\n",
+    "                return CheckpointTuple(\n",
+    "                    config=config,\n",
+    "                    checkpoint=checkpoint,\n",
+    "                    metadata=metadata,\n",
+    "                    parent_config=parent_config,\n",
+    "                )\n",
+    "        except Exception as e:\n",
+    "            logger.error(f\"Failed to get checkpoint tuple: {e}\")\n",
+    "            raise\n",
+    "\n",
+    "    def list(\n",
+    "        self,\n",
+    "        config: Optional[RunnableConfig],\n",
+    "        *,\n",
+    "        filter: Optional[dict[str, Any]] = None,\n",
+    "        before: Optional[RunnableConfig] = None,\n",
+    "        limit: Optional[int] = None,\n",
+    "    ) -> Generator[CheckpointTuple, None, None]:\n",
+    "        thread_id = config[\"configurable\"][\"thread_id\"] if config else \"*\"\n",
+    "        pattern = f\"checkpoint:{thread_id}:*\"\n",
+    "        try:\n",
+    "            with _get_sync_connection(self.sync_connection) as conn:\n",
+    "                keys = conn.keys(pattern)\n",
+    "                if before:\n",
+    "                    keys = [\n",
+    "                        k\n",
+    "                        for k in keys\n",
+    "                        if k.decode().split(\":\")[-1]\n",
+    "                        < before[\"configurable\"][\"thread_ts\"]\n",
+    "                    ]\n",
+    "                keys = sorted(\n",
+    "                    keys, key=lambda k: k.decode().split(\":\")[-1], reverse=True\n",
+    "                )\n",
+    "                if limit:\n",
+    "                    keys = keys[:limit]\n",
+    "                for key in keys:\n",
+    "                    data = conn.hgetall(key)\n",
+    "                    if data and \"checkpoint\" in data and \"metadata\" in data:\n",
+    "                        thread_ts = key.decode().split(\":\")[-1]\n",
+    "                        yield CheckpointTuple(\n",
+    "                            config={\n",
+    "                                \"configurable\": {\n",
+    "                                    \"thread_id\": thread_id,\n",
+    "                                    \"thread_ts\": thread_ts,\n",
+    "                                }\n",
+    "                            },\n",
+    "                            checkpoint=self.serde.loads(data[\"checkpoint\"].decode()),\n",
+    "                            metadata=self.serde.loads(data[\"metadata\"].decode()),\n",
+    "                            parent_config={\n",
+    "                                \"configurable\": {\n",
+    "                                    \"thread_id\": thread_id,\n",
+    "                                    \"thread_ts\": data.get(\"parent_ts\", b\"\").decode(),\n",
+    "                                }\n",
+    "                            }\n",
+    "                            if data.get(\"parent_ts\")\n",
+    "                            else None,\n",
+    "                        )\n",
+    "                        logger.info(\n",
+    "                            f\"Checkpoint listed for thread_id: {thread_id}, ts: {thread_ts}\"\n",
+    "                        )\n",
+    "        except Exception as e:\n",
+    "            logger.error(f\"Failed to list checkpoints: {e}\")\n",
+    "            raise\n",
+    "\n",
+    "    async def alist(\n",
+    "        self,\n",
+    "        config: Optional[RunnableConfig],\n",
+    "        *,\n",
+    "        filter: Optional[dict[str, Any]] = None,\n",
+    "        before: Optional[RunnableConfig] = None,\n",
+    "        limit: Optional[int] = None,\n",
+    "    ) -> AsyncGenerator[CheckpointTuple, None]:\n",
+    "        thread_id = config[\"configurable\"][\"thread_id\"] if config else \"*\"\n",
+    "        pattern = f\"checkpoint:{thread_id}:*\"\n",
+    "        try:\n",
+    "            async with _get_async_connection(self.async_connection) as conn:\n",
+    "                keys = await conn.keys(pattern)\n",
+    "                if before:\n",
+    "                    keys = [\n",
+    "                        k\n",
+    "                        for k in keys\n",
+    "                        if k.decode().split(\":\")[-1]\n",
+    "                        < before[\"configurable\"][\"thread_ts\"]\n",
+    "                    ]\n",
+    "                keys = sorted(\n",
+    "                    keys, key=lambda k: k.decode().split(\":\")[-1], reverse=True\n",
+    "                )\n",
+    "                if limit:\n",
+    "                    keys = keys[:limit]\n",
+    "                for key in keys:\n",
+    "                    data = await conn.hgetall(key)\n",
+    "                    if data and \"checkpoint\" in data and \"metadata\" in data:\n",
+    "                        thread_ts = key.decode().split(\":\")[-1]\n",
+    "                        yield CheckpointTuple(\n",
+    "                            config={\n",
+    "                                \"configurable\": {\n",
+    "                                    \"thread_id\": thread_id,\n",
+    "                                    \"thread_ts\": thread_ts,\n",
+    "                                }\n",
+    "                            },\n",
+    "                            checkpoint=self.serde.loads(data[\"checkpoint\"].decode()),\n",
+    "                            metadata=self.serde.loads(data[\"metadata\"].decode()),\n",
+    "                            parent_config={\n",
+    "                                \"configurable\": {\n",
+    "                                    \"thread_id\": thread_id,\n",
+    "                                    \"thread_ts\": data.get(\"parent_ts\", b\"\").decode(),\n",
+    "                                }\n",
+    "                            }\n",
+    "                            if data.get(\"parent_ts\")\n",
+    "                            else None,\n",
+    "                        )\n",
+    "                        logger.info(\n",
+    "                            f\"Checkpoint listed for thread_id: {thread_id}, ts: {thread_ts}\"\n",
+    "                        )\n",
+    "        except Exception as e:\n",
+    "            logger.error(f\"Failed to list checkpoints: {e}\")\n",
+    "            raise"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1d142495",
+   "metadata": {},
+   "source": [
+    "## Checkpointer implementation"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "456fa19c-93a5-4750-a410-f2d810b964ad",
+   "metadata": {},
+   "source": [
+    "## Setup environment"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "eca9aafb-a155-407a-8036-682a2f1297d7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import getpass\n",
+    "import os\n",
+    "\n",
+    "\n",
+    "def _set_env(var: str):\n",
+    "    if not os.environ.get(var):\n",
+    "        os.environ[var] = getpass.getpass(f\"{var}: \")\n",
+    "\n",
+    "\n",
+    "_set_env(\"OPENAI_API_KEY\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e26b3204-cca2-414c-800e-7e09032445ae",
+   "metadata": {},
+   "source": [
+    "## Setup model and tools for the graph"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "e5213193-5a7d-43e7-aeba-fe732bb1cd7a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from typing import Literal\n",
+    "from langchain_core.runnables import ConfigurableField\n",
+    "from langchain_core.tools import tool\n",
+    "from langchain_openai import ChatOpenAI\n",
+    "from langgraph.prebuilt import create_react_agent\n",
+    "\n",
+    "\n",
+    "@tool\n",
+    "def get_weather(city: Literal[\"nyc\", \"sf\"]):\n",
+    "    \"\"\"Use this to get weather information.\"\"\"\n",
+    "    if city == \"nyc\":\n",
+    "        return \"It might be cloudy in nyc\"\n",
+    "    elif city == \"sf\":\n",
+    "        return \"It's always sunny in sf\"\n",
+    "    else:\n",
+    "        raise AssertionError(\"Unknown city\")\n",
+    "\n",
+    "\n",
+    "tools = [get_weather]\n",
+    "model = ChatOpenAI(model_name=\"gpt-4o\", temperature=0)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e9342c62-dbb4-40f6-9271-7393f1ca48c4",
+   "metadata": {},
+   "source": [
+    "## Use sync connection"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e39fc712-9e1c-4831-9077-dd07b0c13594",
+   "metadata": {},
+   "source": [
+    "### With a connection pool"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "a1710e2f",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "INFO:__main__:Synchronous Redis pool initialized with host=172.25.0.4, port=6379, db=0\n"
+     ]
+    }
+   ],
+   "source": [
+    "sync_pool = initialize_sync_pool(host=\"172.25.0.4\", port=6379, db=0)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "2657c1c4-d8a5-4fe3-8f77-95415a98ed6c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "checkpointer = RedisSaver(sync_connection=sync_pool)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "6d388241-de57-4b4e-af7b-eb1081fb8f36",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "INFO:__main__:Checkpoint retrieved successfully for thread_id: 1, ts: None\n",
+      "INFO:__main__:Checkpoint stored successfully for thread_id: 1, ts: 2024-07-09T08:22:48.417492+00:00\n",
+      "INFO:__main__:Checkpoint stored successfully for thread_id: 1, ts: 2024-07-09T08:22:48.420714+00:00\n",
+      "INFO:httpx:HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:__main__:Checkpoint stored successfully for thread_id: 1, ts: 2024-07-09T08:22:49.458951+00:00\n",
+      "INFO:__main__:Checkpoint stored successfully for thread_id: 1, ts: 2024-07-09T08:22:49.465101+00:00\n",
+      "INFO:httpx:HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:__main__:Checkpoint stored successfully for thread_id: 1, ts: 2024-07-09T08:22:50.084141+00:00\n"
+     ]
+    }
+   ],
+   "source": [
+    "graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)\n",
+    "config = {\"configurable\": {\"thread_id\": \"1\"}}\n",
+    "res = graph.invoke({\"messages\": [(\"human\", \"what's the weather in sf\")]}, config)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "a7e0e7ec-a675-470b-9270-e4bdc59d4a4d",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'messages': [HumanMessage(content=\"what's the weather in sf\", id='64df8e19-0b9f-47f7-928f-4db3255485aa'),\n",
+       "  AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_n2XQOZHfpXpaNaviJakmjo82', 'function': {'arguments': '{\"city\":\"sf\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 14, 'prompt_tokens': 57, 'total_tokens': 71}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_ce0793330f', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-9057607f-6fa7-452b-95c7-f8f9832cb343-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'sf'}, 'id': 'call_n2XQOZHfpXpaNaviJakmjo82'}], usage_metadata={'input_tokens': 57, 'output_tokens': 14, 'total_tokens': 71}),\n",
+       "  ToolMessage(content=\"It's always sunny in sf\", name='get_weather', id='444d80db-8230-440a-b0eb-46a3f4db1006', tool_call_id='call_n2XQOZHfpXpaNaviJakmjo82'),\n",
+       "  AIMessage(content='The weather in San Francisco is currently sunny.', response_metadata={'token_usage': {'completion_tokens': 10, 'prompt_tokens': 84, 'total_tokens': 94}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_d576307f90', 'finish_reason': 'stop', 'logprobs': None}, id='run-676a1ee4-7301-4405-93a4-87af27a92614-0', usage_metadata={'input_tokens': 84, 'output_tokens': 10, 'total_tokens': 94})]}"
+      ]
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "res"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "96efd8b2-97c9-4207-83b2-00131723a75a",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "INFO:__main__:Checkpoint retrieved successfully for thread_id: 1, ts: None\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "{'v': 1,\n",
+       " 'ts': '2024-07-08T12:21:14.392158+00:00',\n",
+       " 'id': '1ef3d249-1acf-60ed-bfff-248e42e4d9f5',\n",
+       " 'channel_values': {'messages': [],\n",
+       "  '__start__': {'messages': [['human', \"what's the weather in sf\"]]}},\n",
+       " 'channel_versions': {'__start__': 1},\n",
+       " 'versions_seen': {},\n",
+       " 'pending_sends': []}"
+      ]
+     },
+     "execution_count": 9,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "checkpointer.get(config)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "967c95c7-e392-4819-bd71-f29e91c68df3",
+   "metadata": {},
+   "source": [
+    "### With a connection"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "b7d3687b",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "INFO:__main__:Checkpoint retrieved successfully for thread_id: 2, ts: None\n",
+      "INFO:__main__:Checkpoint stored successfully for thread_id: 2, ts: 2024-07-09T08:22:50.132262+00:00\n",
+      "INFO:__main__:Checkpoint stored successfully for thread_id: 2, ts: 2024-07-09T08:22:50.135993+00:00\n",
+      "INFO:httpx:HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:__main__:Checkpoint stored successfully for thread_id: 2, ts: 2024-07-09T08:22:50.875540+00:00\n",
+      "INFO:__main__:Checkpoint retrieved successfully for thread_id: 2, ts: None\n"
+     ]
+    }
+   ],
+   "source": [
+    "import redis\n",
+    "\n",
+    "# Initialize the Redis synchronous direct connection\n",
+    "sync_redis_direct = redis.Redis(host=\"172.25.0.4\", port=6379, db=0)\n",
+    "\n",
+    "# Initialize the RedisSaver with the synchronous direct connection\n",
+    "checkpointer = RedisSaver(sync_connection=sync_redis_direct)\n",
+    "\n",
+    "graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)\n",
+    "config = {\"configurable\": {\"thread_id\": \"2\"}}\n",
+    "res = graph.invoke({\"messages\": [(\"human\", \"what's the weather in sf\")]}, config)\n",
+    "\n",
+    "checkpoint_tuple = checkpointer.get_tuple(config)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c0a47d3e-e588-48fc-a5d4-2145dff17e77",
+   "metadata": {},
+   "source": [
+    "## Use async connection"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ee6b6cf7-d8f7-4777-a48d-93b5855fe681",
+   "metadata": {},
+   "source": [
+    "### With a connection pool"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "20cea8b7-8f13-4dc7-a3c9-825040eb4c57",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "INFO:__main__:Asynchronous Redis pool initialized with url=redis://172.25.0.4:6379/0\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Initialize a synchronous Redis connection pool\n",
+    "async_pool = initialize_async_pool(url=\"redis://172.25.0.4:6379/0\")\n",
+    "\n",
+    "checkpointer = RedisSaver(async_connection=async_pool)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "f889dce6-7ec1-4277-b8af-ace7811733fa",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "INFO:__main__:Checkpoint retrieved successfully for thread_id: 3, ts: None\n",
+      "INFO:__main__:Checkpoint stored successfully for thread_id: 3, ts: 2024-07-09T08:22:50.949172+00:00\n",
+      "INFO:__main__:Checkpoint stored successfully for thread_id: 3, ts: 2024-07-09T08:22:50.951824+00:00\n",
+      "INFO:httpx:HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:__main__:Checkpoint stored successfully for thread_id: 3, ts: 2024-07-09T08:22:51.698633+00:00\n",
+      "INFO:__main__:Checkpoint stored successfully for thread_id: 3, ts: 2024-07-09T08:22:51.702156+00:00\n",
+      "INFO:httpx:HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:__main__:Checkpoint stored successfully for thread_id: 3, ts: 2024-07-09T08:22:53.530983+00:00\n"
+     ]
+    }
+   ],
+   "source": [
+    "graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)\n",
+    "config = {\"configurable\": {\"thread_id\": \"3\"}}\n",
+    "res = await graph.ainvoke(\n",
+    "    {\"messages\": [(\"human\", \"what's the weather in nyc\")]}, config\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "id": "ed58c722-1662-4ae2-9bb7-4872158a5b29",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "INFO:__main__:Checkpoint retrieved successfully for thread_id: 3, ts: None\n"
+     ]
+    }
+   ],
+   "source": [
+    "checkpoint_tuple = await checkpointer.aget_tuple(config)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "id": "e0c42044-4de6-4742-8e00-fe295d50c95a",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "CheckpointTuple(config={'configurable': {'thread_id': '3'}}, checkpoint={'v': 1, 'ts': '2024-07-08T12:21:18.866666+00:00', 'id': '1ef3d249-457b-62d3-bfff-b0e787336a7c', 'channel_values': {'messages': [], '__start__': {'messages': [['human', \"what's the weather in nyc\"]]}}, 'channel_versions': {'__start__': 1}, 'versions_seen': {}, 'pending_sends': []}, metadata={'source': 'input', 'step': -1, 'writes': {'messages': [['human', \"what's the weather in nyc\"]]}}, parent_config=None)"
+      ]
+     },
+     "execution_count": 16,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "checkpoint_tuple"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "56552584-9eb8-40df-a6a0-44151018b509",
+   "metadata": {},
+   "source": [
+    "### Use connection"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "id": "a7bf32bd",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "INFO:__main__:Checkpoint retrieved successfully for thread_id: 4, ts: None\n",
+      "INFO:__main__:Checkpoint stored successfully for thread_id: 4, ts: 2024-07-09T08:22:53.585109+00:00\n",
+      "INFO:__main__:Checkpoint stored successfully for thread_id: 4, ts: 2024-07-09T08:22:53.587207+00:00\n",
+      "INFO:httpx:HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:__main__:Checkpoint stored successfully for thread_id: 4, ts: 2024-07-09T08:22:54.932663+00:00\n",
+      "INFO:__main__:Checkpoint stored successfully for thread_id: 4, ts: 2024-07-09T08:22:54.936425+00:00\n",
+      "INFO:httpx:HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
+      "INFO:__main__:Checkpoint stored successfully for thread_id: 4, ts: 2024-07-09T08:22:55.982495+00:00\n"
+     ]
+    }
+   ],
+   "source": [
+    "from redis.asyncio import Redis as AsyncRedis\n",
+    "\n",
+    "async with await AsyncRedis(host=\"172.25.0.4\", port=6379, db=0) as conn:\n",
+    "    checkpointer = RedisSaver(async_connection=conn)\n",
+    "    graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)\n",
+    "    config = {\"configurable\": {\"thread_id\": \"4\"}}\n",
+    "    res = await graph.ainvoke(\n",
+    "        {\"messages\": [(\"human\", \"what's the weather in nyc\")]}, config\n",
+    "    )\n",
+    "    checkpoint_tuples = [c async for c in checkpointer.alist(config)]"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

examples/rag/langgraph_self_rag_local.ipynb

@@ -252,7 +252,7 @@
    "id": "5324ea49-5745-47b5-a0a5-bf58c8babe46",
    "metadata": {},
    "outputs": [],
-   "source": ["### Nodes\n\n\ndef retrieve(state):\n    \"\"\"\n    Retrieve documents\n\n    Args:\n        state (dict): The current graph state\n\n    Returns:\n        state (dict): New key added to state, documents, that contains retrieved documents\n    \"\"\"\n    print(\"---RETRIEVE---\")\n    question = state[\"question\"]\n\n    # Retrieval\n    documents = retriever.get_relevant_documents(question)\n    return {\"documents\": documents, \"question\": question}\n\n\ndef generate(state):\n    \"\"\"\n    Generate answer\n\n    Args:\n        state (dict): The current graph state\n\n    Returns:\n        state (dict): New key added to state, generation, that contains LLM generation\n    \"\"\"\n    print(\"---GENERATE---\")\n    question = state[\"question\"]\n    documents = state[\"documents\"]\n\n    # RAG generation\n    generation = rag_chain.invoke({\"context\": documents, \"question\": question})\n    return {\"documents\": documents, \"question\": question, \"generation\": generation}\n\n\ndef grade_documents(state):\n    \"\"\"\n    Determines whether the retrieved documents are relevant to the question.\n\n    Args:\n        state (dict): The current graph state\n\n    Returns:\n        state (dict): Updates documents key with only filtered relevant documents\n    \"\"\"\n\n    print(\"---CHECK DOCUMENT RELEVANCE TO QUESTION---\")\n    question = state[\"question\"]\n    documents = state[\"documents\"]\n\n    # Score each doc\n    filtered_docs = []\n    for d in documents:\n        score = retrieval_grader.invoke(\n            {\"question\": question, \"document\": d.page_content}\n        )\n        grade = score[\"score\"]\n        if grade == \"yes\":\n            print(\"---GRADE: DOCUMENT RELEVANT---\")\n            filtered_docs.append(d)\n        else:\n            print(\"---GRADE: DOCUMENT NOT RELEVANT---\")\n            continue\n    return {\"documents\": filtered_docs, \"question\": question}\n\n\ndef transform_query(state):\n    \"\"\"\n    Transform the query to produce a better question.\n\n    Args:\n        state (dict): The current graph state\n\n    Returns:\n        state (dict): Updates question key with a re-phrased question\n    \"\"\"\n\n    print(\"---TRANSFORM QUERY---\")\n    question = state[\"question\"]\n    documents = state[\"documents\"]\n\n    # Re-write question\n    better_question = question_rewriter.invoke({\"question\": question})\n    return {\"documents\": documents, \"question\": better_question}\n\n\n### Edges\n\n\ndef decide_to_generate(state):\n    \"\"\"\n    Determines whether to generate an answer, or re-generate a question.\n\n    Args:\n        state (dict): The current graph state\n\n    Returns:\n        str: Binary decision for next node to call\n    \"\"\"\n\n    print(\"---ASSESS GRADED DOCUMENTS---\")\n    state[\"question\"]\n    filtered_documents = state[\"documents\"]\n\n    if not filtered_documents:\n        # All documents have been filtered check_relevance\n        # We will re-generate a new query\n        print(\n            \"---DECISION: ALL DOCUMENTS ARE NOT RELEVANT TO QUESTION, TRANSFORM QUERY---\"\n        )\n        return \"transform_query\"\n    else:\n        # We have relevant documents, so generate answer\n        print(\"---DECISION: GENERATE---\")\n        return \"generate\"\n\n\ndef grade_generation_v_documents_and_question(state):\n    \"\"\"\n    Determines whether the generation is grounded in the document and answers question.\n\n    Args:\n        state (dict): The current graph state\n\n    Returns:\n        str: Decision for next node to call\n    \"\"\"\n\n    print(\"---CHECK HALLUCINATIONS---\")\n    question = state[\"question\"]\n    documents = state[\"documents\"]\n    generation = state[\"generation\"]\n\n    score = hallucination_grader.invoke(\n        {\"documents\": documents, \"generation\": generation}\n    )\n    grade = score[\"score\"]\n\n    # Check hallucination\n    if grade == \"yes\":\n        print(\"---DECISION: GENERATION IS GROUNDED IN DOCUMENTS---\")\n        # Check question-answering\n        print(\"---GRADE GENERATION vs QUESTION---\")\n        score = answer_grader.invoke({\"question\": question, \"generation\": generation})\n        grade = score[\"score\"]\n        if grade == \"yes\":\n            print(\"---DECISION: GENERATION ADDRESSES QUESTION---\")\n            return \"useful\"\n        else:\n            print(\"---DECISION: GENERATION DOES NOT ADDRESS QUESTION---\")\n            return \"not useful\"\n    else:\n        pprint(\"---DECISION: GENERATION IS NOT GROUNDED IN DOCUMENTS, RE-TRY---\")\n        return \"not supported\""]
+   "source": ["### Nodes\n\n\ndef retrieve(state):\n    \"\"\"\n    Retrieve documents\n\n    Args:\n        state (dict): The current graph state\n\n    Returns:\n        state (dict): New key added to state, documents, that contains retrieved documents\n    \"\"\"\n    print(\"---RETRIEVE---\")\n    question = state[\"question\"]\n\n    # Retrieval\n    documents = retriever.get_relevant_documents(question)\n    return {\"documents\": documents, \"question\": question}\n\n\ndef generate(state):\n    \"\"\"\n    Generate answer\n\n    Args:\n        state (dict): The current graph state\n\n    Returns:\n        state (dict): New key added to state, generation, that contains LLM generation\n    \"\"\"\n    print(\"---GENERATE---\")\n    question = state[\"question\"]\n    documents = state[\"documents\"]\n\n    # RAG generation\n    generation = rag_chain.invoke({\"context\": documents, \"question\": question})\n    return {\"documents\": documents, \"question\": question, \"generation\": generation}\n\n\ndef grade_documents(state):\n    \"\"\"\n    Determines whether the retrieved documents are relevant to the question.\n\n    Args:\n        state (dict): The current graph state\n\n    Returns:\n        state (dict): Updates documents key with only filtered relevant documents\n    \"\"\"\n\n    print(\"---CHECK DOCUMENT RELEVANCE TO QUESTION---\")\n    question = state[\"question\"]\n    documents = state[\"documents\"]\n\n    # Score each doc\n    filtered_docs = []\n    for d in documents:\n        score = retrieval_grader.invoke(\n            {\"question\": question, \"document\": d.page_content}\n        )\n        grade = score[\"score\"]\n        if grade == \"yes\":\n            print(\"---GRADE: DOCUMENT RELEVANT---\")\n            filtered_docs.append(d)\n        else:\n            print(\"---GRADE: DOCUMENT NOT RELEVANT---\")\n            continue\n    return {\"documents\": filtered_docs, \"question\": question}\n\n\ndef transform_query(state):\n    \"\"\"\n    Transform the query to produce a better question.\n\n    Args:\n        state (dict): The current graph state\n\n    Returns:\n        state (dict): Updates question key with a re-phrased question\n    \"\"\"\n\n    print(\"---TRANSFORM QUERY---\")\n    question = state[\"question\"]\n    documents = state[\"documents\"]\n\n    # Re-write question\n    better_question = question_rewriter.invoke({\"question\": question})\n    return {\"documents\": documents, \"question\": better_question}\n\n\n### Edges\n\n\ndef decide_to_generate(state):\n    \"\"\"\n    Determines whether to generate an answer, or re-generate a question.\n\n    Args:\n        state (dict): The current graph state\n\n    Returns:\n        str: Binary decision for next node to call\n    \"\"\"\n\n    print(\"---ASSESS GRADED DOCUMENTS---\")\n    state[\"question\"]\n    filtered_documents = state[\"documents\"]\n\n    if not filtered_documents:\n        # All documents have been filtered check_relevance\n        # We will re-generate a new query\n        print(\n            \"---DECISION: ALL DOCUMENTS ARE NOT RELEVANT TO QUESTION, TRANSFORM QUERY---\"\n        )\n        return \"transform_query\"\n    else:\n        # We have relevant documents, so generate answer\n        print(\"---DECISION: GENERATE---\")\n        return \"generate\"\n\n\ndef grade_generation_v_documents_and_question(state):\n    \"\"\"\n    Determines whether the generation is grounded in the document and answers question.\n\n    Args:\n        state (dict): The current graph state\n\n    Returns:\n        str: Decision for next node to call\n    \"\"\"\n\n    print(\"---CHECK HALLUCINATIONS---\")\n    question = state[\"question\"]\n    documents = state[\"documents\"]\n    generation = state[\"generation\"]\n\n    score = hallucination_grader.invoke(\n        {\"documents\": documents, \"generation\": generation}\n    )\n    grade = score[\"score\"]\n\n    # Check hallucination\n    if grade == \"yes\":\n        print(\"---DECISION: GENERATION IS GROUNDED IN DOCUMENTS---\")\n        # Check question-answering\n        print(\"---GRADE GENERATION vs QUESTION---\")\n        score = answer_grader.invoke({\"question\": question, \"generation\": generation})\n        grade = score[\"score\"]\n        if grade == \"yes\":\n            print(\"---DECISION: GENERATION ADDRESSES QUESTION---\")\n            return \"useful\"\n        else:\n            print(\"---DECISION: GENERATION DOES NOT ADDRESS QUESTION---\")\n            return \"not useful\"\n    else:\n        print(\"---DECISION: GENERATION IS NOT GROUNDED IN DOCUMENTS, RE-TRY---\")\n        return \"not supported\""]
   },
   {
    "cell_type": "markdown",

examples/streaming-events-from-within-tools-without-langchain.ipynb

@@ -0,0 +1,352 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "b23ced4e-dc29-43be-9f94-0c36bb181b8a",
+   "metadata": {},
+   "source": [
+    "# How to stream events from within a tool (without LangChain LLMs / tools)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7044eeb8-4074-4f9c-8a62-962488744557",
+   "metadata": {},
+   "source": [
+    "In this example we will stream tokens from within tools that an agent is using. We'll also be using OpenAI client library directly, without using LangChain chat models. We will use a ReAct agent as an example."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a37f60af-43ea-4aa6-847a-df8cc47065f5",
+   "metadata": {},
+   "source": [
+    "## Setup"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "47f79af8-58d8-4a48-8d9a-88823d88701f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%capture --no-stderr\n",
+    "%pip install -U langgraph openai"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "0cf6b41d-7fcb-40b6-9a72-229cdd00a094",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdin",
+     "output_type": "stream",
+     "text": [
+      "OPENAI_API_KEY:  ········\n"
+     ]
+    }
+   ],
+   "source": [
+    "import getpass\n",
+    "import os\n",
+    "\n",
+    "\n",
+    "def _set_env(var: str):\n",
+    "    if not os.environ.get(var):\n",
+    "        os.environ[var] = getpass.getpass(f\"{var}: \")\n",
+    "\n",
+    "\n",
+    "_set_env(\"OPENAI_API_KEY\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e3d02ebb-c2e1-4ef7-b187-810d55139317",
+   "metadata": {},
+   "source": [
+    "## Define model, tools and graph"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3ba684f1-d46b-42e4-95cf-9685209a5992",
+   "metadata": {},
+   "source": [
+    "### Define a node that will call OpenAI API"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "d59234f9-173e-469d-a725-c13e0979663e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from openai import AsyncOpenAI\n",
+    "from langchain_core.language_models.chat_models import ChatGenerationChunk\n",
+    "from langchain_core.messages import AIMessageChunk\n",
+    "from langchain_core.runnables.config import (\n",
+    "    ensure_config,\n",
+    "    get_callback_manager_for_config,\n",
+    ")\n",
+    "\n",
+    "openai_client = AsyncOpenAI()\n",
+    "# define tool schema for openai tool calling\n",
+    "\n",
+    "tool = {\n",
+    "    \"type\": \"function\",\n",
+    "    \"function\": {\n",
+    "        \"name\": \"get_items\",\n",
+    "        \"description\": \"Use this tool to look up which items are in the given place.\",\n",
+    "        \"parameters\": {\n",
+    "            \"type\": \"object\",\n",
+    "            \"properties\": {\"place\": {\"type\": \"string\"}},\n",
+    "            \"required\": [\"place\"],\n",
+    "        },\n",
+    "    },\n",
+    "}\n",
+    "\n",
+    "\n",
+    "async def call_model(state, config=None):\n",
+    "    config = ensure_config(config | {\"tags\": [\"agent_llm\"]})\n",
+    "    callback_manager = get_callback_manager_for_config(config)\n",
+    "    messages = state[\"messages\"]\n",
+    "\n",
+    "    llm_run_manager = callback_manager.on_chat_model_start({}, [messages])[0]\n",
+    "    response = await openai_client.chat.completions.create(\n",
+    "        messages=messages, model=\"gpt-3.5-turbo\", tools=[tool], stream=True\n",
+    "    )\n",
+    "\n",
+    "    response_content = \"\"\n",
+    "    role = None\n",
+    "\n",
+    "    tool_call_id = None\n",
+    "    tool_call_function_name = None\n",
+    "    tool_call_function_arguments = \"\"\n",
+    "    async for chunk in response:\n",
+    "        delta = chunk.choices[0].delta\n",
+    "        if delta.role is not None:\n",
+    "            role = delta.role\n",
+    "\n",
+    "        if delta.content:\n",
+    "            response_content += delta.content\n",
+    "            llm_run_manager.on_llm_new_token(delta.content)\n",
+    "\n",
+    "        if delta.tool_calls:\n",
+    "            # note: for simplicity we're only handling a single tool call here\n",
+    "            if delta.tool_calls[0].function.name is not None:\n",
+    "                tool_call_function_name = delta.tool_calls[0].function.name\n",
+    "                tool_call_id = delta.tool_calls[0].id\n",
+    "\n",
+    "            # note: we're wrapping the tools calls in ChatGenerationChunk so that the events from .astream_events in the graph can render tool calls correctly\n",
+    "            tool_call_chunk = ChatGenerationChunk(\n",
+    "                message=AIMessageChunk(\n",
+    "                    content=\"\",\n",
+    "                    additional_kwargs={\"tool_calls\": [delta.tool_calls[0].dict()]},\n",
+    "                )\n",
+    "            )\n",
+    "            llm_run_manager.on_llm_new_token(\"\", chunk=tool_call_chunk)\n",
+    "            tool_call_function_arguments += delta.tool_calls[0].function.arguments\n",
+    "\n",
+    "    if tool_call_function_name is not None:\n",
+    "        tool_calls = [\n",
+    "            {\n",
+    "                \"id\": tool_call_id,\n",
+    "                \"function\": {\n",
+    "                    \"name\": tool_call_function_name,\n",
+    "                    \"arguments\": tool_call_function_arguments,\n",
+    "                },\n",
+    "                \"type\": \"function\",\n",
+    "            }\n",
+    "        ]\n",
+    "    else:\n",
+    "        tool_calls = None\n",
+    "\n",
+    "    response_message = {\n",
+    "        \"role\": role,\n",
+    "        \"content\": response_content,\n",
+    "        \"tool_calls\": tool_calls,\n",
+    "    }\n",
+    "    return {\"messages\": [response_message]}"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3a3877e8-8ace-40d5-ad04-cbf21c6f3250",
+   "metadata": {},
+   "source": [
+    "### Define our tools and a tool-calling node"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "b90941d8-afe4-42ec-9262-9c3b87c3b1ec",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import json\n",
+    "from langchain_core.callbacks import adispatch_custom_event\n",
+    "\n",
+    "\n",
+    "async def get_items(place: str) -> str:\n",
+    "    \"\"\"Use this tool to look up which items are in the given place.\"\"\"\n",
+    "\n",
+    "    # this can be replaced with any actual streaming logic that you might have\n",
+    "    def stream(place: str):\n",
+    "        if \"bed\" in place:  # For under the bed\n",
+    "            yield from [\"socks\", \"shoes\", \"dust bunnies\"]\n",
+    "        elif \"shelf\" in place:  # For 'shelf'\n",
+    "            yield from [\"books\", \"penciles\", \"pictures\"]\n",
+    "        else:  # if the agent decides to ask about a different place\n",
+    "            yield \"cat snacks\"\n",
+    "\n",
+    "    tokens = []\n",
+    "    for token in stream(place):\n",
+    "        await adispatch_custom_event(\n",
+    "            # this will allow you to filter events by name\n",
+    "            \"tool_call_token_stream\",\n",
+    "            {\n",
+    "                \"function_name\": \"get_items\",\n",
+    "                \"arguments\": {\"place\": place},\n",
+    "                \"tool_output_token\": token,\n",
+    "            },\n",
+    "            # this will allow you to filter events by tags\n",
+    "            config={\"tags\": [\"tool_call\"]},\n",
+    "        )\n",
+    "        tokens.append(token)\n",
+    "\n",
+    "    return \", \".join(tokens)\n",
+    "\n",
+    "\n",
+    "# define mapping to look up functions when running tools\n",
+    "function_name_to_function = {\"get_items\": get_items}\n",
+    "\n",
+    "\n",
+    "async def call_tools(state):\n",
+    "    messages = state[\"messages\"]\n",
+    "\n",
+    "    tool_call = messages[-1][\"tool_calls\"][0]\n",
+    "    function_name = tool_call[\"function\"][\"name\"]\n",
+    "    function_arguments = tool_call[\"function\"][\"arguments\"]\n",
+    "    arguments = json.loads(function_arguments)\n",
+    "\n",
+    "    function_response = await function_name_to_function[function_name](**arguments)\n",
+    "    tool_message = {\n",
+    "        \"tool_call_id\": tool_call[\"id\"],\n",
+    "        \"role\": \"tool\",\n",
+    "        \"name\": function_name,\n",
+    "        \"content\": function_response,\n",
+    "    }\n",
+    "    return {\"messages\": [tool_message]}"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6685898c-9a1c-4803-a492-bd70574ebe38",
+   "metadata": {},
+   "source": [
+    "### Define our graph"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "228260be-1f9a-4195-80e0-9604f8a5dba6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import operator\n",
+    "from typing import Annotated, TypedDict, Literal\n",
+    "\n",
+    "from langgraph.graph import StateGraph, END\n",
+    "\n",
+    "\n",
+    "class State(TypedDict):\n",
+    "    messages: Annotated[list, operator.add]\n",
+    "\n",
+    "\n",
+    "def should_continue(state) -> Literal[\"tools\", END]:\n",
+    "    messages = state[\"messages\"]\n",
+    "    last_message = messages[-1]\n",
+    "    if last_message[\"tool_calls\"]:\n",
+    "        return \"tools\"\n",
+    "    return END\n",
+    "\n",
+    "\n",
+    "workflow = StateGraph(State)\n",
+    "workflow.set_entry_point(\"model\")\n",
+    "workflow.add_node(\"model\", call_model)  # i.e. our \"agent\"\n",
+    "workflow.add_node(\"tools\", call_tools)\n",
+    "workflow.add_conditional_edges(\"model\", should_continue)\n",
+    "workflow.add_edge(\"tools\", \"model\")\n",
+    "graph = workflow.compile()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d046e2ef-f208-4831-ab31-203b2e75a49a",
+   "metadata": {},
+   "source": [
+    "## Stream tokens from within the tool"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "45c96a79-4147-42e3-89fd-d942b2b49f6c",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/Users/vadymbarda/.virtualenvs/langgraph/lib/python3.11/site-packages/langchain_core/_api/beta_decorator.py:87: LangChainBetaWarning: This API is in beta and may change in the future.\n",
+      "  warn_beta(\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Tool token socks\n",
+      "Tool token shoes\n",
+      "Tool token dust bunnies\n"
+     ]
+    }
+   ],
+   "source": [
+    "async for event in graph.astream_events(\n",
+    "    {\"messages\": [{\"role\": \"user\", \"content\": \"what's in the bedroom\"}]}, version=\"v2\"\n",
+    "):\n",
+    "    tags = event.get(\"tags\", [])\n",
+    "    if event[\"event\"] == \"on_custom_event\" and \"tool_call\" in tags:\n",
+    "        print(\"Tool token\", event[\"data\"][\"tool_output_token\"])"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "langgraph",
+   "language": "python",
+   "name": "langgraph"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

examples/streaming-tokens-without-langchain.ipynb

@@ -0,0 +1,370 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "b23ced4e-dc29-43be-9f94-0c36bb181b8a",
+   "metadata": {},
+   "source": [
+    "# How to stream LLM tokens (without LangChain LLMs)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7044eeb8-4074-4f9c-8a62-962488744557",
+   "metadata": {},
+   "source": [
+    "In this example we will stream tokens from the language model powering an agent. We'll be using OpenAI client library directly, without using LangChain chat models. We will also use a ReAct agent as an example."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a37f60af-43ea-4aa6-847a-df8cc47065f5",
+   "metadata": {},
+   "source": [
+    "## Setup"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "47f79af8-58d8-4a48-8d9a-88823d88701f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%capture --no-stderr\n",
+    "%pip install -U langgraph openai"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "0cf6b41d-7fcb-40b6-9a72-229cdd00a094",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "OPENAI_API_KEY:  ········\n"
+     ]
+    }
+   ],
+   "source": [
+    "import getpass\n",
+    "import os\n",
+    "\n",
+    "\n",
+    "def _set_env(var: str):\n",
+    "    if not os.environ.get(var):\n",
+    "        os.environ[var] = getpass.getpass(f\"{var}: \")\n",
+    "\n",
+    "\n",
+    "_set_env(\"OPENAI_API_KEY\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e3d02ebb-c2e1-4ef7-b187-810d55139317",
+   "metadata": {},
+   "source": [
+    "## Define model, tools and graph"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3ba684f1-d46b-42e4-95cf-9685209a5992",
+   "metadata": {},
+   "source": [
+    "### Define a node that will call OpenAI API"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d59234f9-173e-469d-a725-c13e0979663e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from openai import AsyncOpenAI\n",
+    "from langchain_core.language_models.chat_models import ChatGenerationChunk\n",
+    "from langchain_core.messages import AIMessageChunk\n",
+    "from langchain_core.runnables.config import (\n",
+    "    ensure_config,\n",
+    "    get_callback_manager_for_config,\n",
+    ")\n",
+    "\n",
+    "openai_client = AsyncOpenAI()\n",
+    "# define tool schema for openai tool calling\n",
+    "\n",
+    "tool = {\n",
+    "    \"type\": \"function\",\n",
+    "    \"function\": {\n",
+    "        \"name\": \"get_items\",\n",
+    "        \"description\": \"Use this tool to look up which items are in the given place.\",\n",
+    "        \"parameters\": {\n",
+    "            \"type\": \"object\",\n",
+    "            \"properties\": {\"place\": {\"type\": \"string\"}},\n",
+    "            \"required\": [\"place\"],\n",
+    "        },\n",
+    "    },\n",
+    "}\n",
+    "\n",
+    "\n",
+    "async def call_model(state, config=None):\n",
+    "    config = ensure_config(config | {\"tags\": [\"agent_llm\"]})\n",
+    "    callback_manager = get_callback_manager_for_config(config)\n",
+    "    messages = state[\"messages\"]\n",
+    "\n",
+    "    llm_run_manager = callback_manager.on_chat_model_start({}, [messages])[0]\n",
+    "    response = await openai_client.chat.completions.create(\n",
+    "        messages=messages, model=\"gpt-3.5-turbo\", tools=[tool], stream=True\n",
+    "    )\n",
+    "\n",
+    "    response_content = \"\"\n",
+    "    role = None\n",
+    "\n",
+    "    tool_call_id = None\n",
+    "    tool_call_function_name = None\n",
+    "    tool_call_function_arguments = \"\"\n",
+    "    async for chunk in response:\n",
+    "        delta = chunk.choices[0].delta\n",
+    "        if delta.role is not None:\n",
+    "            role = delta.role\n",
+    "\n",
+    "        if delta.content:\n",
+    "            response_content += delta.content\n",
+    "            llm_run_manager.on_llm_new_token(delta.content)\n",
+    "\n",
+    "        if delta.tool_calls:\n",
+    "            # note: for simplicity we're only handling a single tool call here\n",
+    "            if delta.tool_calls[0].function.name is not None:\n",
+    "                tool_call_function_name = delta.tool_calls[0].function.name\n",
+    "                tool_call_id = delta.tool_calls[0].id\n",
+    "\n",
+    "            # note: we're wrapping the tools calls in ChatGenerationChunk so that the events from .astream_events in the graph can render tool calls correctly\n",
+    "            tool_call_chunk = ChatGenerationChunk(\n",
+    "                message=AIMessageChunk(\n",
+    "                    content=\"\",\n",
+    "                    additional_kwargs={\"tool_calls\": [delta.tool_calls[0].dict()]},\n",
+    "                )\n",
+    "            )\n",
+    "            llm_run_manager.on_llm_new_token(\"\", chunk=tool_call_chunk)\n",
+    "            tool_call_function_arguments += delta.tool_calls[0].function.arguments\n",
+    "\n",
+    "    if tool_call_function_name is not None:\n",
+    "        tool_calls = [\n",
+    "            {\n",
+    "                \"id\": tool_call_id,\n",
+    "                \"function\": {\n",
+    "                    \"name\": tool_call_function_name,\n",
+    "                    \"arguments\": tool_call_function_arguments,\n",
+    "                },\n",
+    "                \"type\": \"function\",\n",
+    "            }\n",
+    "        ]\n",
+    "    else:\n",
+    "        tool_calls = None\n",
+    "\n",
+    "    response_message = {\n",
+    "        \"role\": role,\n",
+    "        \"content\": response_content,\n",
+    "        \"tool_calls\": tool_calls,\n",
+    "    }\n",
+    "    return {\"messages\": [response_message]}"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3a3877e8-8ace-40d5-ad04-cbf21c6f3250",
+   "metadata": {},
+   "source": [
+    "### Define our tools and a tool-calling node"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b756ea32",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import json\n",
+    "\n",
+    "\n",
+    "async def get_items(place: str) -> str:\n",
+    "    \"\"\"Use this tool to look up which items are in the given place.\"\"\"\n",
+    "    if \"bed\" in place:  # For under the bed\n",
+    "        return \"socks, shoes and dust bunnies\"\n",
+    "    if \"shelf\" in place:  # For 'shelf'\n",
+    "        return \"books, penciles and pictures\"\n",
+    "    else:  # if the agent decides to ask about a different place\n",
+    "        return \"cat snacks\"\n",
+    "\n",
+    "\n",
+    "# define mapping to look up functions when running tools\n",
+    "function_name_to_function = {\"get_items\": get_items}\n",
+    "\n",
+    "\n",
+    "async def call_tools(state):\n",
+    "    messages = state[\"messages\"]\n",
+    "\n",
+    "    tool_call = messages[-1][\"tool_calls\"][0]\n",
+    "    function_name = tool_call[\"function\"][\"name\"]\n",
+    "    function_arguments = tool_call[\"function\"][\"arguments\"]\n",
+    "    arguments = json.loads(function_arguments)\n",
+    "\n",
+    "    function_response = await function_name_to_function[function_name](**arguments)\n",
+    "    tool_message = {\n",
+    "        \"tool_call_id\": tool_call[\"id\"],\n",
+    "        \"role\": \"tool\",\n",
+    "        \"name\": function_name,\n",
+    "        \"content\": function_response,\n",
+    "    }\n",
+    "    return {\"messages\": [tool_message]}"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6685898c-9a1c-4803-a492-bd70574ebe38",
+   "metadata": {},
+   "source": [
+    "### Define our graph"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "228260be-1f9a-4195-80e0-9604f8a5dba6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import operator\n",
+    "from typing import Annotated, TypedDict, Literal\n",
+    "\n",
+    "from langgraph.graph import StateGraph, END\n",
+    "\n",
+    "\n",
+    "class State(TypedDict):\n",
+    "    messages: Annotated[list, operator.add]\n",
+    "\n",
+    "\n",
+    "def should_continue(state) -> Literal[\"tools\", END]:\n",
+    "    messages = state[\"messages\"]\n",
+    "    last_message = messages[-1]\n",
+    "    if last_message[\"tool_calls\"]:\n",
+    "        return \"tools\"\n",
+    "    return END\n",
+    "\n",
+    "\n",
+    "workflow = StateGraph(State)\n",
+    "workflow.set_entry_point(\"model\")\n",
+    "workflow.add_node(\"model\", call_model)  # i.e. our \"agent\"\n",
+    "workflow.add_node(\"tools\", call_tools)\n",
+    "workflow.add_conditional_edges(\"model\", should_continue)\n",
+    "workflow.add_edge(\"tools\", \"model\")\n",
+    "graph = workflow.compile()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d046e2ef-f208-4831-ab31-203b2e75a49a",
+   "metadata": {},
+   "source": [
+    "## Stream tokens"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "45c96a79-4147-42e3-89fd-d942b2b49f6c",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/Users/vadymbarda/.virtualenvs/langgraph/lib/python3.11/site-packages/langchain_core/_api/beta_decorator.py:87: LangChainBetaWarning: This API is in beta and may change in the future.\n",
+      "  warn_beta(\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "LLM token {'content': '', 'additional_kwargs': {'tool_calls': [{'index': 0, 'id': 'call_xUcx3IPa8GREPOpjHVj5k9Wx', 'function': {'arguments': '', 'name': 'get_items'}, 'type': 'function'}]}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [{'name': 'get_items', 'args': '', 'id': 'call_xUcx3IPa8GREPOpjHVj5k9Wx', 'error': None}], 'usage_metadata': None, 'tool_call_chunks': [{'name': 'get_items', 'args': '', 'id': 'call_xUcx3IPa8GREPOpjHVj5k9Wx', 'index': 0}]}\n",
+      "LLM token {'content': '', 'additional_kwargs': {'tool_calls': [{'index': 0, 'id': None, 'function': {'arguments': '{\"', 'name': None}, 'type': None}]}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [{'name': '', 'args': {}, 'id': None}], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': [{'name': None, 'args': '{\"', 'id': None, 'index': 0}]}\n",
+      "LLM token {'content': '', 'additional_kwargs': {'tool_calls': [{'index': 0, 'id': None, 'function': {'arguments': 'place', 'name': None}, 'type': None}]}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [{'name': None, 'args': 'place', 'id': None, 'error': None}], 'usage_metadata': None, 'tool_call_chunks': [{'name': None, 'args': 'place', 'id': None, 'index': 0}]}\n",
+      "LLM token {'content': '', 'additional_kwargs': {'tool_calls': [{'index': 0, 'id': None, 'function': {'arguments': '\":\"', 'name': None}, 'type': None}]}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [{'name': None, 'args': '\":\"', 'id': None, 'error': None}], 'usage_metadata': None, 'tool_call_chunks': [{'name': None, 'args': '\":\"', 'id': None, 'index': 0}]}\n",
+      "LLM token {'content': '', 'additional_kwargs': {'tool_calls': [{'index': 0, 'id': None, 'function': {'arguments': 'bed', 'name': None}, 'type': None}]}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [{'name': None, 'args': 'bed', 'id': None, 'error': None}], 'usage_metadata': None, 'tool_call_chunks': [{'name': None, 'args': 'bed', 'id': None, 'index': 0}]}\n",
+      "LLM token {'content': '', 'additional_kwargs': {'tool_calls': [{'index': 0, 'id': None, 'function': {'arguments': 'room', 'name': None}, 'type': None}]}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [{'name': None, 'args': 'room', 'id': None, 'error': None}], 'usage_metadata': None, 'tool_call_chunks': [{'name': None, 'args': 'room', 'id': None, 'index': 0}]}\n",
+      "LLM token {'content': '', 'additional_kwargs': {'tool_calls': [{'index': 0, 'id': None, 'function': {'arguments': '\"}', 'name': None}, 'type': None}]}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [{'name': None, 'args': '\"}', 'id': None, 'error': None}], 'usage_metadata': None, 'tool_call_chunks': [{'name': None, 'args': '\"}', 'id': None, 'index': 0}]}\n",
+      "LLM token {'content': 'In', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
+      "LLM token {'content': ' the', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
+      "LLM token {'content': ' bedroom', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
+      "LLM token {'content': ',', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
+      "LLM token {'content': ' you', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
+      "LLM token {'content': ' have', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
+      "LLM token {'content': ' socks', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
+      "LLM token {'content': ',', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
+      "LLM token {'content': ' shoes', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
+      "LLM token {'content': ',', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
+      "LLM token {'content': ' and', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
+      "LLM token {'content': ' some', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
+      "LLM token {'content': ' dust', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
+      "LLM token {'content': ' b', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
+      "LLM token {'content': 'unn', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
+      "LLM token {'content': 'ies', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
+      "LLM token {'content': '.', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
+      "LLM token {'content': ' Is', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
+      "LLM token {'content': ' there', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
+      "LLM token {'content': ' anything', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
+      "LLM token {'content': ' else', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
+      "LLM token {'content': ' you', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
+      "LLM token {'content': ' would', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
+      "LLM token {'content': ' like', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
+      "LLM token {'content': ' to', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
+      "LLM token {'content': ' know', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n",
+      "LLM token {'content': '?', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'AIMessageChunk', 'name': None, 'id': None, 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None, 'tool_call_chunks': []}\n"
+     ]
+    }
+   ],
+   "source": [
+    "async for event in graph.astream_events(\n",
+    "    {\"messages\": [{\"role\": \"user\", \"content\": \"what's in the bedroom\"}]}, version=\"v2\"\n",
+    "):\n",
+    "    tags = event.get(\"tags\", [])\n",
+    "    if event[\"event\"] == \"on_chat_model_stream\" and \"agent_llm\" in tags:\n",
+    "        print(\"LLM token\", event[\"data\"][\"chunk\"].dict())"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "adb0f7bc-6e51-478e-bd32-8f72df072d6c",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "langgraph",
+   "language": "python",
+   "name": "langgraph"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

libs/cli/examples/poetry.lock

@@ -215,13 +215,13 @@ css = ["tinycss2 (>=1.1.0,<1.3)"]
 
 [[package]]
 name = "certifi"
-version = "2024.6.2"
+version = "2024.7.4"
 description = "Python package for providing Mozilla's CA Bundle."
 optional = false
 python-versions = ">=3.6"
 files = [
-    {file = "certifi-2024.6.2-py3-none-any.whl", hash = "sha256:ddc6c8ce995e6987e7faf5e3f1b02b302836a0e5d98ece18392cb1a36c72ad56"},
-    {file = "certifi-2024.6.2.tar.gz", hash = "sha256:3cd43f1c6fa7dedc5899d69d3ad0398fd018ad1a17fba83ddaf78aa46c747516"},
+    {file = "certifi-2024.7.4-py3-none-any.whl", hash = "sha256:c198e21b1289c2ab85ee4e67bb4b4ef3ead0892059901a8d5b622f24a1101e90"},
+    {file = "certifi-2024.7.4.tar.gz", hash = "sha256:5a1e7645bc0ec61a09e26c36f6106dd4cf40c6db3a1fb6352b0244e7fb057c7b"},
 ]
 
 [[package]]
@@ -1105,7 +1105,7 @@ files = [
 
 [[package]]
 name = "langgraph-cli"
-version = "0.1.39"
+version = "0.1.48"
 description = "CLI for interacting with LangGraph API"
 optional = false
 python-versions = "^3.9.0,<4.0"
@@ -1117,7 +1117,7 @@ click = "^8.1.7"
 
 [package.source]
 type = "directory"
-url = ".."
+url = "../../cli"
 
 [[package]]
 name = "langgraph-sdk"

libs/cli/langgraph_cli/cli.py

@@ -452,6 +452,21 @@ def build(
         _build(runner, set, config, config_json, platform, base_image, pull, tag)
 
 
+@OPT_CONFIG
+@click.argument("save_path", type=click.Path(resolve_path=True))
+@cli.command(help="Generate a Dockerfile for langgraph API server")
+@log_command
+def dockerfile(save_path: pathlib.Path, config: pathlib.Path):
+    with open(config) as f:
+        config_json = langgraph_cli.config.validate_config(json.load(f))
+    with open(save_path, "w") as f:
+        f.write(
+            langgraph_cli.config.config_to_docker(
+                config, config_json, "langchain/langgraph-api"
+            )
+        )
+
+
 def prepare_args_and_stdin(
     *,
     capabilities: DockerCapabilities,

libs/cli/langgraph_cli/config.py

@@ -193,7 +193,9 @@ def _update_graph_paths(
 
 def config_to_docker(config_path: pathlib.Path, config: Config, base_image: str):
     # configure pip
-    pip_install = "pip install -c /api/constraints.txt"
+    pip_install = (
+        "PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt"
+    )
     if config.get("pip_config_file"):
         pip_install = f"PIP_CONFIG_FILE=/pipconfig.txt {pip_install}"
     pip_config_file_str = (
@@ -230,8 +232,7 @@ RUN set -ex && \\
                 '[tool.setuptools.package-data]' \\
                 '"*" = ["**/*"]'; do \\
         echo "$line" >> /deps/__outer_{fullpath.name}/pyproject.toml; \\
-    done
-"""
+    done"""
         for fullpath, (relpath, destpath) in local_deps.faux_pkgs.items()
     )
     local_pkgs_str = os.linesep.join(
@@ -239,19 +240,24 @@ RUN set -ex && \\
         for fullpath, relpath in local_deps.real_pkgs.items()
     )
 
+    installs = f"{os.linesep}{os.linesep}".join(
+        filter(
+            None,
+            [
+                pip_config_file_str,
+                pip_pkgs_str,
+                pip_reqs_str,
+                local_pkgs_str,
+                faux_pkgs_str,
+            ],
+        )
+    )
+
     return f"""FROM {base_image}:{config['python_version']}
 
 {os.linesep.join(config["dockerfile_lines"])}
 
-{pip_config_file_str}
-
-{pip_pkgs_str}
-
-{pip_reqs_str}
-
-{local_pkgs_str}
-
-{faux_pkgs_str}
+{installs}
 
 RUN {pip_install} -e /deps/*
 

libs/cli/langgraph_cli/exec.py

@@ -131,21 +131,42 @@ async def monitor_stream(
     if collect:
         ba = bytearray()
 
-    def handle(line: bytes):
+    def handle(line: bytes, overrun: bool):
         nonlocal on_line
         nonlocal display
 
+        if display:
+            sys.stdout.buffer.write(line)
+        if overrun:
+            return
         if collect:
             ba.extend(line)
-        if display:
-            sys.stdout.write(line.decode())
         if on_line:
             if on_line(line.decode()):
                 on_line = None
                 display = True
 
-    async for line in stream:
-        await asyncio.to_thread(handle, line)
+    """Adapted from asyncio.StreamReader.readline() to handle LimitOverrunError."""
+    sep = b"\n"
+    seplen = len(sep)
+    while True:
+        try:
+            line = await stream.readuntil(sep)
+            overrun = False
+        except asyncio.IncompleteReadError as e:
+            line = e.partial
+            overrun = False
+        except asyncio.LimitOverrunError as e:
+            if stream._buffer.startswith(sep, e.consumed):
+                line = stream._buffer[: e.consumed + seplen]
+            else:
+                line = stream._buffer.clear()
+            overrun = True
+            stream._maybe_resume_transport()
+        await asyncio.to_thread(handle, line, overrun)
+        if line == b"":
+            break
+
     if collect:
         return ba
     else:

libs/cli/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph-cli"
-version = "0.1.48"
+version = "0.1.50"
 description = "CLI for interacting with LangGraph API"
 authors = []
 license = "MIT"

libs/cli/tests/unit_tests/test_cli.py

@@ -92,7 +92,7 @@ services:
             dockerfile_inline: |
                 FROM langchain/langgraph-api:3.11
                 ADD . /deps/
-                RUN pip install -c /api/constraints.txt -e /deps/*
+                RUN PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/*
                 ENV LANGSERVE_GRAPHS='{{"agent": "agent.py:graph"}}'
                 WORKDIR /deps/
         

libs/cli/tests/unit_tests/test_config.json

@@ -1,13 +1,15 @@
 {
     "python_version": "3.12",
     "pip_config_file": "pipconfig.txt",
-    "dockerfile_lines": ["ARG meow"],
+    "dockerfile_lines": [
+        "ARG meow=woof"
+    ],
     "dependencies": [
         "langchain_openai",
         "."
     ],
     "graphs": {
-        "agent": "tests/unit_tests/agent.py:graph"
+        "agent": "graphs/agent.py:graph"
     },
     "env": ".env"
-}
+}

libs/cli/tests/unit_tests/test_config.py

@@ -81,7 +81,7 @@ RUN set -ex && \\
                 '"*" = ["**/*"]'; do \\
         echo "$line" >> /deps/__outer_unit_tests/pyproject.toml; \\
     done
-RUN pip install -c /api/constraints.txt -e /deps/*
+RUN PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/*
 ENV LANGSERVE_GRAPHS='{"agent": "/deps/__outer_unit_tests/unit_tests/agent.py:graph"}'
 WORKDIR /deps/__outer_unit_tests/unit_tests\
 """
@@ -113,7 +113,7 @@ RUN set -ex && \\
                 '"*" = ["**/*"]'; do \\
         echo "$line" >> /deps/__outer_unit_tests/pyproject.toml; \\
     done
-RUN PIP_CONFIG_FILE=/pipconfig.txt pip install -c /api/constraints.txt -e /deps/*
+RUN PIP_CONFIG_FILE=/pipconfig.txt PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/*
 ENV LANGSERVE_GRAPHS='{"agent": "/deps/__outer_unit_tests/unit_tests/agent.py:graph"}'
 WORKDIR /deps/__outer_unit_tests/unit_tests\
 """
@@ -163,7 +163,7 @@ RUN set -ex && \\
                 '"*" = ["**/*"]'; do \\
         echo "$line" >> /deps/__outer_graphs/pyproject.toml; \\
     done
-RUN pip install -c /api/constraints.txt -e /deps/*
+RUN PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/*
 ENV LANGSERVE_GRAPHS='{"agent": "/deps/__outer_graphs/src/agent.py:graph"}'\
 """
     assert clean_empty_lines(actual_docker_stdin) == expected_docker_stdin
@@ -192,7 +192,7 @@ dependencies = ["langchain"]"""
     os.remove(pyproject_path)
     expected_docker_stdin = """FROM langchain/langgraph-api:3.11
 ADD . /deps/unit_tests
-RUN pip install -c /api/constraints.txt -e /deps/*
+RUN PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/*
 ENV LANGSERVE_GRAPHS='{"agent": "/deps/unit_tests/graphs/agent.py:graph"}'
 WORKDIR /deps/unit_tests"""
     assert clean_empty_lines(actual_docker_stdin) == expected_docker_stdin
@@ -217,7 +217,7 @@ def test_config_to_docker_end_to_end():
 ARG meow
 ARG foo
 ADD pipconfig.txt /pipconfig.txt
-RUN PIP_CONFIG_FILE=/pipconfig.txt pip install -c /api/constraints.txt langchain langchain_openai
+RUN PIP_CONFIG_FILE=/pipconfig.txt PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt langchain langchain_openai
 ADD ./graphs/ /deps/__outer_graphs/src
 RUN set -ex && \\
     for line in '[project]' \\
@@ -227,7 +227,7 @@ RUN set -ex && \\
                 '"*" = ["**/*"]'; do \\
         echo "$line" >> /deps/__outer_graphs/pyproject.toml; \\
     done
-RUN PIP_CONFIG_FILE=/pipconfig.txt pip install -c /api/constraints.txt -e /deps/*
+RUN PIP_CONFIG_FILE=/pipconfig.txt PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/*
 ENV LANGSERVE_GRAPHS='{"agent": "/deps/__outer_graphs/src/agent.py:graph"}'"""
     assert clean_empty_lines(actual_docker_stdin) == expected_docker_stdin
 
@@ -251,7 +251,7 @@ def test_config_to_compose_simple_config():
                                 '"*" = ["**/*"]'; do \\
                         echo "$line" >> /deps/__outer_unit_tests/pyproject.toml; \\
                     done
-                RUN pip install -c /api/constraints.txt -e /deps/*
+                RUN PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/*
                 ENV LANGSERVE_GRAPHS='{"agent": "/deps/__outer_unit_tests/unit_tests/agent.py:graph"}'
                 WORKDIR /deps/__outer_unit_tests/unit_tests
         """
@@ -281,7 +281,7 @@ def test_config_to_compose_env_vars():
                                 '"*" = ["**/*"]'; do \\
                         echo "$line" >> /deps/__outer_unit_tests/pyproject.toml; \\
                     done
-                RUN pip install -c /api/constraints.txt -e /deps/*
+                RUN PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/*
                 ENV LANGSERVE_GRAPHS='{"agent": "/deps/__outer_unit_tests/unit_tests/agent.py:graph"}'
                 WORKDIR /deps/__outer_unit_tests/unit_tests
         """
@@ -318,7 +318,7 @@ def test_config_to_compose_env_file():
                                 '"*" = ["**/*"]'; do \\
                         echo "$line" >> /deps/__outer_unit_tests/pyproject.toml; \\
                     done
-                RUN pip install -c /api/constraints.txt -e /deps/*
+                RUN PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/*
                 ENV LANGSERVE_GRAPHS='{"agent": "/deps/__outer_unit_tests/unit_tests/agent.py:graph"}'
                 WORKDIR /deps/__outer_unit_tests/unit_tests
         """
@@ -348,7 +348,7 @@ def test_config_to_compose_watch():
                                 '"*" = ["**/*"]'; do \\
                         echo "$line" >> /deps/__outer_unit_tests/pyproject.toml; \\
                     done
-                RUN pip install -c /api/constraints.txt -e /deps/*
+                RUN PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/*
                 ENV LANGSERVE_GRAPHS='{"agent": "/deps/__outer_unit_tests/unit_tests/agent.py:graph"}'
                 WORKDIR /deps/__outer_unit_tests/unit_tests
         
@@ -387,7 +387,7 @@ def test_config_to_compose_end_to_end():
                                 '"*" = ["**/*"]'; do \\
                         echo "$line" >> /deps/__outer_unit_tests/pyproject.toml; \\
                     done
-                RUN pip install -c /api/constraints.txt -e /deps/*
+                RUN PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt -e /deps/*
                 ENV LANGSERVE_GRAPHS='{"agent": "/deps/__outer_unit_tests/unit_tests/agent.py:graph"}'
                 WORKDIR /deps/__outer_unit_tests/unit_tests
         

libs/langgraph/langgraph/_api/deprecation.py

@@ -11,21 +11,21 @@ F = TypeVar("F", bound=Callable[..., Any])
 
 
 def deprecated(
-    version: str, alternative: str, *, removal: str = "", example: str = ""
+    since: str, alternative: str, *, removal: str = "", example: str = ""
 ) -> Callable[[F], F]:
     def decorator(func: F) -> F:
         @functools.wraps(func)
         def wrapper(*args: Any, **kwargs: Any) -> Any:
             removal_str = removal if removal else "a future version"
             message = (
-                f"{func.__name__} is deprecated as of version {version} and will be"
+                f"{func.__name__} is deprecated as of version {since} and will be"
                 f" removed in {removal_str}. Use {alternative} instead.{example}"
             )
             warnings.warn(message, LangGraphDeprecationWarning, stacklevel=2)
             return func(*args, **kwargs)
 
         docstring = (
-            f"**Deprecated**: This function is deprecated as of version {version}. "
+            f"**Deprecated**: This function is deprecated as of version {since}. "
             f"Use `{alternative}` instead."
         )
         if func.__doc__:
@@ -35,3 +35,24 @@ def deprecated(
         return cast(F, wrapper)
 
     return decorator
+
+
+def deprecated_parameter(
+    arg_name: str, since: str, alternative: str, *, removal: str
+) -> Callable[[F], F]:
+    def decorator(func: F) -> F:
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            if arg_name in kwargs:
+                warnings.warn(
+                    f"Parameter '{arg_name}' in function '{func.__name__}' is "
+                    f"deprecated as of version {since} and will be removed in version {removal}. "
+                    f"Use '{alternative}' parameter instead.",
+                    category=LangGraphDeprecationWarning,
+                    stacklevel=2,
+                )
+            return func(*args, **kwargs)
+
+        return cast(F, wrapper)
+
+    return decorator

libs/langgraph/langgraph/channels/any_value.py

@@ -15,6 +15,9 @@ class AnyValue(Generic[Value], BaseChannel[Value, Value, Value]):
     def __init__(self, typ: Type[Value]) -> None:
         self.typ = typ
 
+    def __eq__(self, value: object) -> bool:
+        return isinstance(value, AnyValue)
+
     @property
     def ValueType(self) -> Type[Value]:
         """The type of the value stored in the channel."""

libs/langgraph/langgraph/channels/binop.py

@@ -55,6 +55,12 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]):
         except Exception:
             pass
 
+    def __eq__(self, value: object) -> bool:
+        return (
+            isinstance(value, BinaryOperatorAggregate)
+            and value.operator == self.operator
+        )
+
     @property
     def ValueType(self) -> Type[Value]:
         """The type of the value stored in the channel."""

libs/langgraph/langgraph/channels/context.py

@@ -47,6 +47,13 @@ class Context(Generic[Value], BaseChannel[Value, None, None]):
         self.ctx = ctx
         self.actx = actx
 
+    def __eq__(self, value: object) -> bool:
+        return (
+            isinstance(value, Context)
+            and value.ctx == self.ctx
+            and value.actx == self.actx
+        )
+
     @property
     def ValueType(self) -> Any:
         """The type of the value stored in the channel."""

libs/langgraph/langgraph/channels/dynamic_barrier_value.py

@@ -32,6 +32,9 @@ class DynamicBarrierValue(
         self.names = None
         self.seen = set()
 
+    def __eq__(self, value: object) -> bool:
+        return isinstance(value, DynamicBarrierValue) and value.names == self.names
+
     @property
     def ValueType(self) -> Type[Value]:
         """The type of the value stored in the channel."""

libs/langgraph/langgraph/channels/ephemeral_value.py

@@ -15,6 +15,9 @@ class EphemeralValue(Generic[Value], BaseChannel[Value, Value, Value]):
         self.typ = typ
         self.guard = guard
 
+    def __eq__(self, value: object) -> bool:
+        return isinstance(value, EphemeralValue) and value.guard == self.guard
+
     @property
     def ValueType(self) -> Type[Value]:
         """The type of the value stored in the channel."""

libs/langgraph/langgraph/channels/last_value.py

@@ -14,6 +14,9 @@ class LastValue(Generic[Value], BaseChannel[Value, Value, Value]):
     def __init__(self, typ: Type[Value]) -> None:
         self.typ = typ
 
+    def __eq__(self, value: object) -> bool:
+        return isinstance(value, LastValue)
+
     @property
     def ValueType(self) -> Type[Value]:
         """The type of the value stored in the channel."""

libs/langgraph/langgraph/channels/manager.py

@@ -61,4 +61,6 @@ def create_checkpoint(
         channel_versions=checkpoint["channel_versions"],
         versions_seen=checkpoint["versions_seen"],
         pending_sends=checkpoint.get("pending_sends", []),
+        # checkpoints are saved only at the end of a step, ie. when current tasks should be cleared
+        current_tasks={},
     )

libs/langgraph/langgraph/channels/named_barrier_value.py

@@ -16,6 +16,9 @@ class NamedBarrierValue(Generic[Value], BaseChannel[Value, Value, set[Value]]):
         self.names = names
         self.seen = set()
 
+    def __eq__(self, value: object) -> bool:
+        return isinstance(value, NamedBarrierValue) and value.names == self.names
+
     @property
     def ValueType(self) -> Type[Value]:
         """The type of the value stored in the channel."""

libs/langgraph/langgraph/channels/topic.py

@@ -41,6 +41,13 @@ class Topic(
         self.seen = set[Value]()
         self.values = list[Value]()
 
+    def __eq__(self, value: object) -> bool:
+        return (
+            isinstance(value, Topic)
+            and value.unique == self.unique
+            and value.accumulate == self.accumulate
+        )
+
     @property
     def ValueType(self) -> Any:
         """The type of the value stored in the channel."""

libs/langgraph/langgraph/checkpoint/aiosqlite.py

@@ -2,7 +2,16 @@ import asyncio
 import functools
 from contextlib import AbstractAsyncContextManager
 from types import TracebackType
-from typing import Any, AsyncIterator, Dict, Iterator, Optional, TypeVar
+from typing import (
+    Any,
+    AsyncIterator,
+    Dict,
+    Iterator,
+    Optional,
+    Sequence,
+    Tuple,
+    TypeVar,
+)
 
 import aiosqlite
 from langchain_core.runnables import RunnableConfig
@@ -37,22 +46,29 @@ def not_implemented_sync_method(func: T) -> T:
 class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
     """An asynchronous checkpoint saver that stores checkpoints in a SQLite database.
 
+    This class provides an asynchronous interface for saving and retrieving checkpoints
+    using a SQLite database. It's designed for use in asynchronous environments and
+    offers better performance for I/O-bound operations compared to synchronous alternatives.
+
+    Attributes:
+        conn (aiosqlite.Connection): The asynchronous SQLite database connection.
+        serde (SerializerProtocol): The serializer used for encoding/decoding checkpoints.
+
     Tip:
         Requires the [aiosqlite](https://pypi.org/project/aiosqlite/) package.
         Install it with `pip install aiosqlite`.
 
-    Note:
-        While this class does support asynchronous checkpointing, it is not recommended
-        for production workloads, due to limitations in SQLite's write performance. For
-        production workloads, consider using a more robust database like PostgreSQL.
+    Warning:
+        While this class supports asynchronous checkpointing, it is not recommended
+        for production workloads due to limitations in SQLite's write performance.
+        For production use, consider a more robust database like PostgreSQL.
 
-    !!! Important
+    Tip:
         Remember to **close the database connection** after executing your code,
         otherwise, you may see the graph "hang" after execution (since the program
         will not exit until the connection is closed).
 
-        The easiest way to do this is to use the `async with` statement, as shown in the
-        examples below.
+        The easiest way is to use the `async with` statement as shown in the examples.
 
         ```python
         async with AsyncSqliteSaver.from_conn_string("checkpoints.sqlite") as saver:
@@ -63,12 +79,9 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
                 print(event)
         ```
 
-    Args:
-        conn (aiosqlite.Connection): The asynchronous SQLite database connection.
-        serde (Optional[SerializerProtocol]): The serializer to use for serializing and deserializing checkpoints. Defaults to JsonPlusSerializerCompat.
-
     Examples:
-        Usage within a StateGraph:
+        Usage within StateGraph:
+
         ```pycon
         >>> import asyncio
         >>> import aiosqlite
@@ -86,8 +99,8 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
         >>> asyncio.run(coro)
         Output: 2
         ```
-
         Raw usage:
+
         ```pycon
         >>> import asyncio
         >>> import aiosqlite
@@ -203,6 +216,15 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
                     metadata BLOB,
                     PRIMARY KEY (thread_id, thread_ts)
                 );
+                CREATE TABLE IF NOT EXISTS writes (
+                    thread_id TEXT NOT NULL,
+                    thread_ts TEXT NOT NULL,
+                    task_id TEXT NOT NULL,
+                    idx INTEGER NOT NULL,
+                    channel TEXT NOT NULL,
+                    value BLOB,
+                    PRIMARY KEY (thread_id, thread_ts, task_id, idx)
+                );
                 """
             ):
                 await self.conn.commit()
@@ -224,56 +246,58 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
             Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found.
         """
         await self.setup()
-        if config["configurable"].get("thread_ts"):
-            async with self.conn.execute(
-                "SELECT checkpoint, parent_ts, metadata FROM checkpoints WHERE thread_id = ? AND thread_ts = ?",
-                (
-                    str(config["configurable"]["thread_id"]),
-                    str(config["configurable"]["thread_ts"]),
-                ),
-            ) as cursor:
-                if value := await cursor.fetchone():
-                    return CheckpointTuple(
-                        config,
-                        self.serde.loads(value[0]),
-                        self.serde.loads(value[2]) if value[2] is not None else {},
-                        (
-                            {
-                                "configurable": {
-                                    "thread_id": config["configurable"]["thread_id"],
-                                    "thread_ts": value[1],
-                                }
-                            }
-                            if value[1]
-                            else None
-                        ),
-                    )
-        else:
-            async with self.conn.execute(
-                "SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id = ? ORDER BY thread_ts DESC LIMIT 1",
-                (str(config["configurable"]["thread_id"]),),
-            ) as cursor:
-                if value := await cursor.fetchone():
-                    return CheckpointTuple(
+        async with self.conn.cursor() as cur:
+            # find the latest checkpoint for the thread_id
+            if config["configurable"].get("thread_ts"):
+                await cur.execute(
+                    "SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id = ? AND thread_ts = ?",
+                    (
+                        str(config["configurable"]["thread_id"]),
+                        str(config["configurable"]["thread_ts"]),
+                    ),
+                )
+            else:
+                await cur.execute(
+                    "SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id = ? ORDER BY thread_ts DESC LIMIT 1",
+                    (str(config["configurable"]["thread_id"]),),
+                )
+            # if a checkpoint is found, return it
+            if value := await cur.fetchone():
+                if not config["configurable"].get("thread_ts"):
+                    config = {
+                        "configurable": {
+                            "thread_id": value[0],
+                            "thread_ts": value[1],
+                        }
+                    }
+                # find any pending writes
+                await cur.execute(
+                    "SELECT task_id, channel, value FROM writes WHERE thread_id = ? AND thread_ts = ?",
+                    (
+                        str(config["configurable"]["thread_id"]),
+                        str(config["configurable"]["thread_ts"]),
+                    ),
+                )
+                # deserialize the checkpoint and metadata
+                return CheckpointTuple(
+                    config,
+                    self.serde.loads(value[3]),
+                    self.serde.loads(value[4]) if value[4] is not None else {},
+                    (
                         {
                             "configurable": {
                                 "thread_id": value[0],
-                                "thread_ts": value[1],
+                                "thread_ts": value[2],
                             }
-                        },
-                        self.serde.loads(value[3]),
-                        self.serde.loads(value[4]) if value[4] is not None else {},
-                        (
-                            {
-                                "configurable": {
-                                    "thread_id": value[0],
-                                    "thread_ts": value[2],
-                                }
-                            }
-                            if value[2]
-                            else None
-                        ),
-                    )
+                        }
+                        if value[2]
+                        else None
+                    ),
+                    [
+                        (task_id, channel, self.serde.loads(value))
+                        async for task_id, channel, value in cur
+                    ],
+                )
 
     async def alist(
         self,
@@ -289,12 +313,13 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
         on the provided config. The checkpoints are ordered by timestamp in descending order.
 
         Args:
-            config (RunnableConfig): The config to use for listing the checkpoints.
-            before (Optional[RunnableConfig]): If provided, only checkpoints before the specified timestamp are returned. Defaults to None.
-            limit (Optional[int]): The maximum number of checkpoints to return. Defaults to None.
+            config (Optional[RunnableConfig]): Base configuration for filtering checkpoints.
+            filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata.
+            before (Optional[RunnableConfig]): List checkpoints created before this configuration.
+            limit (Optional[int]): Maximum number of checkpoints to return.
 
         Yields:
-            AsyncIterator[CheckpointTuple]: An asynchronous iterator of checkpoint tuples.
+            AsyncIterator[CheckpointTuple]: An asynchronous iterator of matching checkpoint tuples.
         """
         await self.setup()
         where, param_values = search_where(config, filter, before)
@@ -336,6 +361,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
         Args:
             config (RunnableConfig): The config to associate with the checkpoint.
             checkpoint (Checkpoint): The checkpoint to save.
+            metadata (CheckpointMetadata): Additional metadata to save with the checkpoint.
 
         Returns:
             RunnableConfig: The updated config containing the saved checkpoint's timestamp.
@@ -358,3 +384,35 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
                 "thread_ts": checkpoint["id"],
             }
         }
+
+    async def aput_writes(
+        self,
+        config: RunnableConfig,
+        writes: Sequence[Tuple[str, Any]],
+        task_id: str,
+    ) -> None:
+        """Store intermediate writes linked to a checkpoint asynchronously.
+
+        This method saves intermediate writes associated with a checkpoint to the database.
+
+        Args:
+            config (RunnableConfig): Configuration of the related checkpoint.
+            writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair.
+            task_id (str): Identifier for the task creating the writes.
+        """
+        await self.setup()
+        async with self.conn.executemany(
+            "INSERT OR REPLACE INTO writes (thread_id, thread_ts, task_id, idx, channel, value) VALUES (?, ?, ?, ?, ?, ?)",
+            [
+                (
+                    str(config["configurable"]["thread_id"]),
+                    str(config["configurable"]["thread_ts"]),
+                    task_id,
+                    idx,
+                    channel,
+                    self.serde.dumps(value),
+                )
+                for idx, (channel, value) in enumerate(writes)
+            ],
+        ):
+            await self.conn.commit()

libs/langgraph/langgraph/checkpoint/base.py

@@ -1,5 +1,4 @@
 from abc import ABC
-from collections import defaultdict
 from datetime import datetime, timezone
 from typing import (
     Any,
@@ -10,6 +9,7 @@ from typing import (
     Literal,
     NamedTuple,
     Optional,
+    Tuple,
     TypedDict,
     TypeVar,
     Union,
@@ -24,10 +24,13 @@ from langgraph.serde.base import SerializerProtocol
 from langgraph.serde.jsonplus import JsonPlusSerializer
 
 V = TypeVar("V", int, float, str)
+PendingWrite = Tuple[str, str, Any]
 
 
 # Marked as total=False to allow for future expansion.
 class CheckpointMetadata(TypedDict, total=False):
+    """Metadata associated with a checkpoint."""
+
     source: Literal["input", "loop", "update"]
     """The source of the checkpoint.
     - "input": The checkpoint was created from an input to invoke/stream/batch.
@@ -52,6 +55,10 @@ class CheckpointMetadata(TypedDict, total=False):
     """
 
 
+class TaskInfo(TypedDict):
+    status: Literal["scheduled", "success", "error"]
+
+
 class Checkpoint(TypedDict):
     """State snapshot at a given point in time."""
 
@@ -73,7 +80,7 @@ class Checkpoint(TypedDict):
     The keys are channel names and the values are the logical time step
     at which the channel was last updated.
     """
-    versions_seen: defaultdict[str, dict[str, Union[str, int, float]]]
+    versions_seen: dict[str, dict[str, Union[str, int, float]]]
     """Map from node ID to map from channel name to version seen.
     
     This keeps track of the versions of the channels that each node has seen.
@@ -83,6 +90,8 @@ class Checkpoint(TypedDict):
     pending_sends: List[Send]
     """List of packets sent to nodes but not yet processed.
     Cleared by the next checkpoint."""
+    current_tasks: Dict[str, TaskInfo]
+    """Map from task ID to task info."""
 
 
 def empty_checkpoint() -> Checkpoint:
@@ -92,8 +101,9 @@ def empty_checkpoint() -> Checkpoint:
         ts=datetime.now(timezone.utc).isoformat(),
         channel_values={},
         channel_versions={},
-        versions_seen=defaultdict(dict),
+        versions_seen={},
         pending_sends=[],
+        current_tasks={},
     )
 
 
@@ -104,19 +114,20 @@ def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint:
         id=checkpoint["id"],
         channel_values=checkpoint["channel_values"].copy(),
         channel_versions=checkpoint["channel_versions"].copy(),
-        versions_seen=defaultdict(
-            dict,
-            {k: v.copy() for k, v in checkpoint["versions_seen"].items()},
-        ),
+        versions_seen={k: v.copy() for k, v in checkpoint["versions_seen"].items()},
         pending_sends=checkpoint.get("pending_sends", []).copy(),
+        current_tasks=checkpoint.get("current_tasks", {}).copy(),
     )
 
 
 class CheckpointTuple(NamedTuple):
+    """A tuple containing a checkpoint and its associated data."""
+
     config: RunnableConfig
     checkpoint: Checkpoint
     metadata: CheckpointMetadata
     parent_config: Optional[RunnableConfig] = None
+    pending_writes: Optional[List[PendingWrite]] = None
 
 
 CheckpointThreadId = ConfigurableFieldSpec(
@@ -139,6 +150,19 @@ CheckpointThreadTs = ConfigurableFieldSpec(
 
 
 class BaseCheckpointSaver(ABC):
+    """Base class for creating a graph checkpointer.
+
+    Checkpointers allow LangGraph agents to persist their state
+    within and across multiple interactions.
+
+    Attributes:
+        serde (SerializerProtocol): Serializer for encoding/decoding checkpoints.
+
+    Note:
+        When creating a custom checkpoint saver, consider implementing async
+        versions to avoid blocking the main thread.
+    """
+
     serde: SerializerProtocol = JsonPlusSerializer()
 
     def __init__(
@@ -150,13 +174,37 @@ class BaseCheckpointSaver(ABC):
 
     @property
     def config_specs(self) -> list[ConfigurableFieldSpec]:
+        """Define the configuration options for the checkpoint saver.
+
+        Returns:
+            list[ConfigurableFieldSpec]: List of configuration field specs.
+        """
         return [CheckpointThreadId, CheckpointThreadTs]
 
     def get(self, config: RunnableConfig) -> Optional[Checkpoint]:
+        """Fetch a checkpoint using the given configuration.
+
+        Args:
+            config (RunnableConfig): Configuration specifying which checkpoint to retrieve.
+
+        Returns:
+            Optional[Checkpoint]: The requested checkpoint, or None if not found.
+        """
         if value := self.get_tuple(config):
             return value.checkpoint
 
     def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
+        """Fetch a checkpoint tuple using the given configuration.
+
+        Args:
+            config (RunnableConfig): Configuration specifying which checkpoint to retrieve.
+
+        Returns:
+            Optional[CheckpointTuple]: The requested checkpoint tuple, or None if not found.
+
+        Raises:
+            NotImplementedError: Implement this method in your custom checkpoint saver.
+        """
         raise NotImplementedError
 
     def list(
@@ -167,6 +215,20 @@ class BaseCheckpointSaver(ABC):
         before: Optional[RunnableConfig] = None,
         limit: Optional[int] = None,
     ) -> Iterator[CheckpointTuple]:
+        """List checkpoints that match the given criteria.
+
+        Args:
+            config (Optional[RunnableConfig]): Base configuration for filtering checkpoints.
+            filter (Optional[Dict[str, Any]]): Additional filtering criteria.
+            before (Optional[RunnableConfig]): List checkpoints created before this configuration.
+            limit (Optional[int]): Maximum number of checkpoints to return.
+
+        Returns:
+            Iterator[CheckpointTuple]: Iterator of matching checkpoint tuples.
+
+        Raises:
+            NotImplementedError: Implement this method in your custom checkpoint saver.
+        """
         raise NotImplementedError
 
     def put(
@@ -175,16 +237,68 @@ class BaseCheckpointSaver(ABC):
         checkpoint: Checkpoint,
         metadata: CheckpointMetadata,
     ) -> RunnableConfig:
+        """Store a checkpoint with its configuration and metadata.
+
+        Args:
+            config (RunnableConfig): Configuration for the checkpoint.
+            checkpoint (Checkpoint): The checkpoint to store.
+            metadata (CheckpointMetadata): Additional metadata for the checkpoint.
+
+        Returns:
+            RunnableConfig: Updated configuration after storing the checkpoint.
+
+        Raises:
+            NotImplementedError: Implement this method in your custom checkpoint saver.
+        """
         raise NotImplementedError
 
+    def put_writes(
+        self,
+        config: RunnableConfig,
+        writes: List[Tuple[str, Any]],
+        task_id: str,
+    ) -> None:
+        """Store intermediate writes linked to a checkpoint.
+
+        Args:
+            config (RunnableConfig): Configuration of the related checkpoint.
+            writes (List[Tuple[str, Any]]): List of writes to store.
+            task_id (str): Identifier for the task creating the writes.
+
+        Raises:
+            NotImplementedError: Implement this method in your custom checkpoint saver.
+        """
+        raise NotImplementedError(
+            "This method was added in langgraph 0.1.7. Please update your checkpoint saver to implement it."
+        )
+
     async def aget(self, config: RunnableConfig) -> Optional[Checkpoint]:
+        """Asynchronously fetch a checkpoint using the given configuration.
+
+        Args:
+            config (RunnableConfig): Configuration specifying which checkpoint to retrieve.
+
+        Returns:
+            Optional[Checkpoint]: The requested checkpoint, or None if not found.
+        """
         if value := await self.aget_tuple(config):
             return value.checkpoint
 
     async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
+        """Asynchronously fetch a checkpoint tuple using the given configuration.
+
+        Args:
+            config (RunnableConfig): Configuration specifying which checkpoint to retrieve.
+
+        Returns:
+            Optional[CheckpointTuple]: The requested checkpoint tuple, or None if not found.
+
+        Raises:
+            NotImplementedError: Implement this method in your custom checkpoint saver.
+        """
         raise NotImplementedError
 
-    def alist(
+    async def alist(
         self,
         config: Optional[RunnableConfig],
         *,
@@ -192,6 +306,20 @@ class BaseCheckpointSaver(ABC):
         before: Optional[RunnableConfig] = None,
         limit: Optional[int] = None,
     ) -> AsyncIterator[CheckpointTuple]:
+        """Asynchronously list checkpoints that match the given criteria.
+
+        Args:
+            config (Optional[RunnableConfig]): Base configuration for filtering checkpoints.
+            filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata.
+            before (Optional[RunnableConfig]): List checkpoints created before this configuration.
+            limit (Optional[int]): Maximum number of checkpoints to return.
+
+        Returns:
+            AsyncIterator[CheckpointTuple]: Async iterator of matching checkpoint tuples.
+
+        Raises:
+            NotImplementedError: Implement this method in your custom checkpoint saver.
+        """
         raise NotImplementedError
         yield
 
@@ -201,9 +329,52 @@ class BaseCheckpointSaver(ABC):
         checkpoint: Checkpoint,
         metadata: CheckpointMetadata,
     ) -> RunnableConfig:
+        """Asynchronously store a checkpoint with its configuration and metadata.
+
+        Args:
+            config (RunnableConfig): Configuration for the checkpoint.
+            checkpoint (Checkpoint): The checkpoint to store.
+            metadata (CheckpointMetadata): Additional metadata for the checkpoint.
+
+        Returns:
+            RunnableConfig: Updated configuration after storing the checkpoint.
+
+        Raises:
+            NotImplementedError: Implement this method in your custom checkpoint saver.
+        """
         raise NotImplementedError
 
+    async def aput_writes(
+        self,
+        config: RunnableConfig,
+        writes: List[Tuple[str, Any]],
+        task_id: str,
+    ) -> None:
+        """Asynchronously store intermediate writes linked to a checkpoint.
+
+        Args:
+            config (RunnableConfig): Configuration of the related checkpoint.
+            writes (List[Tuple[str, Any]]): List of writes to store.
+            task_id (str): Identifier for the task creating the writes.
+
+        Raises:
+            NotImplementedError: Implement this method in your custom checkpoint saver.
+        """
+        raise NotImplementedError(
+            "This method was added in langgraph 0.1.7. Please update your checkpoint saver to implement it."
+        )
+
     def get_next_version(self, current: Optional[V], channel: BaseChannel) -> V:
-        """Get the next version of a channel. Default is to use int versions, incrementing by 1. If you override, you can use str/int/float versions,
-        as long as they are monotonically increasing."""
+        """Generate the next version ID for a channel.
+
+        Default is to use integer versions, incrementing by 1. If you override, you can use str/int/float versions,
+        as long as they are monotonically increasing.
+
+        Args:
+            current (Optional[V]): The current version identifier (int, float, or str).
+            channel (BaseChannel): The channel being versioned.
+
+        Returns:
+            V: The next version identifier, which must be increasing.
+        """
         return current + 1 if current is not None else 1

libs/langgraph/langgraph/checkpoint/memory.py

@@ -1,7 +1,7 @@
 import asyncio
 from collections import defaultdict
 from functools import partial
-from typing import Any, AsyncIterator, Dict, Iterator, Optional
+from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Tuple
 
 from langchain_core.runnables import RunnableConfig
 
@@ -53,6 +53,7 @@ class MemorySaver(BaseCheckpointSaver):
     ) -> None:
         super().__init__(serde=serde)
         self.storage = defaultdict(dict)
+        self.writes = defaultdict(list)
 
     def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
         """Get a checkpoint tuple from the in-memory storage.
@@ -72,19 +73,27 @@ class MemorySaver(BaseCheckpointSaver):
         if ts := config["configurable"].get("thread_ts"):
             if saved := self.storage[thread_id].get(ts):
                 checkpoint, metadata = saved
+                writes = self.writes[(thread_id, ts)]
                 return CheckpointTuple(
                     config=config,
                     checkpoint=self.serde.loads(checkpoint),
                     metadata=self.serde.loads(metadata),
+                    pending_writes=[
+                        (id, c, self.serde.loads(v)) for id, c, v in writes
+                    ],
                 )
         else:
             if checkpoints := self.storage[thread_id]:
                 ts = max(checkpoints.keys())
                 checkpoint, metadata = checkpoints[ts]
+                writes = self.writes[(thread_id, ts)]
                 return CheckpointTuple(
                     config={"configurable": {"thread_id": thread_id, "thread_ts": ts}},
                     checkpoint=self.serde.loads(checkpoint),
                     metadata=self.serde.loads(metadata),
+                    pending_writes=[
+                        (id, c, self.serde.loads(v)) for id, c, v in writes
+                    ],
                 )
 
     def list(
@@ -98,15 +107,16 @@ class MemorySaver(BaseCheckpointSaver):
         """List checkpoints from the in-memory storage.
 
         This method retrieves a list of checkpoint tuples from the in-memory storage based
-        on the provided config. The checkpoints are ordered by timestamp in insertion order.
+        on the provided criteria.
 
         Args:
-            config (RunnableConfig): The config to use for listing the checkpoints.
-            before (Optional[RunnableConfig]): If provided, only checkpoints before the specified timestamp are returned. Defaults to None.
-            limit (Optional[int]): The maximum number of checkpoints to return. Defaults to None.
+            config (Optional[RunnableConfig]): Base configuration for filtering checkpoints.
+            filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata.
+            before (Optional[RunnableConfig]): List checkpoints created before this configuration.
+            limit (Optional[int]): Maximum number of checkpoints to return.
 
         Yields:
-            Iterator[CheckpointTuple]: An iterator of checkpoint tuples.
+            Iterator[CheckpointTuple]: An iterator of matching checkpoint tuples.
         """
         thread_ids = (config["configurable"]["thread_id"],) if config else self.storage
         for thread_id in thread_ids:
@@ -149,6 +159,7 @@ class MemorySaver(BaseCheckpointSaver):
         Args:
             config (RunnableConfig): The config to associate with the checkpoint.
             checkpoint (Checkpoint): The checkpoint to save.
+            metadata (CheckpointMetadata): Additional metadata to save with the checkpoint.
 
         Returns:
             RunnableConfig: The updated config containing the saved checkpoint's timestamp.
@@ -168,6 +179,31 @@ class MemorySaver(BaseCheckpointSaver):
             }
         }
 
+    def put_writes(
+        self,
+        config: RunnableConfig,
+        writes: List[Tuple[str, Any]],
+        task_id: str,
+    ) -> RunnableConfig:
+        """Save a list of writes to the in-memory storage.
+
+        This method saves a list of writes to the in-memory storage. The writes are associated
+        with the provided config.
+
+        Args:
+            config (RunnableConfig): The config to associate with the writes.
+            writes (list[tuple[str, Any]]): The writes to save.
+            task_id (str): Identifier for the task creating the writes.
+
+        Returns:
+            RunnableConfig: The updated config containing the saved writes' timestamp.
+        """
+        thread_id = config["configurable"]["thread_id"]
+        ts = config["configurable"]["thread_ts"]
+        self.writes[(thread_id, ts)].extend(
+            [(task_id, c, self.serde.dumps(v)) for c, v in writes]
+        )
+
     async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
         """Asynchronous version of get_tuple.
 
@@ -221,6 +257,36 @@ class MemorySaver(BaseCheckpointSaver):
         checkpoint: Checkpoint,
         metadata: CheckpointMetadata,
     ) -> RunnableConfig:
+        """Asynchronous version of put.
+
+        Args:
+            config (RunnableConfig): The config to associate with the checkpoint.
+            checkpoint (Checkpoint): The checkpoint to save.
+            metadata (CheckpointMetadata): Additional metadata to save with the checkpoint.
+
+        Returns:
+            RunnableConfig: The updated config containing the saved checkpoint's timestamp.
+        """
         return await asyncio.get_running_loop().run_in_executor(
             None, self.put, config, checkpoint, metadata
         )
+
+    async def aput_writes(
+        self,
+        config: RunnableConfig,
+        writes: List[Tuple[str, Any]],
+        task_id: str,
+    ) -> RunnableConfig:
+        """Asynchronous version of put_writes.
+
+        This method is an asynchronous wrapper around put_writes that runs the synchronous
+        method in a separate thread using asyncio.
+
+        Args:
+            config (RunnableConfig): The config to associate with the writes.
+            writes (List[Tuple[str, Any]]): The writes to save, each as a (channel, value) pair.
+            task_id (str): Identifier for the task creating the writes.
+        """
+        return await asyncio.get_running_loop().run_in_executor(
+            None, self.put_writes, config, writes, task_id
+        )

libs/langgraph/langgraph/checkpoint/sqlite.py

@@ -171,6 +171,15 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
                 metadata BLOB,
                 PRIMARY KEY (thread_id, thread_ts)
             );
+            CREATE TABLE IF NOT EXISTS writes (
+                thread_id TEXT NOT NULL,
+                thread_ts TEXT NOT NULL,
+                task_id TEXT NOT NULL,
+                idx INTEGER NOT NULL,
+                channel TEXT NOT NULL,
+                value BLOB,
+                PRIMARY KEY (thread_id, thread_ts, task_id, idx)
+            );
             """
         )
 
@@ -233,56 +242,57 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
             CheckpointTuple(...)
         """  # noqa
         with self.cursor(transaction=False) as cur:
+            # find the latest checkpoint for the thread_id
             if config["configurable"].get("thread_ts"):
                 cur.execute(
-                    "SELECT checkpoint, parent_ts, metadata FROM checkpoints WHERE thread_id = ? AND thread_ts = ?",
+                    "SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id = ? AND thread_ts = ?",
                     (
                         str(config["configurable"]["thread_id"]),
                         str(config["configurable"]["thread_ts"]),
                     ),
                 )
-                if value := cur.fetchone():
-                    return CheckpointTuple(
-                        config,
-                        self.serde.loads(value[0]),
-                        self.serde.loads(value[2]) if value[2] is not None else {},
-                        (
-                            {
-                                "configurable": {
-                                    "thread_id": config["configurable"]["thread_id"],
-                                    "thread_ts": value[1],
-                                }
-                            }
-                            if value[1]
-                            else None
-                        ),
-                    )
             else:
                 cur.execute(
                     "SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id = ? ORDER BY thread_ts DESC LIMIT 1",
                     (str(config["configurable"]["thread_id"]),),
                 )
-                if value := cur.fetchone():
-                    return CheckpointTuple(
+            # if a checkpoint is found, return it
+            if value := cur.fetchone():
+                if not config["configurable"].get("thread_ts"):
+                    config = {
+                        "configurable": {
+                            "thread_id": value[0],
+                            "thread_ts": value[1],
+                        }
+                    }
+                # find any pending writes
+                cur.execute(
+                    "SELECT task_id, channel, value FROM writes WHERE thread_id = ? AND thread_ts = ?",
+                    (
+                        str(config["configurable"]["thread_id"]),
+                        str(config["configurable"]["thread_ts"]),
+                    ),
+                )
+                # deserialize the checkpoint and metadata
+                return CheckpointTuple(
+                    config,
+                    self.serde.loads(value[3]),
+                    self.serde.loads(value[4]) if value[4] is not None else {},
+                    (
                         {
                             "configurable": {
                                 "thread_id": value[0],
-                                "thread_ts": value[1],
+                                "thread_ts": value[2],
                             }
-                        },
-                        self.serde.loads(value[3]),
-                        self.serde.loads(value[4]) if value[4] is not None else {},
-                        (
-                            {
-                                "configurable": {
-                                    "thread_id": value[0],
-                                    "thread_ts": value[2],
-                                }
-                            }
-                            if value[2]
-                            else None
-                        ),
-                    )
+                        }
+                        if value[2]
+                        else None
+                    ),
+                    [
+                        (task_id, channel, self.serde.loads(value))
+                        for task_id, channel, value in cur
+                    ],
+                )
 
     def list(
         self,
@@ -299,6 +309,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
 
         Args:
             config (RunnableConfig): The config to use for listing the checkpoints.
+            filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata. Defaults to None.
             before (Optional[RunnableConfig]): If provided, only checkpoints before the specified timestamp are returned. Defaults to None.
             limit (Optional[int]): The maximum number of checkpoints to return. Defaults to None.
 
@@ -394,6 +405,37 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
             }
         }
 
+    def put_writes(
+        self,
+        config: RunnableConfig,
+        writes: Sequence[Tuple[str, Any]],
+        task_id: str,
+    ) -> None:
+        """Store intermediate writes linked to a checkpoint.
+
+        This method saves intermediate writes associated with a checkpoint to the SQLite database.
+
+        Args:
+            config (RunnableConfig): Configuration of the related checkpoint.
+            writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair.
+            task_id (str): Identifier for the task creating the writes.
+        """
+        with self.lock, self.cursor() as cur:
+            cur.executemany(
+                "INSERT OR REPLACE INTO writes (thread_id, thread_ts, task_id, idx, channel, value) VALUES (?, ?, ?, ?, ?, ?)",
+                [
+                    (
+                        str(config["configurable"]["thread_id"]),
+                        str(config["configurable"]["thread_ts"]),
+                        task_id,
+                        idx,
+                        channel,
+                        self.serde.dumps(value),
+                    )
+                    for idx, (channel, value) in enumerate(writes)
+                ],
+            )
+
     async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
         """Get a checkpoint tuple from the database asynchronously.
 
@@ -435,6 +477,17 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
         raise NotImplementedError(_AIO_ERROR_MSG)
 
     def get_next_version(self, current: Optional[str], channel: BaseChannel) -> str:
+        """Generate the next version ID for a channel.
+
+        This method creates a new version identifier for a channel based on its current version.
+
+        Args:
+            current (Optional[str]): The current version identifier of the channel.
+            channel (BaseChannel): The channel being versioned.
+
+        Returns:
+            str: The next version identifier, which is guaranteed to be monotonically increasing.
+        """
         if current is None:
             current_v = 0
         else:

libs/langgraph/langgraph/constants.py

@@ -1,10 +1,11 @@
 from typing import Any
 
+INPUT = "__input__"
 CONFIG_KEY_SEND = "__pregel_send"
 CONFIG_KEY_READ = "__pregel_read"
 INTERRUPT = "__interrupt__"
 TASKS = "__pregel_tasks"
-RESERVED = {INTERRUPT, TASKS, CONFIG_KEY_SEND, CONFIG_KEY_READ}
+RESERVED = {INTERRUPT, TASKS, CONFIG_KEY_SEND, CONFIG_KEY_READ, INPUT}
 TAG_HIDDEN = "langsmith:hidden"
 
 START = "__start__"

libs/langgraph/langgraph/graph/graph.py

@@ -20,9 +20,8 @@ from typing import (
 from langchain_core.runnables import Runnable
 from langchain_core.runnables.base import RunnableLike
 from langchain_core.runnables.config import RunnableConfig
-from langchain_core.runnables.graph import (
-    Node as RunnableGraphNode,
-)
+from langchain_core.runnables.graph import Graph as DrawableGraph
+from langchain_core.runnables.graph import Node as DrawableNode
 
 from langgraph.channels.ephemeral_value import EphemeralValue
 from langgraph.checkpoint import BaseCheckpointSaver
@@ -32,11 +31,16 @@ from langgraph.pregel import Channel, Pregel
 from langgraph.pregel.read import PregelNode
 from langgraph.pregel.types import All
 from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
-from langgraph.utils import DrawableGraph, RunnableCallable, coerce_to_runnable
+from langgraph.utils import RunnableCallable, coerce_to_runnable
 
 logger = logging.getLogger(__name__)
 
 
+class NodeSpec(NamedTuple):
+    runnable: Runnable
+    metadata: Optional[dict[str, Any]] = None
+
+
 class Branch(NamedTuple):
     path: Runnable[Any, Union[Hashable, list[Hashable]]]
     ends: Optional[dict[Hashable, str]]
@@ -114,7 +118,7 @@ class Branch(NamedTuple):
 
 class Graph:
     def __init__(self) -> None:
-        self.nodes: dict[str, Runnable] = {}
+        self.nodes: dict[str, NodeSpec] = {}
         self.edges = set[tuple[str, str]]()
         self.branches: defaultdict[str, dict[str, Branch]] = defaultdict(dict)
         self.support_multiple_edges = False
@@ -125,15 +129,30 @@ class Graph:
         return self.edges
 
     @overload
-    def add_node(self, node: RunnableLike) -> None:
+    def add_node(
+        self,
+        node: RunnableLike,
+        *,
+        metadata: Optional[dict[str, Any]] = None,
+    ) -> None:
         ...
 
     @overload
-    def add_node(self, node: str, action: RunnableLike) -> None:
+    def add_node(
+        self,
+        node: str,
+        action: RunnableLike,
+        *,
+        metadata: Optional[dict[str, Any]] = None,
+    ) -> None:
         ...
 
     def add_node(
-        self, node: Union[str, RunnableLike], action: Optional[RunnableLike] = None
+        self,
+        node: Union[str, RunnableLike],
+        action: Optional[RunnableLike] = None,
+        *,
+        metadata: Optional[dict[str, Any]] = None,
     ) -> None:
         if self.compiled:
             logger.warning(
@@ -148,7 +167,9 @@ class Graph:
         if node == END or node == START:
             raise ValueError(f"Node `{node}` is reserved.")
 
-        self.nodes[node] = coerce_to_runnable(action, name=node, trace=False)
+        self.nodes[node] = NodeSpec(
+            coerce_to_runnable(action, name=node, trace=False), metadata
+        )
 
     def add_edge(self, start_key: str, end_key: str) -> None:
         if self.compiled:
@@ -196,6 +217,10 @@ class Graph:
 
         Returns:
             None
+
+        Note: Without typehints on the `path` function's return value (e.g., `-> Literal["foo", "__end__"]:`)
+            or a path_map, the graph visualization assumes the edge could transition to any node in the graph.
+
         """  # noqa: E501
         if self.compiled:
             logger.warning(
@@ -385,11 +410,11 @@ class Graph:
 class CompiledGraph(Pregel):
     builder: Graph
 
-    def attach_node(self, key: str, node: Runnable) -> None:
+    def attach_node(self, key: str, node: NodeSpec) -> None:
         self.channels[key] = EphemeralValue(Any)
         self.nodes[key] = (
-            PregelNode(channels=[], triggers=[])
-            | node
+            PregelNode(channels=[], triggers=[], metadata=node.metadata)
+            | node.runnable
             | ChannelWrite([ChannelWriteEntry(key)], tags=[TAG_HIDDEN])
         )
         cast(list[str], self.stream_channels).append(key)
@@ -441,14 +466,22 @@ class CompiledGraph(Pregel):
     ) -> DrawableGraph:
         """Returns a drawable representation of the computation graph."""
         graph = DrawableGraph()
-        start_nodes: dict[str, RunnableGraphNode] = {
+        start_nodes: dict[str, DrawableNode] = {
             START: graph.add_node(self.get_input_schema(config), START)
         }
-        end_nodes: dict[str, RunnableGraphNode] = {
-            END: graph.add_node(self.get_output_schema(config), END)
-        }
+        end_nodes: dict[str, DrawableNode] = {}
 
-        for key, node in self.builder.nodes.items():
+        def add_edge(
+            start: str, end: str, label: Optional[str] = None, conditional: bool = False
+        ) -> None:
+            if end == END and END not in end_nodes:
+                end_nodes[END] = graph.add_node(self.get_output_schema(config), END)
+            return graph.add_edge(
+                start_nodes[start], end_nodes[end], label, conditional
+            )
+
+        for key, n in self.builder.nodes.items():
+            node = n.runnable
             if xray:
                 subgraph = (
                     node.get_graph(
@@ -469,11 +502,11 @@ class CompiledGraph(Pregel):
                     start_nodes[key] = n
                     end_nodes[key] = n
             else:
-                n = graph.add_node(node, key)
+                n = graph.add_node(node, key, metadata=n.metadata)
                 start_nodes[key] = n
                 end_nodes[key] = n
         for start, end in sorted(self.builder._all_edges):
-            graph.add_edge(start_nodes[start], end_nodes[end])
+            add_edge(start, end)
         for start, branches in self.builder.branches.items():
             default_ends = {
                 **{k: k for k in self.builder.nodes if k != start},
@@ -487,13 +520,13 @@ class CompiledGraph(Pregel):
                 else:
                     ends = default_ends
                 for label, end in ends.items():
-                    graph.add_edge(
-                        start_nodes[start],
-                        end_nodes[end],
+                    add_edge(
+                        start,
+                        end,
                         label if label != end else None,
                         conditional=True,
                     )
                     if branch.then is not None:
-                        graph.add_edge(start_nodes[end], end_nodes[branch.then])
+                        add_edge(end, branch.then)
 
         return graph

libs/langgraph/langgraph/graph/message.py

@@ -78,10 +78,11 @@ def add_messages(left: Messages, right: Messages) -> Messages:
     # merge
     left_idx_by_id = {m.id: i for i, m in enumerate(left)}
     merged = left.copy()
+    ids_to_remove = set()
     for m in right:
         if (existing_idx := left_idx_by_id.get(m.id)) is not None:
             if isinstance(m, RemoveMessage):
-                del merged[existing_idx]
+                ids_to_remove.add(m.id)
             else:
                 merged[existing_idx] = m
         else:
@@ -91,6 +92,7 @@ def add_messages(left: Messages, right: Messages) -> Messages:
                 )
 
             merged.append(m)
+    merged = [m for m in merged if m.id not in ids_to_remove]
     return merged
 
 

libs/langgraph/langgraph/graph/state.py

@@ -2,9 +2,10 @@ import logging
 import typing
 import warnings
 from functools import partial
-from inspect import signature
+from inspect import isclass, isfunction, signature
 from typing import (
     Any,
+    NamedTuple,
     Optional,
     Sequence,
     Type,
@@ -17,6 +18,9 @@ from typing import (
 from langchain_core.pydantic_v1 import BaseModel
 from langchain_core.runnables import Runnable, RunnableConfig
 from langchain_core.runnables.base import RunnableLike
+from langchain_core.runnables.utils import (
+    create_model,
+)
 
 from langgraph.channels.base import BaseChannel
 from langgraph.channels.binop import BinaryOperatorAggregate
@@ -28,12 +32,19 @@ from langgraph.channels.named_barrier_value import NamedBarrierValue
 from langgraph.checkpoint import BaseCheckpointSaver
 from langgraph.constants import TAG_HIDDEN
 from langgraph.errors import InvalidUpdateError
-from langgraph.graph.graph import END, START, Branch, CompiledGraph, Graph, Send
+from langgraph.graph.graph import (
+    END,
+    START,
+    Branch,
+    CompiledGraph,
+    Graph,
+    Send,
+)
 from langgraph.managed.base import ManagedValue, is_managed_value
 from langgraph.pregel.read import ChannelRead, PregelNode
-from langgraph.pregel.types import All
+from langgraph.pregel.types import All, RetryPolicy
 from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry
-from langgraph.utils import RunnableCallable
+from langgraph.utils import RunnableCallable, coerce_to_runnable
 
 logger = logging.getLogger(__name__)
 
@@ -50,6 +61,13 @@ def _warn_invalid_state_schema(schema: Union[Type[Any], Any]) -> None:
     )
 
 
+class StateNodeSpec(NamedTuple):
+    runnable: Runnable
+    metadata: dict[str, Any]
+    input: Type[Any]
+    retry_policy: Optional[RetryPolicy]
+
+
 class StateGraph(Graph):
     """A graph whose nodes communicate by reading and writing to a shared state.
     The signature of each node is State -> Partial<State>.
@@ -101,16 +119,39 @@ class StateGraph(Graph):
         >>> print(step1)
         {'x': [0.5, 0.75]}"""
 
+    nodes: dict[str, StateNodeSpec]
+    channels: dict[str, BaseChannel]
+    managed: dict[str, Type[ManagedValue]]
+    schemas: dict[Type[Any], dict[str, Union[BaseChannel, Type[ManagedValue]]]]
+
     def __init__(
-        self, state_schema: Type[Any], config_schema: Optional[Type[Any]] = None
+        self,
+        state_schema: Optional[Type[Any]] = None,
+        config_schema: Optional[Type[Any]] = None,
+        *,
+        input: Optional[Type[Any]] = None,
+        output: Optional[Type[Any]] = None,
     ) -> None:
         super().__init__()
-        _warn_invalid_state_schema(state_schema)
+        if state_schema is None:
+            if input is None or output is None:
+                raise ValueError("Must provide state_schema or input and output")
+            state_schema = input
+        else:
+            if input is None:
+                input = state_schema
+            if output is None:
+                output = state_schema
+        self.schemas = {}
+        self.channels = {}
+        self.managed = {}
         self.schema = state_schema
+        self.input = input
+        self.output = output
+        self._add_schema(state_schema)
+        self._add_schema(input)
+        self._add_schema(output)
         self.config_schema = config_schema
-        self.channels, self.managed = _get_channels(state_schema)
-        if any(isinstance(c, BinaryOperatorAggregate) for c in self.channels.values()):
-            self.support_multiple_edges = True
         self.waiting_edges: set[tuple[tuple[str, ...], str]] = set()
 
     @property
@@ -119,8 +160,44 @@ class StateGraph(Graph):
             (start, end) for starts, end in self.waiting_edges for start in starts
         }
 
+    def _add_schema(self, schema: Type[Any]) -> None:
+        if schema not in self.schemas:
+            _warn_invalid_state_schema(schema)
+            channels, managed = _get_channels(schema)
+            self.schemas[schema] = {**channels, **managed}
+            for key, channel in channels.items():
+                if key in self.channels:
+                    if self.channels[key] != channel:
+                        if isinstance(channel, LastValue):
+                            pass
+                        else:
+                            raise ValueError(
+                                f"Channel '{key}' already exists with a different type"
+                            )
+                else:
+                    self.channels[key] = channel
+            for key, managed in managed.items():
+                if key in self.managed:
+                    if self.managed[key] != managed:
+                        raise ValueError(
+                            f"Managed value '{key}' already exists with a different type"
+                        )
+                else:
+                    self.managed[key] = managed
+            if any(
+                isinstance(c, BinaryOperatorAggregate) for c in self.channels.values()
+            ):
+                self.support_multiple_edges = True
+
     @overload
-    def add_node(self, node: RunnableLike) -> None:
+    def add_node(
+        self,
+        node: RunnableLike,
+        *,
+        metadata: Optional[dict[str, Any]] = None,
+        input: Optional[Type[Any]] = None,
+        retry: Optional[RetryPolicy] = None,
+    ) -> None:
         """Adds a new node to the state graph.
         Will take the name of the function/runnable as the node name.
 
@@ -136,7 +213,15 @@ class StateGraph(Graph):
         ...
 
     @overload
-    def add_node(self, node: str, action: RunnableLike) -> None:
+    def add_node(
+        self,
+        node: str,
+        action: RunnableLike,
+        *,
+        metadata: Optional[dict[str, Any]] = None,
+        input: Optional[Type[Any]] = None,
+        retry: Optional[RetryPolicy] = None,
+    ) -> None:
         """Adds a new node to the state graph.
 
         Args:
@@ -152,7 +237,13 @@ class StateGraph(Graph):
         ...
 
     def add_node(
-        self, node: Union[str, RunnableLike], action: Optional[RunnableLike] = None
+        self,
+        node: Union[str, RunnableLike],
+        action: Optional[RunnableLike] = None,
+        *,
+        metadata: Optional[dict[str, Any]] = None,
+        input: Optional[Type[Any]] = None,
+        retry: Optional[RetryPolicy] = None,
     ) -> None:
         """Adds a new node to the state graph.
 
@@ -161,6 +252,9 @@ class StateGraph(Graph):
         Args:
             node (Union[str, RunnableLike)]: The function or runnable this node will run.
             action (Optional[RunnableLike]): The action associated with the node. (default: None)
+            metadata (Optional[dict[str, Any]]): The metadata associated with the node. (default: None)
+            input (Optional[Type[Any]]): The input schema for the node. (default: the graph's input schema)
+            retry (Optional[RetryPolicy]): The policy for retrying the node. (default: None)
         Raises:
             ValueError: If the key is already being used as a state key.
 
@@ -205,7 +299,36 @@ class StateGraph(Graph):
                 )
         if node in self.channels:
             raise ValueError(f"'{node}' is already being used as a state key")
-        return super().add_node(node, action)
+        if self.compiled:
+            logger.warning(
+                "Adding a node to a graph that has already been compiled. This will "
+                "not be reflected in the compiled graph."
+            )
+        if not isinstance(node, str):
+            action = node
+            node = getattr(action, "name", action.__name__)
+        if node in self.nodes:
+            raise ValueError(f"Node `{node}` already present.")
+        if node == END or node == START:
+            raise ValueError(f"Node `{node}` is reserved.")
+        try:
+            if isfunction(action) and (
+                hints := get_type_hints(action.__call__) or get_type_hints(action)
+            ):
+                if input is None:
+                    input_hint = hints[list(hints.keys())[0]]
+                    if isinstance(input_hint, type) and get_type_hints(input_hint):
+                        input = input_hint
+        except TypeError:
+            pass
+        if input is not None:
+            self._add_schema(input)
+        self.nodes[node] = StateNodeSpec(
+            coerce_to_runnable(action, name=node, trace=False),
+            metadata,
+            input=input or self.schema,
+            retry_policy=retry,
+        )
 
     def add_edge(self, start_key: Union[str, list[str]], end_key: str) -> None:
         """Adds a directed edge from the start node to the end node.
@@ -279,15 +402,14 @@ class StateGraph(Graph):
         )
 
         # prepare output channels
-        state_keys = list(self.channels)
         output_channels = (
-            state_keys[0]
-            if state_keys == ["__root__"]
+            "__root__"
+            if len(self.schemas[self.output]) == 1
+            and "__root__" in self.schemas[self.output]
             else [
                 key
-                for key in state_keys
-                if not isinstance(self.channels[key], Context)
-                and not is_managed_value(self.channels[key])
+                for key, val in self.schemas[self.output].items()
+                if not isinstance(val, Context) and not is_managed_value(val)
             ]
         )
 
@@ -295,7 +417,7 @@ class StateGraph(Graph):
             builder=self,
             config_type=self.config_schema,
             nodes={},
-            channels={**self.channels, START: EphemeralValue(self.schema)},
+            channels={**self.channels, START: EphemeralValue(self.input)},
             input_channels=START,
             stream_mode="updates",
             output_channels=output_channels,
@@ -330,12 +452,47 @@ class CompiledStateGraph(CompiledGraph):
     def get_input_schema(
         self, config: Optional[RunnableConfig] = None
     ) -> type[BaseModel]:
-        return self.get_output_schema(config)
+        if isclass(self.builder.input) and issubclass(self.builder.input, BaseModel):
+            return self.builder.input
+        else:
+            keys = list(self.builder.schemas[self.builder.input].keys())
+            if len(keys) == 1 and keys[0] == "__root__":
+                return create_model(  # type: ignore[call-overload]
+                    self.get_name("Input"),
+                    __root__=(self.channels[keys[0]].UpdateType, None),
+                )
+            else:
+                return create_model(  # type: ignore[call-overload]
+                    self.get_name("Input"),
+                    **{
+                        k: (self.channels[k].UpdateType, None)
+                        for k in self.builder.schemas[self.builder.input]
+                        if k in self.channels
+                        and not isinstance(self.channels[k], Context)
+                    },
+                )
 
-    def attach_node(self, key: str, node: Optional[Runnable]) -> None:
-        state_keys = list(self.builder.channels)
+    def get_output_schema(
+        self, config: Optional[RunnableConfig] = None
+    ) -> type[BaseModel]:
+        if isclass(self.builder.input) and issubclass(self.builder.output, BaseModel):
+            return self.builder.output
 
-        def _get_state_key(input: dict, config: RunnableConfig, *, key: str) -> Any:
+        return super().get_output_schema(config)
+
+    def attach_node(self, key: str, node: Optional[StateNodeSpec]) -> None:
+        if key == START:
+            output_keys = [
+                k
+                for k, v in self.builder.schemas[self.builder.input].items()
+                if not isinstance(v, Context) and not is_managed_value(v)
+            ]
+        else:
+            output_keys = list(self.builder.channels)
+
+        def _get_state_key(
+            input: Union[None, dict, Any], config: RunnableConfig, *, key: str
+        ) -> Any:
             if input is None:
                 return SKIP_WRITE
             elif isinstance(input, dict):
@@ -347,9 +504,9 @@ class CompiledStateGraph(CompiledGraph):
                 raise InvalidUpdateError(f"Expected dict, got {input}")
 
         # state updaters
-        state_write_entries = (
+        write_entries = (
             [ChannelWriteEntry("__root__", skip_none=True)]
-            if state_keys == ["__root__"]
+            if output_keys == ["__root__"]
             else [
                 ChannelWriteEntry(
                     key,
@@ -357,7 +514,7 @@ class CompiledStateGraph(CompiledGraph):
                         _get_state_key, key=key, trace=False, recurse=False
                     ),
                 )
-                for key in state_keys
+                for key in output_keys
             ]
         )
 
@@ -369,37 +526,40 @@ class CompiledStateGraph(CompiledGraph):
                 channels=[START],
                 writers=[
                     ChannelWrite(
-                        state_write_entries,
+                        write_entries,
                         tags=[TAG_HIDDEN],
-                        require_at_least_one_of=state_keys,
+                        require_at_least_one_of=output_keys,
                     ),
                 ],
             )
         else:
+            input_schema = node.input if node else self.builder.schema
+            input_values = {
+                k: v if is_managed_value(v) else k
+                for k, v in self.builder.schemas[input_schema].items()
+            }
+            is_single_input = len(input_values) == 1 and "__root__" in input_values
+
             self.channels[key] = EphemeralValue(Any, guard=False)
             self.nodes[key] = PregelNode(
                 triggers=[],
                 # read state keys and managed values
-                channels=(
-                    state_keys
-                    if state_keys == ["__root__"]
-                    else ({chan: chan for chan in state_keys} | self.builder.managed)
-                ),
+                channels=(list(input_values) if is_single_input else input_values),
                 # coerce state dict to schema class (eg. pydantic model)
                 mapper=(
-                    None
-                    if state_keys == ["__root__"]
-                    else partial(_coerce_state, self.builder.schema)
+                    None if is_single_input else partial(_coerce_state, input_schema)
                 ),
                 writers=[
                     # publish to this channel and state keys
                     ChannelWrite(
-                        [ChannelWriteEntry(key, key)] + state_write_entries,
+                        [ChannelWriteEntry(key, key)] + write_entries,
                         tags=[TAG_HIDDEN],
-                        require_at_least_one_of=state_keys,
+                        require_at_least_one_of=output_keys,
                     ),
                 ],
-            ).pipe(node)
+                metadata=node.metadata,
+                retry_policy=node.retry_policy,
+            ).pipe(node.runnable)
 
     def attach_edge(self, starts: Union[str, Sequence[str]], end: str) -> None:
         if isinstance(starts, str):

libs/langgraph/langgraph/prebuilt/__init__.py

@@ -3,7 +3,7 @@ from langgraph.prebuilt import chat_agent_executor
 from langgraph.prebuilt.agent_executor import create_agent_executor
 from langgraph.prebuilt.chat_agent_executor import create_react_agent
 from langgraph.prebuilt.tool_executor import ToolExecutor, ToolInvocation
-from langgraph.prebuilt.tool_node import ToolNode, tools_condition
+from langgraph.prebuilt.tool_node import InjectedState, ToolNode, tools_condition
 from langgraph.prebuilt.tool_validator import ValidationNode
 
 __all__ = [
@@ -15,4 +15,5 @@ __all__ = [
     "ToolNode",
     "tools_condition",
     "ValidationNode",
+    "InjectedState",
 ]

libs/langgraph/langgraph/prebuilt/chat_agent_executor.py

@@ -1,5 +1,15 @@
 import json
-from typing import Annotated, Callable, Optional, Sequence, TypedDict, Union
+import types
+from typing import (
+    Annotated,
+    Callable,
+    Optional,
+    Sequence,
+    Type,
+    TypedDict,
+    TypeVar,
+    Union,
+)
 
 from langchain_core.language_models import LanguageModelLike
 from langchain_core.messages import (
@@ -12,7 +22,7 @@ from langchain_core.runnables import Runnable, RunnableConfig, RunnableLambda
 from langchain_core.tools import BaseTool
 from langchain_core.utils.function_calling import convert_to_openai_function
 
-from langgraph._api.deprecation import deprecated
+from langgraph._api.deprecation import deprecated, deprecated_parameter
 from langgraph.checkpoint import BaseCheckpointSaver
 from langgraph.graph import END, StateGraph
 from langgraph.graph.graph import CompiledGraph
@@ -34,6 +44,26 @@ class AgentState(TypedDict):
     is_last_step: IsLastStep
 
 
+StateSchema = TypeVar("StateSchema", bound=AgentState)
+StateSchemaType = Type[StateSchema]
+
+STATE_MODIFIER_RUNNABLE_NAME = "StateModifier"
+
+MessagesModifier = Union[
+    SystemMessage,
+    str,
+    Callable[[Sequence[BaseMessage]], Sequence[BaseMessage]],
+    Runnable[Sequence[BaseMessage], Sequence[BaseMessage]],
+]
+
+StateModifier = Union[
+    SystemMessage,
+    str,
+    Callable[[StateSchema], Sequence[BaseMessage]],
+    Runnable[StateSchema, Sequence[BaseMessage]],
+]
+
+
 @deprecated("0.0.44", "create_react_agent", removal="0.2.0")
 def create_function_calling_executor(
     model: LanguageModelLike, tools: Union[ToolExecutor, Sequence[BaseTool]]
@@ -164,10 +194,81 @@ def create_function_calling_executor(
     return workflow.compile()
 
 
+def _get_state_modifier_runnable(state_modifier: Optional[StateModifier]) -> Runnable:
+    state_modifier_runnable: Runnable
+    if state_modifier is None:
+        state_modifier_runnable = RunnableLambda(
+            lambda state: state["messages"], name=STATE_MODIFIER_RUNNABLE_NAME
+        )
+    elif isinstance(state_modifier, str):
+        _system_message: BaseMessage = SystemMessage(content=state_modifier)
+        state_modifier_runnable = RunnableLambda(
+            lambda state: [_system_message] + state["messages"],
+            name=STATE_MODIFIER_RUNNABLE_NAME,
+        )
+    elif isinstance(state_modifier, SystemMessage):
+        state_modifier_runnable = RunnableLambda(
+            lambda state: [state_modifier] + state["messages"],
+            name=STATE_MODIFIER_RUNNABLE_NAME,
+        )
+    elif callable(state_modifier):
+        state_modifier_runnable = RunnableLambda(
+            state_modifier, name=STATE_MODIFIER_RUNNABLE_NAME
+        )
+    elif isinstance(state_modifier, Runnable):
+        state_modifier_runnable = state_modifier
+    else:
+        raise ValueError(
+            f"Got unexpected type for `state_modifier`: {type(state_modifier)}"
+        )
+
+    return state_modifier_runnable
+
+
+def _convert_messages_modifier_to_state_modifier(
+    messages_modifier: MessagesModifier,
+) -> StateModifier:
+    state_modifier: StateModifier
+    if isinstance(messages_modifier, (str, SystemMessage)):
+        return messages_modifier
+    elif isinstance(messages_modifier, types.FunctionType):
+
+        def state_modifier(state: AgentState) -> Sequence[BaseMessage]:
+            return messages_modifier(state["messages"])
+
+        return state_modifier
+    elif isinstance(messages_modifier, Runnable):
+        state_modifier = (lambda state: state["messages"]) | messages_modifier
+        return state_modifier
+    raise ValueError(
+        f"Got unexpected type for `messages_modifier`: {type(messages_modifier)}"
+    )
+
+
+def _get_model_preprocessing_runnable(
+    state_modifier: Optional[StateModifier],
+    messages_modifier: Optional[MessagesModifier],
+) -> Runnable:
+    # Add the state or message modifier, if exists
+    if state_modifier is not None and messages_modifier is not None:
+        raise ValueError(
+            "Expected value for either state_modifier or messages_modifier, got values for both"
+        )
+
+    if state_modifier is None and messages_modifier is not None:
+        state_modifier = _convert_messages_modifier_to_state_modifier(messages_modifier)
+
+    return _get_state_modifier_runnable(state_modifier)
+
+
+@deprecated_parameter("messages_modifier", "0.1.9", "state_modifier", removal="0.2.0")
 def create_react_agent(
     model: LanguageModelLike,
     tools: Union[ToolExecutor, Sequence[BaseTool]],
-    messages_modifier: Optional[Union[SystemMessage, str, Callable, Runnable]] = None,
+    *,
+    state_schema: Optional[StateSchemaType] = None,
+    messages_modifier: Optional[MessagesModifier] = None,
+    state_modifier: Optional[StateModifier] = None,
     checkpointer: Optional[BaseCheckpointSaver] = None,
     interrupt_before: Optional[Sequence[str]] = None,
     interrupt_after: Optional[Sequence[str]] = None,
@@ -178,6 +279,9 @@ def create_react_agent(
     Args:
         model: The `LangChain` chat model that supports tool calling.
         tools: A list of tools or a ToolExecutor instance.
+        state_schema: An optional state schema that defines graph state.
+            Must have `messages` and `is_last_step` keys.
+            Defaults to `AgentState` that defines those two keys.
         messages_modifier: An optional
             messages modifier. This applies to messages BEFORE they are passed into the LLM.
 
@@ -187,6 +291,17 @@ def create_react_agent(
             - str: This is converted to a SystemMessage and added to the beginning of the list of messages.
             - Callable: This function should take in a list of messages and the output is then passed to the language model.
             - Runnable: This runnable should take in a list of messages and the output is then passed to the language model.
+            !!! Warning
+                `messages_modifier` parameter is deprecated as of version 0.1.9 and will be removed in 0.2.0
+        state_modifier: An optional
+            state modifier. This takes full graph state BEFORE the LLM is called and prepares the input to LLM.
+
+            Can take a few different forms:
+
+            - SystemMessage: this is added to the beginning of the list of messages in state["messages"].
+            - str: This is converted to a SystemMessage and added to the beginning of the list of messages in state["messages"].
+            - Callable: This function should take in full graph state and the output is then passed to the language model.
+            - Runnable: This runnable should take in full graph state and the output is then passed to the language model.
         checkpointer: An optional checkpoint saver object. This is useful for persisting
             the state of the graph (e.g., as chat memory).
         interrupt_before: An optional list of node names to interrupt before.
@@ -282,7 +397,7 @@ def create_react_agent(
 
         ```pycon
         >>> system_prompt = "You are a helpful bot named Fred."
-        >>> graph = create_react_agent(model, tools, messages_modifier=system_prompt)
+        >>> graph = create_react_agent(model, tools, state_modifier=system_prompt)
         >>> inputs = {"messages": [("user", "What's your name? And what's the weather in SF?")]}
         >>> for s in graph.stream(inputs, stream_mode="values"):
         ...     message = s["messages"][-1]
@@ -314,11 +429,11 @@ def create_react_agent(
         ...     ("placeholder", "{messages}"),
         ...     ("user", "Remember, always be polite!"),
         ... ])
-        >>> def modify_messages(messages: list):
+        >>> def modify_state_messages(state: AgentState):
         ...     # You can do more complex modifications here
-        ...     return prompt.invoke({"messages": messages})
+        ...     return prompt.invoke({"messages": state["messages"]})
         >>>
-        >>> graph = create_react_agent(model, tools, messages_modifier=modify_messages)
+        >>> graph = create_react_agent(model, tools, state_modifier=modify_state_messages)
         >>> inputs = {"messages": [("user", "What's your name? And what's the weather in SF?")]}
         >>> for s in graph.stream(inputs, stream_mode="values"):
         ...     message = s["messages"][-1]
@@ -328,6 +443,32 @@ def create_react_agent(
         ...         message.pretty_print()
         ```
 
+        Add complex prompt with custom graph state:
+
+        ```pycon
+        >>> from typing import TypedDict
+        >>> prompt = ChatPromptTemplate.from_messages(
+        ...     [
+        ...         ("system", "Today is {today}"),
+        ...         ("placeholder", "{messages}"),
+        ...     ]
+        ... )
+        >>>
+        >>> class CustomState(TypedDict):
+        ...     today: str
+        ...     messages: Annotated[list[BaseMessage], add_messages]
+        ...     is_last_step: str
+        >>>
+        >>> graph = create_react_agent(model, tools, state_schema=CustomState, state_modifier=prompt)
+        >>> inputs = {"messages": [("user", "What's today's date? And what's the weather in SF?")], "today": "July 16, 2004"}
+        >>> for s in graph.stream(inputs, stream_mode="values"):
+        ...     message = s["messages"][-1]
+        ...     if isinstance(message, tuple):
+        ...         print(message)
+        ...     else:
+        ...         message.pretty_print()
+        ```
+
         Add "chat memory" to the graph:
 
         ```pycon
@@ -404,6 +545,12 @@ def create_react_agent(
         ```
     """
 
+    if state_schema is not None:
+        if missing_keys := {"messages", "is_last_step"} - set(
+            state_schema.__annotations__
+        ):
+            raise ValueError(f"Missing required key(s) {missing_keys} in state_schema")
+
     if isinstance(tools, ToolExecutor):
         tool_classes = tools.tools
     else:
@@ -421,28 +568,15 @@ def create_react_agent(
         else:
             return "continue"
 
-    # Add the message modifier, if exists
-    if messages_modifier is None:
-        model_runnable = model
-    elif isinstance(messages_modifier, str):
-        _system_message: BaseMessage = SystemMessage(content=messages_modifier)
-        model_runnable = (lambda messages: [_system_message] + messages) | model
-    elif isinstance(messages_modifier, SystemMessage):
-        model_runnable = (lambda messages: [messages_modifier] + messages) | model
-    elif isinstance(messages_modifier, (Callable, Runnable)):
-        model_runnable = messages_modifier | model
-    else:
-        raise ValueError(
-            f"Got unexpected type for `messages_modifier`: {type(messages_modifier)}"
-        )
+    preprocessor = _get_model_preprocessing_runnable(state_modifier, messages_modifier)
+    model_runnable = preprocessor | model
 
     # Define the function that calls the model
     def call_model(
         state: AgentState,
         config: RunnableConfig,
     ):
-        messages = state["messages"]
-        response = model_runnable.invoke(messages, config)
+        response = model_runnable.invoke(state, config)
         if state["is_last_step"] and response.tool_calls:
             return {
                 "messages": [
@@ -456,8 +590,7 @@ def create_react_agent(
         return {"messages": [response]}
 
     async def acall_model(state: AgentState, config: RunnableConfig):
-        messages = state["messages"]
-        response = await model_runnable.ainvoke(messages, config)
+        response = await model_runnable.ainvoke(state, config)
         if state["is_last_step"] and response.tool_calls:
             return {
                 "messages": [
@@ -471,7 +604,7 @@ def create_react_agent(
         return {"messages": [response]}
 
     # Define a new graph
-    workflow = StateGraph(AgentState)
+    workflow = StateGraph(state_schema or AgentState)
 
     # Define the two nodes we will cycle between
     workflow.add_node("agent", RunnableLambda(call_model, acall_model))

libs/langgraph/langgraph/prebuilt/tool_node.py

@@ -1,30 +1,38 @@
 import asyncio
-import json
-from typing import Any, Callable, Dict, Literal, Optional, Sequence, Union
+from copy import copy
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    Sequence,
+    Tuple,
+    Union,
+    cast,
+)
 
 from langchain_core.messages import AIMessage, AnyMessage, ToolCall, ToolMessage
 from langchain_core.runnables import RunnableConfig
-from langchain_core.runnables.config import get_executor_for_config
-from langchain_core.tools import BaseTool
+from langchain_core.runnables.config import get_config_list, get_executor_for_config
+from langchain_core.tools import BaseTool, InjectedToolArg
 from langchain_core.tools import tool as create_tool
+from typing_extensions import get_args
 
 from langgraph.utils import RunnableCallable
 
-
-def str_output(output: Any) -> str:
-    if isinstance(output, str):
-        return output
-    else:
-        try:
-            return json.dumps(output)
-        except Exception:
-            return str(output)
+INVALID_TOOL_NAME_ERROR_TEMPLATE = (
+    "Error: {requested_tool} is not a valid tool, try one of [{available_tools}]."
+)
+TOOL_CALL_ERROR_TEMPLATE = "Error: {error}\n Please fix your mistakes."
 
 
 class ToolNode(RunnableCallable):
-    """A node that runs the tools requested in the last AIMessage. It can be used
-    either in StateGraph with a "messages" key or in MessageGraph. If multiple
-    tool calls are requested, they will be run in parallel. The output will be
+    """A node that runs the tools called in the last AIMessage.
+
+    It can be used either in StateGraph with a "messages" key or in MessageGraph. If
+    multiple tool calls are requested, they will be run in parallel. The output will be
     a list of ToolMessages, one for each tool call.
 
     The `ToolNode` is roughly analogous to:
@@ -65,39 +73,49 @@ class ToolNode(RunnableCallable):
     def _func(
         self, input: Union[list[AnyMessage], dict[str, Any]], config: RunnableConfig
     ) -> Any:
-        if isinstance(input, list):
-            output_type = "list"
-            message: AnyMessage = input[-1]
-        elif messages := input.get("messages", []):
-            output_type = "dict"
-            message = messages[-1]
-        else:
-            raise ValueError("No message found in input")
-
-        if not isinstance(message, AIMessage):
-            raise ValueError("Last message is not an AIMessage")
-
-        def run_one(call: ToolCall):
-            try:
-                output = self.tools_by_name[call["name"]].invoke(call["args"], config)
-            except Exception as e:
-                if not self.handle_tool_errors:
-                    raise e
-                output = f"Error: {repr(e)}\n Please fix your mistakes."
-            return ToolMessage(
-                content=str_output(output), name=call["name"], tool_call_id=call["id"]
-            )
-
+        tool_calls, output_type = self._parse_input(input)
+        config_list = get_config_list(config, len(tool_calls))
         with get_executor_for_config(config) as executor:
-            outputs = [*executor.map(run_one, message.tool_calls)]
-            if output_type == "list":
-                return outputs
-            else:
-                return {"messages": outputs}
+            outputs = [*executor.map(self._run_one, tool_calls, config_list)]
+        return outputs if output_type == "list" else {"messages": outputs}
 
     async def _afunc(
         self, input: Union[list[AnyMessage], dict[str, Any]], config: RunnableConfig
     ) -> Any:
+        tool_calls, output_type = self._parse_input(input)
+        outputs = await asyncio.gather(
+            *(self._arun_one(call, config) for call in tool_calls)
+        )
+        return outputs if output_type == "list" else {"messages": outputs}
+
+    def _run_one(self, call: ToolCall, config: RunnableConfig) -> ToolMessage:
+        if invalid_tool_message := self._validate_tool_call(call):
+            return invalid_tool_message
+
+        try:
+            input = {**call, **{"type": "tool_call"}}
+            return self.tools_by_name[call["name"]].invoke(input, config)
+        except Exception as e:
+            if not self.handle_tool_errors:
+                raise e
+            content = TOOL_CALL_ERROR_TEMPLATE.format(error=repr(e))
+            return ToolMessage(content, name=call["name"], tool_call_id=call["id"])
+
+    async def _arun_one(self, call: ToolCall, config: RunnableConfig) -> ToolMessage:
+        if invalid_tool_message := self._validate_tool_call(call):
+            return invalid_tool_message
+        try:
+            input = {**call, **{"type": "tool_call"}}
+            return await self.tools_by_name[call["name"]].ainvoke(input, config)
+        except Exception as e:
+            if not self.handle_tool_errors:
+                raise e
+            content = TOOL_CALL_ERROR_TEMPLATE.format(error=repr(e))
+            return ToolMessage(content, name=call["name"], tool_call_id=call["id"])
+
+    def _parse_input(
+        self, input: Union[list[AnyMessage], dict[str, Any]]
+    ) -> Tuple[List[ToolCall], Literal["list", "dict"]]:
         if isinstance(input, list):
             output_type = "list"
             message: AnyMessage = input[-1]
@@ -110,24 +128,54 @@ class ToolNode(RunnableCallable):
         if not isinstance(message, AIMessage):
             raise ValueError("Last message is not an AIMessage")
 
-        async def run_one(call: ToolCall):
-            try:
-                output = await self.tools_by_name[call["name"]].ainvoke(
-                    call["args"], config
-                )
-            except Exception as e:
-                if not self.handle_tool_errors:
-                    raise e
-                output = f"Error: {repr(e)}\n Please fix your mistakes."
-            return ToolMessage(
-                content=str_output(output), name=call["name"], tool_call_id=call["id"]
-            )
+        tool_calls = [
+            self._inject_state(call, input)
+            for call in cast(AIMessage, message).tool_calls
+        ]
+        return tool_calls, output_type
 
-        outputs = await asyncio.gather(*(run_one(call) for call in message.tool_calls))
-        if output_type == "list":
-            return outputs
+    def _validate_tool_call(self, call: ToolCall) -> Optional[ToolMessage]:
+        if (requested_tool := call["name"]) not in self.tools_by_name:
+            content = INVALID_TOOL_NAME_ERROR_TEMPLATE.format(
+                requested_tool=requested_tool,
+                available_tools=", ".join(self.tools_by_name.keys()),
+            )
+            return ToolMessage(content, name=requested_tool, tool_call_id=call["id"])
         else:
-            return {"messages": outputs}
+            return None
+
+    def _inject_state(
+        self, tool_call: ToolCall, input: Union[list[AnyMessage], dict[str, Any]]
+    ) -> ToolCall:
+        if tool_call["name"] not in self.tools_by_name:
+            return tool_call
+        state_args = _get_state_args(self.tools_by_name[tool_call["name"]])
+        if state_args and not isinstance(input, dict):
+            required_fields = list(state_args.values())
+            if (
+                len(required_fields) == 1
+                and required_fields[0] == "messages"
+                or required_fields[0] is None
+            ):
+                input = {"messages": input}
+            else:
+                err_msg = (
+                    f"Invalid input to ToolNode. Tool {tool_call['name']} requires "
+                    f"graph state dict as input."
+                )
+                if any(state_field for state_field in state_args.values()):
+                    required_fields_str = ", ".join(f for f in required_fields if f)
+                    err_msg += f" State should contain fields {required_fields_str}."
+                raise ValueError(err_msg)
+        tool_call_copy: ToolCall = copy(tool_call)
+        tool_call_copy["args"] = {
+            **tool_call_copy["args"],
+            **{
+                tool_arg: cast(dict, input)[state_field] if state_field else input
+                for tool_arg, state_field in state_args.items()
+            },
+        }
+        return tool_call_copy
 
 
 def tools_condition(
@@ -185,3 +233,92 @@ def tools_condition(
     if hasattr(ai_message, "tool_calls") and len(ai_message.tool_calls) > 0:
         return "tools"
     return "__end__"
+
+
+class InjectedState(InjectedToolArg):
+    """Annotation for a Tool arg that is meant to be populated with the graph state.
+
+    Any Tool argument annotated with InjectedState will be hidden from a tool-calling
+    model, so that the model doesn't attempt to generate the argument. If using
+    ToolNode, the appropriate graph state field will be automatically injected into
+    the model-generated tool args.
+
+    Args:
+        field: The key from state to insert. If None, the entire state is expected to
+            be passed in.
+
+    Example:
+        ```python
+        from typing import List
+        from typing_extensions import Annotated, TypedDict
+
+        from langchain_core.messages import BaseMessage, AIMessage
+        from langchain_core.tools import tool
+
+        from langgraph.prebuilt import InjectedState, ToolNode
+
+
+        class AgentState(TypedDict):
+            messages: List[BaseMessage]
+            foo: str
+
+        @tool
+        def state_tool(x: int, state: Annotated[dict, InjectedState]) -> str:
+            '''Do something with state.'''
+            if len(state["messages"]) > 2:
+                return state["foo"] + str(x)
+            else:
+                return "not enough messages"
+
+        @tool
+        def foo_tool(x: int, foo: Annotated[str, InjectedState("foo")]) -> str:
+            '''Do something else with state.'''
+            return foo + str(x + 1)
+
+        node = ToolNode([state_tool, foo_tool])
+
+        tool_call1 = {"name": "state_tool", "args": {"x": 1}, "id": "1", "type": "tool_call"}
+        tool_call2 = {"name": "foo_tool", "args": {"x": 1}, "id": "2", "type": "tool_call"}
+        state = {
+            "messages": [AIMessage("", tool_calls=[tool_call1, tool_call2])],
+            "foo": "bar",
+        }
+        node.invoke(state)
+        ```
+
+        ```pycon
+        [
+            ToolMessage(content='not enough messages', name='state_tool', tool_call_id='1'),
+            ToolMessage(content='bar2', name='foo_tool', tool_call_id='2')
+        ]
+        ```
+    """  # noqa: E501
+
+    def __init__(self, field: Optional[str] = None) -> None:
+        self.field = field
+
+
+def _get_state_args(tool: BaseTool) -> Dict[str, Optional[str]]:
+    full_schema = tool.get_input_schema()
+    tool_args_to_state_fields: Dict = {}
+    for name, type_ in full_schema.__annotations__.items():
+        injections = [
+            type_arg
+            for type_arg in get_args(type_)
+            if isinstance(type_arg, InjectedState)
+            or (isinstance(type_arg, type) and issubclass(type_arg, InjectedState))
+        ]
+        if len(injections) > 1:
+            raise ValueError(
+                "A tool argument should not be annotated with InjectedState more than "
+                f"once. Received arg {name} with annotations {injections}."
+            )
+        elif len(injections) == 1:
+            injection = injections[0]
+            if isinstance(injection, InjectedState) and injection.field:
+                tool_args_to_state_fields[name] = injection.field
+            else:
+                tool_args_to_state_fields[name] = None
+        else:
+            pass
+    return tool_args_to_state_fields

libs/langgraph/langgraph/pregel/algo.py

@@ -0,0 +1,427 @@
+import json
+from collections import defaultdict, deque
+from functools import partial
+from typing import (
+    Any,
+    Callable,
+    Iterator,
+    Literal,
+    Mapping,
+    NamedTuple,
+    Optional,
+    Protocol,
+    Sequence,
+    Union,
+    overload,
+)
+from uuid import UUID, uuid5
+
+from langchain_core.callbacks.manager import AsyncParentRunManager, ParentRunManager
+from langchain_core.runnables.config import (
+    RunnableConfig,
+    merge_configs,
+    patch_config,
+)
+
+from langgraph.channels.base import BaseChannel
+from langgraph.channels.context import Context
+from langgraph.channels.manager import ChannelsManager, create_checkpoint
+from langgraph.checkpoint.base import Checkpoint, copy_checkpoint
+from langgraph.constants import (
+    CONFIG_KEY_READ,
+    CONFIG_KEY_SEND,
+    INTERRUPT,
+    RESERVED,
+    TAG_HIDDEN,
+    TASKS,
+    Send,
+)
+from langgraph.errors import EmptyChannelError, InvalidUpdateError
+from langgraph.managed.base import ManagedValueMapping, is_managed_value
+from langgraph.pregel.io import read_channel, read_channels
+from langgraph.pregel.log import logger
+from langgraph.pregel.read import PregelNode
+from langgraph.pregel.types import All, PregelExecutableTask, PregelTaskDescription
+
+
+class WritesProtocol(Protocol):
+    name: str
+    writes: Sequence[tuple[str, Any]]
+    triggers: Sequence[str]
+
+
+class PregelTaskWrites(NamedTuple):
+    name: str
+    writes: Sequence[tuple[str, Any]]
+    triggers: Sequence[str]
+
+
+def should_interrupt(
+    checkpoint: Checkpoint,
+    interrupt_nodes: Union[All, Sequence[str]],
+    tasks: list[PregelExecutableTask],
+) -> bool:
+    version_type = type(next(iter(checkpoint["channel_versions"].values()), None))
+    null_version = version_type()
+    seen = checkpoint["versions_seen"].get(INTERRUPT, {})
+    return (
+        # interrupt if any channel has been updated since last interrupt
+        any(
+            version > seen.get(chan, null_version)
+            for chan, version in checkpoint["channel_versions"].items()
+        )
+        # and any triggered node is in interrupt_nodes list
+        and any(
+            task.name
+            for task in tasks
+            if (
+                (not task.config or TAG_HIDDEN not in task.config.get("tags"))
+                if interrupt_nodes == "*"
+                else task.name in interrupt_nodes
+            )
+        )
+    )
+
+
+def local_read(
+    checkpoint: Checkpoint,
+    channels: Mapping[str, BaseChannel],
+    task: WritesProtocol,
+    config: RunnableConfig,
+    select: Union[list[str], str],
+    fresh: bool = False,
+) -> Union[dict[str, Any], Any]:
+    if fresh:
+        new_checkpoint = create_checkpoint(copy_checkpoint(checkpoint), channels, -1)
+        context_channels = {k: v for k, v in channels.items() if isinstance(v, Context)}
+        with ChannelsManager(
+            {k: v for k, v in channels.items() if k not in context_channels},
+            new_checkpoint,
+            config,
+        ) as channels:
+            all_channels = {**channels, **context_channels}
+            apply_writes(new_checkpoint, all_channels, [task], None)
+            return read_channels(all_channels, select)
+    else:
+        return read_channels(channels, select)
+
+
+def local_write(
+    commit: Callable[[Sequence[tuple[str, Any]]], None],
+    processes: Mapping[str, PregelNode],
+    channels: Mapping[str, BaseChannel],
+    writes: Sequence[tuple[str, Any]],
+) -> None:
+    for chan, value in writes:
+        if chan == TASKS:
+            if not isinstance(value, Send):
+                raise InvalidUpdateError(
+                    f"Invalid packet type, expected Packet, got {value}"
+                )
+            if value.node not in processes:
+                raise InvalidUpdateError(f"Invalid node name {value.node} in packet")
+        elif chan not in channels:
+            logger.warning(f"Skipping write for channel '{chan}' which has no readers")
+    commit(writes)
+
+
+def increment(current: Optional[int], channel: BaseChannel) -> int:
+    return current + 1 if current is not None else 1
+
+
+def apply_writes(
+    checkpoint: Checkpoint,
+    channels: Mapping[str, BaseChannel],
+    tasks: Sequence[WritesProtocol],
+    get_next_version: Optional[Callable[[int, BaseChannel], int]],
+) -> None:
+    # update seen versions
+    for task in tasks:
+        checkpoint["versions_seen"].setdefault(task.name, {}).update(
+            {
+                chan: checkpoint["channel_versions"][chan]
+                for chan in task.triggers
+                if chan in checkpoint["channel_versions"]
+            }
+        )
+
+    # Find the highest version of all channels
+    if checkpoint["channel_versions"]:
+        max_version = max(checkpoint["channel_versions"].values())
+    else:
+        max_version = None
+    # Consume all channels that were read
+    for chan in {
+        chan for task in tasks for chan in task.triggers if chan not in RESERVED
+    }:
+        if channels[chan].consume():
+            if get_next_version is not None:
+                checkpoint["channel_versions"][chan] = get_next_version(
+                    max_version, channels[chan]
+                )
+
+    # clear pending sends
+    if checkpoint["pending_sends"]:
+        checkpoint["pending_sends"].clear()
+
+    # Group writes by channel
+    pending_writes_by_channel: dict[str, list[Any]] = defaultdict(list)
+    for task in tasks:
+        for chan, val in task.writes:
+            if chan == TASKS:
+                checkpoint["pending_sends"].append(val)
+            else:
+                pending_writes_by_channel[chan].append(val)
+
+    # Find the highest version of all channels
+    if checkpoint["channel_versions"]:
+        max_version = max(checkpoint["channel_versions"].values())
+    else:
+        max_version = None
+
+    # Apply writes to channels
+    updated_channels: set[str] = set()
+    for chan, vals in pending_writes_by_channel.items():
+        if chan in channels:
+            try:
+                updated = channels[chan].update(vals)
+            except InvalidUpdateError as e:
+                raise InvalidUpdateError(
+                    f"Invalid update for channel {chan} with values {vals}"
+                ) from e
+            if updated and get_next_version is not None:
+                checkpoint["channel_versions"][chan] = get_next_version(
+                    max_version, channels[chan]
+                )
+            updated_channels.add(chan)
+
+    # Channels that weren't updated in this step are notified of a new step
+    for chan in channels:
+        if chan not in updated_channels:
+            if channels[chan].update([]) and get_next_version is not None:
+                checkpoint["channel_versions"][chan] = get_next_version(
+                    max_version, channels[chan]
+                )
+
+
+@overload
+def prepare_next_tasks(
+    checkpoint: Checkpoint,
+    processes: Mapping[str, PregelNode],
+    channels: Mapping[str, BaseChannel],
+    managed: ManagedValueMapping,
+    config: RunnableConfig,
+    step: int,
+    for_execution: Literal[False],
+    manager: Literal[None] = None,
+) -> list[PregelTaskDescription]:
+    ...
+
+
+@overload
+def prepare_next_tasks(
+    checkpoint: Checkpoint,
+    processes: Mapping[str, PregelNode],
+    channels: Mapping[str, BaseChannel],
+    managed: ManagedValueMapping,
+    config: RunnableConfig,
+    step: int,
+    for_execution: Literal[True],
+    manager: Union[None, ParentRunManager, AsyncParentRunManager],
+) -> list[PregelExecutableTask]:
+    ...
+
+
+def prepare_next_tasks(
+    checkpoint: Checkpoint,
+    processes: Mapping[str, PregelNode],
+    channels: Mapping[str, BaseChannel],
+    managed: ManagedValueMapping,
+    config: RunnableConfig,
+    step: int,
+    *,
+    for_execution: bool,
+    manager: Union[None, ParentRunManager, AsyncParentRunManager] = None,
+) -> Union[list[PregelTaskDescription], list[PregelExecutableTask]]:
+    tasks: Union[list[PregelTaskDescription], list[PregelExecutableTask]] = []
+    # Consume pending packets
+    for packet in checkpoint["pending_sends"]:
+        if not isinstance(packet, Send):
+            logger.warn(f"Ignoring invalid packet type {type(packet)} in pending sends")
+            continue
+        if for_execution:
+            proc = processes[packet.node]
+            if node := proc.get_node():
+                triggers = [TASKS]
+                metadata = {
+                    "langgraph_step": step,
+                    "langgraph_node": packet.node,
+                    "langgraph_triggers": triggers,
+                    "langgraph_task_idx": len(tasks),
+                }
+                task_id = str(uuid5(UUID(checkpoint["id"]), json.dumps(metadata)))
+                writes = deque()
+                tasks.append(
+                    PregelExecutableTask(
+                        packet.node,
+                        packet.arg,
+                        node,
+                        writes,
+                        patch_config(
+                            merge_configs(
+                                config,
+                                processes[packet.node].config,
+                                {"metadata": metadata},
+                            ),
+                            run_name=packet.node,
+                            callbacks=(
+                                manager.get_child(f"graph:step:{step}")
+                                if manager
+                                else None
+                            ),
+                            configurable={
+                                # deque.extend is thread-safe
+                                CONFIG_KEY_SEND: partial(
+                                    local_write, writes.extend, processes, channels
+                                ),
+                                CONFIG_KEY_READ: partial(
+                                    local_read,
+                                    checkpoint,
+                                    channels,
+                                    PregelTaskWrites(packet.node, writes, triggers),
+                                    config,
+                                ),
+                            },
+                        ),
+                        triggers,
+                        proc.retry_policy,
+                        task_id,
+                    )
+                )
+        else:
+            tasks.append(PregelTaskDescription(packet.node, packet.arg))
+    # Check if any processes should be run in next step
+    # If so, prepare the values to be passed to them
+    version_type = type(next(iter(checkpoint["channel_versions"].values()), None))
+    null_version = version_type()
+    if null_version is None:
+        return tasks
+    for name, proc in processes.items():
+        seen = checkpoint["versions_seen"].get(name, {})
+        # If any of the channels read by this process were updated
+        if triggers := sorted(
+            chan
+            for chan in proc.triggers
+            if not isinstance(
+                read_channel(channels, chan, return_exception=True), EmptyChannelError
+            )
+            and checkpoint["channel_versions"].get(chan, null_version)
+            > seen.get(chan, null_version)
+        ):
+            try:
+                val = next(_proc_input(step, name, proc, managed, channels))
+            except StopIteration:
+                continue
+
+            if for_execution:
+                if node := proc.get_node():
+                    metadata = {
+                        "langgraph_step": step,
+                        "langgraph_node": name,
+                        "langgraph_triggers": triggers,
+                        "langgraph_task_idx": len(tasks),
+                    }
+                    task_id = str(uuid5(UUID(checkpoint["id"]), json.dumps(metadata)))
+                    writes = deque()
+                    tasks.append(
+                        PregelExecutableTask(
+                            name,
+                            val,
+                            node,
+                            writes,
+                            patch_config(
+                                merge_configs(
+                                    config,
+                                    proc.config,
+                                    {"metadata": metadata},
+                                ),
+                                run_name=name,
+                                callbacks=(
+                                    manager.get_child(f"graph:step:{step}")
+                                    if manager
+                                    else None
+                                ),
+                                configurable={
+                                    # deque.extend is thread-safe
+                                    CONFIG_KEY_SEND: partial(
+                                        local_write, writes.extend, processes, channels
+                                    ),
+                                    CONFIG_KEY_READ: partial(
+                                        local_read,
+                                        checkpoint,
+                                        channels,
+                                        PregelTaskWrites(name, writes, triggers),
+                                        config,
+                                    ),
+                                },
+                            ),
+                            triggers,
+                            proc.retry_policy,
+                            task_id,
+                        )
+                    )
+            else:
+                tasks.append(PregelTaskDescription(name, val))
+    return tasks
+
+
+def _proc_input(
+    step: int,
+    name: str,
+    proc: PregelNode,
+    managed: ManagedValueMapping,
+    channels: Mapping[str, BaseChannel],
+) -> Iterator[Any]:
+    # If all trigger channels subscribed by this process are not empty
+    # then invoke the process with the values of all non-empty channels
+    if isinstance(proc.channels, dict):
+        try:
+            val: dict = {
+                k: read_channel(
+                    channels,
+                    chan,
+                    catch=chan not in proc.triggers,
+                )
+                for k, chan in proc.channels.items()
+                if isinstance(chan, str)
+            }
+
+            managed_values = {}
+            for key, chan in proc.channels.items():
+                if is_managed_value(chan):
+                    managed_values[key] = managed[key](
+                        step, PregelTaskDescription(name, val)
+                    )
+
+            val.update(managed_values)
+        except EmptyChannelError:
+            return
+    elif isinstance(proc.channels, list):
+        for chan in proc.channels:
+            try:
+                val = read_channel(channels, chan, catch=False)
+                break
+            except EmptyChannelError:
+                pass
+        else:
+            return
+    else:
+        raise RuntimeError(
+            "Invalid channels type, expected list or dict, got {proc.channels}"
+        )
+
+    # If the process has a mapper, apply it to the value
+    if proc.mapper is not None:
+        val = proc.mapper(val)
+
+    yield val

libs/langgraph/langgraph/pregel/debug.py

@@ -66,7 +66,7 @@ def map_debug_tasks(
     step: int, tasks: list[PregelExecutableTask]
 ) -> Iterator[DebugOutputTask]:
     ts = datetime.now(timezone.utc).isoformat()
-    for name, input, _, _, config, triggers in tasks:
+    for name, input, _, _, config, triggers, _, _ in tasks:
         if config is not None and TAG_HIDDEN in config.get("tags", []):
             continue
 
@@ -91,7 +91,7 @@ def map_debug_task_results(
     stream_channels_list: Sequence[str],
 ) -> Iterator[DebugOutputTaskResult]:
     ts = datetime.now(timezone.utc).isoformat()
-    for name, _, _, writes, config, _ in tasks:
+    for name, _, _, writes, config, _, _, _ in tasks:
         if config is not None and TAG_HIDDEN in config.get("tags", []):
             continue
 
@@ -134,11 +134,11 @@ def print_step_tasks(step: int, next_tasks: list[PregelExecutableTask]) -> None:
     print(
         f"{get_colored_text(f'[{step}:tasks]', color='blue')} "
         + get_bolded_text(
-            f"Starting step {step} with {n_tasks} task{'s' if n_tasks > 1 else ''}:\n"
+            f"Starting step {step} with {n_tasks} task{'s' if n_tasks != 1 else ''}:\n"
         )
         + "\n".join(
-            f"- {get_colored_text(name, 'green')} -> {pformat(val)}"
-            for name, val, _, _, _, _ in next_tasks
+            f"- {get_colored_text(task.name, 'green')} -> {pformat(task.input)}"
+            for task in next_tasks
         )
     )
 
@@ -153,7 +153,7 @@ def print_step_writes(
     print(
         f"{get_colored_text(f'[{step}:writes]', color='blue')} "
         + get_bolded_text(
-            f"Finished step {step} with writes to {len(by_channel)} channel{'s' if len(by_channel) > 1 else ''}:\n"
+            f"Finished step {step} with writes to {len(by_channel)} channel{'s' if len(by_channel) != 1 else ''}:\n"
         )
         + "\n".join(
             f"- {get_colored_text(name, 'yellow')} -> {', '.join(pformat(v) for v in vals)}"

libs/langgraph/langgraph/pregel/executor.py

@@ -0,0 +1,129 @@
+import asyncio
+import concurrent.futures
+import sys
+from contextlib import contextmanager
+from contextvars import copy_context
+from types import TracebackType
+from typing import (
+    AsyncContextManager,
+    Awaitable,
+    Callable,
+    Iterator,
+    Optional,
+    Protocol,
+    TypeVar,
+)
+
+from langchain_core.runnables import RunnableConfig
+from langchain_core.runnables.config import get_executor_for_config
+from typing_extensions import ParamSpec
+
+P = ParamSpec("P")
+T = TypeVar("T")
+
+
+class Submit(Protocol[P, T]):
+    def __call__(
+        self,
+        fn: Callable[P, T],
+        *args: P.args,
+        __name__: Optional[str] = None,
+        __cancel_on_exit__: bool = False,
+        **kwargs: P.kwargs,
+    ) -> concurrent.futures.Future[T]:
+        ...
+
+
+@contextmanager
+def BackgroundExecutor(config: RunnableConfig) -> Iterator[Submit]:
+    tasks: dict[concurrent.futures.Future, bool] = {}
+    with get_executor_for_config(config) as executor:
+
+        def done(task: concurrent.futures.Future) -> None:
+            try:
+                task.result()
+            except BaseException:
+                pass
+            else:
+                tasks.pop(task)
+
+        def submit(
+            fn: Callable[P, T],
+            *args: P.args,
+            __name__: Optional[str] = None,  # currently not used in sync version
+            __cancel_on_exit__: bool = False,
+            **kwargs: P.kwargs,
+        ) -> concurrent.futures.Future:
+            task = executor.submit(fn, *args, **kwargs)
+            tasks[task] = __cancel_on_exit__
+            task.add_done_callback(done)
+            return task
+
+        try:
+            yield submit
+        finally:
+            for task, cancel in tasks.items():
+                if cancel:
+                    task.cancel()
+            # executor waits for all tasks to finish on exit
+    for task in tasks:
+        # the first task to have raised an exception will be re-raised here
+        task.result()
+
+
+class AsyncBackgroundExecutor(AsyncContextManager):
+    def __init__(self) -> None:
+        self.context_not_supported = sys.version_info < (3, 11)
+        self.tasks: dict[asyncio.Task, bool] = {}
+        self.sentinel = object()
+
+    def submit(
+        self,
+        fn: Callable[P, Awaitable[T]],
+        *args: P.args,
+        __name__: Optional[str] = None,
+        __cancel_on_exit__: bool = False,
+        **kwargs: P.kwargs,
+    ) -> asyncio.Task[T]:
+        coro = fn(*args, **kwargs)
+        if self.context_not_supported:
+            task = asyncio.create_task(coro, name=__name__)
+        else:
+            task = asyncio.create_task(coro, name=__name__, context=copy_context())
+        self.tasks[task] = __cancel_on_exit__
+        task.add_done_callback(self.done)
+        return task
+
+    def done(self, task: asyncio.Task) -> None:
+        try:
+            task.result()
+        except BaseException:
+            pass
+        else:
+            self.tasks.pop(task)
+
+    async def __aenter__(self) -> "submit":
+        return self.submit
+
+    async def exit(self) -> None:
+        fut = asyncio.gather(*self.tasks, return_exceptions=True)
+        try:
+            rtns = await asyncio.shield(fut)
+        finally:
+            del self.tasks
+        for rtn in rtns:
+            # if this is ever changed to BaseException, need to ignore CancelledError
+            if isinstance(rtn, Exception):
+                raise rtn
+
+    async def __aexit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> Optional[bool]:
+        for task, cancel in self.tasks.items():
+            if cancel:
+                task.cancel(self.sentinel)
+        # wait for all background tasks to finish, shielded from cancellation
+        await asyncio.shield(self.exit())

libs/langgraph/langgraph/pregel/io.py

@@ -104,9 +104,9 @@ def map_output_updates(
     ]
     if isinstance(output_channels, str):
         if updated := [
-            (node, value)
-            for node, _, _, writes, _, _ in output_tasks
-            for chan, value in writes
+            (task.name, value)
+            for task in output_tasks
+            for chan, value in task.writes
             if chan == output_channels
         ]:
             grouped = defaultdict(list)
@@ -119,11 +119,11 @@ def map_output_updates(
     else:
         if updated := [
             (
-                node,
-                {chan: value for chan, value in writes if chan in output_channels},
+                task.name,
+                {chan: value for chan, value in task.writes if chan in output_channels},
             )
-            for node, _, _, writes, _, _ in output_tasks
-            if any(chan in output_channels for chan, _ in writes)
+            for task in output_tasks
+            if any(chan in output_channels for chan, _ in task.writes)
         ]:
             grouped = defaultdict(list)
             for node, value in updated:

libs/langgraph/langgraph/pregel/loop.py

@@ -0,0 +1,423 @@
+import asyncio
+from collections import deque
+from contextlib import AsyncExitStack, ExitStack
+from types import TracebackType
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    AsyncContextManager,
+    Callable,
+    ContextManager,
+    List,
+    Literal,
+    Mapping,
+    Optional,
+    Sequence,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
+)
+
+from langchain_core.callbacks import AsyncParentRunManager, ParentRunManager
+from langchain_core.runnables import RunnableConfig
+from typing_extensions import Self
+
+from langgraph.channels.base import BaseChannel
+from langgraph.channels.manager import (
+    AsyncChannelsManager,
+    ChannelsManager,
+    create_checkpoint,
+)
+from langgraph.checkpoint.base import (
+    BaseCheckpointSaver,
+    Checkpoint,
+    CheckpointMetadata,
+    CheckpointTuple,
+    PendingWrite,
+    copy_checkpoint,
+    empty_checkpoint,
+)
+from langgraph.constants import INPUT, INTERRUPT
+from langgraph.managed.base import (
+    AsyncManagedValuesManager,
+    ManagedValueMapping,
+    ManagedValuesManager,
+)
+from langgraph.pregel.algo import (
+    PregelTaskWrites,
+    apply_writes,
+    increment,
+    prepare_next_tasks,
+    should_interrupt,
+)
+from langgraph.pregel.debug import map_debug_checkpoint, map_debug_tasks
+from langgraph.pregel.executor import (
+    AsyncBackgroundExecutor,
+    BackgroundExecutor,
+    Submit,
+)
+from langgraph.pregel.io import map_input, map_output_updates, map_output_values, single
+from langgraph.pregel.types import PregelExecutableTask
+
+if TYPE_CHECKING:
+    from langgraph.pregel import Pregel
+
+
+V = TypeVar("V")
+INPUT_DONE = object()
+
+
+class PregelLoop:
+    input: Optional[Any]
+    config: RunnableConfig
+    checkpointer: Optional[BaseCheckpointSaver]
+    checkpointer_get_next_version: Callable[[Optional[V]], V]
+    checkpointer_put_writes: Optional[
+        Callable[[RunnableConfig, Sequence[tuple[str, Any]], str], Any]
+    ]
+    checkpointer_put: Optional[
+        Callable[[RunnableConfig, Checkpoint, CheckpointMetadata], Any]
+    ]
+    graph: "Pregel"
+
+    submit: Submit
+    channels: Mapping[str, BaseChannel]
+    managed: ManagedValueMapping
+    checkpoint: Checkpoint
+    checkpoint_config: RunnableConfig
+    checkpoint_metadata: CheckpointMetadata
+    checkpoint_pending_writes: Optional[List[PendingWrite]]
+
+    step: int
+    status: Literal[
+        "pending", "done", "interrupt_before", "interrupt_after", "out_of_steps"
+    ]
+    tasks: Sequence[PregelExecutableTask]
+    stream: deque[Tuple[str, Any]]
+
+    # public
+
+    def mark_tasks_scheduled(self, tasks: Sequence[PregelExecutableTask]) -> None:
+        """Mark tasks as scheduled, to be used by queue-based executors."""
+        raise NotImplementedError
+
+    def put_writes(self, task_id: str, writes: Sequence[tuple[str, Any]]) -> None:
+        """Put writes for a task, to be read by the next tick."""
+        self.checkpoint_pending_writes.extend((task_id, k, v) for k, v in writes)
+        if self.checkpointer_put_writes is not None:
+            self.submit(
+                self.checkpointer_put_writes,
+                {
+                    **self.checkpoint_config,
+                    "configurable": {
+                        **self.checkpoint_config["configurable"],
+                        "thread_ts": self.checkpoint["id"],
+                    },
+                },
+                writes,
+                task_id,
+            )
+
+    def tick(
+        self,
+        *,
+        output_keys: Union[str, Sequence[str]] = None,
+        interrupt_after: Optional[Sequence[str]] = None,
+        interrupt_before: Optional[Sequence[str]] = None,
+        manager: Union[None, AsyncParentRunManager, ParentRunManager] = None,
+    ) -> bool:
+        """Execute a single iteration of the Pregel loop.
+        Returns True if more iterations are needed."""
+
+        if self.status != "pending":
+            raise RuntimeError("Cannot tick when status is no longer 'pending'")
+
+        if self.input is not INPUT_DONE:
+            self._first()
+        elif all(task.writes for task in self.tasks):
+            writes = [w for t in self.tasks for w in t.writes]
+            # all tasks have finished
+            apply_writes(
+                self.checkpoint,
+                self.channels,
+                self.tasks,
+                self.checkpointer_get_next_version,
+            )
+            # produce values output
+            self.stream.extend(
+                ("values", v)
+                for v in map_output_values(output_keys, writes, self.channels)
+            )
+            # clear pending writes
+            self.checkpoint_pending_writes.clear()
+            # save checkpoint
+            self._put_checkpoint(
+                {
+                    "source": "loop",
+                    "writes": single(
+                        map_output_updates(output_keys, self.tasks)
+                        if self.graph.stream_mode == "updates"
+                        else map_output_values(output_keys, writes, self.channels)
+                    ),
+                }
+            )
+            # after execution, check if we should interrupt
+            if should_interrupt(self.checkpoint, interrupt_after, self.tasks):
+                self.status = "interrupt_after"
+                return False
+        else:
+            return False
+
+        # check if iteration limit is reached
+        if self.step > self.config["recursion_limit"]:
+            self.status = "out_of_steps"
+            return False
+
+        # prepare next tasks
+        self.tasks = prepare_next_tasks(
+            self.checkpoint,
+            self.graph.nodes,
+            self.channels,
+            self.managed,
+            self.config,
+            self.step,
+            for_execution=True,
+            manager=manager,
+        )
+
+        # if no more tasks, we're done
+        if not self.tasks:
+            self.status = "done"
+            return False
+
+        # if there are pending writes from a previous loop, apply them
+        if self.checkpoint_pending_writes:
+            for tid, k, v in self.checkpoint_pending_writes:
+                if task := next((t for t in self.tasks if t.id == tid), None):
+                    task.writes.append((k, v))
+
+        # if all tasks have finished, re-tick
+        if all(task.writes for task in self.tasks):
+            return self.tick()
+
+        # before execution, check if we should interrupt
+        if should_interrupt(self.checkpoint, interrupt_before, self.tasks):
+            self.status = "interrupt_before"
+            return False
+
+        # produce debug output
+        self.stream.extend(("debug", v) for v in map_debug_tasks(self.step, self.tasks))
+
+        return True
+
+    # private
+
+    def _first(self) -> None:
+        # map inputs to channel updates
+        if input_writes := deque(map_input(self.graph.input_channels, self.input)):
+            # discard any unfinished tasks from previous checkpoint
+            discard_tasks = prepare_next_tasks(
+                self.checkpoint,
+                self.graph.nodes,
+                self.channels,
+                self.managed,
+                self.config,
+                self.step,
+                for_execution=True,
+            )
+            # apply input writes
+            apply_writes(
+                self.checkpoint,
+                self.channels,
+                discard_tasks + [PregelTaskWrites(INPUT, input_writes, [])],
+                self.checkpointer_get_next_version,
+            )
+            # save input checkpoint
+            self._put_checkpoint({"source": "input", "writes": self.input})
+        else:
+            # no input is taken as signal to proceed past previous interrupt
+            self.checkpoint["versions_seen"].setdefault(INTERRUPT, {})
+            for k in self.channels:
+                if k in self.checkpoint["channel_versions"]:
+                    version = self.checkpoint["channel_versions"][k]
+                    self.checkpoint["versions_seen"][INTERRUPT][k] = version
+        # done with input
+        self.input = INPUT_DONE
+
+    def _put_checkpoint(
+        self,
+        metadata: CheckpointMetadata,
+    ) -> None:
+        # assign step
+        metadata["step"] = self.step
+        # bail if no checkpointer
+        if self.checkpointer_put is not None:
+            # create new checkpoint
+            self.checkpoint_metadata = metadata
+            self.checkpoint = create_checkpoint(
+                self.checkpoint, self.channels, self.step
+            )
+            # save it, without blocking
+            self.submit(
+                self.checkpointer_put,
+                self.checkpoint_config,
+                copy_checkpoint(self.checkpoint),
+                self.checkpoint_metadata,
+            )
+            self.checkpoint_config = {
+                **self.checkpoint_config,
+                "configurable": {
+                    **self.checkpoint_config["configurable"],
+                    "thread_ts": self.checkpoint["id"],
+                },
+            }
+            # produce debug output
+            self.stream.extend(
+                ("debug", v)
+                for v in map_debug_checkpoint(
+                    self.step,
+                    self.checkpoint_config,
+                    self.channels,
+                    self.graph.stream_channels_asis,
+                    self.checkpoint_metadata,
+                )
+            )
+        # increment step
+        self.step += 1
+
+
+class SyncPregelLoop(PregelLoop, ContextManager):
+    def __init__(
+        self,
+        input: Optional[Any],
+        *,
+        config: RunnableConfig,
+        checkpointer: Optional[BaseCheckpointSaver],
+        graph: "Pregel",
+    ) -> None:
+        self.stream = deque()
+        self.stack = ExitStack()
+        self.input = input
+        self.config = config
+        self.checkpointer = checkpointer
+        self.checkpointer_get_next_version = (
+            checkpointer.get_next_version if checkpointer else increment
+        )
+        self.checkpointer_put_writes = checkpointer.put_writes if checkpointer else None
+        self.checkpointer_put = checkpointer.put if checkpointer else None
+        self.graph = graph
+        # TODO if managed values no longer needs graph we can replace with
+        # managed_specs, channel_specs
+
+    # context manager
+
+    def __enter__(self) -> Self:
+        saved = (
+            self.checkpointer.get_tuple(self.config) if self.checkpointer else None
+        ) or CheckpointTuple(self.config, empty_checkpoint(), {"step": -2}, None, [])
+        self.checkpoint_config = {
+            **self.config,
+            **saved.config,
+            "configurable": {
+                **self.config.get("configurable", {}),
+                **saved.config.get("configurable", {}),
+            },
+        }
+        self.checkpoint = copy_checkpoint(saved.checkpoint)
+        self.checkpoint_metadata = saved.metadata
+        self.checkpoint_pending_writes = saved.pending_writes
+
+        self.submit = self.stack.enter_context(BackgroundExecutor(self.config))
+        self.channels = self.stack.enter_context(
+            ChannelsManager(self.graph.channels, self.checkpoint, self.config)
+        )
+        self.managed = self.stack.enter_context(
+            ManagedValuesManager(
+                self.graph.managed_values_dict, self.config, self.graph
+            )
+        )
+        self.status = "pending"
+        self.step = self.checkpoint_metadata["step"] + 1
+
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> Optional[bool]:
+        del self.graph
+        return self.stack.__exit__(exc_type, exc_value, traceback)
+
+
+class AsyncPregelLoop(PregelLoop, AsyncContextManager):
+    def __init__(
+        self,
+        input: Optional[Any],
+        *,
+        config: RunnableConfig,
+        checkpointer: Optional[BaseCheckpointSaver],
+        graph: "Pregel",
+    ) -> None:
+        self.stream = deque()
+        self.stack = AsyncExitStack()
+        self.input = input
+        self.config = config
+        self.checkpointer = checkpointer
+        self.checkpointer_get_next_version = (
+            checkpointer.get_next_version if checkpointer else increment
+        )
+        self.checkpointer_put_writes = (
+            checkpointer.aput_writes if checkpointer else None
+        )
+        self.checkpointer_put = checkpointer.aput if checkpointer else None
+        self.graph = graph
+        # TODO if managed values no longer needs graph we can replace with
+        # managed_specs, channel_specs
+
+    # context manager
+
+    async def __aenter__(self) -> Self:
+        saved = (
+            await self.checkpointer.aget_tuple(self.config)
+            if self.checkpointer
+            else None
+        ) or CheckpointTuple(self.config, empty_checkpoint(), {"step": -2}, None, [])
+        self.checkpoint_config = {
+            **self.config,
+            **saved.config,
+            "configurable": {
+                **self.config.get("configurable", {}),
+                **saved.config.get("configurable", {}),
+            },
+        }
+        self.checkpoint = copy_checkpoint(saved.checkpoint)
+        self.checkpoint_metadata = saved.metadata
+        self.checkpoint_pending_writes = saved.pending_writes
+
+        self.submit = await self.stack.enter_async_context(AsyncBackgroundExecutor())
+        self.channels = await self.stack.enter_async_context(
+            AsyncChannelsManager(self.graph.channels, self.checkpoint, self.config)
+        )
+        self.managed = await self.stack.enter_async_context(
+            AsyncManagedValuesManager(
+                self.graph.managed_values_dict, self.config, self.graph
+            )
+        )
+        self.status = "pending"
+        self.step = self.checkpoint_metadata["step"] + 1
+
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> Optional[bool]:
+        del self.graph
+        return await asyncio.shield(
+            self.stack.__aexit__(exc_type, exc_value, traceback)
+        )

libs/langgraph/langgraph/pregel/read.py

@@ -16,6 +16,7 @@ from langchain_core.runnables.utils import ConfigurableFieldSpec
 
 from langgraph.constants import CONFIG_KEY_READ
 from langgraph.managed.base import ManagedValueSpec
+from langgraph.pregel.retry import RetryPolicy
 from langgraph.pregel.write import ChannelWrite
 from langgraph.utils import RunnableCallable
 
@@ -112,6 +113,8 @@ class PregelNode(RunnableBindingBase):
 
     kwargs: Mapping[str, Any] = Field(default_factory=dict)
 
+    retry_policy: Optional[RetryPolicy] = None
+
     def get_writers(self) -> list[Runnable]:
         """Get writers with optimizations applied."""
         writers = self.writers.copy()
@@ -151,9 +154,11 @@ class PregelNode(RunnableBindingBase):
         mapper: Optional[Callable[[Any], Any]] = None,
         writers: Optional[list[Runnable]] = None,
         tags: Optional[list[str]] = None,
+        metadata: Optional[Mapping[str, Any]] = None,
         bound: Optional[Runnable[Any, Any]] = None,
         kwargs: Optional[Mapping[str, Any]] = None,
         config: Optional[RunnableConfig] = None,
+        retry_policy: Optional[RetryPolicy] = None,
         **other_kwargs: Any,
     ) -> None:
         super().__init__(
@@ -163,7 +168,10 @@ class PregelNode(RunnableBindingBase):
             writers=writers or [],
             bound=bound or DEFAULT_BOUND,
             kwargs=kwargs or {},
-            config=merge_configs(config, {"tags": tags or []}),
+            retry_policy=retry_policy,
+            config=merge_configs(
+                config, {"tags": tags or [], "metadata": metadata or {}}
+            ),
             **other_kwargs,
         )
 
@@ -177,17 +185,13 @@ class PregelNode(RunnableBindingBase):
         assert isinstance(
             self.channels, dict
         ), "all channels must be named when using .join()"
-        return PregelNode(
-            channels={
-                **self.channels,
-                **{chan: chan for chan in channels},
-            },
-            triggers=self.triggers,
-            mapper=self.mapper,
-            writers=self.writers,
-            bound=self.bound,
-            kwargs=self.kwargs,
-            config=self.config,
+        return self.copy(
+            update=dict(
+                channels={
+                    **self.channels,
+                    **{chan: chan for chan in channels},
+                }
+            ),
         )
 
     def __or__(
@@ -199,36 +203,11 @@ class PregelNode(RunnableBindingBase):
         ],
     ) -> PregelNode:
         if ChannelWrite.is_writer(other):
-            return PregelNode(
-                channels=self.channels,
-                triggers=self.triggers,
-                mapper=self.mapper,
-                writers=[*self.writers, other],
-                bound=self.bound,
-                kwargs=self.kwargs,
-                config=self.config,
-            )
+            return self.copy(update=dict(writers=[*self.writers, other]))
         elif self.bound is DEFAULT_BOUND:
-            return PregelNode(
-                channels=self.channels,
-                triggers=self.triggers,
-                mapper=self.mapper,
-                writers=self.writers,
-                bound=coerce_to_runnable(other),
-                kwargs=self.kwargs,
-                config=self.config,
-            )
+            return self.copy(update=dict(bound=coerce_to_runnable(other)))
         else:
-            return PregelNode(
-                channels=self.channels,
-                triggers=self.triggers,
-                mapper=self.mapper,
-                writers=self.writers,
-                # delegate to __or__ in self.bound
-                bound=self.bound | other,
-                kwargs=self.kwargs,
-                config=self.config,
-            )
+            return self.copy(update=dict(bound=self.bound | other))
 
     def pipe(
         self,

libs/langgraph/langgraph/pregel/retry.py

@@ -2,66 +2,19 @@ import asyncio
 import logging
 import random
 import time
-from typing import Callable, NamedTuple, Optional, Union
+from typing import Optional
 
-import httpx
-import requests
-
-from langgraph.pregel.types import PregelExecutableTask
+from langgraph.pregel.types import PregelExecutableTask, RetryPolicy
 
 logger = logging.getLogger(__name__)
 
 
-def default_retry_on(exc: Exception) -> bool:
-    if isinstance(exc, ConnectionError):
-        return True
-    if isinstance(
-        exc,
-        (
-            ValueError,
-            TypeError,
-            ArithmeticError,
-            ImportError,
-            LookupError,
-            NameError,
-            SyntaxError,
-            RuntimeError,
-            ReferenceError,
-            StopIteration,
-            StopAsyncIteration,
-            OSError,
-        ),
-    ):
-        return False
-    if isinstance(exc, httpx.HTTPStatusError):
-        return 500 <= exc.response.status_code < 600
-    if isinstance(exc, requests.HTTPError):
-        return 500 <= exc.response.status_code < 600 if exc.response else True
-    return True
-
-
-class RetryPolicy(NamedTuple):
-    initial_interval: float = 0.5
-    """Amount of time that must elapse before the first retry occurs. In seconds."""
-    backoff_factor: float = 2.0
-    """Multiplier by which the interval increases after each retry."""
-    max_interval: float = 128.0
-    """Maximum amount of time that may elapse between retries. In seconds."""
-    max_attempts: int = 10
-    """Maximum number of attempts to make before giving up, including the first."""
-    jitter: bool = True
-    """Whether to add random jitter to the interval between retries."""
-    retry_on: Union[
-        tuple[Exception, ...], Callable[[Exception], bool]
-    ] = default_retry_on
-    """List of exceptions that should trigger a retry, or a callable that returns True for exceptions that should trigger a retry."""
-
-
 def run_with_retry(
     task: PregelExecutableTask,
     retry_policy: Optional[RetryPolicy],
 ) -> None:
     """Run a task with retries."""
+    retry_policy = task.retry_policy or retry_policy
     interval = retry_policy.initial_interval if retry_policy else 0
     attempts = 0
     while True:
@@ -106,6 +59,7 @@ async def arun_with_retry(
     stream: bool = False,
 ) -> None:
     """Run a task asynchronously with retries."""
+    retry_policy = task.retry_policy or retry_policy
     interval = retry_policy.initial_interval if retry_policy else 0
     attempts = 0
     while True:

libs/langgraph/langgraph/pregel/types.py

@@ -1,11 +1,61 @@
 from collections import deque
-from typing import Any, Literal, NamedTuple, Optional, Union
+from typing import Any, Callable, Literal, NamedTuple, Optional, Type, Union
 
 from langchain_core.runnables import Runnable, RunnableConfig
 
 from langgraph.checkpoint.base import CheckpointMetadata
 
 
+def default_retry_on(exc: Exception) -> bool:
+    import httpx
+    import requests
+
+    if isinstance(exc, ConnectionError):
+        return True
+    if isinstance(
+        exc,
+        (
+            ValueError,
+            TypeError,
+            ArithmeticError,
+            ImportError,
+            LookupError,
+            NameError,
+            SyntaxError,
+            RuntimeError,
+            ReferenceError,
+            StopIteration,
+            StopAsyncIteration,
+            OSError,
+        ),
+    ):
+        return False
+    if isinstance(exc, httpx.HTTPStatusError):
+        return 500 <= exc.response.status_code < 600
+    if isinstance(exc, requests.HTTPError):
+        return 500 <= exc.response.status_code < 600 if exc.response else True
+    return True
+
+
+class RetryPolicy(NamedTuple):
+    """Configuration for retrying nodes."""
+
+    initial_interval: float = 0.5
+    """Amount of time that must elapse before the first retry occurs. In seconds."""
+    backoff_factor: float = 2.0
+    """Multiplier by which the interval increases after each retry."""
+    max_interval: float = 128.0
+    """Maximum amount of time that may elapse between retries. In seconds."""
+    max_attempts: int = 3
+    """Maximum number of attempts to make before giving up, including the first."""
+    jitter: bool = True
+    """Whether to add random jitter to the interval between retries."""
+    retry_on: Union[
+        Type[Exception], tuple[Type[Exception], ...], Callable[[Exception], bool]
+    ] = default_retry_on
+    """List of exception classes that should trigger a retry, or a callable that returns True for exceptions that should trigger a retry."""
+
+
 class PregelTaskDescription(NamedTuple):
     name: str
     input: Any
@@ -18,6 +68,8 @@ class PregelExecutableTask(NamedTuple):
     writes: deque[tuple[str, Any]]
     config: RunnableConfig
     triggers: list[str]
+    retry_policy: Optional[RetryPolicy]
+    id: str
 
 
 class StateSnapshot(NamedTuple):
@@ -36,3 +88,12 @@ class StateSnapshot(NamedTuple):
 
 
 All = Literal["*"]
+
+StreamMode = Literal["values", "updates", "debug"]
+"""How the stream method should emit outputs.
+
+- 'values': Emit all values of the state for each step.
+- 'updates': Emit only the node name(s) and updates
+    that were returned by the node(s) **after** each step.
+- 'debug': Emit debug events for each step.
+"""

libs/langgraph/langgraph/utils.py

@@ -18,10 +18,17 @@ from langchain_core.runnables.config import (
     run_in_executor,
     var_child_runnable_config,
 )
-from langchain_core.runnables.graph import Edge, Graph, Node, is_uuid
 from langchain_core.runnables.utils import accepts_config
 from typing_extensions import TypeGuard
 
+try:
+    from langchain_core.runnables.config import _set_config_context
+except ImportError:
+    # For forwards compatibility
+    def _set_config_context(context: RunnableConfig) -> None:  # type: ignore
+        """Set the context for the current thread."""
+        var_child_runnable_config.set(context)
+
 
 # Before Python 3.11 native StrEnum is not available
 class StrEnum(str, enum.Enum):
@@ -89,7 +96,7 @@ class RunnableCallable(Runnable):
         else:
             config = merge_configs(self.config, config)
             context = copy_context()
-            context.run(var_child_runnable_config.set, config)
+            context.run(_set_config_context, config)
             if accepts_config(self.func):
                 kwargs["config"] = config
             ret = context.run(self.func, input, **kwargs)
@@ -110,7 +117,7 @@ class RunnableCallable(Runnable):
         else:
             config = merge_configs(self.config, config)
             context = copy_context()
-            context.run(var_child_runnable_config.set, config)
+            context.run(_set_config_context, config)
             if accepts_config(self.afunc):
                 kwargs["config"] = config
             if sys.version_info >= (3, 11):
@@ -124,37 +131,6 @@ class RunnableCallable(Runnable):
         return ret
 
 
-class DrawableGraph(Graph):
-    def extend(
-        self, graph: Graph, prefix: str = ""
-    ) -> tuple[Optional[Node], Optional[Node]]:
-        if all(is_uuid(node.id) for node in graph.nodes.values()):
-            super().extend(graph)
-            return graph.first_node(), graph.last_node()
-
-        new_nodes = {
-            f"{prefix}:{k}": Node(f"{prefix}:{k}", v.data)
-            for k, v in graph.nodes.items()
-        }
-        new_edges = [
-            Edge(
-                f"{prefix}:{edge.source}",
-                f"{prefix}:{edge.target}",
-                edge.data,
-                edge.conditional,
-            )
-            for edge in graph.edges
-        ]
-        self.nodes.update(new_nodes)
-        self.edges.extend(new_edges)
-        first = graph.first_node()
-        last = graph.last_node()
-        return (
-            Node(f"{prefix}:{first.id}", first.data) if first else None,
-            Node(f"{prefix}:{last.id}", last.data) if last else None,
-        )
-
-
 def is_async_callable(
     func: Any,
 ) -> TypeGuard[Callable[..., Awaitable]]:

libs/langgraph/poetry.lock

@@ -1,4 +1,4 @@
-# This file is automatically @generated by Poetry 1.8.3 and should not be changed by hand.
+# This file is automatically @generated by Poetry 1.8.2 and should not be changed by hand.
 
 [[package]]
 name = "aiohttp"
@@ -387,13 +387,13 @@ css = ["tinycss2 (>=1.1.0,<1.3)"]
 
 [[package]]
 name = "certifi"
-version = "2024.6.2"
+version = "2024.7.4"
 description = "Python package for providing Mozilla's CA Bundle."
 optional = false
 python-versions = ">=3.6"
 files = [
-    {file = "certifi-2024.6.2-py3-none-any.whl", hash = "sha256:ddc6c8ce995e6987e7faf5e3f1b02b302836a0e5d98ece18392cb1a36c72ad56"},
-    {file = "certifi-2024.6.2.tar.gz", hash = "sha256:3cd43f1c6fa7dedc5899d69d3ad0398fd018ad1a17fba83ddaf78aa46c747516"},
+    {file = "certifi-2024.7.4-py3-none-any.whl", hash = "sha256:c198e21b1289c2ab85ee4e67bb4b4ef3ead0892059901a8d5b622f24a1101e90"},
+    {file = "certifi-2024.7.4.tar.gz", hash = "sha256:5a1e7645bc0ec61a09e26c36f6106dd4cf40c6db3a1fb6352b0244e7fb057c7b"},
 ]
 
 [[package]]
@@ -747,6 +747,20 @@ files = [
 [package.extras]
 test = ["pytest (>=6)"]
 
+[[package]]
+name = "execnet"
+version = "2.1.1"
+description = "execnet: rapid multi-Python deployment"
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "execnet-2.1.1-py3-none-any.whl", hash = "sha256:26dee51f1b80cebd6d0ca8e74dd8745419761d3bef34163928cbebbdc4749fdc"},
+    {file = "execnet-2.1.1.tar.gz", hash = "sha256:5189b52c6121c24feae288166ab41b32549c7e2348652736540b9e6e7d4e72e3"},
+]
+
+[package.extras]
+testing = ["hatch", "pre-commit", "pytest", "tox"]
+
 [[package]]
 name = "executing"
 version = "2.0.1"
@@ -1746,13 +1760,13 @@ langchain-core = ">=0.2.2rc1,<0.3"
 
 [[package]]
 name = "langchain-core"
-version = "0.2.11"
+version = "0.2.22"
 description = "Building applications with LLMs through composability"
 optional = false
 python-versions = "<4.0,>=3.8.1"
 files = [
-    {file = "langchain_core-0.2.11-py3-none-any.whl", hash = "sha256:c7ca4dc4d88e3c69fd7916c95a7027c2b1a11c2db5a51141c3ceb8afac212208"},
-    {file = "langchain_core-0.2.11.tar.gz", hash = "sha256:7a4661b50604eeb20c3373fbfd8a4f1b74482a6ab4e0f9df11e96821ead8ef0c"},
+    {file = "langchain_core-0.2.22-py3-none-any.whl", hash = "sha256:7731a86440c0958b3186c003fb9b26b2d5a682a6344bda7bfb9174e2898f8b43"},
+    {file = "langchain_core-0.2.22.tar.gz", hash = "sha256:582d6f929a43b830139444e4124123cd415331ad62f25757b1406252958cdcac"},
 ]
 
 [package.dependencies]
@@ -2784,6 +2798,27 @@ files = [
 tomli = {version = ">=2.0.1,<3.0.0", markers = "python_version < \"3.11\""}
 watchdog = ">=2.0.0"
 
+[[package]]
+name = "pytest-xdist"
+version = "3.6.1"
+description = "pytest xdist plugin for distributed testing, most importantly across multiple CPUs"
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "pytest_xdist-3.6.1-py3-none-any.whl", hash = "sha256:9ed4adfb68a016610848639bb7e02c9352d5d9f03d04809919e2dafc3be4cca7"},
+    {file = "pytest_xdist-3.6.1.tar.gz", hash = "sha256:ead156a4db231eec769737f57668ef58a2084a34b2e55c4a8fa20d861107300d"},
+]
+
+[package.dependencies]
+execnet = ">=2.1"
+psutil = {version = ">=3.0", optional = true, markers = "extra == \"psutil\""}
+pytest = ">=7.0.0"
+
+[package.extras]
+psutil = ["psutil (>=3.0)"]
+setproctitle = ["setproctitle"]
+testing = ["filelock"]
+
 [[package]]
 name = "python-dateutil"
 version = "2.9.0.post0"
@@ -4130,4 +4165,4 @@ test = ["big-O", "importlib-resources", "jaraco.functools", "jaraco.itertools",
 [metadata]
 lock-version = "2.0"
 python-versions = ">=3.9.0,<4.0"
-content-hash = "1ec1c06495c9a1a564c5082600a4fb335278341a19d11993b981b72684679948"
+content-hash = "0d877d3879473de43aca1e1d36a8f420ff3f4b140807cb5cc24935d9114947be"

libs/langgraph/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph"
-version = "0.1.5"
+version = "0.1.9"
 description = "Building stateful, multi-actor applications with LLMs"
 authors = []
 license = "MIT"
@@ -9,7 +9,7 @@ repository = "https://www.github.com/langchain-ai/langgraph"
 
 [tool.poetry.dependencies]
 python = ">=3.9.0,<4.0"
-langchain-core = ">=0.2.11,<0.3"
+langchain-core = ">=0.2.22,<0.3"
 
 
 [tool.poetry.group.dev.dependencies]
@@ -31,6 +31,7 @@ langchainhub = "^0.1.14"
 langchain-openai = ">=0.1.2"
 langchain-anthropic = ">=0.1.8"
 dataclasses-json = "^0.6.7"
+pytest-xdist = {extras = ["psutil"], version = "^3.6.1"}
 
 [tool.poetry.group.dev]
 optional = true
@@ -61,7 +62,7 @@ omit = ["tests/*"]
 [tool.pytest-watcher]
 now = true
 delay = 0.1
-runner_args = ["-x", "--ff", "-vv", "--snapshot-update"]
+runner_args = ["-x", "--ff", "-v", "-n", "auto", "--dist", "worksteal", "--snapshot-update", "--tb", "short"]
 patterns = ["*.py"]
 
 [build-system]

libs/langgraph/tests/test_interruption.py

@@ -0,0 +1,76 @@
+from typing import TypedDict
+
+from langgraph.checkpoint.memory import MemorySaver
+from langgraph.graph import END, START, StateGraph
+
+
+def test_interruption_without_state_updates():
+    """Test interruption without state updates. This test confirms that
+    interrupting doesn't require a state key having been updated in the prev step"""
+
+    class State(TypedDict):
+        input: str
+
+    def noop(_state):
+        pass
+
+    builder = StateGraph(State)
+    builder.add_node("step_1", noop)
+    builder.add_node("step_2", noop)
+    builder.add_node("step_3", noop)
+    builder.add_edge(START, "step_1")
+    builder.add_edge("step_1", "step_2")
+    builder.add_edge("step_2", "step_3")
+    builder.add_edge("step_3", END)
+
+    memory = MemorySaver()
+
+    graph = builder.compile(checkpointer=memory, interrupt_after="*")
+
+    initial_input = {"input": "hello world"}
+    thread = {"configurable": {"thread_id": "1"}}
+
+    graph.invoke(initial_input, thread, debug=True)
+    assert graph.get_state(thread).next == ("step_2",)
+
+    graph.invoke(None, thread, debug=True)
+    assert graph.get_state(thread).next == ("step_3",)
+
+    graph.invoke(None, thread, debug=True)
+    assert graph.get_state(thread).next == ()
+
+
+async def test_interruption_without_state_updates_async():
+    """Test interruption without state updates. This test confirms that
+    interrupting doesn't require a state key having been updated in the prev step"""
+
+    class State(TypedDict):
+        input: str
+
+    async def noop(_state):
+        pass
+
+    builder = StateGraph(State)
+    builder.add_node("step_1", noop)
+    builder.add_node("step_2", noop)
+    builder.add_node("step_3", noop)
+    builder.add_edge(START, "step_1")
+    builder.add_edge("step_1", "step_2")
+    builder.add_edge("step_2", "step_3")
+    builder.add_edge("step_3", END)
+
+    memory = MemorySaver()
+
+    graph = builder.compile(checkpointer=memory, interrupt_after="*")
+
+    initial_input = {"input": "hello world"}
+    thread = {"configurable": {"thread_id": "1"}}
+
+    await graph.ainvoke(initial_input, thread, debug=True)
+    assert (await graph.aget_state(thread)).next == ("step_2",)
+
+    await graph.ainvoke(None, thread, debug=True)
+    assert (await graph.aget_state(thread)).next == ("step_3",)
+
+    await graph.ainvoke(None, thread, debug=True)
+    assert (await graph.aget_state(thread)).next == ()

libs/langgraph/tests/test_messages_state.py

@@ -0,0 +1,137 @@
+from uuid import UUID
+
+import pytest
+from langchain_core.messages import (
+    AIMessage,
+    HumanMessage,
+    RemoveMessage,
+    SystemMessage,
+)
+
+from langgraph.graph import add_messages
+
+
+def test_add_single_message():
+    left = [HumanMessage(content="Hello", id="1")]
+    right = AIMessage(content="Hi there!", id="2")
+    result = add_messages(left, right)
+    expected_result = [
+        HumanMessage(content="Hello", id="1"),
+        AIMessage(content="Hi there!", id="2"),
+    ]
+    assert result == expected_result
+
+
+def test_add_multiple_messages():
+    left = [HumanMessage(content="Hello", id="1")]
+    right = [
+        AIMessage(content="Hi there!", id="2"),
+        SystemMessage(content="System message", id="3"),
+    ]
+    result = add_messages(left, right)
+    expected_result = [
+        HumanMessage(content="Hello", id="1"),
+        AIMessage(content="Hi there!", id="2"),
+        SystemMessage(content="System message", id="3"),
+    ]
+    assert result == expected_result
+
+
+def test_update_existing_message():
+    left = [HumanMessage(content="Hello", id="1")]
+    right = HumanMessage(content="Hello again", id="1")
+    result = add_messages(left, right)
+    expected_result = [HumanMessage(content="Hello again", id="1")]
+    assert result == expected_result
+
+
+def test_missing_ids():
+    left = [HumanMessage(content="Hello")]
+    right = [AIMessage(content="Hi there!")]
+    result = add_messages(left, right)
+    assert len(result) == 2
+    assert all(isinstance(m.id, str) and UUID(m.id, version=4) for m in result)
+
+
+def test_remove_message():
+    left = [
+        HumanMessage(content="Hello", id="1"),
+        AIMessage(content="Hi there!", id="2"),
+    ]
+    right = RemoveMessage(id="2")
+    result = add_messages(left, right)
+    expected_result = [HumanMessage(content="Hello", id="1")]
+    assert result == expected_result
+
+
+def test_duplicate_remove_message():
+    left = [
+        HumanMessage(content="Hello", id="1"),
+        AIMessage(content="Hi there!", id="2"),
+    ]
+    right = [RemoveMessage(id="2"), RemoveMessage(id="2")]
+    result = add_messages(left, right)
+    expected_result = [HumanMessage(content="Hello", id="1")]
+    assert result == expected_result
+
+
+def test_remove_nonexistent_message():
+    left = [HumanMessage(content="Hello", id="1")]
+    right = RemoveMessage(id="2")
+    with pytest.raises(
+        ValueError, match="Attempting to delete a message with an ID that doesn't exist"
+    ):
+        add_messages(left, right)
+
+
+def test_mixed_operations():
+    left = [
+        HumanMessage(content="Hello", id="1"),
+        AIMessage(content="Hi there!", id="2"),
+    ]
+    right = [
+        HumanMessage(content="Updated hello", id="1"),
+        RemoveMessage(id="2"),
+        SystemMessage(content="New message", id="3"),
+    ]
+    result = add_messages(left, right)
+    expected_result = [
+        HumanMessage(content="Updated hello", id="1"),
+        SystemMessage(content="New message", id="3"),
+    ]
+    assert result == expected_result
+
+
+def test_empty_inputs():
+    assert add_messages([], []) == []
+    assert add_messages([], [HumanMessage(content="Hello", id="1")]) == [
+        HumanMessage(content="Hello", id="1")
+    ]
+    assert add_messages([HumanMessage(content="Hello", id="1")], []) == [
+        HumanMessage(content="Hello", id="1")
+    ]
+
+
+def test_non_list_inputs():
+    left = HumanMessage(content="Hello", id="1")
+    right = AIMessage(content="Hi there!", id="2")
+    result = add_messages(left, right)
+    expected_result = [
+        HumanMessage(content="Hello", id="1"),
+        AIMessage(content="Hi there!", id="2"),
+    ]
+    assert result == expected_result
+
+
+def test_delete_all():
+    left = [
+        HumanMessage(content="Hello", id="1"),
+        AIMessage(content="Hi there!", id="2"),
+    ]
+    right = [
+        RemoveMessage(id="1"),
+        RemoveMessage(id="2"),
+    ]
+    result = add_messages(left, right)
+    expected_result = []
+    assert result == expected_result

libs/langgraph/tests/test_prebuilt.py

@@ -1,15 +1,11 @@
-from typing import Any, Callable, Dict, List, Optional, Sequence, Type, Union
+from typing import Annotated, Any, Callable, Dict, List, Optional, Sequence, Type, Union
 
 import pytest
-from langchain_core.callbacks import (
-    CallbackManagerForLLMRun,
-)
-from langchain_core.language_models import (
-    BaseChatModel,
-    LanguageModelInput,
-)
+from langchain_core.callbacks import CallbackManagerForLLMRun
+from langchain_core.language_models import BaseChatModel, LanguageModelInput
 from langchain_core.messages import (
     AIMessage,
+    AnyMessage,
     BaseMessage,
     HumanMessage,
     SystemMessage,
@@ -22,7 +18,11 @@ from langchain_core.tools import BaseTool
 from langchain_core.tools import tool as dec_tool
 from pydantic import BaseModel as BaseModelV2
 
+from langgraph.checkpoint.base import BaseCheckpointSaver
 from langgraph.prebuilt import ToolNode, ValidationNode, create_react_agent
+from langgraph.prebuilt.tool_node import InjectedState
+from tests.any_str import AnyStr
+from tests.memory_assert import MemorySaverAssertImmutable
 
 
 class FakeToolCallingModel(BaseChatModel):
@@ -52,57 +52,205 @@ class FakeToolCallingModel(BaseChatModel):
         return self
 
 
-def test_no_modifier():
+@pytest.mark.parametrize(
+    "checkpointer",
+    [
+        MemorySaverAssertImmutable(),
+        None,
+    ],
+    ids=[
+        "memory",
+        "none",
+    ],
+)
+def test_no_modifier(checkpointer: Optional[BaseCheckpointSaver]):
     model = FakeToolCallingModel()
-    agent = create_react_agent(model, [])
+    agent = create_react_agent(model, [], checkpointer=checkpointer)
     inputs = [HumanMessage("hi?")]
-    response = agent.invoke({"messages": inputs})
+    thread = {"configurable": {"thread_id": "123"}}
+    response = agent.invoke({"messages": inputs}, thread, debug=True)
     expected_response = {"messages": inputs + [AIMessage(content="hi?", id="0")]}
     assert response == expected_response
 
+    if checkpointer:
+        saved = checkpointer.get_tuple(thread)
+        assert saved is not None
+        assert saved.checkpoint == {
+            "v": 1,
+            "ts": AnyStr(),
+            "id": AnyStr(),
+            "channel_values": {
+                "messages": [
+                    HumanMessage(content="hi?", id=AnyStr()),
+                    AIMessage(content="hi?", id="0"),
+                ],
+                "agent": "agent",
+            },
+            "channel_versions": {
+                "__start__": 2,
+                "messages": 3,
+                "start:agent": 3,
+                "agent": 3,
+            },
+            "versions_seen": {
+                "__input__": {},
+                "__start__": {"__start__": 1},
+                "agent": {"start:agent": 2},
+            },
+            "pending_sends": [],
+            "current_tasks": {},
+        }
+        assert saved.metadata == {
+            "source": "loop",
+            "writes": {"agent": {"messages": [AIMessage(content="hi?", id="0")]}},
+            "step": 1,
+        }
+        assert saved.pending_writes == []
+
+
+@pytest.mark.parametrize(
+    "checkpointer",
+    [
+        MemorySaverAssertImmutable(),
+        None,
+    ],
+    ids=[
+        "memory",
+        "none",
+    ],
+)
+async def test_no_modifier_async(checkpointer: Optional[BaseCheckpointSaver]):
+    model = FakeToolCallingModel()
+    agent = create_react_agent(model, [], checkpointer=checkpointer)
+    inputs = [HumanMessage("hi?")]
+    thread = {"configurable": {"thread_id": "123"}}
+    response = await agent.ainvoke({"messages": inputs}, thread, debug=True)
+    expected_response = {"messages": inputs + [AIMessage(content="hi?", id="0")]}
+    assert response == expected_response
+
+    if checkpointer:
+        saved = await checkpointer.aget_tuple(thread)
+        assert saved is not None
+        assert saved.checkpoint == {
+            "v": 1,
+            "ts": AnyStr(),
+            "id": AnyStr(),
+            "channel_values": {
+                "messages": [
+                    HumanMessage(content="hi?", id=AnyStr()),
+                    AIMessage(content="hi?", id="0"),
+                ],
+                "agent": "agent",
+            },
+            "channel_versions": {
+                "__start__": 2,
+                "messages": 3,
+                "start:agent": 3,
+                "agent": 3,
+            },
+            "versions_seen": {
+                "__input__": {},
+                "__start__": {"__start__": 1},
+                "agent": {"start:agent": 2},
+            },
+            "pending_sends": [],
+            "current_tasks": {},
+        }
+        assert saved.metadata == {
+            "source": "loop",
+            "writes": {"agent": {"messages": [AIMessage(content="hi?", id="0")]}},
+            "step": 1,
+        }
+        assert saved.pending_writes == []
+
+
+def test_passing_two_modifiers():
+    model = FakeToolCallingModel()
+    with pytest.raises(ValueError):
+        create_react_agent(model, [], messages_modifier="Foo", state_modifier="Bar")
+
 
 def test_system_message_modifier():
     model = FakeToolCallingModel()
     messages_modifier = SystemMessage(content="Foo")
-    agent = create_react_agent(model, [], messages_modifier=messages_modifier)
-    inputs = [HumanMessage("hi?")]
-    response = agent.invoke({"messages": inputs})
-    expected_response = {"messages": inputs + [AIMessage(content="Foo-hi?", id="0")]}
-    assert response == expected_response
+    agent_1 = create_react_agent(model, [], messages_modifier=messages_modifier)
+    agent_2 = create_react_agent(model, [], state_modifier=messages_modifier)
+    for agent in [agent_1, agent_2]:
+        inputs = [HumanMessage("hi?")]
+        response = agent.invoke({"messages": inputs})
+        expected_response = {
+            "messages": inputs + [AIMessage(content="Foo-hi?", id="0")]
+        }
+        assert response == expected_response
 
 
 def test_system_message_string_modifier():
     model = FakeToolCallingModel()
     messages_modifier = "Foo"
-    agent = create_react_agent(model, [], messages_modifier=messages_modifier)
-    inputs = [HumanMessage("hi?")]
-    response = agent.invoke({"messages": inputs})
-    expected_response = {"messages": inputs + [AIMessage(content="Foo-hi?", id="0")]}
-    assert response == expected_response
+    agent_1 = create_react_agent(model, [], messages_modifier=messages_modifier)
+    agent_2 = create_react_agent(model, [], state_modifier=messages_modifier)
+    for agent in [agent_1, agent_2]:
+        inputs = [HumanMessage("hi?")]
+        response = agent.invoke({"messages": inputs})
+        expected_response = {
+            "messages": inputs + [AIMessage(content="Foo-hi?", id="0")]
+        }
+        assert response == expected_response
 
 
-def test_callable_modifier():
+def test_callable_messages_modifier():
     model = FakeToolCallingModel()
 
     def messages_modifier(messages):
-        return [HumanMessage(content="Bar")]
+        modified_message = f"Bar {messages[-1].content}"
+        return [HumanMessage(content=modified_message)]
 
     agent = create_react_agent(model, [], messages_modifier=messages_modifier)
     inputs = [HumanMessage("hi?")]
     response = agent.invoke({"messages": inputs})
-    expected_response = {"messages": inputs + [AIMessage(content="Bar", id="0")]}
+    expected_response = {"messages": inputs + [AIMessage(content="Bar hi?", id="0")]}
     assert response == expected_response
 
 
-def test_runnable_modifier():
+def test_callable_state_modifier():
     model = FakeToolCallingModel()
 
-    messages_modifier = RunnableLambda(lambda x: [HumanMessage(content="Baz")])
+    def state_modifier(state):
+        modified_message = f"Bar {state['messages'][-1].content}"
+        return [HumanMessage(content=modified_message)]
+
+    agent = create_react_agent(model, [], state_modifier=state_modifier)
+    inputs = [HumanMessage("hi?")]
+    response = agent.invoke({"messages": inputs})
+    expected_response = {"messages": inputs + [AIMessage(content="Bar hi?", id="0")]}
+    assert response == expected_response
+
+
+def test_runnable_messages_modifier():
+    model = FakeToolCallingModel()
+
+    messages_modifier = RunnableLambda(
+        lambda messages: [HumanMessage(content=f"Baz {messages[-1].content}")]
+    )
 
     agent = create_react_agent(model, [], messages_modifier=messages_modifier)
     inputs = [HumanMessage("hi?")]
     response = agent.invoke({"messages": inputs})
-    expected_response = {"messages": inputs + [AIMessage(content="Baz", id="0")]}
+    expected_response = {"messages": inputs + [AIMessage(content="Baz hi?", id="0")]}
+    assert response == expected_response
+
+
+def test_runnable_state_modifier():
+    model = FakeToolCallingModel()
+
+    state_modifier = RunnableLambda(
+        lambda state: [HumanMessage(content=f"Baz {state['messages'][-1].content}")]
+    )
+
+    agent = create_react_agent(model, [], state_modifier=state_modifier)
+    inputs = [HumanMessage("hi?")]
+    response = agent.invoke({"messages": inputs})
+    expected_response = {"messages": inputs + [AIMessage(content="Baz hi?", id="0")]}
     assert response == expected_response
 
 
@@ -204,6 +352,31 @@ async def test_tool_node():
             }
         )
 
+    # incorrect tool name
+    result_incorrect_name = ToolNode([tool1, tool2]).invoke(
+        {
+            "messages": [
+                AIMessage(
+                    "hi?",
+                    tool_calls=[
+                        {
+                            "name": "tool3",
+                            "args": {"some_val": 1, "some_other_val": "foo"},
+                            "id": "some 0",
+                        }
+                    ],
+                )
+            ]
+        }
+    )
+    tool_message: ToolMessage = result_incorrect_name["messages"][-1]
+    assert tool_message.type == "tool"
+    assert (
+        tool_message.content
+        == "Error: tool3 is not a valid tool, try one of [tool1, tool2]."
+    )
+    assert tool_message.tool_call_id == "some 0"
+
 
 def my_function(some_val: int, some_other_val: str) -> str:
     return f"{some_val} - {some_other_val}"
@@ -273,3 +446,69 @@ async def test_validation_node(tool_schema: Any, use_message_key: bool):
     if use_message_key:
         result_sync = result_sync["messages"]
     check_results(result_sync)
+
+
+def test_tool_node_inject_state() -> None:
+    def tool1(some_val: int, state: Annotated[dict, InjectedState]) -> str:
+        """Tool 1 docstring."""
+        return state["foo"]
+
+    def tool2(some_val: int, state: Annotated[dict, InjectedState()]) -> str:
+        """Tool 1 docstring."""
+        return state["foo"]
+
+    def tool3(
+        some_val: int,
+        foo: Annotated[str, InjectedState("foo")],
+        msgs: Annotated[List[AnyMessage], InjectedState("messages")],
+    ) -> str:
+        """Tool 1 docstring."""
+        return foo
+
+    def tool4(
+        some_val: int, msgs: Annotated[List[AnyMessage], InjectedState("messages")]
+    ) -> str:
+        """Tool 1 docstring."""
+        return msgs[0].content
+
+    node = ToolNode([tool1, tool2, tool3, tool4])
+    for tool_name in ("tool1", "tool2", "tool3"):
+        tool_call = {
+            "name": tool_name,
+            "args": {"some_val": 1},
+            "id": "some 0",
+            "type": "tool_call",
+        }
+        msg = AIMessage("hi?", tool_calls=[tool_call])
+        result = node.invoke({"messages": [msg], "foo": "bar"})
+        tool_message = result["messages"][-1]
+        assert tool_message.content == "bar"
+
+        if tool_name == "tool3":
+            with pytest.raises(KeyError):
+                node.invoke({"messages": [msg], "notfoo": "bar"})
+
+            with pytest.raises(ValueError):
+                node.invoke([msg])
+        else:
+            tool_message = node.invoke({"messages": [msg], "notfoo": "bar"})[
+                "messages"
+            ][-1]
+            assert "KeyError" in tool_message.content
+            tool_message = node.invoke([msg])[-1]
+            assert "KeyError" in tool_message.content
+
+    tool_call = {
+        "name": "tool4",
+        "args": {"some_val": 1},
+        "id": "some 0",
+        "type": "tool_call",
+    }
+    msg = AIMessage("hi?", tool_calls=[tool_call])
+    result = node.invoke({"messages": [msg]})
+    tool_message = result["messages"][-1]
+    assert tool_message.content == "hi?"
+
+    result = node.invoke([msg])
+    tool_message = result[-1]
+    assert tool_message.content == "hi?"

libs/langgraph/tests/test_pregel.py

@@ -14,6 +14,7 @@ from typing import (
     Literal,
     Optional,
     Sequence,
+    Tuple,
     TypedDict,
     Union,
 )
@@ -37,6 +38,7 @@ from langgraph.channels.context import Context
 from langgraph.channels.last_value import LastValue
 from langgraph.channels.topic import Topic
 from langgraph.checkpoint.base import (
+    BaseCheckpointSaver,
     Checkpoint,
     CheckpointMetadata,
     CheckpointTuple,
@@ -193,6 +195,12 @@ def test_checkpoint_errors() -> None:
         ) -> RunnableConfig:
             raise ValueError("Faulty put")
 
+    class FaultyPutWritesCheckpointer(MemorySaver):
+        def put_writes(
+            self, config: RunnableConfig, writes: List[Tuple[str, Any]], task_id: str
+        ) -> RunnableConfig:
+            raise ValueError("Faulty put_writes")
+
     class FaultyVersionCheckpointer(MemorySaver):
         def get_next_version(self, current: Optional[int], channel: BaseChannel) -> int:
             raise ValueError("Faulty get_next_version")
@@ -200,10 +208,9 @@ def test_checkpoint_errors() -> None:
     def logic(inp: str) -> str:
         return ""
 
-    builder = Graph()
+    builder = StateGraph(Annotated[str, operator.add])
     builder.add_node("agent", logic)
-    builder.set_entry_point("agent")
-    builder.set_finish_point("agent")
+    builder.add_edge(START, "agent")
 
     graph = builder.compile(checkpointer=FaultyGetCheckpointer())
     with pytest.raises(ValueError, match="Faulty get_tuple"):
@@ -217,6 +224,92 @@ def test_checkpoint_errors() -> None:
     with pytest.raises(ValueError, match="Faulty get_next_version"):
         graph.invoke("", {"configurable": {"thread_id": "thread-1"}})
 
+    # add parallel node
+    builder.add_node("parallel", logic)
+    builder.add_edge(START, "parallel")
+    graph = builder.compile(checkpointer=FaultyPutWritesCheckpointer())
+    with pytest.raises(ValueError, match="Faulty put_writes"):
+        graph.invoke("", {"configurable": {"thread_id": "thread-1"}})
+
+
+def test_node_schemas_custom_output() -> None:
+    from langchain_core.messages import HumanMessage
+
+    class State(TypedDict):
+        hello: str
+        bye: str
+        messages: Annotated[list[str], add_messages]
+
+    class Output(TypedDict):
+        messages: list[str]
+
+    class StateForA(TypedDict):
+        hello: str
+        messages: Annotated[list[str], add_messages]
+
+    def node_a(state: StateForA) -> State:
+        assert state == {
+            "hello": "there",
+            "messages": [HumanMessage(content="hello", id=AnyStr())],
+        }
+
+    class StateForB(TypedDict):
+        bye: str
+        now: int
+
+    def node_b(state: StateForB):
+        assert state == {
+            "bye": "world",
+            "now": None,
+        }
+        return {
+            "now": 123,
+            "hello": "again",
+        }
+
+    class StateForC(TypedDict):
+        hello: str
+        now: int
+
+    def node_c(state: StateForC) -> StateForC:
+        assert state == {
+            "hello": "again",
+            "now": 123,
+        }
+
+    builder = StateGraph(State, output=Output)
+    builder.add_node("a", node_a)
+    builder.add_node("b", node_b)
+    builder.add_node("c", node_c)
+    builder.add_edge(START, "a")
+    builder.add_edge("a", "b")
+    builder.add_edge("b", "c")
+    graph = builder.compile()
+
+    assert graph.invoke({"hello": "there", "bye": "world", "messages": "hello"}) == {
+        "messages": [HumanMessage(content="hello", id=AnyStr())],
+    }
+
+    builder = StateGraph(input=State, output=Output)
+    builder.add_node("a", node_a)
+    builder.add_node("b", node_b)
+    builder.add_node("c", node_c)
+    builder.add_edge(START, "a")
+    builder.add_edge("a", "b")
+    builder.add_edge("b", "c")
+    graph = builder.compile()
+
+    assert graph.invoke(
+        {
+            "hello": "there",
+            "bye": "world",
+            "messages": "hello",
+            "now": 345,  # ignored because not in input schema
+        }
+    ) == {
+        "messages": [HumanMessage(content="hello", id=AnyStr())],
+    }
+
 
 def test_reducer_before_first_node() -> None:
     from langchain_core.messages import HumanMessage
@@ -427,10 +520,8 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
 
     assert app.invoke(2) == 4
 
-    assert app.invoke(2, input_keys="inbox") == 3
-
     with pytest.raises(GraphRecursionError):
-        app.invoke(2, {"recursion_limit": 1})
+        app.invoke(2, {"recursion_limit": 1}, debug=1)
 
     graph = Graph()
     graph.add_node("add_one", add_one)
@@ -442,7 +533,7 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
 
     assert gapp.invoke(2) == 4
 
-    for step, values in enumerate(gapp.stream(2), start=1):
+    for step, values in enumerate(gapp.stream(2, debug=1), start=1):
         if step == 1:
             assert values == {
                 "add_one": 3,
@@ -944,6 +1035,122 @@ def test_invoke_checkpoint(mocker: MockerFixture) -> None:
     assert checkpoint["channel_values"].get("total") == 5
 
 
+@pytest.mark.parametrize(
+    "checkpointer",
+    [
+        MemorySaverAssertImmutable(),
+        SqliteSaver.from_conn_string(":memory:"),
+    ],
+    ids=[
+        "memory",
+        "sqlite",
+    ],
+)
+def test_pending_writes_resume(checkpointer: BaseCheckpointSaver) -> None:
+    try:
+
+        class State(TypedDict):
+            value: Annotated[int, operator.add]
+
+        class AwhileMaker:
+            def __init__(self, sleep: float, rtn: Union[Dict, Exception]) -> None:
+                self.sleep = sleep
+                self.rtn = rtn
+                self.reset()
+
+            def __call__(self, input: State) -> Any:
+                self.calls += 1
+                time.sleep(self.sleep)
+                if isinstance(self.rtn, Exception):
+                    raise self.rtn
+                else:
+                    return self.rtn
+
+            def reset(self):
+                self.calls = 0
+
+        one = AwhileMaker(0.2, {"value": 2})
+        two = AwhileMaker(0.6, 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_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)
+
+        # both nodes should have been called once
+        assert one.calls == 1
+        assert two.calls == 2  # two attempts
+
+        # latest checkpoint should be before nodes "one", "two"
+        state = graph.get_state(thread1)
+        assert state is not None
+        assert state.values == {"value": 1}
+        assert state.next == ("one", "two")
+        assert state.metadata == {"source": "loop", "step": 0, "writes": None}
+        # should contain pending write of "one"
+        checkpoint = checkpointer.get_tuple(thread1)
+        assert checkpoint is not None
+        assert checkpoint.pending_writes == [
+            (AnyStr(), "one", "one"),
+            (AnyStr(), "value", 2),
+        ]
+        # both pending writes come from same task
+        assert checkpoint.pending_writes[0][0] == checkpoint.pending_writes[1][0]
+
+        # resume execution
+        with pytest.raises(ConnectionError, match="I'm not good"):
+            graph.invoke(None, thread1)
+
+        # node "one" succeeded previously, so shouldn't be called again
+        assert one.calls == 1
+        # node "two" should have been called once again
+        assert two.calls == 4  # two attempts before + two attempts now
+
+        # confirm no new checkpoints saved
+        state_two = graph.get_state(thread1)
+        assert state_two == state
+
+        # 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}
+    finally:
+        if getattr(checkpointer, "__exit__", None):
+            checkpointer.__exit__(None, None, None)
+
+
+def test_cond_edge_after_send() -> None:
+    class Node:
+        def __init__(self, name: str):
+            self.name = name
+            setattr(self, "__name__", name)
+
+        def __call__(self, state):
+            return state + [self.name]
+
+    def send_for_fun(state):
+        return [Send("2", state)]
+
+    def route_to_three(state) -> Literal["3"]:
+        return "3"
+
+    builder = StateGraph(list)
+    builder.add_node(Node("1"))
+    builder.add_node(Node("2"))
+    builder.add_node(Node("3"))
+    builder.add_edge(START, "1")
+    builder.add_conditional_edges("1", send_for_fun)
+    builder.add_conditional_edges("2", route_to_three)
+    graph = builder.compile()
+
+    assert graph.invoke(["0"]) == ["0", "1", "2", "3"]
+
+
 def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None:
     adder = mocker.Mock(side_effect=lambda x: x["total"] + x["input"])
 
@@ -1330,7 +1537,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
     workflow = Graph()
 
     workflow.add_node("agent", agent)
-    workflow.add_node("tools", execute_tools)
+    workflow.add_node("tools", execute_tools, metadata={"version": 2, "variant": "b"})
 
     workflow.set_entry_point("agent")
 
@@ -1344,6 +1551,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
 
     assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
     assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
+    assert app.get_graph().draw_mermaid() == snapshot
     assert json.dumps(app.get_graph(xray=True).to_json(), indent=2) == snapshot
     assert app.get_graph(xray=True).draw_mermaid(with_styles=False) == snapshot
 
@@ -2110,6 +2318,10 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None:
         intermediate_steps: Annotated[list[tuple[AgentAction, str]], operator.add]
         session: Annotated[httpx.Client, Context(httpx.Client)]
 
+    class ToolState(TypedDict, total=False):
+        agent_outcome: Union[AgentAction, AgentFinish]
+        session: Annotated[httpx.Client, Context(httpx.Client)]
+
     # Assemble the tools
     @tool()
     def search_api(query: str) -> str:
@@ -2148,9 +2360,11 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None:
     agent = prompt | llm | agent_parser
 
     # Define tool execution logic
-    def execute_tools(data: AgentState) -> dict:
+    def execute_tools(data: ToolState) -> dict:
         # check session in data
         assert isinstance(data["session"], httpx.Client)
+        assert "input" not in data
+        assert "intermediate_steps" not in data
         # execute the tool
         agent_action: AgentAction = data.pop("agent_outcome")
         observation = {t.name: t for t in tools}[agent_action.tool].invoke(
@@ -2172,7 +2386,7 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None:
     workflow = StateGraph(AgentState)
 
     workflow.add_node("agent", agent)
-    workflow.add_node("tools", execute_tools)
+    workflow.add_node("tools", execute_tools, input=ToolState)
 
     workflow.set_entry_point("agent")
 
@@ -3012,6 +3226,7 @@ Some examples of past conversations:
                         "name": "search_api",
                         "args": {"query": "query"},
                         "id": "tool_call123",
+                        "type": "tool_call",
                     }
                 ],
             ),
@@ -3023,7 +3238,7 @@ Some examples of past conversations:
             ),
             AIMessage(content="answer", id=AnyStr()),
         ]
-        assert app.invoke(
+        actual = app.invoke(
             {"messages": "what is weather in sf"},
             {
                 "configurable": {
@@ -3031,7 +3246,9 @@ Some examples of past conversations:
                     "expected_examples": [],
                 },
             },
-        ) == {"messages": first_messages}
+        )
+        expected = {"messages": first_messages}
+        assert actual == expected
 
         # get first checkpoint
         chkpnt_tuple_1 = saver.get_tuple({"configurable": {"thread_id": "1"}})
@@ -3062,6 +3279,7 @@ Some examples of past conversations:
                         "name": "search_api",
                         "args": {"query": "query"},
                         "id": "tool_call123",
+                        "type": "tool_call",
                     }
                 ],
             ),
@@ -3121,6 +3339,7 @@ Some examples of past conversations:
                             "name": "search_api",
                             "args": {"query": "query"},
                             "id": "tool_call123",
+                            "type": "tool_call",
                         }
                     ],
                 ),
@@ -5947,6 +6166,18 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
             "my_key": "value ⛰️",
             "market": "DE",
         }
+        assert [c.metadata for c in tool_two.checkpointer.list(thread1)] == [
+            {
+                "source": "loop",
+                "step": 0,
+                "writes": None,
+            },
+            {
+                "source": "input",
+                "step": -1,
+                "writes": {"my_key": "value ⛰️", "market": "DE"},
+            },
+        ]
         assert tool_two.get_state(thread1) == StateSnapshot(
             values={"my_key": "value ⛰️", "market": "DE"},
             next=("tool_two_slow",),
@@ -6656,7 +6887,7 @@ def test_in_one_fan_out_state_graph_waiting_edge(snapshot: SnapshotAssertion) ->
         },
     )
 
-    assert [c for c in app_w_interrupt.stream(None, config)] == [
+    assert [c for c in app_w_interrupt.stream(None, config, debug=1)] == [
         {"qa": {"answer": "doc1,doc2,doc3,doc4,doc5"}},
     ]