Add image arg to up command (#4385)

Using this argument, you can get more customization since you can do
`langgraph build` or directly `docker build` your image and then re-use
the `langgraph up --image my-image` and have it also spin up redis &
postgres for you.

Easier then writing your own compose file

docs/docs/agents/agents.md

@@ -29,7 +29,9 @@ agent = create_react_agent(
 )
 
 # Run the agent
-agent.invoke({"messages": "what is the weather in sf"})
+agent.invoke(
+    {"messages": [{"role": "user", "content": "what is the weather in sf"}]}
+)
 ```
 
 1. Define a tool for the agent to use. Tools can be defined as vanilla Python functions. For more advanced tool usage and customization, check the [tools](./tools.md) page.
@@ -85,7 +87,7 @@ agent = create_react_agent(
 )
 
 agent.invoke(
-    {"messages": "what is the weather in sf"},
+    {"messages": [{"role": "user", "content": "what is the weather in sf"}]}
 )
 ```
 
@@ -113,7 +115,7 @@ agent = create_react_agent(
 )
 
 agent.invoke(
-    {"messages": "what is the weather in sf"},
+    {"messages": [{"role": "user", "content": "what is the weather in sf"}]},
     # highlight-next-line
     config={"configurable": {"user_name": "John Smith"}}
 )
@@ -150,12 +152,12 @@ agent = create_react_agent(
 # highlight-next-line
 config = {"configurable": {"thread_id": "1"}}
 sf_response = agent.invoke(
-    {"messages": "what is the weather in sf"},
+    {"messages": [{"role": "user", "content": "what is the weather in sf"}]},
     # highlight-next-line
     config  # (2)!
 )
 ny_response = agent.invoke(
-    {"messages": "what about new york?"},
+    {"messages": [{"role": "user", "content": "what about new york?"}]},
     # highlight-next-line
     config
 )
@@ -189,7 +191,9 @@ agent = create_react_agent(
     response_format=WeatherResponse  # (1)!
 )
 
-response = agent.invoke({"messages": "what is the weather in sf"})
+response = agent.invoke(
+    {"messages": [{"role": "user", "content": "what is the weather in sf"}]}
+)
 
 # highlight-next-line
 response["structured_response"]

docs/docs/agents/context.md

@@ -36,7 +36,7 @@ for this purpose:
 
 ```python
 agent.invoke(
-    {"messages": "hi!"},
+    {"messages": [{"role": "user", "content": "hi!"}]},
     # highlight-next-line
     config={"configurable": {"user_id": "user_123"}}
 )
@@ -183,7 +183,7 @@ Tools can access context through special parameter **annotations**.
     )
 
     agent.invoke(
-        {"messages": "look up user information"},
+        {"messages": [{"role": "user", "content": "look up user information"}]},
         # highlight-next-line
         config={"configurable": {"user_id": "user_123"}}
     )
@@ -278,7 +278,7 @@ agent = create_react_agent(
 )
 
 agent.invoke(
-    {"messages": "greet the user"},
+    {"messages": [{"role": "user", "content": "greet the user"}]},
     # highlight-next-line
     config={"configurable": {"user_id": "user_123"}}
 )

docs/docs/agents/human-in-the-loop.md

@@ -70,7 +70,7 @@ config = {
 }
 
 for chunk in agent.stream(
-    {"messages": "book a stay at McKittrick hotel"},
+    {"messages": [{"role": "user", "content": "book a stay at McKittrick hotel"}]},
     # highlight-next-line
     config
 ):
@@ -194,7 +194,7 @@ config = {"configurable": {"thread_id": "1"}}
 
 # Run the agent
 for chunk in agent.stream(
-    {"messages": "book a stay at McKittrick hotel"},
+    {"messages": [{"role": "user", "content": "book a stay at McKittrick hotel"}]},
     # highlight-next-line
     config
 ):

docs/docs/agents/mcp.md

@@ -40,8 +40,12 @@ async with MultiServerMCPClient(
         # highlight-next-line
         client.get_tools()
     )
-    math_response = await agent.ainvoke({"messages": "what's (3 + 5) x 12?"})
-    weather_response = await agent.ainvoke({"messages": "what is the weather in nyc?"})
+    math_response = await agent.ainvoke(
+        {"messages": [{"role": "user", "content": "what's (3 + 5) x 12?"}]}
+    )
+    weather_response = await agent.ainvoke(
+        {"messages": [{"role": "user", "content": "what is the weather in nyc?"}]}
+    )
 ```
 
 ## Custom MCP servers

docs/docs/agents/memory.md

@@ -59,14 +59,14 @@ config = {
 }
 
 sf_response = agent.invoke(
-    {"messages": "what is the weather in sf"},
+    {"messages": [{"role": "user", "content": "what is the weather in sf"}]},
     # highlight-next-line
     config
 )
 
 # Continue the conversation using the same thread_id
 ny_response = agent.invoke(
-    {"messages": "what about new york?"},
+    {"messages": [{"role": "user", "content": "what about new york?"}]},
     # highlight-next-line
     config # (4)!
 )
@@ -188,7 +188,7 @@ agent = create_react_agent(
 
 # Run the agent
 agent.invoke(
-    {"messages": "look up user information"},
+    {"messages": [{"role": "user", "content": "look up user information"}]},
     # highlight-next-line
     config={"configurable": {"user_id": "user_123"}}
 )
@@ -206,7 +206,7 @@ agent.invoke(
 ### Writing
 
 ```python title="Example of a tool that updates user information"
-from typing import TypedDict
+from typing_extensions import TypedDict
 
 from langgraph.config import get_store
 from langgraph.prebuilt import create_react_agent
@@ -236,7 +236,7 @@ agent = create_react_agent(
 
 # Run the agent
 agent.invoke(
-    {"messages": "My name is John Smith"},
+    {"messages": [{"role": "user", "content": "My name is John Smith"}]},
     # highlight-next-line
     config={"configurable": {"user_id": "user_123"}} # (6)!
 )

docs/docs/agents/multi-agent.md

@@ -53,12 +53,22 @@ hotel_assistant = create_react_agent(
 supervisor = create_supervisor(
     agents=[flight_assistant, hotel_assistant],
     model=ChatOpenAI(model="gpt-4o"),
-    prompt="You manage a hotel booking assistant and a flight booking assistant. Assign work to them."
+    prompt=(
+        "You manage a hotel booking assistant and a"
+        "flight booking assistant. Assign work to them."
+    )
 ).compile()
 
-for chunk in supervisor.stream({
-    "messages": "book a flight from BOS to JFK and a stay at McKittrick Hotel"
-}):
+for chunk in supervisor.stream(
+    {
+        "messages": [
+            {
+                "role": "user",
+                "content": "book a flight from BOS to JFK and a stay at McKittrick Hotel"
+            }
+        ]
+    }
+):
     print(chunk)
     print("\n")
 ```
@@ -110,9 +120,16 @@ swarm = create_swarm(
     default_active_agent="flight_assistant"
 ).compile()
 
-for chunk in supervisor.stream({
-    "messages": "book a flight from BOS to JFK and a stay at McKittrick Hotel"
-}):
+for chunk in swarm.stream(
+    {
+        "messages": [
+            {
+                "role": "user",
+                "content": "book a flight from BOS to JFK and a stay at McKittrick Hotel"
+            }
+        ]
+    }
+):
     print(chunk)
     print("\n")
 ```
@@ -253,9 +270,16 @@ multi_agent_graph = (
 )
 
 # Run the multi-agent graph
-for chunk in multi_agent_graph.stream({
-    "messages": "book a flight from BOS to JFK and a stay at McKittrick Hotel"
-}):
+for chunk in multi_agent_graph.stream(
+    {
+        "messages": [
+            {
+                "role": "user",
+                "content": "book a flight from BOS to JFK and a stay at McKittrick Hotel"
+            }
+        ]
+    }
+):
     print(chunk)
     print("\n")
 ```

docs/docs/agents/run_agents.md

@@ -18,7 +18,7 @@ Agents can be executed in two primary modes:
     agent = create_react_agent(...)
 
     # highlight-next-line
-    response = agent.invoke({"messages": "what is the weather in sf"})
+    response = agent.invoke({"messages": [{"role": "user", "content": "what is the weather in sf"}]})
     ```
 
 === "Async invocation"
@@ -27,7 +27,7 @@ Agents can be executed in two primary modes:
 
     agent = create_react_agent(...)
     # highlight-next-line
-    response = await agent.ainvoke({"messages": "what is the weather in sf"})
+    response = await agent.ainvoke({"messages": [{"role": "user", "content": "what is the weather in sf"}]})
     ```
 
 ## Inputs and outputs
@@ -82,7 +82,7 @@ Streaming is available in both sync and async modes:
 
     ```python
     for chunk in agent.stream(
-        {"messages": "what is the weather in sf"},
+        {"messages": [{"role": "user", "content": "what is the weather in sf"}]},
         stream_mode="updates"
     ):
         print(chunk)
@@ -92,7 +92,7 @@ Streaming is available in both sync and async modes:
 
     ```python
     async for chunk in agent.astream(
-        {"messages": "what is the weather in sf"},
+        {"messages": [{"role": "user", "content": "what is the weather in sf"}]},
         stream_mode="updates"
     ):
         print(chunk)
@@ -122,7 +122,7 @@ To control agent execution and avoid infinite loops, set a recursion limit. This
 
     try:
         response = agent.invoke(
-            {"messages": "what's the weather in sf"},
+            {"messages": [{"role": "user", "content": "what's the weather in sf"}]},
             # highlight-next-line
             {"recursion_limit": recursion_limit},
         )
@@ -148,7 +148,7 @@ To control agent execution and avoid infinite loops, set a recursion limit. This
 
     try:
         response = agent_with_recursion_limit.invoke(
-            {"messages": "what's the weather in sf"},
+            {"messages": [{"role": "user", "content": "what's the weather in sf"}]},
         )
     except GraphRecursionError:
         print("Agent stopped due to max iterations.")

docs/docs/agents/streaming.md

@@ -35,7 +35,7 @@ For example, if you have an agent that calls a tool once, you should see the fol
     )
     # highlight-next-line
     for chunk in agent.stream(
-        {"messages": "what is the weather in sf"},
+        {"messages": [{"role": "user", "content": "what is the weather in sf"}]},
         # highlight-next-line
         stream_mode="updates"
     ):
@@ -52,7 +52,7 @@ For example, if you have an agent that calls a tool once, you should see the fol
     )
     # highlight-next-line
     async for chunk in agent.astream(
-        {"messages": "what is the weather in sf"},
+        {"messages": [{"role": "user", "content": "what is the weather in sf"}]},
         # highlight-next-line
         stream_mode="updates"
     ):
@@ -73,7 +73,7 @@ To stream tokens as they are produced by the LLM, use `stream_mode="messages"`:
     )
     # highlight-next-line
     for token, metadata in agent.stream(
-        {"messages": "what is the weather in sf"},
+        {"messages": [{"role": "user", "content": "what is the weather in sf"}]},
         # highlight-next-line
         stream_mode="messages"
     ):
@@ -91,7 +91,7 @@ To stream tokens as they are produced by the LLM, use `stream_mode="messages"`:
     )
     # highlight-next-line
     async for token, metadata in agent.astream(
-        {"messages": "what is the weather in sf"},
+        {"messages": [{"role": "user", "content": "what is the weather in sf"}]},
         # highlight-next-line
         stream_mode="messages"
     ):
@@ -125,7 +125,7 @@ To stream updates from tools as they are executed, you can use [get_stream_write
     )
 
     for chunk in agent.stream(
-        {"messages": "what is the weather in sf"},
+        {"messages": [{"role": "user", "content": "what is the weather in sf"}]},
         # highlight-next-line
         stream_mode="custom"
     ):
@@ -154,7 +154,7 @@ To stream updates from tools as they are executed, you can use [get_stream_write
     )
 
     async for chunk in agent.astream(
-        {"messages": "what is the weather in sf"},
+        {"messages": [{"role": "user", "content": "what is the weather in sf"}]},
         # highlight-next-line
         stream_mode="custom"
     ):
@@ -178,7 +178,7 @@ You can specify multiple streaming modes by passing stream mode as a list: `stre
     )
 
     for stream_mode, chunk in agent.stream(
-        {"messages": "what is the weather in sf"},
+        {"messages": [{"role": "user", "content": "what is the weather in sf"}]},
         # highlight-next-line
         stream_mode=["updates", "messages", "custom"]
     ):
@@ -195,7 +195,7 @@ You can specify multiple streaming modes by passing stream mode as a list: `stre
     )
 
     async for stream_mode, chunk in agent.astream(
-        {"messages": "what is the weather in sf"},
+        {"messages": [{"role": "user", "content": "what is the weather in sf"}]},
         # highlight-next-line
         stream_mode=["updates", "messages", "custom"]
     ):

docs/docs/agents/tools.md

@@ -116,7 +116,9 @@ agent = create_react_agent(
     tools=tools
 )
 
-agent.invoke({"messages": "what's 3 + 5 and 4 * 7? make both calculations in parallel"})
+agent.invoke(
+    {"messages": [{"role": "user", "content": "what's 3 + 5 and 4 * 7?"}]}
+)
 ```
 
 ## Return tool results directly
@@ -137,7 +139,9 @@ agent = create_react_agent(
     tools=[add]
 )
 
-agent.invoke({"messages": "what's 3 + 5?"})
+agent.invoke(
+    {"messages": [{"role": "user", "content": "what's 3 + 5?"}]}
+)
 ```
 
 ## Force tool use
@@ -161,7 +165,9 @@ agent = create_react_agent(
     tools=tools
 )
 
-agent.invoke({"messages": "Hi, I am Bob"})
+agent.invoke(
+    {"messages": [{"role": "user", "content": "Hi, I am Bob"}]}
+)
 ```
 
 !!! Warning "Avoid infinite loops"
@@ -191,7 +197,9 @@ By default, the agent will catch all exceptions raised during tool calls and wil
         model="anthropic:claude-3-7-sonnet-latest",
         tools=[multiply]
     )
-    agent.invoke({"messages": "what's 42 x 7?"})
+    agent.invoke(
+        {"messages": [{"role": "user", "content": "what's 42 x 7?"}]}
+    )
     ```
 
 === "Disable error handling"
@@ -215,7 +223,9 @@ By default, the agent will catch all exceptions raised during tool calls and wil
         model="anthropic:claude-3-7-sonnet-latest",
         tools=tool_node
     )
-    agent_no_error_handling.invoke({"messages": "what's 42 x 7?"})
+    agent_no_error_handling.invoke(
+        {"messages": [{"role": "user", "content": "what's 42 x 7?"}]}
+    )
     ```
 
     1. This disables error handling (enabled by default). See all available strategies in the [API reference][langgraph.prebuilt.tool_node.ToolNode].
@@ -243,7 +253,9 @@ By default, the agent will catch all exceptions raised during tool calls and wil
         model="anthropic:claude-3-7-sonnet-latest",
         tools=tool_node
     )
-    agent_custom_error_handling.invoke({"messages": "what's 42 x 7?"})
+    agent_custom_error_handling.invoke(
+        {"messages": [{"role": "user", "content": "what's 42 x 7?"}]}
+    )
     ```
 
     1. This provides a custom message to send to the LLM in case of an exception. See all available strategies in the [API reference][langgraph.prebuilt.tool_node.ToolNode].

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

@@ -1,7 +1,7 @@
 # LangGraph Studio With Local Deployment
 
 !!! warning "Browser Compatibility"
-    Viewing the studio page of a local LangGraph deployment does not work in Safari. Use Chrome instead.
+    Safari blocks `localhost` connections to Studio. To work around this, start the server with `--tunnel` and you’ll be able to access Studio from Safari via a secure tunnel.
 
 ## Setup
 

docs/docs/cloud/reference/cli.md

@@ -10,9 +10,6 @@ The LangGraph command line interface includes commands to build and run a LangGr
     === "Python"
         ```bash
         pip install langgraph-cli
-
-        # Install via Homebrew
-        brew install langgraph-cli
         ```
 
     === "JS"
