langgraph: release 0.1.12 (#1116)

Fix StateGraph.add_edge end_key check

.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]
 

Makefile

@@ -1,6 +1,14 @@
-.PHONY: build-docs serve-docs serve-clean-docs clean-docs codespell
+.PHONY: build-docs serve-docs serve-clean-docs clean-docs codespell build-typedoc
 
-build-docs:
+build-typedoc:
+	cd libs/sdk-js && yarn install --include-dev && yarn typedoc
+	cd libs/sdk-js && yarn --silent concat-md --decrease-title-levels --ignore=js_ts_sdk_ref.md --start-title-level-at 2 docs > ../../docs/docs/cloud/reference/sdk/js_ts_sdk_ref.md 2>/dev/null
+	#  Add links to the monorepo
+	sed -e '1,10s|@langchain/langgraph-sdk|[@langchain/langgraph-sdk](https://github.com/langchain-ai/langgraph/tree/main/libs/sdk-js)|g' docs/docs/cloud/reference/sdk/js_ts_sdk_ref.md > temp_file && mv temp_file docs/docs/cloud/reference/sdk/js_ts_sdk_ref.md
+
+
+
+build-docs: build-typedoc
 	poetry run python docs/_scripts/copy_notebooks.py
 	poetry run python -m mkdocs build --clean -f docs/mkdocs.yml --strict
 
@@ -8,7 +16,7 @@ serve-clean-docs: clean-docs
 	poetry run python docs/_scripts/copy_notebooks.py
 	poetry run python -m mkdocs serve -c -f docs/mkdocs.yml --strict -w ./libs/langgraph
 
-serve-docs:
+serve-docs: build-typedoc
 	poetry run python docs/_scripts/copy_notebooks.py
 	poetry run python -m mkdocs serve -f docs/mkdocs.yml -w ./libs/langgraph --dirty
 

README.md

@@ -73,8 +73,8 @@ def search(query: str):
     """Call to surf the web."""
     # This is a placeholder, but don't tell the LLM that...
     if "sf" in query.lower() or "san francisco" in query.lower():
-        return ["It's 60 degrees and foggy."]
-    return ["It's 90 degrees and sunny."]
+        return "It's 60 degrees and foggy."
+    return "It's 90 degrees and sunny."
 
 
 tools = [search]

docs/.gitignore

@@ -1,2 +1,4 @@
 *.ipynb
 site/
+docs/tutorials/**/*.png
+docs/cloud/reference/sdk/js_ts_sdk_ref.md

docs/_scripts/copy_notebooks.py

@@ -27,6 +27,8 @@ _MANUAL = {
         "streaming-events-from-within-tools-without-langchain.ipynb",
         "streaming-from-final-node.ipynb",
         "persistence.ipynb",
+        "input_output_schema.ipynb",
+        "pass_private_state.ipynb",
         "memory/manage-conversation-history.ipynb",
         "memory/delete-messages.ipynb",
         "memory/add-summary-conversation-history.ipynb",
@@ -38,6 +40,9 @@ _MANUAL = {
         "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",
@@ -53,6 +58,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",
@@ -88,13 +94,13 @@ _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",
         "learning.ipynb",
         "docs/quickstart.ipynb",
         "tutorials/rag-agent-testing.ipynb",
+        "tutorials/rag-agent-testing-local.ipynb",
         "time-travel.ipynb",
         "code_assistant/langgraph_code_assistant_mistral.ipynb",
     ]

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
 

docs/docs/cloud/deployment/setup.md

@@ -19,6 +19,24 @@ After each step, an example file directory is provided to demonstrate how code c
 
 Dependencies can optionally be specified in one of the following files: `pyproject.toml`, `setup.py`, or `requirements.txt`. If none of these files is created, then dependencies can be specified later in the [LangGraph API configuration file](#create-langgraph-api-config).
 
+The dependencies below will be included in the image, you can also use them in your code, as long as with a compatible version range:
+```
+langgraph>=0.1.7
+langchain-core>=0.2.7
+orjson>=3.10.1
+langsmith>=0.1.50
+httpx>=0.27.0
+langchain-core>=0.2.8
+langsmith>=0.1.63
+tenacity>=8.3.0
+uvicorn>=0.29.0
+sse-starlette>=2.1.0
+uvloop>=0.19.0
+httptools>=0.6.1
+jsonschema-rs>=0.18.0
+croniter>=1.0.1
+```
+
 Example `requirements.txt` file:
 ```
 langgraph
@@ -121,4 +139,4 @@ To deploy the LangGraph application to LangGraph Cloud, the code must be uploade
 
 ## Next
 
-After you setup your repo, it's time to [deploy your app](./cloud.md).
+After you setup your repo, it's time to [deploy your app](./cloud.md).

docs/docs/cloud/deployment/setup_pyproject.md

@@ -20,6 +20,24 @@ After each step, an example file directory is provided to demonstrate how code c
 
 Dependencies can optionally be specified in one of the following files: `pyproject.toml`, `setup.py`, or `requirements.txt`. If none of these files is created, then dependencies can be specified later in the [LangGraph API configuration file](#create-langgraph-api-config).
 
+The dependencies below will be included in the image, you can also use them in your code, as long as with a compatible version range:
+```
+langgraph>=0.1.7
+langchain-core>=0.2.7
+orjson>=3.10.1
+langsmith>=0.1.50
+httpx>=0.27.0
+langchain-core>=0.2.8
+langsmith>=0.1.63
+tenacity>=8.3.0
+uvicorn>=0.29.0
+sse-starlette>=2.1.0
+uvloop>=0.19.0
+httptools>=0.6.1
+jsonschema-rs>=0.18.0
+croniter>=1.0.1
+```
+
 Example `pyproject.toml` file:
 
 ```toml
@@ -33,7 +51,7 @@ readme = "README.md"
 
 [tool.poetry.dependencies]
 python = ">=3.9.0,<3.13"
-langgraph = "^0.1.0"
+langgraph = "^0.1.7"
 langchain-fireworks = "^0.1.3"
 
 
@@ -102,7 +120,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 +151,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

@@ -2,6 +2,8 @@
 
 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:
@@ -10,6 +12,12 @@ Install the proper packages:
 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:

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:
@@ -49,6 +52,7 @@ When creating complex graphs, leaving every decision up to the LLM can be danger
 LangGraph Studio is a built-in UI for visualizing, testing, and debugging your agents.
 
 - [How to enter LangGraph Studio](./test_deployment.md)
+- [How to enter LangGraph Studio for local deployment](./test_local_deployment.md)
 - [How to test your graph in LangGraph Studio](./invoke_studio.md)
 - [Interact with threads in LangGraph Studio](./threads_studio.md)
 

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

@@ -0,0 +1,28 @@
+# LangGraph Studio With Local Deployment
+
+!!! warning "Browser Compatibility"
+    Viewing the studio page of a local LangGraph deployment does not work in Safari. Use Chrome instead.
+
+## Setup
+
+Make sure you have setup your app correctly, by creating a compiled graph, a `.env` file with any environment variables, and a `langgraph.json` config file that points to your environment file and compiled graph. See [here](https://langchain-ai.github.io/langgraph/cloud/deployment/setup/) for more detailed instructions.
+
+After you have your app setup, head into the directory with your `langgraph.json` file and call `langgraph up -c langgraph.json --watch` to start the API server in watch mode which means it will restart on code changes, which is ideal for local testing. If the API server start correctly you should see logs that look something like this:
+
+    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
+
+Read this [reference](https://langchain-ai.github.io/langgraph/cloud/reference/cli/#up) to learn about all the options for starting the API server.
+
+## Access Studio
+
+Once you have successfully started the API server, you can access the studio by going to the following URL: `https://smith.langchain.com/studio/?baseUrl=http://127.0.0.1:8123` (see warning above if using Safari).
+
+If everything is working correctly you should see the studio show up looking something like this (with your graph diagram on the left hand side):
+
+![LangGraph Studio](./img/studio_screenshot.png)
+
+## Use the Studio for Testing
+
+To learn about how to use the studio for testing, read the [LangGraph Studio how-tos](https://langchain-ai.github.io/langgraph/cloud/how-tos/#langgraph-studio).

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).
 

docs/docs/concepts/low_level.md

@@ -46,6 +46,9 @@ The first thing you do when you define a graph is define the `State` of the grap
 
 The main documented way to specify the schema of a graph is by using `TypedDict`. However, we also support [using a Pydantic BaseModel](../how-tos/state-model.ipynb) as your graph state to add **default values** and additional data validation.
 
+By default, the graph will have the same input and output schemas. If you want to change this, you can also specify explicit input and output schemas directly. This is useful when you have a lot of keys, and some are explicitly for input and others for output. See the [notebook here](../how-tos/input_output_schema.ipynb) for how to use.
+
+By default, all nodes in the graph will share the same state. This means that they will read and write to the same state channels. It is possible to have nodes write to private state channels inside the graph for internal node communication - see [this notebook](../how-tos/pass_private_state.ipynb) for how to do that.
 ### Reducers
 
 Reducers are key to understanding how updates from nodes are applied to the `State`. Each key in the `State` has its own independent reducer function. If no reducer function is explicitly specified then it is assumed that all updates to that key should override it. Let's take a look at a few examples to understand them better.
@@ -180,7 +183,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.
@@ -190,7 +193,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
@@ -327,6 +330,16 @@ The final thing you specify when calling `update_state` is `as_node`. This updat
 
 The reason this matters is that the next steps in the graph to execute depend on the last node to have given an update, so this can be used to control which node executes next.
 
+## Graph Migrations
+
+LangGraph can easily handle migrations of graph definitions (nodes, edges, and state) even when using a checkpointer to track state.
+
+- For threads at the end of the graph (i.e. not interrupted) you can change the entire topology of the graph (i.e. all nodes and edges, remove, add, rename, etc)
+- For threads currently interrupted, we support all topology changes other than renaming / removing nodes (as that thread could now be about to enter a node that no longer exists) -- if this is a blocker please reach out and we can prioritize a solution.
+- For modifying state, we have full backwards and forwards compatibility for adding and removing keys
+- State keys that are renamed lose their saved state in existing threads
+- State keys whose types change in incompatible ways could currently cause issues in threads with state from before the change -- if this is a blocker please reach out and we can prioritize a solution.
+
 ## Configuration
 
 When creating a graph, you can also mark that certain parts of the graph are configurable. This is commonly done to enable easily switching between models or system prompts. This allows you to create a single "cognitive architecture" (the graph) but have multiple different instance of it.

docs/docs/how-tos/index.md

@@ -54,12 +54,21 @@ These guides show how to use different streaming modes.
 - [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

@@ -152,12 +152,21 @@ nav:
           - 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
+      - State Management:
+          - Use Pydantic model as state: how-tos/state-model.ipynb
+          - Use a context object in state: how-tos/state-context-key.ipynb
+          - Have a separate input and output schema: how-tos/input_output_schema.ipynb
+          - Pass private state between nodes inside the graph: how-tos/pass_private_state.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
@@ -205,6 +214,7 @@ nav:
               - Replay and Branch from Prior States: "cloud/how-tos/human_in_the_loop_time_travel.md"
           - LangGraph Studio:
               - Test Cloud Deployment: "cloud/how-tos/test_deployment.md"
+              - Test Local Deployment: "cloud/how-tos/test_local_deployment.md"
               - Invoke graph in LangGraph Studio: "cloud/how-tos/invoke_studio.md"
               - Interact with threads in LangGraph Studio: "cloud/how-tos/threads_studio.md"
           - Different Types of Runs:

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

@@ -32,10 +32,7 @@
     "scrolled": true
    },
    "outputs": [],
-   "source": [
-    "%%capture --no-stderr\n",
-    "%pip install -U langgraph langchain-community langchain-openai scikit-learn"
-   ]
+   "source": ["%%capture --no-stderr\n%pip install -U langgraph langchain-community langchain-openai scikit-learn"]
   },
   {
    "cell_type": "markdown",
@@ -51,15 +48,7 @@
    "id": "3d1ef253-6b0c-4481-868c-e1fe84f2c8ff",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "import requests\n",
-    "\n",
-    "url = \"https://storage.googleapis.com/benchmarks-artifacts/chinook/Chinook.db\"\n",
-    "response = requests.get(url)\n",
-    "\n",
-    "with open(\"Chinook.db\", \"wb\") as file:\n",
-    "    file.write(response.content)"
-   ]
+   "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",
@@ -88,12 +77,7 @@
      "output_type": "execute_result"
     }
    ],
-   "source": [
-    "from langchain_community.utilities import SQLDatabase\n",
-    "\n",
-    "db = SQLDatabase.from_uri(\"sqlite:///Chinook.db\")\n",
-    "db.get_usable_table_names()"
-   ]
+   "source": ["from langchain_community.utilities import SQLDatabase\n\ndb = SQLDatabase.from_uri(\"sqlite:///Chinook.db\")\ndb.get_usable_table_names()"]
   },
   {
    "cell_type": "markdown",
@@ -112,11 +96,7 @@
    "id": "d9ea4e80-30e6-4d46-b480-35f0be2fb055",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "from langchain_openai import ChatOpenAI\n",
-    "\n",
-    "model = ChatOpenAI(temperature=0, streaming=True, model=\"gpt-4o\")"
-   ]
+   "source": ["from langchain_openai import ChatOpenAI\n\nmodel = ChatOpenAI(temperature=0, streaming=True, model=\"gpt-4o\")"]
   },
   {
    "cell_type": "markdown",
@@ -138,9 +118,7 @@
    "id": "ea958e9f-ab1f-49b5-bd85-16332055297c",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "from langchain_core.messages import HumanMessage, SystemMessage"
-   ]
+   "source": ["from langchain_core.messages import HumanMessage, SystemMessage"]
   },
   {
    "cell_type": "markdown",
@@ -159,12 +137,7 @@
    "id": "975b039a",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "# This tool is given to the agent to look up information about a customer\n",
-    "def 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};\")"
-   ]
+   "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",
@@ -172,20 +145,7 @@
    "id": "1d5fa446",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "customer_prompt = \"\"\"Your job is to help a user update their profile.\n",
