cli: release 0.1.52 (#1516)

feat(cli): add building support for JS

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

@@ -38,6 +38,7 @@ _MANUAL = {
         "visualization.ipynb",
         "state-model.ipynb",
         "subgraph.ipynb",
+        "recursion-limit.ipynb",
         "force-calling-a-tool-first.ipynb",
         "pass-run-time-values-to-tools.ipynb",
         "tool-calling.ipynb",
@@ -56,11 +57,13 @@ _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",
         "human_in_the_loop/review-tool-calls.ipynb",
         "node-retries.ipynb",
+        "react-agent-structured-output.ipynb",
     ],
     "tutorials": [
         "introduction.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

@@ -401,7 +401,17 @@ def node_a(state, config):
     ...
 ```
 
-See [this guide](../how-tos/configuration.ipynb) for a full breakdown on configuration
+See [this guide](../how-tos/configuration.ipynb) for a full breakdown on configuration.
+
+### Recursion Limit
+
+The recursion limit sets the maximum number of [super-steps](#graphs) the graph can execute during a single execution. Once the limit is reached, LangGraph will raise `GraphRecursionError`. By default this value is set to 25 steps. The recursion limit can be set on any graph at runtime, and is passed to `.invoke`/`.stream` via the config dictionary. Importantly, `recursion_limit` is a standalone `config` key and should not be passed inside the `configurable` key as all other user-defined configuration. See the example below:
+
+```python
+graph.invoke(inputs, config={"recursion_limit": 5, "configurable":{"llm": "anthropic"}})
+```
+
+Read [this how-to] to learn more about how the recursion limit works.
 
 ## Breakpoints
 
@@ -433,7 +443,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

@@ -15,7 +15,7 @@ These how-to guides show how to achieve that controllability.
 - [How to create subgraphs](subgraph.ipynb)
 - [How to create branches for parallel execution](branching.ipynb)
 - [How to create map-reduce branches for parallel execution](map-reduce.ipynb)
-
+- [How to control graph recursion limit](recursion-limit.ipynb)
 
 ## Persistence
 
@@ -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)
@@ -78,6 +79,7 @@ These guides show how to use different streaming modes.
 - [How to use a Pydantic model as your state](state-model.ipynb)
 - [How to use a context object in state](state-context-key.ipynb)
 - [How to add node retries](node-retries.ipynb)
+- [How to force function calling agent to structure output](react-agent-structured-output.ipynb)
 
 ## Prebuilt ReAct Agent
 

docs/mkdocs.yml

@@ -129,6 +129,7 @@ nav:
           - Create subgraphs: how-tos/subgraph.ipynb
           - Create branches for parallel execution: how-tos/branching.ipynb
           - Create map-reduce branches for parallel execution: how-tos/map-reduce.ipynb
+          - Control graph recursion limit: how-tos/recursion-limit.ipynb
       - Persistence:
           - Add persistence ("memory"): how-tos/persistence.ipynb
           - Manage conversation history: how-tos/memory/manage-conversation-history.ipynb
@@ -139,6 +140,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
@@ -169,6 +171,7 @@ nav:
           - Visualize your graph: how-tos/visualization.ipynb
           - Add runtime configuration: how-tos/configuration.ipynb
           - Add node retries: how-tos/node-retries.ipynb
+          - How to force function calling agent to structure output: how-tos/react-agent-structured-output.ipynb
       - Prebuilt ReAct Agent:
           - Create a ReAct agent: how-tos/create-react-agent.ipynb
           - Add memory to a ReAct agent: how-tos/create-react-agent-memory.ipynb

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/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/memory/manage-conversation-history.ipynb

@@ -268,8 +268,8 @@
     "\n",
     "\n",
     "def filter_messages(messages: list):\n",
-    "    # This is very simple helper function which only ever uses the last two messages\n",
-    "    return messages[-2:]\n",
+    "    # This is very simple helper function which only ever uses the last message\n",
+    "    return messages[-1:]\n",
     "\n",
     "\n",
     "# Define the function that calls the model\n",
@@ -372,9 +372,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "langgraph-example-dev",
    "language": "python",
-   "name": "python3"
+   "name": "langgraph-example-dev"
   },
   "language_info": {
    "codemirror_mode": {
@@ -386,7 +386,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.1"
+   "version": "3.11.9"
   }
  },
  "nbformat": 4,

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

@@ -122,7 +122,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "DB_URI = \"postgresql://postgres:postgres@localhost:5441/postgres?sslmode=disable\""
+    "DB_URI = \"postgresql://postgres:postgres@localhost:5442/postgres?sslmode=disable\""
    ]
   },
   {
@@ -134,10 +134,9 @@
    "source": [
     "from psycopg.rows import dict_row\n",
     "\n",
-    "connection_kwargs ={\n",
+    "connection_kwargs = {\n",
     "    \"autocommit\": True,\n",
     "    \"prepare_threshold\": 0,\n",
-    "    \"row_factory\": dict_row,\n",
     "}"
    ]
   },
@@ -166,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",
@@ -394,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",
@@ -551,9 +550,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "langgraph-postgres",
+   "display_name": "langgraph",
    "language": "python",
-   "name": "langgraph-postgres"
+   "name": "langgraph"
   },
   "language_info": {
    "codemirror_mode": {

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/rag/langgraph_agentic_rag.ipynb

@@ -20,7 +20,10 @@
    "id": "969fb438",
    "metadata": {},
    "outputs": [],
-   "source": ["%%capture --no-stderr\n%pip install -U --quiet langchain-community tiktoken langchain-openai langchainhub chromadb langchain langgraph langchain-text-splitters"]
+   "source": [
+    "%%capture --no-stderr\n",
+    "%pip install -U --quiet langchain-community tiktoken langchain-openai langchainhub chromadb langchain langgraph langchain-text-splitters"
+   ]
   },
   {
    "cell_type": "code",
@@ -28,7 +31,22 @@
    "id": "e4958a8c",
    "metadata": {},
    "outputs": [],
-   "source": ["import getpass\nimport os\n\n\ndef _set_env(key: str):\n    if key not in os.environ:\n        os.environ[key] = getpass.getpass(f\"{key}:\")\n\n\n_set_env(\"OPENAI_API_KEY\")\n\n# (Optional) For tracing\nos.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\n_set_env(\"LANGCHAIN_API_KEY\")"]
+   "source": [
+    "import getpass\n",
+    "import os\n",
+    "\n",
+    "\n",
+    "def _set_env(key: str):\n",
+    "    if key not in os.environ:\n",
+    "        os.environ[key] = getpass.getpass(f\"{key}:\")\n",
+    "\n",
+    "\n",
+    "_set_env(\"OPENAI_API_KEY\")\n",
+    "\n",
+    "# (Optional) For tracing\n",
+    "os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\n",
+    "_set_env(\"LANGCHAIN_API_KEY\")"
+   ]
   },
   {
    "cell_type": "markdown",
@@ -46,7 +64,34 @@
    "id": "e50c9efe-4abe-42fa-b35a-05eeeede9ec6",
    "metadata": {},
    "outputs": [],
-   "source": ["from langchain_community.document_loaders import WebBaseLoader\nfrom langchain_community.vectorstores import Chroma\nfrom langchain_openai import OpenAIEmbeddings\nfrom langchain_text_splitters import RecursiveCharacterTextSplitter\n\nurls = [\n    \"https://lilianweng.github.io/posts/2023-06-23-agent/\",\n    \"https://lilianweng.github.io/posts/2023-03-15-prompt-engineering/\",\n    \"https://lilianweng.github.io/posts/2023-10-25-adv-attack-llm/\",\n]\n\ndocs = [WebBaseLoader(url).load() for url in urls]\ndocs_list = [item for sublist in docs for item in sublist]\n\ntext_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(\n    chunk_size=100, chunk_overlap=50\n)\ndoc_splits = text_splitter.split_documents(docs_list)\n\n# Add to vectorDB\nvectorstore = Chroma.from_documents(\n    documents=doc_splits,\n    collection_name=\"rag-chroma\",\n    embedding=OpenAIEmbeddings(),\n)\nretriever = vectorstore.as_retriever()"]
+   "source": [
+    "from langchain_community.document_loaders import WebBaseLoader\n",
+    "from langchain_community.vectorstores import Chroma\n",
+    "from langchain_openai import OpenAIEmbeddings\n",
+    "from langchain_text_splitters import RecursiveCharacterTextSplitter\n",
+    "\n",
+    "urls = [\n",
+    "    \"https://lilianweng.github.io/posts/2023-06-23-agent/\",\n",
+    "    \"https://lilianweng.github.io/posts/2023-03-15-prompt-engineering/\",\n",
+    "    \"https://lilianweng.github.io/posts/2023-10-25-adv-attack-llm/\",\n",
+    "]\n",
+    "\n",
+    "docs = [WebBaseLoader(url).load() for url in urls]\n",
+    "docs_list = [item for sublist in docs for item in sublist]\n",
+    "\n",
+    "text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(\n",
+    "    chunk_size=100, chunk_overlap=50\n",
+    ")\n",
+    "doc_splits = text_splitter.split_documents(docs_list)\n",
+    "\n",
+    "# Add to vectorDB\n",
+    "vectorstore = Chroma.from_documents(\n",
+    "    documents=doc_splits,\n",
+    "    collection_name=\"rag-chroma\",\n",
+    "    embedding=OpenAIEmbeddings(),\n",
+    ")\n",
+    "retriever = vectorstore.as_retriever()"
+   ]
   },
   {
    "cell_type": "markdown",
@@ -62,7 +107,17 @@
    "id": "0b97bdd8-d7e3-444d-ac96-5ef4725f9048",
    "metadata": {},
    "outputs": [],
-   "source": ["from langchain.tools.retriever import create_retriever_tool\n\nretriever_tool = create_retriever_tool(\n    retriever,\n    \"retrieve_blog_posts\",\n    \"Search and return information about Lilian Weng blog posts on LLM agents, prompt engineering, and adversarial attacks on LLMs.\",\n)\n\ntools = [retriever_tool]"]
+   "source": [
+    "from langchain.tools.retriever import create_retriever_tool\n",
+    "\n",
+    "retriever_tool = create_retriever_tool(\n",
+    "    retriever,\n",
+    "    \"retrieve_blog_posts\",\n",
+    "    \"Search and return information about Lilian Weng blog posts on LLM agents, prompt engineering, and adversarial attacks on LLMs.\",\n",
+    ")\n",
+    "\n",
+    "tools = [retriever_tool]"
+   ]
   },
   {
    "cell_type": "markdown",
@@ -86,7 +141,19 @@
    "id": "0e378706-47d5-425a-8ba0-57b9acffbd0c",
    "metadata": {},
    "outputs": [],
-   "source": ["from typing import Annotated, Sequence, TypedDict\n\nfrom langchain_core.messages import BaseMessage\n\nfrom langgraph.graph.message import add_messages\n\n\nclass AgentState(TypedDict):\n    # The add_messages function defines how an update should be processed\n    # Default is to replace. add_messages says \"append\"\n    messages: Annotated[Sequence[BaseMessage], add_messages]"]
+   "source": [
+    "from typing import Annotated, Sequence, TypedDict\n",
+    "\n",
+    "from langchain_core.messages import BaseMessage\n",
+    "\n",
+    "from langgraph.graph.message import add_messages\n",
+    "\n",
+    "\n",
+    "class AgentState(TypedDict):\n",
+    "    # The add_messages function defines how an update should be processed\n",
+    "    # Default is to replace. add_messages says \"append\"\n",
+    "    messages: Annotated[Sequence[BaseMessage], add_messages]"
+   ]
   },
   {
    "attachments": {
@@ -129,7 +196,173 @@
      ]
     }
    ],
-   "source": ["from typing import Annotated, Literal, Sequence, TypedDict\n\nfrom langchain import hub\nfrom langchain_core.messages import BaseMessage, HumanMessage\nfrom langchain_core.output_parsers import StrOutputParser\nfrom langchain_core.prompts import PromptTemplate\nfrom langchain_core.pydantic_v1 import BaseModel, Field\nfrom langchain_openai import ChatOpenAI\n\nfrom langgraph.prebuilt import tools_condition\n\n### Edges\n\n\ndef grade_documents(state) -> Literal[\"generate\", \"rewrite\"]:\n    \"\"\"\n    Determines whether the retrieved documents are relevant to the question.\n\n    Args:\n        state (messages): The current state\n\n    Returns:\n        str: A decision for whether the documents are relevant or not\n    \"\"\"\n\n    print(\"---CHECK RELEVANCE---\")\n\n    # Data model\n    class grade(BaseModel):\n        \"\"\"Binary score for relevance check.\"\"\"\n\n        binary_score: str = Field(description=\"Relevance score 'yes' or 'no'\")\n\n    # LLM\n    model = ChatOpenAI(temperature=0, model=\"gpt-4-0125-preview\", streaming=True)\n\n    # LLM with tool and validation\n    llm_with_tool = model.with_structured_output(grade)\n\n    # Prompt\n    prompt = PromptTemplate(\n        template=\"\"\"You are a grader assessing relevance of a retrieved document to a user question. \\n \n        Here is the retrieved document: \\n\\n {context} \\n\\n\n        Here is the user question: {question} \\n\n        If the document contains keyword(s) or semantic meaning related to the user question, grade it as relevant. \\n\n        Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question.\"\"\",\n        input_variables=[\"context\", \"question\"],\n    )\n\n    # Chain\n    chain = prompt | llm_with_tool\n\n    messages = state[\"messages\"]\n    last_message = messages[-1]\n\n    question = messages[0].content\n    docs = last_message.content\n\n    scored_result = chain.invoke({\"question\": question, \"context\": docs})\n\n    score = scored_result.binary_score\n\n    if score == \"yes\":\n        print(\"---DECISION: DOCS RELEVANT---\")\n        return \"generate\"\n\n    else:\n        print(\"---DECISION: DOCS NOT RELEVANT---\")\n        print(score)\n        return \"rewrite\"\n\n\n### Nodes\n\n\ndef agent(state):\n    \"\"\"\n    Invokes the agent model to generate a response based on the current state. Given\n    the question, it will decide to retrieve using the retriever tool, or simply end.\n\n    Args:\n        state (messages): The current state\n\n    Returns:\n        dict: The updated state with the agent response appended to messages\n    \"\"\"\n    print(\"---CALL AGENT---\")\n    messages = state[\"messages\"]\n    model = ChatOpenAI(temperature=0, streaming=True, model=\"gpt-4-turbo\")\n    model = model.bind_tools(tools)\n    response = model.invoke(messages)\n    # We return a list, because this will get added to the existing list\n    return {\"messages\": [response]}\n\n\ndef rewrite(state):\n    \"\"\"\n    Transform the query to produce a better question.\n\n    Args:\n        state (messages): The current state\n\n    Returns:\n        dict: The updated state with re-phrased question\n    \"\"\"\n\n    print(\"---TRANSFORM QUERY---\")\n    messages = state[\"messages\"]\n    question = messages[0].content\n\n    msg = [\n        HumanMessage(\n            content=f\"\"\" \\n \n    Look at the input and try to reason about the underlying semantic intent / meaning. \\n \n    Here is the initial question:\n    \\n ------- \\n\n    {question} \n    \\n ------- \\n\n    Formulate an improved question: \"\"\",\n        )\n    ]\n\n    # Grader\n    model = ChatOpenAI(temperature=0, model=\"gpt-4-0125-preview\", streaming=True)\n    response = model.invoke(msg)\n    return {\"messages\": [response]}\n\n\ndef generate(state):\n    \"\"\"\n    Generate answer\n\n    Args:\n        state (messages): The current state\n\n    Returns:\n         dict: The updated state with re-phrased question\n    \"\"\"\n    print(\"---GENERATE---\")\n    messages = state[\"messages\"]\n    question = messages[0].content\n    last_message = messages[-1]\n\n    question = messages[0].content\n    docs = last_message.content\n\n    # Prompt\n    prompt = hub.pull(\"rlm/rag-prompt\")\n\n    # LLM\n    llm = ChatOpenAI(model_name=\"gpt-3.5-turbo\", temperature=0, streaming=True)\n\n    # Post-processing\n    def format_docs(docs):\n        return \"\\n\\n\".join(doc.page_content for doc in docs)\n\n    # Chain\n    rag_chain = prompt | llm | StrOutputParser()\n\n    # Run\n    response = rag_chain.invoke({\"context\": docs, \"question\": question})\n    return {\"messages\": [response]}\n\n\nprint(\"*\" * 20 + \"Prompt[rlm/rag-prompt]\" + \"*\" * 20)\nprompt = hub.pull(\"rlm/rag-prompt\").pretty_print()  # Show what the prompt looks like"]
+   "source": [
+    "from typing import Annotated, Literal, Sequence, TypedDict\n",
+    "\n",
+    "from langchain import hub\n",
+    "from langchain_core.messages import BaseMessage, HumanMessage\n",
+    "from langchain_core.output_parsers import StrOutputParser\n",
+    "from langchain_core.prompts import PromptTemplate\n",
+    "from langchain_core.pydantic_v1 import BaseModel, Field\n",
+    "from langchain_openai import ChatOpenAI\n",
+    "\n",
+    "from langgraph.prebuilt import tools_condition\n",
+    "\n",
+    "### Edges\n",
+    "\n",
+    "\n",
+    "def grade_documents(state) -> Literal[\"generate\", \"rewrite\"]:\n",
+    "    \"\"\"\n",
+    "    Determines whether the retrieved documents are relevant to the question.\n",
+    "\n",
+    "    Args:\n",
+    "        state (messages): The current state\n",
+    "\n",
+    "    Returns:\n",
+    "        str: A decision for whether the documents are relevant or not\n",
+    "    \"\"\"\n",
+    "\n",
+    "    print(\"---CHECK RELEVANCE---\")\n",
+    "\n",
+    "    # Data model\n",
+    "    class grade(BaseModel):\n",
+    "        \"\"\"Binary score for relevance check.\"\"\"\n",
+    "\n",
+    "        binary_score: str = Field(description=\"Relevance score 'yes' or 'no'\")\n",
+    "\n",
+    "    # LLM\n",
+    "    model = ChatOpenAI(temperature=0, model=\"gpt-4-0125-preview\", streaming=True)\n",
+    "\n",
+    "    # LLM with tool and validation\n",
+    "    llm_with_tool = model.with_structured_output(grade)\n",
+    "\n",
+    "    # Prompt\n",
+    "    prompt = PromptTemplate(\n",
+    "        template=\"\"\"You are a grader assessing relevance of a retrieved document to a user question. \\n \n",
+    "        Here is the retrieved document: \\n\\n {context} \\n\\n\n",
+    "        Here is the user question: {question} \\n\n",
+    "        If the document contains keyword(s) or semantic meaning related to the user question, grade it as relevant. \\n\n",
+    "        Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question.\"\"\",\n",
+    "        input_variables=[\"context\", \"question\"],\n",
+    "    )\n",
+    "\n",
+    "    # Chain\n",
+    "    chain = prompt | llm_with_tool\n",
+    "\n",
+    "    messages = state[\"messages\"]\n",
+    "    last_message = messages[-1]\n",
+    "\n",
+    "    question = messages[0].content\n",
+    "    docs = last_message.content\n",
+    "\n",
+    "    scored_result = chain.invoke({\"question\": question, \"context\": docs})\n",
+    "\n",
+    "    score = scored_result.binary_score\n",
+    "\n",
+    "    if score == \"yes\":\n",
+    "        print(\"---DECISION: DOCS RELEVANT---\")\n",
+    "        return \"generate\"\n",
+    "\n",
+    "    else:\n",
+    "        print(\"---DECISION: DOCS NOT RELEVANT---\")\n",
+    "        print(score)\n",
+    "        return \"rewrite\"\n",
+    "\n",
+    "\n",
+    "### Nodes\n",
+    "\n",
+    "\n",
+    "def agent(state):\n",
+    "    \"\"\"\n",
+    "    Invokes the agent model to generate a response based on the current state. Given\n",
+    "    the question, it will decide to retrieve using the retriever tool, or simply end.\n",
+    "\n",
+    "    Args:\n",
+    "        state (messages): The current state\n",
+    "\n",
+    "    Returns:\n",
+    "        dict: The updated state with the agent response appended to messages\n",
+    "    \"\"\"\n",
+    "    print(\"---CALL AGENT---\")\n",
+    "    messages = state[\"messages\"]\n",
+    "    model = ChatOpenAI(temperature=0, streaming=True, model=\"gpt-4-turbo\")\n",
+    "    model = model.bind_tools(tools)\n",
+    "    response = model.invoke(messages)\n",
+    "    # We return a list, because this will get added to the existing list\n",
+    "    return {\"messages\": [response]}\n",
+    "\n",
+    "\n",
+    "def rewrite(state):\n",
+    "    \"\"\"\n",
+    "    Transform the query to produce a better question.\n",
+    "\n",
+    "    Args:\n",
+    "        state (messages): The current state\n",
+    "\n",
+    "    Returns:\n",
+    "        dict: The updated state with re-phrased question\n",
+    "    \"\"\"\n",
+    "\n",
+    "    print(\"---TRANSFORM QUERY---\")\n",
+    "    messages = state[\"messages\"]\n",
+    "    question = messages[0].content\n",
+    "\n",
+    "    msg = [\n",
+    "        HumanMessage(\n",
+    "            content=f\"\"\" \\n \n",
+    "    Look at the input and try to reason about the underlying semantic intent / meaning. \\n \n",
+    "    Here is the initial question:\n",
+    "    \\n ------- \\n\n",
+    "    {question} \n",
+    "    \\n ------- \\n\n",
+    "    Formulate an improved question: \"\"\",\n",
+    "        )\n",
+    "    ]\n",
+    "\n",
+    "    # Grader\n",
+    "    model = ChatOpenAI(temperature=0, model=\"gpt-4-0125-preview\", streaming=True)\n",
+    "    response = model.invoke(msg)\n",
+    "    return {\"messages\": [response]}\n",
+    "\n",
+    "\n",
+    "def generate(state):\n",
+    "    \"\"\"\n",
+    "    Generate answer\n",
+    "\n",
+    "    Args:\n",
+    "        state (messages): The current state\n",
+    "\n",
+    "    Returns:\n",
+    "         dict: The updated state with re-phrased question\n",
+    "    \"\"\"\n",
+    "    print(\"---GENERATE---\")\n",
+    "    messages = state[\"messages\"]\n",
+    "    question = messages[0].content\n",
+    "    last_message = messages[-1]\n",
+    "\n",
+    "    docs = last_message.content\n",
+    "\n",
+    "    # Prompt\n",
+    "    prompt = hub.pull(\"rlm/rag-prompt\")\n",
+    "\n",
+    "    # LLM\n",
+    "    llm = ChatOpenAI(model_name=\"gpt-3.5-turbo\", temperature=0, streaming=True)\n",
+    "\n",
+    "    # Post-processing\n",
+    "    def format_docs(docs):\n",
+    "        return \"\\n\\n\".join(doc.page_content for doc in docs)\n",
+    "\n",
+    "    # Chain\n",
+    "    rag_chain = prompt | llm | StrOutputParser()\n",
+    "\n",
+    "    # Run\n",
+    "    response = rag_chain.invoke({\"context\": docs, \"question\": question})\n",
+    "    return {\"messages\": [response]}\n",
+    "\n",
+    "\n",
+    "print(\"*\" * 20 + \"Prompt[rlm/rag-prompt]\" + \"*\" * 20)\n",
+    "prompt = hub.pull(\"rlm/rag-prompt\").pretty_print()  # Show what the prompt looks like"
+   ]
   },
   {
    "cell_type": "markdown",
@@ -150,7 +383,48 @@
    "id": "8718a37f-83c2-4f16-9850-e61e0f49c3d4",
    "metadata": {},
    "outputs": [],
-   "source": ["from langgraph.graph import END, StateGraph, START\nfrom langgraph.prebuilt import ToolNode\n\n# Define a new graph\nworkflow = StateGraph(AgentState)\n\n# Define the nodes we will cycle between\nworkflow.add_node(\"agent\", agent)  # agent\nretrieve = ToolNode([retriever_tool])\nworkflow.add_node(\"retrieve\", retrieve)  # retrieval\nworkflow.add_node(\"rewrite\", rewrite)  # Re-writing the question\nworkflow.add_node(\n    \"generate\", generate\n)  # Generating a response after we know the documents are relevant\n# Call agent node to decide to retrieve or not\nworkflow.add_edge(START, \"agent\")\n\n# Decide whether to retrieve\nworkflow.add_conditional_edges(\n    \"agent\",\n    # Assess agent decision\n    tools_condition,\n    {\n        # Translate the condition outputs to nodes in our graph\n        \"tools\": \"retrieve\",\n        END: END,\n    },\n)\n\n# Edges taken after the `action` node is called.\nworkflow.add_conditional_edges(\n    \"retrieve\",\n    # Assess agent decision\n    grade_documents,\n)\nworkflow.add_edge(\"generate\", END)\nworkflow.add_edge(\"rewrite\", \"agent\")\n\n# Compile\ngraph = workflow.compile()"]
+   "source": [
+    "from langgraph.graph import END, StateGraph, START\n",
+    "from langgraph.prebuilt import ToolNode\n",
+    "\n",
+    "# Define a new graph\n",
+    "workflow = StateGraph(AgentState)\n",
+    "\n",
+    "# Define the nodes we will cycle between\n",
+    "workflow.add_node(\"agent\", agent)  # agent\n",
+    "retrieve = ToolNode([retriever_tool])\n",
+    "workflow.add_node(\"retrieve\", retrieve)  # retrieval\n",
+    "workflow.add_node(\"rewrite\", rewrite)  # Re-writing the question\n",
+    "workflow.add_node(\n",
+    "    \"generate\", generate\n",
+    ")  # Generating a response after we know the documents are relevant\n",
+    "# Call agent node to decide to retrieve or not\n",
+    "workflow.add_edge(START, \"agent\")\n",
+    "\n",
+    "# Decide whether to retrieve\n",
+    "workflow.add_conditional_edges(\n",
+    "    \"agent\",\n",
+    "    # Assess agent decision\n",
+    "    tools_condition,\n",
+    "    {\n",
+    "        # Translate the condition outputs to nodes in our graph\n",
+    "        \"tools\": \"retrieve\",\n",
+    "        END: END,\n",
+    "    },\n",
+    ")\n",
+    "\n",
+    "# Edges taken after the `action` node is called.\n",
+    "workflow.add_conditional_edges(\n",
+    "    \"retrieve\",\n",
+    "    # Assess agent decision\n",
+    "    grade_documents,\n",
+    ")\n",
+    "workflow.add_edge(\"generate\", END)\n",
+    "workflow.add_edge(\"rewrite\", \"agent\")\n",
+    "\n",
+    "# Compile\n",
+    "graph = workflow.compile()"
+   ]
   },
   {
    "cell_type": "code",
@@ -169,7 +443,15 @@
      "output_type": "display_data"
     }
    ],
-   "source": ["from IPython.display import Image, display\n\ntry:\n    display(Image(graph.get_graph(xray=True).draw_mermaid_png()))\nexcept Exception:\n    # This requires some extra dependencies and is optional\n    pass"]
+   "source": [
+    "from IPython.display import Image, display\n",
+    "\n",
+    "try:\n",
+    "    display(Image(graph.get_graph(xray=True).draw_mermaid_png()))\n",
+    "except Exception:\n",
+    "    # This requires some extra dependencies and is optional\n",
+    "    pass"
+   ]
   },
   {
    "cell_type": "code",
@@ -203,7 +485,21 @@
      ]
     }
    ],
-   "source": ["import pprint\n\ninputs = {\n    \"messages\": [\n        (\"user\", \"What does Lilian Weng say about the types of agent memory?\"),\n    ]\n}\nfor output in graph.stream(inputs):\n    for key, value in output.items():\n        pprint.pprint(f\"Output from node '{key}':\")\n        pprint.pprint(\"---\")\n        pprint.pprint(value, indent=2, width=80, depth=None)\n    pprint.pprint(\"\\n---\\n\")"]
+   "source": [
+    "import pprint\n",
+    "\n",
+    "inputs = {\n",
+    "    \"messages\": [\n",
+    "        (\"user\", \"What does Lilian Weng say about the types of agent memory?\"),\n",
+    "    ]\n",
+    "}\n",
+    "for output in graph.stream(inputs):\n",
+    "    for key, value in output.items():\n",
+    "        pprint.pprint(f\"Output from node '{key}':\")\n",
+    "        pprint.pprint(\"---\")\n",
+    "        pprint.pprint(value, indent=2, width=80, depth=None)\n",
+    "    pprint.pprint(\"\\n---\\n\")"
+   ]
   },
   {
    "cell_type": "code",
@@ -211,7 +507,7 @@
    "id": "189333cc-5d34-4869-9f9b-741210e1096f",
    "metadata": {},
    "outputs": [],
-   "source": [""]
+   "source": []
   }
  ],
  "metadata": {

examples/react-agent-structured-output.ipynb

@@ -0,0 +1,238 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# How to force function calling agent to structure output\n",
+    "\n",
+    "You might want your agent to return its output in a structured format. For example, if the output of the agent is used by some other downstream software, you may want the output to be in the same structured format every time the agent is invoked to ensure consistency.\n",
+    "\n",
+    "This guide shows how you can do this. We will be using a basic [ReAct agent](https://langchain-ai.github.io/langgraph/how-tos/create-react-agent/) (a model node and a tool-calling node) together with a third node at the end that will format response for the user.\n",
+    "\n",
+    "\n",
+    "## Setup\n",
+    "\n",
+    "### Structured Output\n",
+    "\n",
+    "First we need to define how we want to structure our output. To do this, we will use the `with_structured_output` method from LangChain, which you can read more about [here](https://python.langchain.com/v0.2/docs/how_to/structured_output/)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pydantic import BaseModel, Field\n",
+    "\n",
+    "class WeatherResponse(BaseModel):\n",
+    "    temperature: float = Field(description=\"The temperature in fahrenheit\")\n",
+    "    wind_directon: str = Field(description=\"The direction of the wind in abbreviated form\")\n",
+    "    wind_speed: float = Field(description=\"The speed of the wind in km/h\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Graph State\n",
+    "\n",
+    "We can now define our graph state:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from typing import Annotated, Any\n",
+    "from typing_extensions import TypedDict\n",
+    "from langgraph.graph.message import add_messages\n",
+    "\n",
+    "class AgentState(TypedDict):\n",
+    "    # list of chat messages from user, LLM, and tools\n",
+    "    messages: Annotated[list, add_messages]\n",
+    "    # Final structured response from the agent\n",
+    "    final_response: WeatherResponse"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Tools + Models\n",
+    "\n",
+    "We can now instantiate the tools and models we are going to use in our graph. We are going to use a single tool in this example for finding the weather, and we are going to have two models in our graph, one that does the function calling and one that does the responding."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from typing import Literal\n",
+    "from langchain_core.tools import tool\n",
+    "from langchain_anthropic import ChatAnthropic\n",
+    "\n",
+    "@tool\n",
+    "def get_weather(city: Literal[\"nyc\", \"sf\"]):\n",
+    "    \"\"\"Use this to get weather information.\"\"\"\n",
+    "    if city == \"nyc\":\n",
+    "        return \"It is cloudy in NYC, with 5 mph winds in the North-East direction and a temperature of 70 degrees\"\n",
+    "    elif city == \"sf\":\n",
+    "        return \"It is 75 degrees and sunny in SF, with 3 mph winds in the South-East direction\"\n",
+    "    else:\n",
+    "        raise AssertionError(\"Unknown city\")\n",
+    "    \n",
+    "tools = [get_weather]\n",
+    "    \n",
+    "model = ChatAnthropic(model=\"claude-3-opus-20240229\")\n",
+    "    \n",
+    "model_with_tools = model.bind_tools(tools)\n",
+    "model_with_structured_output = model.with_structured_output(WeatherResponse)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Define Graph\n",
+    "\n",
+    "Now that we have defined our tools and models, we can define our graph."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from langgraph.graph import StateGraph, END\n",
+    "from langgraph.prebuilt import ToolNode\n",
+    "from langchain_core.messages import HumanMessage\n",
+    "\n",
+    "# Define the function that determines whether to continue or not\n",
+    "def should_continue(state: AgentState):\n",
+    "    messages = state[\"messages\"]\n",
+    "    last_message = messages[-1]\n",
+    "    # If there is no function call, then we respond to the user\n",
+    "    if not last_message.tool_calls:\n",
+    "        return \"respond\"\n",
+    "    # Otherwise if there is, we continue\n",
+    "    else:\n",
+    "        return \"continue\"\n",
+    "\n",
+    "# Define the function that calls the model\n",
+    "def call_model(state: AgentState):\n",
+    "    response = model_with_tools.invoke(state['messages'])\n",
+    "    # We return a list, because this will get added to the existing list\n",
+    "    return {\"messages\": [response]}\n",
+    "\n",
+    "# Define the function that responds to the user\n",
+    "def respond(state: AgentState):\n",
+    "    # We call the model with structured output in order to return the same format to the user every time\n",
+    "    # state['messages'][-2] is the last ToolMessage in the convo, which we convert to a HumanMessage for the model to use\n",
+    "    # We could also pass the entire chat history, but this saves tokens since all we care to structure is the output of the tool\n",
+    "    response = model_with_structured_output.invoke([HumanMessage(content=state['messages'][-2].content)])\n",
+    "    # We return the final answer\n",
+    "    return {\"final_response\": response}\n",
+    "\n",
+    "\n",
+    "# Define a new graph\n",
+    "workflow = StateGraph(AgentState)\n",
+    "\n",
+    "# Define the two nodes we will cycle between\n",
+    "workflow.add_node(\"agent\", call_model)\n",
+    "workflow.add_node(\"respond\", respond)\n",
+    "workflow.add_node(\"tools\", ToolNode(tools))\n",
+    "\n",
+    "# Set the entrypoint as `agent`\n",
+    "# This means that this node is the first one called\n",
+    "workflow.set_entry_point(\"agent\")\n",
+    "\n",
+    "# We now add a conditional edge\n",
+    "workflow.add_conditional_edges(\n",
+    "    \"agent\",\n",
+    "    should_continue,\n",
+    "    {\n",
+    "        \"continue\": \"tools\",\n",
+    "        \"respond\": \"respond\",\n",
+    "    },\n",
+    ")\n",
+    "\n",
+    "workflow.add_edge(\"tools\", \"agent\")\n",
+    "workflow.add_edge(\"respond\", END)\n",
+    "graph = workflow.compile()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Usage\n",
+    "\n",
+    "We can now invoke our graph to verify that the output is being structured as desired:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "answer = graph.invoke(input={\"messages\": [(\"human\", \"what's the weather in SF?\")]})['final_response']"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "WeatherResponse(temperature=75.0, wind_directon='SE', wind_speed=3.0)"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "answer"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "As we can see, the agent returned a `WeatherResponse` object as we expected. If would now be easy to use this agent in a more complex software stack without having to worry about the output of the agent not matching the format expected from the next step in the stack."
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

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

@@ -48,7 +48,10 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "%%capture --no-stderr\n%pip install -U langchain_community langchain_openai langgraph wikipedia  scikit-learn  langchain_fireworks\n# We use one or the other search engine below\n%pip install -U duckduckgo tavily-python"
+        "%%capture --no-stderr\n",
+        "%pip install -U langchain_community langchain_openai langgraph wikipedia  scikit-learn  langchain_fireworks\n",
+        "# We use one or the other search engine below\n",
+        "%pip install -U duckduckgo tavily-python"
       ]
     },
     {
@@ -57,7 +60,10 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "# Uncomment if you want to draw the pretty graph diagrams.\n# If you are on MacOS, you will need to run brew install graphviz before installing and update some environment flags\n# ! brew install graphviz\n# !CFLAGS=\"-I $(brew --prefix graphviz)/include\" LDFLAGS=\"-L $(brew --prefix graphviz)/lib\" pip install -U pygraphviz"
+        "# Uncomment if you want to draw the pretty graph diagrams.\n",
+        "# If you are on MacOS, you will need to run brew install graphviz before installing and update some environment flags\n",
+        "# ! brew install graphviz\n",
+        "# !CFLAGS=\"-I $(brew --prefix graphviz)/include\" LDFLAGS=\"-L $(brew --prefix graphviz)/lib\" pip install -U pygraphviz"
       ]
     },
     {
@@ -66,7 +72,21 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "import getpass\nimport os\n\n\ndef _set_env(var: str):\n    if os.environ.get(var):\n        return\n    os.environ[var] = getpass.getpass(var + \":\")\n\n\n# Set for tracing\nos.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\nos.environ[\"LANGCHAIN_PROJECT\"] = \"STORM\"\n_set_env(\"LANGCHAIN_API_KEY\")\n_set_env(\"OPENAI_API_KEY\")"
+        "import getpass\n",
+        "import os\n",
+        "\n",
+        "\n",
+        "def _set_env(var: str):\n",
+        "    if os.environ.get(var):\n",
+        "        return\n",
+        "    os.environ[var] = getpass.getpass(var + \":\")\n",
+        "\n",
+        "\n",
+        "# Set for tracing\n",
+        "os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\n",
+        "os.environ[\"LANGCHAIN_PROJECT\"] = \"STORM\"\n",
+        "_set_env(\"LANGCHAIN_API_KEY\")\n",
+        "_set_env(\"OPENAI_API_KEY\")"
       ]
     },
     {
@@ -84,7 +104,12 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "from langchain_openai import ChatOpenAI\n\nfast_llm = ChatOpenAI(model=\"gpt-3.5-turbo\")\n# Uncomment for a Fireworks model\n# fast_llm = ChatFireworks(model=\"accounts/fireworks/models/firefunction-v1\", max_tokens=32_000)\nlong_context_llm = ChatOpenAI(model=\"gpt-4-turbo-preview\")"
+        "from langchain_openai import ChatOpenAI\n",
+        "\n",
+        "fast_llm = ChatOpenAI(model=\"gpt-3.5-turbo\")\n",
+        "# Uncomment for a Fireworks model\n",
+        "# fast_llm = ChatFireworks(model=\"accounts/fireworks/models/firefunction-v1\", max_tokens=32_000)\n",
+        "long_context_llm = ChatOpenAI(model=\"gpt-4-turbo-preview\")"
       ]
     },
     {
@@ -112,7 +137,64 @@
         }
       ],
       "source": [
-        "from typing import List, Optional\n\nfrom langchain_core.prompts import ChatPromptTemplate\nfrom langchain_core.pydantic_v1 import BaseModel, Field\n\ndirect_gen_outline_prompt = ChatPromptTemplate.from_messages(\n    [\n        (\n            \"system\",\n            \"You are a Wikipedia writer. Write an outline for a Wikipedia page about a user-provided topic. Be comprehensive and specific.\",\n        ),\n        (\"user\", \"{topic}\"),\n    ]\n)\n\n\nclass Subsection(BaseModel):\n    subsection_title: str = Field(..., title=\"Title of the subsection\")\n    description: str = Field(..., title=\"Content of the subsection\")\n\n    @property\n    def as_str(self) -> str:\n        return f\"### {self.subsection_title}\\n\\n{self.description}\".strip()\n\n\nclass Section(BaseModel):\n    section_title: str = Field(..., title=\"Title of the section\")\n    description: str = Field(..., title=\"Content of the section\")\n    subsections: Optional[List[Subsection]] = Field(\n        default=None,\n        title=\"Titles and descriptions for each subsection of the Wikipedia page.\",\n    )\n\n    @property\n    def as_str(self) -> str:\n        subsections = \"\\n\\n\".join(\n            f\"### {subsection.subsection_title}\\n\\n{subsection.description}\"\n            for subsection in self.subsections or []\n        )\n        return f\"## {self.section_title}\\n\\n{self.description}\\n\\n{subsections}\".strip()\n\n\nclass Outline(BaseModel):\n    page_title: str = Field(..., title=\"Title of the Wikipedia page\")\n    sections: List[Section] = Field(\n        default_factory=list,\n        title=\"Titles and descriptions for each section of the Wikipedia page.\",\n    )\n\n    @property\n    def as_str(self) -> str:\n        sections = \"\\n\\n\".join(section.as_str for section in self.sections)\n        return f\"# {self.page_title}\\n\\n{sections}\".strip()\n\n\ngenerate_outline_direct = direct_gen_outline_prompt | fast_llm.with_structured_output(\n    Outline\n)"
+        "from typing import List, Optional\n",
+        "\n",
+        "from langchain_core.prompts import ChatPromptTemplate\n",
+        "from langchain_core.pydantic_v1 import BaseModel, Field\n",
+        "\n",
+        "direct_gen_outline_prompt = ChatPromptTemplate.from_messages(\n",
+        "    [\n",
+        "        (\n",
+        "            \"system\",\n",
+        "            \"You are a Wikipedia writer. Write an outline for a Wikipedia page about a user-provided topic. Be comprehensive and specific.\",\n",
+        "        ),\n",
+        "        (\"user\", \"{topic}\"),\n",
+        "    ]\n",
+        ")\n",
+        "\n",
+        "\n",
+        "class Subsection(BaseModel):\n",
+        "    subsection_title: str = Field(..., title=\"Title of the subsection\")\n",
+        "    description: str = Field(..., title=\"Content of the subsection\")\n",
+        "\n",
+        "    @property\n",
+        "    def as_str(self) -> str:\n",
+        "        return f\"### {self.subsection_title}\\n\\n{self.description}\".strip()\n",
+        "\n",
+        "\n",
+        "class Section(BaseModel):\n",
+        "    section_title: str = Field(..., title=\"Title of the section\")\n",
+        "    description: str = Field(..., title=\"Content of the section\")\n",
+        "    subsections: Optional[List[Subsection]] = Field(\n",
+        "        default=None,\n",
+        "        title=\"Titles and descriptions for each subsection of the Wikipedia page.\",\n",
+        "    )\n",
+        "\n",
+        "    @property\n",
+        "    def as_str(self) -> str:\n",
+        "        subsections = \"\\n\\n\".join(\n",
+        "            f\"### {subsection.subsection_title}\\n\\n{subsection.description}\"\n",
+        "            for subsection in self.subsections or []\n",
+        "        )\n",
+        "        return f\"## {self.section_title}\\n\\n{self.description}\\n\\n{subsections}\".strip()\n",
+        "\n",
+        "\n",
+        "class Outline(BaseModel):\n",
+        "    page_title: str = Field(..., title=\"Title of the Wikipedia page\")\n",
+        "    sections: List[Section] = Field(\n",
+        "        default_factory=list,\n",
+        "        title=\"Titles and descriptions for each section of the Wikipedia page.\",\n",
+        "    )\n",
+        "\n",
+        "    @property\n",
+        "    def as_str(self) -> str:\n",
+        "        sections = \"\\n\\n\".join(section.as_str for section in self.sections)\n",
+        "        return f\"# {self.page_title}\\n\\n{sections}\".strip()\n",
+        "\n",
+        "\n",
+        "generate_outline_direct = direct_gen_outline_prompt | fast_llm.with_structured_output(\n",
+        "    Outline\n",
+        ")"
       ]
     },
     {
@@ -145,7 +227,11 @@
         }
       ],
       "source": [
-        "example_topic = \"Impact of million-plus token context window language models on RAG\"\n\ninitial_outline = generate_outline_direct.invoke({\"topic\": example_topic})\n\nprint(initial_outline.as_str)"
+        "example_topic = \"Impact of million-plus token context window language models on RAG\"\n",
+        "\n",
+        "initial_outline = generate_outline_direct.invoke({\"topic\": example_topic})\n",
+        "\n",
+        "print(initial_outline.as_str)"
       ]
     },
     {
@@ -165,7 +251,25 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "gen_related_topics_prompt = ChatPromptTemplate.from_template(\n    \"\"\"I'm writing a Wikipedia page for a topic mentioned below. Please identify and recommend some Wikipedia pages on closely related subjects. I'm looking for examples that provide insights into interesting aspects commonly associated with this topic, or examples that help me understand the typical content and structure included in Wikipedia pages for similar topics.\n\nPlease list the as many subjects and urls as you can.\n\nTopic of interest: {topic}\n\"\"\"\n)\n\n\nclass RelatedSubjects(BaseModel):\n    topics: List[str] = Field(\n        description=\"Comprehensive list of related subjects as background research.\",\n    )\n\n\nexpand_chain = gen_related_topics_prompt | fast_llm.with_structured_output(\n    RelatedSubjects\n)"
+        "gen_related_topics_prompt = ChatPromptTemplate.from_template(\n",
+        "    \"\"\"I'm writing a Wikipedia page for a topic mentioned below. Please identify and recommend some Wikipedia pages on closely related subjects. I'm looking for examples that provide insights into interesting aspects commonly associated with this topic, or examples that help me understand the typical content and structure included in Wikipedia pages for similar topics.\n",
+        "\n",
+        "Please list the as many subjects and urls as you can.\n",
+        "\n",
+        "Topic of interest: {topic}\n",
+        "\"\"\"\n",
+        ")\n",
+        "\n",
+        "\n",
+        "class RelatedSubjects(BaseModel):\n",
+        "    topics: List[str] = Field(\n",
+        "        description=\"Comprehensive list of related subjects as background research.\",\n",
+        "    )\n",
+        "\n",
+        "\n",
+        "expand_chain = gen_related_topics_prompt | fast_llm.with_structured_output(\n",
+        "    RelatedSubjects\n",
+        ")"
       ]
     },
     {
@@ -185,7 +289,8 @@
         }
       ],
       "source": [
-        "related_subjects = await expand_chain.ainvoke({\"topic\": example_topic})\nrelated_subjects"
+        "related_subjects = await expand_chain.ainvoke({\"topic\": example_topic})\n",
+        "related_subjects"
       ]
     },
     {
@@ -204,7 +309,49 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "class Editor(BaseModel):\n    affiliation: str = Field(\n        description=\"Primary affiliation of the editor.\",\n    )\n    name: str = Field(\n        description=\"Name of the editor.\", pattern=r\"^[a-zA-Z0-9_-]{1,64}$\"\n    )\n    role: str = Field(\n        description=\"Role of the editor in the context of the topic.\",\n    )\n    description: str = Field(\n        description=\"Description of the editor's focus, concerns, and motives.\",\n    )\n\n    @property\n    def persona(self) -> str:\n        return f\"Name: {self.name}\\nRole: {self.role}\\nAffiliation: {self.affiliation}\\nDescription: {self.description}\\n\"\n\n\nclass Perspectives(BaseModel):\n    editors: List[Editor] = Field(\n        description=\"Comprehensive list of editors with their roles and affiliations.\",\n        # Add a pydantic validation/restriction to be at most M editors\n    )\n\n\ngen_perspectives_prompt = ChatPromptTemplate.from_messages(\n    [\n        (\n            \"system\",\n            \"\"\"You need to select a diverse (and distinct) group of Wikipedia editors who will work together to create a comprehensive article on the topic. Each of them represents a different perspective, role, or affiliation related to this topic.\\\n    You can use other Wikipedia pages of related topics for inspiration. For each editor, add a description of what they will focus on.\n\n    Wiki page outlines of related topics for inspiration:\n    {examples}\"\"\",\n        ),\n        (\"user\", \"Topic of interest: {topic}\"),\n    ]\n)\n\ngen_perspectives_chain = gen_perspectives_prompt | ChatOpenAI(\n    model=\"gpt-3.5-turbo\"\n).with_structured_output(Perspectives)"
+        "class Editor(BaseModel):\n",
+        "    affiliation: str = Field(\n",
+        "        description=\"Primary affiliation of the editor.\",\n",
+        "    )\n",
+        "    name: str = Field(\n",
+        "        description=\"Name of the editor.\", pattern=r\"^[a-zA-Z0-9_-]{1,64}$\"\n",
+        "    )\n",
+        "    role: str = Field(\n",
+        "        description=\"Role of the editor in the context of the topic.\",\n",
+        "    )\n",
+        "    description: str = Field(\n",
+        "        description=\"Description of the editor's focus, concerns, and motives.\",\n",
+        "    )\n",
+        "\n",
+        "    @property\n",
+        "    def persona(self) -> str:\n",
+        "        return f\"Name: {self.name}\\nRole: {self.role}\\nAffiliation: {self.affiliation}\\nDescription: {self.description}\\n\"\n",
+        "\n",
+        "\n",
+        "class Perspectives(BaseModel):\n",
+        "    editors: List[Editor] = Field(\n",
+        "        description=\"Comprehensive list of editors with their roles and affiliations.\",\n",
+        "        # Add a pydantic validation/restriction to be at most M editors\n",
+        "    )\n",
+        "\n",
+        "\n",
+        "gen_perspectives_prompt = ChatPromptTemplate.from_messages(\n",
+        "    [\n",
+        "        (\n",
+        "            \"system\",\n",
+        "            \"\"\"You need to select a diverse (and distinct) group of Wikipedia editors who will work together to create a comprehensive article on the topic. Each of them represents a different perspective, role, or affiliation related to this topic.\\\n",
+        "    You can use other Wikipedia pages of related topics for inspiration. For each editor, add a description of what they will focus on.\n",
+        "\n",
+        "    Wiki page outlines of related topics for inspiration:\n",
+        "    {examples}\"\"\",\n",
+        "        ),\n",
+        "        (\"user\", \"Topic of interest: {topic}\"),\n",
+        "    ]\n",
+        ")\n",
+        "\n",
+        "gen_perspectives_chain = gen_perspectives_prompt | ChatOpenAI(\n",
+        "    model=\"gpt-3.5-turbo\"\n",
+        ").with_structured_output(Perspectives)"
       ]
     },
     {
@@ -213,7 +360,37 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "from langchain_community.retrievers import WikipediaRetriever\nfrom langchain_core.runnables import RunnableLambda\nfrom langchain_core.runnables import chain as as_runnable\n\nwikipedia_retriever = WikipediaRetriever(load_all_available_meta=True, top_k_results=1)\n\n\ndef format_doc(doc, max_length=1000):\n    related = \"- \".join(doc.metadata[\"categories\"])\n    return f\"### {doc.metadata['title']}\\n\\nSummary: {doc.page_content}\\n\\nRelated\\n{related}\"[\n        :max_length\n    ]\n\n\ndef format_docs(docs):\n    return \"\\n\\n\".join(format_doc(doc) for doc in docs)\n\n\n@as_runnable\nasync def survey_subjects(topic: str):\n    related_subjects = await expand_chain.ainvoke({\"topic\": topic})\n    retrieved_docs = await wikipedia_retriever.abatch(\n        related_subjects.topics, return_exceptions=True\n    )\n    all_docs = []\n    for docs in retrieved_docs:\n        if isinstance(docs, BaseException):\n            continue\n        all_docs.extend(docs)\n    formatted = format_docs(all_docs)\n    return await gen_perspectives_chain.ainvoke({\"examples\": formatted, \"topic\": topic})"
+        "from langchain_community.retrievers import WikipediaRetriever\n",
+        "from langchain_core.runnables import RunnableLambda\n",
+        "from langchain_core.runnables import chain as as_runnable\n",
+        "\n",
+        "wikipedia_retriever = WikipediaRetriever(load_all_available_meta=True, top_k_results=1)\n",
+        "\n",
+        "\n",
+        "def format_doc(doc, max_length=1000):\n",
+        "    related = \"- \".join(doc.metadata[\"categories\"])\n",
+        "    return f\"### {doc.metadata['title']}\\n\\nSummary: {doc.page_content}\\n\\nRelated\\n{related}\"[\n",
+        "        :max_length\n",
+        "    ]\n",
+        "\n",
+        "\n",
+        "def format_docs(docs):\n",
+        "    return \"\\n\\n\".join(format_doc(doc) for doc in docs)\n",
+        "\n",
+        "\n",
+        "@as_runnable\n",
+        "async def survey_subjects(topic: str):\n",
+        "    related_subjects = await expand_chain.ainvoke({\"topic\": topic})\n",
+        "    retrieved_docs = await wikipedia_retriever.abatch(\n",
+        "        related_subjects.topics, return_exceptions=True\n",
+        "    )\n",
+        "    all_docs = []\n",
+        "    for docs in retrieved_docs:\n",
+        "        if isinstance(docs, BaseException):\n",
+        "            continue\n",
+        "        all_docs.extend(docs)\n",
+        "    formatted = format_docs(all_docs)\n",
+        "    return await gen_perspectives_chain.ainvoke({\"examples\": formatted, \"topic\": topic})"
       ]
     },
     {
@@ -280,7 +457,40 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "from typing import Annotated\n\nfrom langchain_core.messages import AnyMessage\nfrom typing_extensions import TypedDict\n\nfrom langgraph.graph import END, StateGraph, START\n\n\ndef add_messages(left, right):\n    if not isinstance(left, list):\n        left = [left]\n    if not isinstance(right, list):\n        right = [right]\n    return left + right\n\n\ndef update_references(references, new_references):\n    if not references:\n        references = {}\n    references.update(new_references)\n    return references\n\n\ndef update_editor(editor, new_editor):\n    # Can only set at the outset\n    if not editor:\n        return new_editor\n    return editor\n\n\nclass InterviewState(TypedDict):\n    messages: Annotated[List[AnyMessage], add_messages]\n    references: Annotated[Optional[dict], update_references]\n    editor: Annotated[Optional[Editor], update_editor]"
+        "from typing import Annotated\n",
+        "\n",
+        "from langchain_core.messages import AnyMessage\n",
+        "from typing_extensions import TypedDict\n",
+        "\n",
+        "from langgraph.graph import END, StateGraph, START\n",
+        "\n",
+        "\n",
+        "def add_messages(left, right):\n",
+        "    if not isinstance(left, list):\n",
+        "        left = [left]\n",
+        "    if not isinstance(right, list):\n",
+        "        right = [right]\n",
+        "    return left + right\n",
+        "\n",
+        "\n",
+        "def update_references(references, new_references):\n",
+        "    if not references:\n",
+        "        references = {}\n",
+        "    references.update(new_references)\n",
+        "    return references\n",
+        "\n",
+        "\n",
+        "def update_editor(editor, new_editor):\n",
+        "    # Can only set at the outset\n",
+        "    if not editor:\n",
+        "        return new_editor\n",
+        "    return editor\n",
+        "\n",
+        "\n",
+        "class InterviewState(TypedDict):\n",
+        "    messages: Annotated[List[AnyMessage], add_messages]\n",
+        "    references: Annotated[Optional[dict], update_references]\n",
+        "    editor: Annotated[Optional[Editor], update_editor]"
       ]
     },
     {
@@ -298,7 +508,56 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "from langchain_core.messages import AIMessage, HumanMessage, ToolMessage\nfrom langchain_core.prompts import MessagesPlaceholder\n\ngen_qn_prompt = ChatPromptTemplate.from_messages(\n    [\n        (\n            \"system\",\n            \"\"\"You are an experienced Wikipedia writer and want to edit a specific page. \\\nBesides your identity as a Wikipedia writer, you have a specific focus when researching the topic. \\\nNow, you are chatting with an expert to get information. Ask good questions to get more useful information.\n\nWhen you have no more questions to ask, say \"Thank you so much for your help!\" to end the conversation.\\\nPlease only ask one question at a time and don't ask what you have asked before.\\\nYour questions should be related to the topic you want to write.\nBe comprehensive and curious, gaining as much unique insight from the expert as possible.\\\n\nStay true to your specific perspective:\n\n{persona}\"\"\",\n        ),\n        MessagesPlaceholder(variable_name=\"messages\", optional=True),\n    ]\n)\n\n\ndef tag_with_name(ai_message: AIMessage, name: str):\n    ai_message.name = name\n    return ai_message\n\n\ndef swap_roles(state: InterviewState, name: str):\n    converted = []\n    for message in state[\"messages\"]:\n        if isinstance(message, AIMessage) and message.name != name:\n            message = HumanMessage(**message.dict(exclude={\"type\"}))\n        converted.append(message)\n    return {\"messages\": converted}\n\n\n@as_runnable\nasync def generate_question(state: InterviewState):\n    editor = state[\"editor\"]\n    gn_chain = (\n        RunnableLambda(swap_roles).bind(name=editor.name)\n        | gen_qn_prompt.partial(persona=editor.persona)\n        | fast_llm\n        | RunnableLambda(tag_with_name).bind(name=editor.name)\n    )\n    result = await gn_chain.ainvoke(state)\n    return {\"messages\": [result]}"
+        "from langchain_core.messages import AIMessage, HumanMessage, ToolMessage\n",
+        "from langchain_core.prompts import MessagesPlaceholder\n",
+        "\n",
+        "gen_qn_prompt = ChatPromptTemplate.from_messages(\n",
+        "    [\n",
+        "        (\n",
+        "            \"system\",\n",
+        "            \"\"\"You are an experienced Wikipedia writer and want to edit a specific page. \\\n",
+        "Besides your identity as a Wikipedia writer, you have a specific focus when researching the topic. \\\n",
+        "Now, you are chatting with an expert to get information. Ask good questions to get more useful information.\n",
+        "\n",
+        "When you have no more questions to ask, say \"Thank you so much for your help!\" to end the conversation.\\\n",
+        "Please only ask one question at a time and don't ask what you have asked before.\\\n",
+        "Your questions should be related to the topic you want to write.\n",
+        "Be comprehensive and curious, gaining as much unique insight from the expert as possible.\\\n",
+        "\n",
+        "Stay true to your specific perspective:\n",
+        "\n",
+        "{persona}\"\"\",\n",
+        "        ),\n",
+        "        MessagesPlaceholder(variable_name=\"messages\", optional=True),\n",
+        "    ]\n",
+        ")\n",
+        "\n",
+        "\n",
+        "def tag_with_name(ai_message: AIMessage, name: str):\n",
+        "    ai_message.name = name\n",
+        "    return ai_message\n",
+        "\n",
+        "\n",
+        "def swap_roles(state: InterviewState, name: str):\n",
+        "    converted = []\n",
+        "    for message in state[\"messages\"]:\n",
+        "        if isinstance(message, AIMessage) and message.name != name:\n",
+        "            message = HumanMessage(**message.dict(exclude={\"type\"}))\n",
+        "        converted.append(message)\n",
+        "    return {\"messages\": converted}\n",
+        "\n",
+        "\n",
+        "@as_runnable\n",
+        "async def generate_question(state: InterviewState):\n",
+        "    editor = state[\"editor\"]\n",
+        "    gn_chain = (\n",
+        "        RunnableLambda(swap_roles).bind(name=editor.name)\n",
+        "        | gen_qn_prompt.partial(persona=editor.persona)\n",
+        "        | fast_llm\n",
+        "        | RunnableLambda(tag_with_name).bind(name=editor.name)\n",
+        "    )\n",
+        "    result = await gn_chain.ainvoke(state)\n",
+        "    return {\"messages\": [result]}"
       ]
     },
     {
@@ -318,7 +577,17 @@
         }
       ],
       "source": [
-        "messages = [\n    HumanMessage(f\"So you said you were writing an article on {example_topic}?\")\n]\nquestion = await generate_question.ainvoke(\n    {\n        \"editor\": perspectives.editors[0],\n        \"messages\": messages,\n    }\n)\n\nquestion[\"messages\"][0].content"
+        "messages = [\n",
+        "    HumanMessage(f\"So you said you were writing an article on {example_topic}?\")\n",
+        "]\n",
+        "question = await generate_question.ainvoke(\n",
+        "    {\n",
+        "        \"editor\": perspectives.editors[0],\n",
+        "        \"messages\": messages,\n",
+        "    }\n",
+        ")\n",
+        "\n",
+        "question[\"messages\"][0].content"
       ]
     },
     {
@@ -336,7 +605,24 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "class Queries(BaseModel):\n    queries: List[str] = Field(\n        description=\"Comprehensive list of search engine queries to answer the user's questions.\",\n    )\n\n\ngen_queries_prompt = ChatPromptTemplate.from_messages(\n    [\n        (\n            \"system\",\n            \"You are a helpful research assistant. Query the search engine to answer the user's questions.\",\n        ),\n        MessagesPlaceholder(variable_name=\"messages\", optional=True),\n    ]\n)\ngen_queries_chain = gen_queries_prompt | ChatOpenAI(\n    model=\"gpt-3.5-turbo\"\n).with_structured_output(Queries, include_raw=True)"
+        "class Queries(BaseModel):\n",
+        "    queries: List[str] = Field(\n",
+        "        description=\"Comprehensive list of search engine queries to answer the user's questions.\",\n",
+        "    )\n",
+        "\n",
+        "\n",
+        "gen_queries_prompt = ChatPromptTemplate.from_messages(\n",
+        "    [\n",
+        "        (\n",
+        "            \"system\",\n",
+        "            \"You are a helpful research assistant. Query the search engine to answer the user's questions.\",\n",
+        "        ),\n",
+        "        MessagesPlaceholder(variable_name=\"messages\", optional=True),\n",
+        "    ]\n",
+        ")\n",
+        "gen_queries_chain = gen_queries_prompt | ChatOpenAI(\n",
+        "    model=\"gpt-3.5-turbo\"\n",
+        ").with_structured_output(Queries, include_raw=True)"
       ]
     },
     {
@@ -357,7 +643,10 @@
         }
       ],
       "source": [
-        "queries = await gen_queries_chain.ainvoke(\n    {\"messages\": [HumanMessage(content=question[\"messages\"][0].content)]}\n)\nqueries[\"parsed\"].queries"
+        "queries = await gen_queries_chain.ainvoke(\n",
+        "    {\"messages\": [HumanMessage(content=question[\"messages\"][0].content)]}\n",
+        ")\n",
+        "queries[\"parsed\"].queries"
       ]
     },
     {
@@ -366,7 +655,38 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "class AnswerWithCitations(BaseModel):\n    answer: str = Field(\n        description=\"Comprehensive answer to the user's question with citations.\",\n    )\n    cited_urls: List[str] = Field(\n        description=\"List of urls cited in the answer.\",\n    )\n\n    @property\n    def as_str(self) -> str:\n        return f\"{self.answer}\\n\\nCitations:\\n\\n\" + \"\\n\".join(\n            f\"[{i+1}]: {url}\" for i, url in enumerate(self.cited_urls)\n        )\n\n\ngen_answer_prompt = ChatPromptTemplate.from_messages(\n    [\n        (\n            \"system\",\n            \"\"\"You are an expert who can use information effectively. You are chatting with a Wikipedia writer who wants\\\n to write a Wikipedia page on the topic you know. You have gathered the related information and will now use the information to form a response.\n\nMake your response as informative as possible and make sure every sentence is supported by the gathered information.\nEach response must be backed up by a citation from a reliable source, formatted as a footnote, reproducing the URLS after your response.\"\"\",\n        ),\n        MessagesPlaceholder(variable_name=\"messages\", optional=True),\n    ]\n)\n\ngen_answer_chain = gen_answer_prompt | fast_llm.with_structured_output(\n    AnswerWithCitations, include_raw=True\n).with_config(run_name=\"GenerateAnswer\")"
+        "class AnswerWithCitations(BaseModel):\n",
+        "    answer: str = Field(\n",
+        "        description=\"Comprehensive answer to the user's question with citations.\",\n",
+        "    )\n",
+        "    cited_urls: List[str] = Field(\n",
+        "        description=\"List of urls cited in the answer.\",\n",
+        "    )\n",
+        "\n",
+        "    @property\n",
+        "    def as_str(self) -> str:\n",
+        "        return f\"{self.answer}\\n\\nCitations:\\n\\n\" + \"\\n\".join(\n",
+        "            f\"[{i+1}]: {url}\" for i, url in enumerate(self.cited_urls)\n",
+        "        )\n",
+        "\n",
+        "\n",
+        "gen_answer_prompt = ChatPromptTemplate.from_messages(\n",
+        "    [\n",
+        "        (\n",
+        "            \"system\",\n",
+        "            \"\"\"You are an expert who can use information effectively. You are chatting with a Wikipedia writer who wants\\\n",
+        " to write a Wikipedia page on the topic you know. You have gathered the related information and will now use the information to form a response.\n",
+        "\n",
+        "Make your response as informative as possible and make sure every sentence is supported by the gathered information.\n",
+        "Each response must be backed up by a citation from a reliable source, formatted as a footnote, reproducing the URLS after your response.\"\"\",\n",
+        "        ),\n",
+        "        MessagesPlaceholder(variable_name=\"messages\", optional=True),\n",
+        "    ]\n",
+        ")\n",
+        "\n",
+        "gen_answer_chain = gen_answer_prompt | fast_llm.with_structured_output(\n",
+        "    AnswerWithCitations, include_raw=True\n",
+        ").with_config(run_name=\"GenerateAnswer\")"
       ]
     },
     {
@@ -375,7 +695,29 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "from langchain_community.utilities.duckduckgo_search import DuckDuckGoSearchAPIWrapper\nfrom langchain_core.tools import tool\n\n'''\n# Tavily is typically a better search engine, but your free queries are limited\nsearch_engine = TavilySearchResults(max_results=4)\n\n@tool\nasync def search_engine(query: str):\n    \"\"\"Search engine to the internet.\"\"\"\n    results = tavily_search.invoke(query)\n    return [{\"content\": r[\"content\"], \"url\": r[\"url\"]} for r in results]\n'''\n\n# DDG\nsearch_engine = DuckDuckGoSearchAPIWrapper()\n\n\n@tool\nasync def search_engine(query: str):\n    \"\"\"Search engine to the internet.\"\"\"\n    results = DuckDuckGoSearchAPIWrapper()._ddgs_text(query)\n    return [{\"content\": r[\"body\"], \"url\": r[\"href\"]} for r in results]"
+        "from langchain_community.utilities.duckduckgo_search import DuckDuckGoSearchAPIWrapper\n",
+        "from langchain_core.tools import tool\n",
+        "\n",
+        "'''\n",
+        "# Tavily is typically a better search engine, but your free queries are limited\n",
+        "search_engine = TavilySearchResults(max_results=4)\n",
+        "\n",
+        "@tool\n",
+        "async def search_engine(query: str):\n",
+        "    \"\"\"Search engine to the internet.\"\"\"\n",
+        "    results = tavily_search.invoke(query)\n",
+        "    return [{\"content\": r[\"content\"], \"url\": r[\"url\"]} for r in results]\n",
+        "'''\n",
+        "\n",
+        "# DDG\n",
+        "search_engine = DuckDuckGoSearchAPIWrapper()\n",
+        "\n",
+        "\n",
+        "@tool\n",
+        "async def search_engine(query: str):\n",
+        "    \"\"\"Search engine to the internet.\"\"\"\n",
+        "    results = DuckDuckGoSearchAPIWrapper()._ddgs_text(query)\n",
+        "    return [{\"content\": r[\"body\"], \"url\": r[\"href\"]} for r in results]"
       ]
     },
     {
@@ -384,7 +726,43 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "import json\n\nfrom langchain_core.runnables import RunnableConfig\n\n\nasync def gen_answer(\n    state: InterviewState,\n    config: Optional[RunnableConfig] = None,\n    name: str = \"Subject_Matter_Expert\",\n    max_str_len: int = 15000,\n):\n    swapped_state = swap_roles(state, name)  # Convert all other AI messages\n    queries = await gen_queries_chain.ainvoke(swapped_state)\n    query_results = await search_engine.abatch(\n        queries[\"parsed\"].queries, config, return_exceptions=True\n    )\n    successful_results = [\n        res for res in query_results if not isinstance(res, Exception)\n    ]\n    all_query_results = {\n        res[\"url\"]: res[\"content\"] for results in successful_results for res in results\n    }\n    # We could be more precise about handling max token length if we wanted to here\n    dumped = json.dumps(all_query_results)[:max_str_len]\n    ai_message: AIMessage = queries[\"raw\"]\n    tool_call = queries[\"raw\"].additional_kwargs[\"tool_calls\"][0]\n    tool_id = tool_call[\"id\"]\n    tool_message = ToolMessage(tool_call_id=tool_id, content=dumped)\n    swapped_state[\"messages\"].extend([ai_message, tool_message])\n    # Only update the shared state with the final answer to avoid\n    # polluting the dialogue history with intermediate messages\n    generated = await gen_answer_chain.ainvoke(swapped_state)\n    cited_urls = set(generated[\"parsed\"].cited_urls)\n    # Save the retrieved information to a the shared state for future reference\n    cited_references = {k: v for k, v in all_query_results.items() if k in cited_urls}\n    formatted_message = AIMessage(name=name, content=generated[\"parsed\"].as_str)\n    return {\"messages\": [formatted_message], \"references\": cited_references}"
+        "import json\n",
+        "\n",
+        "from langchain_core.runnables import RunnableConfig\n",
+        "\n",
+        "\n",
+        "async def gen_answer(\n",
+        "    state: InterviewState,\n",
+        "    config: Optional[RunnableConfig] = None,\n",
+        "    name: str = \"Subject_Matter_Expert\",\n",
+        "    max_str_len: int = 15000,\n",
+        "):\n",
+        "    swapped_state = swap_roles(state, name)  # Convert all other AI messages\n",
+        "    queries = await gen_queries_chain.ainvoke(swapped_state)\n",
+        "    query_results = await search_engine.abatch(\n",
+        "        queries[\"parsed\"].queries, config, return_exceptions=True\n",
+        "    )\n",
+        "    successful_results = [\n",
+        "        res for res in query_results if not isinstance(res, Exception)\n",
+        "    ]\n",
+        "    all_query_results = {\n",
+        "        res[\"url\"]: res[\"content\"] for results in successful_results for res in results\n",
+        "    }\n",
+        "    # We could be more precise about handling max token length if we wanted to here\n",
+        "    dumped = json.dumps(all_query_results)[:max_str_len]\n",
+        "    ai_message: AIMessage = queries[\"raw\"]\n",
+        "    tool_call = queries[\"raw\"].tool_calls[0]\n",
+        "    tool_id = tool_call[\"id\"]\n",
+        "    tool_message = ToolMessage(tool_call_id=tool_id, content=dumped)\n",
+        "    swapped_state[\"messages\"].extend([ai_message, tool_message])\n",
+        "    # Only update the shared state with the final answer to avoid\n",
+        "    # polluting the dialogue history with intermediate messages\n",
+        "    generated = await gen_answer_chain.ainvoke(swapped_state)\n",
+        "    cited_urls = set(generated[\"parsed\"].cited_urls)\n",
+        "    # Save the retrieved information to a the shared state for future reference\n",
+        "    cited_references = {k: v for k, v in all_query_results.items() if k in cited_urls}\n",
+        "    formatted_message = AIMessage(name=name, content=generated[\"parsed\"].as_str)\n",
+        "    return {\"messages\": [formatted_message], \"references\": cited_references}"
       ]
     },
     {
@@ -404,7 +782,10 @@
         }
       ],
       "source": [
-        "example_answer = await gen_answer(\n    {\"messages\": [HumanMessage(content=question[\"messages\"][0].content)]}\n)\nexample_answer[\"messages\"][-1].content"
+        "example_answer = await gen_answer(\n",
+        "    {\"messages\": [HumanMessage(content=question[\"messages\"][0].content)]}\n",
+        ")\n",
+        "example_answer[\"messages\"][-1].content"
       ]
     },
     {
@@ -423,7 +804,31 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "max_num_turns = 5\n\n\ndef route_messages(state: InterviewState, name: str = \"Subject_Matter_Expert\"):\n    messages = state[\"messages\"]\n    num_responses = len(\n        [m for m in messages if isinstance(m, AIMessage) and m.name == name]\n    )\n    if num_responses >= max_num_turns:\n        return END\n    last_question = messages[-2]\n    if last_question.content.endswith(\"Thank you so much for your help!\"):\n        return END\n    return \"ask_question\"\n\n\nbuilder = StateGraph(InterviewState)\n\nbuilder.add_node(\"ask_question\", generate_question)\nbuilder.add_node(\"answer_question\", gen_answer)\nbuilder.add_conditional_edges(\"answer_question\", route_messages)\nbuilder.add_edge(\"ask_question\", \"answer_question\")\n\nbuilder.add_edge(START, \"ask_question\")\ninterview_graph = builder.compile().with_config(run_name=\"Conduct Interviews\")"
+        "max_num_turns = 5\n",
+        "\n",
+        "\n",
+        "def route_messages(state: InterviewState, name: str = \"Subject_Matter_Expert\"):\n",
+        "    messages = state[\"messages\"]\n",
+        "    num_responses = len(\n",
+        "        [m for m in messages if isinstance(m, AIMessage) and m.name == name]\n",
+        "    )\n",
+        "    if num_responses >= max_num_turns:\n",
+        "        return END\n",
+        "    last_question = messages[-2]\n",
+        "    if last_question.content.endswith(\"Thank you so much for your help!\"):\n",
+        "        return END\n",
+        "    return \"ask_question\"\n",
+        "\n",
+        "\n",
+        "builder = StateGraph(InterviewState)\n",
+        "\n",
+        "builder.add_node(\"ask_question\", generate_question)\n",
+        "builder.add_node(\"answer_question\", gen_answer)\n",
+        "builder.add_conditional_edges(\"answer_question\", route_messages)\n",
+        "builder.add_edge(\"ask_question\", \"answer_question\")\n",
+        "\n",
+        "builder.add_edge(START, \"ask_question\")\n",
+        "interview_graph = builder.compile().with_config(run_name=\"Conduct Interviews\")"
       ]
     },
     {
@@ -444,7 +849,11 @@
         }
       ],
       "source": [
-        "from IPython.display import Image\n\n# Feel free to comment out if you have\n# not installed pygraphviz\nImage(interview_graph.get_graph().draw_png())"
+        "from IPython.display import Image\n",
+        "\n",
+        "# Feel free to comment out if you have\n",
+        "# not installed pygraphviz\n",
+        "Image(interview_graph.get_graph().draw_png())"
       ]
     },
     {
@@ -474,7 +883,23 @@
         }
       ],
       "source": [
-        "final_step = None\n\ninitial_state = {\n    \"editor\": perspectives.editors[0],\n    \"messages\": [\n        AIMessage(\n            content=f\"So you said you were writing an article on {example_topic}?\",\n            name=\"Subject_Matter_Expert\",\n        )\n    ],\n}\nasync for step in interview_graph.astream(initial_state):\n    name = next(iter(step))\n    print(name)\n    print(\"-- \", str(step[name][\"messages\"])[:300])\n    if END in step:\n        final_step = step"
+        "final_step = None\n",
+        "\n",
+        "initial_state = {\n",
+        "    \"editor\": perspectives.editors[0],\n",
+        "    \"messages\": [\n",
+        "        AIMessage(\n",
+        "            content=f\"So you said you were writing an article on {example_topic}?\",\n",
+        "            name=\"Subject_Matter_Expert\",\n",
+        "        )\n",
+        "    ],\n",
+        "}\n",
+        "async for step in interview_graph.astream(initial_state):\n",
+        "    name = next(iter(step))\n",
+        "    print(name)\n",
+        "    print(\"-- \", str(step[name][\"messages\"])[:300])\n",
+        "    if END in step:\n",
+        "        final_step = step"
       ]
     },
     {
@@ -501,7 +926,29 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "refine_outline_prompt = ChatPromptTemplate.from_messages(\n    [\n        (\n            \"system\",\n            \"\"\"You are a Wikipedia writer. You have gathered information from experts and search engines. Now, you are refining the outline of the Wikipedia page. \\\nYou need to make sure that the outline is comprehensive and specific. \\\nTopic you are writing about: {topic} \n\nOld outline:\n\n{old_outline}\"\"\",\n        ),\n        (\n            \"user\",\n            \"Refine the outline based on your conversations with subject-matter experts:\\n\\nConversations:\\n\\n{conversations}\\n\\nWrite the refined Wikipedia outline:\",\n        ),\n    ]\n)\n\n# Using turbo preview since the context can get quite long\nrefine_outline_chain = refine_outline_prompt | long_context_llm.with_structured_output(\n    Outline\n)"
+        "refine_outline_prompt = ChatPromptTemplate.from_messages(\n",
+        "    [\n",
+        "        (\n",
+        "            \"system\",\n",
+        "            \"\"\"You are a Wikipedia writer. You have gathered information from experts and search engines. Now, you are refining the outline of the Wikipedia page. \\\n",
+        "You need to make sure that the outline is comprehensive and specific. \\\n",
+        "Topic you are writing about: {topic} \n",
+        "\n",
+        "Old outline:\n",
+        "\n",
+        "{old_outline}\"\"\",\n",
+        "        ),\n",
+        "        (\n",
+        "            \"user\",\n",
+        "            \"Refine the outline based on your conversations with subject-matter experts:\\n\\nConversations:\\n\\n{conversations}\\n\\nWrite the refined Wikipedia outline:\",\n",
+        "        ),\n",
+        "    ]\n",
+        ")\n",
+        "\n",
+        "# Using turbo preview since the context can get quite long\n",
+        "refine_outline_chain = refine_outline_prompt | long_context_llm.with_structured_output(\n",
+        "    Outline\n",
+        ")"
       ]
     },
     {
@@ -510,7 +957,15 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "refined_outline = refine_outline_chain.invoke(\n    {\n        \"topic\": example_topic,\n        \"old_outline\": initial_outline.as_str,\n        \"conversations\": \"\\n\\n\".join(\n            f\"### {m.name}\\n\\n{m.content}\" for m in final_state[\"messages\"]\n        ),\n    }\n)"
+        "refined_outline = refine_outline_chain.invoke(\n",
+        "    {\n",
+        "        \"topic\": example_topic,\n",
+        "        \"old_outline\": initial_outline.as_str,\n",
+        "        \"conversations\": \"\\n\\n\".join(\n",
+        "            f\"### {m.name}\\n\\n{m.content}\" for m in final_state[\"messages\"]\n",
+        "        ),\n",
+        "    }\n",
+        ")"
       ]
     },
     {
@@ -595,7 +1050,23 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "from langchain_community.vectorstores import SKLearnVectorStore\nfrom langchain_core.documents import Document\nfrom langchain_openai import OpenAIEmbeddings\n\nembeddings = OpenAIEmbeddings(model=\"text-embedding-3-small\")\nreference_docs = [\n    Document(page_content=v, metadata={\"source\": k})\n    for k, v in final_state[\"references\"].items()\n]\n# This really doesn't need to be a vectorstore for this size of data.\n# It could just be a numpy matrix. Or you could store documents\n# across requests if you want.\nvectorstore = SKLearnVectorStore.from_documents(\n    reference_docs,\n    embedding=embeddings,\n)\nretriever = vectorstore.as_retriever(k=10)"
+        "from langchain_community.vectorstores import SKLearnVectorStore\n",
+        "from langchain_core.documents import Document\n",
+        "from langchain_openai import OpenAIEmbeddings\n",
+        "\n",
+        "embeddings = OpenAIEmbeddings(model=\"text-embedding-3-small\")\n",
+        "reference_docs = [\n",
+        "    Document(page_content=v, metadata={\"source\": k})\n",
+        "    for k, v in final_state[\"references\"].items()\n",
+        "]\n",
+        "# This really doesn't need to be a vectorstore for this size of data.\n",
+        "# It could just be a numpy matrix. Or you could store documents\n",
+        "# across requests if you want.\n",
+        "vectorstore = SKLearnVectorStore.from_documents(\n",
+        "    reference_docs,\n",
+        "    embedding=embeddings,\n",
+        ")\n",
+        "retriever = vectorstore.as_retriever(k=10)"
       ]
     },
     {
@@ -636,7 +1107,67 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "class SubSection(BaseModel):\n    subsection_title: str = Field(..., title=\"Title of the subsection\")\n    content: str = Field(\n        ...,\n        title=\"Full content of the subsection. Include [#] citations to the cited sources where relevant.\",\n    )\n\n    @property\n    def as_str(self) -> str:\n        return f\"### {self.subsection_title}\\n\\n{self.content}\".strip()\n\n\nclass WikiSection(BaseModel):\n    section_title: str = Field(..., title=\"Title of the section\")\n    content: str = Field(..., title=\"Full content of the section\")\n    subsections: Optional[List[Subsection]] = Field(\n        default=None,\n        title=\"Titles and descriptions for each subsection of the Wikipedia page.\",\n    )\n    citations: List[str] = Field(default_factory=list)\n\n    @property\n    def as_str(self) -> str:\n        subsections = \"\\n\\n\".join(\n            subsection.as_str for subsection in self.subsections or []\n        )\n        citations = \"\\n\".join([f\" [{i}] {cit}\" for i, cit in enumerate(self.citations)])\n        return (\n            f\"## {self.section_title}\\n\\n{self.content}\\n\\n{subsections}\".strip()\n            + f\"\\n\\n{citations}\".strip()\n        )\n\n\nsection_writer_prompt = ChatPromptTemplate.from_messages(\n    [\n        (\n            \"system\",\n            \"You are an expert Wikipedia writer. Complete your assigned WikiSection from the following outline:\\n\\n\"\n            \"{outline}\\n\\nCite your sources, using the following references:\\n\\n<Documents>\\n{docs}\\n<Documents>\",\n        ),\n        (\"user\", \"Write the full WikiSection for the {section} section.\"),\n    ]\n)\n\n\nasync def retrieve(inputs: dict):\n    docs = await retriever.ainvoke(inputs[\"topic\"] + \": \" + inputs[\"section\"])\n    formatted = \"\\n\".join(\n        [\n            f'<Document href=\"{doc.metadata[\"source\"]}\"/>\\n{doc.page_content}\\n</Document>'\n            for doc in docs\n        ]\n    )\n    return {\"docs\": formatted, **inputs}\n\n\nsection_writer = (\n    retrieve\n    | section_writer_prompt\n    | long_context_llm.with_structured_output(WikiSection)\n)"
+        "class SubSection(BaseModel):\n",
+        "    subsection_title: str = Field(..., title=\"Title of the subsection\")\n",
+        "    content: str = Field(\n",
+        "        ...,\n",
+        "        title=\"Full content of the subsection. Include [#] citations to the cited sources where relevant.\",\n",
+        "    )\n",
+        "\n",
+        "    @property\n",
+        "    def as_str(self) -> str:\n",
+        "        return f\"### {self.subsection_title}\\n\\n{self.content}\".strip()\n",
+        "\n",
+        "\n",
+        "class WikiSection(BaseModel):\n",
+        "    section_title: str = Field(..., title=\"Title of the section\")\n",
+        "    content: str = Field(..., title=\"Full content of the section\")\n",
+        "    subsections: Optional[List[Subsection]] = Field(\n",
+        "        default=None,\n",
+        "        title=\"Titles and descriptions for each subsection of the Wikipedia page.\",\n",
+        "    )\n",
+        "    citations: List[str] = Field(default_factory=list)\n",
+        "\n",
+        "    @property\n",
+        "    def as_str(self) -> str:\n",
+        "        subsections = \"\\n\\n\".join(\n",
+        "            subsection.as_str for subsection in self.subsections or []\n",
+        "        )\n",
+        "        citations = \"\\n\".join([f\" [{i}] {cit}\" for i, cit in enumerate(self.citations)])\n",
+        "        return (\n",
+        "            f\"## {self.section_title}\\n\\n{self.content}\\n\\n{subsections}\".strip()\n",
+        "            + f\"\\n\\n{citations}\".strip()\n",
+        "        )\n",
+        "\n",
+        "\n",
+        "section_writer_prompt = ChatPromptTemplate.from_messages(\n",
+        "    [\n",
+        "        (\n",
+        "            \"system\",\n",
+        "            \"You are an expert Wikipedia writer. Complete your assigned WikiSection from the following outline:\\n\\n\"\n",
+        "            \"{outline}\\n\\nCite your sources, using the following references:\\n\\n<Documents>\\n{docs}\\n<Documents>\",\n",
+        "        ),\n",
+        "        (\"user\", \"Write the full WikiSection for the {section} section.\"),\n",
+        "    ]\n",
+        ")\n",
+        "\n",
+        "\n",
+        "async def retrieve(inputs: dict):\n",
+        "    docs = await retriever.ainvoke(inputs[\"topic\"] + \": \" + inputs[\"section\"])\n",
+        "    formatted = \"\\n\".join(\n",
+        "        [\n",
+        "            f'<Document href=\"{doc.metadata[\"source\"]}\"/>\\n{doc.page_content}\\n</Document>'\n",
+        "            for doc in docs\n",
+        "        ]\n",
+        "    )\n",
+        "    return {\"docs\": formatted, **inputs}\n",
+        "\n",
+        "\n",
+        "section_writer = (\n",
+        "    retrieve\n",
+        "    | section_writer_prompt\n",
+        "    | long_context_llm.with_structured_output(WikiSection)\n",
+        ")"
       ]
     },
     {
@@ -663,7 +1194,14 @@
         }
       ],
       "source": [
-        "section = await section_writer.ainvoke(\n    {\n        \"outline\": refined_outline.as_str,\n        \"section\": refined_outline.sections[1].section_title,\n        \"topic\": example_topic,\n    }\n)\nprint(section.as_str)"
+        "section = await section_writer.ainvoke(\n",
+        "    {\n",
+        "        \"outline\": refined_outline.as_str,\n",
+        "        \"section\": refined_outline.sections[1].section_title,\n",
+        "        \"topic\": example_topic,\n",
+        "    }\n",
+        ")\n",
+        "print(section.as_str)"
       ]
     },
     {
@@ -681,7 +1219,24 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "from langchain_core.output_parsers import StrOutputParser\n\nwriter_prompt = ChatPromptTemplate.from_messages(\n    [\n        (\n            \"system\",\n            \"You are an expert Wikipedia author. Write the complete wiki article on {topic} using the following section drafts:\\n\\n\"\n            \"{draft}\\n\\nStrictly follow Wikipedia format guidelines.\",\n        ),\n        (\n            \"user\",\n            'Write the complete Wiki article using markdown format. Organize citations using footnotes like \"[1]\",'\n            \" avoiding duplicates in the footer. Include URLs in the footer.\",\n        ),\n    ]\n)\n\nwriter = writer_prompt | long_context_llm | StrOutputParser()"
+        "from langchain_core.output_parsers import StrOutputParser\n",
+        "\n",
+        "writer_prompt = ChatPromptTemplate.from_messages(\n",
+        "    [\n",
+        "        (\n",
+        "            \"system\",\n",
+        "            \"You are an expert Wikipedia author. Write the complete wiki article on {topic} using the following section drafts:\\n\\n\"\n",
+        "            \"{draft}\\n\\nStrictly follow Wikipedia format guidelines.\",\n",
+        "        ),\n",
+        "        (\n",
+        "            \"user\",\n",
+        "            'Write the complete Wiki article using markdown format. Organize citations using footnotes like \"[1]\",'\n",
+        "            \" avoiding duplicates in the footer. Include URLs in the footer.\",\n",
+        "        ),\n",
+        "    ]\n",
+        ")\n",
+        "\n",
+        "writer = writer_prompt | long_context_llm | StrOutputParser()"
       ]
     },
     {
@@ -774,7 +1329,8 @@
         }
       ],
       "source": [
-        "for tok in writer.stream({\"topic\": example_topic, \"draft\": section.as_str}):\n    print(tok, end=\"\")"
+        "for tok in writer.stream({\"topic\": example_topic, \"draft\": section.as_str}):\n",
+        "    print(tok, end=\"\")"
       ]
     },
     {
@@ -801,7 +1357,14 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "class ResearchState(TypedDict):\n    topic: str\n    outline: Outline\n    editors: List[Editor]\n    interview_results: List[InterviewState]\n    # The final sections output\n    sections: List[WikiSection]\n    article: str"
+        "class ResearchState(TypedDict):\n",
+        "    topic: str\n",
+        "    outline: Outline\n",
+        "    editors: List[Editor]\n",
+        "    interview_results: List[InterviewState]\n",
+        "    # The final sections output\n",
+        "    sections: List[WikiSection]\n",
+        "    article: str"
       ]
     },
     {
@@ -810,7 +1373,109 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "import asyncio\n\n\nasync def initialize_research(state: ResearchState):\n    topic = state[\"topic\"]\n    coros = (\n        generate_outline_direct.ainvoke({\"topic\": topic}),\n        survey_subjects.ainvoke(topic),\n    )\n    results = await asyncio.gather(*coros)\n    return {\n        **state,\n        \"outline\": results[0],\n        \"editors\": results[1].editors,\n    }\n\n\nasync def conduct_interviews(state: ResearchState):\n    topic = state[\"topic\"]\n    initial_states = [\n        {\n            \"editor\": editor,\n            \"messages\": [\n                AIMessage(\n                    content=f\"So you said you were writing an article on {topic}?\",\n                    name=\"Subject_Matter_Expert\",\n                )\n            ],\n        }\n        for editor in state[\"editors\"]\n    ]\n    # We call in to the sub-graph here to parallelize the interviews\n    interview_results = await interview_graph.abatch(initial_states)\n\n    return {\n        **state,\n        \"interview_results\": interview_results,\n    }\n\n\ndef format_conversation(interview_state):\n    messages = interview_state[\"messages\"]\n    convo = \"\\n\".join(f\"{m.name}: {m.content}\" for m in messages)\n    return f'Conversation with {interview_state[\"editor\"].name}\\n\\n' + convo\n\n\nasync def refine_outline(state: ResearchState):\n    convos = \"\\n\\n\".join(\n        [\n            format_conversation(interview_state)\n            for interview_state in state[\"interview_results\"]\n        ]\n    )\n\n    updated_outline = await refine_outline_chain.ainvoke(\n        {\n            \"topic\": state[\"topic\"],\n            \"old_outline\": state[\"outline\"].as_str,\n            \"conversations\": convos,\n        }\n    )\n    return {**state, \"outline\": updated_outline}\n\n\nasync def index_references(state: ResearchState):\n    all_docs = []\n    for interview_state in state[\"interview_results\"]:\n        reference_docs = [\n            Document(page_content=v, metadata={\"source\": k})\n            for k, v in interview_state[\"references\"].items()\n        ]\n        all_docs.extend(reference_docs)\n    await vectorstore.aadd_documents(all_docs)\n    return state\n\n\nasync def write_sections(state: ResearchState):\n    outline = state[\"outline\"]\n    sections = await section_writer.abatch(\n        [\n            {\n                \"outline\": refined_outline.as_str,\n                \"section\": section.section_title,\n                \"topic\": state[\"topic\"],\n            }\n            for section in outline.sections\n        ]\n    )\n    return {\n        **state,\n        \"sections\": sections,\n    }\n\n\nasync def write_article(state: ResearchState):\n    topic = state[\"topic\"]\n    sections = state[\"sections\"]\n    draft = \"\\n\\n\".join([section.as_str for section in sections])\n    article = await writer.ainvoke({\"topic\": topic, \"draft\": draft})\n    return {\n        **state,\n        \"article\": article,\n    }"
+        "import asyncio\n",
+        "\n",
+        "\n",
+        "async def initialize_research(state: ResearchState):\n",
+        "    topic = state[\"topic\"]\n",
+        "    coros = (\n",
+        "        generate_outline_direct.ainvoke({\"topic\": topic}),\n",
+        "        survey_subjects.ainvoke(topic),\n",
+        "    )\n",
+        "    results = await asyncio.gather(*coros)\n",
+        "    return {\n",
+        "        **state,\n",
+        "        \"outline\": results[0],\n",
+        "        \"editors\": results[1].editors,\n",
+        "    }\n",
+        "\n",
+        "\n",
+        "async def conduct_interviews(state: ResearchState):\n",
+        "    topic = state[\"topic\"]\n",
+        "    initial_states = [\n",
+        "        {\n",
+        "            \"editor\": editor,\n",
+        "            \"messages\": [\n",
+        "                AIMessage(\n",
+        "                    content=f\"So you said you were writing an article on {topic}?\",\n",
+        "                    name=\"Subject_Matter_Expert\",\n",
+        "                )\n",
+        "            ],\n",
+        "        }\n",
+        "        for editor in state[\"editors\"]\n",
+        "    ]\n",
+        "    # We call in to the sub-graph here to parallelize the interviews\n",
+        "    interview_results = await interview_graph.abatch(initial_states)\n",
+        "\n",
+        "    return {\n",
+        "        **state,\n",
+        "        \"interview_results\": interview_results,\n",
+        "    }\n",
+        "\n",
+        "\n",
+        "def format_conversation(interview_state):\n",
+        "    messages = interview_state[\"messages\"]\n",
+        "    convo = \"\\n\".join(f\"{m.name}: {m.content}\" for m in messages)\n",
+        "    return f'Conversation with {interview_state[\"editor\"].name}\\n\\n' + convo\n",
+        "\n",
+        "\n",
+        "async def refine_outline(state: ResearchState):\n",
+        "    convos = \"\\n\\n\".join(\n",
+        "        [\n",
+        "            format_conversation(interview_state)\n",
+        "            for interview_state in state[\"interview_results\"]\n",
+        "        ]\n",
+        "    )\n",
+        "\n",
+        "    updated_outline = await refine_outline_chain.ainvoke(\n",
+        "        {\n",
+        "            \"topic\": state[\"topic\"],\n",
+        "            \"old_outline\": state[\"outline\"].as_str,\n",
+        "            \"conversations\": convos,\n",
+        "        }\n",
+        "    )\n",
+        "    return {**state, \"outline\": updated_outline}\n",
+        "\n",
+        "\n",
+        "async def index_references(state: ResearchState):\n",
+        "    all_docs = []\n",
+        "    for interview_state in state[\"interview_results\"]:\n",
+        "        reference_docs = [\n",
+        "            Document(page_content=v, metadata={\"source\": k})\n",
+        "            for k, v in interview_state[\"references\"].items()\n",
+        "        ]\n",
+        "        all_docs.extend(reference_docs)\n",
+        "    await vectorstore.aadd_documents(all_docs)\n",
+        "    return state\n",
+        "\n",
+        "\n",
+        "async def write_sections(state: ResearchState):\n",
+        "    outline = state[\"outline\"]\n",
+        "    sections = await section_writer.abatch(\n",
+        "        [\n",
+        "            {\n",
+        "                \"outline\": refined_outline.as_str,\n",
+        "                \"section\": section.section_title,\n",
+        "                \"topic\": state[\"topic\"],\n",
+        "            }\n",
+        "            for section in outline.sections\n",
+        "        ]\n",
+        "    )\n",
+        "    return {\n",
+        "        **state,\n",
+        "        \"sections\": sections,\n",
+        "    }\n",
+        "\n",
+        "\n",
+        "async def write_article(state: ResearchState):\n",
+        "    topic = state[\"topic\"]\n",
+        "    sections = state[\"sections\"]\n",
+        "    draft = \"\\n\\n\".join([section.as_str for section in sections])\n",
+        "    article = await writer.ainvoke({\"topic\": topic, \"draft\": draft})\n",
+        "    return {\n",
+        "        **state,\n",
+        "        \"article\": article,\n",
+        "    }"
       ]
     },
     {
@@ -826,7 +1491,27 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "from langgraph.checkpoint.memory import MemorySaver\n\nbuilder_of_storm = StateGraph(ResearchState)\n\nnodes = [\n    (\"init_research\", initialize_research),\n    (\"conduct_interviews\", conduct_interviews),\n    (\"refine_outline\", refine_outline),\n    (\"index_references\", index_references),\n    (\"write_sections\", write_sections),\n    (\"write_article\", write_article),\n]\nfor i in range(len(nodes)):\n    name, node = nodes[i]\n    builder_of_storm.add_node(name, node)\n    if i > 0:\n        builder_of_storm.add_edge(nodes[i - 1][0], name)\n\nbuilder_of_storm.add_edge(START, nodes[0][0])\nbuilder_of_storm.add_edge(nodes[-1][0], END)\nstorm = builder_of_storm.compile(checkpointer=MemorySaver())"
+        "from langgraph.checkpoint.memory import MemorySaver\n",
+        "\n",
+        "builder_of_storm = StateGraph(ResearchState)\n",
+        "\n",
+        "nodes = [\n",
+        "    (\"init_research\", initialize_research),\n",
+        "    (\"conduct_interviews\", conduct_interviews),\n",
+        "    (\"refine_outline\", refine_outline),\n",
+        "    (\"index_references\", index_references),\n",
+        "    (\"write_sections\", write_sections),\n",
+        "    (\"write_article\", write_article),\n",
+        "]\n",
+        "for i in range(len(nodes)):\n",
+        "    name, node = nodes[i]\n",
+        "    builder_of_storm.add_node(name, node)\n",
+        "    if i > 0:\n",
+        "        builder_of_storm.add_edge(nodes[i - 1][0], name)\n",
+        "\n",
+        "builder_of_storm.add_edge(START, nodes[0][0])\n",
+        "builder_of_storm.add_edge(nodes[-1][0], END)\n",
+        "storm = builder_of_storm.compile(checkpointer=MemorySaver())"
       ]
     },
     {
@@ -877,7 +1562,16 @@
         }
       ],
       "source": [
-        "config = {\"configurable\": {\"thread_id\": \"my-thread\"}}\nasync for step in storm.astream(\n    {\n        \"topic\": \"Groq, NVIDIA, Llamma.cpp and the future of LLM Inference\",\n    },\n    config,\n):\n    name = next(iter(step))\n    print(name)\n    print(\"-- \", str(step[name])[:300])"
+        "config = {\"configurable\": {\"thread_id\": \"my-thread\"}}\n",
+        "async for step in storm.astream(\n",
+        "    {\n",
+        "        \"topic\": \"Groq, NVIDIA, Llamma.cpp and the future of LLM Inference\",\n",
+        "    },\n",
+        "    config,\n",
+        "):\n",
+        "    name = next(iter(step))\n",
+        "    print(name)\n",
+        "    print(\"-- \", str(step[name])[:300])"
       ]
     },
     {
@@ -886,7 +1580,8 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "checkpoint = storm.get_state(config)\narticle = checkpoint.values[\"article\"]"
+        "checkpoint = storm.get_state(config)\n",
+        "article = checkpoint.values[\"article\"]"
       ]
     },
     {
@@ -967,7 +1662,10 @@
         }
       ],
       "source": [
-        "from IPython.display import Markdown\n\n# We will down-header the sections to create less confusion in this notebook\nMarkdown(article.replace(\"\\n#\", \"\\n##\"))"
+        "from IPython.display import Markdown\n",
+        "\n",
+        "# We will down-header the sections to create less confusion in this notebook\n",
+        "Markdown(article.replace(\"\\n#\", \"\\n##\"))"
       ]
     },
     {
@@ -975,9 +1673,7 @@
       "execution_count": null,
       "metadata": {},
       "outputs": [],
-      "source": [
-        ""
-      ]
+      "source": []
     }
   ],
   "metadata": {

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/Makefile

@@ -38,11 +38,11 @@ lint_tests: PYTHON_FILES=tests
 lint_tests: MYPY_CACHE=.mypy_cache_test
 
 lint lint_diff lint_package lint_tests:
-	poetry run ruff .
+	poetry run ruff check .
 	[ "$(PYTHON_FILES)" = "" ] || poetry run ruff format $(PYTHON_FILES) --diff
-	[ "$(PYTHON_FILES)" = "" ] || poetry run ruff --select I $(PYTHON_FILES)
+	[ "$(PYTHON_FILES)" = "" ] || poetry run ruff check --select I $(PYTHON_FILES)
 	[ "$(PYTHON_FILES)" = "" ] || mkdir -p $(MYPY_CACHE) || poetry run mypy $(PYTHON_FILES) --cache-dir $(MYPY_CACHE)
 
 format format_diff:
 	poetry run ruff format $(PYTHON_FILES)
-	poetry run ruff --select I --fix $(PYTHON_FILES)
+	poetry run ruff check --select I --fix $(PYTHON_FILES)

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]
@@ -44,7 +48,6 @@ with PostgresSaver.from_conn_string(DB_URI) as checkpointer:
             }
         },
         "pending_sends": [],
-        "current_tasks": {}
     }
 
     # store checkpoint
@@ -87,7 +90,6 @@ async with AsyncPostgresSaver.from_conn_string(DB_URI) as checkpointer:
             }
         },
         "pending_sends": [],
-        "current_tasks": {}
     }
 
     # store checkpoint

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) 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,
@@ -127,30 +143,34 @@ class PostgresSaver(BasePostgresSaver):
         if limit:
             query += f" LIMIT {limit}"
         # if we change this to use .stream() we need to make sure to close the cursor
-        for value in self.conn.execute(query, args, binary=True):
-            yield CheckpointTuple(
-                {
-                    "configurable": {
-                        "thread_id": value["thread_id"],
-                        "checkpoint_ns": value["checkpoint_ns"],
-                        "checkpoint_id": value["checkpoint_id"],
+        with self._cursor() as cur:
+            cur.execute(query, args, binary=True)
+            for value in cur:
+                yield CheckpointTuple(
+                    {
+                        "configurable": {
+                            "thread_id": value["thread_id"],
+                            "checkpoint_ns": value["checkpoint_ns"],
+                            "checkpoint_id": value["checkpoint_id"],
+                        }
+                    },
+                    self._load_checkpoint(
+                        value["checkpoint"],
+                        value["channel_values"],
+                        value["pending_sends"],
+                    ),
+                    self._load_metadata(value["metadata"]),
+                    {
+                        "configurable": {
+                            "thread_id": value["thread_id"],
+                            "checkpoint_ns": value["checkpoint_ns"],
+                            "checkpoint_id": value["parent_checkpoint_id"],
+                        }
                     }
-                },
-                {
-                    **self._load_checkpoint(value["checkpoint"]),
-                    "channel_values": self._load_blobs(value["channel_values"]),
-                },
-                self._load_metadata(value["metadata"]),
-                {
-                    "configurable": {
-                        "thread_id": value["thread_id"],
-                        "checkpoint_ns": value["checkpoint_ns"],
-                        "checkpoint_id": value["parent_checkpoint_id"],
-                    }
-                }
-                if value["parent_checkpoint_id"]
-                else None,
-            )
+                    if value["parent_checkpoint_id"]
+                    else None,
+                    self._load_writes(value["pending_writes"]),
+                )
 
     def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
         """Get a checkpoint tuple from the database.
@@ -198,7 +218,7 @@ class PostgresSaver(BasePostgresSaver):
             where = "WHERE thread_id = %s AND checkpoint_ns = %s ORDER BY checkpoint_id DESC LIMIT 1"
 
         with self._cursor() as cur:
-            cur = self.conn.execute(
+            cur.execute(
                 self.SELECT_SQL + where,
                 args,
                 binary=True,
@@ -213,10 +233,11 @@ class PostgresSaver(BasePostgresSaver):
                             "checkpoint_id": value["checkpoint_id"],
                         }
                     },
-                    {
-                        **self._load_checkpoint(value["checkpoint"]),
-                        "channel_values": self._load_blobs(value["channel_values"]),
-                    },
+                    self._load_checkpoint(
+                        value["checkpoint"],
+                        value["channel_values"],
+                        value["pending_sends"],
+                    ),
                     self._load_metadata(value["metadata"]),
                     {
                         "configurable": {
@@ -317,16 +338,6 @@ class PostgresSaver(BasePostgresSaver):
             task_id (str): Identifier for the task creating the writes.
         """
         with self._cursor(pipeline=True) as cur:
-            cur.execute(
-                self.DELETE_WRITES_SQL,
-                (
-                    config["configurable"]["thread_id"],
-                    config["configurable"]["checkpoint_ns"],
-                    config["configurable"]["checkpoint_id"],
-                    task_id,
-                    len(writes),
-                ),
-            )
             cur.executemany(
                 self.UPSERT_CHECKPOINT_WRITES_SQL,
                 self._dump_writes(
@@ -340,21 +351,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) 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) as cur:
-                yield cur
-        else:
-            with self.lock, self.conn.cursor(binary=True) 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()
@@ -45,7 +64,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
             pipeline (bool): whether to use AsyncPipeline
 
         Returns:
-            PostgresSaver: A new PostgresSaver instance.
+            AsyncPostgresSaver: A new AsyncPostgresSaver instance.
         """
         async with await AsyncConnection.connect(
             conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row
@@ -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) 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,
@@ -110,32 +126,35 @@ class AsyncPostgresSaver(BasePostgresSaver):
         if limit:
             query += f" LIMIT {limit}"
         # if we change this to use .stream() we need to make sure to close the cursor
-        async for value in await self.conn.execute(query, args, binary=True):
-            yield CheckpointTuple(
-                {
-                    "configurable": {
-                        "thread_id": value["thread_id"],
-                        "checkpoint_ns": value["checkpoint_ns"],
-                        "checkpoint_id": value["checkpoint_id"],
-                    }
-                },
-                {
-                    **self._load_checkpoint(value["checkpoint"]),
-                    "channel_values": await asyncio.to_thread(
-                        self._load_blobs, value["channel_values"]
+        async with self._cursor() as cur:
+            await cur.execute(query, args, binary=True)
+            async for value in cur:
+                yield CheckpointTuple(
+                    {
+                        "configurable": {
+                            "thread_id": value["thread_id"],
+                            "checkpoint_ns": value["checkpoint_ns"],
+                            "checkpoint_id": value["checkpoint_id"],
+                        }
+                    },
+                    await asyncio.to_thread(
+                        self._load_checkpoint,
+                        value["checkpoint"],
+                        value["channel_values"],
+                        value["pending_sends"],
                     ),
-                },
-                self._load_metadata(value["metadata"]),
-                {
-                    "configurable": {
-                        "thread_id": value["thread_id"],
-                        "checkpoint_ns": value["checkpoint_ns"],
-                        "checkpoint_id": value["parent_checkpoint_id"],
+                    self._load_metadata(value["metadata"]),
+                    {
+                        "configurable": {
+                            "thread_id": value["thread_id"],
+                            "checkpoint_ns": value["checkpoint_ns"],
+                            "checkpoint_id": value["parent_checkpoint_id"],
+                        }
                     }
-                }
-                if value["parent_checkpoint_id"]
-                else None,
-            )
+                    if value["parent_checkpoint_id"]
+                    else None,
+                    await asyncio.to_thread(self._load_writes, value["pending_writes"]),
+                )
 
     async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
         """Get a checkpoint tuple from the database asynchronously.
@@ -162,7 +181,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
             where = "WHERE thread_id = %s AND checkpoint_ns = %s ORDER BY checkpoint_id DESC LIMIT 1"
 
         async with self._cursor() as cur:
-            cur = await self.conn.execute(
+            await cur.execute(
                 self.SELECT_SQL + where,
                 args,
                 binary=True,
@@ -177,12 +196,12 @@ class AsyncPostgresSaver(BasePostgresSaver):
                             "checkpoint_id": value["checkpoint_id"],
                         }
                     },
-                    {
-                        **self._load_checkpoint(value["checkpoint"]),
-                        "channel_values": await asyncio.to_thread(
-                            self._load_blobs, value["channel_values"]
-                        ),
-                    },
+                    await asyncio.to_thread(
+                        self._load_checkpoint,
+                        value["checkpoint"],
+                        value["channel_values"],
+                        value["pending_sends"],
+                    ),
                     self._load_metadata(value["metadata"]),
                     {
                         "configurable": {
@@ -273,16 +292,6 @@ class AsyncPostgresSaver(BasePostgresSaver):
             task_id (str): Identifier for the task creating the writes.
         """
         async with self._cursor(pipeline=True) as cur:
-            await cur.execute(
-                self.DELETE_WRITES_SQL,
-                (
-                    config["configurable"]["thread_id"],
-                    config["configurable"]["checkpoint_ns"],
-                    config["configurable"]["checkpoint_id"],
-                    task_id,
-                    len(writes),
-                ),
-            )
             await cur.executemany(
                 self.UPSERT_CHECKPOINT_WRITES_SQL,
                 await asyncio.to_thread(
@@ -297,23 +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) 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
-            ) as cur:
-                yield cur
-        else:
-            async with self.lock, self.conn.cursor(binary=True) as cur:
-                yield cur

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

@@ -1,4 +1,3 @@
-from base64 import b64decode, b64encode
 from hashlib import md5
 from typing import Any, List, Optional, Tuple
 
@@ -6,13 +5,14 @@ from langchain_core.runnables import RunnableConfig
 from psycopg.types.json import Jsonb
 
 from langgraph.checkpoint.base import (
+    WRITES_IDX_MAP,
     BaseCheckpointSaver,
     Checkpoint,
     EmptyChannelError,
     get_checkpoint_id,
 )
 from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
-from langgraph.checkpoint.serde.types import ChannelProtocol
+from langgraph.checkpoint.serde.types import TASKS, ChannelProtocol
 
 MetadataInput = Optional[dict[str, Any]]
 
@@ -57,7 +57,7 @@ MIGRATIONS = [
     "ALTER TABLE checkpoint_blobs ALTER COLUMN blob DROP not null;",
 ]
 
-SELECT_SQL = """
+SELECT_SQL = f"""
 select
     thread_id,
     checkpoint,
@@ -81,7 +81,15 @@ select
         where cw.thread_id = checkpoints.thread_id
             and cw.checkpoint_ns = checkpoints.checkpoint_ns
             and cw.checkpoint_id = checkpoints.checkpoint_id
-    ) as pending_writes
+    ) as pending_writes,
+    (
+        select array_agg(array[cw.type::bytea, cw.blob])
+        from checkpoint_writes cw
+        where cw.thread_id = checkpoints.thread_id
+            and cw.checkpoint_ns = checkpoints.checkpoint_ns
+            and cw.checkpoint_id = checkpoints.parent_checkpoint_id
+            and cw.channel = '{TASKS}'
+    ) as pending_sends
 from checkpoints """
 
 UPSERT_CHECKPOINT_BLOBS_SQL = """
@@ -105,15 +113,6 @@ UPSERT_CHECKPOINT_WRITES_SQL = """
     ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) DO NOTHING
 """
 
-DELETE_WRITES_SQL = """
-    DELETE FROM checkpoint_writes
-    WHERE thread_id = %s
-    AND checkpoint_ns = %s
-    AND checkpoint_id = %s
-    AND task_id = %s
-    AND idx >= %s
-"""
-
 
 class BasePostgresSaver(BaseCheckpointSaver):
     SELECT_SQL = SELECT_SQL
@@ -121,29 +120,26 @@ class BasePostgresSaver(BaseCheckpointSaver):
     UPSERT_CHECKPOINT_BLOBS_SQL = UPSERT_CHECKPOINT_BLOBS_SQL
     UPSERT_CHECKPOINTS_SQL = UPSERT_CHECKPOINTS_SQL
     UPSERT_CHECKPOINT_WRITES_SQL = UPSERT_CHECKPOINT_WRITES_SQL
-    DELETE_WRITES_SQL = DELETE_WRITES_SQL
 
     jsonplus_serde = JsonPlusSerializer()
 
-    def _load_checkpoint(self, checkpoint: dict[str, Any]) -> Checkpoint:
-        if len(checkpoint["pending_sends"]) == 2 and all(
-            isinstance(a, str) for a in checkpoint["pending_sends"]
-        ):
-            type, bs = checkpoint["pending_sends"]
-            return {
-                **checkpoint,
-                "pending_sends": self.serde.loads_typed((type, b64decode(bs))),
-            }
-
-        return checkpoint
-
-    def _dump_checkpoint(self, checkpoint: Checkpoint) -> dict[str, Any]:
-        type, bs = self.serde.dumps_typed(checkpoint["pending_sends"])
+    def _load_checkpoint(
+        self,
+        checkpoint: dict[str, Any],
+        channel_values: list[tuple[bytes, bytes, bytes]],
+        pending_sends: list[tuple[bytes, bytes]],
+    ) -> Checkpoint:
         return {
             **checkpoint,
-            "pending_sends": (type, b64encode(bs).decode()),
+            "pending_sends": [
+                self.serde.loads_typed((c.decode(), b)) for c, b in pending_sends or []
+            ],
+            "channel_values": self._load_blobs(channel_values),
         }
 
+    def _dump_checkpoint(self, checkpoint: Checkpoint) -> dict[str, Any]:
+        return {**checkpoint, "pending_sends": []}
+
     def _load_blobs(
         self, blob_values: list[tuple[bytes, bytes, bytes]]
     ) -> dict[str, Any]:
@@ -210,7 +206,7 @@ class BasePostgresSaver(BaseCheckpointSaver):
                 checkpoint_ns,
                 checkpoint_id,
                 task_id,
-                idx,
+                WRITES_IDX_MAP.get(channel, idx),
                 channel,
                 *self.serde.dumps_typed(value),
             )
@@ -264,9 +260,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

@@ -1,4 +1,4 @@
-# This file is automatically @generated by Poetry 1.8.3 and should not be changed by hand.
+# This file is automatically @generated by Poetry 1.8.2 and should not be changed by hand.
 
 [[package]]
 name = "annotated-types"
@@ -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"
@@ -394,6 +394,8 @@ files = [
     {file = "orjson-3.10.6-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:960db0e31c4e52fa0fc3ecbaea5b2d3b58f379e32a95ae6b0ebeaa25b93dfd34"},
     {file = "orjson-3.10.6-cp312-none-win32.whl", hash = "sha256:a6ea7afb5b30b2317e0bee03c8d34c8181bc5a36f2afd4d0952f378972c4efd5"},
     {file = "orjson-3.10.6-cp312-none-win_amd64.whl", hash = "sha256:874ce88264b7e655dde4aeaacdc8fd772a7962faadfb41abe63e2a4861abc3dc"},
+    {file = "orjson-3.10.6-cp313-none-win32.whl", hash = "sha256:efdf2c5cde290ae6b83095f03119bdc00303d7a03b42b16c54517baa3c4ca3d0"},
+    {file = "orjson-3.10.6-cp313-none-win_amd64.whl", hash = "sha256:8e190fe7888e2e4392f52cafb9626113ba135ef53aacc65cd13109eb9746c43e"},
     {file = "orjson-3.10.6-cp38-cp38-macosx_10_15_x86_64.macosx_11_0_arm64.macosx_10_15_universal2.whl", hash = "sha256:66680eae4c4e7fc193d91cfc1353ad6d01b4801ae9b5314f17e11ba55e934183"},
     {file = "orjson-3.10.6-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:caff75b425db5ef8e8f23af93c80f072f97b4fb3afd4af44482905c9f588da28"},
     {file = "orjson-3.10.6-cp38-cp38-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:3722fddb821b6036fd2a3c814f6bd9b57a89dc6337b9924ecd614ebce3271394"},
@@ -822,28 +824,29 @@ use-chardet-on-py3 = ["chardet (>=3.0.2,<6)"]
 
 [[package]]
 name = "ruff"
-version = "0.1.15"
+version = "0.6.2"
 description = "An extremely fast Python linter and code formatter, written in Rust."
 optional = false
 python-versions = ">=3.7"
 files = [
-    {file = "ruff-0.1.15-py3-none-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl", hash = "sha256:5fe8d54df166ecc24106db7dd6a68d44852d14eb0729ea4672bb4d96c320b7df"},
-    {file = "ruff-0.1.15-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:6f0bfbb53c4b4de117ac4d6ddfd33aa5fc31beeaa21d23c45c6dd249faf9126f"},
-    {file = "ruff-0.1.15-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e0d432aec35bfc0d800d4f70eba26e23a352386be3a6cf157083d18f6f5881c8"},
-    {file = "ruff-0.1.15-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:9405fa9ac0e97f35aaddf185a1be194a589424b8713e3b97b762336ec79ff807"},
-    {file = "ruff-0.1.15-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:c66ec24fe36841636e814b8f90f572a8c0cb0e54d8b5c2d0e300d28a0d7bffec"},
-    {file = "ruff-0.1.15-py3-none-manylinux_2_17_ppc64.manylinux2014_ppc64.whl", hash = "sha256:6f8ad828f01e8dd32cc58bc28375150171d198491fc901f6f98d2a39ba8e3ff5"},
-    {file = "ruff-0.1.15-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:86811954eec63e9ea162af0ffa9f8d09088bab51b7438e8b6488b9401863c25e"},
-    {file = "ruff-0.1.15-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:fd4025ac5e87d9b80e1f300207eb2fd099ff8200fa2320d7dc066a3f4622dc6b"},
-    {file = "ruff-0.1.15-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b17b93c02cdb6aeb696effecea1095ac93f3884a49a554a9afa76bb125c114c1"},
-    {file = "ruff-0.1.15-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:ddb87643be40f034e97e97f5bc2ef7ce39de20e34608f3f829db727a93fb82c5"},
-    {file = "ruff-0.1.15-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:abf4822129ed3a5ce54383d5f0e964e7fef74a41e48eb1dfad404151efc130a2"},
-    {file = "ruff-0.1.15-py3-none-musllinux_1_2_i686.whl", hash = "sha256:6c629cf64bacfd136c07c78ac10a54578ec9d1bd2a9d395efbee0935868bf852"},
-    {file = "ruff-0.1.15-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:1bab866aafb53da39c2cadfb8e1c4550ac5340bb40300083eb8967ba25481447"},
-    {file = "ruff-0.1.15-py3-none-win32.whl", hash = "sha256:2417e1cb6e2068389b07e6fa74c306b2810fe3ee3476d5b8a96616633f40d14f"},
-    {file = "ruff-0.1.15-py3-none-win_amd64.whl", hash = "sha256:3837ac73d869efc4182d9036b1405ef4c73d9b1f88da2413875e34e0d6919587"},
-    {file = "ruff-0.1.15-py3-none-win_arm64.whl", hash = "sha256:9a933dfb1c14ec7a33cceb1e49ec4a16b51ce3c20fd42663198746efc0427360"},
-    {file = "ruff-0.1.15.tar.gz", hash = "sha256:f6dfa8c1b21c913c326919056c390966648b680966febcb796cc9d1aaab8564e"},
+    {file = "ruff-0.6.2-py3-none-linux_armv6l.whl", hash = "sha256:5c8cbc6252deb3ea840ad6a20b0f8583caab0c5ef4f9cca21adc5a92b8f79f3c"},
+    {file = "ruff-0.6.2-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:17002fe241e76544448a8e1e6118abecbe8cd10cf68fde635dad480dba594570"},
+    {file = "ruff-0.6.2-py3-none-macosx_11_0_arm64.whl", hash = "sha256:3dbeac76ed13456f8158b8f4fe087bf87882e645c8e8b606dd17b0b66c2c1158"},
+    {file = "ruff-0.6.2-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:094600ee88cda325988d3f54e3588c46de5c18dae09d683ace278b11f9d4d534"},
+    {file = "ruff-0.6.2-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:316d418fe258c036ba05fbf7dfc1f7d3d4096db63431546163b472285668132b"},
+    {file = "ruff-0.6.2-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:d72b8b3abf8a2d51b7b9944a41307d2f442558ccb3859bbd87e6ae9be1694a5d"},
+    {file = "ruff-0.6.2-py3-none-manylinux_2_17_ppc64.manylinux2014_ppc64.whl", hash = "sha256:2aed7e243be68487aa8982e91c6e260982d00da3f38955873aecd5a9204b1d66"},
+    {file = "ruff-0.6.2-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:d371f7fc9cec83497fe7cf5eaf5b76e22a8efce463de5f775a1826197feb9df8"},
+    {file = "ruff-0.6.2-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:a8f310d63af08f583363dfb844ba8f9417b558199c58a5999215082036d795a1"},
+    {file = "ruff-0.6.2-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7db6880c53c56addb8638fe444818183385ec85eeada1d48fc5abe045301b2f1"},
+    {file = "ruff-0.6.2-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:1175d39faadd9a50718f478d23bfc1d4da5743f1ab56af81a2b6caf0a2394f23"},
+    {file = "ruff-0.6.2-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:5b939f9c86d51635fe486585389f54582f0d65b8238e08c327c1534844b3bb9a"},
+    {file = "ruff-0.6.2-py3-none-musllinux_1_2_i686.whl", hash = "sha256:d0d62ca91219f906caf9b187dea50d17353f15ec9bb15aae4a606cd697b49b4c"},
+    {file = "ruff-0.6.2-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:7438a7288f9d67ed3c8ce4d059e67f7ed65e9fe3aa2ab6f5b4b3610e57e3cb56"},
+    {file = "ruff-0.6.2-py3-none-win32.whl", hash = "sha256:279d5f7d86696df5f9549b56b9b6a7f6c72961b619022b5b7999b15db392a4da"},
+    {file = "ruff-0.6.2-py3-none-win_amd64.whl", hash = "sha256:d9f3469c7dd43cd22eb1c3fc16926fb8258d50cb1b216658a07be95dd117b0f2"},
+    {file = "ruff-0.6.2-py3-none-win_arm64.whl", hash = "sha256:f28fcd2cd0e02bdf739297516d5643a945cc7caf09bd9bcb4d932540a5ea4fa9"},
+    {file = "ruff-0.6.2.tar.gz", hash = "sha256:239ee6beb9e91feb8e0ec384204a763f36cb53fb895a1a364618c6abb076b3be"},
 ]
 
 [[package]]
@@ -969,4 +972,4 @@ watchmedo = ["PyYAML (>=3.10)"]
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.9.0,<4.0"
-content-hash = "422b6d716b86db072ea3a612287ad20ff5700c18f22d9e9d59cc4e198514519d"
+content-hash = "b139531e8c6f4e24cea4bdfc29d111d1ea000a6a8a604ba81be4bff0977a2466"

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,10 +12,11 @@ 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"
+ruff = "^0.6.2"
 codespell = "^2.2.0"
 pytest = "^7.2.1"
 anyio = "^4.4.0"
@@ -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/Makefile

@@ -24,11 +24,11 @@ lint_tests: PYTHON_FILES=tests
 lint_tests: MYPY_CACHE=.mypy_cache_test
 
 lint lint_diff lint_package lint_tests:
-	poetry run ruff .
+	poetry run ruff check .
 	[ "$(PYTHON_FILES)" = "" ] || poetry run ruff format $(PYTHON_FILES) --diff
-	[ "$(PYTHON_FILES)" = "" ] || poetry run ruff --select I $(PYTHON_FILES)
+	[ "$(PYTHON_FILES)" = "" ] || poetry run ruff check --select I $(PYTHON_FILES)
 	[ "$(PYTHON_FILES)" = "" ] || mkdir -p $(MYPY_CACHE) || poetry run mypy $(PYTHON_FILES) --cache-dir $(MYPY_CACHE)
 
 format format_diff:
 	poetry run ruff format $(PYTHON_FILES)
-	poetry run ruff --select I --fix $(PYTHON_FILES)
+	poetry run ruff check --select I --fix $(PYTHON_FILES)

libs/checkpoint-sqlite/README.md

@@ -35,7 +35,6 @@ with SqliteSaver.from_conn_string(":memory:") as checkpointer:
             }
         },
         "pending_sends": [],
-        "current_tasks": {}
     }
 
     # store checkpoint
@@ -78,7 +77,6 @@ async with AsyncSqliteSaver.from_conn_string(":memory:") as checkpointer:
             }
         },
         "pending_sends": [],
-        "current_tasks": {}
     }
 
     # store checkpoint
@@ -89,4 +87,4 @@ async with AsyncSqliteSaver.from_conn_string(":memory:") as checkpointer:
 
     # list checkpoints
     [c async for c in checkpointer.alist(read_config)]
-```
+```

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

@@ -1,12 +1,13 @@
 import sqlite3
 import threading
-from contextlib import contextmanager
+from contextlib import closing, contextmanager
 from hashlib import md5
 from typing import Any, AsyncIterator, Dict, Iterator, Optional, Sequence, Tuple
 
 from langchain_core.runnables import RunnableConfig
 
 from langgraph.checkpoint.base import (
+    WRITES_IDX_MAP,
     BaseCheckpointSaver,
     ChannelVersions,
     Checkpoint,
@@ -318,7 +319,7 @@ class SqliteSaver(BaseCheckpointSaver):
         ORDER BY checkpoint_id DESC"""
         if limit:
             query += f" LIMIT {limit}"
-        with self.cursor(transaction=False) as cur:
+        with self.cursor(transaction=False) as cur, closing(self.conn.cursor()) as wcur:
             cur.execute(query, param_values)
             for (
                 thread_id,
@@ -329,6 +330,10 @@ class SqliteSaver(BaseCheckpointSaver):
                 checkpoint,
                 metadata,
             ) in cur:
+                wcur.execute(
+                    "SELECT task_id, channel, type, value FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?",
+                    (thread_id, checkpoint_ns, checkpoint_id),
+                )
                 yield CheckpointTuple(
                     {
                         "configurable": {
@@ -350,6 +355,10 @@ class SqliteSaver(BaseCheckpointSaver):
                         if parent_checkpoint_id
                         else None
                     ),
+                    [
+                        (task_id, channel, self.serde.loads_typed((type, value)))
+                        for task_id, channel, type, value in wcur
+                    ],
                 )
 
     def put(
@@ -424,25 +433,15 @@ class SqliteSaver(BaseCheckpointSaver):
             task_id (str): Identifier for the task creating the writes.
         """
         with self.lock, self.cursor() as cur:
-            cur.execute(
-                "DELETE FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ? AND task_id = ? AND idx >= ?",
-                (
-                    str(config["configurable"]["thread_id"]),
-                    str(config["configurable"]["checkpoint_ns"]),
-                    str(config["configurable"]["checkpoint_id"]),
-                    task_id,
-                    len(writes),
-                ),
-            )
             cur.executemany(
-                "INSERT OR REPLACE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
+                "INSERT OR IGNORE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
                 [
                     (
                         str(config["configurable"]["thread_id"]),
                         str(config["configurable"]["checkpoint_ns"]),
                         str(config["configurable"]["checkpoint_id"]),
                         task_id,
-                        idx,
+                        WRITES_IDX_MAP.get(channel, idx),
                         channel,
                         *self.serde.dumps_typed(value),
                     )

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

@@ -16,6 +16,7 @@ import aiosqlite
 from langchain_core.runnables import RunnableConfig
 
 from langgraph.checkpoint.base import (
+    WRITES_IDX_MAP,
     BaseCheckpointSaver,
     ChannelVersions,
     Checkpoint,
@@ -329,14 +330,14 @@ class AsyncSqliteSaver(BaseCheckpointSaver):
             AsyncIterator[CheckpointTuple]: An asynchronous iterator of matching checkpoint tuples.
         """
         await self.setup()
-        where, param_values = search_where(config, filter, before)
+        where, params = search_where(config, filter, before)
         query = f"""SELECT thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata
         FROM checkpoints
         {where}
         ORDER BY checkpoint_id DESC"""
         if limit:
             query += f" LIMIT {limit}"
-        async with self.conn.execute(query, param_values) as cursor:
+        async with self.conn.execute(query, params) as cur, self.conn.cursor() as wcur:
             async for (
                 thread_id,
                 checkpoint_ns,
@@ -345,7 +346,11 @@ class AsyncSqliteSaver(BaseCheckpointSaver):
                 type,
                 checkpoint,
                 metadata,
-            ) in cursor:
+            ) in cur:
+                await wcur.execute(
+                    "SELECT task_id, channel, type, value FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?",
+                    (thread_id, checkpoint_ns, checkpoint_id),
+                )
                 yield CheckpointTuple(
                     {
                         "configurable": {
@@ -367,6 +372,10 @@ class AsyncSqliteSaver(BaseCheckpointSaver):
                         if parent_checkpoint_id
                         else None
                     ),
+                    [
+                        (task_id, channel, self.serde.loads_typed((type, value)))
+                        async for task_id, channel, type, value in wcur
+                    ],
                 )
 
     async def aput(
@@ -433,25 +442,15 @@ class AsyncSqliteSaver(BaseCheckpointSaver):
         """
         await self.setup()
         async with self.conn.cursor() as cur:
-            await cur.execute(
-                "DELETE FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ? AND task_id = ? AND idx >= ?",
-                (
-                    str(config["configurable"]["thread_id"]),
-                    str(config["configurable"]["checkpoint_ns"]),
-                    str(config["configurable"]["checkpoint_id"]),
-                    task_id,
-                    len(writes),
-                ),
-            )
             await cur.executemany(
-                "INSERT OR REPLACE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
+                "INSERT OR IGNORE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
                 [
                     (
                         str(config["configurable"]["thread_id"]),
                         str(config["configurable"]["checkpoint_ns"]),
                         str(config["configurable"]["checkpoint_id"]),
                         task_id,
-                        idx,
+                        WRITES_IDX_MAP.get(channel, idx),
                         channel,
                         *self.serde.dumps_typed(value),
                     )

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/poetry.lock

@@ -1,4 +1,4 @@
-# This file is automatically @generated by Poetry 1.8.3 and should not be changed by hand.
+# This file is automatically @generated by Poetry 1.8.2 and should not be changed by hand.
 
 [[package]]
 name = "aiosqlite"
@@ -380,6 +380,8 @@ files = [
     {file = "orjson-3.10.6-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:960db0e31c4e52fa0fc3ecbaea5b2d3b58f379e32a95ae6b0ebeaa25b93dfd34"},
     {file = "orjson-3.10.6-cp312-none-win32.whl", hash = "sha256:a6ea7afb5b30b2317e0bee03c8d34c8181bc5a36f2afd4d0952f378972c4efd5"},
     {file = "orjson-3.10.6-cp312-none-win_amd64.whl", hash = "sha256:874ce88264b7e655dde4aeaacdc8fd772a7962faadfb41abe63e2a4861abc3dc"},
+    {file = "orjson-3.10.6-cp313-none-win32.whl", hash = "sha256:efdf2c5cde290ae6b83095f03119bdc00303d7a03b42b16c54517baa3c4ca3d0"},
+    {file = "orjson-3.10.6-cp313-none-win_amd64.whl", hash = "sha256:8e190fe7888e2e4392f52cafb9626113ba135ef53aacc65cd13109eb9746c43e"},
     {file = "orjson-3.10.6-cp38-cp38-macosx_10_15_x86_64.macosx_11_0_arm64.macosx_10_15_universal2.whl", hash = "sha256:66680eae4c4e7fc193d91cfc1353ad6d01b4801ae9b5314f17e11ba55e934183"},
     {file = "orjson-3.10.6-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:caff75b425db5ef8e8f23af93c80f072f97b4fb3afd4af44482905c9f588da28"},
     {file = "orjson-3.10.6-cp38-cp38-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:3722fddb821b6036fd2a3c814f6bd9b57a89dc6337b9924ecd614ebce3271394"},
@@ -707,28 +709,29 @@ use-chardet-on-py3 = ["chardet (>=3.0.2,<6)"]
 
 [[package]]
 name = "ruff"
-version = "0.1.15"
+version = "0.6.2"
 description = "An extremely fast Python linter and code formatter, written in Rust."
 optional = false
 python-versions = ">=3.7"
 files = [
-    {file = "ruff-0.1.15-py3-none-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl", hash = "sha256:5fe8d54df166ecc24106db7dd6a68d44852d14eb0729ea4672bb4d96c320b7df"},
-    {file = "ruff-0.1.15-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:6f0bfbb53c4b4de117ac4d6ddfd33aa5fc31beeaa21d23c45c6dd249faf9126f"},
-    {file = "ruff-0.1.15-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e0d432aec35bfc0d800d4f70eba26e23a352386be3a6cf157083d18f6f5881c8"},
-    {file = "ruff-0.1.15-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:9405fa9ac0e97f35aaddf185a1be194a589424b8713e3b97b762336ec79ff807"},
-    {file = "ruff-0.1.15-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:c66ec24fe36841636e814b8f90f572a8c0cb0e54d8b5c2d0e300d28a0d7bffec"},
-    {file = "ruff-0.1.15-py3-none-manylinux_2_17_ppc64.manylinux2014_ppc64.whl", hash = "sha256:6f8ad828f01e8dd32cc58bc28375150171d198491fc901f6f98d2a39ba8e3ff5"},
-    {file = "ruff-0.1.15-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:86811954eec63e9ea162af0ffa9f8d09088bab51b7438e8b6488b9401863c25e"},
-    {file = "ruff-0.1.15-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:fd4025ac5e87d9b80e1f300207eb2fd099ff8200fa2320d7dc066a3f4622dc6b"},
-    {file = "ruff-0.1.15-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b17b93c02cdb6aeb696effecea1095ac93f3884a49a554a9afa76bb125c114c1"},
-    {file = "ruff-0.1.15-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:ddb87643be40f034e97e97f5bc2ef7ce39de20e34608f3f829db727a93fb82c5"},
-    {file = "ruff-0.1.15-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:abf4822129ed3a5ce54383d5f0e964e7fef74a41e48eb1dfad404151efc130a2"},
-    {file = "ruff-0.1.15-py3-none-musllinux_1_2_i686.whl", hash = "sha256:6c629cf64bacfd136c07c78ac10a54578ec9d1bd2a9d395efbee0935868bf852"},
-    {file = "ruff-0.1.15-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:1bab866aafb53da39c2cadfb8e1c4550ac5340bb40300083eb8967ba25481447"},
-    {file = "ruff-0.1.15-py3-none-win32.whl", hash = "sha256:2417e1cb6e2068389b07e6fa74c306b2810fe3ee3476d5b8a96616633f40d14f"},
-    {file = "ruff-0.1.15-py3-none-win_amd64.whl", hash = "sha256:3837ac73d869efc4182d9036b1405ef4c73d9b1f88da2413875e34e0d6919587"},
-    {file = "ruff-0.1.15-py3-none-win_arm64.whl", hash = "sha256:9a933dfb1c14ec7a33cceb1e49ec4a16b51ce3c20fd42663198746efc0427360"},
-    {file = "ruff-0.1.15.tar.gz", hash = "sha256:f6dfa8c1b21c913c326919056c390966648b680966febcb796cc9d1aaab8564e"},
+    {file = "ruff-0.6.2-py3-none-linux_armv6l.whl", hash = "sha256:5c8cbc6252deb3ea840ad6a20b0f8583caab0c5ef4f9cca21adc5a92b8f79f3c"},
+    {file = "ruff-0.6.2-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:17002fe241e76544448a8e1e6118abecbe8cd10cf68fde635dad480dba594570"},
+    {file = "ruff-0.6.2-py3-none-macosx_11_0_arm64.whl", hash = "sha256:3dbeac76ed13456f8158b8f4fe087bf87882e645c8e8b606dd17b0b66c2c1158"},
+    {file = "ruff-0.6.2-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:094600ee88cda325988d3f54e3588c46de5c18dae09d683ace278b11f9d4d534"},
+    {file = "ruff-0.6.2-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:316d418fe258c036ba05fbf7dfc1f7d3d4096db63431546163b472285668132b"},
+    {file = "ruff-0.6.2-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:d72b8b3abf8a2d51b7b9944a41307d2f442558ccb3859bbd87e6ae9be1694a5d"},
+    {file = "ruff-0.6.2-py3-none-manylinux_2_17_ppc64.manylinux2014_ppc64.whl", hash = "sha256:2aed7e243be68487aa8982e91c6e260982d00da3f38955873aecd5a9204b1d66"},
+    {file = "ruff-0.6.2-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:d371f7fc9cec83497fe7cf5eaf5b76e22a8efce463de5f775a1826197feb9df8"},
+    {file = "ruff-0.6.2-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:a8f310d63af08f583363dfb844ba8f9417b558199c58a5999215082036d795a1"},
+    {file = "ruff-0.6.2-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7db6880c53c56addb8638fe444818183385ec85eeada1d48fc5abe045301b2f1"},
+    {file = "ruff-0.6.2-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:1175d39faadd9a50718f478d23bfc1d4da5743f1ab56af81a2b6caf0a2394f23"},
+    {file = "ruff-0.6.2-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:5b939f9c86d51635fe486585389f54582f0d65b8238e08c327c1534844b3bb9a"},
+    {file = "ruff-0.6.2-py3-none-musllinux_1_2_i686.whl", hash = "sha256:d0d62ca91219f906caf9b187dea50d17353f15ec9bb15aae4a606cd697b49b4c"},
+    {file = "ruff-0.6.2-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:7438a7288f9d67ed3c8ce4d059e67f7ed65e9fe3aa2ab6f5b4b3610e57e3cb56"},
+    {file = "ruff-0.6.2-py3-none-win32.whl", hash = "sha256:279d5f7d86696df5f9549b56b9b6a7f6c72961b619022b5b7999b15db392a4da"},
+    {file = "ruff-0.6.2-py3-none-win_amd64.whl", hash = "sha256:d9f3469c7dd43cd22eb1c3fc16926fb8258d50cb1b216658a07be95dd117b0f2"},
+    {file = "ruff-0.6.2-py3-none-win_arm64.whl", hash = "sha256:f28fcd2cd0e02bdf739297516d5643a945cc7caf09bd9bcb4d932540a5ea4fa9"},
+    {file = "ruff-0.6.2.tar.gz", hash = "sha256:239ee6beb9e91feb8e0ec384204a763f36cb53fb895a1a364618c6abb076b3be"},
 ]
 
 [[package]]
@@ -832,4 +835,4 @@ watchmedo = ["PyYAML (>=3.10)"]
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.9.0"
-content-hash = "e073e1a73cdae1fae8ea46499c39e55980cb61c7c8bd6be774c64f80a627eb31"
+content-hash = "d50ec7c6b55075d19193e080cc95d3b0998b9efa1a0ed407bd6055b5b5e867e0"

libs/checkpoint-sqlite/pyproject.toml

@@ -14,7 +14,7 @@ langgraph-checkpoint = "^1.0.1"
 aiosqlite = "^0.20.0"
 
 [tool.poetry.group.dev.dependencies]
-ruff = "^0.1.4"
+ruff = "^0.6.2"
 codespell = "^2.2.0"
 pytest = "^7.2.1"
 pytest-asyncio = "^0.21.1"

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/Makefile

@@ -24,11 +24,11 @@ lint_tests: PYTHON_FILES=tests
 lint_tests: MYPY_CACHE=.mypy_cache_test
 
 lint lint_diff lint_package lint_tests:
-	poetry run ruff .
+	poetry run ruff check .
 	[ "$(PYTHON_FILES)" = "" ] || poetry run ruff format $(PYTHON_FILES) --diff
-	[ "$(PYTHON_FILES)" = "" ] || poetry run ruff --select I $(PYTHON_FILES)
+	[ "$(PYTHON_FILES)" = "" ] || poetry run ruff check --select I $(PYTHON_FILES)
 	[ "$(PYTHON_FILES)" = "" ] || mkdir -p $(MYPY_CACHE) || poetry run mypy $(PYTHON_FILES) --cache-dir $(MYPY_CACHE)
 
 format format_diff:
 	poetry run ruff format $(PYTHON_FILES)
-	poetry run ruff --select I --fix $(PYTHON_FILES)
+	poetry run ruff check --select I --fix $(PYTHON_FILES)

libs/checkpoint/README.md

@@ -74,7 +74,6 @@ checkpoint = {
       }
     },
     "pending_sends": [],
-    "current_tasks": {}
 }
 
 # store checkpoint

libs/checkpoint/langgraph/checkpoint/base/__init__.py

@@ -1,4 +1,3 @@
-from abc import ABC
 from datetime import datetime, timezone
 from typing import (
     Any,
@@ -22,6 +21,7 @@ from langgraph.checkpoint.base.id import uuid6
 from langgraph.checkpoint.serde.base import SerializerProtocol, maybe_add_typed_methods
 from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
 from langgraph.checkpoint.serde.types import (
+    ERROR,
     ChannelProtocol,
     SendProtocol,
 )
@@ -96,8 +96,6 @@ class Checkpoint(TypedDict):
     pending_sends: List[SendProtocol]
     """List of packets sent to nodes but not yet processed.
     Cleared by the next checkpoint."""
-    current_tasks: Dict[str, TaskInfo]
-    """Map from task ID to task info."""
 
 
 def empty_checkpoint() -> Checkpoint:
@@ -109,7 +107,6 @@ def empty_checkpoint() -> Checkpoint:
         channel_versions={},
         versions_seen={},
         pending_sends=[],
-        current_tasks={},
     )
 
 
@@ -122,7 +119,6 @@ def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint:
         channel_versions=checkpoint["channel_versions"].copy(),
         versions_seen={k: v.copy() for k, v in checkpoint["versions_seen"].items()},
         pending_sends=checkpoint.get("pending_sends", []).copy(),
-        current_tasks=checkpoint.get("current_tasks", {}).copy(),
     )
 
 
@@ -140,6 +136,8 @@ def create_checkpoint(
     else:
         values: dict[str, Any] = {}
         for k, v in channels.items():
+            if k not in checkpoint["channel_versions"]:
+                continue
             try:
                 values[k] = v.checkpoint()
             except EmptyChannelError:
@@ -152,7 +150,6 @@ def create_checkpoint(
         channel_versions=checkpoint["channel_versions"],
         versions_seen=checkpoint["versions_seen"],
         pending_sends=checkpoint.get("pending_sends", []),
-        current_tasks={},
     )
 
 
@@ -194,7 +191,7 @@ CheckpointId = ConfigurableFieldSpec(
 )
 
 
-class BaseCheckpointSaver(ABC):
+class BaseCheckpointSaver:
     """Base class for creating a graph checkpointer.
 
     Checkpointers allow LangGraph agents to persist their state
@@ -437,3 +434,14 @@ def get_checkpoint_id(config: RunnableConfig) -> Optional[str]:
     return config["configurable"].get(
         "checkpoint_id", config["configurable"].get("thread_ts")
     )
+
+
+"""
+Mapping from error type to error index.
+Regular writes just map to their index in the list of writes being saved.
+Special writes (e.g. errors) map to negative indices, to avoid those writes from
+saving regular writes.
+Each Checkpointer implementation should use this mapping in put_writes.
+"""
+WRITES_IDX_MAP = {ERROR: -1}
+# TODO To store scheduled status of tasks, add a special channel here

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

@@ -8,6 +8,7 @@ from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Tuple
 from langchain_core.runnables import RunnableConfig
 
 from langgraph.checkpoint.base import (
+    WRITES_IDX_MAP,
     BaseCheckpointSaver,
     ChannelVersions,
     Checkpoint,
@@ -16,6 +17,7 @@ from langgraph.checkpoint.base import (
     SerializerProtocol,
     get_checkpoint_id,
 )
+from langgraph.checkpoint.serde.types import TASKS
 
 
 class MemorySaver(
@@ -52,6 +54,9 @@ class MemorySaver(
 
     # thread ID ->  checkpoint NS -> checkpoint ID -> checkpoint mapping
     storage: defaultdict[str, dict[str, dict[str, tuple[bytes, bytes, Optional[str]]]]]
+    writes: defaultdict[
+        tuple[str, str, str], dict[tuple[str, int], tuple[str, str, bytes]]
+    ]
 
     def __init__(
         self,
@@ -60,7 +65,7 @@ class MemorySaver(
     ) -> None:
         super().__init__(serde=serde)
         self.storage = defaultdict(lambda: defaultdict(dict))
-        self.writes = defaultdict(list)
+        self.writes = defaultdict(dict)
 
     def __enter__(self) -> "MemorySaver":
         return self
@@ -103,10 +108,23 @@ class MemorySaver(
         if checkpoint_id := get_checkpoint_id(config):
             if saved := self.storage[thread_id][checkpoint_ns].get(checkpoint_id):
                 checkpoint, metadata, parent_checkpoint_id = saved
-                writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)]
+                writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)].values()
+                if parent_checkpoint_id:
+                    sends = [
+                        w[2]
+                        for w in self.writes[
+                            (thread_id, checkpoint_ns, parent_checkpoint_id)
+                        ].values()
+                        if w[1] == TASKS
+                    ]
+                else:
+                    sends = []
                 return CheckpointTuple(
                     config=config,
-                    checkpoint=self.serde.loads_typed(checkpoint),
+                    checkpoint={
+                        **self.serde.loads_typed(checkpoint),
+                        "pending_sends": [self.serde.loads_typed(s) for s in sends],
+                    },
                     metadata=self.serde.loads_typed(metadata),
                     pending_writes=[
                         (id, c, self.serde.loads_typed(v)) for id, c, v in writes
@@ -125,7 +143,17 @@ class MemorySaver(
             if checkpoints := self.storage[thread_id][checkpoint_ns]:
                 checkpoint_id = max(checkpoints.keys())
                 checkpoint, metadata, parent_checkpoint_id = checkpoints[checkpoint_id]
-                writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)]
+                writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)].values()
+                if parent_checkpoint_id:
+                    sends = [
+                        w[2]
+                        for w in self.writes[
+                            (thread_id, checkpoint_ns, parent_checkpoint_id)
+                        ].values()
+                        if w[1] == TASKS
+                    ]
+                else:
+                    sends = []
                 return CheckpointTuple(
                     config={
                         "configurable": {
@@ -134,7 +162,10 @@ class MemorySaver(
                             "checkpoint_id": checkpoint_id,
                         }
                     },
-                    checkpoint=self.serde.loads_typed(checkpoint),
+                    checkpoint={
+                        **self.serde.loads_typed(checkpoint),
+                        "pending_sends": [self.serde.loads_typed(s) for s in sends],
+                    },
                     metadata=self.serde.loads_typed(metadata),
                     pending_writes=[
                         (id, c, self.serde.loads_typed(v)) for id, c, v in writes
@@ -173,57 +204,91 @@ 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
 
-                yield CheckpointTuple(
-                    config={
-                        "configurable": {
-                            "thread_id": thread_id,
-                            "checkpoint_ns": checkpoint_ns,
-                            "checkpoint_id": checkpoint_id,
+                    # 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
+
+                    # 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()
+
+                    if parent_checkpoint_id:
+                        sends = [
+                            w[2]
+                            for w in self.writes[
+                                (thread_id, checkpoint_ns, parent_checkpoint_id)
+                            ].values()
+                            if w[1] == TASKS
+                        ]
+                    else:
+                        sends = []
+
+                    yield CheckpointTuple(
+                        config={
+                            "configurable": {
+                                "thread_id": thread_id,
+                                "checkpoint_ns": checkpoint_ns,
+                                "checkpoint_id": checkpoint_id,
+                            }
+                        },
+                        checkpoint={
+                            **self.serde.loads_typed(checkpoint),
+                            "pending_sends": [self.serde.loads_typed(s) for s in sends],
+                        },
+                        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,
-                )
+                        if parent_checkpoint_id
+                        else None,
+                        pending_writes=[
+                            (id, c, self.serde.loads_typed(v)) for id, c, v in writes
+                        ],
+                    )
 
     def put(
         self,
@@ -246,12 +311,14 @@ class MemorySaver(
         Returns:
             RunnableConfig: The updated config containing the saved checkpoint's timestamp.
         """
+        c = checkpoint.copy()
+        c.pop("pending_sends")
         thread_id = config["configurable"]["thread_id"]
         checkpoint_ns = config["configurable"]["checkpoint_ns"]
         self.storage[thread_id][checkpoint_ns].update(
             {
                 checkpoint["id"]: (
-                    self.serde.dumps_typed(checkpoint),
+                    self.serde.dumps_typed(c),
                     self.serde.dumps_typed(metadata),
                     config["configurable"].get("checkpoint_id"),  # parent
                 )
@@ -287,11 +354,10 @@ class MemorySaver(
         thread_id = config["configurable"]["thread_id"]
         checkpoint_ns = config["configurable"]["checkpoint_ns"]
         checkpoint_id = config["configurable"]["checkpoint_id"]
-        key = (thread_id, checkpoint_ns, checkpoint_id)
-        self.writes[key] = [w for w in self.writes[key] if w[0] != task_id]
-        self.writes[key].extend(
-            [(task_id, c, self.serde.dumps_typed(v)) for c, v in writes]
-        )
+        outer_key = (thread_id, checkpoint_ns, checkpoint_id)
+        for idx, (c, v) in enumerate(writes):
+            inner_key = (task_id, WRITES_IDX_MAP.get(c, idx))
+            self.writes[outer_key][inner_key] = (task_id, c, self.serde.dumps_typed(v))
 
     async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
         """Asynchronous version of get_tuple.
@@ -330,7 +396,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/base.py

@@ -12,17 +12,13 @@ class SerializerProtocol(Protocol):
     Valid implementations include the `pickle`, `json` and `orjson` modules.
     """
 
-    def dumps(self, obj: Any) -> bytes:
-        ...
+    def dumps(self, obj: Any) -> bytes: ...
 
-    def dumps_typed(self, obj: Any) -> tuple[str, bytes]:
-        ...
+    def dumps_typed(self, obj: Any) -> tuple[str, bytes]: ...
 
-    def loads(self, data: bytes) -> Any:
-        ...
+    def loads(self, data: bytes) -> Any: ...
 
-    def loads_typed(self, data: tuple[str, bytes]) -> Any:
-        ...
+    def loads_typed(self, data: tuple[str, bytes]) -> Any: ...
 
 
 class SerializerCompat(SerializerProtocol):

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

@@ -12,6 +12,9 @@ from typing import (
 from langchain_core.runnables import RunnableConfig
 from typing_extensions import Self
 
+ERROR = "__error__"
+TASKS = "__pregel_tasks"
+
 Value = TypeVar("Value")
 Update = TypeVar("Update")
 C = TypeVar("C")
@@ -20,34 +23,26 @@ C = TypeVar("C")
 class ChannelProtocol(Protocol[Value, Update, C]):
     # Mirrors langgraph.channels.base.BaseChannel
     @property
-    def ValueType(self) -> Any:
-        ...
+    def ValueType(self) -> Any: ...
 
     @property
-    def UpdateType(self) -> Any:
-        ...
+    def UpdateType(self) -> Any: ...
 
-    def checkpoint(self) -> Optional[C]:
-        ...
+    def checkpoint(self) -> Optional[C]: ...
 
     def from_checkpoint(
         self, checkpoint: Optional[C], config: RunnableConfig
-    ) -> Generator[Self, None, None]:
-        ...
+    ) -> Generator[Self, None, None]: ...
 
     async def afrom_checkpoint(
         self, checkpoint: Optional[C], config: RunnableConfig
-    ) -> AsyncGenerator[Self, None]:
-        ...
+    ) -> AsyncGenerator[Self, None]: ...
 
-    def update(self, values: Sequence[Update]) -> bool:
-        ...
+    def update(self, values: Sequence[Update]) -> bool: ...
 
-    def get(self) -> Value:
-        ...
+    def get(self) -> Value: ...
 
-    def consume(self) -> bool:
-        ...
+    def consume(self) -> bool: ...
 
 
 @runtime_checkable
@@ -55,12 +50,10 @@ class SendProtocol(Protocol):
     # Mirrors langgraph.constants.Send
     node: str
     arg: Any
+    id: str
 
-    def __hash__(self) -> int:
-        ...
+    def __hash__(self) -> int: ...
 
-    def __repr__(self) -> str:
-        ...
+    def __repr__(self) -> str: ...
 
-    def __eq__(self, value: object) -> bool:
-        ...
+    def __eq__(self, value: object) -> bool: ...

libs/checkpoint/poetry.lock

@@ -1,4 +1,4 @@
-# This file is automatically @generated by Poetry 1.8.3 and should not be changed by hand.
+# This file is automatically @generated by Poetry 1.8.2 and should not be changed by hand.
 
 [[package]]
 name = "annotated-types"
@@ -380,6 +380,8 @@ files = [
     {file = "orjson-3.10.6-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:960db0e31c4e52fa0fc3ecbaea5b2d3b58f379e32a95ae6b0ebeaa25b93dfd34"},
     {file = "orjson-3.10.6-cp312-none-win32.whl", hash = "sha256:a6ea7afb5b30b2317e0bee03c8d34c8181bc5a36f2afd4d0952f378972c4efd5"},
     {file = "orjson-3.10.6-cp312-none-win_amd64.whl", hash = "sha256:874ce88264b7e655dde4aeaacdc8fd772a7962faadfb41abe63e2a4861abc3dc"},
+    {file = "orjson-3.10.6-cp313-none-win32.whl", hash = "sha256:efdf2c5cde290ae6b83095f03119bdc00303d7a03b42b16c54517baa3c4ca3d0"},
+    {file = "orjson-3.10.6-cp313-none-win_amd64.whl", hash = "sha256:8e190fe7888e2e4392f52cafb9626113ba135ef53aacc65cd13109eb9746c43e"},
     {file = "orjson-3.10.6-cp38-cp38-macosx_10_15_x86_64.macosx_11_0_arm64.macosx_10_15_universal2.whl", hash = "sha256:66680eae4c4e7fc193d91cfc1353ad6d01b4801ae9b5314f17e11ba55e934183"},
     {file = "orjson-3.10.6-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:caff75b425db5ef8e8f23af93c80f072f97b4fb3afd4af44482905c9f588da28"},
     {file = "orjson-3.10.6-cp38-cp38-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:3722fddb821b6036fd2a3c814f6bd9b57a89dc6337b9924ecd614ebce3271394"},
@@ -707,28 +709,29 @@ use-chardet-on-py3 = ["chardet (>=3.0.2,<6)"]
 
 [[package]]
 name = "ruff"
-version = "0.1.15"
+version = "0.6.2"
 description = "An extremely fast Python linter and code formatter, written in Rust."
 optional = false
 python-versions = ">=3.7"
 files = [
-    {file = "ruff-0.1.15-py3-none-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl", hash = "sha256:5fe8d54df166ecc24106db7dd6a68d44852d14eb0729ea4672bb4d96c320b7df"},
-    {file = "ruff-0.1.15-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:6f0bfbb53c4b4de117ac4d6ddfd33aa5fc31beeaa21d23c45c6dd249faf9126f"},
-    {file = "ruff-0.1.15-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e0d432aec35bfc0d800d4f70eba26e23a352386be3a6cf157083d18f6f5881c8"},
-    {file = "ruff-0.1.15-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:9405fa9ac0e97f35aaddf185a1be194a589424b8713e3b97b762336ec79ff807"},
-    {file = "ruff-0.1.15-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:c66ec24fe36841636e814b8f90f572a8c0cb0e54d8b5c2d0e300d28a0d7bffec"},
-    {file = "ruff-0.1.15-py3-none-manylinux_2_17_ppc64.manylinux2014_ppc64.whl", hash = "sha256:6f8ad828f01e8dd32cc58bc28375150171d198491fc901f6f98d2a39ba8e3ff5"},
-    {file = "ruff-0.1.15-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:86811954eec63e9ea162af0ffa9f8d09088bab51b7438e8b6488b9401863c25e"},
-    {file = "ruff-0.1.15-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:fd4025ac5e87d9b80e1f300207eb2fd099ff8200fa2320d7dc066a3f4622dc6b"},
-    {file = "ruff-0.1.15-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b17b93c02cdb6aeb696effecea1095ac93f3884a49a554a9afa76bb125c114c1"},
-    {file = "ruff-0.1.15-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:ddb87643be40f034e97e97f5bc2ef7ce39de20e34608f3f829db727a93fb82c5"},
-    {file = "ruff-0.1.15-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:abf4822129ed3a5ce54383d5f0e964e7fef74a41e48eb1dfad404151efc130a2"},
-    {file = "ruff-0.1.15-py3-none-musllinux_1_2_i686.whl", hash = "sha256:6c629cf64bacfd136c07c78ac10a54578ec9d1bd2a9d395efbee0935868bf852"},
-    {file = "ruff-0.1.15-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:1bab866aafb53da39c2cadfb8e1c4550ac5340bb40300083eb8967ba25481447"},
-    {file = "ruff-0.1.15-py3-none-win32.whl", hash = "sha256:2417e1cb6e2068389b07e6fa74c306b2810fe3ee3476d5b8a96616633f40d14f"},
-    {file = "ruff-0.1.15-py3-none-win_amd64.whl", hash = "sha256:3837ac73d869efc4182d9036b1405ef4c73d9b1f88da2413875e34e0d6919587"},
-    {file = "ruff-0.1.15-py3-none-win_arm64.whl", hash = "sha256:9a933dfb1c14ec7a33cceb1e49ec4a16b51ce3c20fd42663198746efc0427360"},
-    {file = "ruff-0.1.15.tar.gz", hash = "sha256:f6dfa8c1b21c913c326919056c390966648b680966febcb796cc9d1aaab8564e"},
+    {file = "ruff-0.6.2-py3-none-linux_armv6l.whl", hash = "sha256:5c8cbc6252deb3ea840ad6a20b0f8583caab0c5ef4f9cca21adc5a92b8f79f3c"},
+    {file = "ruff-0.6.2-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:17002fe241e76544448a8e1e6118abecbe8cd10cf68fde635dad480dba594570"},
+    {file = "ruff-0.6.2-py3-none-macosx_11_0_arm64.whl", hash = "sha256:3dbeac76ed13456f8158b8f4fe087bf87882e645c8e8b606dd17b0b66c2c1158"},
+    {file = "ruff-0.6.2-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:094600ee88cda325988d3f54e3588c46de5c18dae09d683ace278b11f9d4d534"},
+    {file = "ruff-0.6.2-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:316d418fe258c036ba05fbf7dfc1f7d3d4096db63431546163b472285668132b"},
+    {file = "ruff-0.6.2-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:d72b8b3abf8a2d51b7b9944a41307d2f442558ccb3859bbd87e6ae9be1694a5d"},
+    {file = "ruff-0.6.2-py3-none-manylinux_2_17_ppc64.manylinux2014_ppc64.whl", hash = "sha256:2aed7e243be68487aa8982e91c6e260982d00da3f38955873aecd5a9204b1d66"},
+    {file = "ruff-0.6.2-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:d371f7fc9cec83497fe7cf5eaf5b76e22a8efce463de5f775a1826197feb9df8"},
+    {file = "ruff-0.6.2-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:a8f310d63af08f583363dfb844ba8f9417b558199c58a5999215082036d795a1"},
+    {file = "ruff-0.6.2-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7db6880c53c56addb8638fe444818183385ec85eeada1d48fc5abe045301b2f1"},
+    {file = "ruff-0.6.2-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:1175d39faadd9a50718f478d23bfc1d4da5743f1ab56af81a2b6caf0a2394f23"},
+    {file = "ruff-0.6.2-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:5b939f9c86d51635fe486585389f54582f0d65b8238e08c327c1534844b3bb9a"},
+    {file = "ruff-0.6.2-py3-none-musllinux_1_2_i686.whl", hash = "sha256:d0d62ca91219f906caf9b187dea50d17353f15ec9bb15aae4a606cd697b49b4c"},
+    {file = "ruff-0.6.2-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:7438a7288f9d67ed3c8ce4d059e67f7ed65e9fe3aa2ab6f5b4b3610e57e3cb56"},
+    {file = "ruff-0.6.2-py3-none-win32.whl", hash = "sha256:279d5f7d86696df5f9549b56b9b6a7f6c72961b619022b5b7999b15db392a4da"},
+    {file = "ruff-0.6.2-py3-none-win_amd64.whl", hash = "sha256:d9f3469c7dd43cd22eb1c3fc16926fb8258d50cb1b216658a07be95dd117b0f2"},
+    {file = "ruff-0.6.2-py3-none-win_arm64.whl", hash = "sha256:f28fcd2cd0e02bdf739297516d5643a945cc7caf09bd9bcb4d932540a5ea4fa9"},
+    {file = "ruff-0.6.2.tar.gz", hash = "sha256:239ee6beb9e91feb8e0ec384204a763f36cb53fb895a1a364618c6abb076b3be"},
 ]
 
 [[package]]
@@ -847,4 +850,4 @@ watchmedo = ["PyYAML (>=3.10)"]
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.9.0,<4.0"
-content-hash = "3bee25f1adc1349de4358a88037693cec47f6a2b88b809490f9198e6313239f6"
+content-hash = "de79db3dc7701542739b3417e9d3f02c3b41167719603ee4d08b92e23b7443ee"

libs/checkpoint/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph-checkpoint"
-version = "1.0.3"
+version = "1.0.6"
 description = "Library with base interfaces for LangGraph checkpoint savers."
 authors = []
 license = "MIT"
@@ -13,7 +13,7 @@ python = "^3.9.0,<4.0"
 langchain-core = ">=0.2.22,<0.3"
 
 [tool.poetry.group.dev.dependencies]
-ruff = "^0.1.4"
+ruff = "^0.6.2"
 codespell = "^2.2.0"
 pytest = "^7.2.1"
 pytest-asyncio = "^0.21.1"

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/cli/Makefile

@@ -21,11 +21,11 @@ lint_tests: PYTHON_FILES=tests
 lint_tests: MYPY_CACHE=.mypy_cache_test
 
 lint lint_diff lint_package lint_tests:
-	poetry run ruff .
+	poetry run ruff check .
 	[ "$(PYTHON_FILES)" = "" ] || poetry run ruff format $(PYTHON_FILES) --diff
-	[ "$(PYTHON_FILES)" = "" ] || poetry run ruff --select I $(PYTHON_FILES)
+	[ "$(PYTHON_FILES)" = "" ] || poetry run ruff check --select I $(PYTHON_FILES)
 	[ "$(PYTHON_FILES)" = "" ] || mkdir -p $(MYPY_CACHE) || poetry run mypy $(PYTHON_FILES) --cache-dir $(MYPY_CACHE)
 
 format format_diff:
 	poetry run ruff format $(PYTHON_FILES)
-	poetry run ruff --select I --fix $(PYTHON_FILES)
+	poetry run ruff check --select I --fix $(PYTHON_FILES)

libs/cli/examples/graphs/storm.py

@@ -339,7 +339,7 @@ async def gen_answer(
     # We could be more precise about handling max token length if we wanted to here
     dumped = json.dumps(all_query_results)[:max_str_len]
     ai_message: AIMessage = queries["raw"]
-    tool_call = queries["raw"].additional_kwargs["tool_calls"][0]
+    tool_call = queries["raw"].tool_calls[0]
     tool_id = tool_call["id"]
     tool_message = ToolMessage(tool_call_id=tool_id, content=dumped)
     swapped_state["messages"].extend([ai_message, tool_message])

libs/cli/langgraph_cli/cli.py

@@ -372,7 +372,11 @@ def _build(
     pull: bool,
     tag: str,
 ):
-    base_image = base_image or "langchain/langgraph-api"
+    base_image = base_image or (
+        "langchain/langgraphjs-api"
+        if config_json.get("node_version")
+        else "langchain/langgraph-api"
+    )
 
     # pull latest images
     if pull:
@@ -380,7 +384,9 @@ def _build(
             subp_exec(
                 "docker",
                 "pull",
-                f"{base_image}:{config_json['python_version']}",
+                f"{base_image}:{config_json['node_version']}"
+                if config_json.get("node_version")
+                else f"{base_image}:{config_json['python_version']}",
                 verbose=True,
             )
         )
@@ -462,7 +468,11 @@ def dockerfile(save_path: pathlib.Path, config: pathlib.Path):
     with open(save_path, "w") as f:
         f.write(
             langgraph_cli.config.config_to_docker(
-                config, config_json, "langchain/langgraph-api"
+                config,
+                config_json,
+                "langchain/langgraphjs-api"
+                if config_json.get("node_version")
+                else "langchain/langgraph-api",
             )
         )
 
@@ -500,7 +510,9 @@ def prepare_args_and_stdin(
         config_path,
         config,
         watch=watch,
-        base_image="langchain/langgraph-api",
+        base_image="langchain/langgraphjs-api"
+        if config.get("node_version")
+        else "langchain/langgraph-api",
     )
     return args, stdin
 
@@ -527,7 +539,9 @@ def prepare(
             subp_exec(
                 "docker",
                 "pull",
-                f"langchain/langgraph-api:{config['python_version']}",
+                f"langchain/langgraphjs-api:{config['node_version']}"
+                if config.get("node_version")
+                else f"langchain/langgraph-api:{config['python_version']}",
                 verbose=verbose,
             )
         )

libs/cli/langgraph_cli/config.py

@@ -9,6 +9,7 @@ import click
 
 class Config(TypedDict):
     python_version: str
+    node_version: Optional[str]
     pip_config_file: Optional[str]
     dockerfile_lines: list[str]
     dependencies: list[str]
@@ -17,27 +18,46 @@ class Config(TypedDict):
 
 
 def validate_config(config: Config) -> Config:
-    config = {
-        "python_version": config.get("python_version", "3.11"),
-        "pip_config_file": config.get("pip_config_file"),
-        "dockerfile_lines": config.get("dockerfile_lines", []),
-        "dependencies": config.get("dependencies", []),
-        "graphs": config.get("graphs", {}),
-        "env": config.get("env", {}),
-    }
-    if config["python_version"] not in (
-        "3.11",
-        "3.12",
-    ):
-        raise click.UsageError(
-            f"Unsupported Python version: {config['python_version']}. "
-            "Supported versions are 3.11 and 3.12."
-        )
-    if not config["dependencies"]:
-        raise click.UsageError(
-            "No dependencies found in config. "
-            "Add at least one dependency to 'dependencies' list."
-        )
+    config = (
+        {
+            "node_version": config.get("node_version"),
+            "dockerfile_lines": config.get("dockerfile_lines", []),
+            "graphs": config.get("graphs", {}),
+            "env": config.get("env", {}),
+        }
+        if config.get("node_version")
+        else {
+            "python_version": config.get("python_version", "3.11"),
+            "pip_config_file": config.get("pip_config_file"),
+            "dockerfile_lines": config.get("dockerfile_lines", []),
+            "dependencies": config.get("dependencies", []),
+            "graphs": config.get("graphs", {}),
+            "env": config.get("env", {}),
+        }
+    )
+
+    if config.get("node_version"):
+        if config["node_version"] not in ("20",):
+            raise click.UsageError(
+                f"Unsupported Node.js version: {config['node_version']}. "
+                "Currently only `node_version: \"20\"` is supported."
+            )
+
+    if config.get("python_version"):
+        if config["python_version"] not in (
+            "3.11",
+            "3.12",
+        ):
+            raise click.UsageError(
+                f"Unsupported Python version: {config['python_version']}. "
+                "Supported versions are 3.11 and 3.12."
+            )
+        if not config["dependencies"]:
+            raise click.UsageError(
+                "No dependencies found in config. "
+                "Add at least one dependency to 'dependencies' list."
+            )
+
     if not config["graphs"]:
         raise click.UsageError(
             "No graphs found in config. "
@@ -191,7 +211,7 @@ def _update_graph_paths(
             config["graphs"][graph_id] = f"{module_str}:{attr_str}"
 
 
-def config_to_docker(config_path: pathlib.Path, config: Config, base_image: str):
+def python_config_to_docker(config_path: pathlib.Path, config: Config, base_image: str):
     # configure pip
     pip_install = (
         "PYTHONDONTWRITEBYTECODE=1 pip install --no-cache-dir -c /api/constraints.txt"
@@ -266,6 +286,29 @@ ENV LANGSERVE_GRAPHS='{json.dumps(config["graphs"])}'
 {f"WORKDIR {local_deps.working_dir}" if local_deps.working_dir else ""}"""
 
 
+def node_config_to_docker(config_path: pathlib.Path, config: Config, base_image: str):
+    faux_path = f"/deps/{config_path.parent.name}"
+
+    return f"""FROM {base_image}:{config['node_version']}
+
+{os.linesep.join(config["dockerfile_lines"])}
+
+ADD . {faux_path}
+
+RUN cd {faux_path} && yarn install --frozen-lockfile
+
+ENV LANGSERVE_GRAPHS='{json.dumps(config["graphs"])}'
+
+WORKDIR {faux_path}"""
+
+
+def config_to_docker(config_path: pathlib.Path, config: Config, base_image: str):
+    if config.get("node_version"):
+        return node_config_to_docker(config_path, config, base_image)
+
+    return python_config_to_docker(config_path, config, base_image)
+
+
 def config_to_compose(
     config_path: pathlib.Path,
     config: Config,

libs/cli/langgraph_cli/docker.py

@@ -12,6 +12,15 @@ DEFAULT_POSTGRES_URI = (
     "postgres://postgres:postgres@langgraph-postgres:5432/postgres?sslmode=disable"
 )
 
+REDIS = """
+    langgraph-redis:
+        image: redis:6
+        healthcheck:
+            test: redis-cli ping
+            interval: 5s
+            timeout: 1s
+            retries: 5
+"""
 
 DB = """
     langgraph-postgres:
@@ -165,18 +174,22 @@ def compose(
             interval: 5s"""
 
     compose_str = f"""{volumes}services:
+{REDIS}
 {db}
 {debugger_compose(port=debugger_port, base_url=debugger_base_url)}
     langgraph-api:
         ports:
-            - "{port}:8000\""""
+            - "{port}:8000\"
+        depends_on:
+            langgraph-redis:
+                condition: service_healthy"""
     if include_db:
         compose_str += """
-        depends_on:
             langgraph-postgres:
                 condition: service_healthy"""
     compose_str += f"""
         environment:
+            REDIS_URI: redis://langgraph-redis:6379
             POSTGRES_URI: {postgres_uri}
 """
     if capabilities.healthcheck_start_interval:

libs/cli/poetry.lock

@@ -236,28 +236,29 @@ watchdog = ">=0.6.0"
 
 [[package]]
 name = "ruff"
-version = "0.1.6"
+version = "0.6.2"
 description = "An extremely fast Python linter and code formatter, written in Rust."
 optional = false
 python-versions = ">=3.7"
 files = [
-    {file = "ruff-0.1.6-py3-none-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl", hash = "sha256:88b8cdf6abf98130991cbc9f6438f35f6e8d41a02622cc5ee130a02a0ed28703"},
-    {file = "ruff-0.1.6-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:5c549ed437680b6105a1299d2cd30e4964211606eeb48a0ff7a93ef70b902248"},
-    {file = "ruff-0.1.6-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1cf5f701062e294f2167e66d11b092bba7af6a057668ed618a9253e1e90cfd76"},
-    {file = "ruff-0.1.6-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:05991ee20d4ac4bb78385360c684e4b417edd971030ab12a4fbd075ff535050e"},
-    {file = "ruff-0.1.6-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:87455a0c1f739b3c069e2f4c43b66479a54dea0276dd5d4d67b091265f6fd1dc"},
-    {file = "ruff-0.1.6-py3-none-manylinux_2_17_ppc64.manylinux2014_ppc64.whl", hash = "sha256:683aa5bdda5a48cb8266fcde8eea2a6af4e5700a392c56ea5fb5f0d4bfdc0240"},
-    {file = "ruff-0.1.6-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:137852105586dcbf80c1717facb6781555c4e99f520c9c827bd414fac67ddfb6"},
-    {file = "ruff-0.1.6-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:bd98138a98d48a1c36c394fd6b84cd943ac92a08278aa8ac8c0fdefcf7138f35"},
-    {file = "ruff-0.1.6-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3a0cd909d25f227ac5c36d4e7e681577275fb74ba3b11d288aff7ec47e3ae745"},
-    {file = "ruff-0.1.6-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:e8fd1c62a47aa88a02707b5dd20c5ff20d035d634aa74826b42a1da77861b5ff"},
-    {file = "ruff-0.1.6-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:fd89b45d374935829134a082617954120d7a1470a9f0ec0e7f3ead983edc48cc"},
-    {file = "ruff-0.1.6-py3-none-musllinux_1_2_i686.whl", hash = "sha256:491262006e92f825b145cd1e52948073c56560243b55fb3b4ecb142f6f0e9543"},
-    {file = "ruff-0.1.6-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:ea284789861b8b5ca9d5443591a92a397ac183d4351882ab52f6296b4fdd5462"},
-    {file = "ruff-0.1.6-py3-none-win32.whl", hash = "sha256:1610e14750826dfc207ccbcdd7331b6bd285607d4181df9c1c6ae26646d6848a"},
-    {file = "ruff-0.1.6-py3-none-win_amd64.whl", hash = "sha256:4558b3e178145491e9bc3b2ee3c4b42f19d19384eaa5c59d10acf6e8f8b57e33"},
-    {file = "ruff-0.1.6-py3-none-win_arm64.whl", hash = "sha256:03910e81df0d8db0e30050725a5802441c2022ea3ae4fe0609b76081731accbc"},
-    {file = "ruff-0.1.6.tar.gz", hash = "sha256:1b09f29b16c6ead5ea6b097ef2764b42372aebe363722f1605ecbcd2b9207184"},
+    {file = "ruff-0.6.2-py3-none-linux_armv6l.whl", hash = "sha256:5c8cbc6252deb3ea840ad6a20b0f8583caab0c5ef4f9cca21adc5a92b8f79f3c"},
+    {file = "ruff-0.6.2-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:17002fe241e76544448a8e1e6118abecbe8cd10cf68fde635dad480dba594570"},
+    {file = "ruff-0.6.2-py3-none-macosx_11_0_arm64.whl", hash = "sha256:3dbeac76ed13456f8158b8f4fe087bf87882e645c8e8b606dd17b0b66c2c1158"},
+    {file = "ruff-0.6.2-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:094600ee88cda325988d3f54e3588c46de5c18dae09d683ace278b11f9d4d534"},
+    {file = "ruff-0.6.2-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:316d418fe258c036ba05fbf7dfc1f7d3d4096db63431546163b472285668132b"},
+    {file = "ruff-0.6.2-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:d72b8b3abf8a2d51b7b9944a41307d2f442558ccb3859bbd87e6ae9be1694a5d"},
+    {file = "ruff-0.6.2-py3-none-manylinux_2_17_ppc64.manylinux2014_ppc64.whl", hash = "sha256:2aed7e243be68487aa8982e91c6e260982d00da3f38955873aecd5a9204b1d66"},
+    {file = "ruff-0.6.2-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:d371f7fc9cec83497fe7cf5eaf5b76e22a8efce463de5f775a1826197feb9df8"},
+    {file = "ruff-0.6.2-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:a8f310d63af08f583363dfb844ba8f9417b558199c58a5999215082036d795a1"},
+    {file = "ruff-0.6.2-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7db6880c53c56addb8638fe444818183385ec85eeada1d48fc5abe045301b2f1"},
+    {file = "ruff-0.6.2-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:1175d39faadd9a50718f478d23bfc1d4da5743f1ab56af81a2b6caf0a2394f23"},
+    {file = "ruff-0.6.2-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:5b939f9c86d51635fe486585389f54582f0d65b8238e08c327c1534844b3bb9a"},
+    {file = "ruff-0.6.2-py3-none-musllinux_1_2_i686.whl", hash = "sha256:d0d62ca91219f906caf9b187dea50d17353f15ec9bb15aae4a606cd697b49b4c"},
+    {file = "ruff-0.6.2-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:7438a7288f9d67ed3c8ce4d059e67f7ed65e9fe3aa2ab6f5b4b3610e57e3cb56"},
+    {file = "ruff-0.6.2-py3-none-win32.whl", hash = "sha256:279d5f7d86696df5f9549b56b9b6a7f6c72961b619022b5b7999b15db392a4da"},
+    {file = "ruff-0.6.2-py3-none-win_amd64.whl", hash = "sha256:d9f3469c7dd43cd22eb1c3fc16926fb8258d50cb1b216658a07be95dd117b0f2"},
+    {file = "ruff-0.6.2-py3-none-win_arm64.whl", hash = "sha256:f28fcd2cd0e02bdf739297516d5643a945cc7caf09bd9bcb4d932540a5ea4fa9"},
+    {file = "ruff-0.6.2.tar.gz", hash = "sha256:239ee6beb9e91feb8e0ec384204a763f36cb53fb895a1a364618c6abb076b3be"},
 ]
 
 [[package]]
@@ -324,4 +325,4 @@ watchmedo = ["PyYAML (>=3.10)"]
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.9.0,<4.0"
-content-hash = "5efa2f1ed4bd611a45e5d43d7c3fb907a8fa4447e2d1c30ce26b830411e189dd"
+content-hash = "a1b0cc1de3e63b8d419342e311cd1e32b81606e89d3f8cf448be560b5727c202"

libs/cli/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph-cli"
-version = "0.1.50"
+version = "0.1.52"
 description = "CLI for interacting with LangGraph API"
 authors = []
 license = "MIT"
@@ -16,7 +16,7 @@ python = "^3.9.0,<4.0"
 click = "^8.1.7"
 
 [tool.poetry.group.dev.dependencies]
-ruff = "^0.1.4"
+ruff = "^0.6.2"
 codespell = "^2.2.0"
 pytest = "^7.2.1"
 pytest-asyncio = "^0.21.1"
@@ -40,7 +40,7 @@ requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"
 
 [tool.ruff]
-select = [
+lint.select = [
   # pycodestyle
   "E",
   # Pyflakes
@@ -52,4 +52,4 @@ select = [
   # isort
   "I",
 ]
-ignore = [ "E501", "B008" ]
+lint.ignore = [ "E501", "B008" ]

libs/cli/tests/unit_tests/test_cli.py

@@ -45,6 +45,13 @@ def test_prepare_args_and_stdin():
     langgraph-data:
         driver: local
 services:
+    langgraph-redis:
+        image: redis:6
+        healthcheck:
+            test: redis-cli ping
+            interval: 5s
+            timeout: 1s
+            retries: 5
     langgraph-postgres:
         image: postgres:16
         ports:
@@ -76,9 +83,12 @@ services:
         ports:
             - "8000:8000"
         depends_on:
+            langgraph-redis:
+                condition: service_healthy
             langgraph-postgres:
                 condition: service_healthy
         environment:
+            REDIS_URI: redis://langgraph-redis:6379
             POSTGRES_URI: {DEFAULT_POSTGRES_URI}
         healthcheck:
             test: python /api/healthcheck.py

libs/cli/tests/unit_tests/test_config.py

@@ -232,6 +232,31 @@ ENV LANGSERVE_GRAPHS='{"agent": "/deps/__outer_graphs/src/agent.py:graph"}'"""
     assert clean_empty_lines(actual_docker_stdin) == expected_docker_stdin
 
 
+# node.js build used for LangGraph Cloud
+def test_config_to_docker_nodejs():
+    graphs = {"agent": "./graphs/agent.js:graph"}
+    actual_docker_stdin = config_to_docker(
+        PATH_TO_CONFIG,
+        validate_config(
+            {
+                "node_version": "20",
+                "graphs": graphs,
+                "dockerfile_lines": ["ARG meow", "ARG foo"],
+            }
+        ),
+        "langchain/langgraphjs-api",
+    )
+    expected_docker_stdin = """FROM langchain/langgraphjs-api:20
+ARG meow
+ARG foo
+ADD . /deps/unit_tests
+RUN cd /deps/unit_tests && yarn install --frozen-lockfile
+ENV LANGSERVE_GRAPHS='{"agent": "./graphs/agent.js:graph"}'
+WORKDIR /deps/unit_tests"""
+
+    assert clean_empty_lines(actual_docker_stdin) == expected_docker_stdin
+
+
 # config_to_compose
 def test_config_to_compose_simple_config():
     graphs = {"agent": "./agent.py:graph"}

libs/cli/tests/unit_tests/test_docker.py

@@ -20,10 +20,21 @@ def test_compose_with_no_debugger_and_custom_db():
         DEFAULT_DOCKER_CAPABILITIES, port=port, postgres_uri=custom_postgres_uri
     )
     expected_compose_str = f"""services:
+    langgraph-redis:
+        image: redis:6
+        healthcheck:
+            test: redis-cli ping
+            interval: 5s
+            timeout: 1s
+            retries: 5
     langgraph-api:
         ports:
             - "{port}:8000"
+        depends_on:
+            langgraph-redis:
+                condition: service_healthy
         environment:
+            REDIS_URI: redis://langgraph-redis:6379
             POSTGRES_URI: {custom_postgres_uri}"""
     assert clean_empty_lines(actual_compose_str) == expected_compose_str
 
@@ -37,10 +48,21 @@ def test_compose_with_no_debugger_and_custom_db_with_healthcheck():
         postgres_uri=custom_postgres_uri,
     )
     expected_compose_str = f"""services:
+    langgraph-redis:
+        image: redis:6
+        healthcheck:
+            test: redis-cli ping
+            interval: 5s
+            timeout: 1s
+            retries: 5
     langgraph-api:
         ports:
             - "{port}:8000"
+        depends_on:
+            langgraph-redis:
+                condition: service_healthy
         environment:
+            REDIS_URI: redis://langgraph-redis:6379
             POSTGRES_URI: {custom_postgres_uri}
         healthcheck:
             test: python /api/healthcheck.py
@@ -59,10 +81,21 @@ def test_compose_with_debugger_and_custom_db():
         postgres_uri=custom_postgres_uri,
     )
     expected_compose_str = f"""services:
+    langgraph-redis:
+        image: redis:6
+        healthcheck:
+            test: redis-cli ping
+            interval: 5s
+            timeout: 1s
+            retries: 5
     langgraph-api:
         ports:
             - "{port}:8000"
+        depends_on:
+            langgraph-redis:
+                condition: service_healthy
         environment:
+            REDIS_URI: redis://langgraph-redis:6379
             POSTGRES_URI: {custom_postgres_uri}"""
     assert clean_empty_lines(actual_compose_str) == expected_compose_str
 
@@ -74,6 +107,13 @@ def test_compose_with_debugger_and_default_db():
     langgraph-data:
         driver: local
 services:
+    langgraph-redis:
+        image: redis:6
+        healthcheck:
+            test: redis-cli ping
+            interval: 5s
+            timeout: 1s
+            retries: 5
     langgraph-postgres:
         image: postgres:16
         ports:
@@ -94,8 +134,11 @@ services:
         ports:
             - "{port}:8000"
         depends_on:
+            langgraph-redis:
+                condition: service_healthy
             langgraph-postgres:
                 condition: service_healthy
         environment:
+            REDIS_URI: redis://langgraph-redis:6379
             POSTGRES_URI: {DEFAULT_POSTGRES_URI}"""
     assert clean_empty_lines(actual_compose_str) == expected_compose_str

libs/langgraph/Makefile

@@ -51,14 +51,14 @@ lint_tests: PYTHON_FILES=tests
 lint_tests: MYPY_CACHE=.mypy_cache_test
 
 lint lint_diff lint_package lint_tests:
-	poetry run ruff .
+	poetry run ruff check .
 	[ "$(PYTHON_FILES)" = "" ] || poetry run ruff format $(PYTHON_FILES) --diff
-	[ "$(PYTHON_FILES)" = "" ] || poetry run ruff --select I $(PYTHON_FILES)
+	[ "$(PYTHON_FILES)" = "" ] || poetry run ruff check --select I $(PYTHON_FILES)
 	[ "$(PYTHON_FILES)" = "" ] || mkdir -p $(MYPY_CACHE) || poetry run mypy $(PYTHON_FILES) --cache-dir $(MYPY_CACHE)
 
 format format_diff:
 	poetry run ruff format $(PYTHON_FILES)
-	poetry run ruff --select I --fix $(PYTHON_FILES)
+	poetry run ruff check --select I --fix $(PYTHON_FILES)
 
 spell_check:
 	poetry run codespell --toml pyproject.toml

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

@@ -2,9 +2,9 @@ from abc import ABC, abstractmethod
 from contextlib import asynccontextmanager, contextmanager
 from typing import (
     Any,
-    AsyncGenerator,
-    Generator,
+    AsyncIterator,
     Generic,
+    Iterator,
     Optional,
     Sequence,
     TypeVar,
@@ -21,6 +21,8 @@ C = TypeVar("C")
 
 
 class BaseChannel(Generic[Value, Update, C], ABC):
+    key: str = ""
+
     @property
     @abstractmethod
     def ValueType(self) -> Any:
@@ -33,29 +35,45 @@ 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
     def from_checkpoint(
         self, checkpoint: Optional[C], config: RunnableConfig
-    ) -> Generator[Self, None, None]:
+    ) -> Iterator[Self]:
         """Return a new identical channel, optionally initialized from a checkpoint.
         If the checkpoint contains complex data structures, they should be copied."""
 
+    @contextmanager
+    def from_checkpoint_named(
+        self, checkpoint: Optional[C], config: RunnableConfig
+    ) -> Iterator[Self]:
+        with self.from_checkpoint(checkpoint, config) as value:
+            value.key = self.key
+            yield value
+
     @asynccontextmanager
     async def afrom_checkpoint(
         self, checkpoint: Optional[C], config: RunnableConfig
-    ) -> AsyncGenerator[Self, None]:
+    ) -> AsyncIterator[Self]:
         """Return a new identical channel, optionally initialized from a checkpoint.
         If the checkpoint contains complex data structures, they should be copied."""
         with self.from_checkpoint(checkpoint, config) as value:
             yield value
 
+    @asynccontextmanager
+    async def afrom_checkpoint_named(
+        self, checkpoint: Optional[C], config: RunnableConfig
+    ) -> AsyncIterator[Self]:
+        async with self.afrom_checkpoint(checkpoint, config) as value:
+            value.key = self.key
+            yield value
+
     # state methods
 
     @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,122 +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("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/dynamic_barrier_value.py

@@ -69,7 +69,7 @@ class DynamicBarrierValue(
         if wait_for_names := [v for v in values if isinstance(v, WaitForNames)]:
             if len(wait_for_names) > 1:
                 raise InvalidUpdateError(
-                    "Received multiple WaitForNames updates in the same step."
+                    f"At key '{self.key}': Received multiple WaitForNames updates in the same step."
                 )
             self.names = wait_for_names[0].names
             return True

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
@@ -58,7 +52,7 @@ class EphemeralValue(Generic[Value], BaseChannel[Value, Value, Value]):
                 return False
         if len(values) != 1 and self.guard:
             raise InvalidUpdateError(
-                "EphemeralValue can only receive one value per step."
+                f"At key '{self.key}': EphemeralValue(guard=True) can receive only one value per step. Use guard=False if you want to store any one of multiple values."
             )
 
         self.value = values[-1]

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
@@ -52,7 +46,9 @@ class LastValue(Generic[Value], BaseChannel[Value, Value, Value]):
         if len(values) == 0:
             return False
         if len(values) != 1:
-            raise InvalidUpdateError("LastValue can only receive one value per step.")
+            raise InvalidUpdateError(
+                f"At key '{self.key}': Can receive only one value per step. Use an Annotated key to handle multiple values."
+            )
 
         self.value = values[-1]
         return True

libs/langgraph/langgraph/channels/manager.py

@@ -1,39 +0,0 @@
-from contextlib import AsyncExitStack, ExitStack, asynccontextmanager, contextmanager
-from typing import AsyncGenerator, Generator, Mapping
-
-from langchain_core.runnables import RunnableConfig
-
-from langgraph.channels.base import BaseChannel
-from langgraph.checkpoint.base import Checkpoint
-
-
-@contextmanager
-def ChannelsManager(
-    channels: Mapping[str, BaseChannel],
-    checkpoint: Checkpoint,
-    config: RunnableConfig,
-) -> Generator[Mapping[str, BaseChannel], None, None]:
-    """Manage channels for the lifetime of a Pregel invocation (multiple steps)."""
-    with ExitStack() as stack:
-        yield {
-            k: stack.enter_context(
-                v.from_checkpoint(checkpoint["channel_values"].get(k), config)
-            )
-            for k, v in channels.items()
-        }
-
-
-@asynccontextmanager
-async def AsyncChannelsManager(
-    channels: Mapping[str, BaseChannel],
-    checkpoint: Checkpoint,
-    config: RunnableConfig,
-) -> AsyncGenerator[Mapping[str, BaseChannel], None]:
-    """Manage channels for the lifetime of a Pregel invocation (multiple steps)."""
-    async with AsyncExitStack() as stack:
-        yield {
-            k: await stack.enter_async_context(
-                v.afrom_checkpoint(checkpoint["channel_values"].get(k), config)
-            )
-            for k, v in channels.items()
-        }

libs/langgraph/langgraph/channels/named_barrier_value.py

@@ -53,7 +53,9 @@ class NamedBarrierValue(Generic[Value], BaseChannel[Value, Value, set[Value]]):
                     self.seen.add(value)
                     updated = True
             else:
-                raise InvalidUpdateError(f"Value {value} not in {self.names}")
+                raise InvalidUpdateError(
+                    f"At key '{self.key}': Value {value} not in {self.names}"
+                )
         return updated
 
     def get(self) -> Value:

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/channels/untracked_value.py

@@ -49,7 +49,7 @@ class UntrackedValue(Generic[Value], BaseChannel[Value, Value, Value]):
             return False
         if len(values) != 1 and self.guard:
             raise InvalidUpdateError(
-                "UntrackedValue can only receive one value per step."
+                f"At key '{self.key}': UntrackedValue(guard=True) can receive only one value per step. Use guard=False if you want to store any one of multiple values."
             )
 
         self.value = values[-1]

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 (
@@ -102,5 +110,5 @@ class Send:
 
 @dataclass
 class Interrupt:
-    when: Literal["before", "during", "after"]
-    value: Any = None
+    value: Any
+    when: Literal["during"] = "during"

libs/langgraph/langgraph/errors.py

@@ -1,4 +1,4 @@
-from typing import Any
+from typing import Any, Sequence
 
 from langgraph.checkpoint.base import EmptyChannelError
 from langgraph.constants import Interrupt
@@ -32,7 +32,7 @@ class InvalidUpdateError(Exception):
 class GraphInterrupt(Exception):
     """Raised when a subgraph is interrupted."""
 
-    def __init__(self, interrupts: list[Interrupt]) -> None:
+    def __init__(self, interrupts: Sequence[Interrupt] = ()) -> None:
         super().__init__(interrupts)
 
 
@@ -40,7 +40,7 @@ class NodeInterrupt(GraphInterrupt):
     """Raised by a node to interrupt execution."""
 
     def __init__(self, value: Any) -> None:
-        super().__init__([Interrupt("during", value)])
+        super().__init__([Interrupt(value)])
 
 
 class EmptyInputError(Exception):

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,
@@ -140,8 +141,7 @@ class Graph:
         node: RunnableLike,
         *,
         metadata: Optional[dict[str, Any]] = None,
-    ) -> None:
-        ...
+    ) -> None: ...
 
     @overload
     def add_node(
@@ -150,8 +150,7 @@ class Graph:
         action: RunnableLike,
         *,
         metadata: Optional[dict[str, Any]] = None,
-    ) -> None:
-        ...
+    ) -> None: ...
 
     def add_node(
         self,
@@ -160,10 +159,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(
@@ -192,12 +196,14 @@ class Graph:
             raise ValueError("END cannot be a start node")
         if end_key == START:
             raise ValueError("START cannot be an end node")
-        if not self.support_multiple_edges and start_key in set(
+
+        # run this validation only for non-StateGraph graphs
+        if not hasattr(self, "channels") and start_key in set(
             start for start, _ in self.edges
         ):
             raise ValueError(
                 f"Already found path for node '{start_key}'.\n"
-                "For multiple edges, use StateGraph with an annotated state key."
+                "For multiple edges, use StateGraph with an Annotated state key."
             )
 
         self.edges.add((start_key, end_key))

libs/langgraph/langgraph/graph/state.py

@@ -1,3 +1,4 @@
+import inspect
 import logging
 import typing
 import warnings
@@ -5,6 +6,7 @@ from functools import partial
 from inspect import isclass, isfunction, signature
 from typing import (
     Any,
+    Callable,
     NamedTuple,
     Optional,
     Sequence,
@@ -24,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,
@@ -44,7 +49,7 @@ from langgraph.managed.base import (
     ChannelKeyPlaceholder,
     ChannelTypePlaceholder,
     ConfiguredManagedValue,
-    ManagedValue,
+    ManagedValueSpec,
     is_managed_value,
     is_writable_managed_value,
 )
@@ -129,8 +134,8 @@ class StateGraph(Graph):
 
     nodes: dict[str, StateNodeSpec]
     channels: dict[str, BaseChannel]
-    managed: dict[str, Type[ManagedValue]]
-    schemas: dict[Type[Any], dict[str, Union[BaseChannel, Type[ManagedValue]]]]
+    managed: dict[str, ManagedValueSpec]
+    schemas: dict[Type[Any], dict[str, Union[BaseChannel, ManagedValueSpec]]]
 
     def __init__(
         self,
@@ -192,10 +197,6 @@ class StateGraph(Graph):
                         )
                 else:
                     self.managed[key] = managed
-            if any(
-                isinstance(c, BinaryOperatorAggregate) for c in self.channels.values()
-            ):
-                self.support_multiple_edges = True
 
     @overload
     def add_node(
@@ -320,20 +321,27 @@ 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 (
                 hints := get_type_hints(action.__call__) or get_type_hints(action)
             ):
                 if input is None:
-                    input_hint = hints[list(hints.keys())[0]]
-                    if isinstance(input_hint, type) and get_type_hints(input_hint):
-                        input = input_hint
-        except TypeError:
+                    first_parameter_name = next(
+                        iter(inspect.signature(action).parameters.keys())
+                    )
+                    if input_hint := hints.get(first_parameter_name):
+                        if isinstance(input_hint, type) and get_type_hints(input_hint):
+                            input = input_hint
+        except (TypeError, StopIteration):
             pass
         if input is not None:
             self._add_schema(input)
@@ -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)
             ]
         )
 
@@ -442,7 +448,11 @@ class StateGraph(Graph):
             builder=self,
             config_type=self.config_schema,
             nodes={},
-            channels={**self.channels, START: EphemeralValue(self.input)},
+            channels={
+                **self.channels,
+                **self.managed,
+                START: EphemeralValue(self.input),
+            },
             input_channels=START,
             stream_mode="updates",
             output_channels=output_channels,
@@ -497,8 +507,7 @@ class CompiledStateGraph(CompiledGraph):
                     **{
                         k: (self.channels[k].UpdateType, None)
                         for k in self.builder.schemas[self.builder.input]
-                        if k in self.channels
-                        and not isinstance(self.channels[k], Context)
+                        if isinstance(self.channels[k], BaseChannel)
                     },
                 )
 
@@ -519,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) + [
@@ -572,10 +581,7 @@ class CompiledStateGraph(CompiledGraph):
             )
         else:
             input_schema = node.input if node else self.builder.schema
-            input_values = {
-                k: v if is_managed_value(v) else k
-                for k, v in self.builder.schemas[input_schema].items()
-            }
+            input_values = {k: k for k in self.builder.schemas[input_schema]}
             is_single_input = len(input_values) == 1 and "__root__" in input_values
 
             self.channels[key] = EphemeralValue(Any, guard=False)
@@ -649,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 = (
@@ -675,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)),
     )
 
 