@@ -298,6 +295,11 @@ The LangGraph CLI requires a JSON configuration file that follows this [schema](
     | `--no-reload`                 |                  | Disable auto-reload                                                                 |
     | `--n-jobs-per-worker INTEGER` |                  | Number of jobs per worker. Default is 10                                            |
     | `--debug-port INTEGER`        |                  | Port for debugger to listen on                                                      |
+    | `--wait-for-client`           | `False`          | Wait for a debugger client to connect to the debug port before starting the server   |
+    | `--no-browser`                |                  | Skip automatically opening the browser when the server starts                       |
+    | `--studio-url TEXT`           |                  | URL of the LangGraph Studio instance to connect to. Defaults to https://smith.langchain.com |
+    | `--allow-blocking`            | `False`          | Do not raise errors for synchronous I/O blocking operations in your code (added in `0.2.6`)           |
+    | `--tunnel`                    | `False`          | Expose the local server via a public tunnel (Cloudflare) for remote frontend access. This avoids issues with browsers like Safari or networks blocking localhost connections        |
     | `--help`                      |                  | Display command documentation                                                       |
 
 
@@ -321,6 +323,11 @@ The LangGraph CLI requires a JSON configuration file that follows this [schema](
     | `--no-reload`                 |                  | Disable auto-reload                                                                 |
     | `--n-jobs-per-worker INTEGER` |                  | Number of jobs per worker. Default is 10                                            |
     | `--debug-port INTEGER`        |                  | Port for debugger to listen on                                                      |
+    | `--wait-for-client`           | `False`          | Wait for a debugger client to connect to the debug port before starting the server   |
+    | `--no-browser`                |                  | Skip automatically opening the browser when the server starts                       |
+    | `--studio-url TEXT`           |                  | URL of the LangGraph Studio instance to connect to. Defaults to https://smith.langchain.com |
+    | `--allow-blocking`            | `False`          | Do not raise errors for synchronous I/O blocking operations in your code            |
+    | `--tunnel`                    | `False`          | Expose the local server via a public tunnel (Cloudflare) for remote frontend access. This avoids issues with browsers or networks blocking localhost connections        |
     | `--help`                      |                  | Display command documentation                                                       |
 
 ### `build`

docs/docs/stylesheets/logos.css

@@ -0,0 +1,3 @@
+.safari {
+  color: #0070C9;
+}

docs/docs/troubleshooting/errors/index.md

@@ -14,3 +14,4 @@ Errors referenced below will have an `lc_error_code` property corresponding to o
 These guides provide troubleshooting information for errors that are specific to the LangGraph Platform.
 
 - [INVALID_LICENSE](./INVALID_LICENSE.md)
+- [Studio Errors](../studio.md)

docs/docs/troubleshooting/studio.md

@@ -0,0 +1,45 @@
+# Troubleshooting LangGraph Studio
+
+## :fontawesome-brands-safari:{ .safari } Safari connection error with local dev server
+
+Safari blocks plain‑HTTP traffic on localhost. If you start Studio with a vanilla  
+`langgraph dev`, the page may report a "Failed to load assistants" error (or something similar) and the browser DevTools will show network errors.
+
+#### Quick fix — run Studio through a secure Cloudflare tunnel
+
+=== "Python"
+
+    ```shell
+    pip install -U langgraph-cli>=0.2.6            # Python
+    langgraph dev --tunnel
+    ```
+=== "JS"
+
+    ```shell
+    # Requires @langchain/langgraph-cli>=0.0.26
+    npx @langchain/langgraph-cli dev
+    ```
+
+The command prints a URL like:
+
+```shell
+https://smith.langchain.com/studio/?baseUrl=https://hamilton-praise-heart-costumes.trycloudflare.com
+```
+where
+```shell
+?baseUrl=https://hamilton-praise-heart-costumes.trycloudflare.com
+```
+indicates the endpoint where your agent server is exposed.
+
+Open that URL in Safari and Studio should load immediately.
+
+#### Alternative — use a Chromium‑based browser
+
+Chrome, Edge, and Brave allow HTTP on localhost, so a plain `langgraph dev` should work without extra steps.
+
+#### If it’s still not loading
+
+1. Make sure the `baseUrl` query parameter in the studio URL points to the **tunnel URL** NOT to localhost.
+2. Confirm your CLI version with `langgraph --version`.
+
+No other configuration, certificates, or CORS tweaks are required.

docs/mkdocs.yml

@@ -1,5 +1,5 @@
-site_name: ""
-site_description: Build language agents as graphs
+site_name: "LangGraph"
+site_description: Build reliable, stateful AI systems, without giving up control
 site_url: https://langchain-ai.github.io/langgraph/
 repo_url: https://github.com/langchain-ai/langgraph
 edit_uri: edit/main/docs/docs/
@@ -400,6 +400,7 @@ nav:
           - troubleshooting/errors/MULTIPLE_SUBGRAPHS.md
           - troubleshooting/errors/INVALID_CHAT_HISTORY.md
           - troubleshooting/errors/INVALID_LICENSE.md
+          - troubleshooting/studio.md
       - LangGraph Academy Course: https://academy.langchain.com/courses/intro-to-langgraph
 
   - Agents:
@@ -549,3 +550,4 @@ copyright: >
   Copyright &copy; 2025 LangChain, Inc | <a href="#__consent">Consent Preferences</a>
 extra_css:
   - stylesheets/version_admonitions.css
+  - stylesheets/logos.css

libs/cli/langgraph_cli/cli.py

@@ -168,6 +168,13 @@ def cli():
 @OPT_DEBUGGER_BASE_URL
 @OPT_WATCH
 @OPT_POSTGRES_URI
+@click.option(
+    "--image",
+    type=str,
+    default=None,
+    help="Docker image to use for the langgraph-api service. If specified, skips building and uses this image directly."
+    " Useful if you want to test against an image already built using `langgraph build`.",
+)
 @click.option(
     "--wait",
     is_flag=True,
@@ -187,6 +194,7 @@ def up(
     debugger_port: Optional[int],
     debugger_base_url: Optional[str],
     postgres_uri: Optional[str],
+    image: Optional[str],
 ):
     click.secho("Starting LangGraph API server...", fg="green")
     click.secho(
@@ -207,6 +215,7 @@ For production use, requires a license key in env var LANGGRAPH_CLOUD_LICENSE_KE
             debugger_port=debugger_port,
             debugger_base_url=debugger_base_url,
             postgres_uri=postgres_uri,
+            image=image,
         )
         # add up + options
         args.extend(["up", "--remove-orphans"])
@@ -692,6 +701,7 @@ def prepare_args_and_stdin(
     debugger_port: Optional[int] = None,
     debugger_base_url: Optional[str] = None,
     postgres_uri: Optional[str] = None,
+    image: Optional[str] = None,
 ) -> Tuple[List[str], str]:
     assert config_path.exists(), f"Config file not found: {config_path}"
     # prepare args
@@ -701,6 +711,7 @@ def prepare_args_and_stdin(
         debugger_port=debugger_port,
         debugger_base_url=debugger_base_url,
         postgres_uri=postgres_uri,
+        image=image,  # Pass image to compose YAML generator
     )
     args = [
         "--project-directory",
@@ -716,6 +727,7 @@ def prepare_args_and_stdin(
         config,
         watch=watch,
         base_image=langgraph_cli.config.default_base_image(config),
+        image=image,
     )
     return args, stdin
 
@@ -733,6 +745,7 @@ def prepare(
     debugger_port: Optional[int] = None,
     debugger_base_url: Optional[str] = None,
     postgres_uri: Optional[str] = None,
+    image: Optional[str] = None,
 ) -> Tuple[List[str], str]:
     """Prepare the arguments and stdin for running the LangGraph API server."""
     config_json = langgraph_cli.config.validate_config_file(config_path)
@@ -757,5 +770,6 @@ def prepare(
         debugger_port=debugger_port,
         debugger_base_url=debugger_base_url or f"http://127.0.0.1:{port}",
         postgres_uri=postgres_uri,
+        image=image,
     )
     return args, stdin

libs/cli/langgraph_cli/config.py

@@ -1288,6 +1288,7 @@ def config_to_compose(
     config_path: pathlib.Path,
     config: Config,
     base_image: Optional[str] = None,
+    image: Optional[str] = None,
     watch: bool = False,
 ) -> str:
     base_image = base_image or default_base_image(config)
@@ -1314,19 +1315,28 @@ def config_to_compose(
 """
     else:
         watch_str = ""
+    if image:
+        return f"""
+{textwrap.indent(env_vars_str, "            ")}
+        {env_file_str}
+        {watch_str}
+"""
 
-    dockerfile, additional_contexts = config_to_docker(config_path, config, base_image)
+    else:
+        dockerfile, additional_contexts = config_to_docker(
+            config_path, config, base_image
+        )
 
-    additional_contexts_str = "\n".join(
-        f"                - {name}: {path}"
-        for name, path in additional_contexts.items()
-    )
-    if additional_contexts_str:
-        additional_contexts_str = f"""
+        additional_contexts_str = "\n".join(
+            f"                - {name}: {path}"
+            for name, path in additional_contexts.items()
+        )
+        if additional_contexts_str:
+            additional_contexts_str = f"""
             additional_contexts:
 {additional_contexts_str}"""
 
-    return f"""
+        return f"""
 {textwrap.indent(env_vars_str, "            ")}
         {env_file_str}
         pull_policy: build

libs/cli/langgraph_cli/docker.py

@@ -143,6 +143,8 @@ def compose_as_dict(
     debugger_base_url: Optional[str] = None,
     # postgres://user:password@host:port/database?option=value
     postgres_uri: Optional[str] = None,
+    # If you are running against an already-built image, you can pass it here
+    image: Optional[str] = None,
 ) -> dict:
     """Create a docker compose file as a dictionary in YML style."""
     if postgres_uri is None:
@@ -211,6 +213,8 @@ def compose_as_dict(
             "POSTGRES_URI": postgres_uri,
         },
     }
+    if image:
+        services["langgraph-api"]["image"] = image
 
     # If Postgres is included, add it to the dependencies of langgraph-api
     if include_db:
@@ -244,6 +248,7 @@ def compose(
     debugger_base_url: Optional[str] = None,
     # postgres://user:password@host:port/database?option=value
     postgres_uri: Optional[str] = None,
+    image: Optional[str] = None,
 ) -> str:
     """Create a docker compose file as a string."""
     compose_content = compose_as_dict(
@@ -252,6 +257,7 @@ def compose(
         debugger_port=debugger_port,
         debugger_base_url=debugger_base_url,
         postgres_uri=postgres_uri,
+        image=image,
     )
     compose_str = dict_to_yaml(compose_content)
     return compose_str

libs/cli/pyproject.toml

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

libs/cli/tests/unit_tests/cli/test_cli.py

@@ -160,6 +160,110 @@ services:
     assert clean_empty_lines(actual_stdin) == expected_stdin
 
 
+def test_prepare_args_and_stdin_with_image() -> None:
+    # this basically serves as an end-to-end test for using config and docker helpers
+    config_path = pathlib.Path(__file__).parent / "langgraph.json"
+    config = validate_config(
+        Config(dependencies=[".", "../../.."], graphs={"agent": "agent.py:graph"})
+    )
+    port = 8000
+    debugger_port = 8001
+    debugger_graph_url = f"http://127.0.0.1:{port}"
+
+    actual_args, actual_stdin = prepare_args_and_stdin(
+        capabilities=DEFAULT_DOCKER_CAPABILITIES,
+        config_path=config_path,
+        config=config,
+        docker_compose=pathlib.Path("custom-docker-compose.yml"),
+        port=port,
+        debugger_port=debugger_port,
+        debugger_base_url=debugger_graph_url,
+        watch=True,
+        image="my-cool-image",
+    )
+
+    expected_args = [
+        "--project-directory",
+        str(pathlib.Path(__file__).parent.absolute()),
+        "-f",
+        "custom-docker-compose.yml",
+        "-f",
+        "-",
+    ]
+    expected_stdin = f"""volumes:
+    langgraph-data:
+        driver: local
+services:
+    langgraph-redis:
+        image: redis:6
+        healthcheck:
+            test: redis-cli ping
+            interval: 5s
+            timeout: 1s
+            retries: 5
+    langgraph-postgres:
+        image: pgvector/pgvector:pg16
+        ports:
+            - "5433:5432"
+        environment:
+            POSTGRES_DB: postgres
+            POSTGRES_USER: postgres
+            POSTGRES_PASSWORD: postgres
+        command:
+            - postgres
+            - -c
+            - shared_preload_libraries=vector
+        volumes:
+            - langgraph-data:/var/lib/postgresql/data
+        healthcheck:
+            test: pg_isready -U postgres
+            start_period: 10s
+            timeout: 1s
+            retries: 5
+            interval: 60s
+            start_interval: 1s
+    langgraph-debugger:
+        image: langchain/langgraph-debugger
+        restart: on-failure
+        depends_on:
+            langgraph-postgres:
+                condition: service_healthy
+        ports:
+            - "{debugger_port}:3968"
+        environment:
+            VITE_STUDIO_LOCAL_GRAPH_URL: {debugger_graph_url}
+    langgraph-api:
+        ports:
+            - "8000:8000"
+        depends_on:
+            langgraph-redis:
+                condition: service_healthy
+            langgraph-postgres:
+                condition: service_healthy
+        environment:
+            REDIS_URI: redis://langgraph-redis:6379
+            POSTGRES_URI: {DEFAULT_POSTGRES_URI}
+        image: my-cool-image
+        healthcheck:
+            test: python /api/healthcheck.py
+            interval: 60s
+            start_interval: 1s
+            start_period: 10s
+        
+        
+        develop:
+            watch:
+                - path: langgraph.json
+                  action: rebuild
+                - path: .
+                  action: rebuild
+                - path: ../../..
+                  action: rebuild\
+"""
+    assert actual_args == expected_args
+    assert clean_empty_lines(actual_stdin) == expected_stdin
+
+
 def test_version_option() -> None:
     """Test the --version option of the CLI."""
     runner = CliRunner()

libs/langgraph/Makefile

@@ -78,7 +78,7 @@ test_watch_all:
 PYTHON_FILES=.
 MYPY_CACHE=.mypy_cache
 lint format: PYTHON_FILES=.
-lint_diff format_diff: PYTHON_FILES=$(shell git diff --name-only --relative --diff-filter=d main . | grep -E '\.py$$|\.ipynb$$')
+lint_diff format_diff: PYTHON_FILES=$(shell git diff --name-only --relative --diff-filter=d main . | grep -E r'\.py$$|\.ipynb$$')
 lint_package: PYTHON_FILES=langgraph
 lint_tests: PYTHON_FILES=tests
 lint_tests: MYPY_CACHE=.mypy_cache_test

libs/langgraph/bench/pydantic_state.py

@@ -1,7 +1,8 @@
 import operator
+from collections.abc import Sequence
 from functools import partial
 from random import choice
-from typing import Annotated, Optional, Sequence
+from typing import Annotated, Optional
 
 from pydantic import BaseModel, Field, field_validator
 

libs/langgraph/bench/wide_dict.py

@@ -1,7 +1,8 @@
 import operator
+from collections.abc import Sequence
 from functools import partial
 from random import choice
-from typing import Annotated, Optional, Sequence
+from typing import Annotated, Optional
 
 from typing_extensions import TypedDict
 

libs/langgraph/bench/wide_state.py

@@ -1,8 +1,9 @@
 import operator
+from collections.abc import Sequence
 from dataclasses import dataclass, field
 from functools import partial
 from random import choice
-from typing import Annotated, Optional, Sequence
+from typing import Annotated, Optional
 
 from langgraph.constants import END, START
 from langgraph.graph.state import StateGraph

libs/langgraph/langgraph/_api/deprecation.py

@@ -1,6 +1,6 @@
 import functools
 import warnings
-from typing import Any, Callable, Type, TypeVar, Union, cast
+from typing import Any, Callable, TypeVar, Union, cast
 
 
 class LangGraphDeprecationWarning(DeprecationWarning):
@@ -8,7 +8,7 @@ class LangGraphDeprecationWarning(DeprecationWarning):
 
 
 F = TypeVar("F", bound=Callable[..., Any])
-C = TypeVar("C", bound=Type[Any])
+C = TypeVar("C", bound=type[Any])
 
 
 def deprecated(

libs/langgraph/langgraph/channels/any_value.py

@@ -1,4 +1,5 @@
-from typing import Any, Generic, Sequence, Type
+from collections.abc import Sequence
+from typing import Any, Generic
 
 from typing_extensions import Self
 
@@ -21,12 +22,12 @@ class AnyValue(Generic[Value], BaseChannel[Value, Value, Value]):
         return isinstance(value, AnyValue)
 
     @property
-    def ValueType(self) -> Type[Value]:
+    def ValueType(self) -> type[Value]:
         """The type of the value stored in the channel."""
         return self.typ
 
     @property
-    def UpdateType(self) -> Type[Value]:
+    def UpdateType(self) -> type[Value]:
         """The type of the update received by the channel."""
         return self.typ
 

libs/langgraph/langgraph/channels/base.py

@@ -1,5 +1,6 @@
 from abc import ABC, abstractmethod
-from typing import Any, Generic, Sequence, TypeVar
+from collections.abc import Sequence
+from typing import Any, Generic, TypeVar
 
 from typing_extensions import Self
 

libs/langgraph/langgraph/channels/binop.py

@@ -1,5 +1,6 @@
 import collections.abc
-from typing import Callable, Generic, Sequence, Type
+from collections.abc import Sequence
+from typing import Callable, Generic
 
 from typing_extensions import NotRequired, Required, Self
 
@@ -31,7 +32,7 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]):
 
     __slots__ = ("value", "operator")
 
-    def __init__(self, typ: Type[Value], operator: Callable[[Value, Value], Value]):
+    def __init__(self, typ: type[Value], operator: Callable[[Value, Value], Value]):
         super().__init__(typ)
         self.operator = operator
         # special forms from typing or collections.abc are not instantiable
@@ -57,12 +58,12 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]):
         )
 
     @property
-    def ValueType(self) -> Type[Value]:
+    def ValueType(self) -> type[Value]:
         """The type of the value stored in the channel."""
         return self.typ
 
     @property
-    def UpdateType(self) -> Type[Value]:
+    def UpdateType(self) -> type[Value]:
         """The type of the update received by the channel."""
         return self.typ
 

libs/langgraph/langgraph/channels/dynamic_barrier_value.py

@@ -1,4 +1,5 @@
-from typing import Any, Generic, NamedTuple, Optional, Sequence, Type, Union
+from collections.abc import Sequence, Set
+from typing import Any, Generic, NamedTuple, Optional, Union
 
 from typing_extensions import Self
 
@@ -8,11 +9,11 @@ from langgraph.errors import EmptyChannelError, InvalidUpdateError
 
 
 class WaitForNames(NamedTuple):
-    names: set[Any]
+    names: Set[Any]
 
 
 class DynamicBarrierValue(
-    Generic[Value], BaseChannel[Value, Union[Value, WaitForNames], set[Value]]
+    Generic[Value], BaseChannel[Value, Union[Value, WaitForNames], Set[Value]]
 ):
     """A channel that switches between two states
 
@@ -25,10 +26,10 @@ class DynamicBarrierValue(
 
     __slots__ = ("names", "seen")
 
-    names: Optional[set[Value]]
+    names: Optional[Set[Value]]
     seen: set[Value]
 
-    def __init__(self, typ: Type[Value]) -> None:
+    def __init__(self, typ: type[Value]) -> None:
         super().__init__(typ)
         self.names = None
         self.seen = set()
@@ -37,12 +38,12 @@ class DynamicBarrierValue(
         return isinstance(value, DynamicBarrierValue) and value.names == self.names
 
     @property
-    def ValueType(self) -> Type[Value]:
+    def ValueType(self) -> type[Value]:
         """The type of the value stored in the channel."""
         return self.typ
 
     @property
-    def UpdateType(self) -> Type[Value]:
+    def UpdateType(self) -> type[Value]:
         """The type of the update received by the channel."""
         return self.typ
 
@@ -54,11 +55,11 @@ class DynamicBarrierValue(
         empty.seen = self.seen.copy()
         return empty
 
-    def checkpoint(self) -> tuple[Optional[set[Value]], set[Value]]:
+    def checkpoint(self) -> tuple[Optional[Set[Value]], set[Value]]:
         return (self.names, self.seen)
 
     def from_checkpoint(
-        self, checkpoint: tuple[Optional[set[Value]], set[Value]]
+        self, checkpoint: tuple[Optional[Set[Value]], set[Value]]
     ) -> Self:
         empty = self.__class__(self.typ)
         empty.key = self.key

libs/langgraph/langgraph/channels/ephemeral_value.py

@@ -1,4 +1,5 @@
-from typing import Any, Generic, Sequence, Type
+from collections.abc import Sequence
+from typing import Any, Generic
 
 from typing_extensions import Self
 
@@ -21,12 +22,12 @@ class EphemeralValue(Generic[Value], BaseChannel[Value, Value, Value]):
         return isinstance(value, EphemeralValue) and value.guard == self.guard
 
     @property
-    def ValueType(self) -> Type[Value]:
+    def ValueType(self) -> type[Value]:
         """The type of the value stored in the channel."""
         return self.typ
 
     @property
-    def UpdateType(self) -> Type[Value]:
+    def UpdateType(self) -> type[Value]:
         """The type of the update received by the channel."""
         return self.typ
 

libs/langgraph/langgraph/channels/last_value.py

@@ -1,4 +1,5 @@
-from typing import Any, Generic, Sequence, Type
+from collections.abc import Sequence
+from typing import Any, Generic
 
 from typing_extensions import Self
 
@@ -25,12 +26,12 @@ class LastValue(Generic[Value], BaseChannel[Value, Value, Value]):
         return isinstance(value, LastValue)
 
     @property
-    def ValueType(self) -> Type[Value]:
+    def ValueType(self) -> type[Value]:
         """The type of the value stored in the channel."""
         return self.typ
 
     @property
-    def UpdateType(self) -> Type[Value]:
+    def UpdateType(self) -> type[Value]:
         """The type of the update received by the channel."""
         return self.typ
 

libs/langgraph/langgraph/channels/named_barrier_value.py

@@ -1,4 +1,5 @@
-from typing import Generic, Sequence, Type
+from collections.abc import Sequence
+from typing import Generic
 
 from typing_extensions import Self
 
@@ -15,7 +16,7 @@ class NamedBarrierValue(Generic[Value], BaseChannel[Value, Value, set[Value]]):
     names: set[Value]
     seen: set[Value]
 
-    def __init__(self, typ: Type[Value], names: set[Value]) -> None:
+    def __init__(self, typ: type[Value], names: set[Value]) -> None:
         super().__init__(typ)
         self.names = names
         self.seen: set[str] = set()
@@ -24,12 +25,12 @@ class NamedBarrierValue(Generic[Value], BaseChannel[Value, Value, set[Value]]):
         return isinstance(value, NamedBarrierValue) and value.names == self.names
 
     @property
-    def ValueType(self) -> Type[Value]:
+    def ValueType(self) -> type[Value]:
         """The type of the value stored in the channel."""
         return self.typ
 
     @property
-    def UpdateType(self) -> Type[Value]:
+    def UpdateType(self) -> type[Value]:
         """The type of the update received by the channel."""
         return self.typ
 

libs/langgraph/langgraph/channels/topic.py

@@ -1,4 +1,5 @@
-from typing import Any, Generic, Iterator, Sequence, Type, Union
+from collections.abc import Iterator, Sequence
+from typing import Any, Generic, Union
 
 from typing_extensions import Self
 
@@ -28,7 +29,7 @@ class Topic(
 
     __slots__ = ("values", "accumulate")
 
-    def __init__(self, typ: Type[Value], accumulate: bool = False) -> None:
+    def __init__(self, typ: type[Value], accumulate: bool = False) -> None:
         super().__init__(typ)
         # attrs
         self.accumulate = accumulate

libs/langgraph/langgraph/channels/untracked_value.py

@@ -1,4 +1,5 @@
-from typing import Generic, Sequence, Type
+from collections.abc import Sequence
+from typing import Generic
 
 from typing_extensions import Self
 
@@ -12,7 +13,7 @@ class UntrackedValue(Generic[Value], BaseChannel[Value, Value, Value]):
 
     __slots__ = ("value", "guard")
 
-    def __init__(self, typ: Type[Value], guard: bool = True) -> None:
+    def __init__(self, typ: type[Value], guard: bool = True) -> None:
         super().__init__(typ)
         self.guard = guard
         self.value = MISSING
@@ -21,12 +22,12 @@ class UntrackedValue(Generic[Value], BaseChannel[Value, Value, Value]):
         return isinstance(value, UntrackedValue) and value.guard == self.guard
 
     @property
-    def ValueType(self) -> Type[Value]:
+    def ValueType(self) -> type[Value]:
         """The type of the value stored in the channel."""
         return self.typ
 
     @property
-    def UpdateType(self) -> Type[Value]:
+    def UpdateType(self) -> type[Value]:
         """The type of the update received by the channel."""
         return self.typ
 

libs/langgraph/langgraph/constants.py

@@ -1,6 +1,7 @@
 import sys
+from collections.abc import Mapping
 from types import MappingProxyType
-from typing import Any, Literal, Mapping, cast
+from typing import Any, Literal, cast
 
 from langgraph.types import Interrupt, Send  # noqa: F401
 

libs/langgraph/langgraph/errors.py

@@ -1,5 +1,6 @@
+from collections.abc import Sequence
 from enum import Enum
-from typing import Any, Sequence
+from typing import Any
 
 from langgraph.checkpoint.base import EmptyChannelError  # noqa: F401
 from langgraph.types import Command, Interrupt

libs/langgraph/langgraph/func/__init__.py

@@ -2,14 +2,13 @@ import asyncio
 import concurrent.futures
 import functools
 import inspect
+from collections.abc import Awaitable, Sequence
 from dataclasses import dataclass
 from typing import (
     Any,
-    Awaitable,
     Callable,
     Generic,
     Optional,
-    Sequence,
     TypeVar,
     Union,
     get_args,

libs/langgraph/langgraph/graph/branch.py

@@ -1,19 +1,17 @@
+from collections.abc import Awaitable, Hashable, Sequence
 from inspect import (
     isfunction,
     ismethod,
     signature,
 )
+from itertools import zip_longest
 from types import FunctionType
 from typing import (
     Any,
-    Awaitable,
     Callable,
-    Hashable,
     Literal,
     NamedTuple,
     Optional,
-    Sequence,
-    Type,
     Union,
     cast,
     get_args,
@@ -29,12 +27,17 @@ from langchain_core.runnables import (
 
 from langgraph.constants import END, START
 from langgraph.errors import InvalidUpdateError
-from langgraph.pregel.write import ChannelWrite
+from langgraph.pregel.write import PASSTHROUGH, ChannelWrite, ChannelWriteEntry
 from langgraph.types import Send
 from langgraph.utils.runnable import (
     RunnableCallable,
 )
 
+Writer = Callable[
+    [Sequence[Union[str, Send]]],
+    Sequence[Union[ChannelWriteEntry, Send]],
+]
+
 
 def _get_branch_path_input_schema(
     path: Union[
@@ -42,7 +45,7 @@ def _get_branch_path_input_schema(
         Callable[..., Awaitable[Union[Hashable, list[Hashable]]]],
         Runnable[Any, Union[Hashable, list[Hashable]]],
     ],
-) -> Optional[Type[Any]]:
+) -> Optional[type[Any]]:
     input = None
     # detect input schema annotation in the branch callable
     try:
@@ -85,7 +88,7 @@ class Branch(NamedTuple):
     path: Runnable[Any, Union[Hashable, list[Hashable]]]
     ends: Optional[dict[Hashable, str]]
     then: Optional[str] = None
-    input_schema: Optional[Type[Any]] = None
+    input_schema: Optional[type[Any]] = None
 
     @classmethod
     def from_path(
@@ -124,9 +127,7 @@ class Branch(NamedTuple):
 
     def run(
         self,
-        writer: Callable[
-            [Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite]
-        ],
+        writer: Writer,
         reader: Optional[Callable[[RunnableConfig], Any]] = None,
     ) -> RunnableCallable:
         return ChannelWrite.register_writer(
@@ -138,7 +139,15 @@ class Branch(NamedTuple):
                 name=None,
                 trace=False,
                 func_accepts_config=True,
+            ),
+            list(
+                zip_longest(
+                    writer([e for e in self.ends.values() if e != END]),
+                    [str(la) for la, e in self.ends.items() if e != END],
+                )
             )
+            if self.ends
+            else None,
         )
 
     def _route(
@@ -147,9 +156,7 @@ class Branch(NamedTuple):
         config: RunnableConfig,
         *,
         reader: Optional[Callable[[RunnableConfig], Any]],
-        writer: Callable[
-            [Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite]
-        ],
+        writer: Writer,
     ) -> Runnable:
         if reader:
             value = reader(config)
@@ -172,9 +179,7 @@ class Branch(NamedTuple):
         config: RunnableConfig,
         *,
         reader: Optional[Callable[[RunnableConfig], Any]],
-        writer: Callable[
-            [Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite]
-        ],
+        writer: Writer,
     ) -> Runnable:
         if reader:
             value = reader(config)
@@ -193,9 +198,7 @@ class Branch(NamedTuple):
 
     def _finish(
         self,
-        writer: Callable[
-            [Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite]
-        ],
+        writer: Writer,
         input: Any,
         result: Any,
         config: RunnableConfig,
@@ -212,4 +215,18 @@ class Branch(NamedTuple):
             raise ValueError("Branch did not return a valid destination")
         if any(p.node == END for p in destinations if isinstance(p, Send)):
             raise InvalidUpdateError("Cannot send a packet to the END node")
-        return writer(destinations, config) or input
+        entries = writer(destinations)
+        if not entries:
+            return input
+        else:
+            need_passthrough = False
+            for e in entries:
+                if isinstance(e, ChannelWriteEntry):
+                    if e.value is PASSTHROUGH:
+                        need_passthrough = True
+                        break
+            if need_passthrough:
+                return ChannelWrite(entries)
+            else:
+                ChannelWrite.do_write(config, entries)
+                return input

libs/langgraph/langgraph/graph/graph.py

@@ -1,23 +1,17 @@
-import asyncio
 import logging
 from collections import defaultdict
+from collections.abc import Awaitable, Hashable, Sequence
 from typing import (
     Any,
-    Awaitable,
     Callable,
-    Hashable,
     NamedTuple,
     Optional,
-    Sequence,
     Union,
     cast,
     overload,
 )
 
 from langchain_core.runnables import Runnable
-from langchain_core.runnables.config import RunnableConfig
-from langchain_core.runnables.graph import Graph as DrawableGraph
-from langchain_core.runnables.graph import Node as DrawableNode
 from typing_extensions import Self
 
 from langgraph.channels.ephemeral_value import EphemeralValue
@@ -32,7 +26,6 @@ from langgraph.constants import (
 )
 from langgraph.graph.branch import Branch
 from langgraph.pregel import Channel, Pregel
-from langgraph.pregel.protocol import PregelProtocol
 from langgraph.pregel.read import PregelNode
 from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
 from langgraph.types import All, Checkpointer
@@ -182,7 +175,7 @@ class Graph:
         # validate the condition
         if name in self.branches[source]:
             raise ValueError(
-                f"Branch with name `{path.name}` already exists for node " f"`{source}`"
+                f"Branch with name `{path.name}` already exists for node `{source}`"
             )
         # save it
         self.branches[source][name] = Branch.from_path(path, path_map, then, False)
@@ -380,10 +373,10 @@ class CompiledGraph(Pregel):
             cast(list[str], self.nodes[end].channels).append(start)
 
     def attach_branch(self, start: str, name: str, branch: Branch) -> None:
-        def branch_writer(
-            packets: Sequence[Union[str, Send]], config: RunnableConfig
-        ) -> Optional[ChannelWrite]:
-            writes = [
+        def get_writes(
+            packets: Sequence[Union[str, Send]],
+        ) -> Sequence[Union[ChannelWriteEntry, Send]]:
+            return [
                 (
                     ChannelWriteEntry(f"branch:{start}:{name}:{p}" if p != END else END)
                     if not isinstance(p, Send)
@@ -391,14 +384,13 @@ class CompiledGraph(Pregel):
                 )
                 for p in packets
             ]
-            return ChannelWrite(cast(Sequence[Union[ChannelWriteEntry, Send]], writes))
 
         # add hidden start node
         if start == START and start not in self.nodes:
             self.nodes[start] = Channel.subscribe_to(START, tags=[TAG_HIDDEN])
 
         # attach branch writer
-        self.nodes[start] |= branch.run(branch_writer)
+        self.nodes[start] |= branch.run(get_writes)
 
         # attach branch readers
         ends = branch.ends.values() if branch.ends else [node for node in self.nodes]
@@ -408,171 +400,3 @@ class CompiledGraph(Pregel):
                 self.channels[channel_name] = EphemeralValue(Any)
                 self.nodes[end].triggers.append(channel_name)
                 cast(list[str], self.nodes[end].channels).append(channel_name)
-
-    async def aget_graph(
-        self,
-        config: Optional[RunnableConfig] = None,
-        *,
-        xray: Union[int, bool] = False,
-    ) -> DrawableGraph:
-        """Returns a drawable representation of the computation graph."""
-        from langgraph.pregel.remote import RemoteGraph
-
-        # gather subgraphs
-        if xray:
-            subpregels: dict[str, PregelProtocol] = {
-                k: v
-                async for k, v in self.aget_subgraphs()
-                if isinstance(v, (CompiledGraph, RemoteGraph))
-            }
-            subgraphs = {
-                k: v
-                for k, v in zip(
-                    subpregels,
-                    await asyncio.gather(
-                        *(
-                            p.aget_graph(
-                                config,
-                                xray=xray
-                                if isinstance(xray, bool) or xray <= 0
-                                else xray - 1,
-                            )
-                            for p in subpregels.values()
-                        )
-                    ),
-                )
-            }
-        else:
-            subgraphs = {}
-
-        # draw the graph
-        return self._draw_graph(config, subgraphs=subgraphs)
-
-    def get_graph(
-        self,
-        config: Optional[RunnableConfig] = None,
-        *,
-        xray: Union[int, bool] = False,
-    ) -> DrawableGraph:
-        """Returns a drawable representation of the computation graph."""
-        from langgraph.pregel.remote import RemoteGraph
-
-        # gather subgraphs
-        if xray:
-            subgraphs = {
-                k: v.get_graph(
-                    config,
-                    xray=xray if isinstance(xray, bool) or xray <= 0 else xray - 1,
-                )
-                for k, v in self.get_subgraphs()
-                if isinstance(v, (CompiledGraph, RemoteGraph))
-            }
-        else:
-            subgraphs = {}
-
-        # draw the graph
-        return self._draw_graph(config, subgraphs=subgraphs)
-
-    def _draw_graph(
-        self,
-        config: Optional[RunnableConfig] = None,
-        *,
-        subgraphs: dict[str, DrawableGraph] = {},
-    ) -> DrawableGraph:
-        # create the graph
-        graph = DrawableGraph()
-        start_nodes: dict[str, DrawableNode] = {
-            START: graph.add_node(self.get_input_schema(config), START)
-        }
-        end_nodes: dict[str, DrawableNode] = {}
-
-        def add_edge(
-            start: str,
-            end: str,
-            label: Optional[Hashable] = None,
-            conditional: bool = False,
-        ) -> None:
-            if end == END and END not in end_nodes:
-                end_nodes[END] = graph.add_node(self.get_output_schema(config), END)
-            if start not in start_nodes or end not in end_nodes:
-                logger.warning(
-                    f"Could not add edge from '{start}' to '{end}' due to missing nodes"
-                )
-                return
-            return graph.add_edge(
-                start_nodes[start],
-                end_nodes[end],
-                str(label) if label is not None else None,
-                conditional,
-            )
-
-        for key, n in self.builder.nodes.items():
-            node = n.runnable
-            metadata = n.metadata or {}
-            if key in self.interrupt_before_nodes and key in self.interrupt_after_nodes:
-                metadata["__interrupt"] = "before,after"
-            elif key in self.interrupt_before_nodes:
-                metadata["__interrupt"] = "before"
-            elif key in self.interrupt_after_nodes:
-                metadata["__interrupt"] = "after"
-            if key in subgraphs:
-                subgraph = subgraphs[key]
-                subgraph.trim_first_node()
-                subgraph.trim_last_node()
-                if len(subgraph.nodes) >= 1:
-                    e, s = graph.extend(subgraph, prefix=key)
-                    if e is None:
-                        logger.warning(
-                            f"Could not extend subgraph '{key}' due to missing entrypoint"
-                        )
-                        continue
-                    if s is not None:
-                        start_nodes[key] = s
-                    end_nodes[key] = e
-                else:
-                    nn = graph.add_node(node, key, metadata=metadata or None)
-                    start_nodes[key] = nn
-                    end_nodes[key] = nn
-            else:
-                nn = graph.add_node(node, key, metadata=metadata or None)
-                start_nodes[key] = nn
-                end_nodes[key] = nn
-        for start, end in sorted(self.builder._all_edges):
-            add_edge(start, end)
-        for start, branches in self.builder.branches.items():
-            default_ends = {
-                **{k: k for k in self.builder.nodes if k != start},
-                END: END,
-            }
-            for _, branch in branches.items():
-                if branch.ends is not None:
-                    ends = branch.ends
-                elif branch.then is not None:
-                    ends = {k: k for k in default_ends if k not in (END, branch.then)}
-                else:
-                    ends = cast(dict[Hashable, str], default_ends)
-                for label, end in ends.items():
-                    add_edge(
-                        start,
-                        end,
-                        label if label != end else None,
-                        conditional=True,
-                    )
-                    if branch.then is not None:
-                        add_edge(end, branch.then)
-        for key, n in self.builder.nodes.items():
-            if isinstance(n.ends, dict):
-                for end, label in n.ends.items():
-                    add_edge(key, end, label, conditional=True)
-            elif isinstance(n.ends, tuple):
-                for end in n.ends:
-                    add_edge(key, end, conditional=True)
-
-        return graph
-
-    def _repr_mimebundle_(self, **kwargs: Any) -> dict[str, Any]:
-        """Mime bundle used by Jupyter to display the graph"""
-        return {
-            "text/plain": repr(self),
-            "image/png": self.get_graph().draw_mermaid_png(),
-        }

libs/langgraph/langgraph/graph/message.py

@@ -1,5 +1,6 @@
 import uuid
 import warnings
+from collections.abc import Sequence
 from functools import partial
 from typing import (
     Annotated,
@@ -7,7 +8,6 @@ from typing import (
     Callable,
     Literal,
     Optional,
-    Sequence,
     Union,
     cast,
 )

libs/langgraph/langgraph/graph/schema_utils.py

@@ -3,10 +3,10 @@ import logging
 import weakref
 from inspect import isclass
 from typing import (
+    Annotated,
     Any,
     Callable,
     Optional,
-    Type,
     Union,
     get_args,
     get_origin,
@@ -15,14 +15,13 @@ from typing import (
 
 from pydantic import BaseModel
 from pydantic.v1 import BaseModel as BaseModelV1
-from typing_extensions import Annotated
 
 __all__ = ["SchemaCoercionMapper"]
 
 logger = logging.getLogger(__name__)
 
 
-_cache: weakref.WeakKeyDictionary[Type[Any], dict[int, "SchemaCoercionMapper"]] = (
+_cache: weakref.WeakKeyDictionary[type[Any], dict[int, "SchemaCoercionMapper"]] = (
     weakref.WeakKeyDictionary()
 )
 
@@ -32,7 +31,7 @@ class SchemaCoercionMapper:
 
     def __new__(
         cls,
-        schema: Type[Any],
+        schema: type[Any],
         type_hints: Optional[dict[str, Any]] = None,
         *,
         max_depth: int = 12,
@@ -46,7 +45,7 @@ class SchemaCoercionMapper:
 
     def __init__(
         self,
-        schema: Type[Any],
+        schema: type[Any],
         type_hints: Optional[dict[str, Any]] = None,
         *,
         max_depth: int = 12,
@@ -187,7 +186,7 @@ class SchemaCoercionMapper:
                 def dict_coercer(v: Any, d: Any) -> Any:
                     if not isinstance(v, dict):
                         if throw:
-                            raise TypeError("Expected dict, got %s" % type(v))
+                            raise TypeError(f"Expected dict, got {type(v)}")
                     return v
 
                 return dict_coercer
@@ -197,7 +196,7 @@ class SchemaCoercionMapper:
             def dict_coercer(v: Any, d: Any) -> Any:
                 if not isinstance(v, dict):
                     if throw:
-                        raise TypeError("Expected dict, got %s" % type(v))
+                        raise TypeError(f"Expected dict, got {type(v)}")
                     return v
                 return {k_sub(k, d - 1): v_sub(val, d - 1) for k, val in v.items()}
 

libs/langgraph/langgraph/graph/state.py

@@ -3,19 +3,16 @@ import logging
 import typing
 import warnings
 from collections import defaultdict
+from collections.abc import Awaitable, Hashable, Sequence
 from functools import partial
 from inspect import isclass, isfunction, ismethod, signature
 from types import FunctionType
 from typing import (
     Any,
-    Awaitable,
     Callable,
-    Hashable,
     Literal,
     NamedTuple,
     Optional,
-    Sequence,
-    Type,
     Union,
     cast,
     get_args,
@@ -84,7 +81,7 @@ from langgraph.utils.runnable import RunnableLike, coerce_to_runnable
 logger = logging.getLogger(__name__)
 
 
-def _warn_invalid_state_schema(schema: Union[Type[Any], Any]) -> None:
+def _warn_invalid_state_schema(schema: Union[type[Any], Any]) -> None:
     if isinstance(schema, type):
         return
     if typing.get_args(schema):
@@ -108,7 +105,7 @@ def _get_node_name(node: RunnableLike) -> str:
 class StateNodeSpec(NamedTuple):
     runnable: Runnable
     metadata: Optional[dict[str, Any]]
-    input: Type[Any]
+    input: type[Any]
     retry_policy: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]]
     ends: Optional[Union[tuple[str, ...], dict[str, str]]] = EMPTY_SEQ
 
@@ -166,15 +163,15 @@ class StateGraph(Graph):
     nodes: dict[str, StateNodeSpec]  # type: ignore[assignment]
     channels: dict[str, BaseChannel]
     managed: dict[str, ManagedValueSpec]
-    schemas: dict[Type[Any], dict[str, Union[BaseChannel, ManagedValueSpec]]]
+    schemas: dict[type[Any], dict[str, Union[BaseChannel, ManagedValueSpec]]]
 
     def __init__(
         self,
-        state_schema: Optional[Type[Any]] = None,
-        config_schema: Optional[Type[Any]] = None,
+        state_schema: Optional[type[Any]] = None,
+        config_schema: Optional[type[Any]] = None,
         *,
-        input: Optional[Type[Any]] = None,
-        output: Optional[Type[Any]] = None,
+        input: Optional[type[Any]] = None,
+        output: Optional[type[Any]] = None,
     ) -> None:
         super().__init__()
         if state_schema is None:
@@ -195,7 +192,7 @@ class StateGraph(Graph):
         self.schemas = {}
         self.channels = {}
         self.managed = {}
-        self.type_hints: dict[Type[Any], dict[str, Any]] = {}
+        self.type_hints: dict[type[Any], dict[str, Any]] = {}
         self.schema = state_schema
         self.input = input
         self.output = output
@@ -211,7 +208,7 @@ class StateGraph(Graph):
             (start, end) for starts, end in self.waiting_edges for start in starts
         }
 
-    def _add_schema(self, schema: Type[Any], /, allow_managed: bool = True) -> None:
+    def _add_schema(self, schema: type[Any], /, allow_managed: bool = True) -> None:
         if schema not in self.schemas:
             _warn_invalid_state_schema(schema)
             channels, managed, type_hints = _get_channels(schema)
@@ -250,7 +247,7 @@ class StateGraph(Graph):
         node: RunnableLike,
         *,
         metadata: Optional[dict[str, Any]] = None,
-        input: Optional[Type[Any]] = None,
+        input: Optional[type[Any]] = None,
         retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None,
         destinations: Optional[Union[dict[str, str], tuple[str, ...]]] = None,
     ) -> Self:
@@ -275,7 +272,7 @@ class StateGraph(Graph):
         action: RunnableLike,
         *,
         metadata: Optional[dict[str, Any]] = None,
-        input: Optional[Type[Any]] = None,
+        input: Optional[type[Any]] = None,
         retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None,
         destinations: Optional[Union[dict[str, str], tuple[str, ...]]] = None,
     ) -> Self:
@@ -299,7 +296,7 @@ class StateGraph(Graph):
         action: Optional[RunnableLike] = None,
         *,
         metadata: Optional[dict[str, Any]] = None,
-        input: Optional[Type[Any]] = None,
+        input: Optional[type[Any]] = None,
         retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None,
         destinations: Optional[Union[dict[str, str], tuple[str, ...]]] = None,
     ) -> Self:
@@ -527,7 +524,7 @@ class StateGraph(Graph):
         # validate the condition
         if name in self.branches[source]:
             raise ValueError(
-                f"Branch with name `{path.name}` already exists for node " f"`{source}`"
+                f"Branch with name `{path.name}` already exists for node `{source}`"
             )
         # save it
         self.branches[source][name] = Branch.from_path(path, path_map, then, True)
@@ -686,12 +683,12 @@ class StateGraph(Graph):
 
 class CompiledStateGraph(CompiledGraph):
     builder: StateGraph
-    schema_to_mapper: dict[Type[Any], Optional[Callable[[Any], Any]]]
+    schema_to_mapper: dict[type[Any], Optional[Callable[[Any], Any]]]
 
     def __init__(
         self,
         *,
-        schema_to_mapper: dict[Type[Any], Optional[Callable[[Any], Any]]],
+        schema_to_mapper: dict[type[Any], Optional[Callable[[Any], Any]]],
         **kwargs: Any,
     ) -> None:
         super().__init__(**kwargs)
@@ -774,7 +771,12 @@ class CompiledStateGraph(CompiledGraph):
             ChannelWriteTupleEntry(
                 mapper=_get_root if output_keys == ["__root__"] else _get_updates
             ),
-            ChannelWriteTupleEntry(mapper=_control_branch),
+            ChannelWriteTupleEntry(
+                mapper=_control_branch,
+                static=_control_static(node.ends)
+                if node is not None and node.ends is not None
+                else None,
+            ),
         )
 
         # add node and output channel
@@ -840,9 +842,9 @@ class CompiledStateGraph(CompiledGraph):
     def attach_branch(
         self, start: str, name: str, branch: Branch, *, with_reader: bool = True
     ) -> None:
-        def branch_writer(
-            packets: Sequence[Union[str, Send]], config: RunnableConfig
-        ) -> None:
+        def get_writes(
+            packets: Sequence[Union[str, Send]],
+        ) -> Sequence[Union[ChannelWriteEntry, Send]]:
             if filtered := [p for p in packets if p != END]:
                 writes = [
                     (
@@ -857,13 +859,15 @@ class CompiledStateGraph(CompiledGraph):
                         ChannelWriteEntry(
                             f"branch:{start}:{name}::then",
                             WaitForNames(
-                                {p.node if isinstance(p, Send) else p for p in filtered}
+                                frozenset(
+                                    p.node if isinstance(p, Send) else p
+                                    for p in filtered
+                                )
                             ),
                         )
                     )
-                ChannelWrite.do_write(
-                    config, cast(Sequence[Union[Send, ChannelWriteEntry]], writes)
-                )
+                return writes
+            return []
 
         if with_reader:
             # get schema
@@ -891,7 +895,7 @@ class CompiledStateGraph(CompiledGraph):
             reader = None
 
         # attach branch publisher
-        self.nodes[start].writers.append(branch.run(branch_writer, reader))
+        self.nodes[start].writers.append(branch.run(get_writes, reader))
 
         # attach then subscriber
         if branch.then and branch.then != END:
@@ -1015,7 +1019,7 @@ class CompiledStateGraph(CompiledGraph):
 
 
 def _pick_mapper(
-    state_keys: Sequence[str], schema: Type[Any], type_hints: Optional[dict[str, Any]]
+    state_keys: Sequence[str], schema: type[Any], type_hints: Optional[dict[str, Any]]
 ) -> Optional[Callable[[Any], Any]]:
     if state_keys == ["__root__"]:
         return None
@@ -1027,7 +1031,7 @@ def _pick_mapper(
     return partial(_coerce_state, schema)
 
 
-def _coerce_state(schema: Type[Any], input: dict[str, Any]) -> dict[str, Any]:
+def _coerce_state(schema: type[Any], input: dict[str, Any]) -> dict[str, Any]:
     return schema(**input)
 
 
@@ -1059,6 +1063,19 @@ def _control_branch(value: Any) -> Sequence[tuple[str, Any]]:
     return rtn
 
 
+def _control_static(
+    ends: Union[tuple[str, ...], dict[str, str]],
+) -> Sequence[tuple[str, Any, Optional[str]]]:
+    if isinstance(ends, dict):
+        return [
+            (CHANNEL_BRANCH_TO.format(k), None, label)
+            for k, label in ends.items()
+            if k != END
+        ]
+    else:
+        return [(CHANNEL_BRANCH_TO.format(e), None, None) for e in ends if e != END]
+
+
 def _get_root(input: Any) -> Optional[Sequence[tuple[str, Any]]]:
     if isinstance(input, Command):
         if input.graph == Command.PARENT:
@@ -1083,7 +1100,7 @@ def _get_root(input: Any) -> Optional[Sequence[tuple[str, Any]]]:
 
 
 def _get_channels(
-    schema: Type[dict],
+    schema: type[dict],
 ) -> tuple[dict[str, BaseChannel], dict[str, ManagedValueSpec], dict[str, Any]]:
     if not hasattr(schema, "__annotations__"):
         return (
@@ -1137,7 +1154,7 @@ def _get_channel(
     return fallback
 
 
-def _is_field_channel(typ: Type[Any]) -> Optional[BaseChannel]:
+def _is_field_channel(typ: type[Any]) -> Optional[BaseChannel]:
     if hasattr(typ, "__metadata__"):
         meta = typ.__metadata__
         if len(meta) >= 1 and isinstance(meta[-1], BaseChannel):
@@ -1147,7 +1164,7 @@ def _is_field_channel(typ: Type[Any]) -> Optional[BaseChannel]:
     return None
 
 
-def _is_field_binop(typ: Type[Any]) -> Optional[BinaryOperatorAggregate]:
+def _is_field_binop(typ: type[Any]) -> Optional[BinaryOperatorAggregate]:
     if hasattr(typ, "__metadata__"):
         meta = typ.__metadata__
         if len(meta) >= 1 and callable(meta[-1]):
@@ -1168,7 +1185,7 @@ def _is_field_binop(typ: Type[Any]) -> Optional[BinaryOperatorAggregate]:
     return None
 
 
-def _is_field_managed_value(name: str, typ: Type[Any]) -> Optional[ManagedValueSpec]:
+def _is_field_managed_value(name: str, typ: type[Any]) -> Optional[ManagedValueSpec]:
     if hasattr(typ, "__metadata__"):
         meta = typ.__metadata__
         if len(meta) >= 1:
@@ -1186,7 +1203,7 @@ def _is_field_managed_value(name: str, typ: Type[Any]) -> Optional[ManagedValueS
 
 
 def _get_schema(
-    typ: Type,
+    typ: type,
     schemas: dict,
     channels: dict,
     name: str,

libs/langgraph/langgraph/managed/base.py

@@ -1,14 +1,11 @@
 from abc import ABC, abstractmethod
+from collections.abc import AsyncIterator, Iterator, Sequence
 from contextlib import asynccontextmanager, contextmanager
 from inspect import isclass
 from typing import (
     Any,
-    AsyncIterator,
     Generic,
-    Iterator,
     NamedTuple,
-    Sequence,
-    Type,
     TypeVar,
     Union,
 )
@@ -66,11 +63,11 @@ class WritableManagedValue(Generic[V, U], ManagedValue[V], ABC):
 
 
 class ConfiguredManagedValue(NamedTuple):
-    cls: Type[ManagedValue]
+    cls: type[ManagedValue]
     kwargs: dict[str, Any]
 
 
-ManagedValueSpec = Union[Type[ManagedValue], ConfiguredManagedValue]
+ManagedValueSpec = Union[type[ManagedValue], ConfiguredManagedValue]
 
 
 def is_managed_value(value: Any) -> TypeGuard[ManagedValueSpec]:
@@ -79,7 +76,7 @@ def is_managed_value(value: Any) -> TypeGuard[ManagedValueSpec]:
     )
 
 
-def is_readonly_managed_value(value: Any) -> TypeGuard[Type[ManagedValue]]:
+def is_readonly_managed_value(value: Any) -> TypeGuard[type[ManagedValue]]:
     return (
         isclass(value)
         and issubclass(value, ManagedValue)
@@ -90,7 +87,7 @@ def is_readonly_managed_value(value: Any) -> TypeGuard[Type[ManagedValue]]:
     )
 
 
-def is_writable_managed_value(value: Any) -> TypeGuard[Type[WritableManagedValue]]:
+def is_writable_managed_value(value: Any) -> TypeGuard[type[WritableManagedValue]]:
     return (isclass(value) and issubclass(value, WritableManagedValue)) or (
         isinstance(value, ConfiguredManagedValue)
         and issubclass(value.cls, WritableManagedValue)

libs/langgraph/langgraph/managed/context.py

@@ -1,15 +1,16 @@
-from contextlib import asynccontextmanager, contextmanager
+from collections.abc import AsyncIterator, Iterator
+from contextlib import (
+    AbstractAsyncContextManager,
+    AbstractContextManager,
+    asynccontextmanager,
+    contextmanager,
+)
 from inspect import signature
 from typing import (
     Any,
-    AsyncContextManager,
-    AsyncIterator,
     Callable,
-    ContextManager,
     Generic,
-    Iterator,
     Optional,
-    Type,
     Union,
 )
 
@@ -28,15 +29,15 @@ class Context(ManagedValue[V], Generic[V]):
     def of(
         ctx: Union[
             None,
-            Callable[..., ContextManager[V]],
-            Type[ContextManager[V]],
-            Callable[..., AsyncContextManager[V]],
-            Type[AsyncContextManager[V]],
+            Callable[..., AbstractContextManager[V]],
+            type[AbstractContextManager[V]],
+            Callable[..., AbstractAsyncContextManager[V]],
+            type[AbstractAsyncContextManager[V]],
         ] = None,
         actx: Optional[
             Union[
-                Callable[..., AsyncContextManager[V]],
-                Type[AsyncContextManager[V]],
+                Callable[..., AbstractAsyncContextManager[V]],
+                type[AbstractAsyncContextManager[V]],
             ]
         ] = None,
     ) -> ConfiguredManagedValue:
@@ -98,8 +99,10 @@ class Context(ManagedValue[V], Generic[V]):
         self,
         loop: LoopProtocol,
         *,
-        ctx: Union[None, Type[ContextManager[V]], Type[AsyncContextManager[V]]] = None,
-        actx: Optional[Type[AsyncContextManager[V]]] = None,
+        ctx: Union[
+            None, type[AbstractContextManager[V]], type[AbstractAsyncContextManager[V]]
+        ] = None,
+        actx: Optional[type[AbstractAsyncContextManager[V]]] = None,
     ) -> None:
         self.ctx = ctx
         self.actx = actx

libs/langgraph/langgraph/managed/shared_value.py

@@ -1,12 +1,9 @@
 import collections.abc
+from collections.abc import AsyncIterator, Iterator, Sequence
 from contextlib import asynccontextmanager, contextmanager
 from typing import (
     Any,
-    AsyncIterator,
-    Iterator,
     Optional,
-    Sequence,
-    Type,
 )
 
 from typing_extensions import NotRequired, Required, Self
@@ -71,7 +68,7 @@ class SharedValue(WritableManagedValue[Value, Update]):
             yield value
 
     def __init__(
-        self, loop: LoopProtocol, *, typ: Type[Any], scope: str, key: str
+        self, loop: LoopProtocol, *, typ: type[Any], scope: str, key: str
     ) -> None:
         super().__init__(loop)
         if typ := _strip_extras(typ):

libs/langgraph/langgraph/pregel/__init__.py

@@ -6,17 +6,11 @@ import concurrent.futures
 import queue
 import weakref
 from collections import defaultdict, deque
+from collections.abc import AsyncIterator, Iterator, Mapping, Sequence
 from functools import partial
 from typing import (
     Any,
-    AsyncIterator,
     Callable,
-    Dict,
-    Iterator,
-    Mapping,
-    Optional,
-    Sequence,
-    Type,
     Union,
     cast,
     get_type_hints,
@@ -93,6 +87,7 @@ from langgraph.pregel.algo import (
 )
 from langgraph.pregel.checkpoint import create_checkpoint, empty_checkpoint
 from langgraph.pregel.debug import tasks_w_writes
+from langgraph.pregel.draw import draw_graph
 from langgraph.pregel.io import map_input, read_channels
 from langgraph.pregel.loop import AsyncPregelLoop, StreamProtocol, SyncPregelLoop
 from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager
@@ -141,8 +136,8 @@ class Channel:
         cls,
         channels: str,
         *,
-        key: Optional[str] = None,
-        tags: Optional[list[str]] = None,
+        key: str | None = None,
+        tags: list[str] | None = None,
     ) -> PregelNode: ...
 
     @overload
@@ -152,16 +147,16 @@ class Channel:
         channels: Sequence[str],
         *,
         key: None = None,
-        tags: Optional[list[str]] = None,
+        tags: list[str] | None = None,
     ) -> PregelNode: ...
 
     @classmethod
     def subscribe_to(
         cls,
-        channels: Union[str, Sequence[str]],
+        channels: str | Sequence[str],
         *,
-        key: Optional[str] = None,
-        tags: Optional[list[str]] = None,
+        key: str | None = None,
+        tags: list[str] | None = None,
     ) -> PregelNode:
         """Runs process.invoke() each time channels are updated,
         with a dict of the channel values as input."""
@@ -467,7 +462,7 @@ class Pregel(PregelProtocol):
 
     nodes: dict[str, PregelNode]
 
-    channels: dict[str, Union[BaseChannel, ManagedValueSpec]]
+    channels: dict[str, BaseChannel | ManagedValueSpec]
 
     stream_mode: StreamMode = "values"
     """Mode to stream output, defaults to 'values'."""
@@ -476,18 +471,18 @@ class Pregel(PregelProtocol):
     """Whether to force emitting stream events eagerly, automatically turned on
     for stream_mode "messages" and "custom"."""
 
-    output_channels: Union[str, Sequence[str]]
+    output_channels: str | Sequence[str]
 
-    stream_channels: Optional[Union[str, Sequence[str]]] = None
+    stream_channels: str | Sequence[str] | None = None
     """Channels to stream, defaults to all channels not in reserved channels"""
 
-    interrupt_after_nodes: Union[All, Sequence[str]]
+    interrupt_after_nodes: All | Sequence[str]
 
-    interrupt_before_nodes: Union[All, Sequence[str]]
+    interrupt_before_nodes: All | Sequence[str]
 
-    input_channels: Union[str, Sequence[str]]
+    input_channels: str | Sequence[str]
 
-    step_timeout: Optional[float] = None
+    step_timeout: float | None = None
     """Maximum time to wait for a step to complete, in seconds. Defaults to None."""
 
     debug: bool
@@ -496,44 +491,44 @@ class Pregel(PregelProtocol):
     checkpointer: Checkpointer = None
     """Checkpointer used to save and load graph state. Defaults to None."""
 
-    store: Optional[BaseStore] = None
+    store: BaseStore | None = None
     """Memory store to use for SharedValues. Defaults to None."""
 
-    retry_policy: Optional[Sequence[RetryPolicy]] = None
+    retry_policy: Sequence[RetryPolicy] | None = None
     """Retry policies to use when running tasks. Set to None to disable."""
 
-    config_type: Optional[Type[Any]] = None
+    config_type: type[Any] | None = None
 
-    input_model: Optional[Type[BaseModel]] = None
+    input_model: type[BaseModel] | None = None
 
-    config: Optional[RunnableConfig] = None
+    config: RunnableConfig | None = None
 
     name: str = "LangGraph"
 
-    trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None
+    trigger_to_nodes: Mapping[str, Sequence[str]] | None = None
 
     def __init__(
         self,
         *,
         nodes: dict[str, PregelNode],
-        channels: Optional[dict[str, Union[BaseChannel, ManagedValueSpec]]],
+        channels: dict[str, BaseChannel | ManagedValueSpec] | None,
         auto_validate: bool = True,
         stream_mode: StreamMode = "values",
         stream_eager: bool = False,
-        output_channels: Union[str, Sequence[str]],
-        stream_channels: Optional[Union[str, Sequence[str]]] = None,
-        interrupt_after_nodes: Union[All, Sequence[str]] = (),
-        interrupt_before_nodes: Union[All, Sequence[str]] = (),
-        input_channels: Union[str, Sequence[str]],
-        step_timeout: Optional[float] = None,
-        debug: Optional[bool] = None,
-        checkpointer: Optional[BaseCheckpointSaver] = None,
-        store: Optional[BaseStore] = None,
-        retry_policy: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None,
-        config_type: Optional[Type[Any]] = None,
-        input_model: Optional[Type[BaseModel]] = None,
-        config: Optional[RunnableConfig] = None,
-        trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None,
+        output_channels: str | Sequence[str],
+        stream_channels: str | Sequence[str] | None = None,
+        interrupt_after_nodes: All | Sequence[str] = (),
+        interrupt_before_nodes: All | Sequence[str] = (),
+        input_channels: str | Sequence[str],
+        step_timeout: float | None = None,
+        debug: bool | None = None,
+        checkpointer: BaseCheckpointSaver | None = None,
+        store: BaseStore | None = None,
+        retry_policy: RetryPolicy | Sequence[RetryPolicy] | None = None,
+        config_type: type[Any] | None = None,
+        input_model: type[BaseModel] | None = None,
+        config: RunnableConfig | None = None,
+        trigger_to_nodes: Mapping[str, Sequence[str]] | None = None,
         name: str = "LangGraph",
     ) -> None:
         self.nodes = nodes
@@ -562,22 +557,87 @@ class Pregel(PregelProtocol):
             self.validate()
 
     def get_graph(
-        self, config: Optional[RunnableConfig] = None, *, xray: Union[int, bool] = False
+        self, config: RunnableConfig | None = None, *, xray: int | bool = False
     ) -> Graph:
-        raise NotImplementedError
+        """Returns a drawable representation of the computation graph."""
+        # gather subgraphs
+        if xray:
+            subgraphs = {
+                k: v.get_graph(
+                    config,
+                    xray=xray if isinstance(xray, bool) or xray <= 0 else xray - 1,
+                )
+                for k, v in self.get_subgraphs()
+            }
+        else:
+            subgraphs = {}
+
+        return draw_graph(
+            merge_configs(self.config, config),
+            nodes=self.nodes,
+            specs=self.channels,
+            input_channels=self.input_channels,
+            interrupt_after_nodes=self.interrupt_after_nodes,
+            interrupt_before_nodes=self.interrupt_before_nodes,
+            trigger_to_nodes=self.trigger_to_nodes,
+            checkpointer=self.checkpointer,
+            subgraphs=subgraphs,
+        )
 
     async def aget_graph(
-        self, config: Optional[RunnableConfig] = None, *, xray: Union[int, bool] = False
+        self, config: RunnableConfig | None = None, *, xray: int | bool = False
     ) -> Graph:
-        raise NotImplementedError
+        """Returns a drawable representation of the computation graph."""
 
-    def copy(self, update: Optional[dict[str, Any]] = None) -> Self:
+        # gather subgraphs
+        if xray:
+            subpregels: dict[str, PregelProtocol] = {
+                k: v async for k, v in self.aget_subgraphs()
+            }
+            subgraphs = {
+                k: v
+                for k, v in zip(
+                    subpregels,
+                    await asyncio.gather(
+                        *(
+                            p.aget_graph(
+                                config,
+                                xray=xray
+                                if isinstance(xray, bool) or xray <= 0
+                                else xray - 1,
+                            )
+                            for p in subpregels.values()
+                        )
+                    ),
+                )
+            }
+        else:
+            subgraphs = {}
+
+        return draw_graph(
+            merge_configs(self.config, config),
+            nodes=self.nodes,
+            specs=self.channels,
+            input_channels=self.input_channels,
+            interrupt_after_nodes=self.interrupt_after_nodes,
+            interrupt_before_nodes=self.interrupt_before_nodes,
+            trigger_to_nodes=self.trigger_to_nodes,
+            checkpointer=self.checkpointer,
+            subgraphs=subgraphs,
+        )
+
+    def _repr_mimebundle_(self, **kwargs: Any) -> dict[str, Any]:
+        """Mime bundle used by Jupyter to display the graph"""
+        return {
+            "text/plain": repr(self),
+            "image/png": self.get_graph().draw_mermaid_png(),
+        }
+
+    def copy(self, update: dict[str, Any] | None = None) -> Self:
         attrs = {**self.__dict__, **(update or {})}
         return self.__class__(**attrs)
 
-    def with_config(
-        self, config: Optional[RunnableConfig] = None, **kwargs: Any
-    ) -> Self:
+    def with_config(self, config: RunnableConfig | None = None, **kwargs: Any) -> Self:
         return self.copy(
             {"config": merge_configs(self.config, config, cast(RunnableConfig, kwargs))}
         )
@@ -632,9 +692,7 @@ class Pregel(PregelProtocol):
             ]
         ]
 
-    def config_schema(
-        self, *, include: Optional[Sequence[str]] = None
-    ) -> Type[BaseModel]:
+    def config_schema(self, *, include: Sequence[str] | None = None) -> type[BaseModel]:
         # If the config type is not set explicitly, we will try to infer it.
         # If the config type is provided, but isn't directly supported by pydantic
         # (e.g., vanilla python class), we will also delegate to the parent class,
@@ -654,8 +712,8 @@ class Pregel(PregelProtocol):
         return create_model(self.get_name("Config"), field_definitions=fields)
 
     def get_config_jsonschema(
-        self, *, include: Optional[Sequence[str]] = None
-    ) -> Dict[str, Any]:
+        self, *, include: Sequence[str] | None = None
+    ) -> dict[str, Any]:
         schema = self.config_schema(include=include)
         if hasattr(schema, "model_json_schema"):
             return schema.model_json_schema()
@@ -669,9 +727,7 @@ class Pregel(PregelProtocol):
             if isinstance(channel, BaseChannel):
                 return channel.UpdateType
 
-    def get_input_schema(
-        self, config: Optional[RunnableConfig] = None
-    ) -> Type[BaseModel]:
+    def get_input_schema(self, config: RunnableConfig | None = None) -> type[BaseModel]:
         if self.input_model is not None:
             return self.input_model
         config = merge_configs(self.config, config)
@@ -688,8 +744,8 @@ class Pregel(PregelProtocol):
             )
 
     def get_input_jsonschema(
-        self, config: Optional[RunnableConfig] = None
-    ) -> Dict[str, Any]:
+        self, config: RunnableConfig | None = None
+    ) -> dict[str, Any]:
         schema = self.get_input_schema(config)
         if hasattr(schema, "model_json_schema"):
             return schema.model_json_schema()
@@ -704,8 +760,8 @@ class Pregel(PregelProtocol):
                 return channel.ValueType
 
     def get_output_schema(
-        self, config: Optional[RunnableConfig] = None
-    ) -> Type[BaseModel]:
+        self, config: RunnableConfig | None = None
+    ) -> type[BaseModel]:
         config = merge_configs(self.config, config)
         if isinstance(self.output_channels, str):
             return super().get_output_schema(config)
@@ -720,8 +776,8 @@ class Pregel(PregelProtocol):
             )
 
     def get_output_jsonschema(
-        self, config: Optional[RunnableConfig] = None
-    ) -> Dict[str, Any]:
+        self, config: RunnableConfig | None = None
+    ) -> dict[str, Any]:
         schema = self.get_output_schema(config)
         if hasattr(schema, "model_json_schema"):
             return schema.model_json_schema()
@@ -736,13 +792,13 @@ class Pregel(PregelProtocol):
         )
 
     @property
-    def stream_channels_asis(self) -> Union[str, Sequence[str]]:
+    def stream_channels_asis(self) -> str | Sequence[str]:
         return self.stream_channels or [
             k for k in self.channels if isinstance(self.channels[k], BaseChannel)
         ]
 
     def get_subgraphs(
-        self, *, namespace: Optional[str] = None, recurse: bool = False
+        self, *, namespace: str | None = None, recurse: bool = False
     ) -> Iterator[tuple[str, PregelProtocol]]:
         for name, node in self.nodes.items():
             # filter by prefix
@@ -771,7 +827,7 @@ class Pregel(PregelProtocol):
                     )
 
     async def aget_subgraphs(
-        self, *, namespace: Optional[str] = None, recurse: bool = False
+        self, *, namespace: str | None = None, recurse: bool = False
     ) -> AsyncIterator[tuple[str, PregelProtocol]]:
         for name, node in self.get_subgraphs(namespace=namespace, recurse=recurse):
             yield name, node
@@ -783,8 +839,8 @@ class Pregel(PregelProtocol):
     def _prepare_state_snapshot(
         self,
         config: RunnableConfig,
-        saved: Optional[CheckpointTuple],
-        recurse: Optional[BaseCheckpointSaver] = None,
+        saved: CheckpointTuple | None,
+        recurse: BaseCheckpointSaver | None = None,
         apply_pending_writes: bool = False,
     ) -> StateSnapshot:
         if not saved:
@@ -832,7 +888,7 @@ class Pregel(PregelProtocol):
             # get the subgraphs
             subgraphs = dict(self.get_subgraphs())
             parent_ns = saved.config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "")
-            task_states: dict[str, Union[RunnableConfig, StateSnapshot]] = {}
+            task_states: dict[str, RunnableConfig | StateSnapshot] = {}
             for task in next_tasks.values():
                 if task.name not in subgraphs:
                     continue
@@ -899,8 +955,8 @@ class Pregel(PregelProtocol):
     async def _aprepare_state_snapshot(
         self,
         config: RunnableConfig,
-        saved: Optional[CheckpointTuple],
-        recurse: Optional[BaseCheckpointSaver] = None,
+        saved: CheckpointTuple | None,
+        recurse: BaseCheckpointSaver | None = None,
         apply_pending_writes: bool = False,
     ) -> StateSnapshot:
         if not saved:
@@ -951,7 +1007,7 @@ class Pregel(PregelProtocol):
             # get the subgraphs
             subgraphs = {n: g async for n, g in self.aget_subgraphs()}
             parent_ns = saved.config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "")
-            task_states: dict[str, Union[RunnableConfig, StateSnapshot]] = {}
+            task_states: dict[str, RunnableConfig | StateSnapshot] = {}
             for task in next_tasks.values():
                 if task.name not in subgraphs:
                     continue
@@ -1019,7 +1075,7 @@ class Pregel(PregelProtocol):
         self, config: RunnableConfig, *, subgraphs: bool = False
     ) -> StateSnapshot:
         """Get the current state of the graph."""
-        checkpointer: Optional[BaseCheckpointSaver] = ensure_config(config)[CONF].get(
+        checkpointer: BaseCheckpointSaver | None = ensure_config(config)[CONF].get(
             CONFIG_KEY_CHECKPOINTER, self.checkpointer
         )
         if not checkpointer:
@@ -1061,7 +1117,7 @@ class Pregel(PregelProtocol):
         self, config: RunnableConfig, *, subgraphs: bool = False
     ) -> StateSnapshot:
         """Get the current state of the graph."""
-        checkpointer: Optional[BaseCheckpointSaver] = ensure_config(config)[CONF].get(
+        checkpointer: BaseCheckpointSaver | None = ensure_config(config)[CONF].get(
             CONFIG_KEY_CHECKPOINTER, self.checkpointer
         )
         if not checkpointer:
@@ -1103,13 +1159,13 @@ class Pregel(PregelProtocol):
         self,
         config: RunnableConfig,
         *,
-        filter: Optional[Dict[str, Any]] = None,
-        before: Optional[RunnableConfig] = None,
-        limit: Optional[int] = None,
+        filter: dict[str, Any] | None = None,
+        before: RunnableConfig | None = None,
+        limit: int | None = None,
     ) -> Iterator[StateSnapshot]:
         config = ensure_config(config)
         """Get the history of the state of the graph."""
-        checkpointer: Optional[BaseCheckpointSaver] = ensure_config(config)[CONF].get(
+        checkpointer: BaseCheckpointSaver | None = ensure_config(config)[CONF].get(
             CONFIG_KEY_CHECKPOINTER, self.checkpointer
         )
         if not checkpointer:
@@ -1154,13 +1210,13 @@ class Pregel(PregelProtocol):
         self,
         config: RunnableConfig,
         *,
-        filter: Optional[Dict[str, Any]] = None,
-        before: Optional[RunnableConfig] = None,
-        limit: Optional[int] = None,
+        filter: dict[str, Any] | None = None,
+        before: RunnableConfig | None = None,
+        limit: int | None = None,
     ) -> AsyncIterator[StateSnapshot]:
         config = ensure_config(config)
         """Get the history of the state of the graph."""
-        checkpointer: Optional[BaseCheckpointSaver] = ensure_config(config)[CONF].get(
+        checkpointer: BaseCheckpointSaver | None = ensure_config(config)[CONF].get(
             CONFIG_KEY_CHECKPOINTER, self.checkpointer
         )
         if not checkpointer:
@@ -1225,7 +1281,7 @@ class Pregel(PregelProtocol):
             RunnableConfig: The updated config.
         """
 
-        checkpointer: Optional[BaseCheckpointSaver] = ensure_config(config)[CONF].get(
+        checkpointer: BaseCheckpointSaver | None = ensure_config(config)[CONF].get(
             CONFIG_KEY_CHECKPOINTER, self.checkpointer
         )
         if not checkpointer:
@@ -1500,7 +1556,7 @@ class Pregel(PregelProtocol):
                         next_tasks[tid].writes.append((k, v))
                     if tasks := [t for t in next_tasks.values() if t.writes]:
                         apply_writes(checkpoint, channels, tasks, None)
-            valid_updates: list[tuple[str, Optional[dict[str, Any]]]] = []
+            valid_updates: list[tuple[str, dict[str, Any] | None]] = []
             if len(updates) == 1:
                 values, as_node = updates[0]
                 # find last node that updated the state, if not provided
@@ -1639,7 +1695,7 @@ class Pregel(PregelProtocol):
             RunnableConfig: The updated config.
         """
 
-        checkpointer: Optional[BaseCheckpointSaver] = ensure_config(config)[CONF].get(
+        checkpointer: BaseCheckpointSaver | None = ensure_config(config)[CONF].get(
             CONFIG_KEY_CHECKPOINTER, self.checkpointer
         )
         if not checkpointer:
@@ -1914,7 +1970,7 @@ class Pregel(PregelProtocol):
                         next_tasks[tid].writes.append((k, v))
                     if tasks := [t for t in next_tasks.values() if t.writes]:
                         apply_writes(checkpoint, channels, tasks, None)
-            valid_updates: list[tuple[str, Optional[dict[str, Any]]]] = []
+            valid_updates: list[tuple[str, dict[str, Any] | None]] = []
             if len(updates) == 1:
                 values, as_node = updates[0]
                 # find last node that updated the state, if not provided
@@ -2034,8 +2090,8 @@ class Pregel(PregelProtocol):
     def update_state(
         self,
         config: RunnableConfig,
-        values: Optional[Union[dict[str, Any], Any]],
-        as_node: Optional[str] = None,
+        values: dict[str, Any] | Any | None,
+        as_node: str | None = None,
     ) -> RunnableConfig:
         """Update the state of the graph with the given values, as if they came from
         node `as_node`. If `as_node` is not provided, it will be set to the last node
@@ -2047,7 +2103,7 @@ class Pregel(PregelProtocol):
         self,
         config: RunnableConfig,
         values: dict[str, Any] | Any,
-        as_node: Optional[str] = None,
+        as_node: str | None = None,
     ) -> RunnableConfig:
         """Update the state of the graph asynchronously with the given values, as if they came from
         node `as_node`. If `as_node` is not provided, it will be set to the last node
@@ -2059,19 +2115,19 @@ class Pregel(PregelProtocol):
         self,
         config: RunnableConfig,
         *,
-        stream_mode: Optional[Union[StreamMode, list[StreamMode]]],
-        output_keys: Optional[Union[str, Sequence[str]]],
-        interrupt_before: Optional[Union[All, Sequence[str]]],
-        interrupt_after: Optional[Union[All, Sequence[str]]],
-        debug: Optional[bool],
+        stream_mode: StreamMode | list[StreamMode] | None,
+        output_keys: str | Sequence[str] | None,
+        interrupt_before: All | Sequence[str] | None,
+        interrupt_after: All | Sequence[str] | None,
+        debug: bool | None,
     ) -> tuple[
         bool,
         set[StreamMode],
-        Union[str, Sequence[str]],
-        Union[All, Sequence[str]],
-        Union[All, Sequence[str]],
-        Optional[BaseCheckpointSaver],
-        Optional[BaseStore],
+        str | Sequence[str],
+        All | Sequence[str],
+        All | Sequence[str],
+        BaseCheckpointSaver | None,
+        BaseStore | None,
     ]:
         if config["recursion_limit"] < 1:
             raise ValueError("recursion_limit must be at least 1")
@@ -2089,7 +2145,7 @@ class Pregel(PregelProtocol):
             # if being called as a node in another graph, always use values mode
             stream_mode = ["values"]
         if self.checkpointer is False:
-            checkpointer: Optional[BaseCheckpointSaver] = None
+            checkpointer: BaseCheckpointSaver | None = None
         elif CONFIG_KEY_CHECKPOINTER in config.get(CONF, {}):
             checkpointer = config[CONF][CONFIG_KEY_CHECKPOINTER]
         elif self.checkpointer is True:
@@ -2101,7 +2157,7 @@ class Pregel(PregelProtocol):
                 f"Checkpointer requires one or more of the following 'configurable' keys: {[s.id for s in checkpointer.config_specs]}"
             )
         if CONFIG_KEY_STORE in config.get(CONF, {}):
-            store: Optional[BaseStore] = config[CONF][CONFIG_KEY_STORE]
+            store: BaseStore | None = config[CONF][CONFIG_KEY_STORE]
         else:
             store = self.store
         return (
@@ -2116,17 +2172,17 @@ class Pregel(PregelProtocol):
 
     def stream(
         self,
-        input: Union[dict[str, Any], Any],
-        config: Optional[RunnableConfig] = None,
+        input: dict[str, Any] | Any,
+        config: RunnableConfig | None = None,
         *,
-        stream_mode: Optional[Union[StreamMode, list[StreamMode]]] = None,
-        output_keys: Optional[Union[str, Sequence[str]]] = None,
-        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
-        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
-        checkpoint_during: Optional[bool] = None,
-        debug: Optional[bool] = None,
+        stream_mode: StreamMode | list[StreamMode] | None = None,
+        output_keys: str | Sequence[str] | None = None,
+        interrupt_before: All | Sequence[str] | None = None,
+        interrupt_after: All | Sequence[str] | None = None,
+        checkpoint_during: bool | None = None,
+        debug: bool | None = None,
         subgraphs: bool = False,
-    ) -> Iterator[Union[dict[str, Any], Any]]:
+    ) -> Iterator[dict[str, Any] | Any]:
         """Stream graph steps for a single input.
 
         Args:
@@ -2352,7 +2408,7 @@ class Pregel(PregelProtocol):
                 ):
                     # we are careful to have a single waiter live at any one time
                     # because on exit we increment semaphore count by exactly 1
