sdk-py0.1.29

Override with_config to store config in Pregel instance

README.md

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

docs/_scripts/copy_notebooks.py

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

docs/docs/cloud/deployment/graph_rebuild.md

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

docs/docs/cloud/deployment/setup.md

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

docs/docs/cloud/deployment/setup_pyproject.md

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

docs/docs/cloud/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/how-tos/index.md

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

docs/mkdocs.yml

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

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

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

examples/chatbots/information-gather-prompting.ipynb

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

examples/create-react-agent-hitl.ipynb

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

examples/customer-support/customer-support.ipynb

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

examples/human_in_the_loop/review-tool-calls.ipynb

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

examples/input_output_schema.ipynb

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

examples/llm-compiler/LLMCompiler.ipynb

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

examples/many-tools.ipynb

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

examples/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/reflection/reflection.ipynb

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

examples/reflexion/reflexion.ipynb

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

examples/streaming-content.ipynb

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

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

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

examples/streaming-tokens-without-langchain.ipynb

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

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

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

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

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

examples/web-navigation/web_voyager.ipynb

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

libs/checkpoint-postgres/README.md

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

@@ -66,7 +66,7 @@ class PostgresSaver(BasePostgresSaver):
         the first time checkpointer is used.
         """
         with self.lock:
-            with self.conn.cursor(binary=True) as cur:
+            with self.conn.cursor(binary=True, row_factory=dict_row) as cur:
                 try:
                     version = cur.execute(
                         "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
@@ -127,30 +127,33 @@ 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"]),
+                        "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"],
+                        }
                     }
-                },
-                {
-                    **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 +201,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,
@@ -317,16 +320,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(
@@ -345,7 +338,7 @@ class PostgresSaver(BasePostgresSaver):
             # in multiple threads/coroutines, but only one cursor can be
             # used at a time
             try:
-                with self.conn.cursor(binary=True) as cur:
+                with self.conn.cursor(binary=True, row_factory=dict_row) as cur:
                     yield cur
             finally:
                 if pipeline:
@@ -353,8 +346,10 @@ class PostgresSaver(BasePostgresSaver):
         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:
+            with self.lock, self.conn.pipeline(), self.conn.cursor(
+                binary=True, row_factory=dict_row
+            ) as cur:
                 yield cur
         else:
-            with self.lock, self.conn.cursor(binary=True) as cur:
+            with self.lock, self.conn.cursor(binary=True, row_factory=dict_row) as cur:
                 yield cur

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

@@ -64,7 +64,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
         the first time checkpointer is used.
         """
         async with self.lock:
-            async with self.conn.cursor(binary=True) as cur:
+            async with self.conn.cursor(binary=True, row_factory=dict_row) as cur:
                 try:
                     results = await cur.execute(
                         "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
@@ -110,32 +110,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"],
+        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"],
+                        }
+                    },
+                    {
+                        **self._load_checkpoint(value["checkpoint"]),
+                        "channel_values": await asyncio.to_thread(
+                            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"],
+                        }
                     }