-    "\n",
-    "You 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",
-    "\n",
-    "If you are unable to help the user, you can \"\"\"\n",
-    "\n",
-    "\n",
-    "def get_customer_messages(messages):\n",
-    "    return [SystemMessage(content=customer_prompt)] + messages\n",
-    "\n",
-    "\n",
-    "customer_chain = get_customer_messages | model.bind_tools([get_customer_info])"
-   ]
+   "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",
@@ -206,19 +166,7 @@
    "id": "a8604a3b-b484-4b2b-a914-4236cb98c524",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "from langchain_community.vectorstores import SKLearnVectorStore\n",
-    "from langchain_openai import OpenAIEmbeddings\n",
-    "\n",
-    "artists = db._execute(\"select * from Artist\")\n",
-    "songs = db._execute(\"select * from Track\")\n",
-    "artist_retriever = SKLearnVectorStore.from_texts(\n",
-    "    [a[\"Name\"] for a in artists], OpenAIEmbeddings(), metadatas=artists\n",
-    ").as_retriever()\n",
-    "song_retriever = SKLearnVectorStore.from_texts(\n",
-    "    [a[\"Name\"] for a in songs], OpenAIEmbeddings(), metadatas=songs\n",
-    ").as_retriever()"
-   ]
+   "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",
@@ -234,16 +182,7 @@
    "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",
-    "    )"
-   ]
+   "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",
@@ -259,16 +198,7 @@
    "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",
-    "    )"
-   ]
+   "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",
@@ -284,11 +214,7 @@
    "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)"
-   ]
+   "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",
@@ -304,23 +230,7 @@
    "id": "72a14d5c",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "song_system_message = \"\"\"Your job is to help a customer find any songs they are looking for. \n",
-    "\n",
-    "You 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",
-    "\n",
-    "When looking up artists and songs, sometimes the artist/song will not be found. In that case, the tools will return information \\\n",
-    "on similar songs and artists. This is intentional, it is not the tool messing up.\"\"\"\n",
-    "\n",
-    "\n",
-    "def get_song_messages(messages):\n",
-    "    return [SystemMessage(content=song_system_message)] + messages\n",
-    "\n",
-    "\n",
-    "song_recc_chain = get_song_messages | model.bind_tools(\n",
-    "    [get_albums_by_artist, get_tracks_by_artist, check_for_songs]\n",
-    ")"
-   ]
+   "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",
@@ -339,10 +249,7 @@
      "output_type": "execute_result"
     }
    ],
-   "source": [
-    "msgs = [HumanMessage(content=\"hi! can you help me find songs by amy whinehouse?\")]\n",
-    "song_recc_chain.invoke(msgs)"
-   ]
+   "source": ["msgs = [HumanMessage(content=\"hi! can you help me find songs by amy whinehouse?\")]\nsong_recc_chain.invoke(msgs)"]
   },
   {
    "cell_type": "markdown",
@@ -360,32 +267,7 @@
    "id": "73e74268",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "from langchain_core.messages import AIMessage, HumanMessage, SystemMessage\n",
-    "from langchain_core.pydantic_v1 import BaseModel, Field\n",
-    "\n",
-    "\n",
-    "class 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",
-    "\n",
-    "system_message = \"\"\"Your job is to help as a customer service representative for a music store.\n",
-    "\n",
-    "You 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",
-    "\n",
-    "If the user is asking or wants to ask about updating or accessing their information, send them to that route.\n",
-    "If the user is asking or wants to ask about music, send them to that route.\n",
-    "Otherwise, respond.\"\"\"\n",
-    "\n",
-    "\n",
-    "def get_messages(messages):\n",
-    "    return [SystemMessage(content=system_message)] + messages"
-   ]
+   "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",
@@ -393,9 +275,7 @@
    "id": "ddf27314",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "chain = get_messages | model.bind_tools([Router])"
-   ]
+   "source": ["chain = get_messages | model.bind_tools([Router])"]
   },
   {
    "cell_type": "code",
@@ -414,10 +294,7 @@
      "output_type": "execute_result"
     }
    ],
-   "source": [
-    "msgs = [HumanMessage(content=\"hi! can you help me find a good song?\")]\n",
-    "chain.invoke(msgs)"
-   ]
+   "source": ["msgs = [HumanMessage(content=\"hi! can you help me find a good song?\")]\nchain.invoke(msgs)"]
   },
   {
    "cell_type": "code",
@@ -436,10 +313,7 @@
      "output_type": "execute_result"
     }
    ],
-   "source": [
-    "msgs = [HumanMessage(content=\"hi! what's the email you have for me?\")]\n",
-    "chain.invoke(msgs)"
-   ]
+   "source": ["msgs = [HumanMessage(content=\"hi! what's the email you have for me?\")]\nchain.invoke(msgs)"]
   },
   {
    "cell_type": "code",
@@ -447,15 +321,7 @@
    "id": "bd6ddd8b-7500-46a7-811d-3bcb937bda51",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "from langchain_core.messages import AIMessage\n",
-    "\n",
-    "\n",
-    "def add_name(message, name):\n",
-    "    _dict = message.dict()\n",
-    "    _dict[\"name\"] = name\n",
-    "    return AIMessage(**_dict)"
-   ]
+   "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",
@@ -463,45 +329,7 @@
    "id": "27494de5-8345-4c23-bc0e-81e0dd5d47d8",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "import json\n",
-    "\n",
-    "from langgraph.graph import END\n",
-    "\n",
-    "\n",
-    "def _get_last_ai_message(messages):\n",
-    "    for m in messages[::-1]:\n",
-    "        if isinstance(m, AIMessage):\n",
-    "            return m\n",
-    "    return None\n",
-    "\n",
-    "\n",
-    "def _is_tool_call(msg):\n",
-    "    return hasattr(msg, \"additional_kwargs\") and \"tool_calls\" in msg.additional_kwargs\n",
-    "\n",
-    "\n",
-    "def _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\""
-   ]
+   "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",
@@ -509,12 +337,7 @@
    "id": "8aec704a-46fe-4fb3-bdee-11c3bbffc370",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "from langgraph.prebuilt import ToolNode\n",
-    "\n",
-    "tools = [get_albums_by_artist, get_tracks_by_artist, check_for_songs, get_customer_info]\n",
-    "tool_node = ToolNode(tools)"
-   ]
+   "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",
@@ -522,16 +345,7 @@
    "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"
-   ]
+   "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",
@@ -539,13 +353,7 @@
    "id": "fd4dbf98-dbb3-411a-bad6-2bb334072aaf",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "from functools import partial\n",
-    "\n",
-    "general_node = _filter_out_routes | chain | partial(add_name, name=\"general\")\n",
-    "music_node = _filter_out_routes | song_recc_chain | partial(add_name, name=\"music\")\n",
-    "customer_node = _filter_out_routes | customer_chain | partial(add_name, name=\"customer\")"
-   ]
+   "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",
@@ -553,33 +361,7 @@
    "id": "dcade924",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "from langgraph.checkpoint.sqlite import SqliteSaver\n",
-    "\n",
-    "from langgraph.graph import MessageGraph\n",
-    "\n",
-    "memory = SqliteSaver.from_conn_string(\":memory:\")\n",
-    "graph = MessageGraph()\n",
-    "nodes = {\n",
-    "    \"general\": \"general\",\n",
-    "    \"music\": \"music\",\n",
-    "    END: END,\n",
-    "    \"tools\": \"tools\",\n",
-    "    \"customer\": \"customer\",\n",
-    "}\n",
-    "# Define a new graph\n",
-    "workflow = MessageGraph()\n",
-    "workflow.add_node(\"general\", general_node)\n",
-    "workflow.add_node(\"music\", music_node)\n",
-    "workflow.add_node(\"customer\", customer_node)\n",
-    "workflow.add_node(\"tools\", tool_node)\n",
-    "workflow.add_conditional_edges(\"general\", _route, nodes)\n",
-    "workflow.add_conditional_edges(\"tools\", _route, nodes)\n",
-    "workflow.add_conditional_edges(\"music\", _route, nodes)\n",
-    "workflow.add_conditional_edges(\"customer\", _route, nodes)\n",
-    "workflow.set_conditional_entry_point(_route, nodes)\n",
-    "graph = workflow.compile()"
-   ]
+   "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",
@@ -715,27 +497,7 @@
      ]
     }
    ],
-   "source": [
-    "import uuid\n",
-    "\n",
-    "from langchain_core.messages import HumanMessage\n",
-    "\n",
-    "from langgraph.graph.graph import START\n",
-    "\n",
-    "history = []\n",
-    "while 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\")"
-   ]
+   "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/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/human_in_the_loop/edit-graph-state.ipynb

@@ -135,21 +135,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",
@@ -162,7 +167,7 @@
     "# Set up memory\n",
     "memory = MemorySaver()\n",
     "\n",
-    "# Add \n",
+    "# Add\n",
     "graph = builder.compile(checkpointer=memory, interrupt_before=[\"step_2\"])\n",
     "\n",
     "# View\n",

examples/human_in_the_loop/wait-user-input.ipynb

@@ -137,25 +137,30 @@
     "from langgraph.checkpoint.memory import MemorySaver\n",
     "from IPython.display import Image, display\n",
     "\n",
+    "\n",
     "class State(TypedDict):\n",
     "    input: str\n",
     "    user_feedback: str\n",
     "\n",
+    "\n",
     "def step_1(state):\n",
     "    print(\"---Step 1---\")\n",
     "    pass\n",
     "\n",
+    "\n",
     "def human_feedback(state):\n",
     "    print(\"---human_feedback---\")\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(\"human_feedback\", step_2)\n",
+    "builder.add_node(\"human_feedback\", human_feedback)\n",
     "builder.add_node(\"step_3\", step_3)\n",
     "builder.add_edge(START, \"step_1\")\n",
     "builder.add_edge(\"step_1\", \"human_feedback\")\n",
@@ -165,7 +170,7 @@
     "# Set up memory\n",
     "memory = MemorySaver()\n",
     "\n",
-    "# Add \n",
+    "# Add\n",
     "graph = builder.compile(checkpointer=memory, interrupt_before=[\"human_feedback\"])\n",
     "\n",
     "# View\n",
@@ -253,7 +258,7 @@
     "\n",
     "# We now update the state as if we are the human_feedback node\n",
     "graph.update_state(thread, {\"user_feedback\": user_input}, as_node=\"human_feedback\")\n",
-    "    \n",
+    "\n",
     "# We can check the state\n",
     "print(\"--State after update--\")\n",
     "print(graph.get_state(thread))\n",

examples/input_output_schema.ipynb

@@ -0,0 +1,93 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "f262985e-e973-4a27-9c9e-dbb3a06a35b7",
+   "metadata": {},
+   "source": [
+    "# How to define input/output schema for your graph\n",
+    "\n",
+    "By default, `StateGraph` takes in a single schema and all nodes are expected to communicate with that schema. However, it is also possible to define explicit input and output schemas for a graph. This is helpful if you want to draw a distinction between input and output keys.\n",
+    "\n",
+    "In this notebook we'll walk through an example of this. At a high level, in order to do this you simply have to pass in `input=..., output=...` when defining the graph. Let's see an example below!"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "6ec0eb77-874e-443e-8c73-93125b515106",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'answer': 'bye'}"
+      ]
+     },
+     "execution_count": 12,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from langgraph.graph import StateGraph, START, END\n",
+    "from typing import TypedDict\n",
+    "\n",
+    "class InputState(TypedDict):\n",
+    "    question: str\n",
+    "\n",
+    "class OutputState(TypedDict):\n",
+    "    answer: str\n",
+    "\n",
+    "def answer_node(state: InputState):\n",
+    "    return {\"answer\": \"bye\"}\n",
+    "\n",
+    "check = SqliteSaver.from_conn_string(\":memory:\")\n",
+    "graph = StateGraph(input=InputState, output=OutputState)\n",
+    "graph.add_node(answer_node)\n",
+    "graph.add_edge(START, \"answer_node\")\n",
+    "graph.add_edge(\"answer_node\", END)\n",
+    "graph = graph.compile()\n",
+    "\n",
+    "graph.invoke({\"question\": \"hi\"})"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6a68836f-98e1-4684-a8a6-c1473c73460c",
+   "metadata": {},
+   "source": [
+    "Notice that the output of invoke only includes the output schema."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b952a554-f2a4-4be3-81ab-2e08f0f441c2",
+   "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/introduction.ipynb

@@ -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

@@ -110,23 +110,26 @@
     "\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",
+    "    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",
+    "        messages = [SystemMessage(content=system_message)] + state[\"messages\"]\n",
     "    else:\n",
-    "        messages = state['messages']\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",
@@ -145,7 +148,7 @@
     "\n",
     "def summarize_conversation(state: State):\n",
     "    # First, we summarize the conversation\n",
-    "    summary = state.get('summary', '')\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",
@@ -155,17 +158,13 @@
     "        )\n",
     "    else:\n",
     "        summary_message = \"Create a summary of the conversation above:\"\n",
-    "        \n",
-    "    messages = state['messages'] + [HumanMessage(content=summary_message)]\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 {\n",
-    "        \"summary\": response.content,\n",
-    "        \"messages\": delete_messages\n",
-    "    }\n",
-    "    \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",
@@ -212,10 +211,10 @@
    "source": [
     "def print_update(update):\n",
     "    for k, v in update.items():\n",
-    "        for m in v['messages']:\n",
+    "        for m in v[\"messages\"]:\n",
     "            m.pretty_print()\n",
-    "        if 'summary' in v:\n",
-    "            print(v['summary'])"
+    "        if \"summary\" in v:\n",
+    "            print(v[\"summary\"])"
    ]
   },
   {

examples/memory/delete-messages.ipynb

@@ -28,7 +28,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "id": "af4ce0ba-7596-4e5f-8bf8-0b0bd6e62833",
    "metadata": {},
    "outputs": [],
@@ -50,7 +50,15 @@
    "execution_count": 2,
    "id": "c903a1cf-2977-4e2d-ad7d-8b3946821d89",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdin",
+     "output_type": "stream",
+     "text": [
+      "ANTHROPIC_API_KEY:  ········\n"
+     ]
+    }
+   ],
    "source": [
     "import getpass\n",
     "import os\n",
@@ -94,7 +102,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 4,
    "id": "378899a9-3b9a-4748-95b6-eb00e0828677",
    "metadata": {},
    "outputs": [],
@@ -176,7 +184,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 5,
    "id": "57b27553-21be-43e5-ac48-d1d0a3aa0dca",
    "metadata": {},
    "outputs": [
@@ -189,7 +197,7 @@
       "hi! I'm bob\n",
       "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
       "\n",
-      "Hello Bob! It's nice to meet you. How can I assist you today?\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",
@@ -225,26 +233,26 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 6,
    "id": "8a850529-d038-48f7-b5a2-8d4d2923f83a",
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "[HumanMessage(content=\"hi! I'm bob\", id='3e1098f8-2657-42d3-b58a-7c2f46930b8c'),\n",
-       " AIMessage(content=\"Hello Bob! It's nice to meet you. How can I assist you today?\", response_metadata={'id': 'msg_01HT8MUEN4p16wbYv9Xm7kfr', 'model': 'claude-3-haiku-20240307', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 12, 'output_tokens': 20}}, id='run-86348912-72c4-42b0-b3e0-a47c4ebd1e52-0'),\n",
-       " HumanMessage(content=\"what's my name?\", id='9c3ef235-ec5c-4e57-a3b2-c17502de496d'),\n",
-       " AIMessage(content='Your name is Bob, as you introduced yourself at the beginning of our conversation.', response_metadata={'id': 'msg_01LVhb56f6RpAAoxASZrLzmK', 'model': 'claude-3-haiku-20240307', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 40, 'output_tokens': 19}}, id='run-e3d7447f-046a-4dfa-8813-38134dbcd1ef-0')]"
+       "[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": 7,
+     "execution_count": 6,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
-    "messages = app.get_state(config).values['messages']\n",
+    "messages = app.get_state(config).values[\"messages\"]\n",
     "messages"
    ]
   },