-                    waiter: Optional[concurrent.futures.Future] = None
+                    waiter: concurrent.futures.Future | None = None
                     # because sync futures cannot be cancelled, we instead
                     # release the stream semaphore on exit, which will cause
                     # a pending waiter to return immediately
@@ -2403,17 +2459,17 @@ class Pregel(PregelProtocol):
 
     async def astream(
         self,
-        input: Union[dict[str, Any], Any],
-        config: Optional[RunnableConfig] = None,
+        input: dict[str, Any] | Any,
+        config: RunnableConfig | None = None,
         *,
-        stream_mode: Optional[Union[StreamMode, list[StreamMode]]] = None,
-        output_keys: Optional[Union[str, Sequence[str]]] = None,
-        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
-        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
-        checkpoint_during: Optional[bool] = None,
-        debug: Optional[bool] = None,
+        stream_mode: StreamMode | list[StreamMode] | None = None,
+        output_keys: str | Sequence[str] | None = None,
+        interrupt_before: All | Sequence[str] | None = None,
+        interrupt_after: All | Sequence[str] | None = None,
+        checkpoint_during: bool | None = None,
+        debug: bool | None = None,
         subgraphs: bool = False,
-    ) -> AsyncIterator[Union[dict[str, Any], Any]]:
+    ) -> AsyncIterator[dict[str, Any] | Any]:
         """Stream graph steps for a single input.
 
         Args:
@@ -2704,17 +2760,17 @@ class Pregel(PregelProtocol):
 
     def invoke(
         self,
-        input: Union[dict[str, Any], Any],
-        config: Optional[RunnableConfig] = None,
+        input: dict[str, Any] | Any,
+        config: RunnableConfig | None = None,
         *,
         stream_mode: StreamMode = "values",
-        output_keys: Optional[Union[str, Sequence[str]]] = None,
-        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
-        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
-        checkpoint_during: Optional[bool] = None,
-        debug: Optional[bool] = None,
+        output_keys: str | Sequence[str] | None = None,
+        interrupt_before: All | Sequence[str] | None = None,
+        interrupt_after: All | Sequence[str] | None = None,
+        checkpoint_during: bool | None = None,
+        debug: bool | None = None,
         **kwargs: Any,
-    ) -> Union[dict[str, Any], Any]:
+    ) -> dict[str, Any] | Any:
         """Run the graph with a single input and config.
 
         Args:
@@ -2733,7 +2789,7 @@ class Pregel(PregelProtocol):
         """
         output_keys = output_keys if output_keys is not None else self.output_channels
         if stream_mode == "values":
-            latest: Union[dict[str, Any], Any] = None
+            latest: dict[str, Any] | Any = None
         else:
             chunks = []
         for chunk in self.stream(
@@ -2758,17 +2814,17 @@ class Pregel(PregelProtocol):
 
     async def ainvoke(
         self,
-        input: Union[dict[str, Any], Any],
-        config: Optional[RunnableConfig] = None,
+        input: dict[str, Any] | Any,
+        config: RunnableConfig | None = None,
         *,
         stream_mode: StreamMode = "values",
-        output_keys: Optional[Union[str, Sequence[str]]] = None,
-        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
-        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
-        checkpoint_during: Optional[bool] = None,
-        debug: Optional[bool] = None,
+        output_keys: str | Sequence[str] | None = None,
+        interrupt_before: All | Sequence[str] | None = None,
+        interrupt_after: All | Sequence[str] | None = None,
+        checkpoint_during: bool | None = None,
+        debug: bool | None = None,
         **kwargs: Any,
-    ) -> Union[dict[str, Any], Any]:
+    ) -> dict[str, Any] | Any:
         """Asynchronously invoke the graph on a single input.
 
         Args:
@@ -2788,7 +2844,7 @@ class Pregel(PregelProtocol):
 
         output_keys = output_keys if output_keys is not None else self.output_channels
         if stream_mode == "values":