-                },
-                {
-                    **self._load_checkpoint(value["checkpoint"]),
-                    "channel_values": await asyncio.to_thread(
-                        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,
+                    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 +165,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,
@@ -273,16 +276,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(
@@ -302,7 +295,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
             # 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 self.conn.cursor(binary=True, row_factory=dict_row) as cur:
                     yield cur
             finally:
                 if pipeline:
@@ -311,9 +304,11 @@ class AsyncPostgresSaver(BasePostgresSaver):
             # 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
+                binary=True, row_factory=dict_row
             ) as cur:
                 yield cur
         else:
-            async with self.lock, self.conn.cursor(binary=True) as cur:
+            async with self.lock, self.conn.cursor(
+                binary=True, row_factory=dict_row
+            ) as cur:
                 yield cur

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

@@ -6,6 +6,7 @@ from langchain_core.runnables import RunnableConfig
 from psycopg.types.json import Jsonb
 
 from langgraph.checkpoint.base import (
+    WRITES_IDX_MAP,
     BaseCheckpointSaver,
     Checkpoint,
     EmptyChannelError,
@@ -105,15 +106,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,7 +113,6 @@ 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()
 
@@ -210,7 +201,7 @@ class BasePostgresSaver(BaseCheckpointSaver):
                 checkpoint_ns,
                 checkpoint_id,
                 task_id,
-                idx,
+                WRITES_IDX_MAP.get(channel, idx),
                 channel,
                 *self.serde.dumps_typed(value),
             )

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/README.md

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

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

@@ -22,6 +22,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 +97,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 +108,6 @@ def empty_checkpoint() -> Checkpoint:
         channel_versions={},
         versions_seen={},
         pending_sends=[],
-        current_tasks={},
     )
 
 
@@ -122,7 +120,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 +137,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 +151,6 @@ def create_checkpoint(
         channel_versions=checkpoint["channel_versions"],
         versions_seen=checkpoint["versions_seen"],
         pending_sends=checkpoint.get("pending_sends", []),
-        current_tasks={},
     )
 
 
@@ -437,3 +435,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,
@@ -52,6 +53,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 +64,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,7 +107,7 @@ 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()
                 return CheckpointTuple(
                     config=config,
                     checkpoint=self.serde.loads_typed(checkpoint),
@@ -125,7 +129,7 @@ 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()
                 return CheckpointTuple(
                     config={
                         "configurable": {
@@ -204,6 +208,8 @@ class MemorySaver(
                 elif limit is not None:
                     limit -= 1
 
+                writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)].values()
+
                 yield CheckpointTuple(
                     config={
                         "configurable": {
@@ -223,6 +229,9 @@ class MemorySaver(
                     }
                     if parent_checkpoint_id
                     else None,
+                    pending_writes=[
+                        (id, c, self.serde.loads_typed(v)) for id, c, v in writes
+                    ],
                 )
 
     def put(
@@ -287,11 +296,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.

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):
@@ -111,7 +116,7 @@ class JsonPlusSerializer(SerializerProtocol):
                 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,8 @@ from typing import (
 from langchain_core.runnables import RunnableConfig
 from typing_extensions import Self
 
+ERROR = "__error__"
+
 Value = TypeVar("Value")
 Update = TypeVar("Update")
 C = TypeVar("C")

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"

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/cli/pyproject.toml

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

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/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:
@@ -43,19 +45,35 @@ class BaseChannel(Generic[Value, Update, C], ABC):
     @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/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

@@ -58,7 +58,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

@@ -52,7 +52,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/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/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

@@ -11,6 +11,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 +23,7 @@ RESERVED = {
     CONFIG_KEY_RESUMING,
     CONFIG_KEY_TASK_ID,
     INPUT,
+    RUNTIME_PLACEHOLDER,
 }
 TAG_HIDDEN = "langsmith:hidden"
 
@@ -102,5 +104,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

@@ -192,12 +192,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,7 +26,6 @@ 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
@@ -192,10 +193,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(
@@ -330,10 +327,13 @@ class StateGraph(Graph):
                 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 +374,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 +425,14 @@ class StateGraph(Graph):
             else [
                 key
                 for key, val in self.schemas[self.output].items()
-                if not isinstance(val, Context) and not is_managed_value(val)
+                if not is_managed_value(val)
             ]
         )
         stream_channels = (
             "__root__"
             if len(self.channels) == 1 and "__root__" in self.channels
             else [
-                key
-                for key, val in self.channels.items()
-                if not isinstance(val, Context) and not is_managed_value(val)
+                key for key, val in self.channels.items() if not is_managed_value(val)
             ]
         )
 
@@ -502,7 +500,6 @@ class CompiledStateGraph(CompiledGraph):
                         k: (self.channels[k].UpdateType, None)
                         for k in self.builder.schemas[self.builder.input]
                         if isinstance(self.channels[k], BaseChannel)
-                        and not isinstance(self.channels[k], Context)
                     },
                 )
 
@@ -523,7 +520,7 @@ class CompiledStateGraph(CompiledGraph):
             output_keys = [
                 k
                 for k, v in self.builder.schemas[self.builder.input].items()
-                if not isinstance(v, Context) and not is_managed_value(v)
+                if not is_managed_value(v)
             ]
         else:
             output_keys = list(self.builder.channels) + [
@@ -650,7 +647,14 @@ class CompiledStateGraph(CompiledGraph):
                 return ChannelWrite(writes, tags=[TAG_HIDDEN])
 
         # attach branch publisher
-        self.nodes[start] |= branch.run(branch_writer, _get_state_reader(self.builder))
+        schema = (
+            self.builder.nodes[start].input
+            if start in self.builder.nodes
+            else self.builder.schema
+        )
+        self.nodes[start] |= branch.run(
+            branch_writer, _get_state_reader(self.builder, schema)
+        )
 
         # attach branch subscribers
         ends = (
@@ -676,16 +680,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)),
     )
 
 
@@ -719,10 +724,15 @@ def _get_channel(
         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]:

libs/langgraph/langgraph/managed/base.py

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

libs/langgraph/langgraph/managed/context.py

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

libs/langgraph/langgraph/managed/shared_value.py

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

libs/langgraph/langgraph/prebuilt/chat_agent_executor.py

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

libs/langgraph/langgraph/pregel/__init__.py

@@ -38,6 +38,7 @@ from langchain_core.runnables.config import (
     ensure_config,
     get_async_callback_manager_for_config,
     get_callback_manager_for_config,
+    merge_configs,
     patch_config,
 )
 from langchain_core.runnables.utils import (
@@ -51,7 +52,6 @@ from typing_extensions import Self
 from langgraph.channels.base import (
     BaseChannel,
 )
-from langgraph.channels.context import Context
 from langgraph.checkpoint.base import (
     BaseCheckpointSaver,
     copy_checkpoint,
@@ -65,15 +65,14 @@ from langgraph.constants import (
     CONFIG_KEY_SEND,
     ERROR,
     INTERRUPT,
-    Interrupt,
 )
 from langgraph.errors import GraphInterrupt, GraphRecursionError, InvalidUpdateError
 from langgraph.managed.base import ManagedValueSpec
 from langgraph.pregel.algo import (
     apply_writes,
     local_read,
+    local_write,
     prepare_next_tasks,
-    should_interrupt,
 )
 from langgraph.pregel.debug import (
     print_step_checkpoint,
@@ -221,11 +220,18 @@ class Pregel(
 
     config_type: Optional[Type[Any]] = None
 
+    config: Optional[RunnableConfig] = None
+
     name: str = "LangGraph"
 
     class Config:
         arbitrary_types_allowed = True
 
+    def with_config(self, config: RunnableConfig | None = None, **kwargs: Any) -> Self:
+        return self.copy(
+            update={"config": cast(RunnableConfig, {**(config or {}), **kwargs})}
+        )
+
     @classmethod
     def is_lc_serializable(cls) -> bool:
         """Return whether the graph can be serialized by Langchain."""
@@ -299,6 +305,7 @@ class Pregel(
     def get_input_schema(
         self, config: Optional[RunnableConfig] = None
     ) -> Type[BaseModel]:
+        config = merge_configs(self.config, config)
         if isinstance(self.input_channels, str):
             return super().get_input_schema(config)
         else:
@@ -318,6 +325,7 @@ class Pregel(
     def get_output_schema(
         self, config: Optional[RunnableConfig] = None
     ) -> Type[BaseModel]:
+        config = merge_configs(self.config, config)
         if isinstance(self.output_channels, str):
             return super().get_output_schema(config)
         else:
@@ -336,10 +344,7 @@ class Pregel(
     @property
     def stream_channels_asis(self) -> Union[str, Sequence[str]]:
         return self.stream_channels or [
-            k
-            for k in self.channels
-            if isinstance(self.channels[k], BaseChannel)
-            and not isinstance(self.channels[k], Context)
+            k for k in self.channels if isinstance(self.channels[k], BaseChannel)
         ]
 
     def get_state(self, config: RunnableConfig) -> StateSnapshot:
@@ -347,6 +352,7 @@ class Pregel(
         if not self.checkpointer:
             raise ValueError("No checkpointer set")
 
+        config = merge_configs(self.config, config) if self.config else config
         saved = self.checkpointer.get_tuple(config)
         checkpoint = saved.checkpoint if saved else empty_checkpoint()
         config = saved.config if saved else config
@@ -379,6 +385,7 @@ class Pregel(
         if not self.checkpointer:
             raise ValueError("No checkpointer set")
 
+        config = merge_configs(self.config, config) if self.config else config
         saved = await self.checkpointer.aget_tuple(config)
         checkpoint = saved.checkpoint if saved else empty_checkpoint()
 
@@ -427,7 +434,12 @@ class Pregel(
             metadata,
             parent_config,
             pending_writes,
-        ) in self.checkpointer.list(config, before=before, limit=limit, filter=filter):
+        ) in self.checkpointer.list(
+            merge_configs(self.config, config) if self.config else config,
+            before=before,
+            limit=limit,
+            filter=filter,
+        ):
             with ChannelsManager(
                 self.channels, checkpoint, config, skip_context=True
             ) as (channels, managed):
@@ -472,7 +484,12 @@ class Pregel(
             metadata,
             parent_config,
             pending_writes,
-        ) in self.checkpointer.alist(config, before=before, limit=limit, filter=filter):
+        ) in self.checkpointer.alist(
+            merge_configs(self.config, config) if self.config else config,
+            before=before,
+            limit=limit,
+            filter=filter,
+        ):
             async with AsyncChannelsManager(
                 self.channels, checkpoint, config, skip_context=True
             ) as (channels, managed):
@@ -509,6 +526,7 @@ class Pregel(
             raise ValueError("No checkpointer set")
 
         # get last checkpoint
+        config = merge_configs(self.config, config) if self.config else config
         saved = self.checkpointer.get_tuple(config)
         checkpoint = copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint()
         checkpoint_previous_versions = (
@@ -594,9 +612,22 @@ class Pregel(
                     run_name=self.name + "UpdateState",
                     configurable={
                         # deque.extend is thread-safe
-                        CONFIG_KEY_SEND: task.writes.extend,
+                        CONFIG_KEY_SEND: partial(
+                            local_write,
+                            step + 1,
+                            task.writes.extend,
+                            self.nodes,
+                            channels,
+                            managed,
+                        ),
                         CONFIG_KEY_READ: partial(
-                            local_read, checkpoint, channels, task, config
+                            local_read,
+                            step + 1,
+                            checkpoint,
+                            channels,
+                            managed,
+                            task,
+                            config,
                         ),
                     },
                 ),
@@ -606,31 +637,6 @@ class Pregel(
                 checkpoint, channels, [task], self.checkpointer.get_next_version
             ), "Can't write to SharedValues from update_state"
             checkpoint = create_checkpoint(checkpoint, channels, step + 1)
-            # check interrupt before
-            if tasks := should_interrupt(
-                checkpoint,
-                self.interrupt_before_nodes,
-                prepare_next_tasks(
-                    checkpoint,
-                    self.nodes,
-                    channels,
-                    managed,
-                    config,
-                    step + 2,
-                    for_execution=False,
-                ),
-            ):
-                for t in tasks:
-                    self.checkpointer.put_writes(
-                        {
-                            "configurable": {
-                                **checkpoint_config["configurable"],
-                                "checkpoint_id": checkpoint["id"],
-                            }
-                        },
-                        [(INTERRUPT, Interrupt("before"))],
-                        t.id,
-                    )
             return self.checkpointer.put(
                 checkpoint_config,
                 checkpoint,
@@ -654,6 +660,7 @@ class Pregel(
             raise ValueError("No checkpointer set")
 
         # get last checkpoint
+        config = merge_configs(self.config, config) if self.config else config
         saved = await self.checkpointer.aget_tuple(config)
         checkpoint = copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint()
         checkpoint_previous_versions = (
@@ -737,9 +744,22 @@ class Pregel(
                     run_name=self.name + "UpdateState",
                     configurable={
                         # deque.extend is thread-safe
-                        CONFIG_KEY_SEND: task.writes.extend,
+                        CONFIG_KEY_SEND: partial(
+                            local_write,
+                            step + 1,
+                            task.writes.extend,
+                            self.nodes,
+                            channels,
+                            managed,
+                        ),
                         CONFIG_KEY_READ: partial(
-                            local_read, checkpoint, channels, task, config
+                            local_read,
+                            step + 1,
+                            checkpoint,
+                            channels,
+                            managed,
+                            task,
+                            config,
                         ),
                     },
                 ),
@@ -749,35 +769,6 @@ class Pregel(
                 checkpoint, channels, [task], self.checkpointer.get_next_version
             ), "Can't write to SharedValues from update_state"
             checkpoint = create_checkpoint(checkpoint, channels, step + 1)
-            # check interrupt before
-            if tasks := should_interrupt(
-                checkpoint,
-                self.interrupt_before_nodes,
-                prepare_next_tasks(
-                    checkpoint,
-                    self.nodes,
-                    channels,
-                    managed,
-                    config,
-                    step + 2,
-                    for_execution=False,
-                ),
-            ):
-                await asyncio.gather(
-                    *(
-                        self.checkpointer.aput_writes(
-                            {
-                                "configurable": {
-                                    **checkpoint_config["configurable"],
-                                    "checkpoint_id": checkpoint["id"],
-                                }
-                            },
-                            [(INTERRUPT, Interrupt("before"))],
-                            t.id,
-                        )
-                        for t in tasks
-                    )
-                )
             return await self.checkpointer.aput(
                 checkpoint_config,
                 checkpoint,
@@ -918,7 +909,7 @@ class Pregel(
             {'type': 'task_result', 'timestamp': '2024-06-23T...+00:00', 'step': 2, 'payload': {'id': '...', 'name': 'b', 'result': [('alist', ['there'])]}}
             ```
         """
-        config = ensure_config(config)
+        config = ensure_config(merge_configs(self.config, config))
         callback_manager = get_callback_manager_for_config(config)
         run_manager = callback_manager.on_chain_start(
             dumpd(self),
@@ -1158,7 +1149,7 @@ class Pregel(
             {'type': 'task_result', 'timestamp': '2024-06-23T...+00:00', 'step': 2, 'payload': {'id': '...', 'name': 'b', 'result': [('alist', ['there'])]}}
             ```
         """
-        config = ensure_config(config)
+        config = ensure_config(merge_configs(self.config, config))
         callback_manager = get_async_callback_manager_for_config(config)
         run_manager = await callback_manager.on_chain_start(
             dumpd(self),

libs/langgraph/langgraph/pregel/algo.py

@@ -24,7 +24,6 @@ from langchain_core.runnables.config import (
 )
 
 from langgraph.channels.base import BaseChannel
-from langgraph.channels.context import Context
 from langgraph.checkpoint.base import (
     BaseCheckpointSaver,
     Checkpoint,
@@ -95,28 +94,37 @@ def should_interrupt(
 
 
 def local_read(
+    step: int,
     checkpoint: Checkpoint,
     channels: Mapping[str, BaseChannel],
+    managed: ManagedValueMapping,
     task: WritesProtocol,
     config: RunnableConfig,
     select: Union[list[str], str],
     fresh: bool = False,
 ) -> Union[dict[str, Any], Any]:
+    if isinstance(select, str):
+        managed_keys = []
+    else:
+        managed_keys = [k for k in select if k in managed]
+        select = [k for k in select if k not in managed]
     if fresh:
         new_checkpoint = create_checkpoint(copy_checkpoint(checkpoint), channels, -1)
-        context_channels = {k: v for k, v in channels.items() if isinstance(v, Context)}
         with ChannelsManager(channels, new_checkpoint, config, skip_context=True) as (
             channels,
             _,
         ):
-            all_channels = {**channels, **context_channels}
-            apply_writes(new_checkpoint, all_channels, [task], None)
-            return read_channels(all_channels, select)
+            apply_writes(new_checkpoint, channels, [task], None)
+            values = read_channels(channels, select)
     else:
-        return read_channels(channels, select)
+        values = read_channels(channels, select)
+    if managed_keys:
+        values.update({k: managed[k](step) for k in managed_keys})
+    return values
 
 
 def local_write(
+    step: int,
     commit: Callable[[Sequence[tuple[str, Any]]], None],
     processes: Mapping[str, PregelNode],
     channels: Mapping[str, BaseChannel],
@@ -131,6 +139,8 @@ def local_write(
                 )
             if value.node not in processes:
                 raise InvalidUpdateError(f"Invalid node name {value.node} in packet")
+            # replace any runtime values with placeholders
+            managed.replace_runtime_values(step, value.arg)
         elif chan not in channels and chan not in managed:
             logger.warning(f"Skipping write for channel '{chan}' which has no readers")
     commit(writes)
@@ -198,12 +208,7 @@ def apply_writes(
     updated_channels: set[str] = set()
     for chan, vals in pending_writes_by_channel.items():
         if chan in channels:
-            try:
-                updated = channels[chan].update(vals)
-            except InvalidUpdateError as e:
-                raise InvalidUpdateError(
-                    f"Invalid update for channel {chan} with values {vals}"
-                ) from e
+            updated = channels[chan].update(vals)
             if updated and get_next_version is not None:
                 checkpoint["channel_versions"][chan] = get_next_version(
                     max_version, channels[chan]
@@ -298,6 +303,7 @@ def prepare_next_tasks(
         if for_execution:
             proc = processes[packet.node]
             if node := proc.get_node():
+                managed.replace_runtime_placeholders(step, packet.arg)
                 writes = deque()
                 tasks.append(
                     PregelExecutableTask(
@@ -322,6 +328,7 @@ def prepare_next_tasks(
                                 # deque.extend is thread-safe
                                 CONFIG_KEY_SEND: partial(
                                     local_write,
+                                    step,
                                     writes.extend,
                                     processes,
                                     channels,
@@ -329,8 +336,10 @@ def prepare_next_tasks(
                                 ),
                                 CONFIG_KEY_READ: partial(
                                     local_read,
+                                    step,
                                     checkpoint,
                                     channels,
+                                    managed,
                                     PregelTaskWrites(packet.node, writes, triggers),
                                     config,
                                 ),
@@ -417,6 +426,7 @@ def prepare_next_tasks(
                                     # deque.extend is thread-safe
                                     CONFIG_KEY_SEND: partial(
                                         local_write,
+                                        step,
                                         writes.extend,
                                         processes,
                                         channels,
@@ -424,8 +434,10 @@ def prepare_next_tasks(
                                     ),
                                     CONFIG_KEY_READ: partial(
                                         local_read,
+                                        step,
                                         checkpoint,
                                         channels,
+                                        managed,
                                         PregelTaskWrites(name, writes, triggers),
                                         config,
                                     ),

libs/langgraph/langgraph/pregel/executor.py

@@ -120,14 +120,14 @@ class AsyncBackgroundExecutor(AsyncContextManager):
 
     def done(self, task: asyncio.Task) -> None:
         try:
-            task.result()
-        except GraphInterrupt:
-            # This exception is an interruption signal, not an error
-            # so we don't want to re-raise it on exit
-            self.tasks.pop(task)
-        except BaseException:
-            pass
-        else:
+            if exc := task.exception():
+                # This exception is an interruption signal, not an error
+                # so we don't want to re-raise it on exit
+                if isinstance(exc, GraphInterrupt):
+                    self.tasks.pop(task)
+            else:
+                self.tasks.pop(task)
+        except asyncio.CancelledError:
             self.tasks.pop(task)
 
     async def __aenter__(self) -> Submit:
@@ -146,12 +146,13 @@ class AsyncBackgroundExecutor(AsyncContextManager):
         # wait for all tasks to finish
         if self.tasks:
             await asyncio.wait(self.tasks)
-        # re-raise the first exception that occurred in a task
+        # if there's already an exception being raised, don't raise another one
         if exc_type is None:
-            # if there's already an exception being raised, don't raise another one
+            # re-raise the first exception that occurred in a task
             for task in self.tasks:
                 try:
-                    task.result()
+                    if exc := task.exception():
+                        raise exc
                 except asyncio.CancelledError:
                     pass
 

libs/langgraph/langgraph/pregel/loop.py

@@ -41,7 +41,6 @@ from langgraph.constants import (
     ERROR,
     INPUT,
     INTERRUPT,
-    Interrupt,
 )
 from langgraph.errors import EmptyInputError, GraphInterrupt
 from langgraph.managed.base import (
@@ -156,12 +155,10 @@ class PregelLoop:
         self.stream_keys = stream_keys
         self.is_nested = CONFIG_KEY_READ in self.config.get("configurable", {})
 
-    def mark_tasks_scheduled(self, tasks: Sequence[PregelExecutableTask]) -> None:
-        """Mark tasks as scheduled, to be used by queue-based executors."""
-        raise NotImplementedError
-
     def put_writes(self, task_id: str, writes: Sequence[tuple[str, Any]]) -> None:
         """Put writes for a task, to be read by the next tick."""
+        if not writes:
+            return
         self.checkpoint_pending_writes.extend((task_id, k, v) for k, v in writes)
         if self.checkpointer_put_writes is not None:
             self.submit(
@@ -238,13 +235,10 @@ class PregelLoop:
                 }
             )
             # after execution, check if we should interrupt
-            if tasks := should_interrupt(self.checkpoint, interrupt_after, self.tasks):
+            if should_interrupt(self.checkpoint, interrupt_after, self.tasks):
                 self.status = "interrupt_after"
-                interrupts = [(t.id, Interrupt("after")) for t in tasks]
-                for tid, interrupt in interrupts:
-                    self.put_writes(tid, [(INTERRUPT, interrupt)])
                 if self.is_nested:
-                    raise GraphInterrupt([i[1] for i in interrupts])
+                    raise GraphInterrupt()
                 else:
                     return False
         else:
@@ -308,13 +302,10 @@ class PregelLoop:
             )
 
         # before execution, check if we should interrupt
-        if tasks := should_interrupt(self.checkpoint, interrupt_before, self.tasks):
+        if should_interrupt(self.checkpoint, interrupt_before, self.tasks):
             self.status = "interrupt_before"
-            interrupts = [(t.id, Interrupt("before")) for t in tasks]
-            for tid, interrupt in interrupts:
-                self.put_writes(tid, [(INTERRUPT, interrupt)])
             if self.is_nested:
-                raise GraphInterrupt([i[1] for i in interrupts])
+                raise GraphInterrupt()
             else:
                 return False
 

libs/langgraph/langgraph/pregel/manager.py

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

libs/langgraph/langgraph/pregel/read.py

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

libs/langgraph/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph"
-version = "0.2.11"
+version = "0.2.14"
 description = "Building stateful, multi-actor applications with LLMs"
 authors = []
 license = "MIT"

libs/langgraph/tests/any_str.py

@@ -1,3 +1,6 @@
+from typing import Any, Sequence
+
+
 class AnyStr(str):
     def __init__(self) -> None:
         super().__init__()
@@ -9,16 +12,30 @@ class AnyStr(str):
         return hash(str(self))
 
 
-class ExceptionLike:
-    def __init__(self, exc: Exception) -> None:
-        self.exc = exc
+class AnyVersion:
+    def __init__(self) -> None:
+        super().__init__()
+
+    def __eq__(self, other: object) -> bool:
+        return isinstance(other, (str, int, float))
+
+    def __hash__(self) -> int:
+        return hash(str(self))
+
+
+class UnsortedSequence:
+    def __init__(self, *values: Any) -> None:
+        self.seq = values
 
     def __eq__(self, value: object) -> bool:
         return (
-            isinstance(value, Exception)
-            and self.exc.__class__ == value.__class__
-            and str(self.exc) == str(value)
+            isinstance(value, Sequence)
+            and len(self.seq) == len(value)
+            and all(a in value for a in self.seq)
         )
 
     def __hash__(self) -> int:
-        return hash((self.exc.__class__, str(self.exc)))
+        return hash(frozenset(self.seq))
+
+    def __repr__(self) -> str:
+        return repr(self.seq)

libs/langgraph/tests/memory_assert.py

@@ -24,8 +24,6 @@ class NoopSerializer(SerializerProtocol):
 
 
 class MemorySaverAssertImmutable(MemorySaver):
-    serde = NoopSerializer()
-
     storage_for_copies: defaultdict[str, dict[str, dict[str, Checkpoint]]]
 
     def __init__(
@@ -74,15 +72,6 @@ class MemorySaverAssertCheckpointMetadata(MemorySaver):
     should produce a side effect that can be asserted.
     """
 
-    serde = NoopSerializer()
-
-    def __init__(
-        self,
-        *,
-        serde: Optional[SerializerProtocol] = None,
-    ) -> None:
-        super().__init__(serde=serde)
-
     def put(
         self,
         config: RunnableConfig,

libs/langgraph/tests/test_channels.py

@@ -1,14 +1,9 @@
 import operator
-from contextlib import asynccontextmanager, contextmanager
-from typing import AsyncGenerator, Generator, Sequence, Union
+from typing import Sequence, Union
 
-import httpx
 import pytest
-from langchain_core.runnables import RunnableConfig
-from pytest_mock import MockerFixture
 
 from langgraph.channels.binop import BinaryOperatorAggregate
-from langgraph.channels.context import Context
 from langgraph.channels.last_value import LastValue
 from langgraph.channels.topic import Topic
 from langgraph.errors import EmptyChannelError, InvalidUpdateError
@@ -257,80 +252,3 @@ async def test_binop_async() -> None:
         checkpoint, {}
     ) as channel:
         assert channel.get() == 10
-
-
-def test_ctx_manager(mocker: MockerFixture) -> None:
-    setup = mocker.Mock()
-    cleanup = mocker.Mock()
-
-    @contextmanager
-    def an_int() -> Generator[int, None, None]:
-        setup()
-        try:
-            yield 5
-        finally:
-            cleanup()
-
-    with Context(an_int, None).from_checkpoint(None, {}) as channel:
-        assert setup.call_count == 1
-        assert cleanup.call_count == 0
-
-        assert channel.ValueType is None
-        assert channel.UpdateType is None
-
-        assert channel.get() == 5
-
-        with pytest.raises(InvalidUpdateError):
-            channel.update([5])  # type: ignore
-
-    assert setup.call_count == 1
-    assert cleanup.call_count == 1
-
-
-def test_ctx_manager_ctx(mocker: MockerFixture) -> None:
-    with Context(httpx.Client).from_checkpoint(None, {}) as channel:
-        assert channel.ValueType is None
-        assert channel.UpdateType is None
-
-        assert isinstance(channel.get(), httpx.Client)
-
-        with pytest.raises(InvalidUpdateError):
-            channel.update([5])  # type: ignore
-
-        with pytest.raises(EmptyChannelError):
-            channel.checkpoint()
-
-
-async def test_ctx_manager_async(mocker: MockerFixture) -> None:
-    setup = mocker.Mock()
-    cleanup = mocker.Mock()
-
-    @contextmanager
-    def an_int_sync(config: RunnableConfig) -> Generator[int, None, None]:
-        try:
-            yield 5
-        finally:
-            pass
-
-    @asynccontextmanager
-    async def an_int() -> AsyncGenerator[int, None]:
-        setup()
-        try:
-            yield 5
-        finally:
-            cleanup()
-
-    async with Context(an_int_sync, an_int).afrom_checkpoint(None, {}) as channel:
-        assert setup.call_count == 1
-        assert cleanup.call_count == 0
-
-        assert channel.ValueType is None
-        assert channel.UpdateType is None
-
-        assert channel.get() == 5
-
-        with pytest.raises(InvalidUpdateError):
-            channel.update([5])  # type: ignore
-
-    assert setup.call_count == 1
-    assert cleanup.call_count == 1

libs/langgraph/tests/test_state.py

@@ -2,10 +2,11 @@ from typing import Annotated as Annotated2
 from typing import Any
 
 import pytest
+from langchain_core.runnables import RunnableConfig
 from pydantic.v1 import BaseModel
 from typing_extensions import Annotated, TypedDict
 
-from langgraph.graph.state import _warn_invalid_state_schema
+from langgraph.graph.state import StateGraph, _warn_invalid_state_schema
 
 
 class State(BaseModel):
@@ -46,3 +47,42 @@ def test_doesnt_warn_valid_schema(schema: Any):
     # Assert the function does not raise a warning
     with pytest.warns(None):
         _warn_invalid_state_schema(schema)
+
+
+def test_state_schema_with_type_hint():
+    class InputState(TypedDict):
+        question: str
+
+    class OutputState(TypedDict):
+        input_state: InputState
+
+    def complete_hint(state: InputState) -> OutputState:
+        return {"input_state": state}
+
+    def miss_first_hint(state, config: RunnableConfig) -> OutputState:
+        return {"input_state": state}
+
+    def only_return_hint(state, config) -> OutputState:
+        return {"input_state": state}
+
+    def miss_all_hint(state, config):
+        return {"input_state": state}
+
+    graph = StateGraph(input=InputState, output=OutputState)
+    actions = [complete_hint, miss_first_hint, only_return_hint, miss_all_hint]
+
+    for action in actions:
+        graph.add_node(action)
+
+    graph.set_entry_point(actions[0].__name__)
+    for i in range(len(actions) - 1):
+        graph.add_edge(actions[i].__name__, actions[i + 1].__name__)
+    graph.set_finish_point(actions[-1].__name__)
+
+    graph = graph.compile()
+
+    input_state = InputState(question="Hello World!")
+    output_state = OutputState(input_state=input_state)
+    for i, c in enumerate(graph.stream(input_state, stream_mode="updates")):
+        node_name = actions[i].__name__
+        assert c[node_name] == output_state

libs/sdk-js/package.json

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

libs/sdk-js/src/client.mts

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

libs/sdk-js/src/schema.ts

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

libs/sdk-js/yarn.lock

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

libs/sdk-py/langgraph_sdk/client.py

@@ -25,9 +25,11 @@ from langgraph_sdk.schema import (
     Assistant,
     Config,
     Cron,
+    DisconnectMode,
     GraphSchema,
     Metadata,
     MultitaskStrategy,
+    OnCompletionBehavior,
     OnConflictBehavior,
     Run,
     RunCreate,
@@ -963,6 +965,8 @@ class RunsClient:
         interrupt_before: Optional[list[str]] = None,
         interrupt_after: Optional[list[str]] = None,
         feedback_keys: Optional[list[str]] = None,
+        on_disconnect: Optional[DisconnectMode] = None,
+        webhook: Optional[str] = None,
         multitask_strategy: Optional[MultitaskStrategy] = None,
     ) -> AsyncIterator[StreamPart]:
         ...
@@ -980,6 +984,9 @@ class RunsClient:
         interrupt_before: Optional[list[str]] = None,
         interrupt_after: Optional[list[str]] = None,
         feedback_keys: Optional[list[str]] = None,
+        on_disconnect: Optional[DisconnectMode] = None,
+        webhook: Optional[str] = None,
+        on_completion: Optional[OnCompletionBehavior] = None,
     ) -> AsyncIterator[StreamPart]:
         ...
 
@@ -996,8 +1003,10 @@ class RunsClient:
         interrupt_before: Optional[list[str]] = None,
         interrupt_after: Optional[list[str]] = None,
         feedback_keys: Optional[list[str]] = None,
+        on_disconnect: Optional[DisconnectMode] = None,
         webhook: Optional[str] = None,
         multitask_strategy: Optional[MultitaskStrategy] = None,
+        on_completion: Optional[OnCompletionBehavior] = None,
     ) -> AsyncIterator[StreamPart]:
         """Create a run and stream the results.
 
@@ -1019,6 +1028,8 @@ class RunsClient:
             webhook: Webhook to call after LangGraph API call is done.
             multitask_strategy: Multitask strategy to use.
                 Must be one of 'reject', 'interrupt', 'rollback', or 'enqueue'.
+            on_disconnect: The disconnect mode to use.
+                Must be one of 'cancel' or 'continue'.
 
         Returns:
             AsyncIterator[StreamPart]: Asynchronous iterator of stream results.
@@ -1061,6 +1072,8 @@ class RunsClient:
             "webhook": webhook,
             "checkpoint_id": checkpoint_id,
             "multitask_strategy": multitask_strategy,
+            "on_disconnect": on_disconnect,
+            "on_completion": on_completion,
         }
         endpoint = (
             f"/threads/{thread_id}/runs/stream"
@@ -1083,6 +1096,7 @@ class RunsClient:
         interrupt_before: Optional[list[str]] = None,
         interrupt_after: Optional[list[str]] = None,
         webhook: Optional[str] = None,
+        on_completion: Optional[OnCompletionBehavior] = None,
     ) -> Run:
         ...
 
@@ -1116,6 +1130,7 @@ class RunsClient:
         interrupt_after: Optional[list[str]] = None,
         webhook: Optional[str] = None,
         multitask_strategy: Optional[MultitaskStrategy] = None,
+        on_completion: Optional[OnCompletionBehavior] = None,
     ) -> Run:
         """Create a background run.
 
@@ -1129,9 +1144,7 @@ class RunsClient:
             config: The configuration for the assistant.
             checkpoint_id: The checkpoint to start streaming from.
             interrupt_before: Nodes to interrupt immediately before they get executed.
-
             interrupt_after: Nodes to Nodes to interrupt immediately after they get executed.
-
             webhook: Webhook to call after LangGraph API call is done.
             multitask_strategy: Multitask strategy to use.
                 Must be one of 'reject', 'interrupt', 'rollback', or 'enqueue'.
@@ -1214,6 +1227,7 @@ class RunsClient:
             "webhook": webhook,
             "checkpoint_id": checkpoint_id,
             "multitask_strategy": multitask_strategy,
+            "on_completion": on_completion,
         }
         payload = {k: v for k, v in payload.items() if v is not None}
         if thread_id:
@@ -1242,6 +1256,8 @@ class RunsClient:
         checkpoint_id: Optional[str] = None,
         interrupt_before: Optional[list[str]] = None,
         interrupt_after: Optional[list[str]] = None,
+        webhook: Optional[str] = None,
+        on_disconnect: Optional[DisconnectMode] = None,
         multitask_strategy: Optional[MultitaskStrategy] = None,
     ) -> Union[list[dict], dict[str, Any]]:
         ...
@@ -1257,6 +1273,9 @@ class RunsClient:
         config: Optional[Config] = None,
         interrupt_before: Optional[list[str]] = None,
         interrupt_after: Optional[list[str]] = None,
+        webhook: Optional[str] = None,
+        on_disconnect: Optional[DisconnectMode] = None,
+        on_completion: Optional[OnCompletionBehavior] = None,
     ) -> Union[list[dict], dict[str, Any]]:
         ...
 
@@ -1272,7 +1291,9 @@ class RunsClient:
         interrupt_before: Optional[list[str]] = None,
         interrupt_after: Optional[list[str]] = None,
         webhook: Optional[str] = None,
+        on_disconnect: Optional[DisconnectMode] = None,
         multitask_strategy: Optional[MultitaskStrategy] = None,
+        on_completion: Optional[OnCompletionBehavior] = None,
     ) -> Union[list[dict], dict[str, Any]]:
         """Create a run, wait until it finishes and return the final state.
 
@@ -1286,12 +1307,12 @@ class RunsClient:
             config: The configuration for the assistant.
             checkpoint_id: The checkpoint to start streaming from.
             interrupt_before: Nodes to interrupt immediately before they get executed.
-
             interrupt_after: Nodes to Nodes to interrupt immediately after they get executed.
-
             webhook: Webhook to call after LangGraph API call is done.
             multitask_strategy: Multitask strategy to use.
                 Must be one of 'reject', 'interrupt', 'rollback', or 'enqueue'.
+            on_disconnect: The disconnect mode to use.
+                Must be one of 'cancel' or 'continue'.
 
         Returns:
             Union[list[dict], dict[str, Any]]: The output of the run.
@@ -1351,6 +1372,8 @@ class RunsClient:
             "webhook": webhook,
             "checkpoint_id": checkpoint_id,
             "multitask_strategy": multitask_strategy,
+            "on_disconnect": on_disconnect,
+            "on_completion": on_completion,
         }
         endpoint = (
             f"/threads/{thread_id}/runs/wait" if thread_id is not None else "/runs/wait"
@@ -1451,6 +1474,28 @@ class RunsClient:
         """  # noqa: E501
         return await self.http.get(f"/threads/{thread_id}/runs/{run_id}/join")
 
+    def join_stream(self, thread_id: str, run_id: str) -> AsyncIterator[StreamPart]:
+        """Stream output from a run in real-time, until the run is done.
+        Output is not buffered, so any output produced before this call will
+        not be received here.
+
+        Args:
+            thread_id: The thread ID to join.
+            run_id: The run ID to join.
+
+        Returns:
+            None
+
+        Example Usage:
+
+            await client.runs.join(
+                thread_id="thread_id_to_join",
+                run_id="run_id_to_join"
+            )
+
+        """  # noqa: E501
+        return self.http.stream(f"/threads/{thread_id}/runs/{run_id}/stream", "GET")
+
     async def delete(self, thread_id: str, run_id: str) -> None:
         """Delete a run.
 

libs/sdk-py/langgraph_sdk/schema.py

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

libs/sdk-py/pyproject.toml

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