checkpoint-postgres: release 1.0.4 (#1493)

* checkpoint-postgres: allow passing pool

* make psycopg_pool a non-dev dependency

* code review

* lockfile

* move methods

* relax requirements, remove binary

* add binary to dev dependencies

* update readme

README.md

@@ -59,7 +59,7 @@ from langchain_core.messages import HumanMessage
 from langchain_anthropic import ChatAnthropic
 from langchain_core.tools import tool
 from langgraph.checkpoint.memory import MemorySaver
-from langgraph.graph import END, StateGraph, MessagesState
+from langgraph.graph import END, START, StateGraph, MessagesState
 from langgraph.prebuilt import ToolNode
 
 
@@ -107,7 +107,7 @@ workflow.add_node("tools", tool_node)
 
 # Set the entrypoint as `agent`
 # This means that this node is the first one called
-workflow.set_entry_point("agent")
+workflow.add_edge(START, "agent")
 
 # We now add a conditional edge
 workflow.add_conditional_edges(

docs/_scripts/copy_notebooks.py

@@ -56,6 +56,7 @@ _MANUAL = {
         "create-react-agent-memory.ipynb",
         "create-react-agent-hitl.ipynb",
         "human_in_the_loop/breakpoints.ipynb",
+        "human_in_the_loop/dynamic_breakpoints.ipynb",
         "human_in_the_loop/time-travel.ipynb",
         "human_in_the_loop/edit-graph-state.ipynb",
         "human_in_the_loop/wait-user-input.ipynb",

docs/docs/cloud/deployment/graph_rebuild.md

@@ -28,7 +28,7 @@ In the standard LangGraph API configuration, the server uses the compiled graph
 
 ```python
 from langchain_openai import ChatOpenAI
-from langgraph.graph import END, MessageGraph
+from langgraph.graph import END, START, MessageGraph
 
 model = ChatOpenAI(temperature=0)
 
@@ -36,7 +36,7 @@ graph_workflow = MessageGraph()
 
 graph_workflow.add_node("agent", model)
 graph_workflow.add_edge("agent", END)
-graph_workflow.set_entry_point("agent")
+graph_workflow.add_edge(START, "agent")
 
 agent = graph_workflow.compile()
 ```
@@ -60,7 +60,7 @@ To make your graph rebuild on each new run with custom configuration, you need t
 ```python
 from typing import Annotated, TypedDict
 from langchain_openai import ChatOpenAI
-from langgraph.graph import END, MessageGraph
+from langgraph.graph import END, START, MessageGraph
 from langgraph.graph.state import StateGraph
 from langgraph.graph.message import add_messages
 from langgraph.prebuilt import ToolNode
@@ -83,7 +83,7 @@ def make_default_graph():
 
     graph_workflow.add_node("agent", call_model)
     graph_workflow.add_edge("agent", END)
-    graph_workflow.set_entry_point("agent")
+    graph_workflow.add_edge(START, "agent")
 
     agent = graph_workflow.compile()
     return agent
@@ -113,7 +113,7 @@ def make_alternative_graph():
     graph_workflow.add_node("agent", call_model)
     graph_workflow.add_node("tools", tool_node)
     graph_workflow.add_edge("tools", "agent")
-    graph_workflow.set_entry_point("agent")
+    graph_workflow.add_edge(START, "agent")
     graph_workflow.add_conditional_edges("agent", should_continue)
 
     agent = graph_workflow.compile()

docs/docs/cloud/deployment/setup.md

@@ -1,15 +1,14 @@
 # How to Set Up a LangGraph Application for Deployment
 
-A LangGraph application must be configured with a [LangGraph API configuration file](../reference/cli.md#configuration-file) in order to be deployed to LangGraph Cloud (or to be self-hosted). This how-to guide discusses the basic steps to setup a LangGraph application for deployment using `requirements.txt` to specify project dependencies. 
+A LangGraph application must be configured with a [LangGraph API configuration file](../reference/cli.md#configuration-file) in order to be deployed to LangGraph Cloud (or to be self-hosted). This how-to guide discusses the basic steps to setup a LangGraph application for deployment using `requirements.txt` to specify project dependencies.
 
 This walkthrough is based on [this repository](https://github.com/langchain-ai/langgraph-example), which you can play around with to learn more about how to setup your LangGraph application for deployment.
 
 !!! tip "Setup with pyproject.toml"
-    If you prefer using poetry for dependency management, check out [this how-to guide](./setup_pyproject.md) on using `pyproject.toml` for LangGraph Cloud.
+If you prefer using poetry for dependency management, check out [this how-to guide](./setup_pyproject.md) on using `pyproject.toml` for LangGraph Cloud.
 
 !!! tip "Setup with a Monorepo"
-    If you are interested in deploying a graph located inside a monorepo, take a look at [this](https://github.com/langchain-ai/langgraph-example-monorepo) repository for an example of how to do so.
-
+If you are interested in deploying a graph located inside a monorepo, take a look at [this](https://github.com/langchain-ai/langgraph-example-monorepo) repository for an example of how to do so.
 
 The final repo structure will look something like this:
 
@@ -35,23 +34,27 @@ 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.2.0,<0.3.0
+langgraph>=0.2.7,<0.3.0
+langgraph-checkpoint>=1.0.4
 langchain-core>=0.2.27,<0.3.0
 langsmith>=0.1.63
-orjson>=3.10.1
-httpx>=0.27.0
-tenacity>=8.3.0
-uvicorn>=0.29.0
+orjson>=3.9.7
+httpx>=0.25.0
+tenacity>=8.0.0
+uvicorn>=0.26.0
 sse-starlette>=2.1.0
-uvloop>=0.19.0
-httptools>=0.6.1
-jsonschema-rs>=0.18.0
+uvloop>=0.18.0
+httptools>=0.5.0
+jsonschema-rs>=0.16.3
 croniter>=1.0.1
-structlog>=24.4.0
+structlog>=23.1.0
+redis>=5.0.0,<6.0.0
 ```
 
 Example `requirements.txt` file:
+
 ```
 langgraph
 langchain_anthropic
@@ -62,6 +65,7 @@ langchain_openai
 ```
 
 Example file directory:
+
 ```bash
 my-app/
 ├── my_agent # all project code lies within here
@@ -73,6 +77,7 @@ my-app/
 Environment variables can optionally be specified in a file (e.g. `.env`). See the [Environment Variables reference](../reference/env_var.md) to configure additional variables for a deployment.
 
 Example `.env` file:
+
 ```
 MY_ENV_VAR_1=foo
 MY_ENV_VAR_2=bar
@@ -94,12 +99,11 @@ Implement your graphs! Graphs can be defined in a single file or multiple files.
 
 Example `agent.py` file, which shows how to import from other modules you define (code for the modules is not shown here, please see [this repo](https://github.com/langchain-ai/langgraph-example) to see their implementation):
 
-
 ```python
 # my_agent/agent.py
 from typing import TypedDict, Literal
 
-from langgraph.graph import StateGraph, END
+from langgraph.graph import StateGraph, END, START
 from my_agent.utils.nodes import call_model, should_continue, tool_node # import nodes
 from my_agent.utils.state import AgentState # import state
 
@@ -110,7 +114,7 @@ class GraphConfig(TypedDict):
 workflow = StateGraph(AgentState, config_schema=GraphConfig)
 workflow.add_node("agent", call_model)
 workflow.add_node("action", tool_node)
-workflow.set_entry_point("agent")
+workflow.add_edge(START, "agent")
 workflow.add_conditional_edges(
     "agent",
     should_continue,
@@ -125,9 +129,10 @@ 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 (alternatively, you can provide [a function that creates a graph](./graph_rebuild.md)).
+The build process for LangGraph Cloud requires that the `CompiledGraph` object be assigned to a variable at the top-level of a Python module (alternatively, you can provide [a function that creates a graph](./graph_rebuild.md)).
 
 Example file directory:
+
 ```bash
 my-app/
 ├── my_agent # all project code lies within here
@@ -147,6 +152,7 @@ my-app/
 Create a [LangGraph API configuration file](../reference/cli.md#configuration-file) called `langgraph.json`. See the [LangGraph CLI reference](../reference/cli.md#configuration-file) for detailed explanations of each key in the JSON object of the configuration file.
 
 Example `langgraph.json` file:
+
 ```json
 {
   "dependencies": ["./my_agent"],
@@ -160,7 +166,7 @@ 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.
+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:
 

docs/docs/cloud/deployment/setup_pyproject.md

@@ -35,19 +35,21 @@ Dependencies can optionally be specified in one of the following files: `pyproje
 
 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.2.0,<0.3.0
+langgraph>=0.2.7,<0.3.0
+langgraph-checkpoint>=1.0.4
 langchain-core>=0.2.27,<0.3.0
 langsmith>=0.1.63
-orjson>=3.10.1
-httpx>=0.27.0
-tenacity>=8.3.0
-uvicorn>=0.29.0
+orjson>=3.9.7
+httpx>=0.25.0
+tenacity>=8.0.0
+uvicorn>=0.26.0
 sse-starlette>=2.1.0
-uvloop>=0.19.0
-httptools>=0.6.1
-jsonschema-rs>=0.18.0
+uvloop>=0.18.0
+httptools>=0.5.0
+jsonschema-rs>=0.16.3
 croniter>=1.0.1
 structlog>=24.4.0
+redis>=5.0.8,<6.0.0
 ```
 
 Example `pyproject.toml` file:
@@ -109,7 +111,7 @@ Example `agent.py` file, which shows how to import from other modules you define
 # my_agent/agent.py
 from typing import TypedDict, Literal
 
-from langgraph.graph import StateGraph, END
+from langgraph.graph import StateGraph, END, START
 from my_agent.utils.nodes import call_model, should_continue, tool_node # import nodes
 from my_agent.utils.state import AgentState # import state
 
@@ -120,7 +122,7 @@ class GraphConfig(TypedDict):
 workflow = StateGraph(AgentState, config_schema=GraphConfig)
 workflow.add_node("agent", call_model)
 workflow.add_node("action", tool_node)
-workflow.set_entry_point("agent")
+workflow.add_edge(START, "agent")
 workflow.add_conditional_edges(
     "agent",
     should_continue,

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

@@ -34,7 +34,8 @@ First, we need to setup our client so that we can communicate with our hosted gr
     ```bash
     curl --request POST \
       --url <DEPLOYMENT_URL>/threads \
-      --header 'Content-Type: application/json'
+      --header 'Content-Type: application/json' \
+      --data {}
     ```
 
 ## Replay a state
@@ -46,7 +47,7 @@ Before replaying a state - we need to create states to replay from! In order to
 === "Python"
 
     ```python
-    input = { 'messages':[{ "role":"user", "content":"Please search the weather in SF" }] }
+    input = {"messages": [{"role": "user", "content": "Please search the weather in SF"}]}
 
     async for chunk in client.runs.stream(
         thread["thread_id"],
@@ -61,7 +62,7 @@ Before replaying a state - we need to create states to replay from! In order to
 === "Javascript"
 
     ```js
-    const input = {"messages": [{ "role": "human", "content": "Please search the weather in SF"}] }
+    const input = { "messages": [{ "role": "human", "content": "Please search the weather in SF" }] }
 
     const streamResponse = client.runs.stream(
       thread["thread_id"],
@@ -155,17 +156,23 @@ Output:
 
 
 
-To rerun from a state, we need to pass in the `checkpoint_id` into the config of the run like follows:
+To rerun from a state, we need first issue an empty update to the thread state. Then we need to pass in the resulting `checkpoint_id` as follows:
 
 === "Python"
 
     ```python
+    state_to_replay = states[2]
+    updated_config = await client.threads.update_state(
+        thread["thread_id"],
+        {"messages": []},
+        checkpoint_id=state_to_replay["checkpoint_id"]
+    )
     async for chunk in client.runs.stream(
         thread["thread_id"],
         assistant_id, # graph_id
         input=None,
         stream_mode="updates",
-        config={"configurable": {"checkpoint_id": state_to_replay['checkpoint_id']}}
+        checkpoint_id=updated_config["checkpoint_id"]
     ):
         if chunk.data and chunk.event != "metadata": 
             print(chunk.data)
@@ -174,13 +181,15 @@ To rerun from a state, we need to pass in the `checkpoint_id` into the config of
 === "Javascript"
 
     ```js
+    const stateToReplay = states[2];
+    const config = await client.threads.updateState(thread["thread_id"], { values: {"messages": [] }, checkpointId: stateToReplay["checkpoint_id"] });
     const streamResponse = client.runs.stream(
       thread["thread_id"],
       assistantId,
       {
         input: null,
         streamMode: "updates",
-        config: {"configurable": {"checkpoint_id": stateToReplay['checkpoint_id']}},
+        checkpointId: config["checkpoint_id"]
       }
     );
     for await (const chunk of streamResponse) {
@@ -193,38 +202,46 @@ To rerun from a state, we need to pass in the `checkpoint_id` into the config of
 === "CURL"
 
     ```bash
-    curl --request GET --url <DEPLOYMENT_URL>/threads/<THREAD_ID>/history | jq -r '.[2].checkpoint_id' | {  
-    read checkpoint_id   
-    curl --request POST \                                                            
-         --url <DEPLOYMENT_URL>/threads/<THREAD_ID>/runs/stream \
-         --header 'Content-Type: application/json' \
-         --data "{                     
-           \"assistant_id\": \"agent\",
-           \"config\": {\"configurable\": {\"checkpoint_id\": \"$checkpoint_id\"}},
-           \"stream_mode\": [
-             \"updates\"
-           ]   
-         }" | \           
-         sed 's/\r$//' | \
-         awk '      
-         /^event:/ {                                              
-             if (data_content != "" && event_type != "metadata") {
-                 print data_content "\n"
-             }                      
-             sub(/^event: /, "", $0)
-             event_type = $0  
-             data_content = ""
-         }         
-         /^data:/ {                
-             sub(/^data: /, "", $0)
-             data_content = $0
-         }    
-         END {                                                    
-             if (data_content != "" && event_type != "metadata") {
-                 print data_content "\n"
-             }
+    curl --request GET --url <DEPLOYMENT_URL>/threads/<THREAD_ID>/history | jq -c '
+        .[2] as $state_to_replay |
+        {
+            values: { messages: .[2].values.messages[-1] },
+            checkpoint_id: $state_to_replay.checkpoint_id
+        }' | \
+    curl --request POST \
+        --url <DEPLOYMENT_URL>/threads/<THREAD_ID>/state \
+        --header 'Content-Type: application/json' \
+        --data @- | jq .checkpoint_id | \
+    curl --request POST \
+     --url <DEPLOYMENT_URL>/threads/<THREAD_ID>/runs/stream \
+     --header 'Content-Type: application/json' \
+     --data "{
+       \"assistant_id\": \"agent\",
+       \"checkpoint_id\": \"$1\",
+       \"stream_mode\": [
+         \"updates\"
+       ]
+     }" | \
+     sed 's/\r$//' | \
+     awk '
+     /^event:/ {
+         if (data_content != "" && event_type != "metadata") {
+             print data_content "\n"
          }
-         ' 
+         sub(/^event: /, "", $0)
+         event_type = $0
+         data_content = ""
+     }
+     /^data:/ {
+         sub(/^data: /, "", $0)
+         data_content = $0
+     }
+     END {
+         if (data_content != "" && event_type != "metadata") {
+             print data_content "\n"
+         }
+     }
+     '
     ```
 
 Output:
@@ -251,7 +268,7 @@ Let's show how to do this to edit the state at a particular point in time. Let's
     # Let's now update the args for that tool call
     last_message['tool_calls'][0]['args'] = {'query': 'current weather in SF'}
 
-    new_state = await client.threads.update_state(thread['thread_id'],{"messages":[last_message]},checkpoint_id=state_to_replay['checkpoint_id'])
+    config = await client.threads.update_state(thread['thread_id'],{"messages":[last_message]},checkpoint_id=state_to_replay['checkpoint_id'])
     ```
 
 === "Javascript"
@@ -262,9 +279,9 @@ Let's show how to do this to edit the state at a particular point in time. Let's
     let lastMessage = stateToReplay['values']['messages'][-1];
 
     // Let's now update the args for that tool call
-    lastMessage['tool_calls'][0]['args'] = {'query': 'current weather in SF'};
+    lastMessage['tool_calls'][0]['args'] = { 'query': 'current weather in SF' };
 
-    const newState = await client.threads.updateState(thread['thread_id'],{values:{"messages":[lastMessage]},checkpointId:stateToReplay['checkpoint_id']});
+    const config = await client.threads.updateState(thread['thread_id'], { values: { "messages": [lastMessage] }, checkpointId: stateToReplay['checkpoint_id'] });
     ```
 
 === "CURL"
@@ -294,7 +311,7 @@ Now we can rerun our graph with this new config, starting from the `new_state`,
         assistant["assistant_id"], # graph_id
         input=None,
         stream_mode="updates",
-        config={"configurable": {"checkpoint_id": new_state['configurable']['checkpoint_id']}}
+        checkpoint_id=config['checkpoint_id']
     ):
         if chunk.data and chunk.event != "metadata": 
             print(chunk.data)
@@ -309,7 +326,7 @@ Now we can rerun our graph with this new config, starting from the `new_state`,
       {
         input: null,
         streamMode: "updates",
-        config: {"configurable": {"checkpoint_id": newState['configurable']['checkpoint_id']}},
+        checkpointId: config['checkpoint_id'],
       }
     );
     for await (const chunk of streamResponse) {
@@ -323,33 +340,37 @@ Now we can rerun our graph with this new config, starting from the `new_state`,
 
     ```bash
     curl -s --request GET --url <DEPLOYMENT_URL>/threads/<THREAD_ID>/state | \
-    jq -r '.config.configurable.checkpoint_id' | \
-    sh -c '
-        CHECKPOINT_ID="$1"
-        curl --request POST \
-            --url <DEPLOYMENT_URL>/threads/<THREAD_ID>/runs/stream \
-            --header "Content-Type: application/json" \
-            --data "{\"assistant_id\": \"agent\", \"config\": {\"configurable\": {\"checkpoint_id\": \"$CHECKPOINT_ID\"}}, \"stream_mode\": [\"updates\"]}" | \
-        sed "s/\r$//" | \
-        awk "
-        /^event:/ {
-            if (data_content != \"\" && event_type != \"metadata\") {
-                print data_content \"\n\"
-            }
-            sub(/^event: /, \"\", \$0)
-            event_type = \$0
-            data_content = \"\"
-        }
-        /^data:/ {
-            sub(/^data: /, \"\", \$0)
-            data_content = \$0
-        }
-        END {
-            if (data_content != \"\" && event_type != \"metadata\") {
-                print data_content \"\n\"
-            }
-        }"
-    ' _ 
+    jq -c '.checkpoint_id' | \
+    curl --request POST \
+     --url <DEPLOYMENT_URL>/threads/<THREAD_ID>/runs/stream \
+     --header 'Content-Type: application/json' \
+     --data "{
+       \"assistant_id\": \"agent\",
+       \"checkpoint_id\": \"$1\",
+       \"stream_mode\": [
+         \"updates\"
+       ]
+     }" | \
+     sed 's/\r$//' | \
+     awk '
+     /^event:/ {
+         if (data_content != "" && event_type != "metadata") {
+             print data_content "\n"
+         }
+         sub(/^event: /, "", $0)
+         event_type = $0
+         data_content = ""
+     }
+     /^data:/ {
+         sub(/^data: /, "", $0)
+         data_content = $0
+     }
+     END {
+         if (data_content != "" && event_type != "metadata") {
+             print data_content "\n"
+         }
+     }
+     '
     ```
 
 Output:

docs/docs/cloud/reference/api/api_ref.md

@@ -3,3 +3,20 @@
 The LangGraph Cloud API reference is available with each deployment at the `/docs` URL path (e.g. `http://localhost:8124/docs`).
 
 Click <a href="/langgraph/cloud/reference/api/api_ref.html" target="_blank">here</a> to view the API reference.
+
+## Authentication
+
+For deployments to LangGraph Cloud, authentication is required. Pass the `X-Api-Key` header with each request to the LangGraph Cloud API. The value of the header should be set to a valid LangSmith API key for the organization where the API is deployed.
+
+Example `curl` command:
+```shell
+curl --request POST \
+  --url http://localhost:8124/assistants/search \
+  --header 'Content-Type: application/json' \
+  --header 'X-Api-Key: LANGSMITH_API_KEY' \
+  --data '{
+  "metadata": {},
+  "limit": 10,
+  "offset": 0
+}'  
+```

docs/docs/concepts/low_level.md

@@ -433,7 +433,17 @@ LangGraph is built with first class support for streaming. There are several dif
 - [`"updates`](../how-tos/stream-updates.ipynb): This streams the updates to the state after each step of the graph. If multiple updates are made in the same step (e.g. multiple nodes are run) then those updates are streamed separately.
 - `"debug"`: This streams as much information as possible throughout the execution of the graph.
 
-In addition, you can use the [`astream_events`](../how-tos/streaming-events-from-within-tools.ipynb) method to stream back events that happen _inside_ nodes. This is useful for [streaming tokens of LLM calls](../how-tos/streaming-tokens.ipynb).
+In addition, you can use the [`astream_events`](../how-tos/streaming-events-from-within-tools.ipynb) method to stream back events that happen _inside_ nodes. This is useful for [streaming tokens of LLM calls](../how-tos/streaming-tokens.ipynb) among other things.
+
+Under the hood, the compiled graph and each of its nodes are turned into [runnables](https://python.langchain.com/v0.2/docs/concepts/#runnable-interface). This means that as the graph is executed, certain events are emitted along the way and can be seen if you run the graph using `.astream_events`:
+
+* each node (runnable) emits `on_chain_start` when it starts execution, `on_chain_stream` during the node execution and `on_chain_end` when the node finishes. Node events will have the node name in the event's `name` field
+* the graph will emit `on_chain_start` in the beginning of the graph execution, `on_chain_stream` during each node execution and `on_chain_end` when the graph finishes. Graph events will have the `LangGraph` in the event's `name` field
+* LangGraph writes to state channels emit `on_chain_start` and `on_chain_end` events
+
+Additionally, any events that are created inside your nodes (LLM events, tool events, manually emitted events, etc.) will also be visible in the output of `.astream_events`.
+ 
+You can find a detailed table of all callback events and triggers [here](https://python.langchain.com/v0.2/docs/concepts/#callback-events).
 
 !!! warning "ASYNC IN PYTHON<=3.10"
     You may fail to see events being emitted from inside a node when using `.astream_events` in Python <= 3.10. If you're using a Langchain RunnableLambda, a RunnableGenerator, or Tool asynchronously inside your node, you will have to propagate callbacks to these objects manually. This is because LangChain cannot automatically propagate callbacks to child objects in this case. Please see examples [here](../how-tos/streaming-content.ipynb) and [here](../how-tos/streaming-events-from-within-tools.ipynb).

docs/docs/how-tos/index.md

@@ -35,6 +35,7 @@ One of LangGraph's main benefits is that it makes human-in-the-loop workflows ea
 These guides cover common examples of that.
 
 - [How to add breakpoints](human_in_the_loop/breakpoints.ipynb)
+- [How to add dynamic breakpoints](human_in_the_loop/dynamic_breakpoints.ipynb)
 - [How to edit graph state](human_in_the_loop/edit-graph-state.ipynb)
 - [How to wait for user input](human_in_the_loop/wait-user-input.ipynb)
 - [How to view and update past graph state](human_in_the_loop/time-travel.ipynb)

docs/mkdocs.yml

@@ -139,6 +139,7 @@ nav:
           - Create custom checkpointer using Redis: how-tos/persistence_redis.ipynb
       - Human-in-the-loop:
           - Add breakpoints: how-tos/human_in_the_loop/breakpoints.ipynb
+          - Add dynamic breakpoints: how-tos/human_in_the_loop/dynamic_breakpoints.ipynb
           - Wait for user input: how-tos/human_in_the_loop/wait-user-input.ipynb
           - View and update past graph state: how-tos/human_in_the_loop/time-travel.ipynb
           - Edit graph state: how-tos/human_in_the_loop/edit-graph-state.ipynb

examples/chatbot-simulation-evaluation/agent-simulation-evaluation.ipynb

@@ -235,7 +235,7 @@
     "    # Call the chat bot\n",
     "    chat_bot_response = my_chat_bot(messages)\n",
     "    # Respond with an AI Message\n",
-    "    return {\"messages\":[AIMessage(content=chat_bot_response[\"content\"])]}"
+    "    return {\"messages\": [AIMessage(content=chat_bot_response[\"content\"])]}"
    ]
   },
   {
@@ -270,7 +270,7 @@
     "    # Call the simulated user\n",
     "    response = simulated_user.invoke({\"messages\": new_messages})\n",
     "    # This response is an AI message - we need to flip this to be a human message\n",
-    "    return {\"messages\":[HumanMessage(content=response.content)]}"
+    "    return {\"messages\": [HumanMessage(content=response.content)]}"
    ]
   },
   {
@@ -331,6 +331,7 @@
     "class State(TypedDict):\n",
     "    messages: Annotated[list, add_messages]\n",
     "\n",
+    "\n",
     "graph_builder = StateGraph(State)\n",
     "graph_builder.add_node(\"user\", simulated_user_node)\n",
     "graph_builder.add_node(\"chat_bot\", chat_bot_node)\n",

examples/chatbots/information-gather-prompting.ipynb

@@ -79,7 +79,7 @@
     "\n",
     "\n",
     "def info_chain(state):\n",
-    "    messages = get_messages_info(state['messages'])\n",
+    "    messages = get_messages_info(state[\"messages\"])\n",
     "    response = llm_with_tool.invoke(messages)\n",
     "    return {\"messages\": [response]}"
    ]
@@ -126,7 +126,7 @@
     "\n",
     "\n",
     "def prompt_gen_chain(state):\n",
-    "    messages = get_prompt_messages(state['messages'])\n",
+    "    messages = get_prompt_messages(state[\"messages\"])\n",
     "    response = llm.invoke(messages)\n",
     "    return {\"messages\": [response]}"
    ]
@@ -158,7 +158,7 @@
     "\n",
     "\n",
     "def get_state(state) -> Literal[\"add_tool_message\", \"info\", \"__end__\"]:\n",
-    "    messages = state['messages']\n",
+    "    messages = state[\"messages\"]\n",
     "    if isinstance(messages[-1], AIMessage) and messages[-1].tool_calls:\n",
     "        return \"add_tool_message\"\n",
     "    elif not isinstance(messages[-1], HumanMessage):\n",
@@ -190,9 +190,11 @@
     "from typing import Annotated\n",
     "from typing_extensions import TypedDict\n",
     "\n",
+    "\n",
     "class State(TypedDict):\n",
     "    messages: Annotated[list, add_messages]\n",
     "\n",
+    "\n",
     "memory = MemorySaver()\n",
     "workflow = StateGraph(State)\n",
     "workflow.add_node(\"info\", info_chain)\n",
@@ -201,9 +203,14 @@
     "\n",
     "@workflow.add_node\n",
     "def add_tool_message(state: State):\n",
-    "    return {\"messages\": [ToolMessage(\n",
-    "        content=\"Prompt generated!\", tool_call_id=state['messages'][-1].tool_calls[0][\"id\"]\n",
-    "    )]}\n",
+    "    return {\n",
+    "        \"messages\": [\n",
+    "            ToolMessage(\n",
+    "                content=\"Prompt generated!\",\n",
+    "                tool_call_id=state[\"messages\"][-1].tool_calls[0][\"id\"],\n",
+    "            )\n",
+    "        ]\n",
+    "    }\n",
     "\n",
     "\n",
     "workflow.add_conditional_edges(\"info\", get_state)\n",
@@ -364,7 +371,7 @@
     "    for output in graph.stream(\n",
     "        {\"messages\": [HumanMessage(content=user)]}, config=config, stream_mode=\"updates\"\n",
     "    ):\n",
-    "        last_message = next(iter(output.values()))['messages'][-1]\n",
+    "        last_message = next(iter(output.values()))[\"messages\"][-1]\n",
     "        last_message.pretty_print()\n",
     "\n",
     "    if output and \"prompt\" in output:\n",

examples/create-react-agent-hitl.ipynb

@@ -239,7 +239,7 @@
       "name": "python",
       "nbconvert_exporter": "python",
       "pygments_lexer": "ipython3",
-      "version": "3.11.1"
+      "version": "3.12.2"
     }
   },
   "nbformat": 4,

examples/customer-support/customer-support.ipynb

@@ -225,7 +225,14 @@
     "\n",
     "Define the (`fetch_user_flight_information`) tool to let the agent see the current user's flight information.  Then define tools to search for flights and manage the passenger's bookings stored in the SQL database.\n",
     "\n",
-    "We use `ensure_config` to pass in the `passenger_id` in via configurable parameters. The LLM never has to provide these explicitly, they are provided for a given invocation of the graph so that each user cannot access other passengers' booking information."
+    "We the can [access the RunnableConfig](https://python.langchain.com/v0.2/docs/how_to/tool_configure/#inferring-by-parameter-type) for a given run to check the `passenger_id` of the user accessing this application. The LLM never has to provide these explicitly, they are provided for a given invocation of the graph so that each user cannot access other passengers' booking information.\n",
+    "\n",
+    "<div class=\"admonition warning\">\n",
+    "    <p class=\"admonition-title\">Compatibility</p>\n",
+    "    <p>\n",
+    "        This tutorial expects `langchain-core>=0.2.16` to use the injected RunnableConfig. Prior to that, you'd use `ensure_config` to collect the config from context.\n",
+    "    </p>\n",
+    "</div> \n"
    ]
   },
   {
@@ -240,18 +247,17 @@
     "from typing import Optional\n",
     "\n",
     "import pytz\n",
-    "from langchain_core.runnables import ensure_config\n",
+    "from langchain_core.runnables import RunnableConfig\n",
     "\n",
     "\n",
     "@tool\n",
-    "def fetch_user_flight_information() -> list[dict]:\n",
+    "def fetch_user_flight_information(config: RunnableConfig) -> list[dict]:\n",
     "    \"\"\"Fetch all tickets for the user along with corresponding flight information and seat assignments.\n",
     "\n",
     "    Returns:\n",
     "        A list of dictionaries where each dictionary contains the ticket details,\n",
     "        associated flight details, and the seat assignments for each ticket belonging to the user.\n",
     "    \"\"\"\n",
-    "    config = ensure_config()  # Fetch from the context\n",
     "    configuration = config.get(\"configurable\", {})\n",
     "    passenger_id = configuration.get(\"passenger_id\", None)\n",
     "    if not passenger_id:\n",
@@ -328,9 +334,10 @@
     "\n",
     "\n",
     "@tool\n",
-    "def update_ticket_to_new_flight(ticket_no: str, new_flight_id: int) -> str:\n",
+    "def update_ticket_to_new_flight(\n",
+    "    ticket_no: str, new_flight_id: int, *, config: RunnableConfig\n",
+    ") -> str:\n",
     "    \"\"\"Update the user's ticket to a new valid flight.\"\"\"\n",
-    "    config = ensure_config()\n",
     "    configuration = config.get(\"configurable\", {})\n",
     "    passenger_id = configuration.get(\"passenger_id\", None)\n",
     "    if not passenger_id:\n",
@@ -396,9 +403,8 @@
     "\n",
     "\n",
     "@tool\n",
-    "def cancel_ticket(ticket_no: str) -> str:\n",
+    "def cancel_ticket(ticket_no: str, *, config: RunnableConfig) -> str:\n",
     "    \"\"\"Cancel the user's ticket and remove it from the database.\"\"\"\n",
-    "    config = ensure_config()\n",
     "    configuration = config.get(\"configurable\", {})\n",
     "    passenger_id = configuration.get(\"passenger_id\", None)\n",
     "    if not passenger_id:\n",
@@ -4407,7 +4413,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.2"
+   "version": "3.12.2"
   }
  },
  "nbformat": 4,