@@ -258,24 +266,33 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "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': '1ef3d750-5bc4-67c6-8005-9490a1b276f5'}}"
+       "  'thread_ts': '1ef42d00-d9ad-6f24-8005-feb089654def'}}"
       ]
      },
-     "execution_count": 9,
+     "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)})"
    ]
   },
@@ -289,25 +306,25 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 8,
    "id": "8bfe4ffa-e170-43bc-aec4-6e36ac620931",
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "[AIMessage(content=\"Hello Bob! It's nice to meet you. How can I assist you today?\", response_metadata={'id': 'msg_01HT8MUEN4p16wbYv9Xm7kfr', 'model': 'claude-3-haiku-20240307', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 12, 'output_tokens': 20}}, id='run-86348912-72c4-42b0-b3e0-a47c4ebd1e52-0'),\n",
-       " HumanMessage(content=\"what's my name?\", id='9c3ef235-ec5c-4e57-a3b2-c17502de496d'),\n",
-       " AIMessage(content='Your name is Bob, as you introduced yourself at the beginning of our conversation.', response_metadata={'id': 'msg_01LVhb56f6RpAAoxASZrLzmK', 'model': 'claude-3-haiku-20240307', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 40, 'output_tokens': 19}}, id='run-e3d7447f-046a-4dfa-8813-38134dbcd1ef-0')]"
+       "[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": 11,
+     "execution_count": 8,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
-    "messages = app.get_state(config).values['messages']\n",
+    "messages = app.get_state(config).values[\"messages\"]\n",
     "messages"
    ]
   },
@@ -323,7 +340,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 9,
    "id": "bb22ede0-e153-4fd0-a4c0-f9af2f7663b1",
    "metadata": {},
    "outputs": [],
@@ -333,10 +350,11 @@
     "\n",
     "\n",
     "def delete_messages(state):\n",
-    "    messages = state['messages']\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",
@@ -358,7 +376,10 @@
     "\n",
     "\n",
     "workflow.add_edge(START, \"agent\")\n",
-    "workflow.add_conditional_edges(\"agent\", should_continue,)\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",
@@ -376,7 +397,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": 10,
    "id": "3975f34c-c243-40ea-b9d2-424d50a48dc9",
    "metadata": {},
    "outputs": [
@@ -384,21 +405,11 @@
      "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! 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",
-      "You said your name is Bob.\n",
-      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
-      "\n",
-      "You said your name is Bob.\n"
+      "[('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"
      ]
     }
    ],
@@ -408,12 +419,12 @@
     "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",
-    "    event[\"messages\"][-1].pretty_print()\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",
-    "    event[\"messages\"][-1].pretty_print()"
+    "    print([(message.type, message.content) for message in event[\"messages\"]])"
    ]
   },
   {
@@ -426,25 +437,25 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 17,
+   "execution_count": 11,
    "id": "a3e15abb-81d8-4072-9f10-61ae0fd61dac",
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "[AIMessage(content=\"It's nice to meet you, Bob! How can I assist you today?\", response_metadata={'id': 'msg_01QMoxepDiCcKQ6XFgge1QQT', 'model': 'claude-3-haiku-20240307', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 12, 'output_tokens': 19}}, id='run-de13ba05-095d-4fd1-907a-6766ef3bf57b-0'),\n",
-       " HumanMessage(content=\"what's my name?\", id='8292e725-8fc4-487e-a9b6-75f8b136bec2'),\n",
-       " AIMessage(content='You said your name is Bob.', response_metadata={'id': 'msg_01DfWfaxavdMCqtoQRmC3mc4', 'model': 'claude-3-haiku-20240307', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 39, 'output_tokens': 10}}, id='run-28167c82-e126-47e4-854c-623e50c8af22-0')]"
+       "[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": 17,
+     "execution_count": 11,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
-    "messages = app.get_state(config).values['messages']\n",
+    "messages = app.get_state(config).values[\"messages\"]\n",
     "messages"
    ]
   },
@@ -467,9 +478,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "langgraph",
    "language": "python",
-   "name": "python3"
+   "name": "langgraph"
   },
   "language_info": {
    "codemirror_mode": {
@@ -481,7 +492,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.1"
+   "version": "3.11.9"
   }
  },
  "nbformat": 4,

examples/memory/manage-conversation-history.ipynb

@@ -138,7 +138,7 @@
     "\n",
     "# Define the function that calls the model\n",
     "def call_model(state: MessagesState):\n",
-    "    response = model.invoke(state[\"messages\"])\n",
+    "    response = bound_model.invoke(state[\"messages\"])\n",
     "    # We return a list, because this will get added to the existing list\n",
     "    return {\"messages\": response}\n",
     "\n",
@@ -275,7 +275,7 @@
     "# Define the function that calls the model\n",
     "def call_model(state: MessagesState):\n",
     "    messages = filter_messages(state[\"messages\"])\n",
-    "    response = model.invoke(messages)\n",
+    "    response = bound_model.invoke(messages)\n",
     "    # We return a list, because this will get added to the existing list\n",
     "    return {\"messages\": response}\n",
     "\n",

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/pass_private_state.ipynb

@@ -0,0 +1,126 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "47ed5db3-bda5-49e1-bf75-23e08c9a3af0",
+   "metadata": {},
+   "source": [
+    "# How to pass private state\n",
+    "\n",
+    "Oftentimes, you may want nodes to be able to pass state to eachv other that should NOT be part of the main schema of the graph. This is often useful because there may be information that is not needed as input/output (and therefore doesn't really make sense to have in the main schema) but is ABSOLUTELY needed as part of the intermediate working logic.\n",
+    "\n",
+    "Let's take a look at an example below. In this example, we will create a RAG pipeline that:\n",
+    "1. Takes in a user question\n",
+    "2. Uses an LLM to generate a search query\n",
+    "3. Retrieves documents for that generated query\n",
+    "4. Generates a final answer based on those documents\n",
+    "\n",
+    "We will have a separate node for each step. We will only have the `question` and `answer` on the overall state. However, we will need separate states for the `search_query` and the `documents` - we will pass these as private state keys.\n",
+    "\n",
+    "Let's look at an example!"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "3114c3ad-0ade-47ba-9488-53d6f7671578",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'question': 'foo', 'answer': 'fo\\n\\nfo\\n\\nfoo'}"
+      ]
+     },
+     "execution_count": 9,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from langgraph.graph import StateGraph, START, END\n",
+    "from typing import TypedDict\n",
+    "\n",
+    "\n",
+    "# The overall state of the graph\n",
+    "class OverallState(TypedDict):\n",
+    "    question: str\n",
+    "    answer: str\n",
+    "\n",
+    "\n",
+    "# This is what the node that generates the query will return\n",
+    "class QueryOutputState(TypedDict):\n",
+    "    query: str\n",
+    "\n",
+    "\n",
+    "# This is what the node that retrieves the documents will return\n",
+    "class DocumentOutputState(TypedDict):\n",
+    "    docs: list[str]\n",
+    "\n",
+    "\n",
+    "# This is what the node that generates the final answer will take in\n",
+    "class GenerateInputState(OverallState, DocumentOutputState):\n",
+    "    pass\n",
+    "\n",
+    "\n",
+    "# Node to generate query\n",
+    "def generate_query(state: OverallState) -> QueryOutputState:\n",
+    "    # Replace this with real logic\n",
+    "    return {\"query\": state[\"question\"][:2]}\n",
+    "\n",
+    "\n",
+    "# Node to retrieve documents\n",
+    "def retrieve_documents(state: QueryOutputState) -> DocumentOutputState:\n",
+    "    # Replace this with real logic\n",
+    "    return {\"docs\": [state['query']] * 2}\n",
+    "\n",
+    "\n",
+    "# Node to generate answer\n",
+    "def generate(state: GenerateInputState) -> OverallState:\n",
+    "    return {\"answer\": \"\\n\\n\".join(state['docs'] + [state['question']])}\n",
+    "\n",
+    "\n",
+    "graph = StateGraph(OverallState)\n",
+    "graph.add_node(generate_query)\n",
+    "graph.add_node(retrieve_documents)\n",
+    "graph.add_node(generate)\n",
+    "graph.add_edge(START, \"generate_query\")\n",
+    "graph.add_edge(\"generate_query\", \"retrieve_documents\")\n",
+    "graph.add_edge(\"retrieve_documents\", \"generate\")\n",
+    "graph.add_edge(\"generate\", END)\n",
+    "graph = graph.compile()\n",
+    "\n",
+    "graph.invoke({\"question\": \"foo\"})"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3ffc2d8c-717f-42c9-b0aa-15b178a5cc8b",
+   "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/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_mongodb.ipynb

@@ -82,6 +82,7 @@
     "            return pickle.loads(data)\n",
     "        return super().loads(data)\n",
     "\n",
+    "\n",
     "class MongoDBSaver(AbstractContextManager, BaseCheckpointSaver):\n",
     "    \"\"\"A checkpoint saver that stores checkpoints in a MongoDB database.\n",
     "\n",
@@ -278,7 +279,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -309,17 +310,19 @@
     }
    ],
    "source": [
-    "from langgraph.graph import StateGraph\n",
+    "from langgraph.graph import StateGraph, START, END\n",
     "\n",
-    "checkpointer = MongoDBSaver(MongoClient(MONGO_URI), \"checkpoints_db\", \"checkpoints_collection\")\n",
+    "checkpointer = MongoDBSaver(\n",
+    "    MongoClient(MONGO_URI), \"checkpoints_db\", \"checkpoints_collection\"\n",
+    ")\n",
     "builder = StateGraph(int)\n",
     "builder.add_node(\"add_one\", lambda x: x + 1)\n",
-    "builder.set_entry_point(\"add_one\")\n",
-    "builder.set_finish_point(\"add_one\")\n",
+    "builder.add_edge(START, \"add_one\")\n",
+    "builder.add_edge(\"add_one\", END)\n",
     "graph = builder.compile(checkpointer=checkpointer)\n",
     "config = {\"configurable\": {\"thread_id\": \"123\"}}\n",
     "graph.get_state(config)\n",
-    "result = graph.invoke(3,config)\n",
+    "result = graph.invoke(3, config)\n",
     "graph.get_state(config)"
    ]
   },
@@ -572,7 +575,417 @@
     "for doc in collection.find():\n",
     "    print(doc)\n",
     "\n",