-            latest: Union[dict[str, Any], Any] = None
+            latest: dict[str, Any] | Any = None
         else:
             chunks = []
         async for chunk in self.astream(

libs/langgraph/langgraph/pregel/algo.py

@@ -3,19 +3,17 @@ import itertools
 import sys
 import threading
 from collections import defaultdict, deque
+from collections.abc import Iterable, Mapping, Sequence
 from copy import copy
 from functools import partial
 from hashlib import sha1
 from typing import (
     Any,
     Callable,
-    Iterable,
     Literal,
-    Mapping,
     NamedTuple,
     Optional,
     Protocol,
-    Sequence,
     Union,
     cast,
     overload,

libs/langgraph/langgraph/pregel/call.py

@@ -5,7 +5,8 @@ import functools
 import inspect
 import sys
 import types
-from typing import Any, Callable, Generator, Generic, Optional, Sequence, TypeVar, cast
+from collections.abc import Generator, Sequence
+from typing import Any, Callable, Generic, Optional, TypeVar, cast
 
 from langchain_core.runnables import Runnable
 from typing_extensions import ParamSpec
@@ -29,16 +30,12 @@ from langgraph.utils.runnable import (
 def _getattribute(obj: Any, name: str) -> Any:
     for subpath in name.split("."):
         if subpath == "<locals>":
-            raise AttributeError(
-                "Can't get local attribute {!r} on {!r}".format(name, obj)
-            )
+            raise AttributeError(f"Can't get local attribute {name!r} on {obj!r}")
         try:
             parent = obj
             obj = getattr(obj, subpath)
         except AttributeError:
-            raise AttributeError(
-                "Can't get attribute {!r} on {!r}".format(name, obj)
-            ) from None
+            raise AttributeError(f"Can't get attribute {name!r} on {obj!r}") from None
     return obj, parent
 
 

libs/langgraph/langgraph/pregel/checkpoint.py

@@ -1,5 +1,6 @@
+from collections.abc import Mapping
 from datetime import datetime, timezone
-from typing import Mapping, Optional
+from typing import Optional
 
 from langgraph.channels.base import BaseChannel
 from langgraph.checkpoint.base import Checkpoint

libs/langgraph/langgraph/pregel/debug.py

@@ -1,15 +1,12 @@
 from collections import defaultdict
+from collections.abc import Iterable, Iterator, Mapping, Sequence
 from dataclasses import asdict
 from datetime import datetime, timezone
 from pprint import pformat
 from typing import (
     Any,
-    Iterable,
-    Iterator,
     Literal,
-    Mapping,
     Optional,
-    Sequence,
     Union,
 )
 from uuid import UUID

libs/langgraph/langgraph/pregel/draw.py

@@ -0,0 +1,212 @@
+from collections import defaultdict
+from collections.abc import Mapping, Sequence
+from typing import Any, Optional, Union, cast
+
+from langchain_core.runnables.config import RunnableConfig
+from langchain_core.runnables.graph import Graph, Node
+
+from langgraph.channels.base import BaseChannel
+from langgraph.checkpoint.base import BaseCheckpointSaver
+from langgraph.constants import CONF, CONFIG_KEY_SEND, END, INPUT, START
+from langgraph.managed.base import ManagedValueSpec
+from langgraph.pregel.algo import (
+    PregelTaskWrites,
+    apply_writes,
+    increment,
+    prepare_next_tasks,
+)
+from langgraph.pregel.checkpoint import empty_checkpoint
+from langgraph.pregel.io import map_input
+from langgraph.pregel.manager import ChannelsManager
+from langgraph.pregel.read import PregelNode
+from langgraph.pregel.write import ChannelWrite
+from langgraph.types import All, Checkpointer, LoopProtocol
+
+
+def draw_graph(
+    config: RunnableConfig,
+    *,
+    nodes: dict[str, PregelNode],
+    specs: dict[str, Union[BaseChannel, ManagedValueSpec]],
+    input_channels: Union[str, Sequence[str]],
+    interrupt_after_nodes: Union[All, Sequence[str]],
+    interrupt_before_nodes: Union[All, Sequence[str]],
+    trigger_to_nodes: Optional[Mapping[str, Sequence[str]]],
+    checkpointer: Checkpointer,
+    subgraphs: dict[str, Graph],
+) -> Graph:
+    """Get the graph for this Pregel instance.
+
+    Args:
+        config: The configuration to use for the graph.
+        subgraphs: The subgraphs to include in the graph.
+        checkpointer: The checkpointer to use for the graph.
+
+    Returns:
+        The graph for this Pregel instance.
+    """
+    # (src, dest, is_conditional, label)
+    edges: set[tuple[str, str, bool, Optional[str]]] = set()
+
+    step = -1
+    checkpoint = empty_checkpoint()
+    get_next_version = (
+        checkpointer.get_next_version
+        if isinstance(checkpointer, BaseCheckpointSaver)
+        else increment
+    )
+    with ChannelsManager(
+        specs,
+        checkpoint,
+        LoopProtocol(step=step, stop=-1, config=config),
+        skip_context=True,
+    ) as (channels, managed):
+        static_seen: set[Any] = set()
+        sources: dict[str, set[tuple[str, bool, Optional[str]]]] = {}
+        step_sources: dict[str, set[tuple[str, bool, Optional[str]]]] = {}
+        # remove node mappers
+        nodes = {
+            k: v.copy(update={"mapper": None}) if v.mapper is not None else v
+            for k, v in nodes.items()
+        }
+        # apply input writes
+        input_writes = list(map_input(input_channels, {}))
+        _, updated_channels = apply_writes(
+            checkpoint,
+            channels,
+            [
+                PregelTaskWrites((), INPUT, input_writes, []),
+            ],
+            get_next_version,
+        )
+        # prepare first tasks
+        tasks = prepare_next_tasks(
+            checkpoint,
+            [],
+            nodes,
+            channels,
+            managed,
+            config,
+            step,
+            for_execution=True,
+            store=None,
+            checkpointer=None,
+            manager=None,
+            trigger_to_nodes=trigger_to_nodes,
+            updated_channels=updated_channels,
+        )
+        start_tasks = tasks
+        # run the pregel loop
+        while tasks:
+            conditionals: dict[tuple[str, str, Any], Optional[str]] = {}
+            # run task writers
+            for task in tasks.values():
+                for w in task.writers:
+                    # apply regular writes
+                    if isinstance(w, ChannelWrite):
+                        w.invoke(None, task.config)
+                    # apply conditional writes declared for static analysis, only once
+                    if w not in static_seen:
+                        static_seen.add(w)
+                        # apply static writes
+                        if writes := ChannelWrite.get_static_writes(w):
+                            conditionals.update(
+                                {(task.name, *t[:2]): t[2] for t in writes}
+                            )
+                            task.config[CONF][CONFIG_KEY_SEND]([t[:2] for t in writes])
+            # collect sources
+            step_sources = {
+                task.name: {
+                    (
+                        w[0],
+                        (task.name, *w) in conditionals,
+                        conditionals.get((task.name, *w)),
+                    )
+                    for w in task.writes
+                }
+                for task in tasks.values()
+            }
+            sources.update(step_sources)
+            # invert triggers
+            trigger_to_sources: dict[str, set[tuple[str, bool, Optional[str]]]] = (
+                defaultdict(set)
+            )
+            for src, triggers in sources.items():
+                for trigger, cond, label in triggers:
+                    trigger_to_sources[trigger].add((src, cond, label))
+            # apply writes
+            _, updated_channels = apply_writes(
+                checkpoint, channels, tasks.values(), get_next_version
+            )
+            # prepare next tasks
+            tasks = prepare_next_tasks(
+                checkpoint,
+                [],
+                nodes,
+                channels,
+                managed,
+                config,
+                step,
+                for_execution=True,
+                store=None,
+                checkpointer=None,
+                manager=None,
+                trigger_to_nodes=trigger_to_nodes,
+                updated_channels=updated_channels,
+            )
+            # collect edges
+            for task in tasks.values():
+                for trigger in task.triggers:
+                    for src, cond, label in sorted(trigger_to_sources[trigger]):
+                        edges.add((src, task.name, cond, label))
+        # assemble the graph
+        graph = Graph()
+        # add nodes
+        for name, node in nodes.items():
+            metadata = dict(node.metadata or {})
+            if name in interrupt_before_nodes and name in interrupt_after_nodes:
+                metadata["__interrupt"] = "before,after"
+            elif name in interrupt_before_nodes:
+                metadata["__interrupt"] = "before"
+            elif name in interrupt_after_nodes:
+                metadata["__interrupt"] = "after"
+            graph.add_node(node.bound, name, metadata=metadata or None)
+        # add start node
+        if START not in nodes:
+            graph.add_node(None, START)
+            for task in start_tasks.values():
+                graph.add_edge(graph.nodes[START], graph.nodes[task.name])
+        # add discovered edges
+        for src, dest, is_conditional, label in sorted(edges):
+            graph.add_edge(
+                graph.nodes[src],
+                graph.nodes[dest],
+                data=label if label != dest else None,
+                conditional=is_conditional,
+            )
+        # add end edges
+        if step_sources:
+            end = graph.add_node(None, END)
+            termini = {d for _, d, _, _ in edges}.difference(s for s, _, _, _ in edges)
+            for src in sorted(termini.union(step_sources)):
+                graph.add_edge(graph.nodes[src], end, conditional=src not in termini)
+        # replace subgraphs
+        for name, subgraph in subgraphs.items():
+            subgraph.trim_first_node()
+            subgraph.trim_last_node()
+            if (
+                len(subgraph.nodes) > 1
+                and name in graph.nodes
+                and subgraph.first_node()
+                and subgraph.last_node()
+            ):
+                # replace the node with the subgraph
+                graph.nodes.pop(name)
+                first, last = graph.extend(subgraph, prefix=name)
+                for idx, edge in enumerate(graph.edges):
+                    if edge.source == name:
+                        graph.edges[idx] = edge.copy(source=cast(Node, last).id)
+                    elif edge.target == name:
+                        graph.edges[idx] = edge.copy(target=cast(Node, first).id)
+
+        return graph

libs/langgraph/langgraph/pregel/executor.py

@@ -1,15 +1,12 @@
 import asyncio
 import concurrent.futures
 import time
-from contextlib import ExitStack
+from collections.abc import Awaitable, Coroutine
+from contextlib import AbstractAsyncContextManager, AbstractContextManager, ExitStack
 from contextvars import copy_context
 from types import TracebackType
 from typing import (
-    AsyncContextManager,
-    Awaitable,
     Callable,
-    ContextManager,
-    Coroutine,
     Optional,
     Protocol,
     TypeVar,
@@ -40,7 +37,7 @@ class Submit(Protocol[P, T]):
     ) -> concurrent.futures.Future[T]: ...
 
 
-class BackgroundExecutor(ContextManager):
+class BackgroundExecutor(AbstractContextManager):
     """A context manager that runs sync tasks in the background.
     Uses a thread pool executor to delegate tasks to separate threads.
     On exit,
@@ -122,7 +119,7 @@ class BackgroundExecutor(ContextManager):
                     pass
 
 
-class AsyncBackgroundExecutor(AsyncContextManager):
+class AsyncBackgroundExecutor(AbstractAsyncContextManager):
     """A context manager that runs async tasks in the background.
     Uses the current event loop to delegate tasks to asyncio tasks.
     On exit,

libs/langgraph/langgraph/pregel/io.py

@@ -1,5 +1,6 @@
 from collections import Counter
-from typing import Any, Iterator, Literal, Mapping, Optional, Sequence, TypeVar, Union
+from collections.abc import Iterator, Mapping, Sequence
+from typing import Any, Literal, Optional, TypeVar, Union
 from uuid import UUID
 
 from langchain_core.runnables.utils import AddableDict

libs/langgraph/langgraph/pregel/loop.py

@@ -3,21 +3,20 @@ import binascii
 import concurrent.futures
 import dataclasses
 from collections import defaultdict, deque
-from contextlib import AsyncExitStack, ExitStack
+from collections.abc import Iterator, Mapping, Sequence
+from contextlib import (
+    AbstractAsyncContextManager,
+    AbstractContextManager,
+    AsyncExitStack,
+    ExitStack,
+)
 from inspect import signature
 from types import TracebackType
 from typing import (
     Any,
-    AsyncContextManager,
     Callable,
-    ContextManager,
-    Iterator,
-    List,
     Literal,
-    Mapping,
     Optional,
-    Sequence,
-    Type,
     TypeVar,
     Union,
     cast,
@@ -146,7 +145,7 @@ def DuplexStream(*streams: StreamProtocol) -> StreamProtocol:
 
 class PregelLoop(LoopProtocol):
     input: Optional[Any]
-    input_model: Optional[Type[BaseModel]]
+    input_model: Optional[type[BaseModel]]
     checkpointer: Optional[BaseCheckpointSaver]
     nodes: Mapping[str, PregelNode]
     specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]]
@@ -186,7 +185,7 @@ class PregelLoop(LoopProtocol):
     checkpoint_ns: tuple[str, ...]
     checkpoint_config: RunnableConfig
     checkpoint_metadata: CheckpointMetadata
-    checkpoint_pending_writes: List[PendingWrite]
+    checkpoint_pending_writes: list[PendingWrite]
     checkpoint_previous_versions: dict[str, Union[str, float, int]]
     prev_checkpoint_config: Optional[RunnableConfig]
 
@@ -214,7 +213,7 @@ class PregelLoop(LoopProtocol):
         interrupt_after: Union[All, Sequence[str]] = EMPTY_SEQ,
         interrupt_before: Union[All, Sequence[str]] = EMPTY_SEQ,
         manager: Union[None, AsyncParentRunManager, ParentRunManager] = None,
-        input_model: Optional[Type[BaseModel]] = None,
+        input_model: Optional[type[BaseModel]] = None,
         debug: bool = False,
         migrate_checkpoint: Optional[Callable[[Checkpoint], None]] = None,
         trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None,
@@ -491,7 +490,7 @@ class PregelLoop(LoopProtocol):
             if self.input is INPUT_SHOULD_VALIDATE:
                 self.input = INPUT_DONE
                 # validate
-                cast(Type[BaseModel], self.input_model)(
+                cast(type[BaseModel], self.input_model)(
                     **read_channels(self.channels, self.stream_keys)
                 )
             # produce values output
@@ -839,7 +838,7 @@ class PregelLoop(LoopProtocol):
 
     def _suppress_interrupt(
         self,
-        exc_type: Optional[Type[BaseException]],
+        exc_type: Optional[type[BaseException]],
         exc_value: Optional[BaseException],
         traceback: Optional[TracebackType],
     ) -> Optional[bool]:
@@ -945,7 +944,7 @@ class PregelLoop(LoopProtocol):
                 )
 
 
-class SyncPregelLoop(PregelLoop, ContextManager):
+class SyncPregelLoop(PregelLoop, AbstractContextManager):
     def __init__(
         self,
         input: Optional[Any],
@@ -961,7 +960,7 @@ class SyncPregelLoop(PregelLoop, ContextManager):
         interrupt_before: Union[All, Sequence[str]] = EMPTY_SEQ,
         output_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
         stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
-        input_model: Optional[Type[BaseModel]] = None,
+        input_model: Optional[type[BaseModel]] = None,
         debug: bool = False,
         migrate_checkpoint: Optional[Callable[[Checkpoint], None]] = None,
         trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None,
@@ -1087,7 +1086,7 @@ class SyncPregelLoop(PregelLoop, ContextManager):
 
     def __exit__(
         self,
-        exc_type: Optional[Type[BaseException]],
+        exc_type: Optional[type[BaseException]],
         exc_value: Optional[BaseException],
         traceback: Optional[TracebackType],
     ) -> Optional[bool]:
@@ -1095,7 +1094,7 @@ class SyncPregelLoop(PregelLoop, ContextManager):
         return self.stack.__exit__(exc_type, exc_value, traceback)
 
 
-class AsyncPregelLoop(PregelLoop, AsyncContextManager):
+class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
     def __init__(
         self,
         input: Optional[Any],
@@ -1111,7 +1110,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
         manager: Union[None, AsyncParentRunManager, ParentRunManager] = None,
         output_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
         stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
-        input_model: Optional[Type[BaseModel]] = None,
+        input_model: Optional[type[BaseModel]] = None,
         debug: bool = False,
         migrate_checkpoint: Optional[Callable[[Checkpoint], None]] = None,
         trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None,
@@ -1240,7 +1239,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
 
     async def __aexit__(
         self,
-        exc_type: Optional[Type[BaseException]],
+        exc_type: Optional[type[BaseException]],
         exc_value: Optional[BaseException],
         traceback: Optional[TracebackType],
     ) -> Optional[bool]:

libs/langgraph/langgraph/pregel/manager.py

@@ -1,6 +1,7 @@
 import asyncio
+from collections.abc import AsyncIterator, Iterator, Mapping
 from contextlib import AsyncExitStack, ExitStack, asynccontextmanager, contextmanager
-from typing import AsyncIterator, Iterator, Mapping, Union
+from typing import Union
 
 from langgraph.channels.base import BaseChannel
 from langgraph.checkpoint.base import Checkpoint

libs/langgraph/langgraph/pregel/messages.py

@@ -1,12 +1,8 @@
+from collections.abc import AsyncIterator, Iterator, Sequence
 from typing import (
     Any,
-    AsyncIterator,
     Callable,
-    Dict,
-    Iterator,
-    List,
     Optional,
-    Sequence,
     TypeVar,
     Union,
     cast,
@@ -115,13 +111,13 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
 
     def on_chain_start(
         self,
-        serialized: Dict[str, Any],
-        inputs: Dict[str, Any],
+        serialized: dict[str, Any],
+        inputs: dict[str, Any],
         *,
         run_id: UUID,
         parent_run_id: Optional[UUID] = None,
-        tags: Optional[List[str]] = None,
-        metadata: Optional[Dict[str, Any]] = None,
+        tags: Optional[list[str]] = None,
+        metadata: Optional[dict[str, Any]] = None,
         **kwargs: Any,
     ) -> Any:
         if (

libs/langgraph/langgraph/pregel/protocol.py

@@ -1,10 +1,8 @@
 from abc import ABC, abstractmethod
+from collections.abc import AsyncIterator, Iterator, Sequence
 from typing import (
     Any,
-    AsyncIterator,
-    Iterator,
     Optional,
-    Sequence,
     Union,
 )
 

libs/langgraph/langgraph/pregel/read.py

@@ -1,14 +1,10 @@
 from __future__ import annotations
 
+from collections.abc import AsyncIterator, Iterator, Mapping, Sequence
 from functools import cached_property
 from typing import (
     Any,
-    AsyncIterator,
     Callable,
-    Iterator,
-    Mapping,
-    Optional,
-    Sequence,
     Union,
 )
 
@@ -37,11 +33,11 @@ class ChannelRead(RunnableCallable):
     """Implements the logic for reading state from CONFIG_KEY_READ.
     Usable both as a runnable as well as a static method to call imperatively."""
 
-    channel: Union[str, list[str]]
+    channel: str | list[str]
 
     fresh: bool = False
 
-    mapper: Optional[Callable[[Any], Any]] = None
+    mapper: Callable[[Any], Any] | None = None
 
     @property
     def config_specs(self) -> list[ConfigurableFieldSpec]:
@@ -57,11 +53,11 @@ class ChannelRead(RunnableCallable):
 
     def __init__(
         self,
-        channel: Union[str, list[str]],
+        channel: str | list[str],
         *,
         fresh: bool = False,
-        mapper: Optional[Callable[[Any], Any]] = None,
-        tags: Optional[list[str]] = None,
+        mapper: Callable[[Any], Any] | None = None,
+        tags: list[str] | None = None,
     ) -> None:
         super().__init__(
             func=self._read,
@@ -75,9 +71,7 @@ class ChannelRead(RunnableCallable):
         self.mapper = mapper
         self.channel = channel
 
-    def get_name(
-        self, suffix: Optional[str] = None, *, name: Optional[str] = None
-    ) -> str:
+    def get_name(self, suffix: str | None = None, *, name: str | None = None) -> str:
         if name:
             pass
         elif isinstance(self.channel, str):
@@ -100,9 +94,9 @@ class ChannelRead(RunnableCallable):
     def do_read(
         config: RunnableConfig,
         *,
-        select: Union[str, list[str]],
+        select: str | list[str],
         fresh: bool = False,
-        mapper: Optional[Callable[[Any], Any]] = None,
+        mapper: Callable[[Any], Any] | None = None,
     ) -> Any:
         try:
             read: READ_TYPE = config[CONF][CONFIG_KEY_READ]
@@ -125,7 +119,7 @@ class PregelNode(Runnable):
     itself, but instead acts as a container for the components necessary to make
     a PregelExecutableTask for a node."""
 
-    channels: Union[list[str], Mapping[str, str]]
+    channels: list[str] | Mapping[str, str]
     """The channels that will be passed as input to `bound`.
     If a list, the node will be invoked with the first of that isn't empty.
     If a dict, the keys are the names of the channels, and the values are the keys
@@ -135,7 +129,7 @@ class PregelNode(Runnable):
     """If any of these channels is written to, this node will be triggered in
     the next step."""
 
-    mapper: Optional[Callable[[Any], Any]]
+    mapper: Callable[[Any], Any] | None
     """A function to transform the input before passing it to `bound`."""
 
     writers: list[Runnable]
@@ -146,13 +140,13 @@ class PregelNode(Runnable):
     """The main logic of the node. This will be invoked with the input from 
     `channels`."""
 
-    retry_policy: Optional[Sequence[RetryPolicy]]
+    retry_policy: Sequence[RetryPolicy] | None
     """The retry policies to use when invoking the node."""
 
-    tags: Optional[Sequence[str]]
+    tags: Sequence[str] | None
     """Tags to attach to the node for tracing."""
 
-    metadata: Optional[Mapping[str, Any]]
+    metadata: Mapping[str, Any] | None
     """Metadata to attach to the node for tracing."""
 
     subgraphs: Sequence[PregelProtocol]
@@ -161,15 +155,15 @@ class PregelNode(Runnable):
     def __init__(
         self,
         *,
-        channels: Union[list[str], Mapping[str, str]],
+        channels: list[str] | Mapping[str, str],
         triggers: Sequence[str],
-        mapper: Optional[Callable[[Any], Any]] = None,
-        writers: Optional[list[Runnable]] = None,
-        tags: Optional[list[str]] = None,
-        metadata: Optional[Mapping[str, Any]] = None,
-        bound: Optional[Runnable[Any, Any]] = None,
-        retry_policy: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None,
-        subgraphs: Optional[Sequence[PregelProtocol]] = None,
+        mapper: Callable[[Any], Any] | None = None,
+        writers: list[Runnable] | None = None,
+        tags: list[str] | None = None,
+        metadata: Mapping[str, Any] | None = None,
+        bound: Runnable[Any, Any] | None = None,
+        retry_policy: RetryPolicy | Sequence[RetryPolicy] | None = None,
+        subgraphs: Sequence[PregelProtocol] | None = None,
     ) -> None:
         self.channels = channels
         self.triggers = list(triggers)
@@ -219,7 +213,7 @@ class PregelNode(Runnable):
         return writers
 
     @cached_property
-    def node(self) -> Optional[Runnable[Any, Any]]:
+    def node(self) -> Runnable[Any, Any] | None:
         """Get a runnable that combines `bound` and `writers`."""
         writers = self.flat_writers
         if self.bound is DEFAULT_BOUND and not writers:
@@ -262,11 +256,9 @@ class PregelNode(Runnable):
 
     def __or__(
         self,
-        other: Union[
-            Runnable[Any, Other],
-            Callable[[Any], Other],
-            Mapping[str, Runnable[Any, Other] | Callable[[Any], Other]],
-        ],
+        other: Runnable[Any, Other]
+        | Callable[[Any], Other]
+        | Mapping[str, Runnable[Any, Other] | Callable[[Any], Other]],
     ) -> PregelNode:
         if isinstance(other, Runnable) and ChannelWrite.is_writer(other):
             return self.copy(update=dict(writers=[*self.writers, other]))
@@ -278,7 +270,7 @@ class PregelNode(Runnable):
     def pipe(
         self,
         *others: Runnable[Any, Other] | Callable[[Any], Other],
-        name: Optional[str] = None,
+        name: str | None = None,
     ) -> RunnableSerializable[Any, Other]:
         for other in others:
             self = self | other
@@ -286,19 +278,17 @@ class PregelNode(Runnable):
 
     def __ror__(
         self,
-        other: Union[
-            Runnable[Other, Any],
-            Callable[[Any], Other],
-            Mapping[str, Union[Runnable[Other, Any], Callable[[Other], Any]]],
-        ],
+        other: Runnable[Other, Any]
+        | Callable[[Any], Other]
+        | Mapping[str, Runnable[Other, Any] | Callable[[Other], Any]],
     ) -> RunnableSerializable:
         raise NotImplementedError()
 
     def invoke(
         self,
         input: Input,
-        config: Optional[RunnableConfig] = None,
-        **kwargs: Optional[Any],
+        config: RunnableConfig | None = None,
+        **kwargs: Any | None,
     ) -> Any:
         return self.bound.invoke(
             input,
@@ -309,8 +299,8 @@ class PregelNode(Runnable):
     async def ainvoke(
         self,
         input: Input,
-        config: Optional[RunnableConfig] = None,
-        **kwargs: Optional[Any],
+        config: RunnableConfig | None = None,
+        **kwargs: Any | None,
     ) -> Any:
         return await self.bound.ainvoke(
             input,
@@ -321,8 +311,8 @@ class PregelNode(Runnable):
     def stream(
         self,
         input: Input,
-        config: Optional[RunnableConfig] = None,
-        **kwargs: Optional[Any],
+        config: RunnableConfig | None = None,
+        **kwargs: Any | None,
     ) -> Iterator[Any]:
         yield from self.bound.stream(
             input,
@@ -333,8 +323,8 @@ class PregelNode(Runnable):
     async def astream(
         self,
         input: Input,
-        config: Optional[RunnableConfig] = None,
-        **kwargs: Optional[Any],
+        config: RunnableConfig | None = None,
+        **kwargs: Any | None,
     ) -> AsyncIterator[Any]:
         async for item in self.bound.astream(
             input,

libs/langgraph/langgraph/pregel/remote.py

@@ -1,11 +1,9 @@
+from collections.abc import AsyncIterator, Iterator, Sequence
 from dataclasses import asdict
 from typing import (
     Any,
-    AsyncIterator,
-    Iterator,
     Literal,
     Optional,
-    Sequence,
     Union,
     cast,
 )

libs/langgraph/langgraph/pregel/retry.py

@@ -3,8 +3,9 @@ import logging
 import random
 import sys
 import time
+from collections.abc import Sequence
 from dataclasses import replace
-from typing import Any, Optional, Sequence
+from typing import Any, Optional
 
 from langgraph.constants import (
     CONF,

libs/langgraph/langgraph/pregel/runner.py

@@ -3,18 +3,13 @@ import concurrent.futures
 import threading
 import time
 import weakref
+from collections.abc import AsyncIterator, Awaitable, Iterable, Iterator, Sequence
 from functools import partial
 from typing import (
     Any,
-    AsyncIterator,
-    Awaitable,
     Callable,
     Generic,
-    Iterable,
-    Iterator,
     Optional,
-    Sequence,
-    Type,
     TypeVar,
     Union,
     cast,
@@ -72,7 +67,7 @@ class FuturesDict(Generic[F, E], dict[F, Optional[PregelExecutableTask]]):
         callback: weakref.ref[
             Callable[[PregelExecutableTask, Optional[BaseException]], None]
         ],
-        future_type: Type[F],
+        future_type: type[F],
         # used for generic typing, newer py supports FutureDict[...](...)
     ) -> None:
         super().__init__()
@@ -475,7 +470,7 @@ def _exception(
 def _panic_or_proceed(
     futs: Union[set[concurrent.futures.Future], set[asyncio.Future]],
     *,
-    timeout_exc_cls: Type[Exception] = TimeoutError,
+    timeout_exc_cls: type[Exception] = TimeoutError,
     panic: bool = True,
 ) -> None:
     """Cancel remaining tasks if any failed, re-raise exception if panic is True."""

libs/langgraph/langgraph/pregel/validate.py

@@ -1,4 +1,5 @@
-from typing import Any, Mapping, Optional, Sequence, Union
+from collections.abc import Mapping, Sequence
+from typing import Any, Optional, Union
 
 from langgraph.channels.base import BaseChannel
 from langgraph.constants import RESERVED

libs/langgraph/langgraph/pregel/write.py

@@ -1,11 +1,11 @@
 from __future__ import annotations
 
+from collections.abc import Sequence
 from typing import (
     Any,
     Callable,
     NamedTuple,
     Optional,
-    Sequence,
     TypeVar,
     Union,
     cast,
@@ -14,7 +14,7 @@ from typing import (
 from langchain_core.runnables import Runnable, RunnableConfig
 from langchain_core.runnables.utils import ConfigurableFieldSpec
 
-from langgraph.constants import CONF, CONFIG_KEY_SEND, TASKS, Send
+from langgraph.constants import CONF, CONFIG_KEY_SEND, MISSING, TASKS, Send
 from langgraph.errors import InvalidUpdateError
 from langgraph.utils.runnable import RunnableCallable
 
@@ -32,30 +32,32 @@ class ChannelWriteEntry(NamedTuple):
     """Value to write, or PASSTHROUGH to use the input."""
     skip_none: bool = False
     """Whether to skip writing if the value is None."""
-    mapper: Optional[Callable] = None
+    mapper: Callable | None = None
     """Function to transform the value before writing."""
 
 
 class ChannelWriteTupleEntry(NamedTuple):
-    mapper: Callable[[Any], Optional[Sequence[tuple[str, Any]]]]
+    mapper: Callable[[Any], Sequence[tuple[str, Any]] | None]
     """Function to extract tuples from value."""
     value: Any = PASSTHROUGH
     """Value to write, or PASSTHROUGH to use the input."""
+    static: Optional[Sequence[tuple[str, Any, Optional[str]]]] = None
+    """Optional, declared writes for static analysis."""
 
 
 class ChannelWrite(RunnableCallable):
     """Implements the logic for sending writes to CONFIG_KEY_SEND.
     Can be used as a runnable or as a static method to call imperatively."""
 
-    writes: list[Union[ChannelWriteEntry, ChannelWriteTupleEntry, Send]]
+    writes: list[ChannelWriteEntry | ChannelWriteTupleEntry | Send]
     """Sequence of write entries or Send objects to write."""
 
     def __init__(
         self,
-        writes: Sequence[Union[ChannelWriteEntry, ChannelWriteTupleEntry, Send]],
+        writes: Sequence[ChannelWriteEntry | ChannelWriteTupleEntry | Send],
         *,
-        tags: Optional[Sequence[str]] = None,  # ignored
-        require_at_least_one_of: Optional[Sequence[str]] = None,  # ignored
+        tags: Sequence[str] | None = None,  # ignored
+        require_at_least_one_of: Sequence[str] | None = None,  # ignored
     ):
         super().__init__(
             func=self._write,
@@ -68,9 +70,7 @@ class ChannelWrite(RunnableCallable):
             list[Union[ChannelWriteEntry, ChannelWriteTupleEntry, Send]], writes
         )
 
-    def get_name(
-        self, suffix: Optional[str] = None, *, name: Optional[str] = None
-    ) -> str:
+    def get_name(self, suffix: str | None = None, *, name: str | None = None) -> str:
         if not name:
             name = f"ChannelWrite<{','.join(w.channel if isinstance(w, ChannelWriteEntry) else '...' if isinstance(w, ChannelWriteTupleEntry) else w.node for w in self.writes)}>"
         return super().get_name(suffix, name=name)
@@ -120,8 +120,9 @@ class ChannelWrite(RunnableCallable):
     @staticmethod
     def do_write(
         config: RunnableConfig,
-        writes: Sequence[Union[ChannelWriteEntry, ChannelWriteTupleEntry, Send]],
-        require_at_least_one_of: Optional[Sequence[str]] = None,  # ignored
+        writes: Sequence[ChannelWriteEntry | ChannelWriteTupleEntry | Send],
+        allow_passthrough: bool = True,
+        require_at_least_one_of: Sequence[str] | None = None,  # ignored
     ) -> None:
         # validate
         for w in writes:
@@ -130,46 +131,80 @@ class ChannelWrite(RunnableCallable):
                     raise InvalidUpdateError(
                         "Cannot write to the reserved channel TASKS"
                     )
-                if w.value is PASSTHROUGH:
+                if w.value is PASSTHROUGH and not allow_passthrough:
                     raise InvalidUpdateError("PASSTHROUGH value must be replaced")
             if isinstance(w, ChannelWriteTupleEntry):
-                if w.value is PASSTHROUGH:
+                if w.value is PASSTHROUGH and not allow_passthrough:
                     raise InvalidUpdateError("PASSTHROUGH value must be replaced")
-        # assemble writes
-        tuples: list[tuple[str, Any]] = []
-        for w in writes:
-            if isinstance(w, Send):
-                tuples.append((TASKS, w))
-            elif isinstance(w, ChannelWriteTupleEntry):
-                if ww := w.mapper(w.value):
-                    tuples.extend(ww)
-            elif isinstance(w, ChannelWriteEntry):
-                value = w.mapper(w.value) if w.mapper is not None else w.value
-                if value is SKIP_WRITE:
-                    continue
-                if w.skip_none and value is None:
-                    continue
-                tuples.append((w.channel, value))
-            else:
-                raise ValueError(f"Invalid write entry: {w}")
         # if we want to persist writes found before hitting a ParentCommand
         # can move this to a finally block
         write: TYPE_SEND = config[CONF][CONFIG_KEY_SEND]
-        write(tuples)
+        write(_assemble_writes(writes))
 
     @staticmethod
     def is_writer(runnable: Runnable) -> bool:
         """Used by PregelNode to distinguish between writers and other runnables."""
         return (
             isinstance(runnable, ChannelWrite)
-            or getattr(runnable, "_is_channel_writer", False) is True
+            or getattr(runnable, "_is_channel_writer", MISSING) is not MISSING
         )
 
     @staticmethod
-    def register_writer(runnable: R) -> R:
+    def get_static_writes(
+        runnable: Runnable,
+    ) -> Optional[Sequence[tuple[str, Any, Optional[str]]]]:
+        """Used to get conditional writes a writer declares for static analysis."""
+        if isinstance(runnable, ChannelWrite):
+            return [
+                w
+                for entry in runnable.writes
+                if isinstance(entry, ChannelWriteTupleEntry) and entry.static
+                for w in entry.static
+            ] or None
+        elif writes := getattr(runnable, "_is_channel_writer", MISSING):
+            if writes is not MISSING:
+                writes = cast(
+                    Sequence[tuple[Union[ChannelWriteEntry, Send], Optional[str]]],
+                    writes,
+                )
+                entries = [e for e, _ in writes]
+                labels = [la for _, la in writes]
+                return [(*t, la) for t, la in zip(_assemble_writes(entries), labels)]
+
+    @staticmethod
+    def register_writer(
+        runnable: R,
+        static: Optional[
+            Sequence[tuple[Union[ChannelWriteEntry, Send], Optional[str]]]
+        ] = None,
+    ) -> R:
         """Used to mark a runnable as a writer, so that it can be detected by is_writer.
-        Instances of ChannelWrite are automatically marked as writers."""
+        Instances of ChannelWrite are automatically marked as writers.
+        Optionally, a list of declared writes can be passed for static analysis."""
         # using object.__setattr__ to work around objects that override __setattr__
         # eg. pydantic models and dataclasses
-        object.__setattr__(runnable, "_is_channel_writer", True)
+        object.__setattr__(runnable, "_is_channel_writer", static)
         return runnable
+
+
+def _assemble_writes(
+    writes: Sequence[Union[ChannelWriteEntry, ChannelWriteTupleEntry, Send]],
+) -> list[tuple[str, Any]]:
+    """Assembles the writes into a list of tuples."""
+    tuples: list[tuple[str, Any]] = []
+    for w in writes:
+        if isinstance(w, Send):
+            tuples.append((TASKS, w))
+        elif isinstance(w, ChannelWriteTupleEntry):
+            if ww := w.mapper(w.value):
+                tuples.extend(ww)
+        elif isinstance(w, ChannelWriteEntry):
+            value = w.mapper(w.value) if w.mapper is not None else w.value
+            if value is SKIP_WRITE:
+                continue
+            if w.skip_none and value is None:
+                continue
+            tuples.append((w.channel, value))
+        else:
+            raise ValueError(f"Invalid write entry: {w}")
+    return tuples

libs/langgraph/langgraph/types.py

@@ -1,18 +1,16 @@
 import dataclasses
 import sys
 from collections import deque
+from collections.abc import Hashable, Sequence
 from typing import (
     TYPE_CHECKING,
     Any,
     Callable,
     ClassVar,
     Generic,
-    Hashable,
     Literal,
     NamedTuple,
     Optional,
-    Sequence,
-    Type,
     TypeVar,
     Union,
     cast,
@@ -118,7 +116,7 @@ class RetryPolicy(NamedTuple):
     jitter: bool = True
     """Whether to add random jitter to the interval between retries."""
     retry_on: Union[
-        Type[Exception], Sequence[Type[Exception]], Callable[[Exception], bool]
+        type[Exception], Sequence[type[Exception]], Callable[[Exception], bool]
     ] = default_retry_on
     """List of exception classes that should trigger a retry, or a callable that returns True for exceptions that should trigger a retry."""
 

libs/langgraph/langgraph/utils/config.py

@@ -1,6 +1,7 @@
 from collections import ChainMap
+from collections.abc import Sequence
 from os import getenv
-from typing import Any, Optional, Sequence, cast
+from typing import Any, Optional, cast
 
 from langchain_core.callbacks import (
     AsyncCallbackManager,

libs/langgraph/langgraph/utils/fields.py

@@ -1,9 +1,10 @@
 import dataclasses
-from typing import Any, Generator, Optional, Sequence, Type, Union, get_type_hints
+from collections.abc import Generator, Sequence
+from typing import Annotated, Any, Optional, Union, get_type_hints
 
 from pydantic import BaseModel
 from pydantic.v1 import BaseModel as BaseModelV1
-from typing_extensions import Annotated, NotRequired, ReadOnly, Required, get_origin
+from typing_extensions import NotRequired, ReadOnly, Required, get_origin
 
 # NOTE: this is redefined here separately from langgraph.constants
 # to avoid a circular import
@@ -68,7 +69,7 @@ def _is_readonly_type(type_: Any) -> bool:
 _DEFAULT_KEYS: frozenset[str] = frozenset()
 
 
-def get_field_default(name: str, type_: Any, schema: Type[Any]) -> Any:
+def get_field_default(name: str, type_: Any, schema: type[Any]) -> Any:
     """Determine the default value for a field in a state schema.
 
     This is based on:
@@ -115,7 +116,7 @@ def get_field_default(name: str, type_: Any, schema: Type[Any]) -> Any:
 
 
 def get_enhanced_type_hints(
-    type: Type[Any],
+    type: type[Any],
 ) -> Generator[tuple[str, Any, Any, Optional[str]], None, None]:
     """Attempt to extract default values and descriptions from provided type, used for config schema."""
     for name, typ in get_type_hints(type).items():

libs/langgraph/langgraph/utils/future.py

@@ -4,7 +4,8 @@ import contextvars
 import inspect
 import sys
 import types
-from typing import Awaitable, Coroutine, Generator, Optional, TypeVar, Union, cast
+from collections.abc import Awaitable, Coroutine, Generator
+from typing import Optional, TypeVar, Union, cast
 
 T = TypeVar("T")
 AnyFuture = Union[asyncio.Future, concurrent.futures.Future]

libs/langgraph/langgraph/utils/pydantic.py

@@ -1,7 +1,7 @@
 import sys
 import typing
 from dataclasses import is_dataclass
-from typing import Any, Dict, Optional, Union
+from typing import Any, Optional, Union
 
 import typing_extensions
 from pydantic import BaseModel
@@ -11,7 +11,7 @@ from pydantic.v1 import BaseModel as BaseModelV1
 def create_model(
     model_name: str,
     *,
-    field_definitions: Optional[Dict[str, Any]] = None,
+    field_definitions: Optional[dict[str, Any]] = None,
     root: Optional[Any] = None,
 ) -> Union[BaseModel, BaseModelV1]:
     """Create a pydantic model with the given field definitions.

libs/langgraph/langgraph/utils/runnable.py

@@ -2,21 +2,22 @@ import asyncio
 import enum
 import inspect
 import sys
+from collections.abc import (
+    AsyncIterator,
+    Awaitable,
+    Coroutine,
+    Generator,
+    Iterator,
+    Sequence,
+)
 from contextlib import AsyncExitStack, contextmanager
 from contextvars import Context, Token, copy_context
 from functools import partial, wraps
 from typing import (
     Any,
-    AsyncIterator,
-    Awaitable,
     Callable,
-    Coroutine,
-    Generator,
-    Iterator,
     Optional,
     Protocol,
-    Sequence,
-    Tuple,
     Union,
     cast,
 )
@@ -278,7 +279,7 @@ class RunnableCallable(Runnable):
 
         if func_accepts_config is not None:
             self.func_accepts_config = func_accepts_config
-            self.func_accepts: dict[str, Tuple[str, Any]] = {}
+            self.func_accepts: dict[str, tuple[str, Any]] = {}
         else:
             params = inspect.signature(cast(Callable, func or afunc)).parameters
 

libs/langgraph/poetry.lock

@@ -1,4 +1,4 @@
-# This file is automatically @generated by Poetry 2.0.1 and should not be changed by hand.
+# This file is automatically @generated by Poetry 2.0.0 and should not be changed by hand.
 
 [[package]]
 name = "aiosqlite"
@@ -1324,14 +1324,14 @@ files = [
 
 [[package]]
 name = "langchain-core"
-version = "0.3.46"
+version = "0.3.55"
 description = "Building applications with LLMs through composability"
 optional = false
 python-versions = "<4.0,>=3.9"
 groups = ["main", "dev"]
 files = [
-    {file = "langchain_core-0.3.46-py3-none-any.whl", hash = "sha256:28b5689fc347975ea520b5364ab4aee5567e661553bbee5e97cabf4596c28ce0"},
-    {file = "langchain_core-0.3.46.tar.gz", hash = "sha256:5fca010eeb0a427be5aa8a8525e2112995dde790c584cef165be7c5e0ee1c2b5"},
+    {file = "langchain_core-0.3.55-py3-none-any.whl", hash = "sha256:b3cb36bf37755a616158a79866657c6697b43a2f7c69dd723ce425f1c76c1baa"},
+    {file = "langchain_core-0.3.55.tar.gz", hash = "sha256:0f2b3e311621116a83510c70b0ac9d959030a0a457a69483535cff18501fedc9"},
 ]
 
 [package.dependencies]

libs/langgraph/pyproject.toml

@@ -41,11 +41,12 @@ types-requests = "^2.32.0.20240914"
 pycryptodome = "^3.21.0"
 
 [tool.ruff]
-lint.select = [ "E", "F", "I", "TID251" ]
-lint.ignore = [ "E501" ]
+lint.select = [ "E", "F", "I", "TID251", "UP" ]
+lint.ignore = [ "E501", "UP007" ]
 line-length = 88
 indent-width = 4
 extend-include = ["*.ipynb"]
+target-version = "py39"
 
 [tool.ruff.format]
 quote-style = "double"

libs/langgraph/tests/__snapshots__/test_large_cases_async.ambr

@@ -1,151 +0,0 @@
-# serializer version: 1
-# name: test_weather_subgraph[memory]
-  '''
-  %%{init: {'flowchart': {'curve': 'linear'}}}%%
-  graph TD;
-  	__start__([<p>__start__</p>]):::first
-  	router_node(router_node)
-  	normal_llm_node(normal_llm_node)
-  	__end__([<p>__end__</p>]):::last
-  	__start__ --> router_node;
-  	normal_llm_node --> __end__;
-  	weather_graph_weather_node --> __end__;
-  	router_node -.-> normal_llm_node;
-  	router_node -.-> weather_graph_model_node;
-  	router_node -.-> __end__;
-  	subgraph weather_graph
-  	weather_graph_model_node(model_node)
-  	weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
-  	weather_graph_model_node --> weather_graph_weather_node;
-  	end
-  	classDef default fill:#f2f0ff,line-height:1.2
-  	classDef first fill-opacity:0
-  	classDef last fill:#bfb6fc
-  
-  '''
-# ---
-# name: test_weather_subgraph[postgres_aio]
-  '''
-  %%{init: {'flowchart': {'curve': 'linear'}}}%%
-  graph TD;
-  	__start__([<p>__start__</p>]):::first
-  	router_node(router_node)
-  	normal_llm_node(normal_llm_node)
-  	__end__([<p>__end__</p>]):::last
-  	__start__ --> router_node;
-  	normal_llm_node --> __end__;
-  	weather_graph_weather_node --> __end__;
-  	router_node -.-> normal_llm_node;
-  	router_node -.-> weather_graph_model_node;
-  	router_node -.-> __end__;
-  	subgraph weather_graph
-  	weather_graph_model_node(model_node)
-  	weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
-  	weather_graph_model_node --> weather_graph_weather_node;
-  	end
-  	classDef default fill:#f2f0ff,line-height:1.2
-  	classDef first fill-opacity:0
-  	classDef last fill:#bfb6fc
-  
-  '''
-# ---
-# name: test_weather_subgraph[postgres_aio_pipe]
-  '''
-  %%{init: {'flowchart': {'curve': 'linear'}}}%%
-  graph TD;
-  	__start__([<p>__start__</p>]):::first
-  	router_node(router_node)
-  	normal_llm_node(normal_llm_node)
-  	__end__([<p>__end__</p>]):::last
-  	__start__ --> router_node;
-  	normal_llm_node --> __end__;
-  	weather_graph_weather_node --> __end__;
-  	router_node -.-> normal_llm_node;
-  	router_node -.-> weather_graph_model_node;
-  	router_node -.-> __end__;
-  	subgraph weather_graph
-  	weather_graph_model_node(model_node)
-  	weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
-  	weather_graph_model_node --> weather_graph_weather_node;
-  	end
-  	classDef default fill:#f2f0ff,line-height:1.2
-  	classDef first fill-opacity:0
-  	classDef last fill:#bfb6fc
-  
-  '''
-# ---
-# name: test_weather_subgraph[postgres_aio_pool]
-  '''
-  %%{init: {'flowchart': {'curve': 'linear'}}}%%
-  graph TD;
-  	__start__([<p>__start__</p>]):::first
-  	router_node(router_node)
-  	normal_llm_node(normal_llm_node)
-  	__end__([<p>__end__</p>]):::last
-  	__start__ --> router_node;
-  	normal_llm_node --> __end__;
-  	weather_graph_weather_node --> __end__;
-  	router_node -.-> normal_llm_node;
-  	router_node -.-> weather_graph_model_node;
-  	router_node -.-> __end__;
-  	subgraph weather_graph
-  	weather_graph_model_node(model_node)
-  	weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
-  	weather_graph_model_node --> weather_graph_weather_node;
-  	end
-  	classDef default fill:#f2f0ff,line-height:1.2
-  	classDef first fill-opacity:0
-  	classDef last fill:#bfb6fc
-  
-  '''
-# ---
-# name: test_weather_subgraph[postgres_aio_shallow]
-  '''
-  %%{init: {'flowchart': {'curve': 'linear'}}}%%
-  graph TD;
-  	__start__([<p>__start__</p>]):::first
-  	router_node(router_node)
-  	normal_llm_node(normal_llm_node)
-  	__end__([<p>__end__</p>]):::last
-  	__start__ --> router_node;
-  	normal_llm_node --> __end__;
-  	weather_graph_weather_node --> __end__;
-  	router_node -.-> normal_llm_node;
-  	router_node -.-> weather_graph_model_node;
-  	router_node -.-> __end__;
-  	subgraph weather_graph
-  	weather_graph_model_node(model_node)
-  	weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
-  	weather_graph_model_node --> weather_graph_weather_node;
-  	end
-  	classDef default fill:#f2f0ff,line-height:1.2
-  	classDef first fill-opacity:0
-  	classDef last fill:#bfb6fc
-  
-  '''
-# ---
-# name: test_weather_subgraph[sqlite_aio]
-  '''
-  %%{init: {'flowchart': {'curve': 'linear'}}}%%
-  graph TD;
-  	__start__([<p>__start__</p>]):::first
-  	router_node(router_node)
-  	normal_llm_node(normal_llm_node)
-  	__end__([<p>__end__</p>]):::last
-  	__start__ --> router_node;
-  	normal_llm_node --> __end__;
-  	weather_graph_weather_node --> __end__;
-  	router_node -.-> normal_llm_node;
-  	router_node -.-> weather_graph_model_node;
-  	router_node -.-> __end__;
-  	subgraph weather_graph
-  	weather_graph_model_node(model_node)
-  	weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
-  	weather_graph_model_node --> weather_graph_weather_node;
-  	end
-  	classDef default fill:#f2f0ff,line-height:1.2
-  	classDef first fill-opacity:0
-  	classDef last fill:#bfb6fc
-  
-  '''
-# ---

libs/langgraph/tests/__snapshots__/test_pregel_async.ambr

@@ -4,11 +4,11 @@
   graph TD;
   	__start__ --> rewrite_query;
   	analyzer_one --> retriever_one;
-  	qa --> __end__;
   	retriever_one --> qa;
   	retriever_two --> qa;
   	rewrite_query --> analyzer_one;
   	rewrite_query -.-> retriever_two;
+  	qa --> __end__;
   
   '''
 # ---
@@ -120,695 +120,21 @@
     'type': 'object',
   })
 # ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio]
-  '''
-  graph TD;
-  	__start__ --> rewrite_query;
-  	analyzer_one --> retriever_one;
-  	qa --> __end__;
-  	retriever_one --> qa;
-  	retriever_two --> qa;
-  	rewrite_query --> analyzer_one;
-  	rewrite_query -.-> retriever_two;
-  
-  '''
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio].1
-  dict({
-    '$defs': dict({
-      'InnerObject': dict({
-        'properties': dict({
-          'yo': dict({
-            'title': 'Yo',
-            'type': 'integer',
-          }),
-        }),
-        'required': list([
-          'yo',
-        ]),
-        'title': 'InnerObject',
-        'type': 'object',
-      }),
-    }),
-    'properties': dict({
-      'answer': dict({
-        'anyOf': list([
-          dict({
-            'type': 'string',
-          }),
-          dict({
-            'type': 'null',
-          }),
-        ]),
-        'default': None,
-        'title': 'Answer',
-      }),
-      'docs': dict({
-        'items': dict({
-          'type': 'string',
-        }),
-        'title': 'Docs',
-        'type': 'array',
-      }),
-      'inner': dict({
-        '$ref': '#/$defs/InnerObject',
-      }),
-      'query': dict({
-        'title': 'Query',
-        'type': 'string',
-      }),
-    }),
-    'required': list([
-      'query',
-      'inner',
-      'docs',
-    ]),
-    'title': 'State',
-    'type': 'object',
-  })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio].2
-  dict({
-    '$defs': dict({
-      'InnerObject': dict({
-        'properties': dict({
-          'yo': dict({
-            'title': 'Yo',
-            'type': 'integer',
-          }),
-        }),
-        'required': list([
-          'yo',
-        ]),
-        'title': 'InnerObject',
-        'type': 'object',
-      }),
-    }),
-    'properties': dict({
-      'answer': dict({
-        'anyOf': list([
-          dict({
-            'type': 'string',
-          }),
-          dict({
-            'type': 'null',
-          }),
-        ]),
-        'default': None,
-        'title': 'Answer',
-      }),
-      'docs': dict({
-        'items': dict({
-          'type': 'string',
-        }),
-        'title': 'Docs',
-        'type': 'array',
-      }),
-      'inner': dict({
-        '$ref': '#/$defs/InnerObject',
-      }),
-      'query': dict({
-        'title': 'Query',
-        'type': 'string',
-      }),
-    }),
-    'required': list([
-      'query',
-      'inner',
-      'docs',
-    ]),
-    'title': 'State',
-    'type': 'object',
-  })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_pipe]
-  '''
-  graph TD;
-  	__start__ --> rewrite_query;
-  	analyzer_one --> retriever_one;
-  	qa --> __end__;
-  	retriever_one --> qa;
-  	retriever_two --> qa;
-  	rewrite_query --> analyzer_one;
-  	rewrite_query -.-> retriever_two;
-  
-  '''
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_pipe].1
-  dict({
-    '$defs': dict({
-      'InnerObject': dict({
-        'properties': dict({
-          'yo': dict({
-            'title': 'Yo',
-            'type': 'integer',
-          }),
-        }),
-        'required': list([
-          'yo',
-        ]),
-        'title': 'InnerObject',
-        'type': 'object',
-      }),
-    }),
-    'properties': dict({
-      'answer': dict({
-        'anyOf': list([
-          dict({
-            'type': 'string',
-          }),
-          dict({
-            'type': 'null',
-          }),
-        ]),
-        'default': None,
-        'title': 'Answer',
-      }),
-      'docs': dict({
-        'items': dict({
-          'type': 'string',
-        }),
-        'title': 'Docs',
-        'type': 'array',
-      }),
-      'inner': dict({
-        '$ref': '#/$defs/InnerObject',
-      }),
-      'query': dict({
-        'title': 'Query',
-        'type': 'string',
-      }),
-    }),
-    'required': list([
-      'query',
-      'inner',
-      'docs',
-    ]),
-    'title': 'State',
-    'type': 'object',
-  })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_pipe].2
-  dict({
-    '$defs': dict({
-      'InnerObject': dict({
-        'properties': dict({
-          'yo': dict({
-            'title': 'Yo',
-            'type': 'integer',
-          }),
-        }),
-        'required': list([
-          'yo',
-        ]),
-        'title': 'InnerObject',
-        'type': 'object',
-      }),
-    }),
-    'properties': dict({
-      'answer': dict({
-        'anyOf': list([
-          dict({
-            'type': 'string',
-          }),
-          dict({
-            'type': 'null',
-          }),
-        ]),
-        'default': None,
-        'title': 'Answer',
-      }),
-      'docs': dict({
-        'items': dict({
-          'type': 'string',
-        }),
-        'title': 'Docs',
-        'type': 'array',
-      }),
-      'inner': dict({
-        '$ref': '#/$defs/InnerObject',
-      }),
-      'query': dict({
-        'title': 'Query',
-        'type': 'string',
-      }),
-    }),
-    'required': list([
-      'query',
-      'inner',
-      'docs',
-    ]),
-    'title': 'State',
-    'type': 'object',
-  })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_pool]
-  '''
-  graph TD;
-  	__start__ --> rewrite_query;
-  	analyzer_one --> retriever_one;
-  	qa --> __end__;
-  	retriever_one --> qa;
-  	retriever_two --> qa;
-  	rewrite_query --> analyzer_one;
-  	rewrite_query -.-> retriever_two;
-  
-  '''
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_pool].1
-  dict({
-    '$defs': dict({
-      'InnerObject': dict({
-        'properties': dict({
-          'yo': dict({
-            'title': 'Yo',
-            'type': 'integer',
-          }),
-        }),
-        'required': list([
-          'yo',
-        ]),
-        'title': 'InnerObject',
-        'type': 'object',
-      }),
-    }),
-    'properties': dict({
-      'answer': dict({
-        'anyOf': list([
-          dict({
-            'type': 'string',
-          }),
-          dict({
-            'type': 'null',
-          }),
-        ]),
-        'default': None,
-        'title': 'Answer',
-      }),
-      'docs': dict({
-        'items': dict({
-          'type': 'string',
-        }),
-        'title': 'Docs',
-        'type': 'array',
-      }),
-      'inner': dict({
-        '$ref': '#/$defs/InnerObject',
-      }),
-      'query': dict({
-        'title': 'Query',
-        'type': 'string',
-      }),
-    }),
-    'required': list([
-      'query',
-      'inner',
-      'docs',
-    ]),
-    'title': 'State',
-    'type': 'object',
-  })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_pool].2
-  dict({
-    '$defs': dict({
-      'InnerObject': dict({
-        'properties': dict({
-          'yo': dict({
-            'title': 'Yo',
-            'type': 'integer',
-          }),
-        }),
-        'required': list([
-          'yo',
-        ]),
-        'title': 'InnerObject',
-        'type': 'object',
-      }),
-    }),
-    'properties': dict({
-      'answer': dict({
-        'anyOf': list([
-          dict({
-            'type': 'string',
-          }),
-          dict({
-            'type': 'null',
-          }),
-        ]),
-        'default': None,
-        'title': 'Answer',
-      }),
-      'docs': dict({
-        'items': dict({
-          'type': 'string',
-        }),
-        'title': 'Docs',
-        'type': 'array',
-      }),
-      'inner': dict({
-        '$ref': '#/$defs/InnerObject',
-      }),
-      'query': dict({
-        'title': 'Query',
-        'type': 'string',
-      }),
-    }),
-    'required': list([
-      'query',
-      'inner',
-      'docs',
-    ]),
-    'title': 'State',
-    'type': 'object',
-  })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_shallow]
-  '''
-  graph TD;
-  	__start__ --> rewrite_query;
-  	analyzer_one --> retriever_one;
-  	qa --> __end__;
-  	retriever_one --> qa;
-  	retriever_two --> qa;
-  	rewrite_query --> analyzer_one;
-  	rewrite_query -.-> retriever_two;
-  
-  '''
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_shallow].1
-  dict({
-    '$defs': dict({
-      'InnerObject': dict({
-        'properties': dict({
-          'yo': dict({
-            'title': 'Yo',
-            'type': 'integer',
-          }),
-        }),
-        'required': list([
-          'yo',
-        ]),
-        'title': 'InnerObject',
-        'type': 'object',
-      }),
-    }),
-    'properties': dict({
-      'answer': dict({
-        'anyOf': list([
-          dict({
-            'type': 'string',
-          }),
-          dict({
-            'type': 'null',
-          }),
-        ]),
-        'default': None,
-        'title': 'Answer',
-      }),
-      'docs': dict({
-        'items': dict({
-          'type': 'string',
-        }),
-        'title': 'Docs',
-        'type': 'array',
-      }),
-      'inner': dict({
-        '$ref': '#/$defs/InnerObject',
-      }),
-      'query': dict({
-        'title': 'Query',
-        'type': 'string',
-      }),
-    }),
-    'required': list([
-      'query',
-      'inner',
-      'docs',
-    ]),
-    'title': 'State',
-    'type': 'object',
-  })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_shallow].2
-  dict({
-    '$defs': dict({
-      'InnerObject': dict({
-        'properties': dict({
-          'yo': dict({
-            'title': 'Yo',
-            'type': 'integer',
-          }),
-        }),
-        'required': list([
-          'yo',
-        ]),
-        'title': 'InnerObject',
-        'type': 'object',
-      }),
-    }),
-    'properties': dict({
-      'answer': dict({
-        'anyOf': list([
-          dict({
-            'type': 'string',
-          }),
-          dict({
-            'type': 'null',
-          }),
-        ]),
-        'default': None,
-        'title': 'Answer',
-      }),
-      'docs': dict({
-        'items': dict({
-          'type': 'string',
-        }),
-        'title': 'Docs',
-        'type': 'array',
-      }),
-      'inner': dict({
-        '$ref': '#/$defs/InnerObject',
-      }),
-      'query': dict({
-        'title': 'Query',
-        'type': 'string',
-      }),
-    }),
-    'required': list([
-      'query',
-      'inner',
-      'docs',
-    ]),
-    'title': 'State',
-    'type': 'object',
-  })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[sqlite_aio]
-  '''
-  graph TD;
-  	__start__ --> rewrite_query;
-  	analyzer_one --> retriever_one;
-  	qa --> __end__;
-  	retriever_one --> qa;
-  	retriever_two --> qa;
-  	rewrite_query --> analyzer_one;
-  	rewrite_query -.-> retriever_two;
-  
-  '''
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[sqlite_aio].1
-  dict({
-    '$defs': dict({
-      'InnerObject': dict({
-        'properties': dict({
-          'yo': dict({
-            'title': 'Yo',
-            'type': 'integer',
-          }),
-        }),
-        'required': list([
-          'yo',
-        ]),
-        'title': 'InnerObject',
-        'type': 'object',
-      }),
-    }),
-    'properties': dict({
-      'answer': dict({
-        'anyOf': list([
-          dict({
-            'type': 'string',
-          }),
-          dict({
-            'type': 'null',
-          }),
-        ]),
-        'default': None,
-        'title': 'Answer',
-      }),
-      'docs': dict({
-        'items': dict({
-          'type': 'string',
-        }),
-        'title': 'Docs',
-        'type': 'array',
-      }),
-      'inner': dict({
-        '$ref': '#/$defs/InnerObject',
-      }),
-      'query': dict({
-        'title': 'Query',
-        'type': 'string',
-      }),
-    }),
-    'required': list([
-      'query',
-      'inner',
-      'docs',
-    ]),
-    'title': 'State',
-    'type': 'object',
-  })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[sqlite_aio].2
-  dict({
-    '$defs': dict({
-      'InnerObject': dict({
-        'properties': dict({
-          'yo': dict({
-            'title': 'Yo',
-            'type': 'integer',
-          }),
-        }),
-        'required': list([
-          'yo',
-        ]),
-        'title': 'InnerObject',
-        'type': 'object',
-      }),
-    }),
-    'properties': dict({
-      'answer': dict({
-        'anyOf': list([
-          dict({
-            'type': 'string',
-          }),
-          dict({
-            'type': 'null',
-          }),
-        ]),
-        'default': None,
-        'title': 'Answer',
-      }),
-      'docs': dict({
-        'items': dict({
-          'type': 'string',
-        }),
-        'title': 'Docs',
-        'type': 'array',
-      }),
-      'inner': dict({
-        '$ref': '#/$defs/InnerObject',
-      }),
-      'query': dict({
-        'title': 'Query',
-        'type': 'string',
-      }),
-    }),
-    'required': list([
-      'query',
-      'inner',
-      'docs',
-    ]),
-    'title': 'State',
-    'type': 'object',
-  })
-# ---
 # name: test_send_react_interrupt_control[memory]
   '''
-  %%{init: {'flowchart': {'curve': 'linear'}}}%%
+  ---
+  config:
+    flowchart:
+      curve: linear
+  ---
   graph TD;
   	__start__([<p>__start__</p>]):::first
   	agent(agent)
-  	foo([foo]):::last
-  	__start__ --> agent;
-  	agent -.-> foo;
-  	classDef default fill:#f2f0ff,line-height:1.2
-  	classDef first fill-opacity:0
-  	classDef last fill:#bfb6fc
-  
-  '''
-# ---
-# name: test_send_react_interrupt_control[postgres_aio]
-  '''
-  %%{init: {'flowchart': {'curve': 'linear'}}}%%
-  graph TD;
-  	__start__([<p>__start__</p>]):::first
-  	agent(agent)
-  	foo([foo]):::last
-  	__start__ --> agent;
-  	agent -.-> foo;
-  	classDef default fill:#f2f0ff,line-height:1.2
-  	classDef first fill-opacity:0
-  	classDef last fill:#bfb6fc
-  
-  '''
-# ---
-# name: test_send_react_interrupt_control[postgres_aio_pipe]
-  '''
-  %%{init: {'flowchart': {'curve': 'linear'}}}%%
-  graph TD;
-  	__start__([<p>__start__</p>]):::first
-  	agent(agent)
-  	foo([foo]):::last
-  	__start__ --> agent;
-  	agent -.-> foo;
-  	classDef default fill:#f2f0ff,line-height:1.2
-  	classDef first fill-opacity:0
-  	classDef last fill:#bfb6fc
-  
-  '''
-# ---
-# name: test_send_react_interrupt_control[postgres_aio_pool]
-  '''
-  %%{init: {'flowchart': {'curve': 'linear'}}}%%
-  graph TD;
-  	__start__([<p>__start__</p>]):::first
-  	agent(agent)
-  	foo([foo]):::last
-  	__start__ --> agent;
-  	agent -.-> foo;
-  	classDef default fill:#f2f0ff,line-height:1.2
-  	classDef first fill-opacity:0
-  	classDef last fill:#bfb6fc
-  
-  '''
-# ---
-# name: test_send_react_interrupt_control[postgres_aio_shallow]
-  '''
-  %%{init: {'flowchart': {'curve': 'linear'}}}%%
-  graph TD;
-  	__start__([<p>__start__</p>]):::first
-  	agent(agent)
-  	foo([foo]):::last
-  	__start__ --> agent;
-  	agent -.-> foo;
-  	classDef default fill:#f2f0ff,line-height:1.2
-  	classDef first fill-opacity:0
-  	classDef last fill:#bfb6fc
-  
-  '''
-# ---
-# name: test_send_react_interrupt_control[sqlite_aio]
-  '''
-  %%{init: {'flowchart': {'curve': 'linear'}}}%%
-  graph TD;
-  	__start__([<p>__start__</p>]):::first
-  	agent(agent)
-  	foo([foo]):::last
+  	foo(foo)
+  	__end__([<p>__end__</p>]):::last
   	__start__ --> agent;
   	agent -.-> foo;
+  	foo --> __end__;
   	classDef default fill:#f2f0ff,line-height:1.2
   	classDef first fill-opacity:0
   	classDef last fill:#bfb6fc

libs/langgraph/tests/any_str.py

@@ -1,5 +1,6 @@
 import re
-from typing import Any, Sequence, Union
+from collections.abc import Sequence
+from typing import Any, Union
 
 from typing_extensions import Self
 

libs/langgraph/tests/conftest.py

@@ -1,6 +1,7 @@
 import sys
+from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
-from typing import AsyncIterator, Optional
+from typing import Optional
 from uuid import UUID, uuid4
 
 import pytest

libs/langgraph/tests/fake_chat.py

@@ -1,5 +1,6 @@
 import re
-from typing import Any, AsyncIterator, Iterator, List, Optional, cast
+from collections.abc import AsyncIterator, Iterator
+from typing import Any, Optional, cast
 
 from langchain_core.callbacks import (
     AsyncCallbackManagerForLLMRun,
@@ -20,8 +21,8 @@ class FakeChatModel(GenericFakeChatModel):
 
     def _generate(
         self,
-        messages: List[BaseMessage],
-        stop: Optional[List[str]] = None,
+        messages: list[BaseMessage],
+        stop: Optional[list[str]] = None,
         run_manager: Optional[CallbackManagerForLLMRun] = None,
         **kwargs: Any,
     ) -> ChatResult:
@@ -42,8 +43,8 @@ class FakeChatModel(GenericFakeChatModel):
 
     def _stream(
         self,
-        messages: List[BaseMessage],
-        stop: Optional[List[str]] = None,
+        messages: list[BaseMessage],
+        stop: Optional[list[str]] = None,
         run_manager: Optional[CallbackManagerForLLMRun] = None,
         **kwargs: Any,
     ) -> Iterator[ChatGenerationChunk]:
@@ -90,8 +91,8 @@ class FakeChatModel(GenericFakeChatModel):
 
     async def _astream(
         self,
-        messages: List[BaseMessage],
-        stop: Optional[List[str]] = None,
+        messages: list[BaseMessage],
+        stop: Optional[list[str]] = None,
         run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
         **kwargs: Any,
     ) -> AsyncIterator[ChatGenerationChunk]:

libs/langgraph/tests/test_channels.py

@@ -1,5 +1,6 @@
 import operator
-from typing import Sequence, Union
+from collections.abc import Sequence
+from typing import Union
 
 import pytest
 
@@ -33,7 +34,7 @@ def test_last_value() -> None:
 
 def test_topic() -> None:
     channel = Topic(str).from_checkpoint(MISSING)
-    assert channel.ValueType is Sequence[str]
+    assert channel.ValueType == Sequence[str]
     assert channel.UpdateType is Union[str, list[str]]
 
     assert channel.update(["a", "b"])
@@ -57,7 +58,7 @@ def test_topic() -> None:
 
 def test_topic_accumulate() -> None:
     channel = Topic(str, accumulate=True).from_checkpoint(MISSING)
-    assert channel.ValueType is Sequence[str]
+    assert channel.ValueType == Sequence[str]
     assert channel.UpdateType is Union[str, list[str]]
 
     assert channel.update(["a", "b"])

libs/langgraph/tests/test_checkpoint_migration.py

@@ -1575,7 +1575,7 @@ def test_migrate_checkpoints(source: str, target: str) -> None:
         # check that the migrated checkpoint matches the target checkpoint
         assert (
             migrated == target_checkpoint.checkpoint
-        ), "Checkpoint mismatch at index {}".format(idx)
+        ), f"Checkpoint mismatch at index {idx}"
 
 
 @NEEDS_CONTEXTVARS

libs/langgraph/tests/test_io.py

@@ -1,4 +1,4 @@
-from typing import Iterator
+from collections.abc import Iterator
 
 from langgraph.pregel.io import single
 

libs/langgraph/tests/test_large_cases.py

@@ -2,9 +2,10 @@ import json
 import operator
 import re
 import time
+from collections.abc import Iterator
 from contextlib import contextmanager
 from dataclasses import replace
-from typing import Annotated, Any, Iterator, Literal, Optional, Union, cast
+from typing import Annotated, Any, Literal, Optional, Union, cast
 
 import httpx
 import pytest
@@ -587,12 +588,10 @@ def test_conditional_graph(
 
     app = workflow.compile()
 
-    if SHOULD_CHECK_SNAPSHOTS:
+    if SHOULD_CHECK_SNAPSHOTS and checkpointer_name == "memory":
         assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
         assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
         assert app.get_graph().draw_mermaid() == snapshot
-        assert json.dumps(app.get_graph(xray=True).to_json(), indent=2) == snapshot
-        assert app.get_graph(xray=True).draw_mermaid(with_styles=False) == snapshot
 
     assert app.invoke({"input": "what is weather in sf"}) == {
         "input": "what is weather in sf",
@@ -722,10 +721,6 @@ def test_conditional_graph(
     )
     config = {"configurable": {"thread_id": "1"}}
 
-    if SHOULD_CHECK_SNAPSHOTS:
-        assert app_w_interrupt.get_graph().to_json() == snapshot
-        assert app_w_interrupt.get_graph().draw_mermaid() == snapshot
-
     assert [
         c for c in app_w_interrupt.stream({"input": "what is weather in sf"}, config)
     ] == [
@@ -1538,7 +1533,7 @@ def test_conditional_state_graph(
 
     app = workflow.compile()
 
-    if SHOULD_CHECK_SNAPSHOTS:
+    if SHOULD_CHECK_SNAPSHOTS and checkpointer_name == "memory":
         assert json.dumps(app.get_input_schema().model_json_schema()) == snapshot
         assert json.dumps(app.get_output_schema().model_json_schema()) == snapshot
         assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
@@ -3774,7 +3769,7 @@ def test_message_graph(
     # meaning you can use it as you would any other runnable
     app = workflow.compile()
 
-    if SHOULD_CHECK_SNAPSHOTS:
+    if SHOULD_CHECK_SNAPSHOTS and checkpointer_name == "memory":
         assert json.dumps(app.get_input_schema().model_json_schema()) == snapshot
         assert json.dumps(app.get_output_schema().model_json_schema()) == snapshot
         assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
@@ -6234,10 +6229,13 @@ def test_start_branch_then(
     tool_two_graph.add_node("tool_two_slow", tool_two_slow)
     tool_two_graph.add_node("tool_two_fast", tool_two_fast)
     tool_two_graph.set_conditional_entry_point(
-        lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast", then=END
+        lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast",
+        then=END,
+        path_map=["tool_two_slow", "tool_two_fast"],
     )
     tool_two = tool_two_graph.compile()
-    assert tool_two.get_graph().draw_mermaid() == snapshot
+    if checkpointer_name == "memory":
+        assert tool_two.get_graph().draw_mermaid() == snapshot
 
     assert tool_two.invoke({"my_key": "value", "market": "DE"}) == {
         "my_key": "value slow",
@@ -6516,6 +6514,7 @@ def test_branch_then(
     tool_two_graph.add_conditional_edges(
         source="prepare",
         path=lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast",
+        path_map=["tool_two_slow", "tool_two_fast"],
         then="finish",
     )
     tool_two_graph.add_node("prepare", lambda s: {"my_key": " prepared"})
@@ -6523,8 +6522,10 @@ def test_branch_then(
     tool_two_graph.add_node("tool_two_fast", lambda s: {"my_key": " fast"})
     tool_two_graph.add_node("finish", lambda s: {"my_key": " finished"})
     tool_two = tool_two_graph.compile()
-    assert tool_two.get_graph().draw_mermaid(with_styles=False) == snapshot
-    assert tool_two.get_graph().draw_mermaid() == snapshot
+
+    if checkpointer_name == "memory":
+        assert tool_two.get_graph().draw_mermaid(with_styles=False) == snapshot
+        assert tool_two.get_graph().draw_mermaid() == snapshot
 
     assert tool_two.invoke({"my_key": "value", "market": "DE"}, debug=1) == {
         "my_key": "value prepared slow finished",
@@ -9856,7 +9857,9 @@ def test_send_react_interrupt_control(
     builder.add_node(foo)
     builder.add_edge(START, "agent")
     graph = builder.compile()
-    assert graph.get_graph().draw_mermaid() == snapshot
+
+    if checkpointer_name == "memory":
+        assert graph.get_graph().draw_mermaid() == snapshot
 
     assert graph.invoke({"messages": [HumanMessage("hello")]}) == {
         "messages": [
@@ -10187,12 +10190,17 @@ def test_weather_subgraph(
     graph.add_node(normal_llm_node)
     graph.add_node("weather_graph", weather_graph)
     graph.add_edge(START, "router_node")
-    graph.add_conditional_edges("router_node", route_after_prediction)
+    graph.add_conditional_edges(
+        "router_node",
+        route_after_prediction,
+        path_map=["weather_graph", "normal_llm_node"],
+    )
     graph.add_edge("normal_llm_node", END)
     graph.add_edge("weather_graph", END)
     graph = graph.compile(checkpointer=checkpointer)
 
-    assert graph.get_graph(xray=1).draw_mermaid() == snapshot
+    if checkpointer_name == "memory":
+        assert graph.get_graph(xray=1).draw_mermaid() == snapshot
 
     config = {"configurable": {"thread_id": "1"}}
     thread2 = {"configurable": {"thread_id": "2"}}

libs/langgraph/tests/test_large_cases_async.py

@@ -2,11 +2,11 @@ import asyncio
 import operator
 import re
 import sys
+from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
 from typing import (
     Annotated,
     Any,
-    AsyncIterator,
     Literal,
     Optional,
     Union,
@@ -3805,7 +3805,7 @@ async def test_in_one_fan_out_out_one_graph_state() -> None:
         docs: Annotated[list[str], operator.add]
 
     async def rewrite_query(data: State) -> State:
-        return {"query": f'query: {data["query"]}'}
+        return {"query": f"query: {data['query']}"}
 
     async def retriever_one(data: State) -> State:
         await asyncio.sleep(0.1)
@@ -7041,7 +7041,11 @@ async def test_weather_subgraph(
     graph.add_node(normal_llm_node)
     graph.add_node("weather_graph", weather_graph)
     graph.add_edge(START, "router_node")
-    graph.add_conditional_edges("router_node", route_after_prediction)
+    graph.add_conditional_edges(
+        "router_node",
+        route_after_prediction,
+        path_map=["weather_graph", "normal_llm_node"],
+    )
     graph.add_edge("normal_llm_node", END)
     graph.add_edge("weather_graph", END)
 
@@ -7051,8 +7055,6 @@ async def test_weather_subgraph(
     async with awith_checkpointer(checkpointer_name) as checkpointer:
         graph = graph.compile(checkpointer=checkpointer)
 
-        assert graph.get_graph(xray=1).draw_mermaid() == snapshot
-
         config = {"configurable": {"thread_id": "1"}}
         thread2 = {"configurable": {"thread_id": "2"}}
         inputs = {"messages": [{"role": "user", "content": "what's the weather in sf"}]}

libs/langgraph/tests/test_pregel.py

@@ -1,35 +1,24 @@
-import datetime
-import decimal
 import enum
 import functools
 import gc
-import ipaddress
 import json
 import logging
 import operator
-import pathlib
-import re
 import threading
 import time
 import uuid
 import warnings
 from collections import Counter, deque
+from collections.abc import Generator, Iterator, Sequence
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import contextmanager
 from dataclasses import dataclass, field
-from enum import Enum
 from random import randrange
 from typing import (
     Annotated,
     Any,
-    Dict,
-    Generator,
-    Iterator,
-    List,
     Literal,
     Optional,
-    Sequence,
-    Tuple,
     Union,
     get_type_hints,
 )
@@ -249,7 +238,7 @@ def test_checkpoint_errors() -> None:
 
     class FaultyPutWritesCheckpointer(InMemorySaver):
         def put_writes(
-            self, config: RunnableConfig, writes: List[Tuple[str, Any]], task_id: str
+            self, config: RunnableConfig, writes: list[tuple[str, Any]], task_id: str
         ) -> RunnableConfig:
             raise ValueError("Faulty put_writes")
 
@@ -454,7 +443,7 @@ def test_reducer_before_first_node() -> None:
 
     class State(TypedDict):
         hello: str
-        messages: Annotated[List[str], add_messages]
+        messages: Annotated[list[str], add_messages]
 
     def node_a(state: State) -> State:
         assert state == {
@@ -1137,7 +1126,7 @@ def test_pending_writes_resume(
         value: Annotated[int, operator.add]
 
     class AwhileMaker:
-        def __init__(self, sleep: float, rtn: Union[Dict, Exception]) -> None:
+        def __init__(self, sleep: float, rtn: Union[dict, Exception]) -> None:
             self.sleep = sleep
             self.rtn = rtn
             self.reset()
@@ -2153,7 +2142,7 @@ def test_conditional_entrypoint_to_multiple_state_graph(
 
     workflow.add_node("get_weather", get_weather)
     workflow.add_edge("get_weather", END)
-    workflow.set_conditional_entry_point(continue_to_weather)
+    workflow.set_conditional_entry_point(continue_to_weather, path_map=["get_weather"])
 
     app = workflow.compile()
 
@@ -2420,7 +2409,8 @@ def test_in_one_fan_out_state_graph_waiting_edge(
 
     app = workflow.compile()
 
-    assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
+    if checkpointer_name == "memory":
+        assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
 
     assert app.invoke({"query": "what is weather in sf"}) == {
         "query": "analyzed: query: what is weather in sf",
@@ -2566,7 +2556,8 @@ def test_in_one_fan_out_state_graph_waiting_edge_via_branch(
 
     app = workflow.compile()
 
-    assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
+    if checkpointer_name == "memory":
+        assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
 
     assert app.invoke({"query": "what is weather in sf"}, debug=True) == {
         "query": "analyzed: query: what is weather in sf",
@@ -2716,9 +2707,10 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1(
 
     app = workflow.compile()
 
-    assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
-    assert app.get_input_jsonschema() == snapshot
-    assert app.get_output_jsonschema() == snapshot
+    if checkpointer_name == "memory":
+        assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
+        assert app.get_input_jsonschema() == snapshot
+        assert app.get_output_jsonschema() == snapshot
 
     with pytest.raises(ValidationError), assert_ctx_once():
         app.invoke({"query": {}})
@@ -2906,7 +2898,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2(
 
     app = workflow.compile()
 
-    if SHOULD_CHECK_SNAPSHOTS:
+    if SHOULD_CHECK_SNAPSHOTS and checkpointer_name == "memory":
         assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
         assert app.get_input_schema().model_json_schema() == snapshot
         assert app.get_output_schema().model_json_schema() == snapshot
@@ -2970,8 +2962,6 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2(
 
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
 def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_input(
-    snapshot: SnapshotAssertion,
-    mocker: MockerFixture,
     request: pytest.FixtureRequest,
     checkpointer_name: str,
 ) -> None:
@@ -3101,328 +3091,6 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_inp
     }
 
 
-@pytest.mark.parametrize("version", ["v1", "v2"])
-def test_nested_pydantic_models(version: str) -> None:
-    """Test that nested Pydantic models are properly constructed from leaf nodes up."""
-
-    # Define nested Pydantic models
-    # Import necessary modules
-
-    if version == "v1":
-        from pydantic.v1 import (  # type: ignore
-            BaseModel,
-            ByteSize,
-            Field,
-            SecretStr,
-            confloat,
-            conint,
-            conlist,
-            constr,
-        )
-    else:
-        from pydantic import (  # type: ignore
-            BaseModel,
-            ByteSize,
-            Field,
-            SecretStr,
-            confloat,
-            conint,
-            conlist,
-            constr,
-        )
-        from pydantic.v1 import BaseModel as BaseModelV1
-
-        if BaseModel is BaseModelV1:
-            pytest.skip("Cannot test pydantic v2 using installed version < 2")
-
-    class NestedModel(BaseModel):
-        value: int
-        name: str
-
-    # For constrained types
-    PositiveInt = Annotated[int, Field(gt=0)]
-    NonNegativeFloat = Annotated[float, Field(ge=0)]
-
-    # Enum type
-    class UserRole(Enum):
-        ADMIN = "admin"
-        USER = "user"
-        GUEST = "guest"
-
-    # Forward reference model
-    class RecursiveModel(BaseModel):
-        value: str
-        child: Optional["RecursiveModel"] = None
-
-    # Discriminated union models
-    class Cat(BaseModel):
-        pet_type: Literal["cat"]
-        meow: str
-
-    class Dog(BaseModel):
-        pet_type: Literal["dog"]
-        bark: str
-
-    # Cyclic reference model
-    class Person(BaseModel):
-        id: str
-        name: str
-        friends: list[str] = Field(default_factory=list)  # IDs of friends
-
-    if version == "v2":
-        conlist_type = conlist(item_type=int, min_length=2, max_length=5)
-    else:
-        conlist_type = conlist(item_type=int, min_items=2, max_items=5)
-
-    class State(BaseModel):
-        # Basic nested model tests
-        top_level: str
-        auuid: uuid.UUID
-        nested: NestedModel
-        optional_nested: Annotated[Optional[NestedModel], lambda x, y: y, "Foo"]
-        dict_nested: dict[str, NestedModel]
-        simple_str_list: list[str]
-        list_nested: Annotated[
-            Union[dict, list[dict[str, NestedModel]]], lambda x, y: (x or []) + [y]
-        ]
-        tuple_nested: tuple[str, NestedModel]
-        tuple_list_nested: list[tuple[int, NestedModel]]
-        complex_tuple: tuple[str, dict[str, tuple[int, NestedModel]]]
-
-        # Forward reference test
-        recursive: RecursiveModel
-
-        # Discriminated union test
-        pet: Union[Cat, Dog]
-
-        # Cyclic reference test
-        people: dict[str, Person]  # Map of ID -> Person
-
-        # Rich type adapters
-        ip_address: ipaddress.IPv4Address
-        ip_address_v6: ipaddress.IPv6Address
-        amount: decimal.Decimal
-        file_path: pathlib.Path
-        timestamp: datetime.datetime
-        date_only: datetime.date
-        time_only: datetime.time
-        duration: datetime.timedelta
-        immutable_set: frozenset[int]
-        binary_data: bytes
-        pattern: re.Pattern
-        secret: SecretStr
-        file_size: ByteSize
-
-        # Constrained types
-        positive_value: PositiveInt
-        non_negative: NonNegativeFloat
-        limited_string: constr(min_length=3, max_length=10)
-        bounded_int: conint(ge=10, le=100)
-        restricted_float: confloat(gt=0, lt=1)
-        required_list: conlist_type
-
-        # Enum & Literal
-        role: UserRole
-        status: Literal["active", "inactive", "pending"]
-
-        # Annotated & NewType
-        validated_age: Annotated[int, Field(gt=0, lt=120)]
-
-        # Generic containers with validators
-        decimal_list: List[decimal.Decimal]
-        id_tuple: tuple[uuid.UUID, uuid.UUID]
-
-    inputs = {
-        # Basic nested models
-        "top_level": "initial",
-        "auuid": str(uuid.uuid4()),
-        "nested": {"value": 42, "name": "test"},
-        "optional_nested": {"value": 10, "name": "optional"},
-        "dict_nested": {"a": {"value": 5, "name": "a"}},
-        "list_nested": [{"a": {"value": 6, "name": "b"}}],
-        "tuple_nested": ["tuple-key", {"value": 7, "name": "tuple-value"}],
-        "tuple_list_nested": [[1, {"value": 8, "name": "tuple-in-list"}]],
-        "simple_str_list": ["siss", "boom", "bah"],
-        "complex_tuple": [
-            "complex",
-            {"nested": [9, {"value": 10, "name": "deep"}]},
-        ],
-        # Forward reference
-        "recursive": {"value": "parent", "child": {"value": "child", "child": None}},
-        # Discriminated union (using a cat in this case)
-        "pet": {"pet_type": "cat", "meow": "meow!"},
-        # Cyclic references
-        "people": {
-            "1": {
-                "id": "1",
-                "name": "Alice",
-                "friends": ["2", "3"],  # Alice is friends with Bob and Charlie
-            },
-            "2": {
-                "id": "2",
-                "name": "Bob",
-                "friends": ["1"],  # Bob is friends with Alice
-            },
-            "3": {
-                "id": "3",
-                "name": "Charlie",
-                "friends": ["1", "2"],  # Charlie is friends with Alice and Bob
-            },
-        },
-        # Rich type adapters
-        "ip_address": "192.168.1.1",
-        "ip_address_v6": "2001:db8::1",
-        "amount": "123.45",
-        "file_path": "/tmp/test.txt",
-        "timestamp": "2025-04-07T10:58:04",
-        "date_only": "2025-04-07",
-        "time_only": "10:58:04",
-        "duration": 3600,  # seconds
-        "immutable_set": [1, 2, 3, 4],
-        "binary_data": b"hello world",
-        "pattern": "^test$",
-        "secret": "password123",
-        "file_size": 1024,
-        # Constrained types
-        "positive_value": 42,
-        "non_negative": 0.0,
-        "limited_string": "test",
-        "bounded_int": 50,
-        "restricted_float": 0.5,
-        "required_list": [10, 20, 30],
-        # Enum & Literal
-        "role": "admin",
-        "status": "active",
-        # Annotated & NewType
-        "validated_age": 30,
-        # Generic containers with validators
-        "decimal_list": ["10.5", "20.75", "30.25"],
-        "id_tuple": [str(uuid.uuid4()), str(uuid.uuid4())],
-    }
-
-    update = {"top_level": "updated", "nested": {"value": 100, "name": "updated"}}
-
-    expected = State(**inputs)
-
-    def node_fn(state: State) -> dict:
-        # Basic assertions
-        assert isinstance(state.auuid, uuid.UUID)
-        assert state == expected
-
-        # Rich type assertions
-        assert isinstance(state.ip_address, ipaddress.IPv4Address)
-        assert isinstance(state.ip_address_v6, ipaddress.IPv6Address)
-        assert isinstance(state.amount, decimal.Decimal)
-        assert isinstance(state.file_path, pathlib.Path)
-        assert isinstance(state.timestamp, datetime.datetime)
-        assert isinstance(state.date_only, datetime.date)
-        assert isinstance(state.time_only, datetime.time)
-        assert isinstance(state.duration, datetime.timedelta)
-        assert isinstance(state.immutable_set, frozenset)
-        assert isinstance(state.binary_data, bytes)
-        assert isinstance(state.pattern, re.Pattern)
-
-        # Constrained types
-        assert state.positive_value > 0
-        assert state.non_negative >= 0
-        assert 3 <= len(state.limited_string) <= 10
-        assert 10 <= state.bounded_int <= 100
-        assert 0 < state.restricted_float < 1
-        assert 2 <= len(state.required_list) <= 5
-
-        # Enum & Literal
-        assert state.role == UserRole.ADMIN
-        assert state.status == "active"
-
-        # Annotated
-        assert 0 < state.validated_age < 120
-
-        # Generic containers
-        assert len(state.decimal_list) == 3
-        assert len(state.id_tuple) == 2
-
-        return update
-
-    builder = StateGraph(State)
-    builder.add_node("process", node_fn)
-    builder.set_entry_point("process")
-    builder.set_finish_point("process")
-    graph = builder.compile()
-
-    result = graph.invoke(inputs.copy())
-
-    assert result == {**inputs, **update}
-
-    new_inputs = inputs.copy()
-    new_inputs["list_nested"] = {"foo": "bar"}
-    expected = State(**new_inputs)
-    assert {**new_inputs, **update} == graph.invoke(new_inputs.copy())
-
-
-def test_pydantic_state_field_validator():
-    from pydantic import BaseModel, field_validator, model_validator
-
-    class State(BaseModel):
-        name: str
-        text: str = ""
-        only_root: int = 13
-
-        @field_validator("name", mode="after")
-        @classmethod
-        def validate_name(cls, value):
-            if value[0].islower():
-                raise ValueError("Name must start with a capital letter")
-            return "Validated " + value
-
-        @model_validator(mode="before")
-        @classmethod
-        def validate_amodel(cls, values: "State"):
-            return values | {"only_root": 392}
-
-    input_state = {"name": "John"}
-
-    def process_node(state: State):
-        assert State.model_validate(input_state) == state
-        return {"text": "Hello, " + state.name + "!"}
-
-    builder = StateGraph(state_schema=State)
-    builder.add_node("process", process_node)
-    builder.add_edge(START, "process")
-    builder.add_edge("process", END)
-    g = builder.compile()
-    res = g.invoke(input_state)
-    assert res["text"] == "Hello, Validated John!"
-
-
-def test_pydantic_v1_state_root_validator():
-    from pydantic.v1 import BaseModel, root_validator
-
-    class State(BaseModel):
-        name: str
-        text: str = ""
-        only_root: int = 13
-
-        @root_validator(pre=True)
-        @classmethod
-        def validate(cls, values: dict):
-            values["name"] = "Validated " + values["name"]
-            return values | {"only_root": 396}
-
-    input_state = {"name": "John"}
-
-    def process_node(state: State):
-        assert State(**input_state) == state
-        return {"text": "Hello, " + state.name + "!"}
-
-    builder = StateGraph(state_schema=State)
-    builder.add_node("process", process_node)
-    builder.add_edge(START, "process")
-    builder.add_edge("process", END)
-    g = builder.compile()
-    res = g.invoke(input_state)
-    assert res["text"] == "Hello, Validated John!"
-
-
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
 def test_in_one_fan_out_state_graph_waiting_edge_plus_regular(
     request: pytest.FixtureRequest, checkpointer_name: str
@@ -4756,7 +4424,7 @@ def test_xray_lance(snapshot: SnapshotAssertion):
             return f"Name: {self.name}\nRole: {self.role}\nAffiliation: {self.affiliation}\nDescription: {self.description}\n"
 
     class Perspectives(BaseModel):
-        analysts: List[Analyst] = Field(
+        analysts: list[Analyst] = Field(
             description="Comprehensive list of investment analysts with their roles and affiliations.",
         )
 
@@ -4775,15 +4443,15 @@ def test_xray_lance(snapshot: SnapshotAssertion):
         )
 
     class InterviewState(TypedDict):
-        messages: Annotated[List[AnyMessage], add_messages]
+        messages: Annotated[list[AnyMessage], add_messages]
         analyst: Analyst
         section: Section
 
     class ResearchGraphState(TypedDict):
-        analysts: List[Analyst]
+        analysts: list[Analyst]
         topic: str
         max_analysts: int
-        sections: List[Section]
+        sections: list[Section]
         interviews: Annotated[list, operator.add]
 
     # Conditional edge
@@ -4804,7 +4472,9 @@ def test_xray_lance(snapshot: SnapshotAssertion):
     # Flow
     interview_builder.add_edge(START, "ask_question")
     interview_builder.add_edge("ask_question", "answer_question")
-    interview_builder.add_conditional_edges("answer_question", route_messages)
+    interview_builder.add_conditional_edges(
+        "answer_question", route_messages, ["ask_question", END]
+    )
 
     # Set up memory
     memory = InMemorySaver()
@@ -7336,6 +7006,8 @@ def test_node_destinations() -> None:
             Edge(source="__start__", target="child", data=None, conditional=False),
             Edge(source="child", target="node_b", data=None, conditional=True),
             Edge(source="child", target="node_c", data=None, conditional=True),
+            Edge(source="node_b", target="__end__", data=None, conditional=False),
+            Edge(source="node_c", target="__end__", data=None, conditional=False),
         ] == graph.edges
 
         # destinations w/ dicts
@@ -7354,6 +7026,8 @@ def test_node_destinations() -> None:
             Edge(source="__start__", target="child", data=None, conditional=False),
             Edge(source="child", target="node_b", data="foo", conditional=True),
             Edge(source="child", target="node_c", data="bar", conditional=True),
+            Edge(source="node_b", target="__end__", data=None, conditional=False),
+            Edge(source="node_c", target="__end__", data=None, conditional=False),
         ] == graph.edges
 
 
@@ -7756,7 +7430,7 @@ def test_parallel_interrupts(
     class ChildState(BaseModel):
         prompt: str = Field(..., description="What is going to be asked to the user?")
         human_input: Optional[str] = Field(None, description="What the human said")
-        human_inputs: Annotated[List[str], operator.add] = Field(
+        human_inputs: Annotated[list[str], operator.add] = Field(
             default_factory=list, description="All of my messages"
         )
 
@@ -7777,10 +7451,10 @@ def test_parallel_interrupts(
     # --- PARENT GRAPH ---
 
     class ParentState(BaseModel):
-        prompts: List[str] = Field(
+        prompts: list[str] = Field(
             ..., description="What is going to be asked to the user?"
         )
-        human_inputs: Annotated[List[str], operator.add] = Field(
+        human_inputs: Annotated[list[str], operator.add] = Field(
             default_factory=list, description="All of my messages"
         )
 
@@ -7933,7 +7607,7 @@ def test_parallel_interrupts_double(
     class ChildState(BaseModel):
         prompt: str = Field(..., description="What is going to be asked to the user?")
         human_input: Optional[str] = Field(None, description="What the human said")
-        human_inputs: Annotated[List[str], operator.add] = Field(
+        human_inputs: Annotated[list[str], operator.add] = Field(
             default_factory=list, description="All of my messages"
         )
 
@@ -7961,10 +7635,10 @@ def test_parallel_interrupts_double(
     # --- PARENT GRAPH ---
 
     class ParentState(BaseModel):
-        prompts: List[str] = Field(
+        prompts: list[str] = Field(
             ..., description="What is going to be asked to the user?"
         )
-        human_inputs: Annotated[List[str], operator.add] = Field(
+        human_inputs: Annotated[list[str], operator.add] = Field(
             default_factory=list, description="All of my messages"
         )
 

libs/langgraph/tests/test_pregel_async.py

@@ -8,20 +8,15 @@ import random
 import sys
 import uuid
 from collections import Counter, deque
+from collections.abc import AsyncGenerator, AsyncIterator, Generator
 from contextlib import asynccontextmanager, contextmanager
 from dataclasses import replace
 from time import perf_counter
 from typing import (
     Annotated,
     Any,
-    AsyncGenerator,
-    AsyncIterator,
-    Dict,
-    Generator,
-    List,
     Literal,
     Optional,
-    Tuple,
     Union,
 )
 from uuid import UUID
@@ -113,7 +108,7 @@ async def test_checkpoint_errors() -> None:
 
     class FaultyPutWritesCheckpointer(InMemorySaver):
         async def aput_writes(
-            self, config: RunnableConfig, writes: List[Tuple[str, Any]], task_id: str
+            self, config: RunnableConfig, writes: list[tuple[str, Any]], task_id: str
         ) -> RunnableConfig:
             raise ValueError("Faulty put_writes")
 
@@ -1959,7 +1954,7 @@ async def test_pending_writes_resume(
         value: Annotated[int, operator.add]
 
     class AwhileMaker:
-        def __init__(self, sleep: float, rtn: Union[Dict, Exception]) -> None:
+        def __init__(self, sleep: float, rtn: Union[dict, Exception]) -> None:
             self.sleep = sleep
             self.rtn = rtn
             self.reset()
@@ -3610,7 +3605,8 @@ async def test_send_react_interrupt_control(
     builder.add_node(foo)
     builder.add_edge(START, "agent")
     graph = builder.compile()
-    assert graph.get_graph().draw_mermaid() == snapshot
+    if checkpointer_name == "memory":
+        assert graph.get_graph().draw_mermaid() == snapshot
 
     assert await graph.ainvoke({"messages": [HumanMessage("hello")]}) == {
         "messages": [
@@ -3928,22 +3924,29 @@ async def test_max_concurrency_control(checkpointer_name: str) -> None:
     builder.add_edge(START, "1")
     graph = builder.compile()
 
-    assert (
-        graph.get_graph().draw_mermaid()
-        == """%%{init: {'flowchart': {'curve': 'linear'}}}%%
+    if checkpointer_name == "memory":
+        assert (
+            graph.get_graph().draw_mermaid()
+            == """---
+config:
+  flowchart:
+    curve: linear
+---
 graph TD;
 	__start__([<p>__start__</p>]):::first
 	1(1)
 	2(2)
-	3([3]):::last
-	__start__ --> 1;
+	3(3)
+	__end__([<p>__end__</p>]):::last
 	1 -.-> 2;
 	2 -.-> 3;
+	__start__ --> 1;
+	3 --> __end__;
 	classDef default fill:#f2f0ff,line-height:1.2
 	classDef first fill-opacity:0
 	classDef last fill:#bfb6fc
 """
-    )
+        )
 
     assert await graph.ainvoke(["0"], debug=True) == ["0", "1", *range(100), "3"]
     assert node2_max_currently == 100
@@ -4980,7 +4983,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydant
 
     app = workflow.compile()
 
-    if SHOULD_CHECK_SNAPSHOTS:
+    if SHOULD_CHECK_SNAPSHOTS and checkpointer_name == "memory":
         assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
         assert app.get_input_schema().model_json_schema() == snapshot
         assert app.get_output_schema().model_json_schema() == snapshot
@@ -6241,7 +6244,7 @@ async def test_store_injected_async(checkpointer_name: str, store_name: str) ->
         assert result == {"count": N + 1}
         returned_doc = (await the_store.aget(namespace, doc_id)).value
         assert returned_doc == {**doc, "from_thread": thread_1, "some_val": 0}
-        assert len((await the_store.asearch(namespace))) == 1
+        assert len(await the_store.asearch(namespace)) == 1
 
         # Check results after another turn of the same thread
         result = await graph.ainvoke(
@@ -6250,7 +6253,7 @@ async def test_store_injected_async(checkpointer_name: str, store_name: str) ->
         assert result == {"count": (N + 1) * 2}
         returned_doc = (await the_store.aget(namespace, doc_id)).value
         assert returned_doc == {**doc, "from_thread": thread_1, "some_val": N + 1}
-        assert len((await the_store.asearch(namespace))) == 1
+        assert len(await the_store.asearch(namespace)) == 1
 
         # Test with a different thread
         result = await graph.ainvoke(
@@ -6264,7 +6267,7 @@ async def test_store_injected_async(checkpointer_name: str, store_name: str) ->
             "some_val": 0,
         }  # Overwrites the whole doc
         assert (
-            len((await the_store.asearch(namespace))) == 1
+            len(await the_store.asearch(namespace)) == 1
         )  # still overwriting the same one
 
 

libs/langgraph/tests/test_pydantic.py

@@ -1,14 +1,26 @@
+import datetime
+import decimal
+import ipaddress
+import pathlib
+import re
 import sys
-import typing
+import uuid
+from enum import Enum
+from typing import Annotated, Literal, Optional, Union
 
-import pydantic
-import typing_extensions
+import pytest
 
+from langgraph.constants import END, START
+from langgraph.graph.state import StateGraph
 from langgraph.utils.pydantic import is_supported_by_pydantic
 
 
 def test_is_supported_by_pydantic() -> None:
     """Test if types are supported by pydantic."""
+    import typing
+
+    import pydantic
+    import typing_extensions
 
     class TypedDictExtensions(typing_extensions.TypedDict):
         x: int
@@ -41,3 +53,325 @@ def test_is_supported_by_pydantic() -> None:
         assert is_supported_by_pydantic(PydanticModelV1) is False
 
     assert is_supported_by_pydantic(int) is False
+
+
+@pytest.mark.parametrize("version", ["v1", "v2"])
+def test_nested_pydantic_models(version: str) -> None:
+    """Test that nested Pydantic models are properly constructed from leaf nodes up."""
+
+    # Define nested Pydantic models
+    # Import necessary modules
+
+    if version == "v1":
+        from pydantic.v1 import (  # type: ignore
+            BaseModel,
+            ByteSize,
+            Field,
+            SecretStr,
+            confloat,
+            conint,
+            conlist,
+            constr,
+        )
+    else:
+        from pydantic import (  # type: ignore
+            BaseModel,
+            ByteSize,
+            Field,
+            SecretStr,
+            confloat,
+            conint,
+            conlist,
+            constr,
+        )
+        from pydantic.v1 import BaseModel as BaseModelV1
+
+        if BaseModel is BaseModelV1:
+            pytest.skip("Cannot test pydantic v2 using installed version < 2")
+
+    class NestedModel(BaseModel):
+        value: int
+        name: str
+
+    # For constrained types
+    PositiveInt = Annotated[int, Field(gt=0)]
+    NonNegativeFloat = Annotated[float, Field(ge=0)]
+
+    # Enum type
+    class UserRole(Enum):
+        ADMIN = "admin"
+        USER = "user"
+        GUEST = "guest"
+
+    # Forward reference model
+    class RecursiveModel(BaseModel):
+        value: str
+        child: Optional["RecursiveModel"] = None
+
+    # Discriminated union models
+    class Cat(BaseModel):
+        pet_type: Literal["cat"]
+        meow: str
+
+    class Dog(BaseModel):
+        pet_type: Literal["dog"]
+        bark: str
+
+    # Cyclic reference model
+    class Person(BaseModel):
+        id: str
+        name: str
+        friends: list[str] = Field(default_factory=list)  # IDs of friends
+
+    if version == "v2":
+        conlist_type = conlist(item_type=int, min_length=2, max_length=5)
+    else:
+        conlist_type = conlist(item_type=int, min_items=2, max_items=5)
+
+    class State(BaseModel):
+        # Basic nested model tests
+        top_level: str
+        auuid: uuid.UUID
+        nested: NestedModel
+        optional_nested: Annotated[Optional[NestedModel], lambda x, y: y, "Foo"]
+        dict_nested: dict[str, NestedModel]
+        simple_str_list: list[str]
+        list_nested: Annotated[
+            Union[dict, list[dict[str, NestedModel]]], lambda x, y: (x or []) + [y]
+        ]
+        tuple_nested: tuple[str, NestedModel]
+        tuple_list_nested: list[tuple[int, NestedModel]]
+        complex_tuple: tuple[str, dict[str, tuple[int, NestedModel]]]
+
+        # Forward reference test
+        recursive: RecursiveModel
+
+        # Discriminated union test
+        pet: Union[Cat, Dog]
+
+        # Cyclic reference test
+        people: dict[str, Person]  # Map of ID -> Person
+
+        # Rich type adapters
+        ip_address: ipaddress.IPv4Address
+        ip_address_v6: ipaddress.IPv6Address
+        amount: decimal.Decimal
+        file_path: pathlib.Path
+        timestamp: datetime.datetime
+        date_only: datetime.date
+        time_only: datetime.time
+        duration: datetime.timedelta
+        immutable_set: frozenset[int]
+        binary_data: bytes
+        pattern: re.Pattern
+        secret: SecretStr
+        file_size: ByteSize
+
+        # Constrained types
+        positive_value: PositiveInt
+        non_negative: NonNegativeFloat
+        limited_string: constr(min_length=3, max_length=10)
+        bounded_int: conint(ge=10, le=100)
+        restricted_float: confloat(gt=0, lt=1)
+        required_list: conlist_type
+
+        # Enum & Literal
+        role: UserRole
+        status: Literal["active", "inactive", "pending"]
+
+        # Annotated & NewType
+        validated_age: Annotated[int, Field(gt=0, lt=120)]
+
+        # Generic containers with validators
+        decimal_list: list[decimal.Decimal]
+        id_tuple: tuple[uuid.UUID, uuid.UUID]
+
+    inputs = {
+        # Basic nested models
+        "top_level": "initial",
+        "auuid": str(uuid.uuid4()),
+        "nested": {"value": 42, "name": "test"},
+        "optional_nested": {"value": 10, "name": "optional"},
+        "dict_nested": {"a": {"value": 5, "name": "a"}},
+        "list_nested": [{"a": {"value": 6, "name": "b"}}],
+        "tuple_nested": ["tuple-key", {"value": 7, "name": "tuple-value"}],
+        "tuple_list_nested": [[1, {"value": 8, "name": "tuple-in-list"}]],
+        "simple_str_list": ["siss", "boom", "bah"],
+        "complex_tuple": [
+            "complex",
+            {"nested": [9, {"value": 10, "name": "deep"}]},
+        ],
+        # Forward reference
+        "recursive": {"value": "parent", "child": {"value": "child", "child": None}},
+        # Discriminated union (using a cat in this case)
+        "pet": {"pet_type": "cat", "meow": "meow!"},
+        # Cyclic references
+        "people": {
+            "1": {
+                "id": "1",
+                "name": "Alice",
+                "friends": ["2", "3"],  # Alice is friends with Bob and Charlie
+            },
+            "2": {
+                "id": "2",
+                "name": "Bob",
+                "friends": ["1"],  # Bob is friends with Alice
+            },
+            "3": {
+                "id": "3",
+                "name": "Charlie",
+                "friends": ["1", "2"],  # Charlie is friends with Alice and Bob
+            },
+        },
+        # Rich type adapters
+        "ip_address": "192.168.1.1",
+        "ip_address_v6": "2001:db8::1",
+        "amount": "123.45",
+        "file_path": "/tmp/test.txt",
+        "timestamp": "2025-04-07T10:58:04",
+        "date_only": "2025-04-07",
+        "time_only": "10:58:04",
+        "duration": 3600,  # seconds
+        "immutable_set": [1, 2, 3, 4],
+        "binary_data": b"hello world",
+        "pattern": "^test$",
+        "secret": "password123",
+        "file_size": 1024,
+        # Constrained types
+        "positive_value": 42,
+        "non_negative": 0.0,
+        "limited_string": "test",
+        "bounded_int": 50,
+        "restricted_float": 0.5,
+        "required_list": [10, 20, 30],
+        # Enum & Literal
+        "role": "admin",
+        "status": "active",
+        # Annotated & NewType
+        "validated_age": 30,
+        # Generic containers with validators
+        "decimal_list": ["10.5", "20.75", "30.25"],
+        "id_tuple": [str(uuid.uuid4()), str(uuid.uuid4())],
+    }
+
+    update = {"top_level": "updated", "nested": {"value": 100, "name": "updated"}}
+
+    expected = State(**inputs)
+
+    def node_fn(state: State) -> dict:
+        # Basic assertions
+        assert isinstance(state.auuid, uuid.UUID)
+        assert state == expected
+
+        # Rich type assertions
+        assert isinstance(state.ip_address, ipaddress.IPv4Address)
+        assert isinstance(state.ip_address_v6, ipaddress.IPv6Address)
+        assert isinstance(state.amount, decimal.Decimal)
+        assert isinstance(state.file_path, pathlib.Path)
+        assert isinstance(state.timestamp, datetime.datetime)
+        assert isinstance(state.date_only, datetime.date)
+        assert isinstance(state.time_only, datetime.time)
+        assert isinstance(state.duration, datetime.timedelta)
+        assert isinstance(state.immutable_set, frozenset)
+        assert isinstance(state.binary_data, bytes)
+        assert isinstance(state.pattern, re.Pattern)
+
+        # Constrained types
+        assert state.positive_value > 0
+        assert state.non_negative >= 0
+        assert 3 <= len(state.limited_string) <= 10
+        assert 10 <= state.bounded_int <= 100
+        assert 0 < state.restricted_float < 1
+        assert 2 <= len(state.required_list) <= 5
+
+        # Enum & Literal
+        assert state.role == UserRole.ADMIN
+        assert state.status == "active"
+
+        # Annotated
+        assert 0 < state.validated_age < 120
+
+        # Generic containers
+        assert len(state.decimal_list) == 3
+        assert len(state.id_tuple) == 2
+
+        return update
+
+    builder = StateGraph(State)
+    builder.add_node("process", node_fn)
+    builder.set_entry_point("process")
+    builder.set_finish_point("process")
+    graph = builder.compile()
+
+    result = graph.invoke(inputs.copy())
+
+    assert result == {**inputs, **update}
+
+    new_inputs = inputs.copy()
+    new_inputs["list_nested"] = {"foo": "bar"}
+    expected = State(**new_inputs)
+    assert {**new_inputs, **update} == graph.invoke(new_inputs.copy())
+
+
+def test_pydantic_state_field_validator():
+    from pydantic import BaseModel, field_validator, model_validator
+
+    class State(BaseModel):
+        name: str
+        text: str = ""
+        only_root: int = 13
+
+        @field_validator("name", mode="after")
+        @classmethod
+        def validate_name(cls, value):
+            if value[0].islower():
+                raise ValueError("Name must start with a capital letter")
+            return "Validated " + value
+
+        @model_validator(mode="before")
+        @classmethod
+        def validate_amodel(cls, values: "State"):
+            return values | {"only_root": 392}
+
+    input_state = {"name": "John"}
+
+    def process_node(state: State):
+        assert State.model_validate(input_state) == state
+        return {"text": "Hello, " + state.name + "!"}
+
+    builder = StateGraph(state_schema=State)
+    builder.add_node("process", process_node)
+    builder.add_edge(START, "process")
+    builder.add_edge("process", END)
+    g = builder.compile()
+    res = g.invoke(input_state)
+    assert res["text"] == "Hello, Validated John!"
+
+
+def test_pydantic_v1_state_root_validator():
+    from pydantic.v1 import BaseModel, root_validator
+
+    class State(BaseModel):
+        name: str
+        text: str = ""
+        only_root: int = 13
+
+        @root_validator(pre=True)
+        @classmethod
+        def validate(cls, values: dict):
+            values["name"] = "Validated " + values["name"]
+            return values | {"only_root": 396}
+
+    input_state = {"name": "John"}
+
+    def process_node(state: State):
+        assert State(**input_state) == state
+        return {"text": "Hello, " + state.name + "!"}
+
+    builder = StateGraph(state_schema=State)
+    builder.add_node("process", process_node)
+    builder.add_edge(START, "process")
+    builder.add_edge("process", END)
+    g = builder.compile()
+    res = g.invoke(input_state)
+    assert res["text"] == "Hello, Validated John!"

libs/langgraph/tests/test_retry.py

@@ -226,9 +226,10 @@ def test_graph_with_jitter_retry_policy():
     )
 
     # Test graph execution with mocked random and sleep
-    with patch("random.uniform", return_value=0.05) as mock_random, patch(
-        "time.sleep"
-    ) as mock_sleep:
+    with (
+        patch("random.uniform", return_value=0.05) as mock_random,
+        patch("time.sleep") as mock_sleep,
+    ):
         result = graph.invoke({"foo": ""})
 
     # Verify retry behavior
@@ -334,8 +335,9 @@ def test_graph_with_max_attempts_exceeded():
     )
 
     # Test graph execution
-    with patch("time.sleep") as mock_sleep, pytest.raises(
-        ValueError, match="Always fails"
+    with (
+        patch("time.sleep") as mock_sleep,
+        pytest.raises(ValueError, match="Always fails"),
     ):
         graph.invoke({"foo": ""})
 

libs/langgraph/tests/test_state.py

@@ -2,13 +2,13 @@ import inspect
 import operator
 import warnings
 from dataclasses import dataclass, field
+from typing import Annotated, Any, Optional
 from typing import Annotated as Annotated2
-from typing import Any, Optional
 
 import pytest
 from langchain_core.runnables import RunnableConfig, RunnableLambda
 from pydantic.v1 import BaseModel
-from typing_extensions import Annotated, NotRequired, Required, TypedDict
+from typing_extensions import NotRequired, Required, TypedDict
 
 from langgraph.graph.state import StateGraph, _get_node_name, _warn_invalid_state_schema
 from langgraph.managed.shared_value import SharedValue

libs/langgraph/tests/test_tracing_interops.py

@@ -1,7 +1,7 @@
 import json
 import sys
 import time
-from typing import Any, Callable, Tuple, TypeVar
+from typing import Any, Callable, TypeVar
 from unittest.mock import MagicMock
 
 import langsmith as ls
@@ -35,7 +35,7 @@ T = TypeVar("T")
 
 
 def wait_for(
-    condition: Callable[[], Tuple[T, bool]],
+    condition: Callable[[], tuple[T, bool]],
     max_sleep_time: int = 10,
     sleep_time: int = 3,
 ) -> T:

libs/langgraph/tests/test_utils.py

@@ -2,11 +2,10 @@ import functools
 import sys
 import uuid
 from typing import (
+    Annotated,
     Any,
     Callable,
-    Dict,
     ForwardRef,
-    List,
     Literal,
     Optional,
     TypeVar,
@@ -16,7 +15,7 @@ from unittest.mock import patch
 
 import langsmith
 import pytest
-from typing_extensions import Annotated, NotRequired, Required, TypedDict
+from typing_extensions import NotRequired, Required, TypedDict
 
 from langgraph.graph import END, StateGraph
 from langgraph.graph.graph import CompiledGraph
@@ -150,9 +149,9 @@ def test_is_optional_type():
     assert not _is_optional_type(int)
     assert _is_optional_type(Optional[Literal[1, 2, 3]])
     assert not _is_optional_type(Literal[1, 2, 3])
-    assert _is_optional_type(Optional[List[int]])
-    assert _is_optional_type(Optional[Dict[str, int]])
-    assert not _is_optional_type(List[Optional[int]])
+    assert _is_optional_type(Optional[list[int]])
+    assert _is_optional_type(Optional[dict[str, int]])
+    assert not _is_optional_type(list[Optional[int]])
     assert _is_optional_type(Union[Optional[str], Optional[int]])
     assert _is_optional_type(
         Union[
@@ -177,8 +176,8 @@ def test_is_optional_type():
     assert _is_optional_type(Optional[ForwardRef("MyClass")])
     assert not _is_optional_type(ForwardRef("MyClass"))
 
-    assert _is_optional_type(Optional[Union[List[int], Dict[str, Optional[int]]]])
-    assert not _is_optional_type(Union[List[int], Dict[str, Optional[int]]])
+    assert _is_optional_type(Optional[Union[list[int], dict[str, Optional[int]]]])
+    assert not _is_optional_type(Union[list[int], dict[str, Optional[int]]])
 
     assert _is_optional_type(Optional[Callable[[int], str]])
     assert not _is_optional_type(Callable[[int], Optional[str]])