@@ -694,7 +708,7 @@ def _coerce_state(schema: Type[Any], input: dict[str, Any]) -> dict[str, Any]:
 
 def _get_channels(
     schema: Type[dict],
-) -> tuple[dict[str, BaseChannel], dict[str, Type[ManagedValue]]]:
+) -> tuple[dict[str, BaseChannel], dict[str, ManagedValueSpec]]:
     if not hasattr(schema, "__annotations__"):
         return {"__root__": _get_channel("__root__", schema, allow_managed=False)}, {}
 
@@ -711,17 +725,22 @@ def _get_channels(
 
 def _get_channel(
     name: str, annotation: Any, *, allow_managed: bool = True
-) -> Union[BaseChannel, Type[ManagedValue]]:
+) -> Union[BaseChannel, ManagedValueSpec]:
     if manager := _is_field_managed_value(name, annotation):
         if allow_managed:
             return manager
         else:
             raise ValueError(f"This {annotation} not allowed in this position")
     elif channel := _is_field_channel(annotation):
+        channel.key = name
         return channel
     elif channel := _is_field_binop(annotation):
+        channel.key = name
         return channel
-    return LastValue(annotation)
+
+    fallback = LastValue(annotation)
+    fallback.key = name
+    return fallback
 
 
 def _is_field_channel(typ: Type[Any]) -> Optional[BaseChannel]:
@@ -751,7 +770,7 @@ def _is_field_binop(typ: Type[Any]) -> Optional[BinaryOperatorAggregate]:
     return None
 
 
-def _is_field_managed_value(name: str, typ: Type[Any]) -> Optional[Type[ManagedValue]]:
+def _is_field_managed_value(name: str, typ: Type[Any]) -> Optional[ManagedValueSpec]:
     if hasattr(typ, "__metadata__"):
         meta = typ.__metadata__
         if len(meta) >= 1:

libs/langgraph/langgraph/managed/base.py

@@ -1,12 +1,9 @@
-import asyncio
 from abc import ABC, abstractmethod
-from contextlib import AsyncExitStack, ExitStack, asynccontextmanager, contextmanager
+from contextlib import asynccontextmanager, contextmanager
 from inspect import isclass
 from typing import (
     Any,
-    AsyncGenerator,
     AsyncIterator,
-    Generator,
     Generic,
     Iterator,
     NamedTuple,
@@ -19,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
 
@@ -56,18 +58,15 @@ class ManagedValue(ABC, Generic[V]):
                 pass
 
     @abstractmethod
-    def __call__(self, step: int) -> V:
-        ...
+    def __call__(self, step: int) -> V: ...
 
 
 class WritableManagedValue(Generic[V, U], ManagedValue[V], ABC):
     @abstractmethod
-    def update(self, writes: Sequence[U]) -> None:
-        ...
+    def update(self, writes: Sequence[U]) -> None: ...
 
     @abstractmethod
-    async def aupdate(self, writes: Sequence[U]) -> None:
-        ...
+    async def aupdate(self, writes: Sequence[U]) -> None: ...
 
 
 class ConfiguredManagedValue(NamedTuple):
@@ -77,8 +76,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(
@@ -104,50 +101,47 @@ def is_writable_managed_value(value: Any) -> TypeGuard[Type[WritableManagedValue
     )
 
 
-@contextmanager
-def ManagedValuesManager(
-    values: dict[str, ManagedValueSpec],
-    config: RunnableConfig,
-) -> Generator[ManagedValueMapping, None, None]:
-    if values:
-        with ExitStack() as stack:
-            yield {
-                key: stack.enter_context(
-                    value.cls.enter(config, **value.kwargs)
-                    if isinstance(value, ConfiguredManagedValue)
-                    else value.enter(config)
-                )
-                for key, value in values.items()
-            }
-    else:
-        yield {}
-
-
-@asynccontextmanager
-async def AsyncManagedValuesManager(
-    values: dict[str, ManagedValueSpec],
-    config: RunnableConfig,
-) -> AsyncGenerator[ManagedValueMapping, None]:
-    if values:
-        async with AsyncExitStack() as stack:
-            # create enter tasks with reference to spec
-            tasks = {
-                asyncio.create_task(
-                    stack.enter_async_context(
-                        value.cls.aenter(config, **value.kwargs)
-                        if isinstance(value, ConfiguredManagedValue)
-                        else value.aenter(config)
-                    )
-                ): key
-                for key, value in values.items()
-            }
-            # wait for all enter tasks
-            done, _ = await asyncio.wait(tasks, return_when=asyncio.ALL_COMPLETED)
-            # build mapping from spec to result
-            yield {tasks[task]: task.result() for task in done}
-    else:
-        yield {}
-
-
 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/__init__.py

@@ -1,4 +1,5 @@
 """langgraph.prebuilt exposes a higher-level API for creating and executing agents and tools."""
+
 from langgraph.prebuilt.chat_agent_executor import create_react_agent
 from langgraph.prebuilt.tool_executor import ToolExecutor, ToolInvocation
 from langgraph.prebuilt.tool_node import InjectedState, ToolNode, tools_condition