-    "#The checkpoints from both the examples have been saved in the database."
+    "# The checkpoints from both the examples have been saved in the database."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Asynchronous implementation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Async package for MongoDB\n",
+    "%pip install motor"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pickle\n",
+    "from contextlib import AbstractContextManager\n",
+    "from types import TracebackType\n",
+    "from typing import Any, Dict, Optional, AsyncIterator\n",
+    "\n",
+    "from langchain_core.runnables import RunnableConfig\n",
+    "from typing_extensions import Self\n",
+    "\n",
+    "from langgraph.checkpoint.base import (\n",
+    "    BaseCheckpointSaver,\n",
+    "    Checkpoint,\n",
+    "    CheckpointMetadata,\n",
+    "    CheckpointTuple,\n",
+    "    SerializerProtocol,\n",
+    ")\n",
+    "from langgraph.serde.jsonplus import JsonPlusSerializer\n",
+    "from motor.motor_asyncio import AsyncIOMotorClient\n",
+    "\n",
+    "\n",
+    "class JsonPlusSerializerCompat(JsonPlusSerializer):\n",
+    "    \"\"\"A serializer that supports loading pickled checkpoints for backwards compatibility.\n",
+    "\n",
+    "    This serializer extends the JsonPlusSerializer and adds support for loading pickled\n",
+    "    checkpoints. If the input data starts with b\"\\x80\" and ends with b\".\", it is treated\n",
+    "    as a pickled checkpoint and loaded using pickle.loads(). Otherwise, the default\n",
+    "    JsonPlusSerializer behavior is used.\n",
+    "\n",
+    "    Examples:\n",
+    "        >>> import pickle\n",
+    "        >>> from langgraph.checkpoint.sqlite import JsonPlusSerializerCompat\n",
+    "        >>>\n",
+    "        >>> serializer = JsonPlusSerializerCompat()\n",
+    "        >>> pickled_data = pickle.dumps({\"key\": \"value\"})\n",
+    "        >>> loaded_data = serializer.loads(pickled_data)\n",
+    "        >>> print(loaded_data)  # Output: {\"key\": \"value\"}\n",
+    "        >>>\n",
+    "        >>> json_data = '{\"key\": \"value\"}'.encode(\"utf-8\")\n",
+    "        >>> loaded_data = serializer.loads(json_data)\n",
+    "        >>> print(loaded_data)  # Output: {\"key\": \"value\"}\n",
+    "    \"\"\"\n",
+    "\n",
+    "    def loads(self, data: bytes) -> Any:\n",
+    "        if data.startswith(b\"\\x80\") and data.endswith(b\".\"):\n",
+    "            return pickle.loads(data)\n",
+    "        return super().loads(data)\n",
+    "\n",
+    "\n",
+    "class MongoDBSaver(AbstractContextManager, BaseCheckpointSaver):\n",
+    "    \"\"\"A checkpoint saver that stores checkpoints in a MongoDB database.\n",
+    "\n",
+    "    Args:\n",
+    "        client (AsyncIOMotorClient): The Async MongoDB client.\n",
+    "        db_name (str): The name of the database to use.\n",
+    "        collection_name (str): The name of the collection to use.\n",
+    "        serde (Optional[SerializerProtocol]): The serializer to use for serializing and deserializing checkpoints. Defaults to JsonPlusSerializerCompat.\n",
+    "\n",
+    "    Examples:\n",
+    "\n",
+    "        >>> from motor.motor_asyncio import AsyncIOMotorClient\n",
+    "        >>> from langgraph.checkpoint.mongodb import MongoDBSaver\n",
+    "        >>> from langgraph.graph import StateGraph\n",
+    "        >>>\n",
+    "        >>> builder = StateGraph(int)\n",
+    "        >>> builder.add_node(\"add_one\", lambda x: x + 1)\n",
+    "        >>> builder.set_entry_point(\"add_one\")\n",
+    "        >>> builder.set_finish_point(\"add_one\")\n",
+    "        >>> client = AsyncIOMotorClient(\"mongodb://localhost:27017/\")\n",
+    "        >>> memory = MongoDBSaver(client, \"checkpoints\", \"checkpoints\")\n",
+    "        >>> graph = builder.compile(checkpointer=memory)\n",
+    "        >>> config = {\"configurable\": {\"thread_id\": \"1\"}}\n",
+    "        >>> result = graph.ainvoke(3, config)\n",
+    "    \"\"\"\n",
+    "\n",
+    "    serde = JsonPlusSerializerCompat()\n",
+    "\n",
+    "    client: AsyncIOMotorClient\n",
+    "    db_name: str\n",
+    "    collection_name: str\n",
+    "\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        client: AsyncIOMotorClient,\n",
+    "        db_name: str,\n",
+    "        collection_name: str,\n",
+    "        *,\n",
+    "        serde: Optional[SerializerProtocol] = None,\n",
+    "    ) -> None:\n",
+    "        super().__init__(serde=serde)\n",
+    "        self.client = client\n",
+    "        self.db_name = db_name\n",
+    "        self.collection_name = collection_name\n",
+    "        self.collection = client[db_name][collection_name]\n",
+    "\n",
+    "    def __enter__(self) -> Self:\n",
+    "        return self\n",
+    "\n",
+    "    def __exit__(\n",
+    "        self,\n",
+    "        __exc_type: Optional[type[BaseException]],\n",
+    "        __exc_value: Optional[BaseException],\n",
+    "        __traceback: Optional[TracebackType],\n",
+    "    ) -> Optional[bool]:\n",
+    "        return True\n",
+    "\n",
+    "    async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:\n",
+    "        \"\"\"Get a checkpoint tuple from the database.\n",
+    "\n",
+    "        This method retrieves a checkpoint tuple from the MongoDB database based on the\n",
+    "        provided config. If the config contains a \"thread_ts\" key, the checkpoint with\n",
+    "        the matching thread ID and timestamp is retrieved. Otherwise, the latest checkpoint\n",
+    "        for the given thread ID is retrieved.\n",
+    "\n",
+    "        Args:\n",
+    "            config (RunnableConfig): The config to use for retrieving the checkpoint.\n",
+    "\n",
+    "        Returns:\n",
+    "            Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found.\n",
+    "        \"\"\"\n",
+    "        if config[\"configurable\"].get(\"thread_ts\"):\n",
+    "            query = {\n",
+    "                \"thread_id\": config[\"configurable\"][\"thread_id\"],\n",
+    "                \"thread_ts\": config[\"configurable\"][\"thread_ts\"],\n",
+    "            }\n",
+    "        else:\n",
+    "            query = {\"thread_id\": config[\"configurable\"][\"thread_id\"]}\n",
+    "        result = self.collection.find(query).sort(\"thread_ts\", -1).limit(1)\n",
+    "        async for doc in result:\n",
+    "            return CheckpointTuple(\n",
+    "                config,\n",
+    "                self.serde.loads(doc[\"checkpoint\"]),\n",
+    "                self.serde.loads(doc[\"metadata\"]),\n",
+    "                (\n",
+    "                    {\n",
+    "                        \"configurable\": {\n",
+    "                            \"thread_id\": doc[\"thread_id\"],\n",
+    "                            \"thread_ts\": doc[\"parent_ts\"],\n",
+    "                        }\n",
+    "                    }\n",
+    "                    if doc.get(\"parent_ts\")\n",
+    "                    else None\n",
+    "                ),\n",
+    "            )\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",
+    "    ) -> AsyncIterator[CheckpointTuple]:\n",
+    "        \"\"\"List checkpoints from the database.\n",
+    "\n",
+    "        This method retrieves a list of checkpoint tuples from the MongoDB database based\n",
+    "        on the provided config. The checkpoints are ordered by timestamp in descending order.\n",
+    "\n",
+    "        Args:\n",
+    "            config (RunnableConfig): The config to use for listing the checkpoints.\n",
+    "            before (Optional[RunnableConfig]): If provided, only checkpoints before the specified timestamp are returned. Defaults to None.\n",
+    "            limit (Optional[int]): The maximum number of checkpoints to return. Defaults to None.\n",
+    "\n",
+    "        Yields:\n",
+    "            AsyncIterator[CheckpointTuple]: An Async iterator of checkpoint tuples.\n",
+    "        \"\"\"\n",
+    "        query = {}\n",
+    "        if config is not None:\n",
+    "            query[\"thread_id\"] = config[\"configurable\"][\"thread_id\"]\n",
+    "        if filter:\n",
+    "            for key, value in filter.items():\n",
+    "                query[f\"metadata.{key}\"] = value\n",
+    "        if before is not None:\n",
+    "            query[\"thread_ts\"] = {\"$lt\": before[\"configurable\"][\"thread_ts\"]}\n",
+    "        result = self.collection.find(query).sort(\"thread_ts\", -1).limit(limit)\n",
+    "        if limit is not None:\n",
+    "            result = result.limit(limit)\n",
+    "        async for doc in result:\n",
+    "            yield CheckpointTuple(\n",
+    "                {\n",
+    "                    \"configurable\": {\n",
+    "                        \"thread_id\": doc[\"thread_id\"],\n",
+    "                        \"thread_ts\": doc[\"thread_ts\"],\n",
+    "                    }\n",
+    "                },\n",
+    "                self.serde.loads(doc[\"checkpoint\"]),\n",
+    "                self.serde.loads(doc[\"metadata\"]),\n",
+    "                (\n",
+    "                    {\n",
+    "                        \"configurable\": {\n",
+    "                            \"thread_id\": doc[\"thread_id\"],\n",
+    "                            \"thread_ts\": doc[\"parent_ts\"],\n",
+    "                        }\n",
+    "                    }\n",
+    "                    if doc.get(\"parent_ts\")\n",
+    "                    else None\n",
+    "                ),\n",
+    "            )\n",
+    "\n",
+    "    async def aput(\n",
+    "        self,\n",
+    "        config: RunnableConfig,\n",
+    "        checkpoint: Checkpoint,\n",
+    "        metadata: CheckpointMetadata,\n",
+    "    ) -> RunnableConfig:\n",
+    "        \"\"\"Save a checkpoint to the database.\n",
+    "\n",
+    "        This method saves a checkpoint to the MongoDB database. The checkpoint is associated\n",
+    "        with the provided config and its parent config (if any).\n",
+    "\n",
+    "        Args:\n",
+    "            config (RunnableConfig): The config to associate with the checkpoint.\n",
+    "            checkpoint (Checkpoint): The checkpoint to save.\n",
+    "            metadata (Optional[dict[str, Any]]): Additional metadata to save with the checkpoint. Defaults to None.\n",
+    "\n",
+    "        Returns:\n",
+    "            RunnableConfig: The updated config containing the saved checkpoint's timestamp.\n",
+    "        \"\"\"\n",
+    "        doc = {\n",
+    "            \"thread_id\": config[\"configurable\"][\"thread_id\"],\n",
+    "            \"thread_ts\": checkpoint[\"id\"],\n",
+    "            \"checkpoint\": self.serde.dumps(checkpoint),\n",
+    "            \"metadata\": self.serde.dumps(metadata),\n",
+    "        }\n",
+    "        if config[\"configurable\"].get(\"thread_ts\"):\n",
+    "            doc[\"parent_ts\"] = config[\"configurable\"][\"thread_ts\"]\n",
+    "        await self.collection.insert_one(doc)\n",
+    "        return {\n",
+    "            \"configurable\": {\n",
+    "                \"thread_id\": config[\"configurable\"][\"thread_id\"],\n",
+    "                \"thread_ts\": checkpoint[\"id\"],\n",
+    "            }\n",
+    "        }"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Example with basic graph"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from langgraph.graph import StateGraph, START\n",
+    "\n",
+    "checkpointer = MongoDBSaver(\n",
+    "    AsyncIOMotorClient(MONGO_URI), \"checkpoints_db\", \"checkpoints_collection\"\n",
+    ")\n",
+    "builder = StateGraph(int)\n",
+    "builder.add_node(\"add_one\", lambda x: x + 1)\n",
+    "builder.add_edge(START, \"add_one\")\n",
+    "builder.add_edge(\"add_one\", END)\n",
+    "graph = builder.compile(checkpointer=checkpointer)\n",
+    "config = {\"configurable\": {\"thread_id\": \"123\"}}\n",
+    "res = await graph.ainvoke(3, config)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "4"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "res"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'v': 1,\n",
+       " 'ts': '2024-07-10T11:34:28.485660+00:00',\n",
+       " 'id': '1ef3eb05-e0d1-651b-8004-15f129f5f4fb',\n",
+       " 'channel_values': {'__root__': 4, 'add_one': 'add_one'},\n",
+       " 'channel_versions': {'__start__': 5,\n",
+       "  '__root__': 6,\n",
+       "  'start:add_one': 6,\n",
+       "  'add_one': 6},\n",
+       " 'versions_seen': {'__start__': {'__start__': 4},\n",
+       "  'add_one': {'start:add_one': 5}},\n",
+       " 'pending_sends': []}"
+      ]
+     },
+     "execution_count": 9,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "await checkpointer.aget(config)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "CheckpointTuple(config={'configurable': {'thread_id': '123'}}, checkpoint={'v': 1, 'ts': '2024-07-10T11:34:28.485660+00:00', 'id': '1ef3eb05-e0d1-651b-8004-15f129f5f4fb', 'channel_values': {'__root__': 4, 'add_one': 'add_one'}, 'channel_versions': {'__start__': 5, '__root__': 6, 'start:add_one': 6, 'add_one': 6}, 'versions_seen': {'__start__': {'__start__': 4}, 'add_one': {'start:add_one': 5}}, 'pending_sends': []}, metadata={'source': 'loop', 'step': 4, 'writes': {'add_one': 4}}, parent_config={'configurable': {'thread_id': '123', 'thread_ts': '1ef3eb05-e0bd-6c9c-8003-aa9cb0fdedc1'}})"
+      ]
+     },
+     "execution_count": 10,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "await checkpointer.aget_tuple(config)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "CheckpointTuple(config={'configurable': {'thread_id': '123', 'thread_ts': '1ef3eb05-e0d1-651b-8004-15f129f5f4fb'}}, checkpoint={'v': 1, 'ts': '2024-07-10T11:34:28.485660+00:00', 'id': '1ef3eb05-e0d1-651b-8004-15f129f5f4fb', 'channel_values': {'__root__': 4, 'add_one': 'add_one'}, 'channel_versions': {'__start__': 5, '__root__': 6, 'start:add_one': 6, 'add_one': 6}, 'versions_seen': {'__start__': {'__start__': 4}, 'add_one': {'start:add_one': 5}}, 'pending_sends': []}, metadata={'source': 'loop', 'step': 4, 'writes': {'add_one': 4}}, parent_config={'configurable': {'thread_id': '123', 'thread_ts': '1ef3eb05-e0bd-6c9c-8003-aa9cb0fdedc1'}})\n",
+      "CheckpointTuple(config={'configurable': {'thread_id': '123', 'thread_ts': '1ef3eb05-e0bd-6c9c-8003-aa9cb0fdedc1'}}, checkpoint={'v': 1, 'ts': '2024-07-10T11:34:28.477660+00:00', 'id': '1ef3eb05-e0bd-6c9c-8003-aa9cb0fdedc1', 'channel_values': {'__root__': 3, 'start:add_one': '__start__'}, 'channel_versions': {'__start__': 5, '__root__': 5, 'start:add_one': 5, 'add_one': 4}, 'versions_seen': {'__start__': {'__start__': 4}, 'add_one': {'start:add_one': 2}}, 'pending_sends': []}, metadata={'source': 'loop', 'step': 3, 'writes': None}, parent_config={'configurable': {'thread_id': '123', 'thread_ts': '1ef3eb05-e0bb-659e-8002-de83b4764141'}})\n",
+      "CheckpointTuple(config={'configurable': {'thread_id': '123', 'thread_ts': '1ef3eb05-e0bb-659e-8002-de83b4764141'}}, checkpoint={'v': 1, 'ts': '2024-07-10T11:34:28.476662+00:00', 'id': '1ef3eb05-e0bb-659e-8002-de83b4764141', 'channel_values': {'__root__': 4, '__start__': 3}, 'channel_versions': {'__start__': 4, '__root__': 3, 'start:add_one': 3, 'add_one': 4}, 'versions_seen': {'__start__': {'__start__': 1}, 'add_one': {'start:add_one': 2}}, 'pending_sends': []}, metadata={'source': 'input', 'step': 2, 'writes': 3}, parent_config=None)\n"
+     ]
+    }
+   ],
+   "source": [
+    "list = checkpointer.alist(config, limit=3)\n",
+    "async for item in list:\n",
+    "    print(item)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Checkpoints saved in MongoDB"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'_id': ObjectId('668e57930f55bbe62f358531'), 'thread_id': '123', 'thread_ts': '1ef3ea0c-18a5-67a6-bfff-0d85b77e4a09', 'checkpoint': b'{\"v\": 1, \"ts\": \"2024-07-10T09:42:43.453328+00:00\", \"id\": \"1ef3ea0c-18a5-67a6-bfff-0d85b77e4a09\", \"channel_values\": {\"__start__\": 3}, \"channel_versions\": {\"__start__\": 1}, \"versions_seen\": {}, \"pending_sends\": []}', 'metadata': b'{\"source\": \"input\", \"step\": -1, \"writes\": 3}'}\n",
+      "{'_id': ObjectId('668e57930f55bbe62f358532'), 'thread_id': '123', 'thread_ts': '1ef3ea0c-18a7-6ea3-8000-9a52ba553d0c', 'checkpoint': b'{\"v\": 1, \"ts\": \"2024-07-10T09:42:43.454326+00:00\", \"id\": \"1ef3ea0c-18a7-6ea3-8000-9a52ba553d0c\", \"channel_values\": {\"__root__\": 3, \"start:add_one\": \"__start__\"}, \"channel_versions\": {\"__start__\": 2, \"__root__\": 2, \"start:add_one\": 2}, \"versions_seen\": {\"__start__\": {\"__start__\": 1}, \"add_one\": {}}, \"pending_sends\": []}', 'metadata': b'{\"source\": \"loop\", \"step\": 0, \"writes\": null}', 'parent_ts': '1ef3ea0c-18a5-67a6-bfff-0d85b77e4a09'}\n",
+      "{'_id': ObjectId('668e57930f55bbe62f358533'), 'thread_id': '123', 'thread_ts': '1ef3ea0c-18bc-6b54-8001-ef5781939492', 'checkpoint': b'{\"v\": 1, \"ts\": \"2024-07-10T09:42:43.462843+00:00\", \"id\": \"1ef3ea0c-18bc-6b54-8001-ef5781939492\", \"channel_values\": {\"__root__\": 4, \"add_one\": \"add_one\"}, \"channel_versions\": {\"__start__\": 2, \"__root__\": 3, \"start:add_one\": 3, \"add_one\": 3}, \"versions_seen\": {\"__start__\": {\"__start__\": 1}, \"add_one\": {\"start:add_one\": 2}}, \"pending_sends\": []}', 'metadata': b'{\"source\": \"loop\", \"step\": 1, \"writes\": {\"add_one\": 4}}', 'parent_ts': '1ef3ea0c-18a7-6ea3-8000-9a52ba553d0c'}\n",
+      "{'_id': ObjectId('668e71c4171972a41a226373'), 'thread_id': '123', 'thread_ts': '1ef3eb05-e0bb-659e-8002-de83b4764141', 'checkpoint': b'{\"v\": 1, \"ts\": \"2024-07-10T11:34:28.476662+00:00\", \"id\": \"1ef3eb05-e0bb-659e-8002-de83b4764141\", \"channel_values\": {\"__root__\": 4, \"__start__\": 3}, \"channel_versions\": {\"__start__\": 4, \"__root__\": 3, \"start:add_one\": 3, \"add_one\": 4}, \"versions_seen\": {\"__start__\": {\"__start__\": 1}, \"add_one\": {\"start:add_one\": 2}}, \"pending_sends\": []}', 'metadata': b'{\"source\": \"input\", \"step\": 2, \"writes\": 3}'}\n",
+      "{'_id': ObjectId('668e71c4171972a41a226374'), 'thread_id': '123', 'thread_ts': '1ef3eb05-e0bd-6c9c-8003-aa9cb0fdedc1', 'checkpoint': b'{\"v\": 1, \"ts\": \"2024-07-10T11:34:28.477660+00:00\", \"id\": \"1ef3eb05-e0bd-6c9c-8003-aa9cb0fdedc1\", \"channel_values\": {\"__root__\": 3, \"start:add_one\": \"__start__\"}, \"channel_versions\": {\"__start__\": 5, \"__root__\": 5, \"start:add_one\": 5, \"add_one\": 4}, \"versions_seen\": {\"__start__\": {\"__start__\": 4}, \"add_one\": {\"start:add_one\": 2}}, \"pending_sends\": []}', 'metadata': b'{\"source\": \"loop\", \"step\": 3, \"writes\": null}', 'parent_ts': '1ef3eb05-e0bb-659e-8002-de83b4764141'}\n",
+      "{'_id': ObjectId('668e71c4171972a41a226375'), 'thread_id': '123', 'thread_ts': '1ef3eb05-e0d1-651b-8004-15f129f5f4fb', 'checkpoint': b'{\"v\": 1, \"ts\": \"2024-07-10T11:34:28.485660+00:00\", \"id\": \"1ef3eb05-e0d1-651b-8004-15f129f5f4fb\", \"channel_values\": {\"__root__\": 4, \"add_one\": \"add_one\"}, \"channel_versions\": {\"__start__\": 5, \"__root__\": 6, \"start:add_one\": 6, \"add_one\": 6}, \"versions_seen\": {\"__start__\": {\"__start__\": 4}, \"add_one\": {\"start:add_one\": 5}}, \"pending_sends\": []}', 'metadata': b'{\"source\": \"loop\", \"step\": 4, \"writes\": {\"add_one\": 4}}', 'parent_ts': '1ef3eb05-e0bd-6c9c-8003-aa9cb0fdedc1'}\n"
+     ]
+    }
+   ],
+   "source": [
+    "from pymongo import MongoClient\n",
+    "\n",
+    "client = MongoClient(MONGO_URI)\n",
+    "database = client[\"checkpoints_db\"]\n",
+    "collection = database[\"checkpoints_collection\"]\n",
+    "\n",
+    "for doc in collection.find():\n",
+    "    print(doc)"
    ]
   }
  ],

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