examples/human_in_the_loop/review-tool-calls.ipynb

@@ -141,15 +141,18 @@
     "    print(\"----\")\n",
     "    return \"Sunny!\"\n",
     "\n",
-    "model = ChatAnthropic(model_name=\"claude-3-5-sonnet-20240620\").bind_tools([weather_search])\n",
+    "\n",
+    "model = ChatAnthropic(model_name=\"claude-3-5-sonnet-20240620\").bind_tools(\n",
+    "    [weather_search]\n",
+    ")\n",
+    "\n",
     "\n",
     "class State(MessagesState):\n",
     "    \"\"\"Simple state.\"\"\"\n",
     "\n",
+    "\n",
     "def call_llm(state):\n",
-    "    return {\n",
-    "        \"messages\": [model.invoke(state['messages'])]\n",
-    "    }\n",
+    "    return {\"messages\": [model.invoke(state[\"messages\"])]}\n",
     "\n",
     "\n",
     "def human_review_node(state):\n",
@@ -159,28 +162,30 @@
     "def run_tool(state):\n",
     "    new_messages = []\n",
     "    tools = {\"weather_search\": weather_search}\n",
-    "    tool_calls = state['messages'][-1].tool_calls\n",
+    "    tool_calls = state[\"messages\"][-1].tool_calls\n",
     "    for tool_call in tool_calls:\n",
-    "        tool = tools[tool_call['name']]\n",
-    "        result = tool.invoke(tool_call['args'])\n",
-    "        new_messages.append({\n",
-    "            \"role\": \"tool\",\n",
-    "            \"name\": tool_call['name'],\n",
-    "            \"content\": result,\n",
-    "            \"tool_call_id\": tool_call['id']\n",
-    "        })\n",
+    "        tool = tools[tool_call[\"name\"]]\n",
+    "        result = tool.invoke(tool_call[\"args\"])\n",
+    "        new_messages.append(\n",
+    "            {\n",
+    "                \"role\": \"tool\",\n",
+    "                \"name\": tool_call[\"name\"],\n",
+    "                \"content\": result,\n",
+    "                \"tool_call_id\": tool_call[\"id\"],\n",
+    "            }\n",
+    "        )\n",
     "    return {\"messages\": new_messages}\n",
     "\n",
     "\n",
     "def route_after_llm(state) -> Literal[END, \"human_review_node\"]:\n",
-    "    if len(state['messages'][-1].tool_calls) == 0:\n",
+    "    if len(state[\"messages\"][-1].tool_calls) == 0:\n",
     "        return END\n",
     "    else:\n",
     "        return \"human_review_node\"\n",
     "\n",
     "\n",
     "def route_after_human(state) -> Literal[\"run_tool\", \"call_llm\"]:\n",
-    "    if isinstance(state['messages'][-1], AIMessage):\n",
+    "    if isinstance(state[\"messages\"][-1], AIMessage):\n",
     "        return \"run_tool\"\n",
     "    else:\n",
     "        return \"call_llm\"\n",
@@ -460,35 +465,35 @@
     "print(\"Current State:\")\n",
     "print(state.values)\n",
     "print(\"\\nCurrent Tool Call ID:\")\n",
-    "current_content = state.values['messages'][-1].content\n",
-    "current_id = state.values['messages'][-1].id\n",
-    "tool_call_id = state.values['messages'][-1].tool_calls[0]['id']\n",
+    "current_content = state.values[\"messages\"][-1].content\n",
+    "current_id = state.values[\"messages\"][-1].id\n",
+    "tool_call_id = state.values[\"messages\"][-1].tool_calls[0][\"id\"]\n",
     "print(tool_call_id)\n",
     "\n",
     "# We now need to construct a replacement tool call.\n",
     "# We will change the argument to be `San Francisco, USA`\n",
     "# Note that we could change any number of arguments or tool names - it just has to be a valid one\n",
     "new_message = {\n",
-    "    \"role\": \"assistant\", \n",
+    "    \"role\": \"assistant\",\n",
     "    \"content\": current_content,\n",
     "    \"tool_calls\": [\n",
     "        {\n",
     "            \"id\": tool_call_id,\n",
     "            \"name\": \"weather_search\",\n",
-    "            \"args\": {\"city\": \"San Francisco, USA\"}\n",
+    "            \"args\": {\"city\": \"San Francisco, USA\"},\n",
     "        }\n",
     "    ],\n",
     "    # This is important - this needs to be the same as the message you replacing!\n",
     "    # Otherwise, it will show up as a separate message\n",
-    "    \"id\": current_id\n",
+    "    \"id\": current_id,\n",
     "}\n",
     "graph.update_state(\n",
     "    # This is the config which represents this thread\n",
-    "    thread, \n",
+    "    thread,\n",
     "    # This is the updated value we want to push\n",
-    "    {\"messages\": [new_message]}, \n",
+    "    {\"messages\": [new_message]},\n",
     "    # We push this update acting as our human_review_node\n",
-    "    as_node=\"human_review_node\"\n",
+    "    as_node=\"human_review_node\",\n",
     ")\n",
     "\n",
     "# Let's now continue executing from here\n",
@@ -595,26 +600,26 @@
     "print(\"Current State:\")\n",
     "print(state.values)\n",
     "print(\"\\nCurrent Tool Call ID:\")\n",
-    "tool_call_id = state.values['messages'][-1].tool_calls[0]['id']\n",
+    "tool_call_id = state.values[\"messages\"][-1].tool_calls[0][\"id\"]\n",
     "print(tool_call_id)\n",
     "\n",
     "# We now need to construct a replacement tool call.\n",
     "# We will change the argument to be `San Francisco, USA`\n",
     "# Note that we could change any number of arguments or tool names - it just has to be a valid one\n",
     "new_message = {\n",
-    "    \"role\": \"tool\", \n",
+    "    \"role\": \"tool\",\n",
     "    # This is our natural language feedback\n",
     "    \"content\": \"User requested changes: pass in the country as well\",\n",
     "    \"name\": \"weather_search\",\n",
-    "    \"tool_call_id\": tool_call_id\n",
+    "    \"tool_call_id\": tool_call_id,\n",
     "}\n",
     "graph.update_state(\n",
     "    # This is the config which represents this thread\n",
-    "    thread, \n",
+    "    thread,\n",
     "    # This is the updated value we want to push\n",
-    "    {\"messages\": [new_message]}, \n",
+    "    {\"messages\": [new_message]},\n",
     "    # We push this update acting as our human_review_node\n",
-    "    as_node=\"human_review_node\"\n",
+    "    as_node=\"human_review_node\",\n",
     ")\n",
     "\n",
     "# Let's now continue executing from here\n",

examples/input_output_schema.ipynb

@@ -33,15 +33,19 @@
     "from langgraph.graph import StateGraph, START, END\n",
     "from typing import TypedDict\n",
     "\n",
+    "\n",
     "class InputState(TypedDict):\n",
     "    question: str\n",
     "\n",
+    "\n",
     "class OutputState(TypedDict):\n",
     "    answer: str\n",
     "\n",
+    "\n",
     "def answer_node(state: InputState):\n",
     "    return {\"answer\": \"bye\"}\n",
     "\n",
+    "\n",
     "graph = StateGraph(input=InputState, output=OutputState)\n",
     "graph.add_node(answer_node)\n",
     "graph.add_edge(START, \"answer_node\")\n",

examples/llm-compiler/LLMCompiler.ipynb

@@ -526,7 +526,7 @@
     "            \"tasks\": tasks,\n",
     "        }\n",
     "    )\n",
-    "    return {\"messages\":[scheduled_tasks]}"
+    "    return {\"messages\": [scheduled_tasks]}"
    ]
   },
   {
@@ -653,7 +653,7 @@
     "            )\n",
     "        ]\n",
     "    else:\n",
-    "        return {\"messages\":response + [AIMessage(content=decision.action.response)]}\n",
+    "        return {\"messages\": response + [AIMessage(content=decision.action.response)]}\n",
     "\n",
     "\n",
     "def select_recent_messages(state) -> dict:\n",
@@ -726,9 +726,11 @@
     "from langgraph.graph.message import add_messages\n",
     "from typing import Annotated\n",
     "\n",
+    "\n",
     "class State(TypedDict):\n",
     "    messages: Annotated[list, add_messages]\n",
     "\n",
+    "\n",
     "graph_builder = StateGraph(State)\n",
     "\n",
     "# 1.  Define vertices\n",
@@ -794,7 +796,9 @@
     }
    ],
    "source": [
-    "for step in chain.stream({\"messages\":[HumanMessage(content=\"What's the GDP of New York?\")]}):\n",
+    "for step in chain.stream(\n",
+    "    {\"messages\": [HumanMessage(content=\"What's the GDP of New York?\")]}\n",
+    "):\n",
     "    print(step)\n",
     "    print(\"---\")"
    ]

examples/many-tools.ipynb

@@ -328,9 +328,9 @@
     "            \"set more_information_needed False and populate a blank string for the query.\"\n",
     "        )\n",
     "        input_messages = [system] + state[\"messages\"]\n",
-    "        response = llm.bind_tools(\n",
-    "            [QueryForTools], tool_choice=True\n",
-    "        ).invoke(input_messages)\n",
+    "        response = llm.bind_tools([QueryForTools], tool_choice=True).invoke(\n",
+    "            input_messages\n",
+    "        )\n",
     "        query = response.tool_calls[0][\"args\"][\"query\"]\n",
     "    tool_documents = vector_store.similarity_search(query)\n",
     "    if hack_remove_tool_condition:\n",

examples/pass-run-time-values-to-tools.ipynb

@@ -329,6 +329,7 @@
     "\n",
     "tools = [get_context, cite_context_sources]\n",
     "\n",
+    "\n",
     "# Define the function that calls the model\n",
     "def call_model(state, config):\n",
     "    messages = state[\"messages\"]\n",

examples/pass_private_state.ipynb

@@ -72,12 +72,12 @@
     "# Node to retrieve documents\n",
     "def retrieve_documents(state: QueryOutputState) -> DocumentOutputState:\n",
     "    # Replace this with real logic\n",
-    "    return {\"docs\": [state['query']] * 2}\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",
+    "    return {\"answer\": \"\\n\\n\".join(state[\"docs\"] + [state[\"question\"]])}\n",
     "\n",
     "\n",
     "graph = StateGraph(OverallState)\n",

examples/persistence.ipynb

@@ -587,7 +587,7 @@
       "name": "python",
       "nbconvert_exporter": "python",
       "pygments_lexer": "ipython3",
-      "version": "3.11.9"
+      "version": "3.12.2"
     }
   },
   "nbformat": 4,

examples/persistence_mongodb.ipynb

@@ -630,7 +630,7 @@
     "                    upsert=True,\n",
     "                )\n",
     "            )\n",
-    "        await self.db[\"checkpoint_writes\"].bulk_write(operations)\n"
+    "        await self.db[\"checkpoint_writes\"].bulk_write(operations)"
    ]
   },
   {
@@ -685,7 +685,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "with MongoDBSaver.from_conn_info(host=\"localhost\", port=27017, db_name=\"checkpoints\") as checkpointer:\n",
+    "with MongoDBSaver.from_conn_info(\n",
+    "    host=\"localhost\", port=27017, db_name=\"checkpoints\"\n",
+    ") as checkpointer:\n",
     "    graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)\n",
     "    config = {\"configurable\": {\"thread_id\": \"1\"}}\n",
     "    res = graph.invoke({\"messages\": [(\"human\", \"what's the weather in sf\")]}, config)\n",
@@ -796,10 +798,14 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "async with AsyncMongoDBSaver.from_conn_info(host=\"localhost\", port=27017, db_name=\"checkpoints\") as checkpointer:\n",
+    "async with AsyncMongoDBSaver.from_conn_info(\n",
+    "    host=\"localhost\", port=27017, db_name=\"checkpoints\"\n",
+    ") as checkpointer:\n",
     "    graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)\n",
     "    config = {\"configurable\": {\"thread_id\": \"2\"}}\n",
-    "    res = await graph.ainvoke({\"messages\": [(\"human\", \"what's the weather in nyc\")]}, config)\n",
+    "    res = await graph.ainvoke(\n",
+    "        {\"messages\": [(\"human\", \"what's the weather in nyc\")]}, config\n",
+    "    )\n",
     "\n",
     "    latest_checkpoint = await checkpointer.aget(config)\n",
     "    latest_checkpoint_tuple = await checkpointer.aget_tuple(config)\n",

examples/persistence_postgres.ipynb

@@ -134,7 +134,7 @@
    "source": [
     "from psycopg.rows import dict_row\n",
     "\n",
-    "connection_kwargs ={\n",
+    "connection_kwargs = {\n",
     "    \"autocommit\": True,\n",
     "    \"prepare_threshold\": 0,\n",
     "}"
@@ -165,7 +165,7 @@
     "    # Example configuration\n",
     "    conninfo=DB_URI,\n",
     "    max_size=20,\n",
-    "    kwargs=connection_kwargs\n",
+    "    kwargs=connection_kwargs,\n",
     ")\n",
     "\n",
     "with pool.connection() as conn:\n",
@@ -393,7 +393,7 @@
     "    # Example configuration\n",
     "    conninfo=DB_URI,\n",
     "    max_size=20,\n",
-    "    kwargs=connection_kwargs\n",
+    "    kwargs=connection_kwargs,\n",
     ") as pool, pool.connection() as conn:\n",
     "    checkpointer = AsyncPostgresSaver(conn)\n",
     "\n",

examples/persistence_redis.ipynb

@@ -530,7 +530,9 @@
     "\n",
     "    @classmethod\n",
     "    @asynccontextmanager\n",
-    "    async def from_conn_info(cls, *, host: str, port: int, db: int) -> AsyncIterator[\"AsyncRedisSaver\"]:\n",
+    "    async def from_conn_info(\n",
+    "        cls, *, host: str, port: int, db: int\n",
+    "    ) -> AsyncIterator[\"AsyncRedisSaver\"]:\n",
     "        conn = None\n",
     "        try:\n",
     "            conn = AsyncRedis(host=host, port=port, db=db)\n",
@@ -887,10 +889,14 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "async with AsyncRedisSaver.from_conn_info(host=\"localhost\", port=6379, db=0) as checkpointer:\n",
+    "async with AsyncRedisSaver.from_conn_info(\n",
+    "    host=\"localhost\", port=6379, db=0\n",
+    ") as checkpointer:\n",
     "    graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)\n",
     "    config = {\"configurable\": {\"thread_id\": \"2\"}}\n",
-    "    res = await graph.ainvoke({\"messages\": [(\"human\", \"what's the weather in nyc\")]}, config)\n",
+    "    res = await graph.ainvoke(\n",
+    "        {\"messages\": [(\"human\", \"what's the weather in nyc\")]}, config\n",
+    "    )\n",
     "\n",
     "    latest_checkpoint = await checkpointer.aget(config)\n",
     "    latest_checkpoint_tuple = await checkpointer.aget_tuple(config)\n",

examples/reflection/reflection.ipynb

@@ -269,7 +269,7 @@
     "class State(TypedDict):\n",
     "    messages: Annotated[list, add_messages]\n",
     "\n",
-    "    \n",
+    "\n",
     "async def generation_node(state: Sequence[BaseMessage]):\n",
     "    return await generate.ainvoke({\"messages\": state})\n",
     "\n",

examples/reflexion/reflexion.ipynb

@@ -392,6 +392,7 @@
     "class State(TypedDict):\n",
     "    messages: Annotated[list, add_messages]\n",
     "\n",
+    "\n",
     "MAX_ITERATIONS = 5\n",
     "builder = StateGraph(State)\n",
     "builder.add_node(\"draft\", first_responder.respond)\n",

examples/streaming-content.ipynb

@@ -68,7 +68,9 @@
     "    # It's completely optional, but useful if you have many functions with similar names\n",
     "    gen = RunnableGenerator(my_generator).with_config(\n",
     "        tags=[\"should_stream\"],\n",
-    "        callbacks=config.get(\"callbacks\", [])  # <-- Propagate callbacks (Python <= 3.10)\n",
+    "        callbacks=config.get(\n",
+    "            \"callbacks\", []\n",
+    "        ),  # <-- Propagate callbacks (Python <= 3.10)\n",
     "    )\n",
     "    async for message in gen.astream(state):\n",
     "        messages.append(message)\n",

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

@@ -30,10 +30,7 @@
    "id": "47f79af8-58d8-4a48-8d9a-88823d88701f",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "%%capture --no-stderr\n",
-    "%pip install -U langgraph openai"
-   ]
+   "source": ["%%capture --no-stderr\n%pip install -U langgraph openai"]
   },
   {
    "cell_type": "code",
@@ -49,18 +46,7 @@
      ]
     }
    ],