@@ -63,10 +63,13 @@
     "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(obj.__class__, method=\"fromhex\", args=[obj.hex()])\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",
@@ -87,17 +90,25 @@
     "            logger.error(f\"Deserialization error: {e}\")\n",
     "            raise\n",
     "\n",
-    "def initialize_sync_pool(host: str = 'localhost', port: int = 6379, db: int = 0, **kwargs) -> redis.ConnectionPool:\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(f\"Synchronous Redis pool initialized with host={host}, port={port}, db={db}\")\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",
-    "def initialize_async_pool(url: str = \"redis://localhost\", **kwargs) -> AsyncConnectionPool:\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",
@@ -107,8 +118,11 @@
     "        logger.error(f\"Error initializing async pool: {e}\")\n",
     "        raise\n",
     "\n",
+    "\n",
     "@contextmanager\n",
-    "def _get_sync_connection(connection: Union[redis.Redis, redis.ConnectionPool, None]) -> Generator[redis.Redis, None, None]:\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",
@@ -125,8 +139,11 @@
     "        if conn:\n",
     "            conn.close()\n",
     "\n",
+    "\n",
     "@asynccontextmanager\n",
-    "async def _get_async_connection(connection: Union[AsyncRedis, AsyncConnectionPool, None]) -> AsyncGenerator[AsyncRedis, None]:\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",
@@ -143,27 +160,42 @@
     "        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__(self, sync_connection: Optional[Union[redis.Redis, redis.ConnectionPool]] = None, async_connection: Optional[Union[AsyncRedis, AsyncConnectionPool]] = None):\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(self, config: RunnableConfig, checkpoint: Checkpoint, metadata: CheckpointMetadata) -> RunnableConfig:\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(key, mapping={\n",
-    "                    \"checkpoint\": self.serde.dumps(checkpoint),\n",
-    "                    \"metadata\": self.serde.dumps(metadata),\n",
-    "                    \"parent_ts\": parent_ts if parent_ts else \"\"\n",
-    "                })\n",
-    "                logger.info(f\"Checkpoint stored successfully for thread_id: {thread_id}, ts: {checkpoint['ts']}\")\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",
@@ -174,18 +206,28 @@
     "            },\n",
     "        }\n",
     "\n",
-    "    async def aput(self, config: RunnableConfig, checkpoint: Checkpoint, metadata: CheckpointMetadata) -> RunnableConfig:\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(key, mapping={\n",
-    "                    \"checkpoint\": self.serde.dumps(checkpoint),\n",
-    "                    \"metadata\": self.serde.dumps(metadata),\n",
-    "                    \"parent_ts\": parent_ts if parent_ts else \"\"\n",
-    "                })\n",
-    "                logger.info(f\"Checkpoint stored successfully for thread_id: {thread_id}, ts: {checkpoint['ts']}\")\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",
@@ -217,9 +259,20 @@
     "                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 = {\"configurable\": {\"thread_id\": thread_id, \"thread_ts\": parent_ts}} if parent_ts else None\n",
-    "                logger.info(f\"Checkpoint retrieved successfully for thread_id: {thread_id}, ts: {thread_ts}\")\n",
-    "                return CheckpointTuple(config=config, checkpoint=checkpoint, metadata=metadata, parent_config=parent_config)\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",
@@ -245,22 +298,47 @@
     "                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 = {\"configurable\": {\"thread_id\": thread_id, \"thread_ts\": parent_ts}} if parent_ts else None\n",
-    "                logger.info(f\"Checkpoint retrieved successfully for thread_id: {thread_id}, ts: {thread_ts}\")\n",
-    "                return CheckpointTuple(config=config, checkpoint=checkpoint, metadata=metadata, parent_config=parent_config)\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(self, config: Optional[RunnableConfig], *, filter: Optional[dict[str, Any]] = None, before: Optional[RunnableConfig] = None, limit: Optional[int] = None) -> Generator[CheckpointTuple, None, None]:\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 = [k for k in keys if k.decode().split(\":\")[-1] < before[\"configurable\"][\"thread_ts\"]]\n",
-    "                keys = sorted(keys, key=lambda k: k.decode().split(\":\")[-1], reverse=True)\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",
@@ -268,25 +346,53 @@
     "                    if data and \"checkpoint\" in data and \"metadata\" in data:\n",
     "                        thread_ts = key.decode().split(\":\")[-1]\n",
     "                        yield CheckpointTuple(\n",
-    "                            config={\"configurable\": {\"thread_id\": thread_id, \"thread_ts\": thread_ts}},\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={\"configurable\": {\"thread_id\": thread_id, \"thread_ts\": data.get(\"parent_ts\", b\"\").decode()}} if data.get(\"parent_ts\") else None,\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",
-    "                        logger.info(f\"Checkpoint listed for thread_id: {thread_id}, ts: {thread_ts}\")\n",
     "        except Exception as e:\n",
     "            logger.error(f\"Failed to list checkpoints: {e}\")\n",
     "            raise\n",
     "\n",
-    "    async def alist(self, config: Optional[RunnableConfig], *, filter: Optional[dict[str, Any]] = None, before: Optional[RunnableConfig] = None, limit: Optional[int] = None) -> AsyncGenerator[CheckpointTuple, None]:\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 = [k for k in keys if k.decode().split(\":\")[-1] < before[\"configurable\"][\"thread_ts\"]]\n",
-    "                keys = sorted(keys, key=lambda k: k.decode().split(\":\")[-1], reverse=True)\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",
@@ -294,15 +400,29 @@
     "                    if data and \"checkpoint\" in data and \"metadata\" in data:\n",
     "                        thread_ts = key.decode().split(\":\")[-1]\n",
     "                        yield CheckpointTuple(\n",
-    "                            config={\"configurable\": {\"thread_id\": thread_id, \"thread_ts\": thread_ts}},\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={\"configurable\": {\"thread_id\": thread_id, \"thread_ts\": data.get(\"parent_ts\", b\"\").decode()}} if data.get(\"parent_ts\") else None,\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",
-    "                        logger.info(f\"Checkpoint listed for thread_id: {thread_id}, ts: {thread_ts}\")\n",
     "        except Exception as e:\n",
     "            logger.error(f\"Failed to list checkpoints: {e}\")\n",
-    "            raise\n"
+    "            raise"
    ]
   },
   {
@@ -538,7 +658,7 @@
     "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",
+    "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",
@@ -582,7 +702,7 @@
    ],
    "source": [
     "# Initialize a synchronous Redis connection pool\n",
-    "async_pool = initialize_async_pool(url='redis://172.25.0.4:6379/0')\n",
+    "async_pool = initialize_async_pool(url=\"redis://172.25.0.4:6379/0\")\n",
     "\n",
     "checkpointer = RedisSaver(async_connection=async_pool)"
    ]
@@ -687,7 +807,7 @@
    "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",
+    "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",

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

@@ -42,7 +42,7 @@
    "metadata": {},
    "outputs": [
     {
-     "name": "stdout",
+     "name": "stdin",
      "output_type": "stream",
      "text": [
       "OPENAI_API_KEY:  ········\n"
@@ -80,7 +80,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 3,
    "id": "d59234f9-173e-469d-a725-c13e0979663e",
    "metadata": {},
    "outputs": [],
@@ -88,40 +88,36 @@
     "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 ensure_config, get_callback_manager_for_config\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\": {\n",
-    "        \"place\": {\n",
-    "          \"type\": \"string\"\n",
-    "        }\n",
-    "      },\n",
-    "      \"required\": [\n",
-    "        \"place\"\n",
-    "      ]\n",
-    "    }\n",
-    "  }\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",
+    "\n",
     "    llm_run_manager = callback_manager.on_chat_model_start({}, [messages])[0]\n",
     "    response = await openai_client.chat.completions.create(\n",
-    "        messages=messages,\n",
-    "        model=\"gpt-3.5-turbo\",\n",
-    "        tools=[tool],\n",
-    "        stream=True\n",
+    "        messages=messages, model=\"gpt-3.5-turbo\", tools=[tool], stream=True\n",
     "    )\n",
     "\n",
     "    response_content = \"\"\n",
@@ -147,7 +143,10 @@
     "\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(content=\"\", additional_kwargs={\"tool_calls\": [delta.tool_calls[0].dict()]})\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",
@@ -156,8 +155,11 @@
     "        tool_calls = [\n",
     "            {\n",
     "                \"id\": tool_call_id,\n",
-    "                \"function\": {\"name\": tool_call_function_name, \"arguments\": tool_call_function_arguments},\n",
-    "                \"type\": \"function\"\n",
+    "                \"function\": {\n",
+    "                    \"name\": tool_call_function_name,\n",
+    "                    \"arguments\": tool_call_function_arguments,\n",
+    "                },\n",
+    "                \"type\": \"function\",\n",
     "            }\n",
     "        ]\n",
     "    else:\n",
@@ -166,7 +168,7 @@
     "    response_message = {\n",
     "        \"role\": role,\n",
     "        \"content\": response_content,\n",
-    "        \"tool_calls\": tool_calls\n",
+    "        \"tool_calls\": tool_calls,\n",
     "    }\n",
     "    return {\"messages\": [response_message]}"
    ]
@@ -181,18 +183,17 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 4,
    "id": "b90941d8-afe4-42ec-9262-9c3b87c3b1ec",
    "metadata": {},
    "outputs": [],
    "source": [
     "import json\n",
-    "from langchain_core.runnables import RunnableLambda\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",
-    "    # NOTE: we need to define a special langchain runnable that we'll be using for logging the streaming outputs from within a tool\n",
-    "    tool_logger = RunnableLambda(lambda inputs: inputs).with_config({\"tags\": [\"tool_call\"]})\n",
     "\n",
     "    # this can be replaced with any actual streaming logic that you might have\n",
     "    def stream(place: str):\n",
@@ -205,15 +206,25 @@
     "\n",
     "    tokens = []\n",
     "    for token in stream(place):\n",
-    "        tool_logger.invoke(token)\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 = {\n",
-    "    \"get_items\": get_items\n",
-    "}\n",
+    "function_name_to_function = {\"get_items\": get_items}\n",
+    "\n",
     "\n",
     "async def call_tools(state):\n",
     "    messages = state[\"messages\"]\n",
@@ -222,17 +233,15 @@
     "    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",
+    "\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 {\n",
-    "        \"messages\": [tool_message]\n",
-    "    }"
+    "    return {\"messages\": [tool_message]}"
    ]
   },
   {
@@ -245,7 +254,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 5,
    "id": "228260be-1f9a-4195-80e0-9604f8a5dba6",
    "metadata": {},
    "outputs": [],
@@ -255,16 +264,19 @@
     "\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",
+    "    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",
@@ -284,10 +296,18 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "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",
@@ -299,10 +319,12 @@
     }
    ],
    "source": [
-    "async for event in graph.astream_events({\"messages\": [{\"role\": \"user\", \"content\": \"what's in the bedroom\"}]}, version=\"v2\"):\n",
+    "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_chain_end\" and \"tool_call\" in tags:\n",
-    "        print(\"Tool token\", event[\"data\"][\"output\"])"
+    "    if event[\"event\"] == \"on_custom_event\" and \"tool_call\" in tags:\n",
+    "        print(\"Tool token\", event[\"data\"][\"tool_output_token\"])"
    ]
   }
  ],

examples/streaming-tokens-without-langchain.ipynb

@@ -88,40 +88,36 @@
     "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 ensure_config, get_callback_manager_for_config\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\": {\n",
-    "        \"place\": {\n",
-    "          \"type\": \"string\"\n",
-    "        }\n",
-    "      },\n",
-    "      \"required\": [\n",
-    "        \"place\"\n",
-    "      ]\n",
-    "    }\n",
-    "  }\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",
+    "\n",
     "    llm_run_manager = callback_manager.on_chat_model_start({}, [messages])[0]\n",
     "    response = await openai_client.chat.completions.create(\n",
-    "        messages=messages,\n",
-    "        model=\"gpt-3.5-turbo\",\n",
-    "        tools=[tool],\n",
-    "        stream=True\n",
+    "        messages=messages, model=\"gpt-3.5-turbo\", tools=[tool], stream=True\n",
     "    )\n",
     "\n",
     "    response_content = \"\"\n",
@@ -147,7 +143,10 @@
     "\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(content=\"\", additional_kwargs={\"tool_calls\": [delta.tool_calls[0].dict()]})\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",
@@ -156,8 +155,11 @@
     "        tool_calls = [\n",
     "            {\n",
     "                \"id\": tool_call_id,\n",
-    "                \"function\": {\"name\": tool_call_function_name, \"arguments\": tool_call_function_arguments},\n",
-    "                \"type\": \"function\"\n",
+    "                \"function\": {\n",
+    "                    \"name\": tool_call_function_name,\n",
+    "                    \"arguments\": tool_call_function_arguments,\n",
+    "                },\n",
+    "                \"type\": \"function\",\n",
     "            }\n",
     "        ]\n",
     "    else:\n",
@@ -166,7 +168,7 @@
     "    response_message = {\n",
     "        \"role\": role,\n",
     "        \"content\": response_content,\n",
-    "        \"tool_calls\": tool_calls\n",
+    "        \"tool_calls\": tool_calls,\n",
     "    }\n",
     "    return {\"messages\": [response_message]}"
    ]
@@ -188,6 +190,7 @@
    "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",
@@ -197,10 +200,10 @@
     "    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 = {\n",
-    "    \"get_items\": get_items\n",
-    "}\n",
+    "function_name_to_function = {\"get_items\": get_items}\n",
+    "\n",
     "\n",
     "async def call_tools(state):\n",
     "    messages = state[\"messages\"]\n",
@@ -209,17 +212,15 @@
     "    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",
+    "\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 {\n",
-    "        \"messages\": [tool_message]\n",
-    "    }"
+    "    return {\"messages\": [tool_message]}"
    ]
   },
   {
@@ -242,16 +243,19 @@
     "\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",
+    "    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",
@@ -325,7 +329,9 @@
     }
    ],
    "source": [
-    "async for event in graph.astream_events({\"messages\": [{\"role\": \"user\", \"content\": \"what's in the bedroom\"}]}, version=\"v2\"):\n",
+    "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())"

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

@@ -232,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(
@@ -241,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_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/langgraph/README.md

@@ -73,8 +73,8 @@ def search(query: str):
     """Call to surf the web."""
     # This is a placeholder, but don't tell the LLM that...
     if "sf" in query.lower() or "san francisco" in query.lower():
-        return ["It's 60 degrees and foggy."]
-    return ["It's 90 degrees and sunny."]
+        return "It's 60 degrees and foggy."
+    return "It's 90 degrees and sunny."
 
 
 tools = [search]

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,14 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]):
         except Exception:
             pass
 
+    def __eq__(self, value: object) -> bool:
+        return isinstance(value, BinaryOperatorAggregate) and (
+            value.operator is self.operator
+            if value.operator.__name__ != "<lambda>"
+            and self.operator.__name__ != "<lambda>"
+            else True
+        )
+
     @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

@@ -1,6 +1,6 @@
 from contextlib import AsyncExitStack, ExitStack, asynccontextmanager, contextmanager
 from datetime import datetime, timezone
-from typing import Any, AsyncGenerator, Generator, Mapping
+from typing import Any, AsyncGenerator, Generator, Mapping, Optional
 
 from langchain_core.runnables import RunnableConfig
 
@@ -43,7 +43,11 @@ async def AsyncChannelsManager(
 
 
 def create_checkpoint(
-    checkpoint: Checkpoint, channels: Mapping[str, BaseChannel], step: int
+    checkpoint: Checkpoint,
+    channels: Mapping[str, BaseChannel],
+    step: int,
+    *,
+    id: Optional[str] = None,
 ) -> Checkpoint:
     """Create a checkpoint for the given channels."""
     ts = datetime.now(timezone.utc).isoformat()
@@ -56,9 +60,11 @@ def create_checkpoint(
     return Checkpoint(
         v=1,
         ts=ts,
-        id=str(uuid6(clock_seq=step)),
+        id=id or str(uuid6(clock_seq=step)),
         channel_values=values,
         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

@@ -46,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:
@@ -72,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
@@ -95,8 +99,8 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
         >>> asyncio.run(coro)
         Output: 2
         ```
-
         Raw usage:
+
         ```pycon
         >>> import asyncio
         >>> import aiosqlite
@@ -309,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)
@@ -356,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.
@@ -385,6 +391,15 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
         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 (?, ?, ?, ?, ?, ?)",

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,
@@ -25,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.
@@ -53,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."""
 
@@ -74,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.
@@ -84,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:
@@ -93,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={},
     )
 
 
@@ -105,20 +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[Tuple[str, str, Any]]] = None
+    pending_writes: Optional[List[PendingWrite]] = None
 
 
 CheckpointThreadId = ConfigurableFieldSpec(
@@ -141,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__(
@@ -152,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(
@@ -169,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(
@@ -177,6 +237,19 @@ 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(
@@ -185,18 +258,47 @@ class BaseCheckpointSaver(ABC):
         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 checkpointer to implement it."
+            "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],
         *,
@@ -204,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
 
@@ -213,6 +329,19 @@ 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(
@@ -221,11 +350,31 @@ class BaseCheckpointSaver(ABC):
         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 checkpointer to implement it."
+            "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

@@ -44,7 +44,7 @@ class MemorySaver(BaseCheckpointSaver):
             asyncio.run(coro)  # Output: 2
     """
 
-    storage: defaultdict[str, dict[str, tuple[bytes, bytes]]]
+    storage: defaultdict[str, dict[str, tuple[bytes, bytes, Optional[str]]]]
 
     def __init__(
         self,
@@ -72,7 +72,7 @@ class MemorySaver(BaseCheckpointSaver):
         thread_id = config["configurable"]["thread_id"]
         if ts := config["configurable"].get("thread_ts"):
             if saved := self.storage[thread_id].get(ts):
-                checkpoint, metadata = saved
+                checkpoint, metadata, parent_ts = saved
                 writes = self.writes[(thread_id, ts)]
                 return CheckpointTuple(
                     config=config,
@@ -81,11 +81,19 @@ class MemorySaver(BaseCheckpointSaver):
                     pending_writes=[
                         (id, c, self.serde.loads(v)) for id, c, v in writes
                     ],
+                    parent_config={
+                        "configurable": {
+                            "thread_id": thread_id,
+                            "thread_ts": parent_ts,
+                        }
+                    }
+                    if parent_ts
+                    else None,
                 )
         else:
             if checkpoints := self.storage[thread_id]:
                 ts = max(checkpoints.keys())
-                checkpoint, metadata = checkpoints[ts]
+                checkpoint, metadata, parent_ts = checkpoints[ts]
                 writes = self.writes[(thread_id, ts)]
                 return CheckpointTuple(
                     config={"configurable": {"thread_id": thread_id, "thread_ts": ts}},
@@ -94,6 +102,14 @@ class MemorySaver(BaseCheckpointSaver):
                     pending_writes=[
                         (id, c, self.serde.loads(v)) for id, c, v in writes
                     ],
+                    parent_config={
+                        "configurable": {
+                            "thread_id": thread_id,
+                            "thread_ts": parent_ts,
+                        }
+                    }
+                    if parent_ts
+                    else None,
                 )
 
     def list(
@@ -107,19 +123,22 @@ 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:
-            for ts, (checkpoint, metadata_b) in self.storage[thread_id].items():
+            for ts, (checkpoint, metadata_b, parent_ts) in sorted(
+                self.storage[thread_id].items(), key=lambda x: x[0], reverse=True
+            ):
                 # filter by thread_ts
                 if before and ts >= before["configurable"]["thread_ts"]:
                     continue
@@ -142,6 +161,14 @@ class MemorySaver(BaseCheckpointSaver):
                     config={"configurable": {"thread_id": thread_id, "thread_ts": ts}},
                     checkpoint=self.serde.loads(checkpoint),
                     metadata=metadata,
+                    parent_config={
+                        "configurable": {
+                            "thread_id": thread_id,
+                            "thread_ts": parent_ts,
+                        }
+                    }
+                    if parent_ts
+                    else None,
                 )
 
     def put(
@@ -158,6 +185,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.
@@ -167,6 +195,7 @@ class MemorySaver(BaseCheckpointSaver):
                 checkpoint["id"]: (
                     self.serde.dumps(checkpoint),
                     self.serde.dumps(metadata),
+                    config["configurable"].get("thread_ts"),  # parent
                 )
             }
         )
@@ -191,6 +220,7 @@ class MemorySaver(BaseCheckpointSaver):
         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.
@@ -254,6 +284,16 @@ 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
         )
@@ -264,6 +304,16 @@ class MemorySaver(BaseCheckpointSaver):
         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

@@ -309,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.
 
@@ -410,6 +411,15 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
         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 (?, ?, ?, ?, ?, ?)",
@@ -467,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,21 @@
 from typing import Any
 
+INPUT = "__input__"
 CONFIG_KEY_SEND = "__pregel_send"
 CONFIG_KEY_READ = "__pregel_read"
+CONFIG_KEY_CHECKPOINTER = "__pregel_checkpointer"
+CONFIG_KEY_RESUMING = "__pregel_resuming"
 INTERRUPT = "__interrupt__"
 TASKS = "__pregel_tasks"
-RESERVED = {INTERRUPT, TASKS, CONFIG_KEY_SEND, CONFIG_KEY_READ}
+RESERVED = {
+    INTERRUPT,
+    TASKS,
+    CONFIG_KEY_SEND,
+    CONFIG_KEY_READ,
+    CONFIG_KEY_CHECKPOINTER,
+    CONFIG_KEY_RESUMING,
+    INPUT,
+}
 TAG_HIDDEN = "langsmith:hidden"
 
 START = "__start__"

libs/langgraph/langgraph/errors.py

@@ -28,3 +28,15 @@ class InvalidUpdateError(Exception):
     """Raised when attempting to update a channel with an invalid sequence of updates."""
 
     pass
+
+
+class GraphInterrupt(Exception):
+    """Raised when a subgraph is interrupted."""
+
+    pass
+
+
+class EmptyInputError(Exception):
+    """Raised when graph receives an empty input."""
+
+    pass

libs/langgraph/langgraph/graph/graph.py

@@ -217,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(
@@ -465,11 +469,19 @@ class CompiledGraph(Pregel):
         start_nodes: dict[str, DrawableNode] = {
             START: graph.add_node(self.get_input_schema(config), START)
         }
-        end_nodes: dict[str, DrawableNode] = {
-            END: graph.add_node(self.get_output_schema(config), END)
-        }
+        end_nodes: dict[str, DrawableNode] = {}
 
-        for key, (node, metadata) 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(
@@ -490,11 +502,11 @@ class CompiledGraph(Pregel):
                     start_nodes[key] = n
                     end_nodes[key] = n
             else:
-                n = graph.add_node(node, key, metadata=metadata)
+                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},
@@ -508,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/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
@@ -34,14 +38,13 @@ from langgraph.graph.graph import (
     Branch,
     CompiledGraph,
     Graph,
-    NodeSpec,
     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__)
 
@@ -58,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>.
@@ -109,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
@@ -127,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.
 
@@ -144,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:
@@ -160,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.
 
@@ -169,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.
 
@@ -213,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.
@@ -243,8 +358,8 @@ class StateGraph(Graph):
                 raise ValueError("END cannot be a start node")
             if start not in self.nodes:
                 raise ValueError(f"Need to add_node `{start}` first")
-        if end_key == END:
-            raise ValueError("END cannot be an end node")
+        if end_key == START:
+            raise ValueError("START cannot be an end node")
         if end_key not in self.nodes:
             raise ValueError(f"Need to add_node `{end_key}` first")
 
@@ -256,7 +371,7 @@ class StateGraph(Graph):
         interrupt_before: Optional[Union[All, Sequence[str]]] = None,
         interrupt_after: Optional[Union[All, Sequence[str]]] = None,
         debug: bool = False,
-    ) -> CompiledGraph:
+    ) -> "CompiledStateGraph":
         """Compiles the state graph into a `CompiledGraph` object.
 
         The compiled graph implements the `Runnable` interface and can be invoked,
@@ -271,7 +386,7 @@ class StateGraph(Graph):
             debug (bool): A flag indicating whether to enable debug mode.
 
         Returns:
-            CompiledGraph: The compiled state graph.
+            CompiledStateGraph: The compiled state graph.
         """
         # assign default values
         interrupt_before = interrupt_before or []