-   "source": [
-    "import getpass\n",
-    "import os\n",
-    "\n",
-    "\n",
-    "def _set_env(var: str):\n",
-    "    if not os.environ.get(var):\n",
-    "        os.environ[var] = getpass.getpass(f\"{var}: \")\n",
-    "\n",
-    "\n",
-    "_set_env(\"OPENAI_API_KEY\")"
-   ]
+   "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(\"OPENAI_API_KEY\")"]
   },
   {
    "cell_type": "markdown",
@@ -84,94 +70,7 @@
    "id": "d59234f9-173e-469d-a725-c13e0979663e",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "from openai import AsyncOpenAI\n",
-    "from langchain_core.language_models.chat_models import ChatGenerationChunk\n",
-    "from langchain_core.messages import AIMessageChunk\n",
-    "from langchain_core.runnables.config import (\n",
-    "    ensure_config,\n",
-    "    get_callback_manager_for_config,\n",
-    ")\n",
-    "\n",
-    "openai_client = AsyncOpenAI()\n",
-    "# define tool schema for openai tool calling\n",
-    "\n",
-    "tool = {\n",
-    "    \"type\": \"function\",\n",
-    "    \"function\": {\n",
-    "        \"name\": \"get_items\",\n",
-    "        \"description\": \"Use this tool to look up which items are in the given place.\",\n",
-    "        \"parameters\": {\n",
-    "            \"type\": \"object\",\n",
-    "            \"properties\": {\"place\": {\"type\": \"string\"}},\n",
-    "            \"required\": [\"place\"],\n",
-    "        },\n",
-    "    },\n",
-    "}\n",
-    "\n",
-    "\n",
-    "async def call_model(state, config=None):\n",
-    "    config = ensure_config(config | {\"tags\": [\"agent_llm\"]})\n",
-    "    callback_manager = get_callback_manager_for_config(config)\n",
-    "    messages = state[\"messages\"]\n",
-    "\n",
-    "    llm_run_manager = callback_manager.on_chat_model_start({}, [messages])[0]\n",
-    "    response = await openai_client.chat.completions.create(\n",
-    "        messages=messages, model=\"gpt-3.5-turbo\", tools=[tool], stream=True\n",
-    "    )\n",
-    "\n",
-    "    response_content = \"\"\n",
-    "    role = None\n",
-    "\n",
-    "    tool_call_id = None\n",
-    "    tool_call_function_name = None\n",
-    "    tool_call_function_arguments = \"\"\n",
-    "    async for chunk in response:\n",
-    "        delta = chunk.choices[0].delta\n",
-    "        if delta.role is not None:\n",
-    "            role = delta.role\n",
-    "\n",
-    "        if delta.content:\n",
-    "            response_content += delta.content\n",
-    "            llm_run_manager.on_llm_new_token(delta.content)\n",
-    "\n",
-    "        if delta.tool_calls:\n",
-    "            # note: for simplicity we're only handling a single tool call here\n",
-    "            if delta.tool_calls[0].function.name is not None:\n",
-    "                tool_call_function_name = delta.tool_calls[0].function.name\n",
-    "                tool_call_id = delta.tool_calls[0].id\n",
-    "\n",
-    "            # note: we're wrapping the tools calls in ChatGenerationChunk so that the events from .astream_events in the graph can render tool calls correctly\n",
-    "            tool_call_chunk = ChatGenerationChunk(\n",
-    "                message=AIMessageChunk(\n",
-    "                    content=\"\",\n",
-    "                    additional_kwargs={\"tool_calls\": [delta.tool_calls[0].dict()]},\n",
-    "                )\n",
-    "            )\n",
-    "            llm_run_manager.on_llm_new_token(\"\", chunk=tool_call_chunk)\n",
-    "            tool_call_function_arguments += delta.tool_calls[0].function.arguments\n",
-    "\n",
-    "    if tool_call_function_name is not None:\n",
-    "        tool_calls = [\n",
-    "            {\n",
-    "                \"id\": tool_call_id,\n",
-    "                \"function\": {\n",
-    "                    \"name\": tool_call_function_name,\n",
-    "                    \"arguments\": tool_call_function_arguments,\n",
-    "                },\n",
-    "                \"type\": \"function\",\n",
-    "            }\n",
-    "        ]\n",
-    "    else:\n",
-    "        tool_calls = None\n",
-    "\n",
-    "    response_message = {\n",
-    "        \"role\": role,\n",
-    "        \"content\": response_content,\n",
-    "        \"tool_calls\": tool_calls,\n",
-    "    }\n",
-    "    return {\"messages\": [response_message]}"
-   ]
+   "source": ["from openai import AsyncOpenAI\nfrom langchain_core.language_models.chat_models import ChatGenerationChunk\nfrom langchain_core.messages import AIMessageChunk\nfrom langchain_core.runnables.config import (\n    ensure_config,\n    get_callback_manager_for_config,\n)\n\nopenai_client = AsyncOpenAI()\n# define tool schema for openai tool calling\n\ntool = {\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\nasync def call_model(state, config=None):\n    config = ensure_config(config | {\"tags\": [\"agent_llm\"]})\n    callback_manager = get_callback_manager_for_config(config)\n    messages = state[\"messages\"]\n\n    llm_run_manager = callback_manager.on_chat_model_start({}, [messages])[0]\n    response = await openai_client.chat.completions.create(\n        messages=messages, model=\"gpt-3.5-turbo\", tools=[tool], stream=True\n    )\n\n    response_content = \"\"\n    role = None\n\n    tool_call_id = None\n    tool_call_function_name = None\n    tool_call_function_arguments = \"\"\n    async for chunk in response:\n        delta = chunk.choices[0].delta\n        if delta.role is not None:\n            role = delta.role\n\n        if delta.content:\n            response_content += delta.content\n            llm_run_manager.on_llm_new_token(delta.content)\n\n        if delta.tool_calls:\n            # note: for simplicity we're only handling a single tool call here\n            if delta.tool_calls[0].function.name is not None:\n                tool_call_function_name = delta.tool_calls[0].function.name\n                tool_call_id = delta.tool_calls[0].id\n\n            # note: we're wrapping the tools calls in ChatGenerationChunk so that the events from .astream_events in the graph can render tool calls correctly\n            tool_call_chunk = ChatGenerationChunk(\n                message=AIMessageChunk(\n                    content=\"\",\n                    additional_kwargs={\"tool_calls\": [delta.tool_calls[0].dict()]},\n                )\n            )\n            llm_run_manager.on_llm_new_token(\"\", chunk=tool_call_chunk)\n            tool_call_function_arguments += delta.tool_calls[0].function.arguments\n\n    if tool_call_function_name is not None:\n        tool_calls = [\n            {\n                \"id\": tool_call_id,\n                \"function\": {\n                    \"name\": tool_call_function_name,\n                    \"arguments\": tool_call_function_arguments,\n                },\n                \"type\": \"function\",\n            }\n        ]\n    else:\n        tool_calls = None\n\n    response_message = {\n        \"role\": role,\n        \"content\": response_content,\n        \"tool_calls\": tool_calls,\n    }\n    return {\"messages\": [response_message]}"]
   },
   {
    "cell_type": "markdown",
@@ -187,62 +86,7 @@
    "id": "b90941d8-afe4-42ec-9262-9c3b87c3b1ec",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "import json\n",
-    "from langchain_core.callbacks import adispatch_custom_event\n",
-    "\n",
-    "\n",
-    "async def get_items(place: str) -> str:\n",
-    "    \"\"\"Use this tool to look up which items are in the given place.\"\"\"\n",
-    "\n",
-    "    # this can be replaced with any actual streaming logic that you might have\n",
-    "    def stream(place: str):\n",
-    "        if \"bed\" in place:  # For under the bed\n",
-    "            yield from [\"socks\", \"shoes\", \"dust bunnies\"]\n",
-    "        elif \"shelf\" in place:  # For 'shelf'\n",
-    "            yield from [\"books\", \"penciles\", \"pictures\"]\n",
-    "        else:  # if the agent decides to ask about a different place\n",
-    "            yield \"cat snacks\"\n",
-    "\n",
-    "    tokens = []\n",
-    "    for token in stream(place):\n",
-    "        await adispatch_custom_event(\n",
-    "            # this will allow you to filter events by name\n",
-    "            \"tool_call_token_stream\",\n",
-    "            {\n",
-    "                \"function_name\": \"get_items\",\n",
-    "                \"arguments\": {\"place\": place},\n",
-    "                \"tool_output_token\": token,\n",
-    "            },\n",
-    "            # this will allow you to filter events by tags\n",
-    "            config={\"tags\": [\"tool_call\"]},\n",
-    "        )\n",
-    "        tokens.append(token)\n",
-    "\n",
-    "    return \", \".join(tokens)\n",
-    "\n",
-    "\n",
-    "# define mapping to look up functions when running tools\n",
-    "function_name_to_function = {\"get_items\": get_items}\n",
-    "\n",
-    "\n",
-    "async def call_tools(state):\n",
-    "    messages = state[\"messages\"]\n",
-    "\n",
-    "    tool_call = messages[-1][\"tool_calls\"][0]\n",
-    "    function_name = tool_call[\"function\"][\"name\"]\n",
-    "    function_arguments = tool_call[\"function\"][\"arguments\"]\n",
-    "    arguments = json.loads(function_arguments)\n",
-    "\n",
-    "    function_response = await function_name_to_function[function_name](**arguments)\n",
-    "    tool_message = {\n",
-    "        \"tool_call_id\": tool_call[\"id\"],\n",
-    "        \"role\": \"tool\",\n",
-    "        \"name\": function_name,\n",
-    "        \"content\": function_response,\n",
-    "    }\n",
-    "    return {\"messages\": [tool_message]}"
-   ]
+   "source": ["import json\nfrom langchain_core.callbacks import adispatch_custom_event\n\n\nasync def get_items(place: str) -> str:\n    \"\"\"Use this tool to look up which items are in the given place.\"\"\"\n\n    # this can be replaced with any actual streaming logic that you might have\n    def stream(place: str):\n        if \"bed\" in place:  # For under the bed\n            yield from [\"socks\", \"shoes\", \"dust bunnies\"]\n        elif \"shelf\" in place:  # For 'shelf'\n            yield from [\"books\", \"penciles\", \"pictures\"]\n        else:  # if the agent decides to ask about a different place\n            yield \"cat snacks\"\n\n    tokens = []\n    for token in stream(place):\n        await adispatch_custom_event(\n            # this will allow you to filter events by name\n            \"tool_call_token_stream\",\n            {\n                \"function_name\": \"get_items\",\n                \"arguments\": {\"place\": place},\n                \"tool_output_token\": token,\n            },\n            # this will allow you to filter events by tags\n            config={\"tags\": [\"tool_call\"]},\n        )\n        tokens.append(token)\n\n    return \", \".join(tokens)\n\n\n# define mapping to look up functions when running tools\nfunction_name_to_function = {\"get_items\": get_items}\n\n\nasync def call_tools(state):\n    messages = state[\"messages\"]\n\n    tool_call = messages[-1][\"tool_calls\"][0]\n    function_name = tool_call[\"function\"][\"name\"]\n    function_arguments = tool_call[\"function\"][\"arguments\"]\n    arguments = json.loads(function_arguments)\n\n    function_response = await function_name_to_function[function_name](**arguments)\n    tool_message = {\n        \"tool_call_id\": tool_call[\"id\"],\n        \"role\": \"tool\",\n        \"name\": function_name,\n        \"content\": function_response,\n    }\n    return {\"messages\": [tool_message]}"]
   },
   {
    "cell_type": "markdown",
@@ -258,33 +102,7 @@
    "id": "228260be-1f9a-4195-80e0-9604f8a5dba6",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "import operator\n",
-    "from typing import Annotated, TypedDict, Literal\n",
-    "\n",
-    "from langgraph.graph import StateGraph, END\n",
-    "\n",
-    "\n",
-    "class State(TypedDict):\n",
-    "    messages: Annotated[list, operator.add]\n",
-    "\n",
-    "\n",
-    "def should_continue(state) -> Literal[\"tools\", END]:\n",
-    "    messages = state[\"messages\"]\n",
-    "    last_message = messages[-1]\n",
-    "    if last_message[\"tool_calls\"]:\n",
-    "        return \"tools\"\n",
-    "    return END\n",
-    "\n",
-    "\n",
-    "workflow = StateGraph(State)\n",
-    "workflow.set_entry_point(\"model\")\n",
-    "workflow.add_node(\"model\", call_model)  # i.e. our \"agent\"\n",
-    "workflow.add_node(\"tools\", call_tools)\n",
-    "workflow.add_conditional_edges(\"model\", should_continue)\n",
-    "workflow.add_edge(\"tools\", \"model\")\n",
-    "graph = workflow.compile()"
-   ]
+   "source": ["import operator\nfrom typing import Annotated, TypedDict, Literal\n\nfrom langgraph.graph import StateGraph, END, START\n\n\nclass State(TypedDict):\n    messages: Annotated[list, operator.add]\n\n\ndef should_continue(state) -> Literal[\"tools\", END]:\n    messages = state[\"messages\"]\n    last_message = messages[-1]\n    if last_message[\"tool_calls\"]:\n        return \"tools\"\n    return END\n\n\nworkflow = StateGraph(State)\nworkflow.add_edge(START, \"model\")\nworkflow.add_node(\"model\", call_model)  # i.e. our \"agent\"\nworkflow.add_node(\"tools\", call_tools)\nworkflow.add_conditional_edges(\"model\", should_continue)\nworkflow.add_edge(\"tools\", \"model\")\ngraph = workflow.compile()"]
   },
   {
    "cell_type": "markdown",
@@ -318,14 +136,7 @@
      ]
     }
    ],
-   "source": [
-    "async for event in graph.astream_events(\n",
-    "    {\"messages\": [{\"role\": \"user\", \"content\": \"what's in the bedroom\"}]}, version=\"v2\"\n",
-    "):\n",
-    "    tags = event.get(\"tags\", [])\n",
-    "    if event[\"event\"] == \"on_custom_event\" and \"tool_call\" in tags:\n",
-    "        print(\"Tool token\", event[\"data\"][\"tool_output_token\"])"
-   ]
+   "source": ["async for event in graph.astream_events(\n    {\"messages\": [{\"role\": \"user\", \"content\": \"what's in the bedroom\"}]}, version=\"v2\"\n):\n    tags = event.get(\"tags\", [])\n    if event[\"event\"] == \"on_custom_event\" and \"tool_call\" in tags:\n        print(\"Tool token\", event[\"data\"][\"tool_output_token\"])"]
   }
  ],
  "metadata": {

examples/streaming-tokens-without-langchain.ipynb

@@ -30,10 +30,7 @@
    "id": "47f79af8-58d8-4a48-8d9a-88823d88701f",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "%%capture --no-stderr\n",
-    "%pip install -U langgraph openai"
-   ]
+   "source": ["%%capture --no-stderr\n%pip install -U langgraph openai"]
   },
   {
    "cell_type": "code",
@@ -49,18 +46,7 @@
      ]
     }
    ],
-   "source": [
-    "import getpass\n",
-    "import os\n",
-    "\n",
-    "\n",
-    "def _set_env(var: str):\n",
-    "    if not os.environ.get(var):\n",
-    "        os.environ[var] = getpass.getpass(f\"{var}: \")\n",
-    "\n",
-    "\n",
-    "_set_env(\"OPENAI_API_KEY\")"
-   ]
+   "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(\"OPENAI_API_KEY\")"]
   },
   {
    "cell_type": "markdown",
@@ -84,94 +70,7 @@
    "id": "d59234f9-173e-469d-a725-c13e0979663e",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "from openai import AsyncOpenAI\n",
-    "from langchain_core.language_models.chat_models import ChatGenerationChunk\n",
-    "from langchain_core.messages import AIMessageChunk\n",
-    "from langchain_core.runnables.config import (\n",
-    "    ensure_config,\n",
-    "    get_callback_manager_for_config,\n",
-    ")\n",
-    "\n",
-    "openai_client = AsyncOpenAI()\n",
-    "# define tool schema for openai tool calling\n",
-    "\n",
-    "tool = {\n",
-    "    \"type\": \"function\",\n",
-    "    \"function\": {\n",
-    "        \"name\": \"get_items\",\n",
-    "        \"description\": \"Use this tool to look up which items are in the given place.\",\n",
-    "        \"parameters\": {\n",
-    "            \"type\": \"object\",\n",
-    "            \"properties\": {\"place\": {\"type\": \"string\"}},\n",
-    "            \"required\": [\"place\"],\n",
-    "        },\n",
-    "    },\n",
-    "}\n",
-    "\n",
-    "\n",
-    "async def call_model(state, config=None):\n",
-    "    config = ensure_config(config | {\"tags\": [\"agent_llm\"]})\n",
-    "    callback_manager = get_callback_manager_for_config(config)\n",
-    "    messages = state[\"messages\"]\n",
-    "\n",
-    "    llm_run_manager = callback_manager.on_chat_model_start({}, [messages])[0]\n",
-    "    response = await openai_client.chat.completions.create(\n",
-    "        messages=messages, model=\"gpt-3.5-turbo\", tools=[tool], stream=True\n",
-    "    )\n",
-    "\n",
-    "    response_content = \"\"\n",
-    "    role = None\n",
-    "\n",
-    "    tool_call_id = None\n",
-    "    tool_call_function_name = None\n",
-    "    tool_call_function_arguments = \"\"\n",
-    "    async for chunk in response:\n",
-    "        delta = chunk.choices[0].delta\n",
-    "        if delta.role is not None:\n",
-    "            role = delta.role\n",
-    "\n",
-    "        if delta.content:\n",
-    "            response_content += delta.content\n",
-    "            llm_run_manager.on_llm_new_token(delta.content)\n",
-    "\n",
-    "        if delta.tool_calls:\n",
-    "            # note: for simplicity we're only handling a single tool call here\n",
-    "            if delta.tool_calls[0].function.name is not None:\n",
-    "                tool_call_function_name = delta.tool_calls[0].function.name\n",
-    "                tool_call_id = delta.tool_calls[0].id\n",
-    "\n",
-    "            # note: we're wrapping the tools calls in ChatGenerationChunk so that the events from .astream_events in the graph can render tool calls correctly\n",
-    "            tool_call_chunk = ChatGenerationChunk(\n",
-    "                message=AIMessageChunk(\n",
-    "                    content=\"\",\n",
-    "                    additional_kwargs={\"tool_calls\": [delta.tool_calls[0].dict()]},\n",
-    "                )\n",
-    "            )\n",
-    "            llm_run_manager.on_llm_new_token(\"\", chunk=tool_call_chunk)\n",
-    "            tool_call_function_arguments += delta.tool_calls[0].function.arguments\n",
-    "\n",
-    "    if tool_call_function_name is not None:\n",
-    "        tool_calls = [\n",
-    "            {\n",
-    "                \"id\": tool_call_id,\n",
-    "                \"function\": {\n",
-    "                    \"name\": tool_call_function_name,\n",
-    "                    \"arguments\": tool_call_function_arguments,\n",
-    "                },\n",
-    "                \"type\": \"function\",\n",
-    "            }\n",
-    "        ]\n",
-    "    else:\n",
-    "        tool_calls = None\n",
-    "\n",
-    "    response_message = {\n",
-    "        \"role\": role,\n",
-    "        \"content\": response_content,\n",
-    "        \"tool_calls\": tool_calls,\n",
-    "    }\n",
-    "    return {\"messages\": [response_message]}"
-   ]
+   "source": ["from openai import AsyncOpenAI\nfrom langchain_core.language_models.chat_models import ChatGenerationChunk\nfrom langchain_core.messages import AIMessageChunk\nfrom langchain_core.runnables.config import (\n    ensure_config,\n    get_callback_manager_for_config,\n)\n\nopenai_client = AsyncOpenAI()\n# define tool schema for openai tool calling\n\ntool = {\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\nasync def call_model(state, config=None):\n    config = ensure_config(config | {\"tags\": [\"agent_llm\"]})\n    callback_manager = get_callback_manager_for_config(config)\n    messages = state[\"messages\"]\n\n    llm_run_manager = callback_manager.on_chat_model_start({}, [messages])[0]\n    response = await openai_client.chat.completions.create(\n        messages=messages, model=\"gpt-3.5-turbo\", tools=[tool], stream=True\n    )\n\n    response_content = \"\"\n    role = None\n\n    tool_call_id = None\n    tool_call_function_name = None\n    tool_call_function_arguments = \"\"\n    async for chunk in response:\n        delta = chunk.choices[0].delta\n        if delta.role is not None:\n            role = delta.role\n\n        if delta.content:\n            response_content += delta.content\n            llm_run_manager.on_llm_new_token(delta.content)\n\n        if delta.tool_calls:\n            # note: for simplicity we're only handling a single tool call here\n            if delta.tool_calls[0].function.name is not None:\n                tool_call_function_name = delta.tool_calls[0].function.name\n                tool_call_id = delta.tool_calls[0].id\n\n            # note: we're wrapping the tools calls in ChatGenerationChunk so that the events from .astream_events in the graph can render tool calls correctly\n            tool_call_chunk = ChatGenerationChunk(\n                message=AIMessageChunk(\n                    content=\"\",\n                    additional_kwargs={\"tool_calls\": [delta.tool_calls[0].dict()]},\n                )\n            )\n            llm_run_manager.on_llm_new_token(\"\", chunk=tool_call_chunk)\n            tool_call_function_arguments += delta.tool_calls[0].function.arguments\n\n    if tool_call_function_name is not None:\n        tool_calls = [\n            {\n                \"id\": tool_call_id,\n                \"function\": {\n                    \"name\": tool_call_function_name,\n                    \"arguments\": tool_call_function_arguments,\n                },\n                \"type\": \"function\",\n            }\n        ]\n    else:\n        tool_calls = None\n\n    response_message = {\n        \"role\": role,\n        \"content\": response_content,\n        \"tool_calls\": tool_calls,\n    }\n    return {\"messages\": [response_message]}"]
   },
   {
    "cell_type": "markdown",
@@ -187,41 +86,7 @@
    "id": "b756ea32",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "import json\n",
-    "\n",
-    "\n",
-    "async def get_items(place: str) -> str:\n",
-    "    \"\"\"Use this tool to look up which items are in the given place.\"\"\"\n",
-    "    if \"bed\" in place:  # For under the bed\n",
-    "        return \"socks, shoes and dust bunnies\"\n",
-    "    if \"shelf\" in place:  # For 'shelf'\n",
-    "        return \"books, penciles and pictures\"\n",
-    "    else:  # if the agent decides to ask about a different place\n",
-    "        return \"cat snacks\"\n",
-    "\n",
-    "\n",
-    "# define mapping to look up functions when running tools\n",
-    "function_name_to_function = {\"get_items\": get_items}\n",
-    "\n",
-    "\n",
-    "async def call_tools(state):\n",
-    "    messages = state[\"messages\"]\n",
-    "\n",
-    "    tool_call = messages[-1][\"tool_calls\"][0]\n",
-    "    function_name = tool_call[\"function\"][\"name\"]\n",
-    "    function_arguments = tool_call[\"function\"][\"arguments\"]\n",
-    "    arguments = json.loads(function_arguments)\n",
-    "\n",
-    "    function_response = await function_name_to_function[function_name](**arguments)\n",
-    "    tool_message = {\n",
-    "        \"tool_call_id\": tool_call[\"id\"],\n",
-    "        \"role\": \"tool\",\n",
-    "        \"name\": function_name,\n",
-    "        \"content\": function_response,\n",
-    "    }\n",
-    "    return {\"messages\": [tool_message]}"
-   ]
+   "source": ["import json\n\n\nasync def get_items(place: str) -> str:\n    \"\"\"Use this tool to look up which items are in the given place.\"\"\"\n    if \"bed\" in place:  # For under the bed\n        return \"socks, shoes and dust bunnies\"\n    if \"shelf\" in place:  # For 'shelf'\n        return \"books, penciles and pictures\"\n    else:  # if the agent decides to ask about a different place\n        return \"cat snacks\"\n\n\n# define mapping to look up functions when running tools\nfunction_name_to_function = {\"get_items\": get_items}\n\n\nasync def call_tools(state):\n    messages = state[\"messages\"]\n\n    tool_call = messages[-1][\"tool_calls\"][0]\n    function_name = tool_call[\"function\"][\"name\"]\n    function_arguments = tool_call[\"function\"][\"arguments\"]\n    arguments = json.loads(function_arguments)\n\n    function_response = await function_name_to_function[function_name](**arguments)\n    tool_message = {\n        \"tool_call_id\": tool_call[\"id\"],\n        \"role\": \"tool\",\n        \"name\": function_name,\n        \"content\": function_response,\n    }\n    return {\"messages\": [tool_message]}"]
   },
   {
    "cell_type": "markdown",
@@ -237,33 +102,7 @@
    "id": "228260be-1f9a-4195-80e0-9604f8a5dba6",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "import operator\n",
-    "from typing import Annotated, TypedDict, Literal\n",
-    "\n",
-    "from langgraph.graph import StateGraph, END\n",
-    "\n",
-    "\n",
-    "class State(TypedDict):\n",
-    "    messages: Annotated[list, operator.add]\n",
-    "\n",
-    "\n",
-    "def should_continue(state) -> Literal[\"tools\", END]:\n",
-    "    messages = state[\"messages\"]\n",
-    "    last_message = messages[-1]\n",
-    "    if last_message[\"tool_calls\"]:\n",
-    "        return \"tools\"\n",
-    "    return END\n",
-    "\n",
-    "\n",
-    "workflow = StateGraph(State)\n",
-    "workflow.set_entry_point(\"model\")\n",
-    "workflow.add_node(\"model\", call_model)  # i.e. our \"agent\"\n",
-    "workflow.add_node(\"tools\", call_tools)\n",
-    "workflow.add_conditional_edges(\"model\", should_continue)\n",
-    "workflow.add_edge(\"tools\", \"model\")\n",
-    "graph = workflow.compile()"
-   ]
+   "source": ["import operator\nfrom typing import Annotated, TypedDict, Literal\n\nfrom langgraph.graph import StateGraph, END, START\n\n\nclass State(TypedDict):\n    messages: Annotated[list, operator.add]\n\n\ndef should_continue(state) -> Literal[\"tools\", END]:\n    messages = state[\"messages\"]\n    last_message = messages[-1]\n    if last_message[\"tool_calls\"]:\n        return \"tools\"\n    return END\n\n\nworkflow = StateGraph(State)\nworkflow.add_edge(START, \"model\")\nworkflow.add_node(\"model\", call_model)  # i.e. our \"agent\"\nworkflow.add_node(\"tools\", call_tools)\nworkflow.add_conditional_edges(\"model\", should_continue)\nworkflow.add_edge(\"tools\", \"model\")\ngraph = workflow.compile()"]
   },
   {
    "cell_type": "markdown",
@@ -328,14 +167,7 @@
      ]
     }
    ],
-   "source": [
-    "async for event in graph.astream_events(\n",
-    "    {\"messages\": [{\"role\": \"user\", \"content\": \"what's in the bedroom\"}]}, version=\"v2\"\n",
-    "):\n",
-    "    tags = event.get(\"tags\", [])\n",
-    "    if event[\"event\"] == \"on_chat_model_stream\" and \"agent_llm\" in tags:\n",
-    "        print(\"LLM token\", event[\"data\"][\"chunk\"].dict())"
-   ]
+   "source": ["async for event in graph.astream_events(\n    {\"messages\": [{\"role\": \"user\", \"content\": \"what's in the bedroom\"}]}, version=\"v2\"\n):\n    tags = event.get(\"tags\", [])\n    if event[\"event\"] == \"on_chat_model_stream\" and \"agent_llm\" in tags:\n        print(\"LLM token\", event[\"data\"][\"chunk\"].dict())"]
   },
   {
    "cell_type": "code",
@@ -343,7 +175,7 @@
    "id": "adb0f7bc-6e51-478e-bd32-8f72df072d6c",
    "metadata": {},
    "outputs": [],
-   "source": []
+   "source": [""]
   }
  ],
  "metadata": {

examples/tutorials/rag-agent-testing-local.ipynb

@@ -169,9 +169,7 @@
     "from langchain_core.output_parsers import JsonOutputParser\n",
     "\n",
     "# JSON\n",
-    "llm = ChatOllama(model=\"llama3.1\", \n",
-    "                 format=\"json\", \n",
-    "                 temperature=0)\n",
+    "llm = ChatOllama(model=\"llama3.1\", format=\"json\", temperature=0)\n",
     "\n",
     "\n",
     "prompt = PromptTemplate(\n",
@@ -210,6 +208,7 @@
     "from IPython.display import Image, display\n",
     "from langgraph.graph import START, END, StateGraph\n",
     "\n",
+    "\n",
     "class GraphState(TypedDict):\n",
     "    \"\"\"\n",
     "    Represents the state of our graph.\n",
@@ -356,7 +355,7 @@
     "workflow.add_node(\"web_search\", web_search)  # web search\n",
     "\n",
     "# Build graph\n",
-    "workflow.set_entry_point(\"retrieve\")\n",
+    "workflow.add_edge(START, retrieve)\n",
     "workflow.add_edge(\"retrieve\", \"grade_documents\")\n",
     "workflow.add_conditional_edges(\n",
     "    \"grade_documents\",\n",
@@ -381,21 +380,22 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import uuid \n",
+    "import uuid\n",
+    "\n",
     "\n",
     "def predict_custom_agent_answer(example: dict):\n",
-    "    \n",
     "    config = {\"configurable\": {\"thread_id\": str(uuid.uuid4())}}\n",
-    "    \n",
+    "\n",
     "    state_dict = custom_graph.invoke(\n",
     "        {\"question\": example[\"input\"], \"steps\": []}, config\n",
     "    )\n",
-    "    \n",
+    "\n",
     "    return {\"response\": state_dict[\"generation\"], \"steps\": state_dict[\"steps\"]}\n",
     "\n",
+    "\n",
     "example = {\"input\": \"What are the types of agent memory?\"}\n",
-    "#response = predict_custom_agent_answer(example)\n",
-    "#response"
+    "# response = predict_custom_agent_answer(example)\n",
+    "# response"
    ]
   },
   {
@@ -544,6 +544,7 @@
     "    \"generate_answer\",\n",
     "]\n",
     "\n",
+    "\n",
     "def check_trajectory_custom(root_run: Run, example: Example) -> dict:\n",
     "    \"\"\"\n",
     "    Check if all expected tools are called in exact order and without any additional tool calls.\n",

examples/tutorials/tool-calling-agent-local.ipynb

@@ -134,6 +134,7 @@
     "        for d in web_results\n",
     "    ]\n",
     "\n",
+    "\n",
     "# Tool list\n",
     "tools = [retrieve_documents, web_search]"
    ]
@@ -152,9 +153,11 @@
     "from langgraph.graph.message import AnyMessage, add_messages\n",
     "from typing_extensions import TypedDict\n",
     "\n",
+    "\n",
     "class State(TypedDict):\n",
     "    messages: Annotated[list[AnyMessage], add_messages]\n",
     "\n",
+    "\n",
     "class Assistant:\n",
     "    def __init__(self, runnable: Runnable):\n",
     "        \"\"\"\n",
@@ -291,6 +294,7 @@
    "source": [
     "import uuid\n",
     "\n",
+    "\n",
     "def predict_react_agent_answer(example: dict):\n",
     "    \"\"\"Use this for answer evaluation\"\"\"\n",
     "\n",

examples/web-navigation/web_voyager.ipynb

@@ -457,13 +457,13 @@
    "source": [
     "from langchain_core.runnables import RunnableLambda\n",
     "\n",
-    "from langgraph.graph import END, StateGraph\n",
+    "from langgraph.graph import END, START, StateGraph\n",
     "\n",
     "graph_builder = StateGraph(AgentState)\n",
     "\n",
     "\n",
     "graph_builder.add_node(\"agent\", agent)\n",
-    "graph_builder.set_entry_point(\"agent\")\n",
+    "graph_builder.add_edge(START, \"agent\")\n",
     "\n",
     "graph_builder.add_node(\"update_scratchpad\", update_scratchpad)\n",
     "graph_builder.add_edge(\"update_scratchpad\", \"agent\")\n",

libs/checkpoint-postgres/README.md

@@ -2,6 +2,10 @@
 
 Implementation of LangGraph CheckpointSaver that uses Postgres.
 
+## Dependencies
+
+By default `langgraph-checkpoint-postgres` installs `psycopg` (Psycopg 3) without any extras. However, you can choose a specific installation that best suits your needs [here](https://www.psycopg.org/psycopg3/docs/basic/install.html) (for example, `psycopg[binary]`).
+
 ## Usage
 
 > [!IMPORTANT]

libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py

@@ -1,12 +1,13 @@
 import threading
 from contextlib import contextmanager
-from typing import Any, Iterator, List, Optional
+from typing import Any, Iterator, List, Optional, Union
 
 from langchain_core.runnables import RunnableConfig
 from psycopg import Connection, Cursor, Pipeline
 from psycopg.errors import UndefinedTable
 from psycopg.rows import dict_row
 from psycopg.types.json import Jsonb
+from psycopg_pool import ConnectionPool
 
 from langgraph.checkpoint.base import (
     ChannelVersions,
@@ -21,16 +22,32 @@ from langgraph.checkpoint.postgres.base import (
 from langgraph.checkpoint.serde.base import SerializerProtocol
 
 
+@contextmanager
+def _get_connection(conn: Union[Connection, ConnectionPool]) -> Iterator[Connection]:
+    if isinstance(conn, Connection):
+        yield conn
+    elif isinstance(conn, ConnectionPool):
+        with conn.connection() as conn:
+            yield conn
+    else:
+        raise TypeError(f"Invalid connection type: {type(conn)}")
+
+
 class PostgresSaver(BasePostgresSaver):
     lock: threading.Lock
 
     def __init__(
         self,
-        conn: Connection,
+        conn: Union[Connection, ConnectionPool],
         pipe: Optional[Pipeline] = None,
         serde: Optional[SerializerProtocol] = None,
     ) -> None:
         super().__init__(serde=serde)
+        if isinstance(conn, ConnectionPool) and pipe is not None:
+            raise ValueError(
+                "Pipeline should be used only with a single Connection, not ConnectionPool."
+            )
+
         self.conn = conn
         self.pipe = pipe
         self.lock = threading.Lock()
@@ -65,22 +82,21 @@ class PostgresSaver(BasePostgresSaver):
         already exist and runs database migrations. It MUST be called directly by the user
         the first time checkpointer is used.
         """
-        with self.lock:
-            with self.conn.cursor(binary=True, row_factory=dict_row) as cur:
-                try:
-                    version = cur.execute(
-                        "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
-                    ).fetchone()["v"]
-                except UndefinedTable:
-                    version = -1
-                for v, migration in zip(
-                    range(version + 1, len(self.MIGRATIONS)),
-                    self.MIGRATIONS[version + 1 :],
-                ):
-                    cur.execute(migration)
-                    cur.execute(f"INSERT INTO checkpoint_migrations (v) VALUES ({v})")
-            if self.pipe:
-                self.pipe.sync()
+        with self._cursor() as cur:
+            try:
+                version = cur.execute(
+                    "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
+                ).fetchone()["v"]
+            except UndefinedTable:
+                version = -1
+            for v, migration in zip(
+                range(version + 1, len(self.MIGRATIONS)),
+                self.MIGRATIONS[version + 1 :],
+            ):
+                cur.execute(migration)
+                cur.execute(f"INSERT INTO checkpoint_migrations (v) VALUES ({v})")
+        if self.pipe:
+            self.pipe.sync()
 
     def list(
         self,
@@ -333,23 +349,24 @@ class PostgresSaver(BasePostgresSaver):
 
     @contextmanager
     def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor]:
-        if self.pipe:
-            # a connection in pipeline mode can be used concurrently
-            # in multiple threads/coroutines, but only one cursor can be
-            # used at a time
-            try:
-                with self.conn.cursor(binary=True, row_factory=dict_row) as cur:
+        with _get_connection(self.conn) as conn:
+            if self.pipe:
+                # a connection in pipeline mode can be used concurrently
+                # in multiple threads/coroutines, but only one cursor can be
+                # used at a time
+                try:
+                    with conn.cursor(binary=True, row_factory=dict_row) as cur:
+                        yield cur
+                finally:
+                    if pipeline:
+                        self.pipe.sync()
+            elif pipeline:
+                # a connection not in pipeline mode can only be used by one
+                # thread/coroutine at a time, so we acquire a lock
+                with self.lock, conn.pipeline(), conn.cursor(
+                    binary=True, row_factory=dict_row
+                ) as cur:
+                    yield cur
+            else:
+                with self.lock, conn.cursor(binary=True, row_factory=dict_row) as cur:
                     yield cur
-            finally:
-                if pipeline:
-                    self.pipe.sync()
-        elif pipeline:
-            # a connection not in pipeline mode can only be used by one
-            # thread/coroutine at a time, so we acquire a lock
-            with self.lock, self.conn.pipeline(), self.conn.cursor(
-                binary=True, row_factory=dict_row
-            ) as cur:
-                yield cur
-        else:
-            with self.lock, self.conn.cursor(binary=True, row_factory=dict_row) as cur:
-                yield cur

libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py

@@ -1,12 +1,13 @@
 import asyncio
 from contextlib import asynccontextmanager
-from typing import Any, AsyncIterator, Optional
+from typing import Any, AsyncIterator, Optional, Union
 
 from langchain_core.runnables import RunnableConfig
 from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline
 from psycopg.errors import UndefinedTable
 from psycopg.rows import dict_row
 from psycopg.types.json import Jsonb
+from psycopg_pool import AsyncConnectionPool
 
 from langgraph.checkpoint.base import (
     ChannelVersions,
@@ -19,16 +20,34 @@ from langgraph.checkpoint.postgres.base import BasePostgresSaver
 from langgraph.checkpoint.serde.base import SerializerProtocol
 
 
+@asynccontextmanager
+async def _get_connection(
+    conn: Union[AsyncConnection, AsyncConnectionPool],
+) -> AsyncIterator[AsyncConnection]:
+    if isinstance(conn, AsyncConnection):
+        yield conn
+    elif isinstance(conn, AsyncConnectionPool):
+        async with conn.connection() as conn:
+            yield conn
+    else:
+        raise TypeError(f"Invalid connection type: {type(conn)}")
+
+
 class AsyncPostgresSaver(BasePostgresSaver):
     lock: asyncio.Lock
 
     def __init__(
         self,
-        conn: AsyncConnection,
+        conn: Union[AsyncConnection, AsyncConnectionPool],
         pipe: Optional[AsyncPipeline] = None,
         serde: Optional[SerializerProtocol] = None,
     ) -> None:
         super().__init__(serde=serde)
+        if isinstance(conn, AsyncConnectionPool) and pipe is not None:
+            raise ValueError(
+                "Pipeline should be used only with a single AsyncConnection, not AsyncConnectionPool."
+            )
+
         self.conn = conn
         self.pipe = pipe
         self.lock = asyncio.Lock()
@@ -63,25 +82,22 @@ class AsyncPostgresSaver(BasePostgresSaver):
         already exist and runs database migrations. It MUST be called directly by the user
         the first time checkpointer is used.
         """
-        async with self.lock:
-            async with self.conn.cursor(binary=True, row_factory=dict_row) as cur:
-                try:
-                    results = await cur.execute(
-                        "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
-                    )
-                    version = (await results.fetchone())["v"]
-                except UndefinedTable:
-                    version = -1
-                for v, migration in zip(
-                    range(version + 1, len(self.MIGRATIONS)),
-                    self.MIGRATIONS[version + 1 :],
-                ):
-                    await cur.execute(migration)
-                    await cur.execute(
-                        f"INSERT INTO checkpoint_migrations (v) VALUES ({v})"
-                    )
-            if self.pipe:
-                await self.pipe.sync()
+        async with self._cursor() as cur:
+            try:
+                results = await cur.execute(
+                    "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
+                )
+                version = (await results.fetchone())["v"]
+            except UndefinedTable:
+                version = -1
+            for v, migration in zip(
+                range(version + 1, len(self.MIGRATIONS)),
+                self.MIGRATIONS[version + 1 :],
+            ):
+                await cur.execute(migration)
+                await cur.execute(f"INSERT INTO checkpoint_migrations (v) VALUES ({v})")
+        if self.pipe:
+            await self.pipe.sync()
 
     async def alist(
         self,
@@ -290,25 +306,26 @@ class AsyncPostgresSaver(BasePostgresSaver):
 
     @asynccontextmanager
     async def _cursor(self, *, pipeline: bool = False) -> AsyncIterator[AsyncCursor]:
-        if self.pipe:
-            # a connection in pipeline mode can be used concurrently
-            # in multiple threads/coroutines, but only one cursor can be
-            # used at a time
-            try:
-                async with self.conn.cursor(binary=True, row_factory=dict_row) as cur:
+        async with _get_connection(self.conn) as conn:
+            if self.pipe:
+                # a connection in pipeline mode can be used concurrently
+                # in multiple threads/coroutines, but only one cursor can be
+                # used at a time
+                try:
+                    async with conn.cursor(binary=True, row_factory=dict_row) as cur:
+                        yield cur
+                finally:
+                    if pipeline:
+                        await self.pipe.sync()
+            elif pipeline:
+                # a connection not in pipeline mode can only be used by one
+                # thread/coroutine at a time, so we acquire a lock
+                async with self.lock, conn.pipeline(), conn.cursor(
+                    binary=True, row_factory=dict_row
+                ) as cur:
+                    yield cur
+            else:
+                async with self.lock, conn.cursor(
+                    binary=True, row_factory=dict_row
+                ) as cur:
                     yield cur
-            finally:
-                if pipeline:
-                    await self.pipe.sync()
-        elif pipeline:
-            # a connection not in pipeline mode can only be used by one
-            # thread/coroutine at a time, so we acquire a lock
-            async with self.lock, self.conn.pipeline(), self.conn.cursor(
-                binary=True, row_factory=dict_row
-            ) as cur:
-                yield cur
-        else:
-            async with self.lock, self.conn.cursor(
-                binary=True, row_factory=dict_row
-            ) as cur:
-                yield cur

libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py

@@ -255,9 +255,13 @@ class BasePostgresSaver(BaseCheckpointSaver):
         if config:
             wheres.append("thread_id = %s ")
             param_values.append(config["configurable"]["thread_id"])
-            checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
-            wheres.append("checkpoint_ns = %s")
-            param_values.append(checkpoint_ns)
+            if checkpoint_ns := config["configurable"].get("checkpoint_ns"):
+                wheres.append("checkpoint_ns = %s")
+                param_values.append(checkpoint_ns)
+
+            if checkpoint_id := get_checkpoint_id(config):
+                wheres.append("checkpoint_id = %s ")
+                param_values.append(checkpoint_id)
 
         # construct predicate for metadata filter
         if filter:

libs/checkpoint-postgres/poetry.lock

@@ -266,7 +266,7 @@ tenacity = ">=8.1.0,<8.4.0 || >8.4.0,<9.0.0"
 
 [[package]]
 name = "langgraph-checkpoint"
-version = "1.0.1"
+version = "1.0.6"
 description = "Library with base interfaces for LangGraph checkpoint savers."
 optional = false
 python-versions = "^3.9.0,<4.0"
@@ -969,4 +969,4 @@ watchmedo = ["PyYAML (>=3.10)"]
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.9.0,<4.0"
-content-hash = "422b6d716b86db072ea3a612287ad20ff5700c18f22d9e9d59cc4e198514519d"
+content-hash = "cfa417cbaf126fa847242ef81e9af3e4d13e93f25631bb4c007cdc69926ff3b8"

libs/checkpoint-postgres/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph-checkpoint-postgres"
-version = "1.0.3"
+version = "1.0.4"
 description = "Library with a Postgres implementation of LangGraph checkpoint saver."
 authors = []
 license = "MIT"
@@ -12,7 +12,8 @@ packages = [{ include = "langgraph" }]
 python = "^3.9.0,<4.0"
 langgraph-checkpoint = "^1.0.1"
 orjson = ">=3.10.1"
-psycopg = {extras = ["binary"], version = ">=3.1.19"}
+psycopg = "^3.0.0"
+psycopg-pool = "^3.0.0"
 
 [tool.poetry.group.dev.dependencies]
 ruff = "^0.1.4"
@@ -23,7 +24,7 @@ pytest-asyncio = "^0.21.1"
 pytest-mock = "^3.11.1"
 pytest-watch = "^4.2.0"
 mypy = "^1.10.0"
-psycopg-pool = "^3.2.2"
+psycopg = {extras = ["binary"], version = ">=3.0.0"}
 langgraph-checkpoint = {path = "../checkpoint", develop = true}
 
 [tool.pytest.ini_options]

libs/checkpoint-postgres/tests/test_async.py

@@ -87,29 +87,15 @@ class TestAsyncPostgresSaver:
             search_results_4 = [c async for c in saver.alist(None, filter=query_4)]
             assert len(search_results_4) == 0
 
-            # search by config (defaults to root graph checkpoints)
+            # search by config (defaults to checkpoints across all namespaces)
             search_results_5 = [
                 c
                 async for c in saver.alist({"configurable": {"thread_id": "thread-2"}})
             ]
-            assert len(search_results_5) == 1
-            assert search_results_5[0].config["configurable"]["checkpoint_ns"] == ""
-
-            # search by config and checkpoint_ns
-            search_results_6 = [
-                c
-                async for c in saver.alist(
-                    {
-                        "configurable": {
-                            "thread_id": "thread-2",
-                            "checkpoint_ns": "inner",
-                        }
-                    }
-                )
-            ]
-            assert len(search_results_6) == 1
-            assert (
-                search_results_6[0].config["configurable"]["checkpoint_ns"] == "inner"
-            )
+            assert len(search_results_5) == 2
+            assert {
+                search_results_5[0].config["configurable"]["checkpoint_ns"],
+                search_results_5[1].config["configurable"]["checkpoint_ns"],
+            } == {"", "inner"}
 
             # TODO: test before and limit params

libs/checkpoint-postgres/tests/test_sync.py

@@ -88,27 +88,14 @@ class TestPostgresSaver:
             search_results_4 = list(saver.list(None, filter=query_4))
             assert len(search_results_4) == 0
 
-            # search by config (defaults to root graph checkpoints)
+            # search by config (defaults to checkpoints across all namespaces)
             search_results_5 = list(
                 saver.list({"configurable": {"thread_id": "thread-2"}})
             )
-            assert len(search_results_5) == 1
-            assert search_results_5[0].config["configurable"]["checkpoint_ns"] == ""
-
-            # search by config and checkpoint_ns
-            search_results_6 = list(
-                saver.list(
-                    {
-                        "configurable": {
-                            "thread_id": "thread-2",
-                            "checkpoint_ns": "inner",
-                        }
-                    }
-                )
-            )
-            assert len(search_results_6) == 1
-            assert (
-                search_results_6[0].config["configurable"]["checkpoint_ns"] == "inner"
-            )
+            assert len(search_results_5) == 2
+            assert {
+                search_results_5[0].config["configurable"]["checkpoint_ns"],
+                search_results_5[1].config["configurable"]["checkpoint_ns"],
+            } == {"", "inner"}
 
             # TODO: test before and limit params

libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/utils.py

@@ -70,9 +70,13 @@ def search_where(
     if config is not None:
         wheres.append("thread_id = ?")
         param_values.append(config["configurable"]["thread_id"])
-        checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
-        wheres.append("checkpoint_ns = ?")
-        param_values.append(checkpoint_ns)
+        if checkpoint_ns := config["configurable"].get("checkpoint_ns"):
+            wheres.append("checkpoint_ns = ?")
+            param_values.append(checkpoint_ns)
+
+        if checkpoint_id := get_checkpoint_id(config):
+            wheres.append("checkpoint_id = ?")
+            param_values.append(checkpoint_id)
 
     # construct predicate for metadata filter
     if filter:

libs/checkpoint-sqlite/tests/test_aiosqlite.py

@@ -84,29 +84,15 @@ class TestAsyncSqliteSaver:
             search_results_4 = [c async for c in saver.alist(None, filter=query_4)]
             assert len(search_results_4) == 0
 
-            # search by config (defaults to root graph checkpoints)
+            # search by config (defaults to checkpoints across all namespaces)
             search_results_5 = [
                 c
                 async for c in saver.alist({"configurable": {"thread_id": "thread-2"}})
             ]
-            assert len(search_results_5) == 1
-            assert search_results_5[0].config["configurable"]["checkpoint_ns"] == ""
-
-            # search by config and checkpoint_ns
-            search_results_6 = [
-                c
-                async for c in saver.alist(
-                    {
-                        "configurable": {
-                            "thread_id": "thread-2",
-                            "checkpoint_ns": "inner",
-                        }
-                    }
-                )
-            ]
-            assert len(search_results_6) == 1
-            assert (
-                search_results_6[0].config["configurable"]["checkpoint_ns"] == "inner"
-            )
+            assert len(search_results_5) == 2
+            assert {
+                search_results_5[0].config["configurable"]["checkpoint_ns"],
+                search_results_5[1].config["configurable"]["checkpoint_ns"],
+            } == {"", "inner"}
 
             # TODO: test before and limit params

libs/checkpoint-sqlite/tests/test_sqlite.py

@@ -87,28 +87,15 @@ class TestSqliteSaver:
             search_results_4 = list(saver.list(None, filter=query_4))
             assert len(search_results_4) == 0
 
-            # search by config (defaults to root graph checkpoints)
+            # search by config (defaults to checkpoints across all namespaces)
             search_results_5 = list(
                 saver.list({"configurable": {"thread_id": "thread-2"}})
             )
-            assert len(search_results_5) == 1
-            assert search_results_5[0].config["configurable"]["checkpoint_ns"] == ""
-
-            # search by config and checkpoint_ns
-            search_results_6 = list(
-                saver.list(
-                    {
-                        "configurable": {
-                            "thread_id": "thread-2",
-                            "checkpoint_ns": "inner",
-                        }
-                    }
-                )
-            )
-            assert len(search_results_6) == 1
-            assert (
-                search_results_6[0].config["configurable"]["checkpoint_ns"] == "inner"
-            )
+            assert len(search_results_5) == 2
+            assert {
+                search_results_5[0].config["configurable"]["checkpoint_ns"],
+                search_results_5[1].config["configurable"]["checkpoint_ns"],
+            } == {"", "inner"}
 
             # TODO: test before and limit params
 

libs/checkpoint/langgraph/checkpoint/memory/__init__.py

@@ -177,62 +177,77 @@ class MemorySaver(
             Iterator[CheckpointTuple]: An iterator of matching checkpoint tuples.
         """
         thread_ids = (config["configurable"]["thread_id"],) if config else self.storage
-        checkpoint_ns = (
-            config["configurable"].get("checkpoint_ns", "") if config else ""
+        config_checkpoint_ns = (
+            config["configurable"].get("checkpoint_ns") if config else None
         )
+        config_checkpoint_id = get_checkpoint_id(config) if config else None
         for thread_id in thread_ids:
-            for checkpoint_id, (checkpoint, metadata_b, parent_checkpoint_id) in sorted(
-                self.storage[thread_id][checkpoint_ns].items(),
-                key=lambda x: x[0],
-                reverse=True,
-            ):
-                # filter by checkpoint ID
-                if (
-                    before
-                    and (before_checkpoint_id := get_checkpoint_id(before))
-                    and checkpoint_id >= before_checkpoint_id
-                ):
+            for checkpoint_ns in self.storage[thread_id].keys():
+                if config_checkpoint_ns and checkpoint_ns != config_checkpoint_ns:
                     continue
 
-                # filter by metadata
-                metadata = self.serde.loads_typed(metadata_b)
-                if filter and not all(
-                    query_value == metadata[query_key]
-                    for query_key, query_value in filter.items()
+                for checkpoint_id, (
+                    checkpoint,
+                    metadata_b,
+                    parent_checkpoint_id,
+                ) in sorted(
+                    self.storage[thread_id][checkpoint_ns].items(),
+                    key=lambda x: x[0],
+                    reverse=True,
                 ):
-                    continue
+                    # filter by checkpoint ID from config
+                    if config_checkpoint_id and checkpoint_id != config_checkpoint_id:
+                        continue
 
-                # limit search results
-                if limit is not None and limit <= 0:
-                    break
-                elif limit is not None:
-                    limit -= 1
+                    # filter by checkpoint ID from `before` config
+                    if (
+                        before
+                        and (before_checkpoint_id := get_checkpoint_id(before))
+                        and checkpoint_id >= before_checkpoint_id
+                    ):
+                        continue
 
-                writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)].values()
+                    # filter by metadata
+                    metadata = self.serde.loads_typed(metadata_b)
+                    if filter and not all(
+                        query_value == metadata.get(query_key)
+                        for query_key, query_value in filter.items()
+                    ):
+                        continue
 
-                yield CheckpointTuple(
-                    config={
-                        "configurable": {
-                            "thread_id": thread_id,
-                            "checkpoint_ns": checkpoint_ns,
-                            "checkpoint_id": checkpoint_id,
+                    # limit search results
+                    if limit is not None and limit <= 0:
+                        break
+                    elif limit is not None:
+                        limit -= 1
+
+                    writes = self.writes[
+                        (thread_id, checkpoint_ns, checkpoint_id)
+                    ].values()
+
+                    yield CheckpointTuple(
+                        config={
+                            "configurable": {
+                                "thread_id": thread_id,
+                                "checkpoint_ns": checkpoint_ns,
+                                "checkpoint_id": checkpoint_id,
+                            }
+                        },
+                        checkpoint=self.serde.loads_typed(checkpoint),
+                        metadata=metadata,
+                        parent_config={
+                            "configurable": {
+                                "thread_id": thread_id,
+                                "checkpoint_ns": checkpoint_ns,
+                                "checkpoint_id": parent_checkpoint_id,
+                            }
                         }
-                    },
-                    checkpoint=self.serde.loads_typed(checkpoint),
-                    metadata=metadata,
-                    parent_config={
-                        "configurable": {
-                            "thread_id": thread_id,
-                            "checkpoint_ns": checkpoint_ns,
-                            "checkpoint_id": parent_checkpoint_id,
-                        }
-                    }
-                    if parent_checkpoint_id
-                    else None,
-                    pending_writes=[
-                        (id, c, self.serde.loads_typed(v)) for id, c, v in writes
-                    ],
-                )
+                        if parent_checkpoint_id
+                        else None,
+                        pending_writes=[
+                            (id, c, self.serde.loads_typed(v)) for id, c, v in writes
+                        ],
+                    )
 
     def put(
         self,
@@ -338,7 +353,14 @@ class MemorySaver(
         """
         loop = asyncio.get_running_loop()
         iter = await loop.run_in_executor(
-            None, partial(self.list, before=before, limit=limit, filter=filter), config
+            None,
+            partial(
+                self.list,
+                before=before,
+                limit=limit,
+                filter=filter,
+            ),
+            config,
         )
         while True:
             # handling StopIteration exception inside coroutine won't work

libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py

@@ -7,6 +7,7 @@ import re
 from collections import deque
 from datetime import date, datetime, time, timedelta, timezone
 from enum import Enum
+from inspect import isclass
 from ipaddress import (
     IPv4Address,
     IPv4Interface,
@@ -50,9 +51,13 @@ class JsonPlusSerializer(SerializerProtocol):
         if isinstance(obj, Serializable):
             return obj.to_json()
         elif hasattr(obj, "model_dump") and callable(obj.model_dump):
-            return self._encode_constructor_args(obj.__class__, kwargs=obj.model_dump())
+            return self._encode_constructor_args(
+                obj.__class__, method="model_construct", kwargs=obj.model_dump()
+            )
         elif hasattr(obj, "dict") and callable(obj.dict):
-            return self._encode_constructor_args(obj.__class__, kwargs=obj.dict())
+            return self._encode_constructor_args(
+                obj.__class__, method="construct", kwargs=obj.dict()
+            )
         elif isinstance(obj, pathlib.Path):
             return self._encode_constructor_args(pathlib.Path, args=obj.parts)
         elif isinstance(obj, re.Pattern):
@@ -104,14 +109,14 @@ class JsonPlusSerializer(SerializerProtocol):
             return self._encode_constructor_args(obj.__class__, args=[obj.value])
         elif isinstance(obj, SendProtocol):
             return self._encode_constructor_args(
-                obj.__class__, kwargs={"node": obj.node, "arg": obj.arg}
+                obj.__class__, kwargs={"node": obj.node, "arg": obj.arg, "id": obj.id}
             )
         elif isinstance(obj, (bytes, bytearray)):
             return self._encode_constructor_args(
                 obj.__class__, method="fromhex", args=[obj.hex()]
             )
         elif isinstance(obj, BaseException):
-            return self._encode_constructor_args(obj.__class__, args=obj.args)
+            return repr(obj)
         else:
             raise TypeError(
                 f"Object of type {obj.__class__.__name__} is not JSON serializable"
@@ -135,6 +140,8 @@ class JsonPlusSerializer(SerializerProtocol):
                     method = getattr(cls, value["method"])
                 else:
                     method = cls
+                if isclass(method) and issubclass(method, BaseException):
+                    return None
                 if value["args"] and value["kwargs"]:
                     return method(*value["args"], **value["kwargs"])
                 elif value["args"]:
@@ -143,7 +150,7 @@ class JsonPlusSerializer(SerializerProtocol):
                     return method(**value["kwargs"])
                 else:
                     return method()
-            except (ImportError, AttributeError):
+            except (ImportError, AttributeError, TypeError):
                 return None
 
         return LC_REVIVER(value)

libs/checkpoint/langgraph/checkpoint/serde/types.py

@@ -57,6 +57,7 @@ class SendProtocol(Protocol):
     # Mirrors langgraph.constants.Send
     node: str
     arg: Any
+    id: str
 
     def __hash__(self) -> int:
         ...

libs/checkpoint/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph-checkpoint"
-version = "1.0.4"
+version = "1.0.6"
 description = "Library with base interfaces for LangGraph checkpoint savers."
 authors = []
 license = "MIT"

libs/checkpoint/tests/test_jsonplus.py

@@ -122,7 +122,7 @@ def test_serde_jsonplus() -> None:
 
     assert dumped == (
         "json",
-        b"""{"path": {"lc": 2, "type": "constructor", "id": ["pathlib", "Path"], "method": null, "args": ["foo", "bar"], "kwargs": {}}, "re": {"lc": 2, "type": "constructor", "id": ["re", "compile"], "method": null, "args": ["foo", 48], "kwargs": {}}, "decimal": {"lc": 2, "type": "constructor", "id": ["decimal", "Decimal"], "method": null, "args": ["1.10101"], "kwargs": {}}, "ip4": {"lc": 2, "type": "constructor", "id": ["ipaddress", "IPv4Address"], "method": null, "args": ["192.168.0.1"], "kwargs": {}}, "deque": {"lc": 2, "type": "constructor", "id": ["collections", "deque"], "method": null, "args": [[1, 2, 3]], "kwargs": {}}, "tzn": {"lc": 2, "type": "constructor", "id": ["zoneinfo", "ZoneInfo"], "method": null, "args": ["America/New_York"], "kwargs": {}}, "date": {"lc": 2, "type": "constructor", "id": ["datetime", "date"], "method": null, "args": [2024, 4, 19], "kwargs": {}}, "time": {"lc": 2, "type": "constructor", "id": ["datetime", "time"], "method": null, "args": [23, 4, 57, 51022, {"lc": 2, "type": "constructor", "id": ["datetime", "timezone"], "method": null, "args": [{"lc": 2, "type": "constructor", "id": ["datetime", "timedelta"], "method": null, "args": [0, 86340, 0], "kwargs": {}}], "kwargs": {}}], "kwargs": {"fold": 0}}, "uid": {"lc": 2, "type": "constructor", "id": ["uuid", "UUID"], "method": null, "args": ["00000000000000000000000000000001"], "kwargs": {}}, "timestamp": {"lc": 2, "type": "constructor", "id": ["datetime", "datetime"], "method": "fromisoformat", "args": ["2024-04-19T23:04:57.051022+23:59"], "kwargs": {}}, "my_slotted_class": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyDataclassWSlots"], "method": null, "args": [], "kwargs": {"foo": "bar", "bar": 2}}, "my_dataclass": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyDataclass"], "method": null, "args": [], "kwargs": {"foo": "foo", "bar": 1}}, "my_enum": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyEnum"], "method": null, "args": ["foo"], "kwargs": {}}, "my_pydantic": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyPydantic"], "method": null, "args": [], "kwargs": {"foo": "foo", "bar": 1}}, "my_funny_pydantic": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyFunnyPydantic"], "method": null, "args": [], "kwargs": {"foo": "foo", "bar": 1}}, "person": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "Person"], "method": null, "args": [], "kwargs": {"name": "foo"}}, "a_bool": true, "a_none": null, "a_str": "foo", "a_str_nuc": "foo\\u0000", "a_str_uc": "foo \xe2\x9b\xb0\xef\xb8\x8f", "a_str_ucuc": "foo \xe2\x9b\xb0\xef\xb8\x8f\\u0000", "a_str_ucucuc": "foo \\\\u26f0\\\\ufe0f", "text": ["Hello", "Python", "Surrogate", "Example", "String", "With", "Surrogates", "Embedded", "In", "The", "Text", "\xe6\x94\xb6\xe8\x8a\xb1\xf0\x9f\x99\x84\xc2\xb7\xe5\x88\xb0"], "an_int": 1, "a_float": 1.1, "runnable_map": {"lc": 1, "type": "constructor", "id": ["langchain", "schema", "runnable", "RunnableParallel"], "kwargs": {"steps__": {}}, "name": "RunnableParallel<>", "graph": {"nodes": [{"id": 0, "type": "schema", "data": "Parallel<>Input"}, {"id": 1, "type": "schema", "data": "Parallel<>Output"}], "edges": []}}}""",
+        b"""{"path": {"lc": 2, "type": "constructor", "id": ["pathlib", "Path"], "method": null, "args": ["foo", "bar"], "kwargs": {}}, "re": {"lc": 2, "type": "constructor", "id": ["re", "compile"], "method": null, "args": ["foo", 48], "kwargs": {}}, "decimal": {"lc": 2, "type": "constructor", "id": ["decimal", "Decimal"], "method": null, "args": ["1.10101"], "kwargs": {}}, "ip4": {"lc": 2, "type": "constructor", "id": ["ipaddress", "IPv4Address"], "method": null, "args": ["192.168.0.1"], "kwargs": {}}, "deque": {"lc": 2, "type": "constructor", "id": ["collections", "deque"], "method": null, "args": [[1, 2, 3]], "kwargs": {}}, "tzn": {"lc": 2, "type": "constructor", "id": ["zoneinfo", "ZoneInfo"], "method": null, "args": ["America/New_York"], "kwargs": {}}, "date": {"lc": 2, "type": "constructor", "id": ["datetime", "date"], "method": null, "args": [2024, 4, 19], "kwargs": {}}, "time": {"lc": 2, "type": "constructor", "id": ["datetime", "time"], "method": null, "args": [23, 4, 57, 51022, {"lc": 2, "type": "constructor", "id": ["datetime", "timezone"], "method": null, "args": [{"lc": 2, "type": "constructor", "id": ["datetime", "timedelta"], "method": null, "args": [0, 86340, 0], "kwargs": {}}], "kwargs": {}}], "kwargs": {"fold": 0}}, "uid": {"lc": 2, "type": "constructor", "id": ["uuid", "UUID"], "method": null, "args": ["00000000000000000000000000000001"], "kwargs": {}}, "timestamp": {"lc": 2, "type": "constructor", "id": ["datetime", "datetime"], "method": "fromisoformat", "args": ["2024-04-19T23:04:57.051022+23:59"], "kwargs": {}}, "my_slotted_class": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyDataclassWSlots"], "method": null, "args": [], "kwargs": {"foo": "bar", "bar": 2}}, "my_dataclass": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyDataclass"], "method": null, "args": [], "kwargs": {"foo": "foo", "bar": 1}}, "my_enum": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyEnum"], "method": null, "args": ["foo"], "kwargs": {}}, "my_pydantic": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyPydantic"], "method": "model_construct", "args": [], "kwargs": {"foo": "foo", "bar": 1}}, "my_funny_pydantic": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyFunnyPydantic"], "method": "construct", "args": [], "kwargs": {"foo": "foo", "bar": 1}}, "person": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "Person"], "method": null, "args": [], "kwargs": {"name": "foo"}}, "a_bool": true, "a_none": null, "a_str": "foo", "a_str_nuc": "foo\\u0000", "a_str_uc": "foo \xe2\x9b\xb0\xef\xb8\x8f", "a_str_ucuc": "foo \xe2\x9b\xb0\xef\xb8\x8f\\u0000", "a_str_ucucuc": "foo \\\\u26f0\\\\ufe0f", "text": ["Hello", "Python", "Surrogate", "Example", "String", "With", "Surrogates", "Embedded", "In", "The", "Text", "\xe6\x94\xb6\xe8\x8a\xb1\xf0\x9f\x99\x84\xc2\xb7\xe5\x88\xb0"], "an_int": 1, "a_float": 1.1, "runnable_map": {"lc": 1, "type": "constructor", "id": ["langchain", "schema", "runnable", "RunnableParallel"], "kwargs": {"steps__": {}}, "name": "RunnableParallel<>", "graph": {"nodes": [{"id": 0, "type": "schema", "data": "Parallel<>Input"}, {"id": 1, "type": "schema", "data": "Parallel<>Output"}], "edges": []}}}""",
     )
 
     assert serde.loads_typed(dumped) == {

libs/checkpoint/tests/test_memory.py

@@ -82,26 +82,20 @@ class TestMemorySaver:
         assert search_results_2[0].metadata == self.metadata_2
 
         search_results_3 = list(self.memory_saver.list(None, filter=query_3))
-        assert len(search_results_3) == 2
+        assert len(search_results_3) == 3
 
         search_results_4 = list(self.memory_saver.list(None, filter=query_4))
         assert len(search_results_4) == 0
 
-        # search by config (defaults to root graph checkpoints)
+        # search by config (defaults to checkpoints across all namespaces)
         search_results_5 = list(
             self.memory_saver.list({"configurable": {"thread_id": "thread-2"}})
         )
-        assert len(search_results_5) == 1
-        assert search_results_5[0].config["configurable"]["checkpoint_ns"] == ""
-
-        # search by config and checkpoint_ns
-        search_results_6 = list(
-            self.memory_saver.list(
-                {"configurable": {"thread_id": "thread-2", "checkpoint_ns": "inner"}}
-            )
-        )
-        assert len(search_results_6) == 1
-        assert search_results_6[0].config["configurable"]["checkpoint_ns"] == "inner"
+        assert len(search_results_5) == 2
+        assert {
+            search_results_5[0].config["configurable"]["checkpoint_ns"],
+            search_results_5[1].config["configurable"]["checkpoint_ns"],
+        } == {"", "inner"}
 
         # TODO: test before and limit params
 
@@ -110,6 +104,7 @@ class TestMemorySaver:
         # save checkpoints
         self.memory_saver.put(self.config_1, self.chkpnt_1, self.metadata_1, {})
         self.memory_saver.put(self.config_2, self.chkpnt_2, self.metadata_2, {})
+        self.memory_saver.put(self.config_3, self.chkpnt_3, self.metadata_3, {})
 
         # call method / assertions
         query_1: CheckpointMetadata = {"source": "input"}  # search by 1 key
@@ -135,7 +130,7 @@ class TestMemorySaver:
         search_results_3 = [
             c async for c in self.memory_saver.alist(None, filter=query_3)
         ]
-        assert len(search_results_3) == 2
+        assert len(search_results_3) == 3
 
         search_results_4 = [
             c async for c in self.memory_saver.alist(None, filter=query_4)

libs/langgraph/README.md

@@ -59,7 +59,7 @@ from langchain_core.messages import HumanMessage
 from langchain_anthropic import ChatAnthropic
 from langchain_core.tools import tool
 from langgraph.checkpoint.memory import MemorySaver
-from langgraph.graph import END, StateGraph, MessagesState
+from langgraph.graph import END, START, StateGraph, MessagesState
 from langgraph.prebuilt import ToolNode
 
 
@@ -107,7 +107,7 @@ workflow.add_node("tools", tool_node)
 
 # Set the entrypoint as `agent`
 # This means that this node is the first one called
-workflow.set_entry_point("agent")
+workflow.add_edge(START, "agent")
 
 # We now add a conditional edge
 workflow.add_conditional_edges(

libs/langgraph/langgraph/channels/any_value.py

@@ -28,12 +28,6 @@ class AnyValue(Generic[Value], BaseChannel[Value, Value, Value]):
         """The type of the update received by the channel."""
         return self.typ
 
-    def checkpoint(self) -> Value:
-        try:
-            return self.value
-        except AttributeError:
-            raise EmptyChannelError()
-
     @contextmanager
     def from_checkpoint(
         self, checkpoint: Optional[Value], config: RunnableConfig

libs/langgraph/langgraph/channels/base.py

@@ -35,11 +35,11 @@ class BaseChannel(Generic[Value, Update, C], ABC):
 
     # serialize/deserialize methods
 
-    @abstractmethod
     def checkpoint(self) -> Optional[C]:
         """Return a serializable representation of the channel's current state.
         Raises EmptyChannelError if the channel is empty (never updated yet),
         or doesn't support checkpoints."""
+        return self.get()
 
     @contextmanager
     @abstractmethod

libs/langgraph/langgraph/channels/binop.py

@@ -73,12 +73,6 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]):
         """The type of the update received by the channel."""
         return self.typ
 
-    def checkpoint(self) -> Value:
-        try:
-            return self.value
-        except AttributeError:
-            raise EmptyChannelError()
-
     @contextmanager
     def from_checkpoint(
         self, checkpoint: Optional[Value], config: RunnableConfig

libs/langgraph/langgraph/channels/context.py

@@ -1,124 +1,5 @@
-from contextlib import asynccontextmanager, contextmanager
-from inspect import signature
-from typing import (
-    Any,
-    AsyncContextManager,
-    AsyncGenerator,
-    ContextManager,
-    Generator,
-    Generic,
-    Optional,
-    Sequence,
-    Type,
-    Union,
-)
+from langgraph.managed.context import Context as ContextManagedValue
 
-from langchain_core.runnables import RunnableConfig
-from typing_extensions import Self
+Context = ContextManagedValue.of
 
-from langgraph.channels.base import BaseChannel, Value
-from langgraph.errors import EmptyChannelError, InvalidUpdateError
-
-
-class Context(Generic[Value], BaseChannel[Value, None, None]):
-    """Exposes the value of a context manager, for the duration of an invocation.
-    Context manager is entered before the first step, and exited after the last step.
-    Optionally, provide an equivalent async context manager, which will be used
-    instead for async invocations.
-
-    ```python
-    import httpx
-
-    client = Channels.Context(httpx.Client, httpx.AsyncClient)
-    ```
-    """
-
-    value: Value
-
-    def __init__(
-        self,
-        ctx: Union[
-            None, Type[ContextManager[Value]], Type[AsyncContextManager[Value]]
-        ] = None,
-        actx: Optional[Type[AsyncContextManager[Value]]] = None,
-    ) -> None:
-        if ctx is None and actx is None:
-            raise ValueError("Must provide either sync or async context manager.")
-        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."""
-        return None
-
-    @property
-    def UpdateType(self) -> Type[None]:
-        """The type of the update received by the channel."""
-        return None
-
-    def checkpoint(self) -> None:
-        raise EmptyChannelError()
-
-    @contextmanager
-    def from_checkpoint(
-        self, checkpoint: None, config: RunnableConfig
-    ) -> Generator[Self, None, None]:
-        if self.ctx is None:
-            raise ValueError("Cannot enter sync context manager.")
-
-        empty = self.__class__(ctx=self.ctx, actx=self.actx)
-        ctx = (
-            self.ctx(config)
-            if signature(self.ctx).parameters.get("config")
-            else self.ctx()
-        )
-        with ctx as value:
-            empty.value = value
-            yield empty
-
-    @asynccontextmanager
-    async def afrom_checkpoint(
-        self, checkpoint: None, config: RunnableConfig
-    ) -> AsyncGenerator[Self, None]:
-        empty = self.__class__(ctx=self.ctx, actx=self.actx)
-        if self.actx is not None:
-            ctx = (
-                self.actx(config)
-                if signature(self.actx).parameters.get("config")
-                else self.actx()
-            )
-        else:
-            ctx = (
-                self.ctx(config)
-                if signature(self.ctx).parameters.get("config")
-                else self.ctx()
-            )
-        if hasattr(ctx, "__aenter__"):
-            async with ctx as value:
-                empty.value = value
-                yield empty
-        else:
-            with ctx as value:
-                empty.value = value
-                yield empty
-
-    def update(self, values: Sequence[None]) -> bool:
-        if values:
-            raise InvalidUpdateError(
-                f"At key '{self.key}': Context channel does not accept writes."
-            )
-        return False
-
-    def get(self) -> Value:
-        try:
-            return self.value
-        except AttributeError:
-            raise EmptyChannelError()
+__all__ = ["Context"]

libs/langgraph/langgraph/channels/ephemeral_value.py

@@ -28,12 +28,6 @@ class EphemeralValue(Generic[Value], BaseChannel[Value, Value, Value]):
         """The type of the update received by the channel."""
         return self.typ
 
-    def checkpoint(self) -> Value:
-        try:
-            return self.value
-        except AttributeError:
-            raise EmptyChannelError()
-
     @contextmanager
     def from_checkpoint(
         self, checkpoint: Optional[Value], config: RunnableConfig

libs/langgraph/langgraph/channels/last_value.py

@@ -27,12 +27,6 @@ class LastValue(Generic[Value], BaseChannel[Value, Value, Value]):
         """The type of the update received by the channel."""
         return self.typ
 
-    def checkpoint(self) -> Value:
-        try:
-            return self.value
-        except AttributeError:
-            raise EmptyChannelError()
-
     @contextmanager
     def from_checkpoint(
         self, checkpoint: Optional[Value], config: RunnableConfig

libs/langgraph/langgraph/channels/topic.py

@@ -30,23 +30,15 @@ class Topic(
         accumulate: Whether to accumulate values across steps. If False, the channel will be emptied after each step.
     """
 
-    def __init__(
-        self, typ: Type[Value], unique: bool = False, accumulate: bool = False
-    ) -> None:
+    def __init__(self, typ: Type[Value], accumulate: bool = False) -> None:
         # attrs
         self.typ = typ
-        self.unique = unique
         self.accumulate = accumulate
         # state
-        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
-        )
+        return isinstance(value, Topic) and value.accumulate == self.accumulate
 
     @property
     def ValueType(self) -> Any:
@@ -59,18 +51,20 @@ class Topic(
         return Union[self.typ, list[self.typ]]  # type: ignore[name-defined]
 
     def checkpoint(self) -> tuple[set[Value], list[Value]]:
-        return (self.seen, self.values)
+        return self.values
 
     @contextmanager
     def from_checkpoint(
         self,
-        checkpoint: Optional[tuple[set[Value], list[Value]]],
+        checkpoint: Optional[list[Value]],
         config: RunnableConfig,
     ) -> Generator[Self, None, None]:
-        empty = self.__class__(self.typ, self.unique, self.accumulate)
+        empty = self.__class__(self.typ, self.accumulate)
         if checkpoint is not None:
-            empty.seen = checkpoint[0].copy()
-            empty.values = checkpoint[1].copy()
+            if isinstance(checkpoint, tuple):
+                empty.values = checkpoint[1].copy()
+            else:
+                empty.values = checkpoint.copy()
         try:
             yield empty
         finally:
@@ -81,13 +75,7 @@ class Topic(
         if not self.accumulate:
             self.values = list[Value]()
         if flat_values := flatten(values):
-            if self.unique:
-                for value in flat_values:
-                    if value not in self.seen:
-                        self.seen.add(value)
-                        self.values.append(value)
-            else:
-                self.values.extend(flat_values)
+            self.values.extend(flat_values)
         return self.values != current
 
     def get(self) -> Sequence[Value]:

libs/langgraph/langgraph/constants.py

@@ -1,5 +1,6 @@
 from dataclasses import dataclass
-from typing import Any, Literal
+from typing import Any, Literal, Optional
+from uuid import uuid4
 
 INPUT = "__input__"
 CONFIG_KEY_SEND = "__pregel_send"
@@ -11,6 +12,7 @@ CONFIG_KEY_TASK_ID = "__pregel_task_id"
 INTERRUPT = "__interrupt__"
 ERROR = "__error__"
 TASKS = "__pregel_tasks"
+RUNTIME_PLACEHOLDER = "__pregel_runtime_placeholder__"
 RESERVED = {
     INTERRUPT,
     ERROR,
@@ -22,6 +24,7 @@ RESERVED = {
     CONFIG_KEY_RESUMING,
     CONFIG_KEY_TASK_ID,
     INPUT,
+    RUNTIME_PLACEHOLDER,
 }
 TAG_HIDDEN = "langsmith:hidden"
 
@@ -29,6 +32,7 @@ START = "__start__"
 END = "__end__"
 
 CHECKPOINT_NAMESPACE_SEPARATOR = "|"
+SEND_CHECKPOINT_NAMESPACE_SEPARATOR = ":"
 
 
 class Send:
@@ -47,6 +51,7 @@ class Send:
     Attributes:
         node (str): The name of the target node to send the message to.
         arg (Any): The state or message to send to the target node.
+        id (str): ID associated with the Send.
 
     Examples:
         >>> from typing import Annotated
@@ -74,23 +79,26 @@ class Send:
 
     node: str
     arg: Any
+    id: Optional[str]
 
-    def __init__(self, /, node: str, arg: Any) -> None:
+    def __init__(self, /, node: str, arg: Any, id: Optional[str] = None) -> None:
         """
         Initialize a new instance of the Send class.
 
         Args:
             node (str): The name of the target node to send the message to.
             arg (Any): The state or message to send to the target node.
+            id (str): ID associated with the Send.
         """
         self.node = node
         self.arg = arg
+        self.id = id or str(uuid4())
 
     def __hash__(self) -> int:
-        return hash((self.node, self.arg))
+        return hash((self.node, self.arg, self.id))
 
     def __repr__(self) -> str:
-        return f"Send(node={self.node!r}, arg={self.arg!r})"
+        return f"Send(node={self.node!r}, arg={self.arg!r}, id={self.id!r})"
 
     def __eq__(self, value: object) -> bool:
         return (

libs/langgraph/langgraph/graph/graph.py

@@ -28,6 +28,7 @@ from langgraph.checkpoint.base import BaseCheckpointSaver
 from langgraph.constants import (
     CHECKPOINT_NAMESPACE_SEPARATOR,
     END,
+    SEND_CHECKPOINT_NAMESPACE_SEPARATOR,
     START,
     TAG_HIDDEN,
     Send,
@@ -160,10 +161,15 @@ class Graph:
         *,
         metadata: Optional[dict[str, Any]] = None,
     ) -> None:
-        if isinstance(node, str) and CHECKPOINT_NAMESPACE_SEPARATOR in node:
-            raise ValueError(
-                f"'{CHECKPOINT_NAMESPACE_SEPARATOR}' is a reserved character and is not allowed in the node names."
-            )
+        if isinstance(node, str):
+            for character in (
+                CHECKPOINT_NAMESPACE_SEPARATOR,
+                SEND_CHECKPOINT_NAMESPACE_SEPARATOR,
+            ):
+                if character in node:
+                    raise ValueError(
+                        f"'{character}' is a reserved character and is not allowed in the node names."
+                    )
 
         if self.compiled:
             logger.warning(

libs/langgraph/langgraph/graph/state.py

@@ -6,6 +6,7 @@ from functools import partial
 from inspect import isclass, isfunction, signature
 from typing import (
     Any,
+    Callable,
     NamedTuple,
     Optional,
     Sequence,
@@ -25,13 +26,16 @@ from langchain_core.runnables.utils import (
 
 from langgraph.channels.base import BaseChannel
 from langgraph.channels.binop import BinaryOperatorAggregate
-from langgraph.channels.context import Context
 from langgraph.channels.dynamic_barrier_value import DynamicBarrierValue, WaitForNames
 from langgraph.channels.ephemeral_value import EphemeralValue
 from langgraph.channels.last_value import LastValue
 from langgraph.channels.named_barrier_value import NamedBarrierValue
 from langgraph.checkpoint.base import BaseCheckpointSaver
-from langgraph.constants import CHECKPOINT_NAMESPACE_SEPARATOR, TAG_HIDDEN
+from langgraph.constants import (
+    CHECKPOINT_NAMESPACE_SEPARATOR,
+    SEND_CHECKPOINT_NAMESPACE_SEPARATOR,
+    TAG_HIDDEN,
+)
 from langgraph.errors import InvalidUpdateError
 from langgraph.graph.graph import (
     END,
@@ -317,10 +321,14 @@ class StateGraph(Graph):
         if node == END or node == START:
             raise ValueError(f"Node `{node}` is reserved.")
 
-        if CHECKPOINT_NAMESPACE_SEPARATOR in node:
-            raise ValueError(
-                f"'{CHECKPOINT_NAMESPACE_SEPARATOR}' is a reserved character and is not allowed in the node names."
-            )
+        for character in (
+            CHECKPOINT_NAMESPACE_SEPARATOR,
+            SEND_CHECKPOINT_NAMESPACE_SEPARATOR,
+        ):
+            if character in node:
+                raise ValueError(
+                    f"'{character}' is a reserved character and is not allowed in the node names."
+                )
 
         try:
             if isfunction(action) and (
@@ -374,7 +382,7 @@ class StateGraph(Graph):
                 raise ValueError(f"Need to add_node `{start}` first")
         if end_key == START:
             raise ValueError("START cannot be an end node")
-        if end_key not in self.nodes:
+        if end_key != END and end_key not in self.nodes:
             raise ValueError(f"Need to add_node `{end_key}` first")
 
         self.waiting_edges.add((tuple(start_key), end_key))
@@ -425,16 +433,14 @@ class StateGraph(Graph):
             else [
                 key
                 for key, val in self.schemas[self.output].items()
-                if not isinstance(val, Context) and not is_managed_value(val)
+                if not is_managed_value(val)
             ]
         )
         stream_channels = (
             "__root__"
             if len(self.channels) == 1 and "__root__" in self.channels
             else [
-                key
-                for key, val in self.channels.items()
-                if not isinstance(val, Context) and not is_managed_value(val)
+                key for key, val in self.channels.items() if not is_managed_value(val)
             ]
         )
 
@@ -502,7 +508,6 @@ class CompiledStateGraph(CompiledGraph):
                         k: (self.channels[k].UpdateType, None)
                         for k in self.builder.schemas[self.builder.input]
                         if isinstance(self.channels[k], BaseChannel)
-                        and not isinstance(self.channels[k], Context)
                     },
                 )
 
@@ -523,7 +528,7 @@ class CompiledStateGraph(CompiledGraph):
             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)
+                if not is_managed_value(v)
             ]
         else:
             output_keys = list(self.builder.channels) + [
@@ -650,7 +655,14 @@ class CompiledStateGraph(CompiledGraph):
                 return ChannelWrite(writes, tags=[TAG_HIDDEN])
 
         # attach branch publisher
-        self.nodes[start] |= branch.run(branch_writer, _get_state_reader(self.builder))
+        schema = (
+            self.builder.nodes[start].input
+            if start in self.builder.nodes
+            else self.builder.schema
+        )
+        self.nodes[start] |= branch.run(
+            branch_writer, _get_state_reader(self.builder, schema)
+        )
 
         # attach branch subscribers
         ends = (
@@ -676,16 +688,17 @@ class CompiledStateGraph(CompiledGraph):
                     )
 
 
-def _get_state_reader(graph: StateGraph) -> ChannelRead:
-    state_keys = list(graph.channels)
+def _get_state_reader(
+    builder: StateGraph, schema: Type[Any]
+) -> Callable[[RunnableConfig], Any]:
+    state_keys = list(builder.channels)
+    select = list(builder.schemas[schema])
     return partial(
         ChannelRead.do_read,
-        channel=state_keys[0] if state_keys == ["__root__"] else state_keys,
+        select=select[0] if select == ["__root__"] else select,
         fresh=True,
         # coerce state dict to schema class (eg. pydantic model)
-        mapper=(
-            None if state_keys == ["__root__"] else partial(_coerce_state, graph.schema)
-        ),
+        mapper=(None if state_keys == ["__root__"] else partial(_coerce_state, schema)),
     )
 
 

libs/langgraph/langgraph/managed/base.py

@@ -16,11 +16,16 @@ from typing import (
 from langchain_core.runnables import RunnableConfig
 from typing_extensions import Self, TypeGuard
 
+from langgraph.constants import RUNTIME_PLACEHOLDER
+
 V = TypeVar("V")
 U = TypeVar("U")
 
 
 class ManagedValue(ABC, Generic[V]):
+    runtime: bool = False
+    """Whether the managed value is always created at runtime, ie. never stored."""
+
     def __init__(self, config: RunnableConfig) -> None:
         self.config = config
 
@@ -74,8 +79,6 @@ class ConfiguredManagedValue(NamedTuple):
 
 ManagedValueSpec = Union[Type[ManagedValue], ConfiguredManagedValue]
 
-ManagedValueMapping = dict[str, ManagedValue]
-
 
 def is_managed_value(value: Any) -> TypeGuard[ManagedValueSpec]:
     return (isclass(value) and issubclass(value, ManagedValue)) or isinstance(
@@ -103,3 +106,45 @@ def is_writable_managed_value(value: Any) -> TypeGuard[Type[WritableManagedValue
 
 ChannelKeyPlaceholder = object()
 ChannelTypePlaceholder = object()
+
+
+class ManagedValueMapping(dict[str, ManagedValue]):
+    def replace_runtime_values(self, step: int, values: Union[dict[str, Any], Any]):
+        if not self or not values:
+            return
+        if all(not mv.runtime for mv in self.values()):
+            return
+        if isinstance(values, dict):
+            for key, value in values.items():
+                for chan, mv in self.items():
+                    if mv.runtime and mv(step) is value:
+                        values[key] = {RUNTIME_PLACEHOLDER: chan}
+        elif hasattr(values, "__dir__") and callable(values.__dir__):
+            for key in dir(values):
+                try:
+                    value = getattr(values, key)
+                    for chan, mv in self.items():
+                        if mv.runtime and mv(step) is value:
+                            setattr(values, key, {RUNTIME_PLACEHOLDER: chan})
+                except AttributeError:
+                    pass
+
+    def replace_runtime_placeholders(
+        self, step: int, values: Union[dict[str, Any], Any]
+    ):
+        if not self or not values:
+            return
+        if all(not mv.runtime for mv in self.values()):
+            return
+        if isinstance(values, dict):
+            for key, value in values.items():
+                if isinstance(value, dict) and RUNTIME_PLACEHOLDER in value:
+                    values[key] = self[value[RUNTIME_PLACEHOLDER]](step)
+        elif hasattr(values, "__dir__") and callable(values.__dir__):
+            for key in dir(values):
+                try:
+                    value = getattr(values, key)
+                    if isinstance(value, dict) and RUNTIME_PLACEHOLDER in value:
+                        setattr(values, key, self[value[RUNTIME_PLACEHOLDER]](step))
+                except AttributeError:
+                    pass

libs/langgraph/langgraph/managed/context.py

@@ -0,0 +1,87 @@
+from contextlib import asynccontextmanager, contextmanager
+from inspect import signature
+from typing import (
+    Any,
+    AsyncContextManager,
+    AsyncIterator,
+    ContextManager,
+    Iterator,
+    Optional,
+    Type,
+    Union,
+)
+
+from langchain_core.runnables import RunnableConfig
+from typing_extensions import Self
+
+from langgraph.managed.base import ConfiguredManagedValue, ManagedValue, V
+
+
+class Context(ManagedValue):
+    runtime = True
+
+    value: V
+
+    @staticmethod
+    def of(
+        ctx: Union[None, Type[ContextManager[V]], Type[AsyncContextManager[V]]] = None,
+        actx: Optional[Type[AsyncContextManager[V]]] = None,
+    ) -> ConfiguredManagedValue:
+        if ctx is None and actx is None:
+            raise ValueError("Must provide either sync or async context manager.")
+        return ConfiguredManagedValue(Context, {"ctx": ctx, "actx": actx})
+
+    @classmethod
+    @contextmanager
+    def enter(cls, config: RunnableConfig, **kwargs: Any) -> Iterator[Self]:
+        with super().enter(config, **kwargs) as self:
+            if self.ctx is None:
+                raise ValueError(
+                    "Synchronous context manager not found. Please initialize Context value with a sync context manager, or invoke your graph asynchronously."
+                )
+            ctx = (
+                self.ctx(config)
+                if signature(self.ctx).parameters.get("config")
+                else self.ctx()
+            )
+            with ctx as v:
+                self.value = v
+                yield self
+
+    @classmethod
+    @asynccontextmanager
+    async def aenter(cls, config: RunnableConfig, **kwargs: Any) -> AsyncIterator[Self]:
+        async with super().aenter(config, **kwargs) as self:
+            if self.actx is not None:
+                ctx = (
+                    self.actx(config)
+                    if signature(self.actx).parameters.get("config")
+                    else self.actx()
+                )
+            else:
+                ctx = (
+                    self.ctx(config)
+                    if signature(self.ctx).parameters.get("config")
+                    else self.ctx()
+                )
+            if hasattr(ctx, "__aenter__"):
+                async with ctx as v:
+                    self.value = v
+                    yield self
+            else:
+                with ctx as v:
+                    self.value = v
+                    yield self
+
+    def __init__(
+        self,
+        config: RunnableConfig,
+        *,
+        ctx: Union[None, Type[ContextManager[V]], Type[AsyncContextManager[V]]] = None,
+        actx: Optional[Type[AsyncContextManager[V]]] = None,
+    ) -> None:
+        self.ctx = ctx
+        self.actx = actx
+
+    def __call__(self, step: int) -> V:
+        return self.value

libs/langgraph/langgraph/managed/shared_value.py

@@ -81,7 +81,6 @@ class SharedValue(WritableManagedValue[Value, Update]):
             ):
                 raise ValueError("SharedValue must be a dict")
         self.scope = scope
-        self.config = config
         self.value: Value = {}
         self.store: BaseStore = config["configurable"].get(CONFIG_KEY_STORE)
         if self.store is None:

libs/langgraph/langgraph/prebuilt/chat_agent_executor.py

@@ -130,7 +130,7 @@ def _get_model_preprocessing_runnable(
 @deprecated_parameter("messages_modifier", "0.1.9", "state_modifier", removal="0.3.0")
 def create_react_agent(
     model: LanguageModelLike,
-    tools: Union[ToolExecutor, Sequence[BaseTool]],
+    tools: Union[ToolExecutor, Sequence[BaseTool], ToolNode],
     *,
     state_schema: Optional[StateSchemaType] = None,
     messages_modifier: Optional[MessagesModifier] = None,
@@ -144,7 +144,7 @@ def create_react_agent(
 
     Args:
         model: The `LangChain` chat model that supports tool calling.
-        tools: A list of tools or a ToolExecutor instance.
+        tools: A list of tools, a ToolExecutor, or a ToolNode 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.
@@ -419,8 +419,13 @@ def create_react_agent(
 
     if isinstance(tools, ToolExecutor):
         tool_classes = tools.tools
+        tool_node = ToolNode(tool_classes)
+    elif isinstance(tools, ToolNode):
+        tool_classes = tools.tools_by_name.values()
+        tool_node = tools
     else:
         tool_classes = tools
+        tool_node = ToolNode(tool_classes)
     model = model.bind_tools(tool_classes)
 
     # Define the function that determines whether to continue or not
@@ -474,7 +479,7 @@ def create_react_agent(
 
     # Define the two nodes we will cycle between
     workflow.add_node("agent", RunnableLambda(call_model, acall_model))
-    workflow.add_node("tools", ToolNode(tool_classes))
+    workflow.add_node("tools", tool_node)
 
     # Set the entrypoint as `agent`
     # This means that this node is the first one called

libs/langgraph/langgraph/pregel/__init__.py

@@ -38,6 +38,7 @@ from langchain_core.runnables.config import (
     ensure_config,
     get_async_callback_manager_for_config,
     get_callback_manager_for_config,
+    merge_configs,
     patch_config,
 )
 from langchain_core.runnables.utils import (
@@ -51,30 +52,38 @@ from typing_extensions import Self
 from langgraph.channels.base import (
     BaseChannel,
 )
-from langgraph.channels.context import Context
 from langgraph.checkpoint.base import (
     BaseCheckpointSaver,
+    CheckpointTuple,
     copy_checkpoint,
     create_checkpoint,
     empty_checkpoint,
 )
 from langgraph.constants import (
+    CHECKPOINT_NAMESPACE_SEPARATOR,
     CONFIG_KEY_CHECKPOINTER,
     CONFIG_KEY_READ,
     CONFIG_KEY_RESUMING,
     CONFIG_KEY_SEND,
     ERROR,
     INTERRUPT,
+    SEND_CHECKPOINT_NAMESPACE_SEPARATOR,
 )
 from langgraph.errors import GraphInterrupt, GraphRecursionError, InvalidUpdateError
 from langgraph.managed.base import ManagedValueSpec
-from langgraph.pregel.algo import apply_writes, local_read, prepare_next_tasks
+from langgraph.pregel.algo import (
+    apply_writes,
+    local_read,
+    local_write,
+    prepare_next_tasks,
+)
 from langgraph.pregel.debug import (
     print_step_checkpoint,
     print_step_tasks,
     print_step_writes,
     tasks_w_writes,
 )
+from langgraph.pregel.get_state import assemble_state_snapshot_hierarchy
 from langgraph.pregel.io import read_channels
 from langgraph.pregel.loop import AsyncPregelLoop, SyncPregelLoop
 from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager
@@ -86,7 +95,9 @@ from langgraph.pregel.types import (
     StateSnapshot,
     StreamMode,
 )
-from langgraph.pregel.utils import get_new_channel_versions
+from langgraph.pregel.utils import (
+    get_new_channel_versions,
+)
 from langgraph.pregel.validate import validate_graph, validate_keys
 from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
 from langgraph.store.base import BaseStore
@@ -173,6 +184,181 @@ class Channel:
         )
 
 
+def _get_checkpoint_ns_to_graph(
+    graph: Pregel,
+    checkpoint_ns_to_graph: Optional[dict[str, Pregel]] = None,
+    checkpoint_ns: str = "",
+    max_depth: int = 10,
+) -> Pregel:
+    if checkpoint_ns_to_graph is None:
+        checkpoint_ns_to_graph = {}
+
+    if max_depth <= 0:
+        raise RecursionError(
+            "Reached maximum recursion depth while building checkpoint NS -> graph mapping."
+        )
+
+    for node_name, node in graph.nodes.items():
+        new_checkpoint_ns = (
+            f"{checkpoint_ns}{CHECKPOINT_NAMESPACE_SEPARATOR}{node_name}"
+            if checkpoint_ns
+            else node_name
+        )
+        if isinstance(node.bound, Pregel):
+            _get_checkpoint_ns_to_graph(
+                node.bound, checkpoint_ns_to_graph, new_checkpoint_ns, max_depth - 1
+            )
+        elif isinstance(node.bound, RunnableSequence):
+            for runnable in node.bound.steps:
+                if isinstance(runnable, Pregel):
+                    _get_checkpoint_ns_to_graph(
+                        runnable,
+                        checkpoint_ns_to_graph,
+                        new_checkpoint_ns,
+                        max_depth - 1,
+                    )
+
+    checkpoint_ns_to_graph[checkpoint_ns] = graph
+    return checkpoint_ns_to_graph
+
+
+def _prepare_state_snapshot(
+    config: RunnableConfig,
+    checkpoint_ns_to_graph: dict[str, Pregel],
+    checkpoint_tuples: Iterator[CheckpointTuple],
+) -> StateSnapshot:
+    checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
+    checkpoint_id = config["configurable"].get("checkpoint_id")
+    checkpoint_ns_to_state_snapshots: dict[str, StateSnapshot] = {}
+    for saved in checkpoint_tuples:
+        saved_checkpoint_ns = saved.config["configurable"]["checkpoint_ns"]
+        saved_checkpoint_id = saved.config["configurable"]["checkpoint_id"]
+        if checkpoint_id and saved_checkpoint_id != checkpoint_id:
+            continue
+
+        graph_checkpoint_ns = saved_checkpoint_ns.split(
+            SEND_CHECKPOINT_NAMESPACE_SEPARATOR
+        )[0]
+        graph = checkpoint_ns_to_graph.get(graph_checkpoint_ns)
+        if graph is None:
+            continue
+
+        with ChannelsManager(
+            graph.channels, saved.checkpoint, saved.config, skip_context=True
+        ) as (
+            channels,
+            managed,
+        ):
+            next_tasks = prepare_next_tasks(
+                saved.checkpoint,
+                graph.nodes,
+                channels,
+                managed,
+                saved.config,
+                saved.metadata.get("step", -1) + 1,
+                for_execution=False,
+            )
+            state_snapshot = StateSnapshot(
+                read_channels(channels, graph.stream_channels_asis),
+                tuple(t.name for t in next_tasks),
+                saved.config,
+                saved.metadata,
+                saved.checkpoint["ts"],
+                saved.parent_config,
+                tasks_w_writes(next_tasks, saved.pending_writes),
+            )
+
+        checkpoint_ns_to_state_snapshots[saved_checkpoint_ns] = state_snapshot
+
+    if not checkpoint_ns_to_state_snapshots:
+        return StateSnapshot(
+            values={},
+            next=(),
+            config=config,
+            metadata=None,
+            created_at=None,
+            parent_config=None,
+            tasks=(),
+        )
+
+    state_snapshot = assemble_state_snapshot_hierarchy(
+        checkpoint_ns, checkpoint_ns_to_state_snapshots
+    )
+    return state_snapshot
+
+
+async def _prepare_state_snapshot_async(
+    config: RunnableConfig,
+    checkpoint_ns_to_graph: dict[str, Pregel],
+    checkpoint_tuples: AsyncIterator[CheckpointTuple],
+) -> StateSnapshot:
+    checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
+    checkpoint_id = config["configurable"].get("checkpoint_id")
+    checkpoint_ns_to_state_snapshots: dict[str, StateSnapshot] = {}
+    async for saved in checkpoint_tuples:
+        saved_checkpoint_ns = saved.config["configurable"]["checkpoint_ns"]
+        saved_checkpoint_id = saved.config["configurable"]["checkpoint_id"]
+        if checkpoint_id and saved_checkpoint_id != checkpoint_id:
+            continue
+
+        graph_checkpoint_ns = saved_checkpoint_ns.split(
+            SEND_CHECKPOINT_NAMESPACE_SEPARATOR
+        )[0]
+        graph = checkpoint_ns_to_graph.get(graph_checkpoint_ns)
+        if graph is None:
+            continue
+
+        async with AsyncChannelsManager(
+            graph.channels, saved.checkpoint, saved.config, skip_context=True
+        ) as (channels, managed):
+            next_tasks = prepare_next_tasks(
+                saved.checkpoint,
+                graph.nodes,
+                channels,
+                managed,
+                saved.config,
+                saved.metadata.get("step", -1) + 1,
+                for_execution=False,
+            )
+            state_snapshot = StateSnapshot(
+                read_channels(channels, graph.stream_channels_asis),
+                tuple(t.name for t in next_tasks),
+                saved.config,
+                saved.metadata,
+                saved.checkpoint["ts"],
+                saved.parent_config,
+                tasks_w_writes(next_tasks, saved.pending_writes),
+            )
+
+        checkpoint_ns_to_state_snapshots[saved_checkpoint_ns] = state_snapshot
+
+    if not checkpoint_ns_to_state_snapshots:
+        return StateSnapshot(
+            values={},
+            next=(),
+            config=config,
+            metadata=None,
+            created_at=None,
+            parent_config=None,
+            tasks=(),
+        )
+
+    state_snapshot = assemble_state_snapshot_hierarchy(
+        checkpoint_ns, checkpoint_ns_to_state_snapshots
+    )
+    return state_snapshot
+
+
+def _has_nested_interrupts(
+    graph: Pregel,
+) -> bool:
+    for child in graph.subgraphs:
+        if child.interrupt_after_nodes or child.interrupt_before_nodes:
+            return True
+    else:
+        return False
+
+
 class Pregel(
     RunnableSerializable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]
 ):
@@ -215,11 +401,18 @@ class Pregel(
 
     config_type: Optional[Type[Any]] = None
 
+    config: Optional[RunnableConfig] = None
+
     name: str = "LangGraph"
 
     class Config:
         arbitrary_types_allowed = True
 
+    def with_config(self, config: RunnableConfig | None = None, **kwargs: Any) -> Self:
+        return self.copy(
+            update={"config": cast(RunnableConfig, {**(config or {}), **kwargs})}
+        )
+
     @classmethod
     def is_lc_serializable(cls) -> bool:
         """Return whether the graph can be serialized by Langchain."""
@@ -293,6 +486,7 @@ class Pregel(
     def get_input_schema(
         self, config: Optional[RunnableConfig] = None
     ) -> Type[BaseModel]:
+        config = merge_configs(self.config, config)
         if isinstance(self.input_channels, str):
             return super().get_input_schema(config)
         else:
@@ -312,6 +506,7 @@ class Pregel(
     def get_output_schema(
         self, config: Optional[RunnableConfig] = None
     ) -> Type[BaseModel]:
+        config = merge_configs(self.config, config)
         if isinstance(self.output_channels, str):
             return super().get_output_schema(config)
         else:
@@ -330,74 +525,64 @@ class Pregel(
     @property
     def stream_channels_asis(self) -> Union[str, Sequence[str]]:
         return self.stream_channels or [
-            k
-            for k in self.channels
-            if isinstance(self.channels[k], BaseChannel)
-            and not isinstance(self.channels[k], Context)
+            k for k in self.channels if isinstance(self.channels[k], BaseChannel)
         ]
 
+    @property
+    def subgraphs(self) -> Iterator[Pregel]:
+        for node in self.nodes.values():
+            if isinstance(node.bound, Pregel):
+                yield node.bound
+                yield from node.bound.subgraphs
+            elif isinstance(node.bound, RunnableSequence):
+                for runnable in node.bound.steps:
+                    if isinstance(runnable, Pregel):
+                        yield runnable
+                        yield from runnable.subgraphs
+
     def get_state(self, config: RunnableConfig) -> StateSnapshot:
         """Get the current state of the graph."""
         if not self.checkpointer:
             raise ValueError("No checkpointer set")
 
+        config = merge_configs(self.config, config) if self.config else config
         saved = self.checkpointer.get_tuple(config)
-        checkpoint = saved.checkpoint if saved else empty_checkpoint()
-        config = saved.config if saved else config
-        with ChannelsManager(self.channels, checkpoint, config, skip_context=True) as (
-            channels,
-            managed,
-        ):
-            next_tasks = prepare_next_tasks(
-                checkpoint,
-                self.nodes,
-                channels,
-                managed,
-                config,
-                saved.metadata.get("step", -1) + 1 if saved else -1,
-                for_execution=False,
-            )
+        checkpoint_config = saved.config if saved else config
+        checkpoint_ns_to_graph: dict[str, Pregel] = _get_checkpoint_ns_to_graph(self)
 
-            return StateSnapshot(
-                read_channels(channels, self.stream_channels_asis),
-                tuple(t.name for t in next_tasks),
-                saved.config if saved else config,
-                saved.metadata if saved else None,
-                saved.checkpoint["ts"] if saved else None,
-                saved.parent_config if saved else None,
-                tasks_w_writes(next_tasks, saved.pending_writes if saved else None),
-            )
+        # we only lookup subgraph checkpoints if we actually have subgraphs
+        if len(set(checkpoint_ns_to_graph)) == 1:
+            checkpoint_tuples = (saved,)
+        else:
+            checkpoint_tuples = self.checkpointer.list(saved.config)
+
+        return _prepare_state_snapshot(
+            checkpoint_config, checkpoint_ns_to_graph, checkpoint_tuples
+        )
 
     async def aget_state(self, config: RunnableConfig) -> StateSnapshot:
         """Get the current state of the graph."""
         if not self.checkpointer:
             raise ValueError("No checkpointer set")
 
+        config = merge_configs(self.config, config) if self.config else config
         saved = await self.checkpointer.aget_tuple(config)
-        checkpoint = saved.checkpoint if saved else empty_checkpoint()
+        checkpoint_config = saved.config if saved else config
+        checkpoint_ns_to_graph: dict[str, Pregel] = _get_checkpoint_ns_to_graph(self)
 
-        config = saved.config if saved else config
-        async with AsyncChannelsManager(
-            self.channels, checkpoint, config, skip_context=True
-        ) as (channels, managed):
-            next_tasks = prepare_next_tasks(
-                checkpoint,
-                self.nodes,
-                channels,
-                managed,
-                config,
-                saved.metadata.get("step", -1) + 1 if saved else -1,
-                for_execution=False,
-            )
-            return StateSnapshot(
-                read_channels(channels, self.stream_channels_asis),
-                tuple(t.name for t in next_tasks),
-                saved.config if saved else config,
-                saved.metadata if saved else None,
-                saved.checkpoint["ts"] if saved else None,
-                saved.parent_config if saved else None,
-                tasks_w_writes(next_tasks, saved.pending_writes if saved else None),
-            )
+        # we only lookup subgraph checkpoints if we actually have subgraphs
+        if len(set(checkpoint_ns_to_graph)) == 1:
+
+            async def alist_checkpoints():
+                yield saved
+
+            checkpoint_tuples = alist_checkpoints()
+        else:
+            checkpoint_tuples = self.checkpointer.alist(saved.config)
+
+        return await _prepare_state_snapshot_async(
+            checkpoint_config, checkpoint_ns_to_graph, checkpoint_tuples
+        )
 
     def get_state_history(
         self,
@@ -415,34 +600,30 @@ class Pregel(
             and signature(self.checkpointer.list).parameters.get("filter") is None
         ):
             raise ValueError("Checkpointer does not support filtering")
-        for (
-            config,
-            checkpoint,
-            metadata,
-            parent_config,
-            pending_writes,
-        ) in self.checkpointer.list(config, before=before, limit=limit, filter=filter):
-            with ChannelsManager(
-                self.channels, checkpoint, config, skip_context=True
-            ) as (channels, managed):
-                next_tasks = prepare_next_tasks(
-                    checkpoint,
-                    self.nodes,
-                    channels,
-                    managed,
-                    config,
-                    metadata.get("step", -1) + 1,
-                    for_execution=False,
-                )
-                yield StateSnapshot(
-                    read_channels(channels, self.stream_channels_asis),
-                    tuple(t.name for t in next_tasks),
-                    config,
-                    metadata,
-                    checkpoint["ts"],
-                    parent_config,
-                    tasks_w_writes(next_tasks, pending_writes),
-                )
+
+        checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
+        checkpoint_ns_to_graph = _get_checkpoint_ns_to_graph(self)
+        # find all matching checkpoint tuples for parent and subgraphs
+        checkpoint_tuples = [
+            checkpoint_tuple
+            for checkpoint_tuple in self.checkpointer.list(
+                merge_configs(self.config, config) if self.config else config,
+                before=before,
+                limit=limit,
+                filter=filter,
+            )
+        ]
+        for checkpoint_tuple in checkpoint_tuples:
+            if (
+                checkpoint_tuple.config["configurable"]["checkpoint_ns"]
+                != checkpoint_ns
+            ):
+                continue
+
+            state_snapshot = _prepare_state_snapshot(
+                checkpoint_tuple.config, checkpoint_ns_to_graph, iter(checkpoint_tuples)
+            )
+            yield state_snapshot
 
     async def aget_state_history(
         self,
@@ -460,34 +641,36 @@ class Pregel(
             and signature(self.checkpointer.list).parameters.get("filter") is None
         ):
             raise ValueError("Checkpointer does not support filtering")
-        async for (
-            config,
-            checkpoint,
-            metadata,
-            parent_config,
-            pending_writes,
-        ) in self.checkpointer.alist(config, before=before, limit=limit, filter=filter):
-            async with AsyncChannelsManager(
-                self.channels, checkpoint, config, skip_context=True
-            ) as (channels, managed):
-                next_tasks = prepare_next_tasks(
-                    checkpoint,
-                    self.nodes,
-                    channels,
-                    managed,
-                    config,
-                    metadata.get("step", -1) + 1,
-                    for_execution=False,
-                )
-                yield StateSnapshot(
-                    read_channels(channels, self.stream_channels_asis),
-                    tuple(t.name for t in next_tasks),
-                    config,
-                    metadata,
-                    checkpoint["ts"],
-                    parent_config,
-                    tasks_w_writes(next_tasks, pending_writes),
-                )
+
+        checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
+        checkpoint_ns_to_graph = _get_checkpoint_ns_to_graph(self)
+        # find all matching checkpoint tuples for parent and subgraphs
+        checkpoint_tuples = [
+            checkpoint_tuple
+            async for checkpoint_tuple in self.checkpointer.alist(
+                merge_configs(self.config, config) if self.config else config,
+                before=before,
+                limit=limit,
+                filter=filter,
+            )
+        ]
+
+        # turn matching checkpoint tuples into an async iterator
+        async def alist_checkpoints() -> AsyncIterator[CheckpointTuple]:
+            for checkpoint_tuple in checkpoint_tuples:
+                yield checkpoint_tuple
+
+        for checkpoint_tuple in checkpoint_tuples:
+            if (
+                checkpoint_tuple.config["configurable"]["checkpoint_ns"]
+                != checkpoint_ns
+            ):
+                continue
+
+            state_snapshot = await _prepare_state_snapshot_async(
+                checkpoint_tuple.config, checkpoint_ns_to_graph, alist_checkpoints()
+            )
+            yield state_snapshot
 
     def update_state(
         self,
@@ -503,6 +686,7 @@ class Pregel(
             raise ValueError("No checkpointer set")
 
         # get last checkpoint
+        config = merge_configs(self.config, config) if self.config else config
         saved = self.checkpointer.get_tuple(config)
         checkpoint = copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint()
         checkpoint_previous_versions = (
@@ -564,7 +748,7 @@ class Pregel(
         # update channels
         with ChannelsManager(self.channels, checkpoint, config) as (
             channels,
-            _,
+            managed,
         ):
             # create task to run all writers of the chosen node
             writers = self.nodes[as_node].get_writers()
@@ -588,9 +772,22 @@ class Pregel(
                     run_name=self.name + "UpdateState",
                     configurable={
                         # deque.extend is thread-safe
-                        CONFIG_KEY_SEND: task.writes.extend,
+                        CONFIG_KEY_SEND: partial(
+                            local_write,
+                            step + 1,
+                            task.writes.extend,
+                            self.nodes,
+                            channels,
+                            managed,
+                        ),
                         CONFIG_KEY_READ: partial(
-                            local_read, checkpoint, channels, task, config
+                            local_read,
+                            step + 1,
+                            checkpoint,
+                            channels,
+                            managed,
+                            task,
+                            config,
                         ),
                     },
                 ),
@@ -623,6 +820,7 @@ class Pregel(
             raise ValueError("No checkpointer set")
 
         # get last checkpoint
+        config = merge_configs(self.config, config) if self.config else config
         saved = await self.checkpointer.aget_tuple(config)
         checkpoint = copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint()
         checkpoint_previous_versions = (
@@ -682,7 +880,7 @@ class Pregel(
         # update channels, acting as the chosen node
         async with AsyncChannelsManager(self.channels, checkpoint, config) as (
             channels,
-            _,
+            managed,
         ):
             # create task to run all writers of the chosen node
             writers = self.nodes[as_node].get_writers()
@@ -706,9 +904,22 @@ class Pregel(
                     run_name=self.name + "UpdateState",
                     configurable={
                         # deque.extend is thread-safe
-                        CONFIG_KEY_SEND: task.writes.extend,
+                        CONFIG_KEY_SEND: partial(
+                            local_write,
+                            step + 1,
+                            task.writes.extend,
+                            self.nodes,
+                            channels,
+                            managed,
+                        ),
                         CONFIG_KEY_READ: partial(
-                            local_read, checkpoint, channels, task, config
+                            local_read,
+                            step + 1,
+                            checkpoint,
+                            channels,
+                            managed,
+                            task,
+                            config,
                         ),
                     },
                 ),
@@ -765,7 +976,7 @@ class Pregel(
         if (
             config is not None
             and config.get("configurable", {}).get(CONFIG_KEY_CHECKPOINTER)
-            and (interrupt_after or interrupt_before)
+            and (interrupt_after or interrupt_before or _has_nested_interrupts(self))
         ):
             checkpointer: Optional[BaseCheckpointSaver] = config["configurable"][
                 CONFIG_KEY_CHECKPOINTER
@@ -858,7 +1069,7 @@ class Pregel(
             {'type': 'task_result', 'timestamp': '2024-06-23T...+00:00', 'step': 2, 'payload': {'id': '...', 'name': 'b', 'result': [('alist', ['there'])]}}
             ```
         """
-        config = ensure_config(config)
+        config = ensure_config(merge_configs(self.config, config))
         callback_manager = get_callback_manager_for_config(config)
         run_manager = callback_manager.on_chain_start(
             dumpd(self),
@@ -972,6 +1183,7 @@ class Pregel(
                                     )
                                 else:
                                     loop.put_writes(task.id, [(ERROR, exc)])
+
                             else:
                                 # save task writes to checkpointer
                                 loop.put_writes(task.id, task.writes)
@@ -1098,7 +1310,7 @@ class Pregel(
             {'type': 'task_result', 'timestamp': '2024-06-23T...+00:00', 'step': 2, 'payload': {'id': '...', 'name': 'b', 'result': [('alist', ['there'])]}}
             ```
         """
-        config = ensure_config(config)
+        config = ensure_config(merge_configs(self.config, config))
         callback_manager = get_async_callback_manager_for_config(config)
         run_manager = await callback_manager.on_chain_start(
             dumpd(self),
@@ -1212,6 +1424,7 @@ class Pregel(
                         )
                         if not done:
                             break  # timed out
+
                         for fut in done:
                             task = futures.pop(fut)
                             if exc := _exception(fut):
@@ -1222,6 +1435,7 @@ class Pregel(
                                     )
                                 else:
                                     loop.put_writes(task.id, [(ERROR, exc)])
+
                             else:
                                 # save task writes to checkpointer
                                 loop.put_writes(task.id, task.writes)
@@ -1429,11 +1643,3 @@ def _panic_or_proceed(
             inflight.pop().cancel()
         # raise timeout error
         raise timeout_exc_cls(f"Timed out at step {step}")
-
-
-def _with_mode(mode: StreamMode, on: bool, iter: Iterator[Any]) -> Iterator[Any]:
-    if on:
-        for chunk in iter:
-            yield (mode, chunk)
-    else:
-        yield from iter

libs/langgraph/langgraph/pregel/algo.py

@@ -24,7 +24,6 @@ from langchain_core.runnables.config import (
 )
 
 from langgraph.channels.base import BaseChannel
-from langgraph.channels.context import Context
 from langgraph.checkpoint.base import (
     BaseCheckpointSaver,
     Checkpoint,
@@ -95,28 +94,37 @@ def should_interrupt(
 
 
 def local_read(
+    step: int,
     checkpoint: Checkpoint,
     channels: Mapping[str, BaseChannel],
+    managed: ManagedValueMapping,
     task: WritesProtocol,
     config: RunnableConfig,
     select: Union[list[str], str],
     fresh: bool = False,
 ) -> Union[dict[str, Any], Any]:
+    if isinstance(select, str):
+        managed_keys = []
+    else:
+        managed_keys = [k for k in select if k in managed]
+        select = [k for k in select if k not in managed]
     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(channels, new_checkpoint, config, skip_context=True) as (
             channels,
             _,
         ):
-            all_channels = {**channels, **context_channels}
-            apply_writes(new_checkpoint, all_channels, [task], None)
-            return read_channels(all_channels, select)
+            apply_writes(new_checkpoint, channels, [task], None)
+            values = read_channels(channels, select)
     else:
-        return read_channels(channels, select)
+        values = read_channels(channels, select)
+    if managed_keys:
+        values.update({k: managed[k](step) for k in managed_keys})
+    return values
 
 
 def local_write(
+    step: int,
     commit: Callable[[Sequence[tuple[str, Any]]], None],
     processes: Mapping[str, PregelNode],
     channels: Mapping[str, BaseChannel],
@@ -131,6 +139,8 @@ def local_write(
                 )
             if value.node not in processes:
                 raise InvalidUpdateError(f"Invalid node name {value.node} in packet")
+            # replace any runtime values with placeholders
+            managed.replace_runtime_values(step, value.arg)
         elif chan not in channels and chan not in managed:
             logger.warning(f"Skipping write for channel '{chan}' which has no readers")
     commit(writes)
@@ -283,9 +293,9 @@ def prepare_next_tasks(
             "langgraph_task_idx": len(tasks),
         }
         checkpoint_ns = (
-            f"{parent_ns}{CHECKPOINT_NAMESPACE_SEPARATOR}{packet.node}"
+            f"{parent_ns}{CHECKPOINT_NAMESPACE_SEPARATOR}{packet.node}:{packet.id}"
             if parent_ns
-            else packet.node
+            else f"{packet.node}:{packet.id}"
         )
         task_id = str(
             uuid5(UUID(checkpoint["id"]), json.dumps((checkpoint_ns, metadata)))
@@ -293,6 +303,7 @@ def prepare_next_tasks(
         if for_execution:
             proc = processes[packet.node]
             if node := proc.get_node():
+                managed.replace_runtime_placeholders(step, packet.arg)
                 writes = deque()
                 tasks.append(
                     PregelExecutableTask(
@@ -317,6 +328,7 @@ def prepare_next_tasks(
                                 # deque.extend is thread-safe
                                 CONFIG_KEY_SEND: partial(
                                     local_write,
+                                    step,
                                     writes.extend,
                                     processes,
                                     channels,
@@ -324,13 +336,17 @@ def prepare_next_tasks(
                                 ),
                                 CONFIG_KEY_READ: partial(
                                     local_read,
+                                    step,
                                     checkpoint,
                                     channels,
+                                    managed,
                                     PregelTaskWrites(packet.node, writes, triggers),
                                     config,
                                 ),
-                                # in Send we can't checkpoint nested graphs
-                                # as they could be running in parallel
+                                CONFIG_KEY_CHECKPOINTER: checkpointer,
+                                CONFIG_KEY_RESUMING: is_resuming,
+                                "checkpoint_id": checkpoint["id"],
+                                "checkpoint_ns": checkpoint_ns,
                             },
                         ),
                         triggers,
@@ -412,6 +428,7 @@ def prepare_next_tasks(
                                     # deque.extend is thread-safe
                                     CONFIG_KEY_SEND: partial(
                                         local_write,
+                                        step,
                                         writes.extend,
                                         processes,
                                         channels,
@@ -419,8 +436,10 @@ def prepare_next_tasks(
                                     ),
                                     CONFIG_KEY_READ: partial(
                                         local_read,
+                                        step,
                                         checkpoint,
                                         channels,
+                                        managed,
                                         PregelTaskWrites(name, writes, triggers),
                                         config,
                                     ),

libs/langgraph/langgraph/pregel/get_state.py

@@ -0,0 +1,36 @@
+from langgraph.constants import CHECKPOINT_NAMESPACE_SEPARATOR
+from langgraph.pregel.types import StateSnapshot
+
+
+def assemble_state_snapshot_hierarchy(
+    root_checkpoint_ns: str,
+    checkpoint_ns_to_state_snapshots: dict[str, StateSnapshot],
+) -> StateSnapshot:
+    checkpoint_ns_list_to_visit = sorted(
+        checkpoint_ns_to_state_snapshots.keys(),
+        key=lambda x: len(x.split(CHECKPOINT_NAMESPACE_SEPARATOR)),
+    )
+    while checkpoint_ns_list_to_visit:
+        checkpoint_ns = checkpoint_ns_list_to_visit.pop()
+        state_snapshot = checkpoint_ns_to_state_snapshots[checkpoint_ns]
+        *path, subgraph_node = checkpoint_ns.split(CHECKPOINT_NAMESPACE_SEPARATOR)
+        parent_checkpoint_ns = CHECKPOINT_NAMESPACE_SEPARATOR.join(path)
+        if subgraph_node and (
+            parent_state_snapshot := checkpoint_ns_to_state_snapshots.get(
+                parent_checkpoint_ns
+            )
+        ):
+            parent_subgraph_snapshots = {
+                **(parent_state_snapshot.subgraph_state_snapshots or {}),
+                subgraph_node: state_snapshot,
+            }
+            checkpoint_ns_to_state_snapshots[
+                parent_checkpoint_ns
+            ] = checkpoint_ns_to_state_snapshots[parent_checkpoint_ns]._replace(
+                subgraph_state_snapshots=parent_subgraph_snapshots
+            )
+
+    state_snapshot = checkpoint_ns_to_state_snapshots.pop(root_checkpoint_ns, None)
+    if state_snapshot is None:
+        raise ValueError(f"Missing checkpoint for checkpoint NS '{root_checkpoint_ns}'")
+    return state_snapshot

libs/langgraph/langgraph/pregel/manager.py

@@ -5,8 +5,6 @@ from typing import AsyncIterator, Iterator, Mapping, Optional, Union
 from langchain_core.runnables import RunnableConfig, patch_config
 
 from langgraph.channels.base import BaseChannel
-from langgraph.channels.context import Context
-from langgraph.channels.last_value import LastValue
 from langgraph.checkpoint.base import Checkpoint
 from langgraph.constants import CONFIG_KEY_STORE
 from langgraph.managed.base import (
@@ -14,6 +12,7 @@ from langgraph.managed.base import (
     ManagedValueMapping,
     ManagedValueSpec,
 )
+from langgraph.managed.context import Context
 from langgraph.store.base import BaseStore
 
 
@@ -31,10 +30,12 @@ def ChannelsManager(
     channel_specs: Mapping[str, BaseChannel] = {}
     managed_specs: Mapping[str, ManagedValueSpec] = {}
     for k, v in specs.items():
-        if skip_context and isinstance(v, Context):
-            channel_specs[k] = LastValue(None)
-        elif isinstance(v, BaseChannel):
+        if isinstance(v, BaseChannel):
             channel_specs[k] = v
+        elif (
+            skip_context and isinstance(v, ConfiguredManagedValue) and v.cls is Context
+        ):
+            managed_specs[k] = Context.of(noop_context)
         else:
             managed_specs[k] = v
     with ExitStack() as stack:
@@ -45,14 +46,16 @@ def ChannelsManager(
                 )
                 for k, v in channel_specs.items()
             },
-            {
-                key: stack.enter_context(
-                    value.cls.enter(config_for_managed, **value.kwargs)
-                    if isinstance(value, ConfiguredManagedValue)
-                    else value.enter(config_for_managed)
-                )
-                for key, value in managed_specs.items()
-            },
+            ManagedValueMapping(
+                {
+                    key: stack.enter_context(
+                        value.cls.enter(config_for_managed, **value.kwargs)
+                        if isinstance(value, ConfiguredManagedValue)
+                        else value.enter(config_for_managed)
+                    )
+                    for key, value in managed_specs.items()
+                }
+            ),
         )
 
 
@@ -70,10 +73,12 @@ async def AsyncChannelsManager(
     channel_specs: Mapping[str, BaseChannel] = {}
     managed_specs: Mapping[str, ManagedValueSpec] = {}
     for k, v in specs.items():
-        if skip_context and isinstance(v, Context):
-            channel_specs[k] = LastValue(None)
-        elif isinstance(v, BaseChannel):
+        if isinstance(v, BaseChannel):
             channel_specs[k] = v
+        elif (
+            skip_context and isinstance(v, ConfiguredManagedValue) and v.cls is Context
+        ):
+            managed_specs[k] = Context.of(noop_context)
         else:
             managed_specs[k] = v
     async with AsyncExitStack() as stack:
@@ -102,5 +107,10 @@ async def AsyncChannelsManager(
                 for k, v in channel_specs.items()
             },
             # managed: build mapping from spec to result
-            {tasks[task]: task.result() for task in done},
+            ManagedValueMapping({tasks[task]: task.result() for task in done}),
         )
+
+
+@contextmanager
+def noop_context() -> Iterator[None]:
+    yield None

libs/langgraph/langgraph/pregel/read.py

@@ -67,19 +67,19 @@ class ChannelRead(RunnableCallable):
 
     def _read(self, _: Any, config: RunnableConfig) -> Any:
         return self.do_read(
-            config, channel=self.channel, fresh=self.fresh, mapper=self.mapper
+            config, select=self.channel, fresh=self.fresh, mapper=self.mapper
         )
 
     async def _aread(self, _: Any, config: RunnableConfig) -> Any:
         return self.do_read(
-            config, channel=self.channel, fresh=self.fresh, mapper=self.mapper
+            config, select=self.channel, fresh=self.fresh, mapper=self.mapper
         )
 
     @staticmethod
     def do_read(
         config: RunnableConfig,
         *,
-        channel: Union[str, list[str]],
+        select: Union[str, list[str]],
         fresh: bool = False,
         mapper: Optional[Callable[[Any], Any]] = None,
     ) -> Any:
@@ -91,9 +91,9 @@ class ChannelRead(RunnableCallable):
                 "Make sure to call in the context of a Pregel process"
             )
         if mapper:
-            return mapper(read(channel, fresh))
+            return mapper(read(select, fresh))
         else:
-            return read(channel, fresh)
+            return read(select, fresh)
 
 
 DEFAULT_BOUND: RunnablePassthrough = RunnablePassthrough()

libs/langgraph/langgraph/pregel/types.py

@@ -92,6 +92,8 @@ class StateSnapshot(NamedTuple):
     """Config used to fetch the parent snapshot, if any"""
     tasks: tuple[PregelTask, ...]
     """Tasks to execute in this step. If already attempted, may contain an error."""
+    subgraph_state_snapshots: Optional[dict[str, "StateSnapshot"]] = None
+    """State snapshots of subgraphs represented as a mapping from checkpoint namespace (`checkpoint_ns`) to snapshot."""
 
 
 All = Literal["*"]

libs/langgraph/poetry.lock

@@ -1832,7 +1832,7 @@ types-requests = ">=2.31.0.2,<3.0.0.0"
 
 [[package]]
 name = "langgraph-checkpoint"
-version = "1.0.2"
+version = "1.0.6"
 description = "Library with base interfaces for LangGraph checkpoint savers."
 optional = false
 python-versions = "^3.9.0,<4.0"
@@ -1848,7 +1848,7 @@ url = "../checkpoint"
 
 [[package]]
 name = "langgraph-checkpoint-postgres"
-version = "1.0.0"
+version = "1.0.3"
 description = "Library with a Postgres implementation of LangGraph checkpoint saver."
 optional = false
 python-versions = "^3.9.0,<4.0"
@@ -1858,7 +1858,8 @@ develop = true
 [package.dependencies]
 langgraph-checkpoint = "^1.0.1"
 orjson = ">=3.10.1"
-psycopg = {version = ">=3.1.19", extras = ["binary"]}
+psycopg = "^3.0.0"
+psycopg-pool = "^3.0.0"
 
 [package.source]
 type = "directory"
@@ -2638,6 +2639,20 @@ files = [
     {file = "psycopg_binary-3.2.1-cp39-cp39-win_amd64.whl", hash = "sha256:921f0c7f39590763d64a619de84d1b142587acc70fd11cbb5ba8fa39786f3073"},
 ]
 
+[[package]]
+name = "psycopg-pool"
+version = "3.2.2"
+description = "Connection Pool for Psycopg"
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "psycopg_pool-3.2.2-py3-none-any.whl", hash = "sha256:273081d0fbfaced4f35e69200c89cb8fbddfe277c38cc86c235b90a2ec2c8153"},
+    {file = "psycopg_pool-3.2.2.tar.gz", hash = "sha256:9e22c370045f6d7f2666a5ad1b0caf345f9f1912195b0b25d0d3bcc4f3a7389c"},
+]
+
+[package.dependencies]
+typing-extensions = ">=4.4"
+
 [[package]]
 name = "ptyprocess"
 version = "0.7.0"
@@ -4294,4 +4309,4 @@ test = ["big-O", "importlib-resources", "jaraco.functools", "jaraco.itertools",
 [metadata]
 lock-version = "2.0"
 python-versions = ">=3.9.0,<4.0"
-content-hash = "4ef9e25016072ce08554c8ff8d091104fb59da7cdab9a128312bd787e6f35146"
+content-hash = "b4d234e851639ecb33b6507b5e396b8b87aae6395c327441d04a7195eae72582"

libs/langgraph/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph"
-version = "0.2.12"
+version = "0.2.14"
 description = "Building stateful, multi-actor applications with LLMs"
 authors = []
 license = "MIT"
@@ -35,6 +35,7 @@ pytest-repeat = "^0.9.3"
 langgraph-checkpoint = {path = "../checkpoint", develop = true}
 langgraph-checkpoint-sqlite = {path = "../checkpoint-sqlite", develop = true}
 langgraph-checkpoint-postgres = {path = "../checkpoint-postgres", develop = true}
+psycopg = {extras = ["binary"], version = ">=3.0.0"}
 
 [tool.poetry.group.dev]
 optional = true

libs/langgraph/tests/any_str.py

@@ -23,24 +23,6 @@ class AnyVersion:
         return hash(str(self))
 
 
-class ExceptionLike:
-    def __init__(self, exc: Exception) -> None:
-        self.exc = exc
-
-    def __eq__(self, value: object) -> bool:
-        return (
-            isinstance(value, Exception)
-            and self.exc.__class__ == value.__class__
-            and str(self.exc) == str(value)
-        )
-
-    def __hash__(self) -> int:
-        return hash((self.exc.__class__, str(self.exc)))
-
-    def __repr__(self) -> str:
-        return str(self.exc)
-
-
 class UnsortedSequence:
     def __init__(self, *values: Any) -> None:
         self.seq = values

libs/langgraph/tests/conftest.py

@@ -7,6 +7,7 @@ from uuid import UUID, uuid4
 
 import pytest
 from psycopg import AsyncConnection, Connection
+from psycopg_pool import AsyncConnectionPool, ConnectionPool
 from pytest_mock import MockerFixture
 
 from langgraph.checkpoint.postgres import PostgresSaver
@@ -122,6 +123,26 @@ def checkpointer_postgres_pipe():
             conn.execute(f"DROP DATABASE {database}")
 
 
+@pytest.fixture(scope="function")
+def checkpointer_postgres_pool():
+    database = f"test_{uuid4().hex[:16]}"
+    # create unique db
+    with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
+        conn.execute(f"CREATE DATABASE {database}")
+    try:
+        # yield checkpointer
+        with ConnectionPool(
+            DEFAULT_POSTGRES_URI + database, max_size=10, kwargs={"autocommit": True}
+        ) as pool:
+            checkpointer = PostgresSaver(pool)
+            checkpointer.setup()
+            yield checkpointer
+    finally:
+        # drop unique db
+        with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
+            conn.execute(f"DROP DATABASE {database}")
+
+
 @pytest.fixture(scope="function")
 def checkpointer_postgres_aio():
     if sys.version_info < (3, 10):
@@ -185,3 +206,51 @@ async def _checkpointer_postgres_aio_pipe():
             DEFAULT_POSTGRES_URI, autocommit=True
         ) as conn:
             await conn.execute(f"DROP DATABASE {database}")
+
+
+@pytest.fixture(scope="function")
+def checkpointer_postgres_aio_pool():
+    if sys.version_info < (3, 10):
+        pytest.skip("Async Postgres tests require Python 3.10+")
+    with agen_to_gen(_checkpointer_postgres_aio_pool()) as checkpointer:
+        yield checkpointer
+
+
+@asynccontextmanager
+async def _checkpointer_postgres_aio_pool():
+    database = f"test_{uuid4().hex[:16]}"
+    # create unique db
+    async with await AsyncConnection.connect(
+        DEFAULT_POSTGRES_URI, autocommit=True
+    ) as conn:
+        await conn.execute(f"CREATE DATABASE {database}")
+    try:
+        # yield checkpointer
+        async with AsyncConnectionPool(
+            DEFAULT_POSTGRES_URI + database, max_size=10, kwargs={"autocommit": True}
+        ) as pool:
+            checkpointer = AsyncPostgresSaver(pool)
+            await checkpointer.setup()
+            yield checkpointer
+    finally:
+        # drop unique db
+        async with await AsyncConnection.connect(
+            DEFAULT_POSTGRES_URI, autocommit=True
+        ) as conn:
+            await conn.execute(f"DROP DATABASE {database}")
+
+
+ALL_CHECKPOINTERS_SYNC = [
+    "memory",
+    "sqlite",
+    "postgres",
+    "postgres_pipe",
+    "postgres_pool",
+]
+ALL_CHECKPOINTERS_ASYNC = [
+    "memory",
+    "sqlite_aio",
+    "postgres_aio",
+    "postgres_aio_pipe",
+    "postgres_aio_pool",
+]

libs/langgraph/tests/memory_assert.py

@@ -24,8 +24,6 @@ class NoopSerializer(SerializerProtocol):
 
 
 class MemorySaverAssertImmutable(MemorySaver):
-    serde = NoopSerializer()
-
     storage_for_copies: defaultdict[str, dict[str, dict[str, Checkpoint]]]
 
     def __init__(

libs/langgraph/tests/test_channels.py

@@ -1,14 +1,9 @@
 import operator
-from contextlib import asynccontextmanager, contextmanager
-from typing import AsyncGenerator, Generator, Sequence, Union
+from typing import Sequence, Union
 
-import httpx
 import pytest
-from langchain_core.runnables import RunnableConfig
-from pytest_mock import MockerFixture
 
 from langgraph.channels.binop import BinaryOperatorAggregate
-from langgraph.channels.context import Context
 from langgraph.channels.last_value import LastValue
 from langgraph.channels.topic import Topic
 from langgraph.errors import EmptyChannelError, InvalidUpdateError
@@ -96,50 +91,6 @@ async def test_topic_async() -> None:
         assert channel.get() == ["e"]
 
 
-def test_topic_unique() -> None:
-    with Topic(str, unique=True).from_checkpoint(None, {}) as channel:
-        assert channel.ValueType is Sequence[str]
-        assert channel.UpdateType is Union[str, list[str]]
-
-        assert channel.update(["a", "b"])
-        assert channel.get() == ["a", "b"]
-        assert channel.update(["b", ["c", "d"], "d"])
-        assert channel.get() == ["c", "d"], "de-dupes from current and previous steps"
-        assert channel.update([])
-        with pytest.raises(EmptyChannelError):
-            channel.get()
-        assert not channel.update([]), "channel already empty"
-        assert channel.update(["e"])
-        assert channel.get() == ["e"]
-        checkpoint = channel.checkpoint()
-    with Topic(str, unique=True).from_checkpoint(checkpoint, {}) as channel:
-        assert channel.get() == ["e"]
-        assert channel.update(["d", "f"])
-        assert channel.get() == ["f"], "de-dupes from checkpoint"
-
-
-async def test_topic_unique_async() -> None:
-    async with Topic(str, unique=True).afrom_checkpoint(None, {}) as channel:
-        assert channel.ValueType is Sequence[str]
-        assert channel.UpdateType is Union[str, list[str]]
-
-        assert channel.update(["a", "b"])
-        assert channel.get() == ["a", "b"]
-        assert channel.update(["b", ["c", "d"], "d"])
-        assert channel.get() == ["c", "d"], "de-dupes from current and previous steps"
-        assert channel.update([])
-        with pytest.raises(EmptyChannelError):
-            channel.get()
-        assert not channel.update([]), "channel already empty"
-        assert channel.update(["e"])
-        assert channel.get() == ["e"]
-        checkpoint = channel.checkpoint()
-    async with Topic(str, unique=True).afrom_checkpoint(checkpoint, {}) as channel:
-        assert channel.get() == ["e"]
-        assert channel.update(["d", "f"])
-        assert channel.get() == ["f"], "de-dupes from checkpoint"
-
-
 def test_topic_accumulate() -> None:
     with Topic(str, accumulate=True).from_checkpoint(None, {}) as channel:
         assert channel.ValueType is Sequence[str]
@@ -176,49 +127,6 @@ async def test_topic_accumulate_async() -> None:
         assert channel.get() == ["a", "b", "b", "c", "d", "d", "e"]
 
 
-def test_topic_unique_accumulate() -> None:
-    with Topic(str, unique=True, accumulate=True).from_checkpoint(None, {}) as channel:
-        assert channel.ValueType is Sequence[str]
-        assert channel.UpdateType is Union[str, list[str]]
-
-        assert channel.update(["a", "b"])
-        assert channel.get() == ["a", "b"]
-        assert channel.update(["b", ["c", "d"], "d"])
-        assert channel.get() == ["a", "b", "c", "d"]
-        assert not channel.update(["c"]), "no new values"
-        assert not channel.update([])
-        assert channel.get() == ["a", "b", "c", "d"]
-        checkpoint = channel.checkpoint()
-    with Topic(str, unique=True, accumulate=True).from_checkpoint(
-        checkpoint, {}
-    ) as channel:
-        assert channel.get() == ["a", "b", "c", "d"]
-        assert channel.update(["d", "e"])
-        assert channel.get() == ["a", "b", "c", "d", "e"]
-
-
-async def test_topic_unique_accumulate_async() -> None:
-    async with Topic(str, unique=True, accumulate=True).afrom_checkpoint(
-        None, {}
-    ) as channel:
-        assert channel.ValueType is Sequence[str]
-        assert channel.UpdateType is Union[str, list[str]]
-
-        channel.update(["a", "b"])
-        assert channel.get() == ["a", "b"]
-        channel.update(["b", ["c", "d"], "d"])
-        assert channel.get() == ["a", "b", "c", "d"]
-        channel.update([])
-        assert channel.get() == ["a", "b", "c", "d"]
-        checkpoint = channel.checkpoint()
-    async with Topic(str, unique=True, accumulate=True).afrom_checkpoint(
-        checkpoint, {}
-    ) as channel:
-        assert channel.get() == ["a", "b", "c", "d"]
-        channel.update(["d", "e"])
-        assert channel.get() == ["a", "b", "c", "d", "e"]
-
-
 def test_binop() -> None:
     with BinaryOperatorAggregate(int, operator.add).from_checkpoint(
         None, {}
@@ -257,80 +165,3 @@ async def test_binop_async() -> None:
         checkpoint, {}
     ) as channel:
         assert channel.get() == 10
-
-
-def test_ctx_manager(mocker: MockerFixture) -> None:
-    setup = mocker.Mock()
-    cleanup = mocker.Mock()
-
-    @contextmanager
-    def an_int() -> Generator[int, None, None]:
-        setup()
-        try:
-            yield 5
-        finally:
-            cleanup()
-
-    with Context(an_int, None).from_checkpoint(None, {}) as channel:
-        assert setup.call_count == 1
-        assert cleanup.call_count == 0
-
-        assert channel.ValueType is None
-        assert channel.UpdateType is None
-
-        assert channel.get() == 5
-
-        with pytest.raises(InvalidUpdateError):
-            channel.update([5])  # type: ignore
-
-    assert setup.call_count == 1
-    assert cleanup.call_count == 1
-
-
-def test_ctx_manager_ctx(mocker: MockerFixture) -> None:
-    with Context(httpx.Client).from_checkpoint(None, {}) as channel:
-        assert channel.ValueType is None
-        assert channel.UpdateType is None
-
-        assert isinstance(channel.get(), httpx.Client)
-
-        with pytest.raises(InvalidUpdateError):
-            channel.update([5])  # type: ignore
-
-        with pytest.raises(EmptyChannelError):
-            channel.checkpoint()
-
-
-async def test_ctx_manager_async(mocker: MockerFixture) -> None:
-    setup = mocker.Mock()
-    cleanup = mocker.Mock()
-
-    @contextmanager
-    def an_int_sync(config: RunnableConfig) -> Generator[int, None, None]:
-        try:
-            yield 5
-        finally:
-            pass
-
-    @asynccontextmanager
-    async def an_int() -> AsyncGenerator[int, None]:
-        setup()
-        try:
-            yield 5
-        finally:
-            cleanup()
-
-    async with Context(an_int_sync, an_int).afrom_checkpoint(None, {}) as channel:
-        assert setup.call_count == 1
-        assert cleanup.call_count == 0
-
-        assert channel.ValueType is None
-        assert channel.UpdateType is None
-
-        assert channel.get() == 5
-
-        with pytest.raises(InvalidUpdateError):
-            channel.update([5])  # type: ignore
-
-    assert setup.call_count == 1
-    assert cleanup.call_count == 1

libs/sdk-js/package.json

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

libs/sdk-js/src/client.mts

@@ -24,6 +24,7 @@ import {
 
 interface ClientConfig {
   apiUrl?: string;
+  apiKey?: string;
   callerOptions?: AsyncCallerParams;
   timeoutMs?: number;
   defaultHeaders?: Record<string, string | null | undefined>;
@@ -48,6 +49,9 @@ class BaseClient {
     this.timeoutMs = config?.timeoutMs || 12_000;
     this.apiUrl = config?.apiUrl || "http://localhost:8123";
     this.defaultHeaders = config?.defaultHeaders || {};
+    if (config?.apiKey != null) {
+      this.defaultHeaders["X-Api-Key"] = config.apiKey;
+    }
   }
 
   protected prepareFetchOptions(
@@ -555,6 +559,7 @@ export class RunsClient extends BaseClient {
       assistant_id: assistantId,
       interrupt_before: payload?.interruptBefore,
       interrupt_after: payload?.interruptAfter,
+      checkpoint_id: payload?.checkpointId,
     };
     if (payload?.multitaskStrategy != null) {
       json["multitask_strategy"] = payload?.multitaskStrategy;
@@ -571,6 +576,7 @@ export class RunsClient extends BaseClient {
     );
 
     let parser: EventSourceParser;
+    let onEndEvent: () => void;
     const textDecoder = new TextDecoder();
 
     const stream: ReadableStream<{ event: string; data: any }> = (
@@ -594,9 +600,17 @@ export class RunsClient extends BaseClient {
               });
             }
           });
+          onEndEvent = () => {
+            ctrl.enqueue({ event: "end", data: undefined });
+          };
         },
         async transform(chunk) {
-          parser.feed(textDecoder.decode(chunk));
+          const payload = textDecoder.decode(chunk);
+          parser.feed(payload);
+
+          // eventsource-parser will ignore events
+          // that are not terminated by a newline
+          if (payload.trim() === "event: end") onEndEvent();
         },
       }),
     );
@@ -625,6 +639,7 @@ export class RunsClient extends BaseClient {
       interrupt_before: payload?.interruptBefore,
       interrupt_after: payload?.interruptAfter,
       webhook: payload?.webhook,
+      checkpoint_id: payload?.checkpointId,
     };
     if (payload?.multitaskStrategy != null) {
       json["multitask_strategy"] = payload?.multitaskStrategy;
@@ -668,6 +683,7 @@ export class RunsClient extends BaseClient {
       assistant_id: assistantId,
       interrupt_before: payload?.interruptBefore,
       interrupt_after: payload?.interruptAfter,
+      checkpoint_id: payload?.checkpointId,
     };
     if (payload?.multitaskStrategy != null) {
       json["multitask_strategy"] = payload?.multitaskStrategy;

libs/sdk-js/src/schema.ts

@@ -2,6 +2,16 @@ import type { JSONSchema7 } from "json-schema";
 
 type Optional<T> = T | null | undefined;
 
+type RunStatus =
+  | "pending"
+  | "running"
+  | "error"
+  | "success"
+  | "timeout"
+  | "interrupted";
+
+type ThreadStatus = "idle" | "busy" | "interrupted";
+
 export interface Config {
   /**
    * Tags for this call and any sub-calls (eg. a Chain calling an LLM).
@@ -80,6 +90,7 @@ export interface Thread {
   created_at: string;
   updated_at: string;
   metadata: Metadata;
+  status: ThreadStatus;
 }
 
 export interface Cron {
@@ -112,12 +123,6 @@ export interface Run {
   assistant_id: string;
   created_at: string;
   updated_at: string;
-  status:
-    | "pending"
-    | "running"
-    | "error"
-    | "success"
-    | "timeout"
-    | "interrupted";
+  status: RunStatus;
   metadata: Metadata;
 }

libs/sdk-js/src/types.mts

@@ -30,6 +30,11 @@ interface RunsInvokePayload {
    */
   config?: Config;
 
+  /**
+   * Checkpoint ID for when creating a new run.
+   */
+  checkpointId?: string;
+
   /**
    * Interrupt execution before entering these nodes.
    */

libs/sdk-js/yarn.lock

@@ -1219,9 +1219,9 @@ micromark@^2.11.3, micromark@~2.11.0, micromark@~2.11.3:
     parse-entities "^2.0.0"
 
 micromatch@^4.0.4:
-  version "4.0.7"
-  resolved "https://registry.yarnpkg.com/micromatch/-/micromatch-4.0.7.tgz#33e8190d9fe474a9895525f5618eee136d46c2e5"
-  integrity sha512-LPP/3KorzCwBxfeUuZmaR6bG2kdeHSbe0P2tY3FLRU4vYrjYz5hI4QZwV0njUx3jeuKe67YukQ1LSPZBKDqO/Q==
+  version "4.0.8"
+  resolved "https://registry.yarnpkg.com/micromatch/-/micromatch-4.0.8.tgz#d66fa18f3a47076789320b9b1af32bd86d9fa202"
+  integrity sha512-PXwfBhYu0hBCPw8Dn0E+WDYb7af3dSLVWKi3HGv84IdF4TyFoC0ysxFd0Goxw7nSv4T/PzEJQxsYsEiFCKo2BA==
   dependencies:
     braces "^3.0.3"
     picomatch "^2.3.1"

libs/sdk-py/langgraph_sdk/client.py

@@ -25,9 +25,11 @@ from langgraph_sdk.schema import (
     Assistant,
     Config,
     Cron,
+    DisconnectMode,
     GraphSchema,
     Metadata,
     MultitaskStrategy,
+    OnCompletionBehavior,
     OnConflictBehavior,
     Run,
     RunCreate,
@@ -851,15 +853,14 @@ class ThreadsClient:
             thread_id: The ID of the thread to update.
             values: The values to update to the state.
             as_node: Update the state as if this node had just executed.
-
-            checkpoint_id: The ID of the checkpoint to get the state of.
+            checkpoint_id: The ID of the checkpoint to update the state of.
 
         Returns:
             None
 
         Example Usage:
 
-            await client.threads.get_state(
+            await client.threads.update_state(
                 thread_id="my_thread_id",
                 values={"messages":[{"role": "user", "content": "hello!"}]},
                 as_node="my_node",
@@ -963,6 +964,8 @@ class RunsClient:
         interrupt_before: Optional[list[str]] = None,
         interrupt_after: Optional[list[str]] = None,
         feedback_keys: Optional[list[str]] = None,
+        on_disconnect: Optional[DisconnectMode] = None,
+        webhook: Optional[str] = None,
         multitask_strategy: Optional[MultitaskStrategy] = None,
     ) -> AsyncIterator[StreamPart]:
         ...
@@ -980,6 +983,9 @@ class RunsClient:
         interrupt_before: Optional[list[str]] = None,
         interrupt_after: Optional[list[str]] = None,
         feedback_keys: Optional[list[str]] = None,
+        on_disconnect: Optional[DisconnectMode] = None,
+        webhook: Optional[str] = None,
+        on_completion: Optional[OnCompletionBehavior] = None,
     ) -> AsyncIterator[StreamPart]:
         ...
 
@@ -996,8 +1002,10 @@ class RunsClient:
         interrupt_before: Optional[list[str]] = None,
         interrupt_after: Optional[list[str]] = None,
         feedback_keys: Optional[list[str]] = None,
+        on_disconnect: Optional[DisconnectMode] = None,
         webhook: Optional[str] = None,
         multitask_strategy: Optional[MultitaskStrategy] = None,
+        on_completion: Optional[OnCompletionBehavior] = None,
     ) -> AsyncIterator[StreamPart]:
         """Create a run and stream the results.
 
@@ -1019,6 +1027,8 @@ class RunsClient:
             webhook: Webhook to call after LangGraph API call is done.
             multitask_strategy: Multitask strategy to use.
                 Must be one of 'reject', 'interrupt', 'rollback', or 'enqueue'.
+            on_disconnect: The disconnect mode to use.
+                Must be one of 'cancel' or 'continue'.
 
         Returns:
             AsyncIterator[StreamPart]: Asynchronous iterator of stream results.
@@ -1061,6 +1071,8 @@ class RunsClient:
             "webhook": webhook,
             "checkpoint_id": checkpoint_id,
             "multitask_strategy": multitask_strategy,
+            "on_disconnect": on_disconnect,
+            "on_completion": on_completion,
         }
         endpoint = (
             f"/threads/{thread_id}/runs/stream"
@@ -1083,6 +1095,7 @@ class RunsClient:
         interrupt_before: Optional[list[str]] = None,
         interrupt_after: Optional[list[str]] = None,
         webhook: Optional[str] = None,
+        on_completion: Optional[OnCompletionBehavior] = None,
     ) -> Run:
         ...
 
@@ -1116,6 +1129,7 @@ class RunsClient:
         interrupt_after: Optional[list[str]] = None,
         webhook: Optional[str] = None,
         multitask_strategy: Optional[MultitaskStrategy] = None,
+        on_completion: Optional[OnCompletionBehavior] = None,
     ) -> Run:
         """Create a background run.
 
@@ -1129,9 +1143,7 @@ class RunsClient:
             config: The configuration for the assistant.
             checkpoint_id: The checkpoint to start streaming from.
             interrupt_before: Nodes to interrupt immediately before they get executed.
-
             interrupt_after: Nodes to Nodes to interrupt immediately after they get executed.
-
             webhook: Webhook to call after LangGraph API call is done.
             multitask_strategy: Multitask strategy to use.
                 Must be one of 'reject', 'interrupt', 'rollback', or 'enqueue'.
@@ -1214,6 +1226,7 @@ class RunsClient:
             "webhook": webhook,
             "checkpoint_id": checkpoint_id,
             "multitask_strategy": multitask_strategy,
+            "on_completion": on_completion,
         }
         payload = {k: v for k, v in payload.items() if v is not None}
         if thread_id:
@@ -1242,6 +1255,8 @@ class RunsClient:
         checkpoint_id: Optional[str] = None,
         interrupt_before: Optional[list[str]] = None,
         interrupt_after: Optional[list[str]] = None,
+        webhook: Optional[str] = None,
+        on_disconnect: Optional[DisconnectMode] = None,
         multitask_strategy: Optional[MultitaskStrategy] = None,
     ) -> Union[list[dict], dict[str, Any]]:
         ...
@@ -1257,6 +1272,9 @@ class RunsClient:
         config: Optional[Config] = None,
         interrupt_before: Optional[list[str]] = None,
         interrupt_after: Optional[list[str]] = None,
+        webhook: Optional[str] = None,
+        on_disconnect: Optional[DisconnectMode] = None,
+        on_completion: Optional[OnCompletionBehavior] = None,
     ) -> Union[list[dict], dict[str, Any]]:
         ...
 
@@ -1272,7 +1290,9 @@ class RunsClient:
         interrupt_before: Optional[list[str]] = None,
         interrupt_after: Optional[list[str]] = None,
         webhook: Optional[str] = None,
+        on_disconnect: Optional[DisconnectMode] = None,
         multitask_strategy: Optional[MultitaskStrategy] = None,
+        on_completion: Optional[OnCompletionBehavior] = None,
     ) -> Union[list[dict], dict[str, Any]]:
         """Create a run, wait until it finishes and return the final state.
 
@@ -1286,12 +1306,12 @@ class RunsClient:
             config: The configuration for the assistant.
             checkpoint_id: The checkpoint to start streaming from.
             interrupt_before: Nodes to interrupt immediately before they get executed.
-
             interrupt_after: Nodes to Nodes to interrupt immediately after they get executed.
-
             webhook: Webhook to call after LangGraph API call is done.
             multitask_strategy: Multitask strategy to use.
                 Must be one of 'reject', 'interrupt', 'rollback', or 'enqueue'.
+            on_disconnect: The disconnect mode to use.
+                Must be one of 'cancel' or 'continue'.
 
         Returns:
             Union[list[dict], dict[str, Any]]: The output of the run.
@@ -1351,6 +1371,8 @@ class RunsClient:
             "webhook": webhook,
             "checkpoint_id": checkpoint_id,
             "multitask_strategy": multitask_strategy,
+            "on_disconnect": on_disconnect,
+            "on_completion": on_completion,
         }
         endpoint = (
             f"/threads/{thread_id}/runs/wait" if thread_id is not None else "/runs/wait"
@@ -1451,6 +1473,28 @@ class RunsClient:
         """  # noqa: E501
         return await self.http.get(f"/threads/{thread_id}/runs/{run_id}/join")
 
+    def join_stream(self, thread_id: str, run_id: str) -> AsyncIterator[StreamPart]:
+        """Stream output from a run in real-time, until the run is done.
+        Output is not buffered, so any output produced before this call will
+        not be received here.
+
+        Args:
+            thread_id: The thread ID to join.
+            run_id: The run ID to join.
+
+        Returns:
+            None
+
+        Example Usage:
+
+            await client.runs.join(
+                thread_id="thread_id_to_join",
+                run_id="run_id_to_join"
+            )
+
+        """  # noqa: E501
+        return self.http.stream(f"/threads/{thread_id}/runs/{run_id}/stream", "GET")
+
     async def delete(self, thread_id: str, run_id: str) -> None:
         """Delete a run.
 

libs/sdk-py/langgraph_sdk/schema.py

@@ -9,10 +9,14 @@ ThreadStatus = Literal["idle", "busy", "interrupted"]
 
 StreamMode = Literal["values", "messages", "updates", "events", "debug"]
 
+DisconnectMode = Literal["cancel", "continue"]
+
 MultitaskStrategy = Literal["reject", "interrupt", "rollback", "enqueue"]
 
 OnConflictBehavior = Literal["raise", "do_nothing"]
 
+OnCompletionBehavior = Literal["delete", "keep"]
+
 All = Literal["*"]
 
 

libs/sdk-py/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph-sdk"
-version = "0.1.28"
+version = "0.1.29"
 description = "SDK for interacting with LangGraph API"
 authors = []
 license = "MIT"