@@ -287,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)
             ]
         )
 
@@ -303,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,
@@ -338,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[NodeSpec]) -> 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):
@@ -355,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,
@@ -365,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
             ]
         )
 
@@ -377,37 +526,39 @@ 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,
                     ),
                 ],
                 metadata=node.metadata,
+                retry_policy=node.retry_policy,
             ).pipe(node.runnable)
 
     def attach_edge(self, starts: Union[str, Sequence[str]], end: str) -> None:

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,15 +1,33 @@
 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
 
+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."
+
 
 def str_output(output: Any) -> str:
     if isinstance(output, str):
@@ -22,9 +40,10 @@ def str_output(output: Any) -> str:
 
 
 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 +84,59 @@ 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"}}
+            tool_message: ToolMessage = self.tools_by_name[call["name"]].invoke(
+                input, config
+            )
+            # TODO: handle this properly in core
+            tool_message.content = str_output(tool_message.content)
+            return tool_message
+        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"}}
+            tool_message: ToolMessage = await self.tools_by_name[call["name"]].ainvoke(
+                input, config
+            )
+            # TODO: handle this properly in core
+            tool_message.content = str_output(tool_message.content)
+            return tool_message
+        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 +149,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 +254,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,447 @@
+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 BaseCheckpointSaver, Checkpoint, copy_checkpoint
+from langgraph.constants import (
+    CONFIG_KEY_CHECKPOINTER,
+    CONFIG_KEY_READ,
+    CONFIG_KEY_RESUMING,
+    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],
+    is_resuming: bool = False,
+    checkpointer: Literal[None] = None,
+    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],
+    is_resuming: bool,
+    checkpointer: Optional[BaseCheckpointSaver],
+    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,
+    is_resuming: bool = False,
+    checkpointer: Optional[BaseCheckpointSaver] = None,
+    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,
+                                ),
+                                # in Send we can't checkpoint nested graphs
+                                # as they could be running in parallel
+                            },
+                        ),
+                        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)))
+                    if parent_thread_id := config.get("configurable", {}).get(
+                        "thread_id"
+                    ):
+                        thread_id: Optional[str] = f"{parent_thread_id}-{name}"
+                    else:
+                        thread_id = None
+                    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,
+                                    ),
+                                    CONFIG_KEY_CHECKPOINTER: checkpointer,
+                                    CONFIG_KEY_RESUMING: is_resuming,
+                                    "thread_id": thread_id,
+                                    "thread_ts": checkpoint["id"],
+                                },
+                            ),
+                            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,18 +66,19 @@ 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
 
+        metadata = config["metadata"].copy()
+        metadata.pop("thread_ts", None)
+
         yield {
             "type": "task",
             "timestamp": ts,
             "step": step,
             "payload": {
-                "id": str(
-                    uuid5(TASK_NAMESPACE, json.dumps((name, step, config["metadata"])))
-                ),
+                "id": str(uuid5(TASK_NAMESPACE, json.dumps((name, step, metadata)))),
                 "name": name,
                 "input": input,
                 "triggers": triggers,
@@ -91,18 +92,19 @@ 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
 
+        metadata = config["metadata"].copy()
+        metadata.pop("thread_ts", None)
+
         yield {
             "type": "task_result",
             "timestamp": ts,
             "step": step,
             "payload": {
-                "id": str(
-                    uuid5(TASK_NAMESPACE, json.dumps((name, step, config["metadata"])))
-                ),
+                "id": str(uuid5(TASK_NAMESPACE, json.dumps((name, step, metadata)))),
                 "name": name,
                 "result": [w for w in writes if w[0] in stream_channels_list],
             },
@@ -134,11 +136,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 +155,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,168 @@
+import asyncio
+import concurrent.futures
+import sys
+from contextlib import ExitStack
+from contextvars import copy_context
+from types import TracebackType
+from typing import (
+    AsyncContextManager,
+    Awaitable,
+    Callable,
+    ContextManager,
+    Optional,
+    Protocol,
+    TypeVar,
+)
+
+from langchain_core.runnables import RunnableConfig
+from langchain_core.runnables.config import get_executor_for_config
+from typing_extensions import ParamSpec
+
+from langgraph.errors import GraphInterrupt
+
+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]:
+        ...
+
+
+class BackgroundExecutor(ContextManager):
+    def __init__(self, config: RunnableConfig) -> None:
+        self.stack = ExitStack()
+        self.executor = self.stack.enter_context(get_executor_for_config(config))
+        self.tasks: dict[concurrent.futures.Future, bool] = {}
+
+    def submit(
+        self,
+        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[T]:
+        task = self.executor.submit(fn, *args, **kwargs)
+        self.tasks[task] = __cancel_on_exit__
+        task.add_done_callback(self.done)
+        return task
+
+    def done(self, task: concurrent.futures.Future) -> None:
+        try:
+            task.result()
+        except GraphInterrupt:
+            # This exception is an interruption signal, not an error
+            # so we don't want to re-raise it on exit
+            self.tasks.pop(task)
+        except BaseException:
+            pass
+        else:
+            self.tasks.pop(task)
+
+    def __enter__(self) -> "submit":
+        return self.submit
+
+    def __exit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> Optional[bool]:
+        # cancel all tasks that should be cancelled
+        for task, cancel in self.tasks.items():
+            if cancel:
+                task.cancel()
+        # wait for all tasks to finish
+        if tasks := {t for t in self.tasks if not t.done()}:
+            concurrent.futures.wait(tasks)
+        # shutdown the executor
+        self.stack.__exit__(exc_type, exc_value, traceback)
+        # re-raise the first exception that occurred in a task
+        if exc_type is None:
+            # if there's already an exception being raised, don't raise another one
+            for task in self.tasks:
+                try:
+                    task.result()
+                except concurrent.futures.CancelledError:
+                    pass
+
+
+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 GraphInterrupt:
+            # This exception is an interruption signal, not an error
+            # so we don't want to re-raise it on exit
+            self.tasks.pop(task)
+        except BaseException:
+            pass
+        else:
+            self.tasks.pop(task)
+
+    async def __aenter__(self) -> Submit:
+        return self.submit
+
+    async def exit(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> None:
+        # cancel all tasks that should be cancelled
+        for task, cancel in self.tasks.items():
+            if cancel:
+                task.cancel(self.sentinel)
+        # wait for all tasks to finish
+        if self.tasks:
+            await asyncio.wait(self.tasks)
+        # re-raise the first exception that occurred in a task
+        if exc_type is None:
+            # if there's already an exception being raised, don't raise another one
+            for task in self.tasks:
+                try:
+                    task.result()
+                except asyncio.CancelledError:
+                    pass
+
+    async def __aexit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> Optional[bool]:
+        # we cannot use `await` outside of asyncio.shield, as this code can run
+        # after owning task is cancelled, so pulling async logic to separate method
+
+        # wait for all background tasks to finish, shielded from cancellation
+        await asyncio.shield(self.exit(exc_type, exc_value, traceback))

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,513 @@
+import asyncio
+import concurrent.futures
+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 CONFIG_KEY_READ, CONFIG_KEY_RESUMING, INPUT, INTERRUPT
+from langgraph.errors import EmptyInputError, GraphInterrupt
+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()
+INPUT_RESUMING = object()
+EMPTY_SEQ = ()
+
+
+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_after_previous: Optional[
+        Callable[
+            [
+                Optional[concurrent.futures.Future],
+                RunnableConfig,
+                Sequence[tuple[str, Any]],
+                str,
+            ],
+            Any,
+        ]
+    ]
+    graph: "Pregel"
+
+    submit: Submit
+    channels: Mapping[str, BaseChannel]
+    managed: ManagedValueMapping
+    checkpoint: Checkpoint
+    checkpoint_config: RunnableConfig
+    checkpoint_metadata: CheckpointMetadata
+    checkpoint_pending_writes: List[PendingWrite]
+
+    step: int
+    status: Literal[
+        "pending", "done", "interrupt_before", "interrupt_after", "out_of_steps"
+    ]
+    tasks: Sequence[PregelExecutableTask]
+    stream: deque[Tuple[str, Any]]
+    is_nested: bool
+
+    # public
+
+    def __init__(
+        self,
+        input: Optional[Any],
+        *,
+        config: RunnableConfig,
+        checkpointer: Optional[BaseCheckpointSaver],
+        graph: "Pregel",
+    ) -> None:
+        self.stream = deque()
+        self.input = input
+        self.config = config
+        self.checkpointer = checkpointer
+        self.graph = graph
+        # TODO if managed values no longer needs graph we can replace with
+        # managed_specs, channel_specs
+        self.is_nested = CONFIG_KEY_READ in self.config.get("configurable", {})
+
+    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]] = EMPTY_SEQ,
+        interrupt_after: Sequence[str] = EMPTY_SEQ,
+        interrupt_before: Sequence[str] = EMPTY_SEQ,
+        manager: Union[None, AsyncParentRunManager, ParentRunManager] = None,
+    ) -> bool:
+        """Execute a single iteration of the Pregel loop.
+        Returns True if more iterations are needed."""
+
+        if self.status != "pending":
+            raise RuntimeError("Cannot tick when status is no longer 'pending'")
+
+        if self.input not in (INPUT_DONE, INPUT_RESUMING):
+            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"
+                if self.is_nested:
+                    raise GraphInterrupt(self)
+                else:
+                    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,
+            checkpointer=self.checkpointer,
+            is_resuming=self.input is INPUT_RESUMING,
+        )
+
+        # 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(
+                output_keys=output_keys,
+                interrupt_after=interrupt_after,
+                interrupt_before=interrupt_before,
+                manager=manager,
+            )
+
+        # before execution, check if we should interrupt
+        if should_interrupt(self.checkpoint, interrupt_before, self.tasks):
+            self.status = "interrupt_before"
+            if self.is_nested:
+                raise GraphInterrupt()
+            else:
+                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:
+        # resuming from previous checkpoint requires
+        # - finding a previous checkpoint
+        # - receiving None input (outer graph) or RESUMING flag (subgraph)
+        is_resuming = bool(self.checkpoint["channel_versions"]) and bool(
+            self.config.get("configurable", {}).get(CONFIG_KEY_RESUMING)
+            or self.input is None
+        )
+
+        # proceed past previous checkpoint
+        if is_resuming:
+            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
+        # map inputs to channel updates
+        elif 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,
+                manager=None,
+            )
+            # 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:
+            raise EmptyInputError(f"Received no input for {self.graph.input_channels}")
+        # done with input
+        self.input = INPUT_RESUMING if is_resuming else INPUT_DONE
+
+    def _put_checkpoint(self, metadata: CheckpointMetadata) -> None:
+        # assign step
+        metadata["step"] = self.step
+        # bail if no checkpointer
+        if self._checkpointer_put_after_previous is not None:
+            # create new checkpoint
+            self.checkpoint_metadata = metadata
+            self.checkpoint = create_checkpoint(
+                self.checkpoint,
+                self.channels,
+                self.step,
+                # child graphs keep at most one checkpoint per parent checkpoint
+                # this is achieved by writing child checkpoints as progress is made
+                # (so that error recovery / resuming from interrupt don't lose work)
+                # but doing so always with an id equal to that of the parent checkpoint
+                id=self.config["configurable"]["thread_ts"] if self.is_nested else None,
+            )
+            # save it, without blocking
+            # if there's a previous checkpoint save in progress, wait for it
+            # ensuring checkpointers receive checkpoints in order
+            self._put_checkpoint_fut = self.submit(
+                self._checkpointer_put_after_previous,
+                getattr(self, "_put_checkpoint_fut", None),
+                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
+
+    def _suppress_interrupt(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> Optional[bool]:
+        if exc_type is GraphInterrupt and not self.is_nested:
+            return True
+
+
+class SyncPregelLoop(PregelLoop, ContextManager):
+    def __init__(
+        self,
+        input: Optional[Any],
+        *,
+        config: RunnableConfig,
+        checkpointer: Optional[BaseCheckpointSaver],
+        graph: "Pregel",
+    ) -> None:
+        super().__init__(input, config=config, checkpointer=checkpointer, graph=graph)
+        self.stack = ExitStack()
+        self.stack.push(self._suppress_interrupt)
+        if checkpointer:
+            self.checkpointer_get_next_version = checkpointer.get_next_version
+            self.checkpointer_put_writes = checkpointer.put_writes
+        else:
+            self.checkpointer_get_next_version = increment
+            self._checkpointer_put_after_previous = None
+            self.checkpointer_put_writes = None
+
+    def _checkpointer_put_after_previous(
+        self,
+        prev: Optional[concurrent.futures.Future],
+        config: RunnableConfig,
+        checkpoint: Checkpoint,
+        metadata: CheckpointMetadata,
+    ) -> RunnableConfig:
+        try:
+            if prev is not None:
+                prev.result()
+        finally:
+            self.checkpointer.put(config, checkpoint, metadata)
+
+    # 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 or []
+
+        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]:
+        # unwind stack
+        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:
+        super().__init__(input, config=config, checkpointer=checkpointer, graph=graph)
+        self.stack = AsyncExitStack()
+        self.stack.push(self._suppress_interrupt)
+        if checkpointer:
+            self.checkpointer_get_next_version = checkpointer.get_next_version
+            self.checkpointer_put_writes = checkpointer.aput_writes
+        else:
+            self.checkpointer_get_next_version = increment
+            self._checkpointer_put_after_previous = None
+            self.checkpointer_put_writes = None
+
+    async def _checkpointer_put_after_previous(
+        self,
+        prev: Optional[asyncio.Task],
+        config: RunnableConfig,
+        checkpoint: Checkpoint,
+        metadata: CheckpointMetadata,
+    ) -> RunnableConfig:
+        try:
+            if prev is not None:
+                await prev
+        finally:
+            await self.checkpointer.aput(config, checkpoint, metadata)
+
+    # 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 or []
+
+        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]:
+        # unwind stack
+        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()
@@ -155,6 +158,7 @@ class PregelNode(RunnableBindingBase):
         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__(
@@ -164,6 +168,7 @@ class PregelNode(RunnableBindingBase):
             writers=writers or [],
             bound=bound or DEFAULT_BOUND,
             kwargs=kwargs or {},
+            retry_policy=retry_policy,
             config=merge_configs(
                 config, {"tags": tags or [], "metadata": metadata or {}}
             ),
@@ -180,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__(
@@ -202,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,7 @@ class PregelExecutableTask(NamedTuple):
     writes: deque[tuple[str, Any]]
     config: RunnableConfig
     triggers: list[str]
+    retry_policy: Optional[RetryPolicy]
     id: str
 
 
@@ -37,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/poetry.lock

@@ -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.15"
+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.15-py3-none-any.whl", hash = "sha256:3bf7afaef96d7c1af0d9d223833bdee5fafc46755dc10f9c7576a85d4f6c5240"},
-    {file = "langchain_core-0.2.15.tar.gz", hash = "sha256:ce03ab0a5c45b4ebfe5475eb07bf081cd21218421ff4cf26b8d2e5573ae2bd42"},
+    {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]
@@ -2769,6 +2783,20 @@ pytest = ">=6.2.5"
 [package.extras]
 dev = ["pre-commit", "pytest-asyncio", "tox"]
 
+[[package]]
+name = "pytest-repeat"
+version = "0.9.3"
+description = "pytest plugin for repeating tests"
+optional = false
+python-versions = ">=3.7"
+files = [
+    {file = "pytest_repeat-0.9.3-py3-none-any.whl", hash = "sha256:26ab2df18226af9d5ce441c858f273121e92ff55f5bb311d25755b8d7abdd8ed"},
+    {file = "pytest_repeat-0.9.3.tar.gz", hash = "sha256:ffd3836dfcd67bb270bec648b330e20be37d2966448c4148c4092d1e8aba8185"},
+]
+
+[package.dependencies]
+pytest = "*"
+
 [[package]]
 name = "pytest-watcher"
 version = "0.4.2"
@@ -2784,6 +2812,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 +4179,4 @@ test = ["big-O", "importlib-resources", "jaraco.functools", "jaraco.itertools",
 [metadata]
 lock-version = "2.0"
 python-versions = ">=3.9.0,<4.0"
-content-hash = "19250230952cb11ee6b5c820ae0a2b589ce520a59e0723204dcee0c46c3b739e"
+content-hash = "18b26895b05f2f7cdcd08d59ba164ac788e0e8005e3ec5d7089469b4f3a96aea"

libs/langgraph/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph"
-version = "0.1.8"
+version = "0.1.12"
 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.15,<0.3"
+langchain-core = ">=0.2.22,<0.3"
 
 
 [tool.poetry.group.dev.dependencies]
@@ -31,6 +31,8 @@ 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"}
+pytest-repeat = "^0.9.3"
 
 [tool.poetry.group.dev]
 optional = true
@@ -61,7 +63,7 @@ omit = ["tests/*"]
 [tool.pytest-watcher]
 now = true
 delay = 0.1
-runner_args = ["-x", "--ff", "-vv", "--snapshot-update"]
+runner_args = ["--ff", "-v", "-n", "auto", "--dist", "worksteal", "--snapshot-update", "--tb", "short"]
 patterns = ["*.py"]
 
 [build-system]

libs/langgraph/tests/memory_assert.py

@@ -30,9 +30,11 @@ class MemorySaverAssertImmutable(MemorySaver):
         self,
         *,
         serde: Optional[SerializerProtocol] = None,
+        put_sleep: Optional[float] = None,
     ) -> None:
         super().__init__(serde=serde)
         self.storage_for_copies = defaultdict(dict)
+        self.put_sleep = put_sleep
 
     def put(
         self,
@@ -40,6 +42,10 @@ class MemorySaverAssertImmutable(MemorySaver):
         checkpoint: Checkpoint,
         metadata: Optional[CheckpointMetadata] = None,
     ) -> None:
+        if self.put_sleep:
+            import time
+
+            time.sleep(self.put_sleep)
         # assert checkpoint hasn't been modified since last written
         thread_id = config["configurable"]["thread_id"]
         if saved := super().get(config):
@@ -85,7 +91,7 @@ class MemorySaverAssertCheckpointMetadata(MemorySaver):
         configurable = config["configurable"].copy()
 
         # remove thread_ts to make testing simpler
-        configurable.pop("thread_ts", None)
+        thread_ts = configurable.pop("thread_ts", None)
 
         self.storage[config["configurable"]["thread_id"]].update(
             {
@@ -93,6 +99,7 @@ class MemorySaverAssertCheckpointMetadata(MemorySaver):
                     self.serde.dumps(checkpoint),
                     # merge configurable fields and metadata
                     self.serde.dumps({**configurable, **metadata}),
+                    thread_ts,
                 )
             }
         )

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_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,11 +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.prebuilt import (
-    ToolNode,
-    ValidationNode,
-    create_react_agent,
-)
+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):
@@ -56,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
 
 
@@ -123,6 +267,13 @@ async def test_tool_node():
             raise ValueError("Test error")
         return f"tool2: {some_val} - {some_other_val}"
 
+    async def tool3(some_val: int, some_other_val: str) -> str:
+        """Tool 3 docstring."""
+        return [
+            {"key_1": some_val, "key_2": "foo"},
+            {"key_1": some_other_val, "key_2": "baz"},
+        ]
+
     result = ToolNode([tool1]).invoke(
         {
             "messages": [
@@ -208,6 +359,56 @@ 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"
+
+    # list of dicts tool content
+    result3 = await ToolNode([tool3]).ainvoke(
+        {
+            "messages": [
+                AIMessage(
+                    "hi?",
+                    tool_calls=[
+                        {
+                            "name": "tool3",
+                            "args": {"some_val": 2, "some_other_val": "bar"},
+                            "id": "some 0",
+                        }
+                    ],
+                )
+            ]
+        }
+    )
+    tool_message: ToolMessage = result3["messages"][-1]
+    assert tool_message.type == "tool"
+    assert (
+        tool_message.content
+        == '[{"key_1": 2, "key_2": "foo"}, {"key_1": "bar", "key_2": "baz"}]'
+    )
+    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}"
@@ -277,3 +478,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/sdk-js/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@langchain/langgraph-sdk",
-  "version": "0.0.1",
+  "version": "0.0.2",
   "description": "Client library for interacting with the LangGraph API",
   "type": "module",
   "packageManager": "yarn@1.22.19",

libs/sdk-js/scripts/create-entrypoints.js

@@ -38,9 +38,19 @@ const updateConfig = () => {
     ...json,
     typedocOptions: {
       ...json.typedocOptions,
-      entryPoints: [...Object.keys(entrypoints)].map(
-        (key) => `src/${entrypoints[key]}.ts`,
-      ),
+      entryPoints: [...Object.keys(entrypoints)].map((key) => {
+        const basePath = `src/${entrypoints[key]}`;
+        if (fs.existsSync(`${basePath}.mts`)) {
+          return `${basePath}.mts`;
+        } else if (fs.existsSync(`${basePath}.ts`)) {
+          return `${basePath}.ts`;
+        } else {
+          console.warn(
+            `Warning: Neither ${basePath}.mts nor ${basePath}.ts found for entrypoint ${key}`,
+          );
+          return `${basePath}.ts`; // Default to .ts if neither exists
+        }
+      }),
     },
   }));
 

libs/sdk-js/src/client.mts

@@ -19,6 +19,7 @@ import {
   RunsWaitPayload,
   StreamEvent,
   CronsCreatePayload,
+  OnConflictBehavior,
 } from "./types.mjs";
 
 interface ClientConfig {
@@ -106,7 +107,7 @@ class BaseClient {
   }
 }
 
-class CronsClient extends BaseClient {
+export class CronsClient extends BaseClient {
   /**
    *
    * @param threadId The ID of the thread.
@@ -194,7 +195,7 @@ class CronsClient extends BaseClient {
   }
 }
 
-class AssistantsClient extends BaseClient {
+export class AssistantsClient extends BaseClient {
   /**
    * Get an assistant by ID.
    *
@@ -301,7 +302,7 @@ class AssistantsClient extends BaseClient {
   }
 }
 
-class ThreadsClient extends BaseClient {
+export class ThreadsClient extends BaseClient {
   /**
    * Get a thread by ID.
    *
@@ -323,10 +324,16 @@ class ThreadsClient extends BaseClient {
      * Metadata for the thread.
      */
     metadata?: Metadata;
+    threadId?: string;
+    ifExists?: OnConflictBehavior;
   }): Promise<Thread> {
     return this.fetch<Thread>(`/threads`, {
       method: "POST",
-      json: { metadata: payload?.metadata },
+      json: {
+        metadata: payload?.metadata,
+        thread_id: payload?.threadId,
+        if_exists: payload?.ifExists,
+      },
     });
   }
 
@@ -489,7 +496,7 @@ class ThreadsClient extends BaseClient {
   }
 }
 
-class RunsClient extends BaseClient {
+export class RunsClient extends BaseClient {
   stream(
     threadId: null,
     assistantId: string,

libs/sdk-js/src/types.mts

@@ -2,6 +2,7 @@ import { Config, Metadata } from "./schema.js";
 
 export type StreamMode = "values" | "messages" | "updates" | "events" | "debug";
 export type MultitaskStrategy = "reject" | "interrupt" | "rollback" | "enqueue";
+export type OnConflictBehavior = "raise" | "do_nothing";
 export type StreamEvent =
   | "events"
   | "metadata"

libs/sdk-js/tsconfig.json

@@ -2,7 +2,11 @@
   "extends": "@tsconfig/recommended",
   "compilerOptions": {
     "target": "ES2021",
-    "lib": ["ES2021", "ES2022.Object", "DOM"],
+    "lib": [
+      "ES2021",
+      "ES2022.Object",
+      "DOM"
+    ],
     "module": "NodeNext",
     "moduleResolution": "nodenext",
     "esModuleInterop": true,
@@ -17,19 +21,26 @@
     "strict": true,
     "outDir": "dist"
   },
-  "include": ["src/**/*"],
-  "exclude": ["node_modules", "dist", "coverage"],
+  "include": [
+    "src/**/*"
+  ],
+  "exclude": [
+    "node_modules",
+    "dist",
+    "coverage"
+  ],
   "includeVersion": true,
   "typedocOptions": {
     "entryPoints": [
-      "src/client.mts",
-      "src/schema.ts",
-      "src/types.mts"
+      "src/client.mts"
     ],
     "readme": "none",
     "out": "docs",
     "plugin": [
       "typedoc-plugin-markdown"
-    ]
+    ],
+    "excludePrivate": true,
+    "excludeProtected": true,
+    "excludeExternals": false
   }
 }

libs/sdk-py/langgraph_sdk/client.py

@@ -29,6 +29,7 @@ from langgraph_sdk.schema import (
     GraphSchema,
     Metadata,
     MultitaskStrategy,
+    OnConflictBehavior,
     Run,
     StreamMode,
     Thread,
@@ -256,6 +257,7 @@ class AssistantsClient:
         *,
         metadata: Metadata = None,
         assistant_id: Optional[str] = None,
+        if_exists: Optional[OnConflictBehavior] = None,
     ) -> Assistant:
         """Create a new assistant."""
         payload: Dict[str, Any] = {
@@ -267,6 +269,8 @@ class AssistantsClient:
             payload["metadata"] = metadata
         if assistant_id:
             payload["assistant_id"] = assistant_id
+        if if_exists:
+            payload["if_exists"] = if_exists
         return await self.http.post("/assistants", json=payload)
 
     async def update(
@@ -344,6 +348,7 @@ class ThreadsClient:
         *,
         metadata: Metadata = None,
         thread_id: Optional[str] = None,
+        if_exists: Optional[OnConflictBehavior] = None,
     ) -> Thread:
         """Create a new thread."""
         payload: Dict[str, Any] = {}
@@ -351,6 +356,8 @@ class ThreadsClient:
             payload["thread_id"] = thread_id
         if metadata:
             payload["metadata"] = metadata
+        if if_exists:
+            payload["if_exists"] = if_exists
         return await self.http.post("/threads", json=payload)
 
     async def update(self, thread_id: str, *, metadata: dict[str, Any]) -> Thread:
@@ -385,6 +392,10 @@ class ThreadsClient:
             json=payload,
         )
 
+    async def copy(self, thread_id: str) -> None:
+        """Copy a thread."""
+        return await self.http.post(f"/threads/{thread_id}/copy", json=None)
+
     async def get_state(
         self, thread_id: str, checkpoint_id: Optional[str] = None
     ) -> ThreadState: