checkpoint-sqlite: release 2.0.1 (#2184)

feat(sdk): bump SDK to js@0.0.18 and py@0.1.34

.github/workflows/_test.yml

@@ -35,6 +35,7 @@ jobs:
           cache-key: test-${{ inputs.working-directory }}
       - name: Login to Docker Hub
         uses: docker/login-action@v3
+        if: ${{ !github.event.pull_request.head.repo.fork }}
         with:
           username: ${{ secrets.DOCKERHUB_USERNAME }}
           password: ${{ secrets.DOCKERHUB_RO_TOKEN }}

.github/workflows/_test_langgraph.yml

@@ -37,6 +37,7 @@ jobs:
           cache-key: test-langgraph
       - name: Login to Docker Hub
         uses: docker/login-action@v3
+        if: ${{ !github.event.pull_request.head.repo.fork }}
         with:
           username: ${{ secrets.DOCKERHUB_USERNAME }}
           password: ${{ secrets.DOCKERHUB_RO_TOKEN }}

.github/workflows/_test_scheduler_kafka.yml

@@ -29,6 +29,7 @@ jobs:
           cache-key: test-scheduler-kafka
       - name: Login to Docker Hub
         uses: docker/login-action@v3
+        if: ${{ !github.event.pull_request.head.repo.fork }}
         with:
           username: ${{ secrets.DOCKERHUB_USERNAME }}
           password: ${{ secrets.DOCKERHUB_RO_TOKEN }}

.github/workflows/ci.yml

@@ -31,6 +31,7 @@ jobs:
             "libs/cli",
             "libs/checkpoint",
             "libs/checkpoint-sqlite",
+            "libs/checkpoint-duckdb",
             "libs/checkpoint-postgres",
             "libs/scheduler-kafka",
           ]
@@ -47,6 +48,7 @@ jobs:
             "libs/cli",
             "libs/checkpoint",
             "libs/checkpoint-sqlite",
+            "libs/checkpoint-duckdb",
             "libs/checkpoint-postgres"
           ]
     uses: ./.github/workflows/_test.yml

.github/workflows/deploy_docs.yml

@@ -81,6 +81,7 @@ jobs:
               --check-links-ignore "https://x.com/.*" \
               --check-links-ignore "https://github\.com/.*" \
               --check-links-ignore "/.*\.(ipynb|html)$" \
+              --check-links-ignore "https://python\.langchain\.com/.*" \
               --check-links $(find docs/site -name "index.html" | grep -v 'storm/index.html')
           else
             echo "Fetching changes from origin/main..."

docs/docs/concepts/low_level.md

@@ -20,7 +20,7 @@ A super-step can be considered a single iteration over the graph nodes. Nodes th
 
 ### StateGraph
 
-The `StateGraph` class is the main graph class to uses. This is parameterized by a user defined `State` object.
+The `StateGraph` class is the main graph class to use. This is parameterized by a user defined `State` object.
 
 ### MessageGraph
 

docs/docs/concepts/memory.md

@@ -183,9 +183,9 @@ store = InMemoryStore()
 user_id = "my-user"
 application_context = "chitchat"
 namespace = (user_id, application_context)
-store.put(namespace, key="a-memory", {"rules": ["User likes short, direct language", "User only speaks English & python"], "my-key": "my-value"})
+store.put(namespace, "a-memory", {"rules": ["User likes short, direct language", "User only speaks English & python"], "my-key": "my-value"})
 # get the "memory" by ID
-item = store.get(namespace)
+item = store.get(namespace, "a-memory")
 # list "memories" within this namespace, filtering on content equivalence
 items = store.search(namespace, filter={"my-key": "my-value"})
 ```

docs/docs/concepts/multi_agent.md

@@ -23,7 +23,7 @@ There are several ways to connect agents in a multi-agent system:
 - **Network**: each agent can communicate with [every other agent](https://langchain-ai.github.io/langgraph/tutorials/multi_agent/multi-agent-collaboration/). Any agent can decide which other agent to call next.
 - **Supervisor**: each agent communicates with a single [supervisor](https://langchain-ai.github.io/langgraph/tutorials/multi_agent/agent_supervisor/) agent. Supervisor agent makes decisions on which agent should be called next.
 - **Supervisor (tool-calling)**: this is a special case of supervisor architecture. Individual agents can be represented as tools. In this case, a supervisor agent uses a tool-calling LLM to decide which of the agent tools to call, as well as the arguments to pass to those agents.
-- **Hierarchical**: you can define a multi-agent system with a supervisor of supervisors. This is a generalization of the supervisor architecture and allows for more complex control flows.
+- **Hierarchical**: you can define a multi-agent system with [a supervisor of supervisors](https://langchain-ai.github.io/langgraph/tutorials/multi_agent/hierarchical_agent_teams/). This is a generalization of the supervisor architecture and allows for more complex control flows.
 - **Custom multi-agent workflow**: each agent communicates with only a subset of agents. Parts of the flow are deterministic, and only some agents can decide which other agents to call next.
 
 ### Network
@@ -47,13 +47,20 @@ from langgraph.graph import StateGraph, MessagesState, START
 model = ChatOpenAI()
 
 class AgentState(MessagesState):
-    next: Literal["agent_1", "agent_2"]
+    next: Literal["agent_1", "agent_2", "__end__"]
 
 def supervisor(state: AgentState):
+    # you can pass relevant parts of the state to the LLM (e.g., state["messages"])
+    # to determine which agent to call next. a common pattern is to call the model
+    # with a structured output (e.g. force it to return an output with a "next_agent" field)
     response = model.invoke(...)
+    # the "next" key will be used by the conditional edges to route execution
+    # to the appropriate agent
     return {"next": response["next_agent"]}
 
 def agent_1(state: AgentState):
+    # you can pass relevant parts of the state to the LLM (e.g., state["messages"])
+    # and add any additional logic (different models, custom prompts, structured output, etc.)
     response = model.invoke(...)
     return {"messages": [response]}
 
@@ -68,6 +75,7 @@ builder.add_node(agent_2)
 
 builder.add_edge(START, "supervisor")
 # route to one of the agents or exit based on the supervisor's decisiion
+# if the supervisor returns "__end__", the graph will finish execution
 builder.add_conditional_edges("supervisor", lambda state: state["next"])
 builder.add_edge("agent_1", "supervisor")
 builder.add_edge("agent_2", "supervisor")
@@ -88,18 +96,117 @@ from langgraph.prebuilt import InjectedState, create_react_agent
 
 model = ChatOpenAI()
 
+# this is the agent function that will be called as tool
+# notice that you can pass the state to the tool via InjectedState annotation
 def agent_1(state: Annotated[dict, InjectedState]):
-    tool_message = ...
-    return {"messages": [tool_message]}
+    # you can pass relevant parts of the state to the LLM (e.g., state["messages"])
+    # and add any additional logic (different models, custom prompts, structured output, etc.)
+    response = model.invoke(...)
+    # return the LLM response as a string (expected tool response format)
+    # this will be automatically turned to ToolMessage
+    # by the prebuilt create_react_agent (supervisor)
+    return response.content
 
 def agent_2(state: Annotated[dict, InjectedState]):
-    tool_message = ...
-    return {"messages": [tool_message]}
+    response = model.invoke(...)
+    return response.content
 
 tools = [agent_1, agent_2]
+# the simplest way to build a supervisor w/ tool-calling is to use prebuilt ReAct agent graph
+# that consists of a tool-calling LLM node (i.e. supervisor) and a tool-executing node
 supervisor = create_react_agent(model, tools)
 ```
 
+### Hierarchical
+
+As you add more agents to your system, it might become too hard for the supervisor to manage all of them. The supervisor might start making poor decisions about which agent to call next, the context might become too complex for a single supervisor to keep track of. In other words, you end up with the same problems that motivated the multi-agent architecture in the first place.
+
+To address this, you can design your system _hierarchically_. For example, you can create separate, specialized teams of agents managed by individual supervisors, and a top-level supervisor to manage the teams.
+
+```python
+from typing import Literal
+from langchain_openai import ChatOpenAI
+from langgraph.graph import StateGraph, MessagesState, START
+
+model = ChatOpenAI()
+
+# define team 1 (same as the single supervisor example above)
+class Team1State(MessagesState):
+    next: Literal["team_1_agent_1", "team_1_agent_2", "__end__"]
+
+def team_1_supervisor(state: Team1State):
+    response = model.invoke(...)
+    return {"next": response["next_agent"]}
+
+def team_1_agent_1(state: Team1State):
+    response = model.invoke(...)
+    return {"messages": [response]}
+
+def team_1_agent_2(state: Team1State):
+    response = model.invoke(...)
+    return {"messages": [response]}
+
+team_1_builder = StateGraph(Team1State)
+team_1_builder.add_node(team_1_supervisor)
+team_1_builder.add_node(team_1_agent_1)
+team_1_builder.add_node(team_1_agent_2)
+team_1_builder.add_edge(START, "team_1_supervisor")
+# route to one of the agents or exit based on the supervisor's decisiion
+# if the supervisor returns "__end__", the graph will finish execution
+team_1_builder.add_conditional_edges("team_1_supervisor", lambda state: state["next"])
+team_1_builder.add_edge("team_1_agent_1", "team_1_supervisor")
+team_1_builder.add_edge("team_1_agent_2", "team_1_supervisor")
+
+team_1_graph = team_1_builder.compile()
+
+# define team 2 (same as the single supervisor example above)
+class Team2State(MessagesState):
+    next: Literal["team_2_agent_1", "team_2_agent_2", "__end__"]
+
+def team_2_supervisor(state: Team2State):
+    ...
+
+def team_2_agent_1(state: Team2State):
+    ...
+
+def team_2_agent_2(state: Team2State):
+    ...
+
+team_2_builder = StateGraph(Team2State)
+...
+team_2_graph = team_2_builder.compile()
+
+
+# define top-level supervisor
+
+class TopLevelState(MessagesState):
+    next: Literal["team_1", "team_2", "__end__"]
+
+builder = StateGraph(TopLevelState)
+def top_level_supervisor(state: TopLevelState):
+    # you can pass relevant parts of the state to the LLM (e.g., state["messages"])
+    # to determine which team to call next. a common pattern is to call the model
+    # with a structured output (e.g. force it to return an output with a "next_team" field)
+    response = model.invoke(...)
+    # the "next" key will be used by the conditional edges to route execution
+    # to the appropriate team
+    return {"next": response["next_team"]}
+
+builder = StateGraph(TopLevelState)
+builder.add_node(top_level_supervisor)
+builder.add_node(team_1_graph)
+builder.add_node(team_2_graph)
+
+builder.add_edge(START, "top_level_supervisor")
+# route to one of the teams or exit based on the supervisor's decision
+# if the top-level supervisor returns "__end__", the graph will finish execution
+builder.add_conditional_edges("top_level_supervisor", lambda state: state["next"])
+builder.add_edge("team_1_graph", "top_level_supervisor")
+builder.add_edge("team_2_graph", "top_level_supervisor")
+
+graph = builder.compile()
+```
+
 ### Custom multi-agent workflow
 
 In this architecture we add individual agents as graph nodes and define the order in which agents are called ahead of time, in a custom workflow. In LangGraph the workflow can be defined in two ways:

docs/docs/how-tos/branching.ipynb

@@ -89,7 +89,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 2,
    "id": "09372b8b-edea-4b9d-9ec3-3d93ce1ba819",
    "metadata": {},
    "outputs": [],
@@ -132,7 +132,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 3,
    "id": "66f52a20",
    "metadata": {},
    "outputs": [
@@ -163,7 +163,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 4,
    "id": "38846b01",
    "metadata": {},
    "outputs": [
@@ -183,7 +183,7 @@
        "{'aggregate': [\"I'm A\", \"I'm B\", \"I'm C\", \"I'm D\"]}"
       ]
      },
-     "execution_count": 5,
+     "execution_count": 4,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -220,7 +220,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 5,
    "id": "259a7704-5aa0-4e4c-aeef-cca04e8be0ff",
    "metadata": {},
    "outputs": [],
@@ -238,6 +238,15 @@
     "    aggregate: Annotated[list, operator.add]\n",
     "\n",
     "\n",
+    "class ReturnNodeValue:\n",
+    "    def __init__(self, node_secret: str):\n",
+    "        self._value = node_secret\n",
+    "\n",
+    "    def __call__(self, state: State) -> Any:\n",
+    "        print(f\"Adding {self._value} to {state['aggregate']}\")\n",
+    "        return {\"aggregate\": [self._value]}\n",
+    "\n",
+    "\n",
     "builder = StateGraph(State)\n",
     "builder.add_node(\"a\", ReturnNodeValue(\"I'm A\"))\n",
     "builder.add_edge(START, \"a\")\n",
@@ -255,7 +264,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 6,
    "id": "83320227-8ab3-44c0-b6cf-064a7a425b9f",
    "metadata": {},
    "outputs": [
@@ -278,7 +287,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 7,
    "id": "3f971fa3-29e4-466f-a85e-2863bfecf7fe",
    "metadata": {},
    "outputs": [
@@ -299,7 +308,7 @@
        "{'aggregate': [\"I'm A\", \"I'm B\", \"I'm C\", \"I'm B2\", \"I'm D\"]}"
       ]
      },
-     "execution_count": 8,
+     "execution_count": 7,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -322,7 +331,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 8,
    "id": "95f5e026",
    "metadata": {},
    "outputs": [],
@@ -341,6 +350,15 @@
     "    which: str\n",
     "\n",
     "\n",
+    "class ReturnNodeValue:\n",
+    "    def __init__(self, node_secret: str):\n",
+    "        self._value = node_secret\n",
+    "\n",
+    "    def __call__(self, state: State) -> Any:\n",
+    "        print(f\"Adding {self._value} to {state['aggregate']}\")\n",
+    "        return {\"aggregate\": [self._value]}\n",
+    "\n",
+    "\n",
     "builder = StateGraph(State)\n",
     "builder.add_node(\"a\", ReturnNodeValue(\"I'm A\"))\n",
     "builder.add_edge(START, \"a\")\n",
@@ -372,7 +390,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 9,
    "id": "1d0e6c56",
    "metadata": {},
    "outputs": [
@@ -395,7 +413,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 10,
    "id": "7134f652",
    "metadata": {},
    "outputs": [
@@ -415,7 +433,7 @@
        "{'aggregate': [\"I'm A\", \"I'm B\", \"I'm C\", \"I'm E\"], 'which': 'bc'}"
       ]
      },
-     "execution_count": 11,
+     "execution_count": 10,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -426,7 +444,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 11,
    "id": "b130e694",
    "metadata": {},
    "outputs": [
@@ -446,7 +464,7 @@
        "{'aggregate': [\"I'm A\", \"I'm C\", \"I'm D\", \"I'm E\"], 'which': 'cd'}"
       ]
      },
-     "execution_count": 12,
+     "execution_count": 11,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -471,7 +489,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 12,
    "id": "836bc12d",
    "metadata": {},
    "outputs": [],
@@ -564,7 +582,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 13,
    "id": "932c497e",
    "metadata": {},
    "outputs": [
@@ -587,7 +605,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": 14,
    "id": "933b3afd",
    "metadata": {},
    "outputs": [
@@ -608,7 +626,7 @@
        " 'which': 'bc'}"
       ]
      },
-     "execution_count": 15,
+     "execution_count": 14,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -619,7 +637,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": 15,
    "id": "e30531bf",
    "metadata": {},
    "outputs": [
@@ -640,7 +658,7 @@
        " 'which': 'cd'}"
       ]
      },
-     "execution_count": 16,
+     "execution_count": 15,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -666,7 +684,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.4"
+   "version": "3.12.3"
   }
  },
  "nbformat": 4,

libs/checkpoint-duckdb/Makefile

@@ -0,0 +1,35 @@
+.PHONY: test test_watch lint format
+
+######################
+# TESTING AND COVERAGE
+######################
+
+test:
+	poetry run pytest tests
+
+test_watch:
+	poetry run ptw .
+
+######################
+# LINTING AND FORMATTING
+######################
+
+# Define a variable for Python and notebook files.
+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_package: PYTHON_FILES=langgraph
+lint_tests: PYTHON_FILES=tests
+lint_tests: MYPY_CACHE=.mypy_cache_test
+
+lint lint_diff lint_package lint_tests:
+	poetry run ruff check .
+	[ "$(PYTHON_FILES)" = "" ] || poetry run ruff format $(PYTHON_FILES) --diff
+	[ "$(PYTHON_FILES)" = "" ] || poetry run ruff check --select I $(PYTHON_FILES)
+	[ "$(PYTHON_FILES)" = "" ] || mkdir -p $(MYPY_CACHE)
+	[ "$(PYTHON_FILES)" = "" ] || poetry run mypy $(PYTHON_FILES) --cache-dir $(MYPY_CACHE)
+
+format format_diff:
+	poetry run ruff format $(PYTHON_FILES)
+	poetry run ruff check --select I --fix $(PYTHON_FILES)

libs/checkpoint-duckdb/README.md

@@ -0,0 +1,95 @@
+# LangGraph Checkpoint DuckDB
+
+Implementation of LangGraph CheckpointSaver that uses DuckDB.
+
+## Usage
+
+> [!IMPORTANT]
+> When using DuckDB checkpointers for the first time, make sure to call `.setup()` method on them to create required tables. See example below.
+
+```python
+from langgraph.checkpoint.duckdb import DuckDBSaver
+
+write_config = {"configurable": {"thread_id": "1", "checkpoint_ns": ""}}
+read_config = {"configurable": {"thread_id": "1"}}
+
+with DuckDBSaver.from_conn_string(":memory:") as checkpointer:
+    # call .setup() the first time you're using the checkpointer
+    checkpointer.setup()
+    checkpoint = {
+        "v": 1,
+        "ts": "2024-07-31T20:14:19.804150+00:00",
+        "id": "1ef4f797-8335-6428-8001-8a1503f9b875",
+        "channel_values": {
+            "my_key": "meow",
+            "node": "node"
+        },
+        "channel_versions": {
+            "__start__": 2,
+            "my_key": 3,
+            "start:node": 3,
+            "node": 3
+        },
+        "versions_seen": {
+            "__input__": {},
+            "__start__": {
+            "__start__": 1
+            },
+            "node": {
+            "start:node": 2
+            }
+        },
+        "pending_sends": [],
+    }
+
+    # store checkpoint
+    checkpointer.put(write_config, checkpoint, {}, {})
+
+    # load checkpoint
+    checkpointer.get(read_config)
+
+    # list checkpoints
+    list(checkpointer.list(read_config))
+```
+
+### Async
+
+```python
+from langgraph.checkpoint.duckdb.aio import AsyncDuckDBSaver
+
+async with AsyncDuckDBSaver.from_conn_string(":memory:") as checkpointer:
+    checkpoint = {
+        "v": 1,
+        "ts": "2024-07-31T20:14:19.804150+00:00",
+        "id": "1ef4f797-8335-6428-8001-8a1503f9b875",
+        "channel_values": {
+            "my_key": "meow",
+            "node": "node"
+        },
+        "channel_versions": {
+            "__start__": 2,
+            "my_key": 3,
+            "start:node": 3,
+            "node": 3
+        },
+        "versions_seen": {
+            "__input__": {},
+            "__start__": {
+            "__start__": 1
+            },
+            "node": {
+            "start:node": 2
+            }
+        },
+        "pending_sends": [],
+    }
+
+    # store checkpoint
+    await checkpointer.aput(write_config, checkpoint, {}, {})
+
+    # load checkpoint
+    await checkpointer.aget(read_config)
+
+    # list checkpoints
+    [c async for c in checkpointer.alist(read_config)]
+```

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

@@ -0,0 +1,356 @@
+import threading
+from contextlib import contextmanager
+from typing import Any, Iterator, Optional, Sequence
+
+from langchain_core.runnables import RunnableConfig
+
+import duckdb
+from langgraph.checkpoint.base import (
+    WRITES_IDX_MAP,
+    ChannelVersions,
+    Checkpoint,
+    CheckpointMetadata,
+    CheckpointTuple,
+    get_checkpoint_id,
+)
+from langgraph.checkpoint.duckdb.base import BaseDuckDBSaver
+from langgraph.checkpoint.serde.base import SerializerProtocol
+
+
+class DuckDBSaver(BaseDuckDBSaver):
+    lock: threading.Lock
+
+    def __init__(
+        self,
+        conn: duckdb.DuckDBPyConnection,
+        serde: Optional[SerializerProtocol] = None,
+    ) -> None:
+        super().__init__(serde=serde)
+
+        self.conn = conn
+        self.lock = threading.Lock()
+
+    @classmethod
+    @contextmanager
+    def from_conn_string(cls, conn_string: str) -> Iterator["DuckDBSaver"]:
+        """Create a new DuckDBSaver instance from a connection string.
+
+        Args:
+            conn_string (str): The DuckDB connection info string.
+
+        Returns:
+            DuckDBSaver: A new DuckDBSaver instance.
+        """
+        with duckdb.connect(conn_string) as conn:
+            yield DuckDBSaver(conn)
+
+    def setup(self) -> None:
+        """Set up the checkpoint database asynchronously.
+
+        This method creates the necessary tables in the DuckDB database if they don't
+        already exist and runs database migrations. It MUST be called directly by the user
+        the first time checkpointer is used.
+        """
+        with self.lock, self.conn.cursor() as cur:
+            try:
+                row = cur.execute(
+                    "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
+                ).fetchone()
+                if row is None:
+                    version = -1
+                else:
+                    version = row[0]
+            except duckdb.CatalogException:
+                version = -1
+            for v, migration in zip(
+                range(version + 1, len(self.MIGRATIONS)),
+                self.MIGRATIONS[version + 1 :],
+            ):
+                cur.execute(migration)
+                cur.execute("INSERT INTO checkpoint_migrations (v) VALUES (?)", [v])
+
+    def list(
+        self,
+        config: Optional[RunnableConfig],
+        *,
+        filter: Optional[dict[str, Any]] = None,
+        before: Optional[RunnableConfig] = None,
+        limit: Optional[int] = None,
+    ) -> Iterator[CheckpointTuple]:
+        """List checkpoints from the database.
+
+        This method retrieves a list of checkpoint tuples from the DuckDB database based
+        on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first).
+
+        Args:
+            config (RunnableConfig): The config to use for listing the checkpoints.
+            filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata. Defaults to None.
+            before (Optional[RunnableConfig]): If provided, only checkpoints before the specified checkpoint ID are returned. Defaults to None.
+            limit (Optional[int]): The maximum number of checkpoints to return. Defaults to None.
+
+        Yields:
+            Iterator[CheckpointTuple]: An iterator of checkpoint tuples.
+
+        Examples:
+            >>> from langgraph.checkpoint.duckdb import DuckDBSaver
+            >>> with DuckDBSaver.from_conn_string(":memory:") as memory:
+            ... # Run a graph, then list the checkpoints
+            >>>     config = {"configurable": {"thread_id": "1"}}
+            >>>     checkpoints = list(memory.list(config, limit=2))
+            >>> print(checkpoints)
+            [CheckpointTuple(...), CheckpointTuple(...)]
+
+            >>> config = {"configurable": {"thread_id": "1"}}
+            >>> before = {"configurable": {"checkpoint_id": "1ef4f797-8335-6428-8001-8a1503f9b875"}}
+            >>> with DuckDBSaver.from_conn_string(":memory:") as memory:
+            ... # Run a graph, then list the checkpoints
+            >>>     checkpoints = list(memory.list(config, before=before))
+            >>> print(checkpoints)
+            [CheckpointTuple(...), ...]
+        """
+        where, args = self._search_where(config, filter, before)
+        query = self.SELECT_SQL + where + " ORDER BY checkpoint_id DESC"
+        if limit:
+            query += f" LIMIT {limit}"
+        # if we change this to use .stream() we need to make sure to close the cursor
+        with self._cursor() as cur:
+            cur.execute(query, args)
+            for value in cur.fetchall():
+                (
+                    thread_id,
+                    checkpoint,
+                    checkpoint_ns,
+                    checkpoint_id,
+                    parent_checkpoint_id,
+                    metadata,
+                    channel_values,
+                    pending_writes,
+                    pending_sends,
+                ) = value
+                yield CheckpointTuple(
+                    {
+                        "configurable": {
+                            "thread_id": thread_id,
+                            "checkpoint_ns": checkpoint_ns,
+                            "checkpoint_id": checkpoint_id,
+                        }
+                    },
+                    self._load_checkpoint(
+                        checkpoint,
+                        channel_values,
+                        pending_sends,
+                    ),
+                    self._load_metadata(metadata),
+                    (
+                        {
+                            "configurable": {
+                                "thread_id": thread_id,
+                                "checkpoint_ns": checkpoint_ns,
+                                "checkpoint_id": parent_checkpoint_id,
+                            }
+                        }
+                        if parent_checkpoint_id
+                        else None
+                    ),
+                    self._load_writes(pending_writes),
+                )
+
+    def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
+        """Get a checkpoint tuple from the database.
+
+        This method retrieves a checkpoint tuple from the DuckDB database based on the
+        provided config. If the config contains a "checkpoint_id" key, the checkpoint with
+        the matching thread ID and timestamp is retrieved. Otherwise, the latest checkpoint
+        for the given thread ID is retrieved.
+
+        Args:
+            config (RunnableConfig): The config to use for retrieving the checkpoint.
+
+        Returns:
+            Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found.
+
+        Examples:
+
+            Basic:
+            >>> config = {"configurable": {"thread_id": "1"}}
+            >>> checkpoint_tuple = memory.get_tuple(config)
+            >>> print(checkpoint_tuple)
+            CheckpointTuple(...)
+
+            With timestamp:
+
+            >>> config = {
+            ...    "configurable": {
+            ...        "thread_id": "1",
+            ...        "checkpoint_ns": "",
+            ...        "checkpoint_id": "1ef4f797-8335-6428-8001-8a1503f9b875",
+            ...    }
+            ... }
+            >>> checkpoint_tuple = memory.get_tuple(config)
+            >>> print(checkpoint_tuple)
+            CheckpointTuple(...)
+        """  # noqa
+        thread_id = config["configurable"]["thread_id"]
+        checkpoint_id = get_checkpoint_id(config)
+        checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
+        if checkpoint_id:
+            args: tuple[Any, ...] = (thread_id, checkpoint_ns, checkpoint_id)
+            where = "WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?"
+        else:
+            args = (thread_id, checkpoint_ns)
+            where = "WHERE thread_id = ? AND checkpoint_ns = ? ORDER BY checkpoint_id DESC LIMIT 1"
+
+        with self._cursor() as cur:
+            cur.execute(
+                self.SELECT_SQL + where,
+                args,
+            )
+
+            value = cur.fetchone()
+            if value:
+                (
+                    thread_id,
+                    checkpoint,
+                    checkpoint_ns,
+                    checkpoint_id,
+                    parent_checkpoint_id,
+                    metadata,
+                    channel_values,
+                    pending_writes,
+                    pending_sends,
+                ) = value
+                return CheckpointTuple(
+                    {
+                        "configurable": {
+                            "thread_id": thread_id,
+                            "checkpoint_ns": checkpoint_ns,
+                            "checkpoint_id": checkpoint_id,
+                        }
+                    },
+                    self._load_checkpoint(
+                        checkpoint,
+                        channel_values,
+                        pending_sends,
+                    ),
+                    self._load_metadata(metadata),
+                    (
+                        {
+                            "configurable": {
+                                "thread_id": thread_id,
+                                "checkpoint_ns": checkpoint_ns,
+                                "checkpoint_id": parent_checkpoint_id,
+                            }
+                        }
+                        if parent_checkpoint_id
+                        else None
+                    ),
+                    self._load_writes(pending_writes),
+                )
+
+    def put(
+        self,
+        config: RunnableConfig,
+        checkpoint: Checkpoint,
+        metadata: CheckpointMetadata,
+        new_versions: ChannelVersions,
+    ) -> RunnableConfig:
+        """Save a checkpoint to the database.
+
+        This method saves a checkpoint to the DuckDB database. The checkpoint is associated
+        with the provided config and its parent config (if any).
+
+        Args:
+            config (RunnableConfig): The config to associate with the checkpoint.
+            checkpoint (Checkpoint): The checkpoint to save.
+            metadata (CheckpointMetadata): Additional metadata to save with the checkpoint.
+            new_versions (ChannelVersions): New channel versions as of this write.
+
+        Returns:
+            RunnableConfig: Updated configuration after storing the checkpoint.
+
+        Examples:
+
+            >>> from langgraph.checkpoint.duckdb import DuckDBSaver
+            >>> with DuckDBSaver.from_conn_string(":memory:") as memory:
+            >>>     config = {"configurable": {"thread_id": "1", "checkpoint_ns": ""}}
+            >>>     checkpoint = {"ts": "2024-05-04T06:32:42.235444+00:00", "id": "1ef4f797-8335-6428-8001-8a1503f9b875", "channel_values": {"key": "value"}}
+            >>>     saved_config = memory.put(config, checkpoint, {"source": "input", "step": 1, "writes": {"key": "value"}}, {})
+            >>> print(saved_config)
+            {'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef4f797-8335-6428-8001-8a1503f9b875'}}
+        """
+        configurable = config["configurable"].copy()
+        thread_id = configurable.pop("thread_id")
+        checkpoint_ns = configurable.pop("checkpoint_ns")
+        checkpoint_id = configurable.pop(
+            "checkpoint_id", configurable.pop("thread_ts", None)
+        )
+
+        copy = checkpoint.copy()
+        next_config = {
+            "configurable": {
+                "thread_id": thread_id,
+                "checkpoint_ns": checkpoint_ns,
+                "checkpoint_id": checkpoint["id"],
+            }
+        }
+        checkpoint_blobs = self._dump_blobs(
+            thread_id,
+            checkpoint_ns,
+            copy.pop("channel_values"),  # type: ignore[misc]
+            new_versions,
+        )
+        with self._cursor() as cur:
+            if checkpoint_blobs:
+                cur.executemany(self.UPSERT_CHECKPOINT_BLOBS_SQL, checkpoint_blobs)
+            cur.execute(
+                self.UPSERT_CHECKPOINTS_SQL,
+                (
+                    thread_id,
+                    checkpoint_ns,
+                    checkpoint["id"],
+                    checkpoint_id,
+                    self._dump_checkpoint(copy),
+                    self._dump_metadata(metadata),
+                ),
+            )
+        return next_config
+
+    def put_writes(
+        self,
+        config: RunnableConfig,
+        writes: Sequence[tuple[str, Any]],
+        task_id: str,
+    ) -> None:
+        """Store intermediate writes linked to a checkpoint.
+
+        This method saves intermediate writes associated with a checkpoint to the DuckDB database.
+
+        Args:
+            config (RunnableConfig): Configuration of the related checkpoint.
+            writes (List[Tuple[str, Any]]): List of writes to store.
+            task_id (str): Identifier for the task creating the writes.
+        """
+        query = (
+            self.UPSERT_CHECKPOINT_WRITES_SQL
+            if all(w[0] in WRITES_IDX_MAP for w in writes)
+            else self.INSERT_CHECKPOINT_WRITES_SQL
+        )
+        with self._cursor() as cur:
+            cur.executemany(
+                query,
+                self._dump_writes(
+                    config["configurable"]["thread_id"],
+                    config["configurable"]["checkpoint_ns"],
+                    config["configurable"]["checkpoint_id"],
+                    task_id,
+                    writes,
+                ),
+            )
+
+    @contextmanager
+    def _cursor(self) -> Iterator[duckdb.DuckDBPyConnection]:
+        with self.lock, self.conn.cursor() as cur:
+            yield cur
+
+
+__all__ = ["DuckDBSaver", "Conn"]

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

@@ -0,0 +1,431 @@
+import asyncio
+from contextlib import asynccontextmanager
+from typing import Any, AsyncIterator, Iterator, Optional, Sequence
+
+from langchain_core.runnables import RunnableConfig
+
+import duckdb
+from langgraph.checkpoint.base import (
+    WRITES_IDX_MAP,
+    ChannelVersions,
+    Checkpoint,
+    CheckpointMetadata,
+    CheckpointTuple,
+    get_checkpoint_id,
+)
+from langgraph.checkpoint.duckdb.base import BaseDuckDBSaver
+from langgraph.checkpoint.serde.base import SerializerProtocol
+
+
+class AsyncDuckDBSaver(BaseDuckDBSaver):
+    lock: asyncio.Lock
+
+    def __init__(
+        self,
+        conn: duckdb.DuckDBPyConnection,
+        serde: Optional[SerializerProtocol] = None,
+    ) -> None:
+        super().__init__(serde=serde)
+        self.conn = conn
+        self.lock = asyncio.Lock()
+        self.loop = asyncio.get_running_loop()
+
+    @classmethod
+    @asynccontextmanager
+    async def from_conn_string(
+        cls,
+        conn_string: str,
+    ) -> AsyncIterator["AsyncDuckDBSaver"]:
+        """Create a new AsyncDuckDBSaver instance from a connection string.
+
+        Args:
+            conn_string (str): The DuckDB connection info string.
+
+        Returns:
+            AsyncDuckDBSaver: A new AsyncDuckDBSaver instance.
+        """
+        with duckdb.connect(conn_string) as conn:
+            yield AsyncDuckDBSaver(conn)
+
+    async def setup(self) -> None:
+        """Set up the checkpoint database asynchronously.
+
+        This method creates the necessary tables in the DuckDB database if they don't
+        already exist and runs database migrations. It MUST be called directly by the user
+        the first time checkpointer is used.
+        """
+        async with self.lock:
+            with self.conn.cursor() as cur:
+                try:
+                    await asyncio.to_thread(
+                        cur.execute,
+                        "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1",
+                    )
+                    row = await asyncio.to_thread(cur.fetchone)
+                    if row is None:
+                        version = -1
+                    else:
+                        version = row[0]
+                except duckdb.CatalogException:
+                    version = -1
+                for v, migration in zip(
+                    range(version + 1, len(self.MIGRATIONS)),
+                    self.MIGRATIONS[version + 1 :],
+                ):
+                    await asyncio.to_thread(cur.execute, migration)
+                    await asyncio.to_thread(
+                        cur.execute,
+                        "INSERT INTO checkpoint_migrations (v) VALUES (?)",
+                        [v],
+                    )
+
+    async def alist(
+        self,
+        config: Optional[RunnableConfig],
+        *,
+        filter: Optional[dict[str, Any]] = None,
+        before: Optional[RunnableConfig] = None,
+        limit: Optional[int] = None,
+    ) -> AsyncIterator[CheckpointTuple]:
+        """List checkpoints from the database asynchronously.
+
+        This method retrieves a list of checkpoint tuples from the DuckDB database based
+        on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first).
+
+        Args:
+            config (Optional[RunnableConfig]): Base configuration for filtering checkpoints.
+            filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata.
+            before (Optional[RunnableConfig]): If provided, only checkpoints before the specified checkpoint ID are returned. Defaults to None.
+            limit (Optional[int]): Maximum number of checkpoints to return.
+
+        Yields:
+            AsyncIterator[CheckpointTuple]: An asynchronous iterator of matching checkpoint tuples.
+        """
+        where, args = self._search_where(config, filter, before)
+        query = self.SELECT_SQL + where + " ORDER BY checkpoint_id DESC"
+        if limit:
+            query += f" LIMIT {limit}"
+        # if we change this to use .stream() we need to make sure to close the cursor
+        async with self._cursor() as cur:
+            await asyncio.to_thread(cur.execute, query, args)
+            results = await asyncio.to_thread(cur.fetchall)
+            for value in results:
+                (
+                    thread_id,
+                    checkpoint,
+                    checkpoint_ns,
+                    checkpoint_id,
+                    parent_checkpoint_id,
+                    metadata,
+                    channel_values,
+                    pending_writes,
+                    pending_sends,
+                ) = value
+                yield CheckpointTuple(
+                    {
+                        "configurable": {
+                            "thread_id": thread_id,
+                            "checkpoint_ns": checkpoint_ns,
+                            "checkpoint_id": checkpoint_id,
+                        }
+                    },
+                    await asyncio.to_thread(
+                        self._load_checkpoint,
+                        checkpoint,
+                        channel_values,
+                        pending_sends,
+                    ),
+                    self._load_metadata(metadata),
+                    (
+                        {
+                            "configurable": {
+                                "thread_id": thread_id,
+                                "checkpoint_ns": checkpoint_ns,
+                                "checkpoint_id": parent_checkpoint_id,
+                            }
+                        }
+                        if parent_checkpoint_id
+                        else None
+                    ),
+                    await asyncio.to_thread(self._load_writes, pending_writes),
+                )
+
+    async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
+        """Get a checkpoint tuple from the database asynchronously.
+
+        This method retrieves a checkpoint tuple from the DuckDBdatabase based on the
+        provided config. If the config contains a "checkpoint_id" key, the checkpoint with
+        the matching thread ID and "checkpoint_id" is retrieved. Otherwise, the latest checkpoint
+        for the given thread ID is retrieved.
+
+        Args:
+            config (RunnableConfig): The config to use for retrieving the checkpoint.
+
+        Returns:
+            Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found.
+        """
+        thread_id = config["configurable"]["thread_id"]
+        checkpoint_id = get_checkpoint_id(config)
+        checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
+        if checkpoint_id:
+            args: tuple[Any, ...] = (thread_id, checkpoint_ns, checkpoint_id)
+            where = "WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?"
+        else:
+            args = (thread_id, checkpoint_ns)
+            where = "WHERE thread_id = ? AND checkpoint_ns = ? ORDER BY checkpoint_id DESC LIMIT 1"
+
+        async with self._cursor() as cur:
+            await asyncio.to_thread(
+                cur.execute,
+                self.SELECT_SQL + where,
+                args,
+            )
+
+            value = await asyncio.to_thread(cur.fetchone)
+            if value:
+                (
+                    thread_id,
+                    checkpoint,
+                    checkpoint_ns,
+                    checkpoint_id,
+                    parent_checkpoint_id,
+                    metadata,
+                    channel_values,
+                    pending_writes,
+                    pending_sends,
+                ) = value
+                return CheckpointTuple(
+                    {
+                        "configurable": {
+                            "thread_id": thread_id,
+                            "checkpoint_ns": checkpoint_ns,
+                            "checkpoint_id": checkpoint_id,
+                        }
+                    },
+                    await asyncio.to_thread(
+                        self._load_checkpoint,
+                        checkpoint,
+                        channel_values,
+                        pending_sends,
+                    ),
+                    self._load_metadata(metadata),
+                    (
+                        {
+                            "configurable": {
+                                "thread_id": thread_id,
+                                "checkpoint_ns": checkpoint_ns,
+                                "checkpoint_id": parent_checkpoint_id,
+                            }
+                        }
+                        if parent_checkpoint_id
+                        else None
+                    ),
+                    await asyncio.to_thread(self._load_writes, pending_writes),
+                )
+
+    async def aput(
+        self,
+        config: RunnableConfig,
+        checkpoint: Checkpoint,
+        metadata: CheckpointMetadata,
+        new_versions: ChannelVersions,
+    ) -> RunnableConfig:
+        """Save a checkpoint to the database asynchronously.
+
+        This method saves a checkpoint to the DuckDB database. The checkpoint is associated
+        with the provided config and its parent config (if any).
+
+        Args:
+            config (RunnableConfig): The config to associate with the checkpoint.
+            checkpoint (Checkpoint): The checkpoint to save.
+            metadata (CheckpointMetadata): Additional metadata to save with the checkpoint.
+            new_versions (ChannelVersions): New channel versions as of this write.
+
+        Returns:
+            RunnableConfig: Updated configuration after storing the checkpoint.
+        """
+        configurable = config["configurable"].copy()
+        thread_id = configurable.pop("thread_id")
+        checkpoint_ns = configurable.pop("checkpoint_ns")
+        checkpoint_id = configurable.pop(
+            "checkpoint_id", configurable.pop("thread_ts", None)
+        )
+
+        copy = checkpoint.copy()
+        next_config = {
+            "configurable": {
+                "thread_id": thread_id,
+                "checkpoint_ns": checkpoint_ns,
+                "checkpoint_id": checkpoint["id"],
+            }
+        }
+
+        checkpoint_blobs = await asyncio.to_thread(
+            self._dump_blobs,
+            thread_id,
+            checkpoint_ns,
+            copy.pop("channel_values"),  # type: ignore[misc]
+            new_versions,
+        )
+        async with self._cursor() as cur:
+            if checkpoint_blobs:
+                await asyncio.to_thread(
+                    cur.executemany, self.UPSERT_CHECKPOINT_BLOBS_SQL, checkpoint_blobs
+                )
+            await asyncio.to_thread(
+                cur.execute,
+                self.UPSERT_CHECKPOINTS_SQL,
+                (
+                    thread_id,
+                    checkpoint_ns,
+                    checkpoint["id"],
+                    checkpoint_id,
+                    self._dump_checkpoint(copy),
+                    self._dump_metadata(metadata),
+                ),
+            )
+
+        return next_config
+
+    async def aput_writes(
+        self,
+        config: RunnableConfig,
+        writes: Sequence[tuple[str, Any]],
+        task_id: str,
+    ) -> None:
+        """Store intermediate writes linked to a checkpoint asynchronously.
+
+        This method saves intermediate writes associated with a checkpoint to the database.
+
+        Args:
+            config (RunnableConfig): Configuration of the related checkpoint.
+            writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair.
+            task_id (str): Identifier for the task creating the writes.
+        """
+        query = (
+            self.UPSERT_CHECKPOINT_WRITES_SQL
+            if all(w[0] in WRITES_IDX_MAP for w in writes)
+            else self.INSERT_CHECKPOINT_WRITES_SQL
+        )
+        params = await asyncio.to_thread(
+            self._dump_writes,
+            config["configurable"]["thread_id"],
+            config["configurable"]["checkpoint_ns"],
+            config["configurable"]["checkpoint_id"],
+            task_id,
+            writes,
+        )
+        async with self._cursor() as cur:
+            await asyncio.to_thread(cur.executemany, query, params)
+
+    @asynccontextmanager
+    async def _cursor(self) -> AsyncIterator[duckdb.DuckDBPyConnection]:
+        async with self.lock:
+            with self.conn.cursor() as cur:
+                yield cur
+
+    def list(
+        self,
+        config: Optional[RunnableConfig],
+        *,
+        filter: Optional[dict[str, Any]] = None,
+        before: Optional[RunnableConfig] = None,
+        limit: Optional[int] = None,
+    ) -> Iterator[CheckpointTuple]:
+        """List checkpoints from the database.
+
+        This method retrieves a list of checkpoint tuples from the DuckDB database based
+        on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first).
+
+        Args:
+            config (Optional[RunnableConfig]): Base configuration for filtering checkpoints.
+            filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata.
+            before (Optional[RunnableConfig]): If provided, only checkpoints before the specified checkpoint ID are returned. Defaults to None.
+            limit (Optional[int]): Maximum number of checkpoints to return.
+
+        Yields:
+            Iterator[CheckpointTuple]: An iterator of matching checkpoint tuples.
+        """
+        aiter_ = self.alist(config, filter=filter, before=before, limit=limit)
+        while True:
+            try:
+                yield asyncio.run_coroutine_threadsafe(
+                    anext(aiter_),
+                    self.loop,
+                ).result()
+            except StopAsyncIteration:
+                break
+
+    def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
+        """Get a checkpoint tuple from the database.
+
+        This method retrieves a checkpoint tuple from the DuckDB database based on the
+        provided config. If the config contains a "checkpoint_id" key, the checkpoint with
+        the matching thread ID and "checkpoint_id" is retrieved. Otherwise, the latest checkpoint
+        for the given thread ID is retrieved.
+
+        Args:
+            config (RunnableConfig): The config to use for retrieving the checkpoint.
+
+        Returns:
+            Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found.
+        """
+        try:
+            # check if we are in the main thread, only bg threads can block
+            # we don't check in other methods to avoid the overhead
+            if asyncio.get_running_loop() is self.loop:
+                raise asyncio.InvalidStateError(
+                    "Synchronous calls to AsyncDuckDBSaver are only allowed from a "
+                    "different thread. From the main thread, use the async interface."
+                    "For example, use `await checkpointer.aget_tuple(...)` or `await "
+                    "graph.ainvoke(...)`."
+                )
+        except RuntimeError:
+            pass
+        return asyncio.run_coroutine_threadsafe(
+            self.aget_tuple(config), self.loop
+        ).result()
+
+    def put(
+        self,
+        config: RunnableConfig,
+        checkpoint: Checkpoint,
+        metadata: CheckpointMetadata,
+        new_versions: ChannelVersions,
+    ) -> RunnableConfig:
+        """Save a checkpoint to the database.
+
+        This method saves a checkpoint to the DuckDB database. The checkpoint is associated
+        with the provided config and its parent config (if any).
+
+        Args:
+            config (RunnableConfig): The config to associate with the checkpoint.
+            checkpoint (Checkpoint): The checkpoint to save.
+            metadata (CheckpointMetadata): Additional metadata to save with the checkpoint.
+            new_versions (ChannelVersions): New channel versions as of this write.
+
+        Returns:
+            RunnableConfig: Updated configuration after storing the checkpoint.
+        """
+        return asyncio.run_coroutine_threadsafe(
+            self.aput(config, checkpoint, metadata, new_versions), self.loop
+        ).result()
+
+    def put_writes(
+        self,
+        config: RunnableConfig,
+        writes: Sequence[tuple[str, Any]],
+        task_id: str,
+    ) -> None:
+        """Store intermediate writes linked to a checkpoint.
+
+        This method saves intermediate writes associated with a checkpoint to the database.
+
+        Args:
+            config (RunnableConfig): Configuration of the related checkpoint.
+            writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair.
+            task_id (str): Identifier for the task creating the writes.
+        """
+        return asyncio.run_coroutine_threadsafe(
+            self.aput_writes(config, writes, task_id), self.loop
+        ).result()

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

@@ -0,0 +1,290 @@
+import json
+import random
+from typing import Any, List, Optional, Sequence, Tuple, cast
+
+from langchain_core.runnables import RunnableConfig
+
+from langgraph.checkpoint.base import (
+    WRITES_IDX_MAP,
+    BaseCheckpointSaver,
+    ChannelVersions,
+    Checkpoint,
+    CheckpointMetadata,
+    get_checkpoint_id,
+)
+from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
+from langgraph.checkpoint.serde.types import TASKS, ChannelProtocol
+
+MetadataInput = Optional[dict[str, Any]]
+
+"""
+To add a new migration, add a new string to the MIGRATIONS list.
+The position of the migration in the list is the version number.
+"""
+MIGRATIONS = [
+    """CREATE TABLE IF NOT EXISTS checkpoint_migrations (
+    v INTEGER PRIMARY KEY
+);""",
+    """CREATE TABLE IF NOT EXISTS checkpoints (
+    thread_id TEXT NOT NULL,
+    checkpoint_ns TEXT NOT NULL DEFAULT '',
+    checkpoint_id TEXT NOT NULL,
+    parent_checkpoint_id TEXT,
+    type TEXT,
+    checkpoint JSON NOT NULL,
+    metadata JSON NOT NULL DEFAULT '{}',
+    PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id)
+);""",
+    """CREATE TABLE IF NOT EXISTS checkpoint_blobs (
+    thread_id TEXT NOT NULL,
+    checkpoint_ns TEXT NOT NULL DEFAULT '',
+    channel TEXT NOT NULL,
+    version TEXT NOT NULL,
+    type TEXT NOT NULL,
+    blob BLOB,
+    PRIMARY KEY (thread_id, checkpoint_ns, channel, version)
+);""",
+    """CREATE TABLE IF NOT EXISTS checkpoint_writes (
+    thread_id TEXT NOT NULL,
+    checkpoint_ns TEXT NOT NULL DEFAULT '',
+    checkpoint_id TEXT NOT NULL,
+    task_id TEXT NOT NULL,
+    idx INTEGER NOT NULL,
+    channel TEXT NOT NULL,
+    type TEXT,
+    blob BLOB NOT NULL,
+    PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx)
+);""",
+]
+
+SELECT_SQL = f"""
+select
+    thread_id,
+    checkpoint,
+    checkpoint_ns,
+    checkpoint_id,
+    parent_checkpoint_id,
+    metadata,
+    (
+        select array_agg(array[bl.channel::bytea, bl.type::bytea, bl.blob])
+        from (
+            SELECT unnest(json_keys(json_extract(checkpoint, '$.channel_versions'))) as key
+        ) cv
+        inner join checkpoint_blobs bl
+            on bl.thread_id = checkpoints.thread_id
+            and bl.checkpoint_ns = checkpoints.checkpoint_ns
+            and bl.channel = cv.key
+            and bl.version = json_extract_string(checkpoint, '$.channel_versions.' || cv.key)
+    ) as channel_values,
+    (
+        select
+        array_agg(array[cw.task_id::blob, cw.channel::blob, cw.type::blob, cw.blob])
+        from checkpoint_writes cw
+        where cw.thread_id = checkpoints.thread_id
+            and cw.checkpoint_ns = checkpoints.checkpoint_ns
+            and cw.checkpoint_id = checkpoints.checkpoint_id
+    ) as pending_writes,
+    (
+        select array_agg(array[cw.type::blob, cw.blob])
+        from checkpoint_writes cw
+        where cw.thread_id = checkpoints.thread_id
+            and cw.checkpoint_ns = checkpoints.checkpoint_ns
+            and cw.checkpoint_id = checkpoints.parent_checkpoint_id
+            and cw.channel = '{TASKS}'
+    ) as pending_sends
+from checkpoints """
+
+UPSERT_CHECKPOINT_BLOBS_SQL = """
+    INSERT INTO checkpoint_blobs (thread_id, checkpoint_ns, channel, version, type, blob)
+    VALUES (?, ?, ?, ?, ?, ?)
+    ON CONFLICT (thread_id, checkpoint_ns, channel, version) DO NOTHING
+"""
+
+UPSERT_CHECKPOINTS_SQL = """
+    INSERT INTO checkpoints (thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, checkpoint, metadata)
+    VALUES (?, ?, ?, ?, ?, ?)
+    ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id)
+    DO UPDATE SET
+        checkpoint = EXCLUDED.checkpoint,
+        metadata = EXCLUDED.metadata;
+"""
+
+UPSERT_CHECKPOINT_WRITES_SQL = """
+    INSERT INTO checkpoint_writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, blob)
+    VALUES (?, ?, ?, ?, ?, ?, ?, ?)
+    ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) DO UPDATE SET
+        channel = EXCLUDED.channel,
+        type = EXCLUDED.type,
+        blob = EXCLUDED.blob;
+"""
+
+INSERT_CHECKPOINT_WRITES_SQL = """
+    INSERT INTO checkpoint_writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, blob)
+    VALUES (?, ?, ?, ?, ?, ?, ?, ?)
+    ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) DO NOTHING
+"""
+
+
+class BaseDuckDBSaver(BaseCheckpointSaver[str]):
+    SELECT_SQL = SELECT_SQL
+    MIGRATIONS = MIGRATIONS
+    UPSERT_CHECKPOINT_BLOBS_SQL = UPSERT_CHECKPOINT_BLOBS_SQL
+    UPSERT_CHECKPOINTS_SQL = UPSERT_CHECKPOINTS_SQL
+    UPSERT_CHECKPOINT_WRITES_SQL = UPSERT_CHECKPOINT_WRITES_SQL
+    INSERT_CHECKPOINT_WRITES_SQL = INSERT_CHECKPOINT_WRITES_SQL
+
+    jsonplus_serde = JsonPlusSerializer()
+
+    def _load_checkpoint(
+        self,
+        checkpoint_json_str: str,
+        channel_values: list[tuple[bytes, bytes, bytes]],
+        pending_sends: list[tuple[bytes, bytes]],
+    ) -> Checkpoint:
+        checkpoint = json.loads(checkpoint_json_str)
+        return {
+            **checkpoint,
+            "pending_sends": [
+                self.serde.loads_typed((c.decode(), b)) for c, b in pending_sends or []
+            ],
+            "channel_values": self._load_blobs(channel_values),
+        }
+
+    def _dump_checkpoint(self, checkpoint: Checkpoint) -> dict[str, Any]:
+        return {**checkpoint, "pending_sends": []}
+
+    def _load_blobs(
+        self, blob_values: list[tuple[bytes, bytes, bytes]]
+    ) -> dict[str, Any]:
+        if not blob_values:
+            return {}
+        return {
+            k.decode(): self.serde.loads_typed((t.decode(), v))
+            for k, t, v in blob_values
+            if t.decode() != "empty"
+        }
+
+    def _dump_blobs(
+        self,
+        thread_id: str,
+        checkpoint_ns: str,
+        values: dict[str, Any],
+        versions: ChannelVersions,
+    ) -> list[tuple[str, str, str, str, str, Optional[bytes]]]:
+        if not versions:
+            return []
+
+        return [
+            (
+                thread_id,
+                checkpoint_ns,
+                k,
+                cast(str, ver),
+                *(
+                    self.serde.dumps_typed(values[k])
+                    if k in values
+                    else ("empty", None)
+                ),
+            )
+            for k, ver in versions.items()
+        ]
+
+    def _load_writes(
+        self, writes: list[tuple[bytes, bytes, bytes, bytes]]
+    ) -> list[tuple[str, str, Any]]:
+        return (
+            [
+                (
+                    tid.decode(),
+                    channel.decode(),
+                    self.serde.loads_typed((t.decode(), v)),
+                )
+                for tid, channel, t, v in writes
+            ]
+            if writes
+            else []
+        )
+
+    def _dump_writes(
+        self,
+        thread_id: str,
+        checkpoint_ns: str,
+        checkpoint_id: str,
+        task_id: str,
+        writes: Sequence[tuple[str, Any]],
+    ) -> list[tuple[str, str, str, str, int, str, str, bytes]]:
+        return [
+            (
+                thread_id,
+                checkpoint_ns,
+                checkpoint_id,
+                task_id,
+                WRITES_IDX_MAP.get(channel, idx),
+                channel,
+                *self.serde.dumps_typed(value),
+            )
+            for idx, (channel, value) in enumerate(writes)
+        ]
+
+    def _load_metadata(self, metadata_json_str: str) -> CheckpointMetadata:
+        return self.jsonplus_serde.loads(metadata_json_str.encode())
+
+    def _dump_metadata(self, metadata: CheckpointMetadata) -> str:
+        serialized_metadata = self.jsonplus_serde.dumps(metadata)
+        # NOTE: we're using JSON serializer (not msgpack), so we need to remove null characters before writing
+        return serialized_metadata.decode().replace("\\u0000", "")
+
+    def get_next_version(self, current: Optional[str], channel: ChannelProtocol) -> str:
+        if current is None:
+            current_v = 0
+        elif isinstance(current, int):
+            current_v = current
+        else:
+            current_v = int(current.split(".")[0])
+        next_v = current_v + 1
+        next_h = random.random()
+        return f"{next_v:032}.{next_h:016}"
+
+    def _search_where(
+        self,
+        config: Optional[RunnableConfig],
+        filter: MetadataInput,
+        before: Optional[RunnableConfig] = None,
+    ) -> Tuple[str, List[Any]]:
+        """Return WHERE clause predicates for alist() given config, filter, before.
+
+        This method returns a tuple of a string and a tuple of values. The string
+        is the parametered WHERE clause predicate (including the WHERE keyword):
+        "WHERE column1 = $1 AND column2 IS $2". The list of values contains the
+        values for each of the corresponding parameters.
+        """
+        wheres = []
+        param_values = []
+
+        # construct predicate for config filter
+        if config:
+            wheres.append("thread_id = ?")
+            param_values.append(config["configurable"]["thread_id"])
+            checkpoint_ns = config["configurable"].get("checkpoint_ns")
+            if checkpoint_ns is not None:
+                wheres.append("checkpoint_ns = ?")
+                param_values.append(checkpoint_ns)
+
+            if checkpoint_id := get_checkpoint_id(config):
+                wheres.append("checkpoint_id = ?")
+                param_values.append(checkpoint_id)
+
+        # construct predicate for metadata filter
+        if filter:
+            wheres.append("json_contains(metadata, ?)")
+            param_values.append(json.dumps(filter))
+
+        # construct predicate for `before`
+        if before is not None:
+            wheres.append("checkpoint_id < ?")
+            param_values.append(get_checkpoint_id(before))
+
+        return (
+            "WHERE " + " AND ".join(wheres) if wheres else "",
+            param_values,
+        )

libs/checkpoint-duckdb/langgraph/store/duckdb/__init__.py

@@ -0,0 +1,4 @@
+from langgraph.store.duckdb.aio import AsyncDuckDBStore
+from langgraph.store.duckdb.base import DuckDBStore
+
+__all__ = ["AsyncDuckDBStore", "DuckDBStore"]

libs/checkpoint-duckdb/langgraph/store/duckdb/aio.py

@@ -0,0 +1,195 @@
+import asyncio
+import logging
+from contextlib import asynccontextmanager
+from typing import (
+    AsyncIterator,
+    Iterable,
+    Sequence,
+    cast,
+)
+
+import duckdb
+from langgraph.store.base import GetOp, ListNamespacesOp, Op, PutOp, Result, SearchOp
+from langgraph.store.base.batch import AsyncBatchedBaseStore
+from langgraph.store.duckdb.base import (
+    BaseDuckDBStore,
+    _convert_ns,
+    _group_ops,
+    _row_to_item,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class AsyncDuckDBStore(AsyncBatchedBaseStore, BaseDuckDBStore):
+    def __init__(
+        self,
+        conn: duckdb.DuckDBPyConnection,
+    ) -> None:
+        super().__init__()
+        self.conn = conn
+        self.loop = asyncio.get_running_loop()
+
+    async def abatch(self, ops: Iterable[Op]) -> list[Result]:
+        grouped_ops, num_ops = _group_ops(ops)
+        results: list[Result] = [None] * num_ops
+
+        tasks = []
+
+        if GetOp in grouped_ops:
+            tasks.append(
+                self._batch_get_ops(
+                    cast(Sequence[tuple[int, GetOp]], grouped_ops[GetOp]), results
+                )
+            )
+
+        if PutOp in grouped_ops:
+            tasks.append(
+                self._batch_put_ops(
+                    cast(Sequence[tuple[int, PutOp]], grouped_ops[PutOp])
+                )
+            )
+
+        if SearchOp in grouped_ops:
+            tasks.append(
+                self._batch_search_ops(
+                    cast(Sequence[tuple[int, SearchOp]], grouped_ops[SearchOp]),
+                    results,
+                )
+            )
+
+        if ListNamespacesOp in grouped_ops:
+            tasks.append(
+                self._batch_list_namespaces_ops(
+                    cast(
+                        Sequence[tuple[int, ListNamespacesOp]],
+                        grouped_ops[ListNamespacesOp],
+                    ),
+                    results,
+                )
+            )
+
+        await asyncio.gather(*tasks)
+
+        return results
+
+    def batch(self, ops: Iterable[Op]) -> list[Result]:
+        return asyncio.run_coroutine_threadsafe(self.abatch(ops), self.loop).result()
+
+    async def _batch_get_ops(
+        self,
+        get_ops: Sequence[tuple[int, GetOp]],
+        results: list[Result],
+    ) -> None:
+        cursors = []
+        for query, params, namespace, items in self._get_batch_GET_ops_queries(get_ops):
+            cur = self.conn.cursor()
+            await asyncio.to_thread(cur.execute, query, params)
+            cursors.append((cur, namespace, items))
+
+        for cur, namespace, items in cursors:
+            rows = await asyncio.to_thread(cur.fetchall)
+            key_to_row = {row[1]: row for row in rows}
+            for idx, key in items:
+                row = key_to_row.get(key)
+                if row:
+                    results[idx] = _row_to_item(namespace, row)
+                else:
+                    results[idx] = None
+
+    async def _batch_put_ops(
+        self,
+        put_ops: Sequence[tuple[int, PutOp]],
+    ) -> None:
+        queries = self._get_batch_PUT_queries(put_ops)
+        for query, params in queries:
+            cur = self.conn.cursor()
+            await asyncio.to_thread(cur.execute, query, params)
+
+    async def _batch_search_ops(
+        self,
+        search_ops: Sequence[tuple[int, SearchOp]],
+        results: list[Result],
+    ) -> None:
+        queries = self._get_batch_search_queries(search_ops)
+        cursors: list[tuple[duckdb.DuckDBPyConnection, int]] = []
+
+        for (query, params), (idx, _) in zip(queries, search_ops):
+            cur = self.conn.cursor()
+            await asyncio.to_thread(cur.execute, query, params)
+            cursors.append((cur, idx))
+
+        for cur, idx in cursors:
+            rows = await asyncio.to_thread(cur.fetchall)
+            items = [_row_to_item(_convert_ns(row[0]), row) for row in rows]
+            results[idx] = items
+
+    async def _batch_list_namespaces_ops(
+        self,
+        list_ops: Sequence[tuple[int, ListNamespacesOp]],
+        results: list[Result],
+    ) -> None:
+        queries = self._get_batch_list_namespaces_queries(list_ops)
+        cursors: list[tuple[duckdb.DuckDBPyConnection, int]] = []
+        for (query, params), (idx, _) in zip(queries, list_ops):
+            cur = self.conn.cursor()
+            await asyncio.to_thread(cur.execute, query, params)
+            cursors.append((cur, idx))
+
+        for cur, idx in cursors:
+            rows = cast(list[tuple], await asyncio.to_thread(cur.fetchall))
+            namespaces = [_convert_ns(row[0]) for row in rows]
+            results[idx] = namespaces
+
+    @classmethod
+    @asynccontextmanager
+    async def from_conn_string(
+        cls,
+        conn_string: str,
+    ) -> AsyncIterator["AsyncDuckDBStore"]:
+        """Create a new AsyncDuckDBStore instance from a connection string.
+
+        Args:
+            conn_string (str): The DuckDB connection info string.
+
+        Returns:
+            AsyncDuckDBStore: A new AsyncDuckDBStore instance.
+        """
+        with duckdb.connect(conn_string) as conn:
+            yield AsyncDuckDBStore(conn)
+
+    async def setup(self) -> None:
+        """Set up the store database asynchronously.
+
+        This method creates the necessary tables in the DuckDB database if they don't
+        already exist and runs database migrations. It is called automatically when needed and should not be called
+        directly by the user.
+        """
+        cur = self.conn.cursor()
+        try:
+            await asyncio.to_thread(
+                cur.execute, "SELECT v FROM store_migrations ORDER BY v DESC LIMIT 1"
+            )
+            row = await asyncio.to_thread(cur.fetchone)
+            if row is None:
+                version = -1
+            else:
+                version = row[0]
+        except duckdb.CatalogException:
+            version = -1
+            # Create store_migrations table if it doesn't exist
+            await asyncio.to_thread(
+                cur.execute,
+                """
+                CREATE TABLE IF NOT EXISTS store_migrations (
+                    v INTEGER PRIMARY KEY
+                )
+                """,
+            )
+        for v, migration in enumerate(
+            self.MIGRATIONS[version + 1 :], start=version + 1
+        ):
+            await asyncio.to_thread(cur.execute, migration)
+            await asyncio.to_thread(
+                cur.execute, "INSERT INTO store_migrations (v) VALUES (?)", (v,)
+            )

libs/checkpoint-duckdb/langgraph/store/duckdb/base.py

@@ -0,0 +1,391 @@
+import asyncio
+import json
+import logging
+from collections import defaultdict
+from contextlib import contextmanager
+from typing import (
+    Any,
+    Generic,
+    Iterable,
+    Iterator,
+    Sequence,
+    TypeVar,
+    Union,
+    cast,
+)
+
+import duckdb
+from langgraph.store.base import (
+    BaseStore,
+    GetOp,
+    Item,
+    ListNamespacesOp,
+    Op,
+    PutOp,
+    Result,
+    SearchOp,
+)
+
+logger = logging.getLogger(__name__)
+
+
+MIGRATIONS = [
+    """
+CREATE TABLE IF NOT EXISTS store (
+    prefix TEXT NOT NULL,
+    key TEXT NOT NULL,
+    value JSON NOT NULL,
+    created_at TIMESTAMP DEFAULT now(),
+    updated_at TIMESTAMP DEFAULT now(),
+    PRIMARY KEY (prefix, key)
+);
+""",
+    """
+CREATE INDEX IF NOT EXISTS store_prefix_idx ON store (prefix);
+""",
+]
+
+C = TypeVar("C", bound=duckdb.DuckDBPyConnection)
+
+
+class BaseDuckDBStore(Generic[C]):
+    MIGRATIONS = MIGRATIONS
+    conn: C
+
+    def _get_batch_GET_ops_queries(
+        self,
+        get_ops: Sequence[tuple[int, GetOp]],
+    ) -> list[tuple[str, tuple, tuple[str, ...], list]]:
+        namespace_groups = defaultdict(list)
+        for idx, op in get_ops:
+            namespace_groups[op.namespace].append((idx, op.key))
+        results = []
+        for namespace, items in namespace_groups.items():
+            _, keys = zip(*items)
+            keys_to_query = ",".join(["?"] * len(keys))
+            query = f"""
+                SELECT prefix, key, value, created_at, updated_at
+                FROM store
+                WHERE prefix = ? AND key IN ({keys_to_query})
+            """
+            params = (_namespace_to_text(namespace), *keys)
+            results.append((query, params, namespace, items))
+        return results
+
+    def _get_batch_PUT_queries(
+        self,
+        put_ops: Sequence[tuple[int, PutOp]],
+    ) -> list[tuple[str, Sequence]]:
+        inserts: list[PutOp] = []
+        deletes: list[PutOp] = []
+        for _, op in put_ops:
+            if op.value is None:
+                deletes.append(op)
+            else:
+                inserts.append(op)
+
+        queries: list[tuple[str, Sequence]] = []
+
+        if deletes:
+            namespace_groups: dict[tuple[str, ...], list[str]] = defaultdict(list)
+            for op in deletes:
+                namespace_groups[op.namespace].append(op.key)
+            for namespace, keys in namespace_groups.items():
+                placeholders = ",".join(["?"] * len(keys))
+                query = (
+                    f"DELETE FROM store WHERE prefix = ? AND key IN ({placeholders})"
+                )
+                params = (_namespace_to_text(namespace), *keys)
+                queries.append((query, params))
+        if inserts:
+            values = []
+            insertion_params = []
+            for op in inserts:
+                values.append("(?, ?, ?, now(), now())")
+                insertion_params.extend(
+                    [
+                        _namespace_to_text(op.namespace),
+                        op.key,
+                        json.dumps(op.value),
+                    ]
+                )
+            values_str = ",".join(values)
+            query = f"""
+                INSERT INTO store (prefix, key, value, created_at, updated_at)
+                VALUES {values_str}
+                ON CONFLICT (prefix, key) DO UPDATE
+                SET value = EXCLUDED.value, updated_at = now()
+            """
+            queries.append((query, insertion_params))
+
+        return queries
+
+    def _get_batch_search_queries(
+        self,
+        search_ops: Sequence[tuple[int, SearchOp]],
+    ) -> list[tuple[str, Sequence]]:
+        queries: list[tuple[str, Sequence]] = []
+        for _, op in search_ops:
+            query = """
+                SELECT prefix, key, value, created_at, updated_at
+                FROM store
+                WHERE prefix LIKE ?
+            """
+            params: list = [f"{_namespace_to_text(op.namespace_prefix)}%"]
+
+            if op.filter:
+                filter_conditions = []
+                for key, value in op.filter.items():
+                    filter_conditions.append(f"json_extract(value, '$.{key}') = ?")
+                    params.append(json.dumps(value))
+                query += " AND " + " AND ".join(filter_conditions)
+
+            query += " ORDER BY updated_at DESC LIMIT ? OFFSET ?"
+            params.extend([op.limit, op.offset])
+
+            queries.append((query, params))
+        return queries
+
+    def _get_batch_list_namespaces_queries(
+        self,
+        list_ops: Sequence[tuple[int, ListNamespacesOp]],
+    ) -> list[tuple[str, Sequence]]:
+        queries: list[tuple[str, Sequence]] = []
+        for _, op in list_ops:
+            query = """
+                WITH split_prefix AS (
+                    SELECT
+                        prefix,
+                        string_split(prefix, '.') AS parts
+                    FROM store
+                )
+                SELECT DISTINCT ON (truncated_prefix)
+                    CASE
+                        WHEN ? IS NOT NULL THEN
+                        array_to_string(array_slice(parts, 1, ?), '.')
+                        ELSE prefix
+                    END AS truncated_prefix,
+                    prefix
+                FROM split_prefix
+            """
+            params: list[Any] = [op.max_depth, op.max_depth]
+
+            conditions = []
+            if op.match_conditions:
+                for condition in op.match_conditions:
+                    if condition.match_type == "prefix":
+                        conditions.append("prefix LIKE ?")
+                        params.append(
+                            f"{_namespace_to_text(condition.path, handle_wildcards=True)}%"
+                        )
+                    elif condition.match_type == "suffix":
+                        conditions.append("prefix LIKE ?")
+                        params.append(
+                            f"%{_namespace_to_text(condition.path, handle_wildcards=True)}"
+                        )
+                    else:
+                        logger.warning(
+                            f"Unknown match_type in list_namespaces: {condition.match_type}"
+                        )
+
+            if conditions:
+                query += " WHERE " + " AND ".join(conditions)
+
+            query += " ORDER BY prefix LIMIT ? OFFSET ?"
+            params.extend([op.limit, op.offset])
+            queries.append((query, params))
+
+        return queries
+
+
+class DuckDBStore(BaseStore, BaseDuckDBStore[duckdb.DuckDBPyConnection]):
+    def __init__(
+        self,
+        conn: duckdb.DuckDBPyConnection,
+    ) -> None:
+        super().__init__()
+        self.conn = conn
+
+    def batch(self, ops: Iterable[Op]) -> list[Result]:
+        grouped_ops, num_ops = _group_ops(ops)
+        results: list[Result] = [None] * num_ops
+
+        if GetOp in grouped_ops:
+            self._batch_get_ops(
+                cast(Sequence[tuple[int, GetOp]], grouped_ops[GetOp]), results
+            )
+
+        if PutOp in grouped_ops:
+            self._batch_put_ops(cast(Sequence[tuple[int, PutOp]], grouped_ops[PutOp]))
+
+        if SearchOp in grouped_ops:
+            self._batch_search_ops(
+                cast(Sequence[tuple[int, SearchOp]], grouped_ops[SearchOp]),
+                results,
+            )
+
+        if ListNamespacesOp in grouped_ops:
+            self._batch_list_namespaces_ops(
+                cast(
+                    Sequence[tuple[int, ListNamespacesOp]],
+                    grouped_ops[ListNamespacesOp],
+                ),
+                results,
+            )
+
+        return results
+
+    async def abatch(self, ops: Iterable[Op]) -> list[Result]:
+        return await asyncio.get_running_loop().run_in_executor(None, self.batch, ops)
+
+    def _batch_get_ops(
+        self,
+        get_ops: Sequence[tuple[int, GetOp]],
+        results: list[Result],
+    ) -> None:
+        cursors = []
+        for query, params, namespace, items in self._get_batch_GET_ops_queries(get_ops):
+            cur = self.conn.cursor()
+            cur.execute(query, params)
+            cursors.append((cur, namespace, items))
+
+        for cur, namespace, items in cursors:
+            rows = cur.fetchall()
+            key_to_row = {row[1]: row for row in rows}
+            for idx, key in items:
+                row = key_to_row.get(key)
+                if row:
+                    results[idx] = _row_to_item(namespace, row)
+                else:
+                    results[idx] = None
+
+    def _batch_put_ops(
+        self,
+        put_ops: Sequence[tuple[int, PutOp]],
+    ) -> None:
+        queries = self._get_batch_PUT_queries(put_ops)
+        for query, params in queries:
+            cur = self.conn.cursor()
+            cur.execute(query, params)
+
+    def _batch_search_ops(
+        self,
+        search_ops: Sequence[tuple[int, SearchOp]],
+        results: list[Result],
+    ) -> None:
+        queries = self._get_batch_search_queries(search_ops)
+        cursors: list[tuple[duckdb.DuckDBPyConnection, int]] = []
+
+        for (query, params), (idx, _) in zip(queries, search_ops):
+            cur = self.conn.cursor()
+            cur.execute(query, params)
+            cursors.append((cur, idx))
+
+        for cur, idx in cursors:
+            rows = cur.fetchall()
+            items = [_row_to_item(_convert_ns(row[0]), row) for row in rows]
+            results[idx] = items
+
+    def _batch_list_namespaces_ops(
+        self,
+        list_ops: Sequence[tuple[int, ListNamespacesOp]],
+        results: list[Result],
+    ) -> None:
+        queries = self._get_batch_list_namespaces_queries(list_ops)
+        cursors: list[tuple[duckdb.DuckDBPyConnection, int]] = []
+        for (query, params), (idx, _) in zip(queries, list_ops):
+            cur = self.conn.cursor()
+            cur.execute(query, params)
+            cursors.append((cur, idx))
+
+        for cur, idx in cursors:
+            rows = cast(list[dict], cur.fetchall())
+            namespaces = [_convert_ns(row[0]) for row in rows]
+            results[idx] = namespaces
+
+    @classmethod
+    @contextmanager
+    def from_conn_string(
+        cls,
+        conn_string: str,
+    ) -> Iterator["DuckDBStore"]:
+        """Create a new BaseDuckDBStore instance from a connection string.
+
+        Args:
+            conn_string (str): The DuckDB connection info string.
+
+        Returns:
+            DuckDBStore: A new DuckDBStore instance.
+        """
+        with duckdb.connect(conn_string) as conn:
+            yield cls(conn=conn)
+
+    def setup(self) -> None:
+        """Set up the store database.
+
+        This method creates the necessary tables in the DuckDB database if they don't
+        already exist and runs database migrations. It is called automatically when needed and should not be called
+        directly by the user.
+        """
+        with self.conn.cursor() as cur:
+            try:
+                cur.execute("SELECT v FROM store_migrations ORDER BY v DESC LIMIT 1")
+                row = cast(dict, cur.fetchone())
+                if row is None:
+                    version = -1
+                else:
+                    version = row["v"]
+            except duckdb.CatalogException:
+                version = -1
+                # Create store_migrations table if it doesn't exist
+                cur.execute(
+                    """
+                    CREATE TABLE IF NOT EXISTS store_migrations (
+                        v INTEGER PRIMARY KEY
+                    )
+                """
+                )
+            for v, migration in enumerate(
+                self.MIGRATIONS[version + 1 :], start=version + 1
+            ):
+                cur.execute(migration)
+                cur.execute("INSERT INTO store_migrations (v) VALUES (?)", (v,))
+
+
+def _namespace_to_text(
+    namespace: tuple[str, ...], handle_wildcards: bool = False
+) -> str:
+    """Convert namespace tuple to text string."""
+    if handle_wildcards:
+        namespace = tuple("%" if val == "*" else val for val in namespace)
+    return ".".join(namespace)
+
+
+def _row_to_item(
+    namespace: tuple[str, ...],
+    row: tuple,
+) -> Item:
+    """Convert a row from the database into an Item."""
+    _, key, val, created_at, updated_at = row
+    return Item(
+        value=val if isinstance(val, dict) else json.loads(val),
+        key=key,
+        namespace=namespace,
+        created_at=created_at,
+        updated_at=updated_at,
+    )
+
+
+def _group_ops(ops: Iterable[Op]) -> tuple[dict[type, list[tuple[int, Op]]], int]:
+    grouped_ops: dict[type, list[tuple[int, Op]]] = defaultdict(list)
+    tot = 0
+    for idx, op in enumerate(ops):
+        grouped_ops[type(op)].append((idx, op))
+        tot += 1
+    return grouped_ops, tot
+
+
+def _convert_ns(namespace: Union[str, list]) -> tuple[str, ...]:
+    if isinstance(namespace, list):
+        return tuple(namespace)
+    return tuple(namespace.split("."))

libs/checkpoint-duckdb/pyproject.toml

@@ -0,0 +1,60 @@
+[tool.poetry]
+name = "langgraph-checkpoint-duckdb"
+version = "2.0.1"
+description = "Library with a DuckDB implementation of LangGraph checkpoint saver."
+authors = []
+license = "MIT"
+readme = "README.md"
+repository = "https://www.github.com/langchain-ai/langgraph"
+packages = [{ include = "langgraph" }]
+
+[tool.poetry.dependencies]
+python = "^3.9.0,<4.0"
+langgraph-checkpoint = "^2.0.2"
+duckdb = ">=1.1.2"
+
+[tool.poetry.group.dev.dependencies]
+ruff = "^0.6.2"
+codespell = "^2.2.0"
+pytest = "^7.2.1"
+anyio = "^4.4.0"
+pytest-asyncio = "^0.21.1"
+pytest-mock = "^3.11.1"
+pytest-watch = "^4.2.0"
+mypy = "^1.10.0"
+langgraph-checkpoint = {path = "../checkpoint", develop = true}
+
+[tool.pytest.ini_options]
+# --strict-markers will raise errors on unknown marks.
+# https://docs.pytest.org/en/7.1.x/how-to/mark.html#raising-errors-on-unknown-marks
+#
+# https://docs.pytest.org/en/7.1.x/reference/reference.html
+# --strict-config       any warnings encountered while parsing the `pytest`
+#                       section of the configuration file raise errors.
+addopts = "--strict-markers --strict-config --durations=5 -vv"
+asyncio_mode = "auto"
+
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.ruff]
+lint.select = [
+  "E",  # pycodestyle
+  "F",  # Pyflakes
+  "UP", # pyupgrade
+  "B",  # flake8-bugbear
+  "I",  # isort
+]
+lint.ignore = ["E501", "B008", "UP007", "UP006"]
+
+[tool.mypy]
+# https://mypy.readthedocs.io/en/stable/config_file.html
+disallow_untyped_defs = "True"
+explicit_package_bases = "True"
+warn_no_return = "False"
+warn_unused_ignores = "True"
+warn_redundant_casts = "True"
+allow_redefinition = "True"
+disable_error_code = "typeddict-item, return-value"

libs/checkpoint-duckdb/tests/test_async.py

@@ -0,0 +1,112 @@
+from typing import Any
+
+import pytest
+from langchain_core.runnables import RunnableConfig
+
+from langgraph.checkpoint.base import (
+    Checkpoint,
+    CheckpointMetadata,
+    create_checkpoint,
+    empty_checkpoint,
+)
+from langgraph.checkpoint.duckdb.aio import AsyncDuckDBSaver
+
+
+class TestAsyncDuckDBSaver:
+    @pytest.fixture(autouse=True)
+    async def setup(self) -> None:
+        # objects for test setup
+        self.config_1: RunnableConfig = {
+            "configurable": {
+                "thread_id": "thread-1",
+                # for backwards compatibility testing
+                "thread_ts": "1",
+                "checkpoint_ns": "",
+            }
+        }
+        self.config_2: RunnableConfig = {
+            "configurable": {
+                "thread_id": "thread-2",
+                "checkpoint_id": "2",
+                "checkpoint_ns": "",
+            }
+        }
+        self.config_3: RunnableConfig = {
+            "configurable": {
+                "thread_id": "thread-2",
+                "checkpoint_id": "2-inner",
+                "checkpoint_ns": "inner",
+            }
+        }
+
+        self.chkpnt_1: Checkpoint = empty_checkpoint()
+        self.chkpnt_2: Checkpoint = create_checkpoint(self.chkpnt_1, {}, 1)
+        self.chkpnt_3: Checkpoint = empty_checkpoint()
+
+        self.metadata_1: CheckpointMetadata = {
+            "source": "input",
+            "step": 2,
+            "writes": {},
+            "score": 1,
+        }
+        self.metadata_2: CheckpointMetadata = {
+            "source": "loop",
+            "step": 1,
+            "writes": {"foo": "bar"},
+            "score": None,
+        }
+        self.metadata_3: CheckpointMetadata = {}
+
+    async def test_asearch(self) -> None:
+        async with AsyncDuckDBSaver.from_conn_string(":memory:") as saver:
+            await saver.setup()
+            await saver.aput(self.config_1, self.chkpnt_1, self.metadata_1, {})
+            await saver.aput(self.config_2, self.chkpnt_2, self.metadata_2, {})
+            await saver.aput(self.config_3, self.chkpnt_3, self.metadata_3, {})
+
+            # call method / assertions
+            query_1 = {"source": "input"}  # search by 1 key
+            query_2 = {
+                "step": 1,
+                "writes": {"foo": "bar"},
+            }  # search by multiple keys
+            query_3: dict[str, Any] = {}  # search by no keys, return all checkpoints
+            query_4 = {"source": "update", "step": 1}  # no match
+
+            search_results_1 = [c async for c in saver.alist(None, filter=query_1)]
+            assert len(search_results_1) == 1
+            assert search_results_1[0].metadata == self.metadata_1
+
+            search_results_2 = [c async for c in saver.alist(None, filter=query_2)]
+            assert len(search_results_2) == 1
+            assert search_results_2[0].metadata == self.metadata_2
+
+            search_results_3 = [c async for c in saver.alist(None, filter=query_3)]
+            assert len(search_results_3) == 3
+
+            search_results_4 = [c async for c in saver.alist(None, filter=query_4)]
+            assert len(search_results_4) == 0
+
+            # search by config (defaults to checkpoints across all namespaces)
+            search_results_5 = [
+                c
+                async for c in saver.alist({"configurable": {"thread_id": "thread-2"}})
+            ]
+            assert len(search_results_5) == 2
+            assert {
+                search_results_5[0].config["configurable"]["checkpoint_ns"],
+                search_results_5[1].config["configurable"]["checkpoint_ns"],
+            } == {"", "inner"}
+
+            # TODO: test before and limit params
+
+    async def test_null_chars(self) -> None:
+        async with AsyncDuckDBSaver.from_conn_string(":memory:") as saver:
+            await saver.setup()
+            config = await saver.aput(
+                self.config_1, self.chkpnt_1, {"my_key": "\x00abc"}, {}
+            )
+            assert (await saver.aget_tuple(config)).metadata["my_key"] == "abc"  # type: ignore
+            assert [c async for c in saver.alist(None, filter={"my_key": "abc"})][
+                0
+            ].metadata["my_key"] == "abc"

libs/checkpoint-duckdb/tests/test_async_store.py

@@ -0,0 +1,517 @@
+# type: ignore
+import uuid
+from datetime import datetime
+from typing import Any
+from unittest.mock import MagicMock
+
+import pytest
+
+from langgraph.store.base import GetOp, Item, ListNamespacesOp, PutOp, SearchOp
+from langgraph.store.duckdb import AsyncDuckDBStore
+
+
+class MockCursor:
+    def __init__(self, fetch_result: Any) -> None:
+        self.fetch_result = fetch_result
+        self.execute = MagicMock()
+        self.fetchall = MagicMock(return_value=self.fetch_result)
+
+
+class MockConnection:
+    def __init__(self) -> None:
+        self.cursor = MagicMock()
+
+
+@pytest.fixture
+def mock_connection() -> MockConnection:
+    return MockConnection()
+
+
+@pytest.fixture
+async def store(mock_connection: MockConnection) -> AsyncDuckDBStore:
+    duck_db_store = AsyncDuckDBStore(mock_connection)
+    await duck_db_store.setup()
+    return duck_db_store
+
+
+async def test_abatch_order(store: AsyncDuckDBStore) -> None:
+    mock_connection = store.conn
+    mock_get_cursor = MockCursor(
+        [
+            (
+                "test.foo",
+                "key1",
+                '{"data": "value1"}',
+                datetime.now(),
+                datetime.now(),
+            ),
+            (
+                "test.bar",
+                "key2",
+                '{"data": "value2"}',
+                datetime.now(),
+                datetime.now(),
+            ),
+        ]
+    )
+    mock_search_cursor = MockCursor(
+        [
+            (
+                "test.foo",
+                "key1",
+                '{"data": "value1"}',
+                datetime.now(),
+                datetime.now(),
+            ),
+        ]
+    )
+    mock_list_namespaces_cursor = MockCursor(
+        [
+            ("test",),
+        ]
+    )
+
+    failures = []
+
+    def cursor_side_effect() -> Any:
+        cursor = MagicMock()
+
+        def execute_side_effect(query: str, *params: Any) -> None:
+            # My super sophisticated database.
+            if "WHERE prefix = ? AND key" in query:
+                cursor.fetchall = mock_get_cursor.fetchall
+            elif "SELECT prefix, key, value" in query:
+                cursor.fetchall = mock_search_cursor.fetchall
+            elif "SELECT DISTINCT ON (truncated_prefix)" in query:
+                cursor.fetchall = mock_list_namespaces_cursor.fetchall
+            elif "INSERT INTO " in query:
+                pass
+            else:
+                e = ValueError(f"Unmatched query: {query}")
+                failures.append(e)
+                raise e
+
+        cursor.execute = MagicMock(side_effect=execute_side_effect)
+        return cursor
+
+    mock_connection.cursor.side_effect = cursor_side_effect  # type: ignore
+
+    ops = [
+        GetOp(namespace=("test",), key="key1"),
+        PutOp(namespace=("test",), key="key2", value={"data": "value2"}),
+        SearchOp(
+            namespace_prefix=("test",), filter={"data": "value1"}, limit=10, offset=0
+        ),
+        ListNamespacesOp(match_conditions=None, max_depth=None, limit=10, offset=0),
+        GetOp(namespace=("test",), key="key3"),
+    ]
+    results = await store.abatch(ops)
+    assert not failures
+    assert len(results) == 5
+    assert isinstance(results[0], Item)
+    assert isinstance(results[0].value, dict)
+    assert results[0].value == {"data": "value1"}
+    assert results[0].key == "key1"
+    assert results[1] is None
+    assert isinstance(results[2], list)
+    assert len(results[2]) == 1
+    assert isinstance(results[3], list)
+    assert results[3] == [("test",)]
+    assert results[4] is None
+
+    ops_reordered = [
+        SearchOp(namespace_prefix=("test",), filter=None, limit=5, offset=0),
+        GetOp(namespace=("test",), key="key2"),
+        ListNamespacesOp(match_conditions=None, max_depth=None, limit=5, offset=0),
+        PutOp(namespace=("test",), key="key3", value={"data": "value3"}),
+        GetOp(namespace=("test",), key="key1"),
+    ]
+
+    results_reordered = await store.abatch(ops_reordered)
+    assert not failures
+    assert len(results_reordered) == 5
+    assert isinstance(results_reordered[0], list)
+    assert len(results_reordered[0]) == 1
+    assert isinstance(results_reordered[1], Item)
+    assert results_reordered[1].value == {"data": "value2"}
+    assert results_reordered[1].key == "key2"
+    assert isinstance(results_reordered[2], list)
+    assert results_reordered[2] == [("test",)]
+    assert results_reordered[3] is None
+    assert isinstance(results_reordered[4], Item)
+    assert results_reordered[4].value == {"data": "value1"}
+    assert results_reordered[4].key == "key1"
+
+
+async def test_batch_get_ops(store: AsyncDuckDBStore) -> None:
+    mock_connection = store.conn
+    mock_cursor = MockCursor(
+        [
+            (
+                "test.foo",
+                "key1",
+                '{"data": "value1"}',
+                datetime.now(),
+                datetime.now(),
+            ),
+            (
+                "test.bar",
+                "key2",
+                '{"data": "value2"}',
+                datetime.now(),
+                datetime.now(),
+            ),
+        ]
+    )
+    mock_connection.cursor.return_value = mock_cursor
+
+    ops = [
+        GetOp(namespace=("test",), key="key1"),
+        GetOp(namespace=("test",), key="key2"),
+        GetOp(namespace=("test",), key="key3"),
+    ]
+
+    results = await store.abatch(ops)
+
+    assert len(results) == 3
+    assert results[0] is not None
+    assert results[1] is not None
+    assert results[2] is None
+    assert results[0].key == "key1"
+    assert results[1].key == "key2"
+
+
+async def test_batch_put_ops(store: AsyncDuckDBStore) -> None:
+    mock_connection = store.conn
+    mock_cursor = MockCursor([])
+    mock_connection.cursor.return_value = mock_cursor
+
+    ops = [
+        PutOp(namespace=("test",), key="key1", value={"data": "value1"}),
+        PutOp(namespace=("test",), key="key2", value={"data": "value2"}),
+        PutOp(namespace=("test",), key="key3", value=None),
+    ]
+
+    results = await store.abatch(ops)
+
+    assert len(results) == 3
+    assert all(result is None for result in results)
+    assert mock_cursor.execute.call_count == 2
+
+
+async def test_batch_search_ops(store: AsyncDuckDBStore) -> None:
+    mock_connection = store.conn
+    mock_cursor = MockCursor(
+        [
+            (
+                "test.foo",
+                "key1",
+                '{"data": "value1"}',
+                datetime.now(),
+                datetime.now(),
+            ),
+            (
+                "test.bar",
+                "key2",
+                '{"data": "value2"}',
+                datetime.now(),
+                datetime.now(),
+            ),
+        ]
+    )
+    mock_connection.cursor.return_value = mock_cursor
+
+    ops = [
+        SearchOp(
+            namespace_prefix=("test",), filter={"data": "value1"}, limit=10, offset=0
+        ),
+        SearchOp(namespace_prefix=("test",), filter=None, limit=5, offset=0),
+    ]
+
+    results = await store.abatch(ops)
+
+    assert len(results) == 2
+    assert len(results[0]) == 2
+    assert len(results[1]) == 2
+
+
+async def test_batch_list_namespaces_ops(store: AsyncDuckDBStore) -> None:
+    mock_connection = store.conn
+    mock_cursor = MockCursor([("test.namespace1",), ("test.namespace2",)])
+    mock_connection.cursor.return_value = mock_cursor
+
+    ops = [ListNamespacesOp(match_conditions=None, max_depth=None, limit=10, offset=0)]
+
+    results = await store.abatch(ops)
+
+    assert len(results) == 1
+    assert results[0] == [("test", "namespace1"), ("test", "namespace2")]
+
+
+# The following use the actual DB connection
+
+
+async def test_basic_store_ops() -> None:
+    async with AsyncDuckDBStore.from_conn_string(":memory:") as store:
+        await store.setup()
+        namespace = ("test", "documents")
+        item_id = "doc1"
+        item_value = {"title": "Test Document", "content": "Hello, World!"}
+
+        await store.aput(namespace, item_id, item_value)
+        item = await store.aget(namespace, item_id)
+
+        assert item
+        assert item.namespace == namespace
+        assert item.key == item_id
+        assert item.value == item_value
+
+        updated_value = {
+            "title": "Updated Test Document",
+            "content": "Hello, LangGraph!",
+        }
+        await store.aput(namespace, item_id, updated_value)
+        updated_item = await store.aget(namespace, item_id)
+
+        assert updated_item.value == updated_value
+        assert updated_item.updated_at > item.updated_at
+        different_namespace = ("test", "other_documents")
+        item_in_different_namespace = await store.aget(different_namespace, item_id)
+        assert item_in_different_namespace is None
+
+        new_item_id = "doc2"
+        new_item_value = {"title": "Another Document", "content": "Greetings!"}
+        await store.aput(namespace, new_item_id, new_item_value)
+
+        search_results = await store.asearch(["test"], limit=10)
+        items = search_results
+        assert len(items) == 2
+        assert any(item.key == item_id for item in items)
+        assert any(item.key == new_item_id for item in items)
+
+        namespaces = await store.alist_namespaces(prefix=["test"])
+        assert ("test", "documents") in namespaces
+
+        await store.adelete(namespace, item_id)
+        await store.adelete(namespace, new_item_id)
+        deleted_item = await store.aget(namespace, item_id)
+        assert deleted_item is None
+
+        deleted_item = await store.aget(namespace, new_item_id)
+        assert deleted_item is None
+
+        empty_search_results = await store.asearch(["test"], limit=10)
+        assert len(empty_search_results) == 0
+
+
+async def test_list_namespaces() -> None:
+    async with AsyncDuckDBStore.from_conn_string(":memory:") as store:
+        await store.setup()
+        test_pref = str(uuid.uuid4())
+        test_namespaces = [
+            (test_pref, "test", "documents", "public", test_pref),
+            (test_pref, "test", "documents", "private", test_pref),
+            (test_pref, "test", "images", "public", test_pref),
+            (test_pref, "test", "images", "private", test_pref),
+            (test_pref, "prod", "documents", "public", test_pref),
+            (
+                test_pref,
+                "prod",
+                "documents",
+                "some",
+                "nesting",
+                "public",
+                test_pref,
+            ),
+            (test_pref, "prod", "documents", "private", test_pref),
+        ]
+
+        for namespace in test_namespaces:
+            await store.aput(namespace, "dummy", {"content": "dummy"})
+
+        prefix_result = await store.alist_namespaces(prefix=[test_pref, "test"])
+        assert len(prefix_result) == 4
+        assert all([ns[1] == "test" for ns in prefix_result])
+
+        specific_prefix_result = await store.alist_namespaces(
+            prefix=[test_pref, "test", "documents"]
+        )
+        assert len(specific_prefix_result) == 2
+        assert all([ns[1:3] == ("test", "documents") for ns in specific_prefix_result])
+
+        suffix_result = await store.alist_namespaces(suffix=["public", test_pref])
+        assert len(suffix_result) == 4
+        assert all(ns[-2] == "public" for ns in suffix_result)
+
+        prefix_suffix_result = await store.alist_namespaces(
+            prefix=[test_pref, "test"], suffix=["public", test_pref]
+        )
+        assert len(prefix_suffix_result) == 2
+        assert all(
+            ns[1] == "test" and ns[-2] == "public" for ns in prefix_suffix_result
+        )
+
+        wildcard_prefix_result = await store.alist_namespaces(
+            prefix=[test_pref, "*", "documents"]
+        )
+        assert len(wildcard_prefix_result) == 5
+        assert all(ns[2] == "documents" for ns in wildcard_prefix_result)
+
+        wildcard_suffix_result = await store.alist_namespaces(
+            suffix=["*", "public", test_pref]
+        )
+        assert len(wildcard_suffix_result) == 4
+        assert all(ns[-2] == "public" for ns in wildcard_suffix_result)
+        wildcard_single = await store.alist_namespaces(
+            suffix=["some", "*", "public", test_pref]
+        )
+        assert len(wildcard_single) == 1
+        assert wildcard_single[0] == (
+            test_pref,
+            "prod",
+            "documents",
+            "some",
+            "nesting",
+            "public",
+            test_pref,
+        )
+
+        max_depth_result = await store.alist_namespaces(max_depth=3)
+        assert all([len(ns) <= 3 for ns in max_depth_result])
+        max_depth_result = await store.alist_namespaces(
+            max_depth=4, prefix=[test_pref, "*", "documents"]
+        )
+        assert (
+            len(set(tuple(res) for res in max_depth_result))
+            == len(max_depth_result)
+            == 5
+        )
+
+        limit_result = await store.alist_namespaces(prefix=[test_pref], limit=3)
+        assert len(limit_result) == 3
+
+        offset_result = await store.alist_namespaces(prefix=[test_pref], offset=3)
+        assert len(offset_result) == len(test_namespaces) - 3
+
+        empty_prefix_result = await store.alist_namespaces(prefix=[test_pref])
+        assert len(empty_prefix_result) == len(test_namespaces)
+        assert set(tuple(ns) for ns in empty_prefix_result) == set(
+            tuple(ns) for ns in test_namespaces
+        )
+
+        for namespace in test_namespaces:
+            await store.adelete(namespace, "dummy")
+
+
+async def test_search():
+    async with AsyncDuckDBStore.from_conn_string(":memory:") as store:
+        await store.setup()
+        test_namespaces = [
+            ("test_search", "documents", "user1"),
+            ("test_search", "documents", "user2"),
+            ("test_search", "reports", "department1"),
+            ("test_search", "reports", "department2"),
+        ]
+        test_items = [
+            {"title": "Doc 1", "author": "John Doe", "tags": ["important"]},
+            {"title": "Doc 2", "author": "Jane Smith", "tags": ["draft"]},
+            {"title": "Report A", "author": "John Doe", "tags": ["final"]},
+            {"title": "Report B", "author": "Alice Johnson", "tags": ["draft"]},
+        ]
+        empty = await store.asearch(
+            (
+                "scoped",
+                "assistant_id",
+                "shared",
+                "6c5356f6-63ab-4158-868d-cd9fd14c736e",
+            ),
+            limit=10,
+            offset=0,
+        )
+        assert len(empty) == 0
+
+        for namespace, item in zip(test_namespaces, test_items):
+            await store.aput(namespace, f"item_{namespace[-1]}", item)
+
+        docs_result = await store.asearch(["test_search", "documents"])
+        assert len(docs_result) == 2
+        assert all([item.namespace[1] == "documents" for item in docs_result]), [
+            item.namespace for item in docs_result
+        ]
+
+        reports_result = await store.asearch(["test_search", "reports"])
+        assert len(reports_result) == 2
+        assert all(item.namespace[1] == "reports" for item in reports_result)
+
+        limited_result = await store.asearch(["test_search"], limit=2)
+        assert len(limited_result) == 2
+        offset_result = await store.asearch(["test_search"])
+        assert len(offset_result) == 4
+
+        offset_result = await store.asearch(["test_search"], offset=2)
+        assert len(offset_result) == 2
+        assert all(item not in limited_result for item in offset_result)
+
+        john_doe_result = await store.asearch(
+            ["test_search"], filter={"author": "John Doe"}
+        )
+        assert len(john_doe_result) == 2
+        assert all(item.value["author"] == "John Doe" for item in john_doe_result)
+
+        draft_result = await store.asearch(["test_search"], filter={"tags": ["draft"]})
+        assert len(draft_result) == 2
+        assert all("draft" in item.value["tags"] for item in draft_result)
+
+        page1 = await store.asearch(["test_search"], limit=2, offset=0)
+        page2 = await store.asearch(["test_search"], limit=2, offset=2)
+        all_items = page1 + page2
+        assert len(all_items) == 4
+        assert len(set(item.key for item in all_items)) == 4
+        empty = await store.asearch(
+            (
+                "scoped",
+                "assistant_id",
+                "shared",
+                "again",
+                "maybe",
+                "some-long",
+                "6be5cb0e-2eb4-42e6-bb6b-fba3c269db25",
+            ),
+            limit=10,
+            offset=0,
+        )
+        assert len(empty) == 0
+
+        # Test with a namespace beginning with a number (like a UUID)
+        uuid_namespace = (str(uuid.uuid4()), "documents")
+        uuid_item_id = "uuid_doc"
+        uuid_item_value = {
+            "title": "UUID Document",
+            "content": "This document has a UUID namespace.",
+        }
+
+        # Insert the item with the UUID namespace
+        await store.aput(uuid_namespace, uuid_item_id, uuid_item_value)
+
+        # Retrieve the item to verify it was stored correctly
+        retrieved_item = await store.aget(uuid_namespace, uuid_item_id)
+        assert retrieved_item is not None
+        assert retrieved_item.namespace == uuid_namespace
+        assert retrieved_item.key == uuid_item_id
+        assert retrieved_item.value == uuid_item_value
+
+        # Search for the item using the UUID namespace
+        search_result = await store.asearch([uuid_namespace[0]])
+        assert len(search_result) == 1
+        assert search_result[0].key == uuid_item_id
+        assert search_result[0].value == uuid_item_value
+
+        # Clean up: delete the item with the UUID namespace
+        await store.adelete(uuid_namespace, uuid_item_id)
+
+        # Verify the item was deleted
+        deleted_item = await store.aget(uuid_namespace, uuid_item_id)
+        assert deleted_item is None
+
+        for namespace in test_namespaces:
+            await store.adelete(namespace, f"item_{namespace[-1]}")

libs/checkpoint-duckdb/tests/test_store.py

@@ -0,0 +1,457 @@
+# type: ignore
+import uuid
+from datetime import datetime
+from typing import Any
+from unittest.mock import MagicMock
+
+import pytest
+
+from langgraph.store.base import GetOp, Item, ListNamespacesOp, PutOp, SearchOp
+from langgraph.store.duckdb import DuckDBStore
+
+
+class MockCursor:
+    def __init__(self, fetch_result: Any) -> None:
+        self.fetch_result = fetch_result
+        self.execute = MagicMock()
+        self.fetchall = MagicMock(return_value=self.fetch_result)
+
+
+class MockConnection:
+    def __init__(self) -> None:
+        self.cursor = MagicMock()
+
+
+@pytest.fixture
+def mock_connection() -> MockConnection:
+    return MockConnection()
+
+
+@pytest.fixture
+def store(mock_connection: MockConnection) -> DuckDBStore:
+    duck_db_store = DuckDBStore(mock_connection)
+    duck_db_store.setup()
+    return duck_db_store
+
+
+def test_batch_order(store: DuckDBStore) -> None:
+    mock_connection = store.conn
+    mock_get_cursor = MockCursor(
+        [
+            (
+                "test.foo",
+                "key1",
+                '{"data": "value1"}',
+                datetime.now(),
+                datetime.now(),
+            ),
+            (
+                "test.bar",
+                "key2",
+                '{"data": "value2"}',
+                datetime.now(),
+                datetime.now(),
+            ),
+        ]
+    )
+    mock_search_cursor = MockCursor(
+        [
+            (
+                "test.foo",
+                "key1",
+                '{"data": "value1"}',
+                datetime.now(),
+                datetime.now(),
+            ),
+        ]
+    )
+    mock_list_namespaces_cursor = MockCursor(
+        [
+            ("test",),
+        ]
+    )
+
+    failures = []
+
+    def cursor_side_effect() -> Any:
+        cursor = MagicMock()
+
+        def execute_side_effect(query: str, *params: Any) -> None:
+            # My super sophisticated database.
+            if "WHERE prefix = ? AND key" in query:
+                cursor.fetchall = mock_get_cursor.fetchall
+            elif "SELECT prefix, key, value" in query:
+                cursor.fetchall = mock_search_cursor.fetchall
+            elif "SELECT DISTINCT ON (truncated_prefix)" in query:
+                cursor.fetchall = mock_list_namespaces_cursor.fetchall
+            elif "INSERT INTO " in query:
+                pass
+            else:
+                e = ValueError(f"Unmatched query: {query}")
+                failures.append(e)
+                raise e
+
+        cursor.execute = MagicMock(side_effect=execute_side_effect)
+        return cursor
+
+    mock_connection.cursor.side_effect = cursor_side_effect
+
+    ops = [
+        GetOp(namespace=("test",), key="key1"),
+        PutOp(namespace=("test",), key="key2", value={"data": "value2"}),
+        SearchOp(
+            namespace_prefix=("test",), filter={"data": "value1"}, limit=10, offset=0
+        ),
+        ListNamespacesOp(match_conditions=None, max_depth=None, limit=10, offset=0),
+        GetOp(namespace=("test",), key="key3"),
+    ]
+    results = store.batch(ops)
+    assert not failures
+    assert len(results) == 5
+    assert isinstance(results[0], Item)
+    assert isinstance(results[0].value, dict)
+    assert results[0].value == {"data": "value1"}
+    assert results[0].key == "key1"
+    assert results[1] is None
+    assert isinstance(results[2], list)
+    assert len(results[2]) == 1
+    assert isinstance(results[3], list)
+    assert results[3] == [("test",)]
+    assert results[4] is None
+
+    ops_reordered = [
+        SearchOp(namespace_prefix=("test",), filter=None, limit=5, offset=0),
+        GetOp(namespace=("test",), key="key2"),
+        ListNamespacesOp(match_conditions=None, max_depth=None, limit=5, offset=0),
+        PutOp(namespace=("test",), key="key3", value={"data": "value3"}),
+        GetOp(namespace=("test",), key="key1"),
+    ]
+
+    results_reordered = store.batch(ops_reordered)
+    assert not failures
+    assert len(results_reordered) == 5
+    assert isinstance(results_reordered[0], list)
+    assert len(results_reordered[0]) == 1
+    assert isinstance(results_reordered[1], Item)
+    assert results_reordered[1].value == {"data": "value2"}
+    assert results_reordered[1].key == "key2"
+    assert isinstance(results_reordered[2], list)
+    assert results_reordered[2] == [("test",)]
+    assert results_reordered[3] is None
+    assert isinstance(results_reordered[4], Item)
+    assert results_reordered[4].value == {"data": "value1"}
+    assert results_reordered[4].key == "key1"
+
+
+def test_batch_get_ops(store: DuckDBStore) -> None:
+    mock_connection = store.conn
+    mock_cursor = MockCursor(
+        [
+            (
+                "test.foo",
+                "key1",
+                '{"data": "value1"}',
+                datetime.now(),
+                datetime.now(),
+            ),
+            (
+                "test.bar",
+                "key2",
+                '{"data": "value2"}',
+                datetime.now(),
+                datetime.now(),
+            ),
+        ]
+    )
+    mock_connection.cursor.return_value = mock_cursor
+
+    ops = [
+        GetOp(namespace=("test",), key="key1"),
+        GetOp(namespace=("test",), key="key2"),
+        GetOp(namespace=("test",), key="key3"),
+    ]
+
+    results = store.batch(ops)
+
+    assert len(results) == 3
+    assert results[0] is not None
+    assert results[1] is not None
+    assert results[2] is None
+    assert results[0].key == "key1"
+    assert results[1].key == "key2"
+
+
+def test_batch_put_ops(store: DuckDBStore) -> None:
+    mock_connection = store.conn
+    mock_cursor = MockCursor([])
+    mock_connection.cursor.return_value = mock_cursor
+
+    ops = [
+        PutOp(namespace=("test",), key="key1", value={"data": "value1"}),
+        PutOp(namespace=("test",), key="key2", value={"data": "value2"}),
+        PutOp(namespace=("test",), key="key3", value=None),
+    ]
+
+    results = store.batch(ops)
+
+    assert len(results) == 3
+    assert all(result is None for result in results)
+    assert mock_cursor.execute.call_count == 2
+
+
+def test_batch_search_ops(store: DuckDBStore) -> None:
+    mock_connection = store.conn
+    mock_cursor = MockCursor(
+        [
+            (
+                "test.foo",
+                "key1",
+                '{"data": "value1"}',
+                datetime.now(),
+                datetime.now(),
+            ),
+            (
+                "test.bar",
+                "key2",
+                '{"data": "value2"}',
+                datetime.now(),
+                datetime.now(),
+            ),
+        ]
+    )
+    mock_connection.cursor.return_value = mock_cursor
+
+    ops = [
+        SearchOp(
+            namespace_prefix=("test",), filter={"data": "value1"}, limit=10, offset=0
+        ),
+        SearchOp(namespace_prefix=("test",), filter=None, limit=5, offset=0),
+    ]
+
+    results = store.batch(ops)
+
+    assert len(results) == 2
+    assert len(results[0]) == 2
+    assert len(results[1]) == 2
+
+
+def test_batch_list_namespaces_ops(store: DuckDBStore) -> None:
+    mock_connection = store.conn
+    mock_cursor = MockCursor([("test.namespace1",), ("test.namespace2",)])
+    mock_connection.cursor.return_value = mock_cursor
+
+    ops = [ListNamespacesOp(match_conditions=None, max_depth=None, limit=10, offset=0)]
+
+    results = store.batch(ops)
+
+    assert len(results) == 1
+    assert results[0] == [("test", "namespace1"), ("test", "namespace2")]
+
+
+def test_basic_store_ops() -> None:
+    with DuckDBStore.from_conn_string(":memory:") as store:
+        store.setup()
+        namespace = ("test", "documents")
+        item_id = "doc1"
+        item_value = {"title": "Test Document", "content": "Hello, World!"}
+
+        store.put(namespace, item_id, item_value)
+        item = store.get(namespace, item_id)
+
+        assert item
+        assert item.namespace == namespace
+        assert item.key == item_id
+        assert item.value == item_value
+
+        updated_value = {
+            "title": "Updated Test Document",
+            "content": "Hello, LangGraph!",
+        }
+        store.put(namespace, item_id, updated_value)
+        updated_item = store.get(namespace, item_id)
+
+        assert updated_item.value == updated_value
+        assert updated_item.updated_at > item.updated_at
+        different_namespace = ("test", "other_documents")
+        item_in_different_namespace = store.get(different_namespace, item_id)
+        assert item_in_different_namespace is None
+
+        new_item_id = "doc2"
+        new_item_value = {"title": "Another Document", "content": "Greetings!"}
+        store.put(namespace, new_item_id, new_item_value)
+
+        search_results = store.search(["test"], limit=10)
+        items = search_results
+        assert len(items) == 2
+        assert any(item.key == item_id for item in items)
+        assert any(item.key == new_item_id for item in items)
+
+        namespaces = store.list_namespaces(prefix=["test"])
+        assert ("test", "documents") in namespaces
+
+        store.delete(namespace, item_id)
+        store.delete(namespace, new_item_id)
+        deleted_item = store.get(namespace, item_id)
+        assert deleted_item is None
+
+        deleted_item = store.get(namespace, new_item_id)
+        assert deleted_item is None
+
+        empty_search_results = store.search(["test"], limit=10)
+        assert len(empty_search_results) == 0
+
+
+def test_list_namespaces() -> None:
+    with DuckDBStore.from_conn_string(":memory:") as store:
+        store.setup()
+        test_pref = str(uuid.uuid4())
+        test_namespaces = [
+            (test_pref, "test", "documents", "public", test_pref),
+            (test_pref, "test", "documents", "private", test_pref),
+            (test_pref, "test", "images", "public", test_pref),
+            (test_pref, "test", "images", "private", test_pref),
+            (test_pref, "prod", "documents", "public", test_pref),
+            (
+                test_pref,
+                "prod",
+                "documents",
+                "some",
+                "nesting",
+                "public",
+                test_pref,
+            ),
+            (test_pref, "prod", "documents", "private", test_pref),
+        ]
+
+        for namespace in test_namespaces:
+            store.put(namespace, "dummy", {"content": "dummy"})
+
+        prefix_result = store.list_namespaces(prefix=[test_pref, "test"])
+        assert len(prefix_result) == 4
+        assert all([ns[1] == "test" for ns in prefix_result])
+
+        specific_prefix_result = store.list_namespaces(
+            prefix=[test_pref, "test", "documents"]
+        )
+        assert len(specific_prefix_result) == 2
+        assert all([ns[1:3] == ("test", "documents") for ns in specific_prefix_result])
+
+        suffix_result = store.list_namespaces(suffix=["public", test_pref])
+        assert len(suffix_result) == 4
+        assert all(ns[-2] == "public" for ns in suffix_result)
+
+        prefix_suffix_result = store.list_namespaces(
+            prefix=[test_pref, "test"], suffix=["public", test_pref]
+        )
+        assert len(prefix_suffix_result) == 2
+        assert all(
+            ns[1] == "test" and ns[-2] == "public" for ns in prefix_suffix_result
+        )
+
+        wildcard_prefix_result = store.list_namespaces(
+            prefix=[test_pref, "*", "documents"]
+        )
+        assert len(wildcard_prefix_result) == 5
+        assert all(ns[2] == "documents" for ns in wildcard_prefix_result)
+
+        wildcard_suffix_result = store.list_namespaces(
+            suffix=["*", "public", test_pref]
+        )
+        assert len(wildcard_suffix_result) == 4
+        assert all(ns[-2] == "public" for ns in wildcard_suffix_result)
+        wildcard_single = store.list_namespaces(
+            suffix=["some", "*", "public", test_pref]
+        )
+        assert len(wildcard_single) == 1
+        assert wildcard_single[0] == (
+            test_pref,
+            "prod",
+            "documents",
+            "some",
+            "nesting",
+            "public",
+            test_pref,
+        )
+
+        max_depth_result = store.list_namespaces(max_depth=3)
+        assert all([len(ns) <= 3 for ns in max_depth_result])
+
+        max_depth_result = store.list_namespaces(
+            max_depth=4, prefix=[test_pref, "*", "documents"]
+        )
+        assert (
+            len(set(tuple(res) for res in max_depth_result))
+            == len(max_depth_result)
+            == 5
+        )
+
+        limit_result = store.list_namespaces(prefix=[test_pref], limit=3)
+        assert len(limit_result) == 3
+
+        offset_result = store.list_namespaces(prefix=[test_pref], offset=3)
+        assert len(offset_result) == len(test_namespaces) - 3
+
+        empty_prefix_result = store.list_namespaces(prefix=[test_pref])
+        assert len(empty_prefix_result) == len(test_namespaces)
+        assert set(tuple(ns) for ns in empty_prefix_result) == set(
+            tuple(ns) for ns in test_namespaces
+        )
+
+        for namespace in test_namespaces:
+            store.delete(namespace, "dummy")
+
+
+def test_search():
+    with DuckDBStore.from_conn_string(":memory:") as store:
+        store.setup()
+        test_namespaces = [
+            ("test_search", "documents", "user1"),
+            ("test_search", "documents", "user2"),
+            ("test_search", "reports", "department1"),
+            ("test_search", "reports", "department2"),
+        ]
+        test_items = [
+            {"title": "Doc 1", "author": "John Doe", "tags": ["important"]},
+            {"title": "Doc 2", "author": "Jane Smith", "tags": ["draft"]},
+            {"title": "Report A", "author": "John Doe", "tags": ["final"]},
+            {"title": "Report B", "author": "Alice Johnson", "tags": ["draft"]},
+        ]
+
+        for namespace, item in zip(test_namespaces, test_items):
+            store.put(namespace, f"item_{namespace[-1]}", item)
+
+        docs_result = store.search(["test_search", "documents"])
+        assert len(docs_result) == 2
+        assert all(
+            [item.namespace[1] == "documents" for item in docs_result]
+        ), docs_result
+
+        reports_result = store.search(["test_search", "reports"])
+        assert len(reports_result) == 2
+        assert all(item.namespace[1] == "reports" for item in reports_result)
+
+        limited_result = store.search(["test_search"], limit=2)
+        assert len(limited_result) == 2
+        offset_result = store.search(["test_search"])
+        assert len(offset_result) == 4
+
+        offset_result = store.search(["test_search"], offset=2)
+        assert len(offset_result) == 2
+        assert all(item not in limited_result for item in offset_result)
+
+        john_doe_result = store.search(["test_search"], filter={"author": "John Doe"})
+        assert len(john_doe_result) == 2
+        assert all(item.value["author"] == "John Doe" for item in john_doe_result)
+
+        draft_result = store.search(["test_search"], filter={"tags": ["draft"]})
+        assert len(draft_result) == 2
+        assert all("draft" in item.value["tags"] for item in draft_result)
+
+        page1 = store.search(["test_search"], limit=2, offset=0)
+        page2 = store.search(["test_search"], limit=2, offset=2)
+        all_items = page1 + page2
+        assert len(all_items) == 4
+        assert len(set(item.key for item in all_items)) == 4
+
+        for namespace in test_namespaces:
+            store.delete(namespace, f"item_{namespace[-1]}")

libs/checkpoint-duckdb/tests/test_sync.py

@@ -0,0 +1,111 @@
+from typing import Any
+
+import pytest
+from langchain_core.runnables import RunnableConfig
+
+from langgraph.checkpoint.base import (
+    Checkpoint,
+    CheckpointMetadata,
+    create_checkpoint,
+    empty_checkpoint,
+)
+from langgraph.checkpoint.duckdb import DuckDBSaver
+
+
+class TestDuckDBSaver:
+    @pytest.fixture(autouse=True)
+    def setup(self) -> None:
+        # objects for test setup
+        self.config_1: RunnableConfig = {
+            "configurable": {
+                "thread_id": "thread-1",
+                # for backwards compatibility testing
+                "thread_ts": "1",
+                "checkpoint_ns": "",
+            }
+        }
+        self.config_2: RunnableConfig = {
+            "configurable": {
+                "thread_id": "thread-2",
+                "checkpoint_id": "2",
+                "checkpoint_ns": "",
+            }
+        }
+        self.config_3: RunnableConfig = {
+            "configurable": {
+                "thread_id": "thread-2",
+                "checkpoint_id": "2-inner",
+                "checkpoint_ns": "inner",
+            }
+        }
+
+        self.chkpnt_1: Checkpoint = empty_checkpoint()
+        self.chkpnt_2: Checkpoint = create_checkpoint(self.chkpnt_1, {}, 1)
+        self.chkpnt_3: Checkpoint = empty_checkpoint()
+
+        self.metadata_1: CheckpointMetadata = {
+            "source": "input",
+            "step": 2,
+            "writes": {},
+            "score": 1,
+        }
+        self.metadata_2: CheckpointMetadata = {
+            "source": "loop",
+            "step": 1,
+            "writes": {"foo": "bar"},
+            "score": None,
+        }
+        self.metadata_3: CheckpointMetadata = {}
+
+    def test_search(self) -> None:
+        with DuckDBSaver.from_conn_string(":memory:") as saver:
+            saver.setup()
+            # save checkpoints
+            saver.put(self.config_1, self.chkpnt_1, self.metadata_1, {})
+            saver.put(self.config_2, self.chkpnt_2, self.metadata_2, {})
+            saver.put(self.config_3, self.chkpnt_3, self.metadata_3, {})
+
+            # call method / assertions
+            query_1 = {"source": "input"}  # search by 1 key
+            query_2 = {
+                "step": 1,
+                "writes": {"foo": "bar"},
+            }  # search by multiple keys
+            query_3: dict[str, Any] = {}  # search by no keys, return all checkpoints
+            query_4 = {"source": "update", "step": 1}  # no match
+
+            search_results_1 = list(saver.list(None, filter=query_1))
+            assert len(search_results_1) == 1
+            assert search_results_1[0].metadata == self.metadata_1
+
+            search_results_2 = list(saver.list(None, filter=query_2))
+            assert len(search_results_2) == 1
+            assert search_results_2[0].metadata == self.metadata_2
+
+            search_results_3 = list(saver.list(None, filter=query_3))
+            assert len(search_results_3) == 3
+
+            search_results_4 = list(saver.list(None, filter=query_4))
+            assert len(search_results_4) == 0
+
+            # search by config (defaults to checkpoints across all namespaces)
+            search_results_5 = list(
+                saver.list({"configurable": {"thread_id": "thread-2"}})
+            )
+            assert len(search_results_5) == 2
+            assert {
+                search_results_5[0].config["configurable"]["checkpoint_ns"],
+                search_results_5[1].config["configurable"]["checkpoint_ns"],
+            } == {"", "inner"}
+
+            # TODO: test before and limit params
+
+    def test_null_chars(self) -> None:
+        with DuckDBSaver.from_conn_string(":memory:") as saver:
+            saver.setup()
+            config = saver.put(self.config_1, self.chkpnt_1, {"my_key": "\x00abc"}, {})
+            assert saver.get_tuple(config).metadata["my_key"] == "abc"  # type: ignore
+            assert (
+                list(saver.list(None, filter={"my_key": "abc"}))[0].metadata["my_key"]  # type: ignore
+                == "abc"
+            )

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

@@ -287,7 +287,7 @@ class PostgresSaver(BasePostgresSaver):
             >>> DB_URI = "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable"
             >>> with PostgresSaver.from_conn_string(DB_URI) as memory:
             >>>     config = {"configurable": {"thread_id": "1", "checkpoint_ns": ""}}
-            >>>     checkpoint = {"ts": "2024-05-04T06:32:42.235444+00:00", "id": "1ef4f797-8335-6428-8001-8a1503f9b875", "data": {"key": "value"}}
+            >>>     checkpoint = {"ts": "2024-05-04T06:32:42.235444+00:00", "id": "1ef4f797-8335-6428-8001-8a1503f9b875", "channel_values": {"key": "value"}}
             >>>     saved_config = memory.put(config, checkpoint, {"source": "input", "step": 1, "writes": {"key": "value"}}, {})
             >>> print(saved_config)
             {'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef4f797-8335-6428-8001-8a1503f9b875'}}

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

@@ -249,7 +249,7 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
         filter: MetadataInput,
         before: Optional[RunnableConfig] = None,
     ) -> Tuple[str, List[Any]]:
-        """Return WHERE clause predicates for alist() given config, filter, cursor.
+        """Return WHERE clause predicates for alist() given config, filter, before.
 
         This method returns a tuple of a string and a tuple of values. The string
         is the parametered WHERE clause predicate (including the WHERE keyword):

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

@@ -44,7 +44,6 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[AsyncConnectio
         super().__init__()
         self._deserializer = deserializer
         self.conn = conn
-        self.conn = conn
         self.loop = asyncio.get_running_loop()
 
     async def abatch(self, ops: Iterable[Op]) -> list[Result]:

libs/checkpoint-postgres/poetry.lock

@@ -1,4 +1,4 @@
-# This file is automatically @generated by Poetry 1.8.2 and should not be changed by hand.
+# This file is automatically @generated by Poetry 1.8.3 and should not be changed by hand.
 
 [[package]]
 name = "annotated-types"
@@ -324,7 +324,7 @@ typing-extensions = ">=4.7"
 
 [[package]]
 name = "langgraph-checkpoint"
-version = "2.0.0"
+version = "2.0.2"
 description = "Library with base interfaces for LangGraph checkpoint savers."
 optional = false
 python-versions = "^3.9.0,<4.0"
@@ -1116,4 +1116,4 @@ watchmedo = ["PyYAML (>=3.10)"]
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.9.0,<4.0"
-content-hash = "8f763cd1727287f8c8b5ad2b4d8df00fb446e68d0cd4e88c278e4007969b83fd"
+content-hash = "6bd85ce8ee1192995c1ff03d5fa65af8ee7872214d71b84559a6192cadf82be6"

libs/checkpoint-postgres/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph-checkpoint-postgres"
-version = "2.0.1"
+version = "2.0.2"
 description = "Library with a Postgres implementation of LangGraph checkpoint saver."
 authors = []
 license = "MIT"
@@ -10,7 +10,7 @@ packages = [{ include = "langgraph" }]
 
 [tool.poetry.dependencies]
 python = "^3.9.0,<4.0"
-langgraph-checkpoint = "^2.0.0"
+langgraph-checkpoint = "^2.0.2"
 orjson = ">=3.10.1"
 psycopg = "^3.0.0"
 psycopg-pool = "^3.0.0"

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

@@ -389,7 +389,7 @@ class SqliteSaver(BaseCheckpointSaver[str]):
             >>> from langgraph.checkpoint.sqlite import SqliteSaver
             >>> with SqliteSaver.from_conn_string(":memory:") as memory:
             >>>     config = {"configurable": {"thread_id": "1", "checkpoint_ns": ""}}
-            >>>     checkpoint = {"ts": "2024-05-04T06:32:42.235444+00:00", "id": "1ef4f797-8335-6428-8001-8a1503f9b875", "data": {"key": "value"}}
+            >>>     checkpoint = {"ts": "2024-05-04T06:32:42.235444+00:00", "id": "1ef4f797-8335-6428-8001-8a1503f9b875", "channel_values": {"key": "value"}}
             >>>     saved_config = memory.put(config, checkpoint, {"source": "input", "step": 1, "writes": {"key": "value"}}, {})
             >>> print(saved_config)
             {'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef4f797-8335-6428-8001-8a1503f9b875'}}

libs/checkpoint-sqlite/poetry.lock

@@ -332,7 +332,7 @@ typing-extensions = ">=4.7"
 
 [[package]]
 name = "langgraph-checkpoint"
-version = "2.0.0"
+version = "2.0.2"
 description = "Library with base interfaces for LangGraph checkpoint savers."
 optional = false
 python-versions = "^3.9.0,<4.0"
@@ -1001,4 +1001,4 @@ watchmedo = ["PyYAML (>=3.10)"]
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.9.0"
-content-hash = "e0091cc2deab4de99a6bc4eb262b0040b771a9659dd3638ac1c4a225a1f11dc2"
+content-hash = "927b49b9ba72a301980237d7adc2e73cdacfbe127a174c7488136a9af9372796"

libs/checkpoint-sqlite/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph-checkpoint-sqlite"
-version = "2.0.0"
+version = "2.0.1"
 description = "Library with a SQLite implementation of LangGraph checkpoint saver."
 authors = []
 license = "MIT"
@@ -10,7 +10,7 @@ packages = [{ include = "langgraph" }]
 
 [tool.poetry.dependencies]
 python = "^3.9.0"
-langgraph-checkpoint = "^2.0.0"
+langgraph-checkpoint = "^2.0.2"
 aiosqlite = "^0.20.0"
 
 [tool.poetry.group.dev.dependencies]

libs/checkpoint/pyproject.toml

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

libs/langgraph/langgraph/graph/graph.py

@@ -514,6 +514,14 @@ class CompiledGraph(Pregel):
                 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:
+        return self.get_graph(config, xray=xray)
+
     def get_graph(
         self,
         config: Optional[RunnableConfig] = None,

libs/langgraph/langgraph/graph/state.py

@@ -3,7 +3,7 @@ import logging
 import typing
 import warnings
 from functools import partial
-from inspect import isclass, isfunction, signature
+from inspect import isclass, isfunction, ismethod, signature
 from typing import (
     Any,
     Callable,
@@ -338,19 +338,8 @@ class StateGraph(Graph):
                     f"'{character}' is a reserved character and is not allowed in the node names."
                 )
 
-        try:
-            if isfunction(action) and (
-                hints := get_type_hints(action.__call__) or get_type_hints(action)
-            ):
-                if input is None:
-                    first_parameter_name = next(
-                        iter(inspect.signature(action).parameters.keys())
-                    )
-                    if input_hint := hints.get(first_parameter_name):
-                        if isinstance(input_hint, type) and get_type_hints(input_hint):
-                            input = input_hint
-        except (TypeError, StopIteration):
-            pass
+        if input is None:
+            input = _get_input_schema_from_type_hint(action)
         if input is not None:
             self._add_schema(input)
         self.nodes[cast(str, node)] = StateNodeSpec(
@@ -834,3 +823,21 @@ def _get_schema(
                     if k in channels and isinstance(channels[k], BaseChannel)
                 },
             )
+
+
+def _get_input_schema_from_type_hint(
+    action: Optional[RunnableLike],
+) -> Optional[Type[Any]]:
+    if not isfunction(action) and not ismethod(getattr(action, "__call__", None)):
+        return None
+    action = cast(Callable, action)
+
+    try:
+        hints = get_type_hints(getattr(action, "__call__")) or get_type_hints(action)
+        first_parameter_name = next(iter(inspect.signature(action).parameters.keys()))
+        input_hint = hints.get(first_parameter_name)
+        if isinstance(input_hint, type) and get_type_hints(input_hint):
+            return input_hint
+    except (TypeError, StopIteration):
+        pass
+    return None

libs/langgraph/langgraph/prebuilt/tool_node.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import inspect
 import json
 from copy import copy
 from typing import (
@@ -16,6 +17,7 @@ from typing import (
     Type,
     Union,
     cast,
+    get_type_hints,
 )
 
 from langchain_core.messages import (
@@ -34,6 +36,7 @@ from langchain_core.tools import BaseTool, InjectedToolArg
 from langchain_core.tools import tool as create_tool
 from typing_extensions import Annotated, get_args, get_origin
 
+from langgraph.errors import GraphInterrupt
 from langgraph.store.base import BaseStore
 from langgraph.utils.runnable import RunnableCallable
 
@@ -67,13 +70,96 @@ def msg_content_output(output: Any) -> str | List[dict]:
             return str(output)
 
 
+def _handle_tool_error(
+    e: Exception,
+    *,
+    flag: Union[
+        bool,
+        str,
+        Callable[..., str],
+        tuple[type[Exception], ...],
+    ],
+) -> str:
+    if isinstance(flag, (bool, tuple)):
+        content = TOOL_CALL_ERROR_TEMPLATE.format(error=repr(e))
+    elif isinstance(flag, str):
+        content = flag
+    elif callable(flag):
+        content = flag(e)
+    else:
+        raise ValueError(
+            f"Got unexpected type of `handle_tool_error`. Expected bool, str "
+            f"or callable. Received: {flag}"
+        )
+    return content
+
+
+def _infer_handled_types(handler: Callable[..., str]) -> tuple[type[Exception]]:
+    sig = inspect.signature(handler)
+    params = list(sig.parameters.values())
+    if params:
+        # If it's a method, the first argument is typically 'self' or 'cls'
+        if params[0].name in ["self", "cls"] and len(params) == 2:
+            first_param = params[1]
+        else:
+            first_param = params[0]
+
+        type_hints = get_type_hints(handler)
+        if first_param.name in type_hints:
+            origin = get_origin(first_param.annotation)
+            if origin is Union:
+                args = get_args(first_param.annotation)
+                if all(issubclass(arg, Exception) for arg in args):
+                    return tuple(args)
+                else:
+                    raise ValueError(
+                        "All types in the error handler error annotation must be Exception types. "
+                        "For example, `def custom_handler(e: Union[ValueError, TypeError])`. "
+                        f"Got '{first_param.annotation}' instead."
+                    )
+
+            exception_type = type_hints[first_param.name]
+            if Exception in exception_type.__mro__:
+                return (exception_type,)
+            else:
+                raise ValueError(
+                    f"Arbitrary types are not supported in the error handler signature. "
+                    "Please annotate the error with either a specific Exception type or a union of Exception types. "
+                    "For example, `def custom_handler(e: ValueError)` or `def custom_handler(e: Union[ValueError, TypeError])`. "
+                    f"Got '{exception_type}' instead."
+                )
+
+    # If no type information is available, return (Exception,) for backwards compatibility.
+    return (Exception,)
+
+
 class ToolNode(RunnableCallable):
     """A node that runs the tools called in the last AIMessage.
 
-    It can be used either in StateGraph with a "messages" key or in MessageGraph. If
-    multiple tool calls are requested, they will be run in parallel. The output will be
+    It can be used either in StateGraph with a "messages" state key (or a custom key passed via ToolNode's 'messages_key').
+    If multiple tool calls are requested, they will be run in parallel. The output will be
     a list of ToolMessages, one for each tool call.
 
+    Args:
+        tools: A sequence of tools that can be invoked by the ToolNode.
+        name: The name of the ToolNode in the graph. Defaults to "tools".
+        tags: Optional tags to associate with the node. Defaults to None.
+        handle_tool_errors: How to handle tool errors raised by tools inside the node. Defaults to True.
+            Must be one of the following:
+
+            - True: all errors will be caught and
+                a ToolMessage with a default error message (TOOL_CALL_ERROR_TEMPLATE) will be returned.
+            - str: all errors will be caught and
+                a ToolMessage with the string value of 'handle_tool_errors' will be returned.
+            - tuple[type[Exception], ...]: exceptions in the tuple will be caught and
+                a ToolMessage with a default error message (TOOL_CALL_ERROR_TEMPLATE) will be returned.
+            - Callable[..., str]: exceptions from the signature of the callable will be caught and
+                a ToolMessage with the string value of the result of the 'handle_tool_errors' callable will be returned.
+            - False: none of the errors raised by the tools will be caught
+        messages_key: The state key in the input that contains the list of messages.
+            The same key will be used for the output from the ToolNode.
+            Defaults to "messages".
+
     The `ToolNode` is roughly analogous to:
 
     ```python
@@ -101,13 +187,17 @@ class ToolNode(RunnableCallable):
         *,
         name: str = "tools",
         tags: Optional[list[str]] = None,
-        handle_tool_errors: Optional[bool] = True,
+        handle_tool_errors: Union[
+            bool, str, Callable[..., str], tuple[type[Exception], ...]
+        ] = True,
+        messages_key: str = "messages",
     ) -> None:
         super().__init__(self._func, self._afunc, name=name, tags=tags, trace=False)
         self.tools_by_name: Dict[str, BaseTool] = {}
         self.tool_to_state_args: Dict[str, Dict[str, Optional[str]]] = {}
         self.tool_to_store_arg: Dict[str, Optional[str]] = {}
         self.handle_tool_errors = handle_tool_errors
+        self.messages_key = messages_key
         for tool_ in tools:
             if not isinstance(tool_, BaseTool):
                 tool_ = cast(BaseTool, create_tool(tool_))
@@ -131,7 +221,7 @@ class ToolNode(RunnableCallable):
         with get_executor_for_config(config) as executor:
             outputs = [*executor.map(self._run_one, tool_calls, config_list)]
         # TypedDict, pydantic, dataclass, etc. should all be able to load from dict
-        return outputs if output_type == "list" else {"messages": outputs}
+        return outputs if output_type == "list" else {self.messages_key: outputs}
 
     def invoke(
         self, input: Input, config: Optional[RunnableConfig] = None, **kwargs: Any
@@ -163,7 +253,7 @@ class ToolNode(RunnableCallable):
             *(self._arun_one(call, config) for call in tool_calls)
         )
         # TypedDict, pydantic, dataclass, etc. should all be able to load from dict
-        return outputs if output_type == "list" else {"messages": outputs}
+        return outputs if output_type == "list" else {self.messages_key: outputs}
 
     def _run_one(self, call: ToolCall, config: RunnableConfig) -> ToolMessage:
         if invalid_tool_message := self._validate_tool_call(call):
@@ -178,15 +268,38 @@ class ToolNode(RunnableCallable):
                 Union[str, list], msg_content_output(tool_message.content)
             )
             return tool_message
+        # GraphInterrupt is a special exception that will always be raised.
+        # It can be triggered in the following scenarios:
+        # (1) a NodeInterrupt is raised inside a tool
+        # (2) a NodeInterrupt is raised inside a graph node for a graph called as a tool
+        # (3) a GraphInterrupt is raised when a subgraph is interrupted inside a graph called as a tool
+        # (2 and 3 can happen in a "supervisor w/ tools" multi-agent architecture)
+        except GraphInterrupt as e:
+            raise e
         except Exception as e:
-            if not self.handle_tool_errors:
+            if isinstance(self.handle_tool_errors, tuple):
+                handled_types: tuple = self.handle_tool_errors
+            elif callable(self.handle_tool_errors):
+                handled_types = _infer_handled_types(self.handle_tool_errors)
+            else:
+                # default behavior is catching all exceptions
+                handled_types = (Exception,)
+
+            # Unhandled
+            if not self.handle_tool_errors or not isinstance(e, handled_types):
                 raise e
-            content = TOOL_CALL_ERROR_TEMPLATE.format(error=repr(e))
-            return ToolMessage(content, name=call["name"], tool_call_id=call["id"])
+            # Handled
+            else:
+                content = _handle_tool_error(e, flag=self.handle_tool_errors)
+
+        return ToolMessage(
+            content=content, name=call["name"], tool_call_id=call["id"], status="error"
+        )
 
     async def _arun_one(self, call: ToolCall, config: RunnableConfig) -> ToolMessage:
         if invalid_tool_message := self._validate_tool_call(call):
             return invalid_tool_message
+
         try:
             input = {**call, **{"type": "tool_call"}}
             tool_message: ToolMessage = await self.tools_by_name[call["name"]].ainvoke(
@@ -196,11 +309,33 @@ class ToolNode(RunnableCallable):
                 Union[str, list], msg_content_output(tool_message.content)
             )
             return tool_message
+        # GraphInterrupt is a special exception that will always be raised.
+        # It can be triggered in the following scenarios:
+        # (1) a NodeInterrupt is raised inside a tool
+        # (2) a NodeInterrupt is raised inside a graph node for a graph called as a tool
+        # (3) a GraphInterrupt is raised when a subgraph is interrupted inside a graph called as a tool
+        # (2 and 3 can happen in a "supervisor w/ tools" multi-agent architecture)
+        except GraphInterrupt as e:
+            raise e
         except Exception as e:
-            if not self.handle_tool_errors:
+            if isinstance(self.handle_tool_errors, tuple):
+                handled_types: tuple = self.handle_tool_errors
+            elif callable(self.handle_tool_errors):
+                handled_types = _infer_handled_types(self.handle_tool_errors)
+            else:
+                # default behavior is catching all exceptions
+                handled_types = (Exception,)
+
+            # Unhandled
+            if not self.handle_tool_errors or not isinstance(e, handled_types):
                 raise e
-            content = TOOL_CALL_ERROR_TEMPLATE.format(error=repr(e))
-            return ToolMessage(content, name=call["name"], tool_call_id=call["id"])
+            # Handled
+            else:
+                content = _handle_tool_error(e, flag=self.handle_tool_errors)
+
+        return ToolMessage(
+            content=content, name=call["name"], tool_call_id=call["id"], status="error"
+        )
 
     def _parse_input(
         self,
@@ -214,10 +349,10 @@ class ToolNode(RunnableCallable):
         if isinstance(input, list):
             output_type = "list"
             message: AnyMessage = input[-1]
-        elif isinstance(input, dict) and (messages := input.get("messages", [])):
+        elif isinstance(input, dict) and (messages := input.get(self.messages_key, [])):
             output_type = "dict"
             message = messages[-1]
-        elif messages := getattr(input, "messages", None):
+        elif messages := getattr(input, self.messages_key, None):
             # Assume dataclass-like state that can coerce from dict
             output_type = "dict"
             message = messages[-1]
@@ -238,7 +373,9 @@ class ToolNode(RunnableCallable):
                 requested_tool=requested_tool,
                 available_tools=", ".join(self.tools_by_name.keys()),
             )
-            return ToolMessage(content, name=requested_tool, tool_call_id=call["id"])
+            return ToolMessage(
+                content, name=requested_tool, tool_call_id=call["id"], status="error"
+            )
         else:
             return None
 
@@ -256,10 +393,10 @@ class ToolNode(RunnableCallable):
             required_fields = list(state_args.values())
             if (
                 len(required_fields) == 1
-                and required_fields[0] == "messages"
+                and required_fields[0] == self.messages_key
                 or required_fields[0] is None
             ):
-                input = {"messages": input}
+                input = {self.messages_key: input}
             else:
                 err_msg = (
                     f"Invalid input to ToolNode. Tool {tool_call['name']} requires "
@@ -325,6 +462,7 @@ class ToolNode(RunnableCallable):
 
 def tools_condition(
     state: Union[list[AnyMessage], dict[str, Any], BaseModel],
+    messages_key: str = "messages",
 ) -> Literal["tools", "__end__"]:
     """Use in the conditional_edge to route to the ToolNode if the last message
 
@@ -377,9 +515,9 @@ def tools_condition(
     """
     if isinstance(state, list):
         ai_message = state[-1]
-    elif isinstance(state, dict) and (messages := state.get("messages", [])):
+    elif isinstance(state, dict) and (messages := state.get(messages_key, [])):
         ai_message = messages[-1]
-    elif messages := getattr(state, "messages", []):
+    elif messages := getattr(state, messages_key, []):
         ai_message = messages[-1]
     else:
         raise ValueError(f"No messages found in input state to tool_edge: {state}")

libs/langgraph/langgraph/pregel/__init__.py

@@ -25,7 +25,6 @@ from uuid import UUID, uuid5
 
 from langchain_core.globals import get_debug
 from langchain_core.runnables import (
-    Runnable,
     RunnableSequence,
 )
 from langchain_core.runnables.base import Input, Output
@@ -34,6 +33,7 @@ from langchain_core.runnables.config import (
     get_async_callback_manager_for_config,
     get_callback_manager_for_config,
 )
+from langchain_core.runnables.graph import Graph
 from langchain_core.runnables.utils import (
     ConfigurableFieldSpec,
     get_unique_config_specs,
@@ -86,6 +86,7 @@ from langgraph.pregel.io import read_channels
 from langgraph.pregel.loop import AsyncPregelLoop, StreamProtocol, SyncPregelLoop
 from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager
 from langgraph.pregel.messages import StreamMessagesHandler
+from langgraph.pregel.protocol import PregelProtocol
 from langgraph.pregel.read import PregelNode
 from langgraph.pregel.retry import RetryPolicy
 from langgraph.pregel.runner import PregelRunner
@@ -179,7 +180,7 @@ class Channel:
         )
 
 
-class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
+class Pregel(PregelProtocol):
     nodes: dict[str, PregelNode]
 
     channels: dict[str, Union[BaseChannel, ManagedValueSpec]]
@@ -259,6 +260,16 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
         if auto_validate:
             self.validate()
 
+    def get_graph(
+        self, config: RunnableConfig | None = None, *, xray: int | bool = False
+    ) -> Graph:
+        raise NotImplementedError
+
+    async def aget_graph(
+        self, config: RunnableConfig | None = None, *, xray: int | bool = False
+    ) -> Graph:
+        raise NotImplementedError
+
     def copy(self, update: dict[str, Any] | None = None) -> Self:
         attrs = {**self.__dict__, **(update or {})}
         return self.__class__(**attrs)

libs/langgraph/langgraph/pregel/executor.py

@@ -118,11 +118,17 @@ class AsyncBackgroundExecutor(AsyncContextManager):
     - re-raises the first exception from tasks with `__reraise_on_exit__=True`
       ignoring CancelledError"""
 
-    def __init__(self) -> None:
+    def __init__(self, config: RunnableConfig) -> None:
         self.context_not_supported = sys.version_info < (3, 11)
         self.tasks: dict[asyncio.Task, tuple[bool, bool]] = {}
         self.sentinel = object()
         self.loop = asyncio.get_running_loop()
+        if max_concurrency := config.get("max_concurrency"):
+            self.semaphore: Optional[asyncio.Semaphore] = asyncio.Semaphore(
+                max_concurrency
+            )
+        else:
+            self.semaphore = None
 
     def submit(  # type: ignore[valid-type]
         self,
@@ -134,6 +140,8 @@ class AsyncBackgroundExecutor(AsyncContextManager):
         **kwargs: P.kwargs,
     ) -> asyncio.Task[T]:
         coro = cast(Coroutine[None, None, T], fn(*args, **kwargs))
+        if self.semaphore:
+            coro = gated(self.semaphore, coro)
         if self.context_not_supported:
             task = self.loop.create_task(coro, name=__name__)
         else:
@@ -183,3 +191,9 @@ class AsyncBackgroundExecutor(AsyncContextManager):
                         raise exc
                 except asyncio.CancelledError:
                     pass
+
+
+async def gated(semaphore: asyncio.Semaphore, coro: Coroutine[None, None, T]) -> T:
+    """A coroutine that waits for a semaphore before running another coroutine."""
+    async with semaphore:
+        return await coro

libs/langgraph/langgraph/pregel/loop.py

@@ -847,7 +847,9 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
             else []
         )
 
-        self.submit = await self.stack.enter_async_context(AsyncBackgroundExecutor())
+        self.submit = await self.stack.enter_async_context(
+            AsyncBackgroundExecutor(self.config)
+        )
         self.channels, self.managed = await self.stack.enter_async_context(
             AsyncChannelsManager(self.specs, self.checkpoint, self)
         )

libs/langgraph/langgraph/pregel/protocol.py

@@ -1,27 +1,29 @@
+from abc import ABC, abstractmethod
 from typing import (
     Any,
     AsyncIterator,
     Iterator,
     Optional,
-    Protocol,
     Sequence,
     Union,
-    runtime_checkable,
 )
 
-from langchain_core.runnables import RunnableConfig
+from langchain_core.runnables import Runnable, RunnableConfig
 from langchain_core.runnables.graph import Graph as DrawableGraph
 from typing_extensions import Self
 
 from langgraph.pregel.types import All, StateSnapshot, StreamMode
 
 
-@runtime_checkable
-class PregelProtocol(Protocol):
+class PregelProtocol(
+    Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]], ABC
+):
+    @abstractmethod
     def with_config(
         self, config: Optional[RunnableConfig] = None, **kwargs: Any
     ) -> Self: ...
 
+    @abstractmethod
     def get_graph(
         self,
         config: Optional[RunnableConfig] = None,
@@ -29,6 +31,7 @@ class PregelProtocol(Protocol):
         xray: Union[int, bool] = False,
     ) -> DrawableGraph: ...
 
+    @abstractmethod
     async def aget_graph(
         self,
         config: Optional[RunnableConfig] = None,
@@ -36,22 +39,17 @@ class PregelProtocol(Protocol):
         xray: Union[int, bool] = False,
     ) -> DrawableGraph: ...
 
-    def get_subgraphs(
-        self, namespace: Optional[str] = None, recurse: bool = False
-    ) -> Iterator[tuple[str, "PregelProtocol"]]: ...
-
-    def aget_subgraphs(
-        self, namespace: Optional[str] = None, recurse: bool = False
-    ) -> AsyncIterator[tuple[str, "PregelProtocol"]]: ...
-
+    @abstractmethod
     def get_state(
         self, config: RunnableConfig, *, subgraphs: bool = False
     ) -> StateSnapshot: ...
 
+    @abstractmethod
     async def aget_state(
         self, config: RunnableConfig, *, subgraphs: bool = False
     ) -> StateSnapshot: ...
 
+    @abstractmethod
     def get_state_history(
         self,
         config: RunnableConfig,
@@ -61,6 +59,7 @@ class PregelProtocol(Protocol):
         limit: Optional[int] = None,
     ) -> Iterator[StateSnapshot]: ...
 
+    @abstractmethod
     def aget_state_history(
         self,
         config: RunnableConfig,
@@ -70,6 +69,7 @@ class PregelProtocol(Protocol):
         limit: Optional[int] = None,
     ) -> AsyncIterator[StateSnapshot]: ...
 
+    @abstractmethod
     def update_state(
         self,
         config: RunnableConfig,
@@ -77,6 +77,7 @@ class PregelProtocol(Protocol):
         as_node: Optional[str] = None,
     ) -> RunnableConfig: ...
 
+    @abstractmethod
     async def aupdate_state(
         self,
         config: RunnableConfig,
@@ -84,6 +85,7 @@ class PregelProtocol(Protocol):
         as_node: Optional[str] = None,
     ) -> RunnableConfig: ...
 
+    @abstractmethod
     def stream(
         self,
         input: Union[dict[str, Any], Any],
@@ -95,6 +97,7 @@ class PregelProtocol(Protocol):
         subgraphs: bool = False,
     ) -> Iterator[Union[dict[str, Any], Any]]: ...
 
+    @abstractmethod
     def astream(
         self,
         input: Union[dict[str, Any], Any],
@@ -106,6 +109,7 @@ class PregelProtocol(Protocol):
         subgraphs: bool = False,
     ) -> AsyncIterator[Union[dict[str, Any], Any]]: ...
 
+    @abstractmethod
     def invoke(
         self,
         input: Union[dict[str, Any], Any],
@@ -115,6 +119,7 @@ class PregelProtocol(Protocol):
         interrupt_after: Optional[Union[All, Sequence[str]]] = None,
     ) -> Union[dict[str, Any], Any]: ...
 
+    @abstractmethod
     async def ainvoke(
         self,
         input: Union[dict[str, Any], Any],

libs/langgraph/langgraph/pregel/remote.py

@@ -2,6 +2,7 @@ from typing import (
     Any,
     AsyncIterator,
     Iterator,
+    Literal,
     Optional,
     Sequence,
     Union,
@@ -9,7 +10,7 @@ from typing import (
 )
 
 import orjson
-from langchain_core.runnables import Runnable, RunnableConfig
+from langchain_core.runnables import RunnableConfig
 from langchain_core.runnables.graph import (
     Edge as DrawableEdge,
 )
@@ -19,7 +20,6 @@ from langchain_core.runnables.graph import (
 from langchain_core.runnables.graph import (
     Node as DrawableNode,
 )
-from langchain_core.runnables.schema import StandardStreamEvent, StreamEvent
 from langgraph_sdk.client import (
     LangGraphClient,
     SyncLangGraphClient,
@@ -27,32 +27,45 @@ from langgraph_sdk.client import (
     get_sync_client,
 )
 from langgraph_sdk.schema import Checkpoint, ThreadState
+from langgraph_sdk.schema import StreamMode as StreamModeSDK
 from typing_extensions import Self
 
 from langgraph.checkpoint.base import CheckpointMetadata
+from langgraph.constants import INTERRUPT
+from langgraph.errors import GraphInterrupt
 from langgraph.pregel.protocol import PregelProtocol
 from langgraph.pregel.types import All, PregelTask, StateSnapshot, StreamMode
 from langgraph.types import Interrupt
 from langgraph.utils.config import merge_configs
 
 
-class RemoteGraph(PregelProtocol, Runnable):
+class RemoteException(Exception):
+    """Exception raised when an error occurs in the remote graph."""
+
+    pass
+
+
+class RemoteGraph(PregelProtocol):
+    name: str
+
     def __init__(
         self,
-        graph_id: str,
-        config: Optional[RunnableConfig] = None,
+        name: str,  # graph_id
+        /,
+        *,
         url: Optional[str] = None,
         api_key: Optional[str] = None,
         headers: Optional[dict[str, str]] = None,
         client: Optional[LangGraphClient] = None,
         sync_client: Optional[SyncLangGraphClient] = None,
+        config: Optional[RunnableConfig] = None,
     ):
         """Specify `url`, `api_key`, and/or `headers` to create default sync and async clients.
 
         If `client` or `sync_client` are provided, they will be used instead of the default clients.
         See `LangGraphClient` and `SyncLangGraphClient` for details on the default clients.
         """
-        self.graph_id = graph_id
+        self.name = name
         self.config = config
         self.client = client or get_client(url=url, api_key=api_key, headers=headers)
         self.sync_client = sync_client or get_sync_client(
@@ -61,7 +74,7 @@ class RemoteGraph(PregelProtocol, Runnable):
 
     def copy(self, update: dict[str, Any]) -> Self:
         attrs = {**self.__dict__, **update}
-        return self.__class__(**attrs)
+        return self.__class__(attrs.pop("name"), **attrs)
 
     def with_config(
         self, config: Optional[RunnableConfig] = None, **kwargs: Any
@@ -91,7 +104,7 @@ class RemoteGraph(PregelProtocol, Runnable):
         xray: Union[int, bool] = False,
     ) -> DrawableGraph:
         graph = self.sync_client.assistants.get_graph(
-            assistant_id=self.graph_id,
+            assistant_id=self.name,
             xray=xray,
         )
         return DrawableGraph(
@@ -106,7 +119,7 @@ class RemoteGraph(PregelProtocol, Runnable):
         xray: Union[int, bool] = False,
     ) -> DrawableGraph:
         graph = await self.client.assistants.get_graph(
-            assistant_id=self.graph_id,
+            assistant_id=self.name,
             xray=xray,
         )
         return DrawableGraph(
@@ -114,30 +127,6 @@ class RemoteGraph(PregelProtocol, Runnable):
             edges=[DrawableEdge(**edge) for edge in graph["edges"]],
         )
 
-    def get_subgraphs(
-        self, namespace: Optional[str] = None, recurse: bool = False
-    ) -> Iterator[tuple[str, "PregelProtocol"]]:
-        subgraphs = self.sync_client.assistants.get_subgraphs(
-            assistant_id=self.graph_id,
-            namespace=namespace,
-            recurse=recurse,
-        )
-        for namespace, graph_schema in subgraphs.items():
-            remote_subgraph = self.copy({"graph_id": graph_schema["graph_id"]})
-            yield (namespace, remote_subgraph)
-
-    async def aget_subgraphs(
-        self, namespace: Optional[str] = None, recurse: bool = False
-    ) -> AsyncIterator[tuple[str, "PregelProtocol"]]:
-        subgraphs = await self.client.assistants.get_subgraphs(
-            assistant_id=self.graph_id,
-            namespace=namespace,
-            recurse=recurse,
-        )
-        for namespace, graph_schema in subgraphs.items():
-            remote_subgraph = self.copy({"graph_id": graph_schema["graph_id"]})
-            yield (namespace, remote_subgraph)
-
     def _create_state_snapshot(self, state: ThreadState) -> StateSnapshot:
         tasks = []
         for task in state["tasks"]:
@@ -250,7 +239,11 @@ class RemoteGraph(PregelProtocol, Runnable):
             if k not in reserved_configurable_keys and not k.startswith("__pregel_")
         }
 
-        return {"configurable": new_configurable}
+        return {
+            "tags": config.get("tags") or [],
+            "metadata": config.get("metadata") or {},
+            "configurable": new_configurable,
+        }
 
     def get_state(
         self, config: RunnableConfig, *, subgraphs: bool = False
@@ -348,6 +341,37 @@ class RemoteGraph(PregelProtocol, Runnable):
         )
         return self._get_config(response["checkpoint"])
 
+    def _get_stream_modes(
+        self,
+        stream_mode: Optional[Union[StreamMode, list[StreamMode]]],
+        default: StreamMode = "updates",
+    ) -> tuple[list[StreamModeSDK], bool, bool]:
+        """Return a tuple of the final list of stream modes sent to the
+        remote graph and a boolean flag indicating if stream mode 'updates'
+        was present in the original list of stream modes.
+
+        'updates' mode is added to the list of stream modes so that interrupts
+        can be detected in the remote graph.
+        """
+        updated_stream_modes: list[StreamMode] = []
+        req_updates = False
+        req_single = True
+        # coerce to list, or add default stream mode
+        if stream_mode:
+            if isinstance(stream_mode, str):
+                updated_stream_modes.append(stream_mode)
+            else:
+                req_single = False
+                updated_stream_modes.extend(stream_mode)
+        else:
+            updated_stream_modes.append(default)
+        # add 'updates' mode if not present
+        if "updates" in updated_stream_modes:
+            req_updates = True
+        else:
+            updated_stream_modes.append("updates")
+        return (updated_stream_modes, req_updates, req_single)
+
     def stream(
         self,
         input: Union[dict[str, Any], Any],
@@ -360,18 +384,40 @@ class RemoteGraph(PregelProtocol, Runnable):
     ) -> Iterator[Union[dict[str, Any], Any]]:
         merged_config = merge_configs(self.config, config)
         sanitized_config = self._sanitize_config(merged_config)
+        stream_modes, req_updates, req_single = self._get_stream_modes(stream_mode)
 
         for chunk in self.sync_client.runs.stream(
-            thread_id=sanitized_config["configurable"]["thread_id"],
-            assistant_id=self.graph_id,
+            thread_id=sanitized_config["configurable"].get("thread_id"),
+            assistant_id=self.name,
             input=input,
             config=sanitized_config,
-            stream_mode=stream_mode,  # type: ignore
-            interrupt_before=interrupt_before,  # type: ignore
-            interrupt_after=interrupt_after,  # type: ignore
+            stream_mode=stream_modes,
+            interrupt_before=interrupt_before,
+            interrupt_after=interrupt_after,
             stream_subgraphs=subgraphs,
+            if_not_exists="create",
         ):
-            yield chunk
+            if chunk.event.startswith("updates"):
+                if isinstance(chunk.data, dict) and INTERRUPT in chunk.data:
+                    raise GraphInterrupt(chunk.data[INTERRUPT])
+                if not req_updates:
+                    continue
+            elif chunk.event.startswith("error"):
+                raise RemoteException(chunk.data)
+            if subgraphs:
+                if "|" in chunk.event:
+                    mode, ns_ = chunk.event.split("|", 1)
+                    ns = tuple(ns_.split("|"))
+                else:
+                    mode, ns = chunk.event, ()
+                if req_single:
+                    yield ns, chunk.data
+                else:
+                    yield ns, mode, chunk.data
+            elif req_single:
+                yield chunk.data
+            else:
+                yield chunk
 
     async def astream(
         self,
@@ -385,47 +431,56 @@ class RemoteGraph(PregelProtocol, Runnable):
     ) -> AsyncIterator[Union[dict[str, Any], Any]]:
         merged_config = merge_configs(self.config, config)
         sanitized_config = self._sanitize_config(merged_config)
+        stream_modes, req_updates, req_single = self._get_stream_modes(stream_mode)
 
         async for chunk in self.client.runs.stream(
-            thread_id=sanitized_config["configurable"]["thread_id"],
-            assistant_id=self.graph_id,
+            thread_id=sanitized_config["configurable"].get("thread_id"),
+            assistant_id=self.name,
             input=input,
             config=sanitized_config,
-            stream_mode=stream_mode if stream_mode else "values",  # type: ignore
-            interrupt_before=interrupt_before,  # type: ignore
-            interrupt_after=interrupt_after,  # type: ignore
+            stream_mode=stream_modes,
+            interrupt_before=interrupt_before,
+            interrupt_after=interrupt_after,
             stream_subgraphs=subgraphs,
+            if_not_exists="create",
         ):
-            yield chunk
+            if chunk.event.startswith("updates"):
+                if isinstance(chunk.data, dict) and INTERRUPT in chunk.data:
+                    raise GraphInterrupt(chunk.data[INTERRUPT])
+                if not req_updates:
+                    continue
+            elif chunk.event.startswith("error"):
+                raise RemoteException(chunk.data)
+            if subgraphs:
+                if "|" in chunk.event:
+                    mode, ns_ = chunk.event.split("|", 1)
+                    ns = tuple(ns_.split("|"))
+                else:
+                    mode, ns = chunk.event, ()
+                if req_single:
+                    yield ns, chunk.data
+                else:
+                    yield ns, mode, chunk.data
+            elif req_single:
+                yield chunk.data
+            else:
+                yield chunk
 
     async def astream_events(
         self,
         input: Any,
         config: Optional[RunnableConfig] = None,
+        *,
+        version: Literal["v1", "v2"],
+        include_names: Optional[Sequence[All]] = None,
+        include_types: Optional[Sequence[All]] = None,
+        include_tags: Optional[Sequence[All]] = None,
+        exclude_names: Optional[Sequence[All]] = None,
+        exclude_types: Optional[Sequence[All]] = None,
+        exclude_tags: Optional[Sequence[All]] = None,
         **kwargs: Any,
-    ) -> AsyncIterator[StreamEvent]:
-        merged_config = merge_configs(self.config, config)
-        sanitized_config = self._sanitize_config(merged_config)
-
-        # manually add 'events' to stream modes list
-        stream_mode: list[str] = kwargs.get("stream_mode", [])
-        if "events" not in stream_mode:
-            stream_mode.append("events")
-
-        async for chunk in self.client.runs.stream(
-            thread_id=sanitized_config["configurable"]["thread_id"],
-            assistant_id=self.graph_id,
-            input=input,
-            config=sanitized_config,
-            stream_mode=stream_mode,  # type: ignore
-            interrupt_before=kwargs.get("interrupt_before"),
-            interrupt_after=kwargs.get("interrupt_after"),
-            stream_subgraphs=kwargs.get("subgraphs", False),
-        ):
-            yield StandardStreamEvent(
-                event=chunk.event,
-                data=chunk.data,
-            )
+    ) -> AsyncIterator[dict[str, Any]]:
+        raise NotImplementedError
 
     def invoke(
         self,
@@ -439,12 +494,13 @@ class RemoteGraph(PregelProtocol, Runnable):
         sanitized_config = self._sanitize_config(merged_config)
 
         return self.sync_client.runs.wait(
-            thread_id=sanitized_config["configurable"]["thread_id"],
-            assistant_id=self.graph_id,
+            thread_id=sanitized_config["configurable"].get("thread_id"),
+            assistant_id=self.name,
             input=input,
             config=sanitized_config,
-            interrupt_before=interrupt_before,  # type: ignore
-            interrupt_after=interrupt_after,  # type: ignore
+            interrupt_before=interrupt_before,
+            interrupt_after=interrupt_after,
+            if_not_exists="create",
         )
 
     async def ainvoke(
@@ -459,10 +515,11 @@ class RemoteGraph(PregelProtocol, Runnable):
         sanitized_config = self._sanitize_config(merged_config)
 
         return await self.client.runs.wait(
-            thread_id=sanitized_config["configurable"]["thread_id"],
-            assistant_id=self.graph_id,
+            thread_id=sanitized_config["configurable"].get("thread_id"),
+            assistant_id=self.name,
             input=input,
             config=sanitized_config,
-            interrupt_before=interrupt_before,  # type: ignore
-            interrupt_after=interrupt_after,  # type: ignore
+            interrupt_before=interrupt_before,
+            interrupt_after=interrupt_after,
+            if_not_exists="create",
         )

libs/langgraph/langgraph/pregel/utils.py

@@ -4,6 +4,7 @@ from langchain_core.runnables import RunnableLambda, RunnableSequence
 from langchain_core.runnables.utils import get_function_nonlocals
 
 from langgraph.checkpoint.base import ChannelVersions
+from langgraph.pregel.protocol import PregelProtocol
 from langgraph.utils.runnable import Runnable, RunnableCallable, RunnableSeq
 
 
@@ -32,9 +33,9 @@ def find_subgraph_pregel(candidate: Runnable) -> Optional[Runnable]:
 
     for c in candidates:
         if (
-            isinstance(c, Pregel)
+            isinstance(c, PregelProtocol)
             # subgraphs that disabled checkpointing are not considered
-            and c.checkpointer is not False
+            and (not isinstance(c, Pregel) or c.checkpointer is not False)
         ):
             return c
         elif isinstance(c, RunnableSequence) or isinstance(c, RunnableSeq):

libs/langgraph/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph"
-version = "0.2.38"
+version = "0.2.39"
 description = "Building stateful, multi-actor applications with LLMs"
 authors = []
 license = "MIT"
@@ -27,11 +27,12 @@ jupyter = "^1.0.0"
 pytest-xdist = {extras = ["psutil"], version = "^3.6.1"}
 pytest-repeat = "^0.9.3"
 langgraph-checkpoint = {path = "../checkpoint", develop = true}
+langgraph-checkpoint-duckdb = {path = "../checkpoint-duckdb", develop = true}
 langgraph-checkpoint-sqlite = {path = "../checkpoint-sqlite", develop = true}
 langgraph-checkpoint-postgres = {path = "../checkpoint-postgres", develop = true}
 langgraph-sdk = {path = "../sdk-py", develop = true}
 psycopg = {extras = ["binary"], version = ">=3.0.0", python = ">=3.10"}
-uvloop = "0.21.0"
+uvloop = "0.21.0beta1"
 pyperf = "^2.7.0"
 py-spy = "^0.3.14"
 types-requests = "^2.32.0.20240914"

libs/langgraph/tests/__snapshots__/test_pregel_async.ambr

@@ -221,19 +221,6 @@
                 +---------+                
   '''
 # ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2
-  '''
-  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.1
   dict({
     '$defs': dict({
@@ -342,6 +329,127 @@
     'type': 'object',
   })
 # ---
+# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[duckdb_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[duckdb_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[duckdb_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[memory]
   '''
   graph TD;
@@ -1194,6 +1302,31 @@
    +---------+   
   '''
 # ---
+# name: test_weather_subgraph[duckdb_aio]
+  '''
+  %%{init: {'flowchart': {'curve': 'linear'}}}%%
+  graph TD;
+  	__start__([<p>__start__</p>]):::first
+  	router_node(router_node)
+  	normal_llm_node(normal_llm_node)
+  	weather_graph_model_node(model_node)
+  	weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
+  	__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 --> 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[memory]
   '''
   %%{init: {'flowchart': {'curve': 'linear'}}}%%

libs/langgraph/tests/conftest.py

@@ -11,11 +11,14 @@ from psycopg_pool import AsyncConnectionPool, ConnectionPool
 from pytest_mock import MockerFixture
 
 from langgraph.checkpoint.base import BaseCheckpointSaver
+from langgraph.checkpoint.duckdb import DuckDBSaver
+from langgraph.checkpoint.duckdb.aio import AsyncDuckDBSaver
 from langgraph.checkpoint.postgres import PostgresSaver
 from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
 from langgraph.checkpoint.sqlite import SqliteSaver
 from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
 from langgraph.store.base import BaseStore
+from langgraph.store.duckdb import AsyncDuckDBStore, DuckDBStore
 from langgraph.store.memory import InMemoryStore
 from langgraph.store.postgres import AsyncPostgresStore, PostgresStore
 from tests.memory_assert import MemorySaverAssertImmutable
@@ -61,6 +64,20 @@ async def _checkpointer_sqlite_aio():
         yield checkpointer
 
 
+@pytest.fixture(scope="function")
+def checkpointer_duckdb():
+    with DuckDBSaver.from_conn_string(":memory:") as checkpointer:
+        checkpointer.setup()
+        yield checkpointer
+
+
+@asynccontextmanager
+async def _checkpointer_duckdb_aio():
+    async with AsyncDuckDBSaver.from_conn_string(":memory:") as checkpointer:
+        await checkpointer.setup()
+        yield checkpointer
+
+
 @pytest.fixture(scope="function")
 def checkpointer_postgres():
     database = f"test_{uuid4().hex[:16]}"
@@ -212,6 +229,9 @@ async def awith_checkpointer(
     elif checkpointer_name == "sqlite_aio":
         async with _checkpointer_sqlite_aio() as checkpointer:
             yield checkpointer
+    elif checkpointer_name == "duckdb_aio":
+        async with _checkpointer_duckdb_aio() as checkpointer:
+            yield checkpointer
     elif checkpointer_name == "postgres_aio":
         async with _checkpointer_postgres_aio() as checkpointer:
             yield checkpointer
@@ -247,6 +267,13 @@ async def _store_postgres_aio():
             await conn.execute(f"DROP DATABASE {database}")
 
 
+@asynccontextmanager
+async def _store_duckdb_aio():
+    async with AsyncDuckDBStore.from_conn_string(":memory:") as store:
+        await store.setup()
+        yield store
+
+
 @pytest.fixture(scope="function")
 def store_postgres():
     database = f"test_{uuid4().hex[:16]}"
@@ -264,6 +291,13 @@ def store_postgres():
             conn.execute(f"DROP DATABASE {database}")
 
 
+@pytest.fixture(scope="function")
+def store_duckdb():
+    with DuckDBStore.from_conn_string(":memory:") as store:
+        store.setup()
+        yield store
+
+
 @pytest.fixture(scope="function")
 def store_in_memory():
     yield InMemoryStore()
@@ -278,6 +312,9 @@ async def awith_store(store_name: Optional[str]) -> AsyncIterator[BaseStore]:
     elif store_name == "postgres_aio":
         async with _store_postgres_aio() as store:
             yield store
+    elif store_name == "duckdb_aio":
+        async with _store_duckdb_aio() as store:
+            yield store
     else:
         raise NotImplementedError(f"Unknown store {store_name}")
 
@@ -285,6 +322,7 @@ async def awith_store(store_name: Optional[str]) -> AsyncIterator[BaseStore]:
 ALL_CHECKPOINTERS_SYNC = [
     "memory",
     "sqlite",
+    "duckdb",
     "postgres",
     "postgres_pipe",
     "postgres_pool",
@@ -292,6 +330,7 @@ ALL_CHECKPOINTERS_SYNC = [
 ALL_CHECKPOINTERS_ASYNC = [
     "memory",
     "sqlite_aio",
+    "duckdb_aio",
     "postgres_aio",
     "postgres_aio_pipe",
     "postgres_aio_pool",
@@ -300,5 +339,5 @@ ALL_CHECKPOINTERS_ASYNC_PLUS_NONE = [
     *ALL_CHECKPOINTERS_ASYNC,
     None,
 ]
-ALL_STORES_SYNC = ["in_memory", "postgres"]
-ALL_STORES_ASYNC = ["in_memory", "postgres_aio"]
+ALL_STORES_SYNC = ["in_memory", "postgres", "duckdb"]
+ALL_STORES_ASYNC = ["in_memory", "postgres_aio", "duckdb_aio"]

libs/langgraph/tests/test_prebuilt.py

@@ -1,5 +1,6 @@
 import dataclasses
 import json
+from functools import partial
 from typing import (
     Annotated,
     Any,
@@ -28,25 +29,39 @@ from langchain_core.messages import (
 )
 from langchain_core.outputs import ChatGeneration, ChatResult
 from langchain_core.runnables import Runnable, RunnableLambda
-from langchain_core.tools import BaseTool
+from langchain_core.tools import BaseTool, ToolException
 from langchain_core.tools import tool as dec_tool
-from pydantic import BaseModel
+from pydantic import BaseModel, ValidationError
 from pydantic.v1 import BaseModel as BaseModelV1
+from pydantic.v1 import ValidationError as ValidationErrorV1
 from typing_extensions import TypedDict
 
 from langgraph.checkpoint.base import BaseCheckpointSaver
-from langgraph.graph import START, MessagesState, StateGraph
-from langgraph.prebuilt import ToolNode, ValidationNode, create_react_agent
-from langgraph.prebuilt.tool_node import InjectedState, InjectedStore
+from langgraph.checkpoint.memory import MemorySaver
+from langgraph.errors import NodeInterrupt
+from langgraph.graph import START, MessagesState, StateGraph, add_messages
+from langgraph.prebuilt import (
+    ToolNode,
+    ValidationNode,
+    create_react_agent,
+    tools_condition,
+)
+from langgraph.prebuilt.tool_node import (
+    TOOL_CALL_ERROR_TEMPLATE,
+    InjectedState,
+    InjectedStore,
+    _infer_handled_types,
+)
 from langgraph.store.base import BaseStore
 from langgraph.store.memory import InMemoryStore
+from langgraph.types import Interrupt
 from tests.conftest import (
     ALL_CHECKPOINTERS_ASYNC,
     ALL_CHECKPOINTERS_SYNC,
     IS_LANGCHAIN_CORE_030_OR_GREATER,
     awith_checkpointer,
 )
-from tests.messages import _AnyIdHumanMessage
+from tests.messages import _AnyIdHumanMessage, _AnyIdToolMessage
 
 pytestmark = pytest.mark.anyio
 
@@ -362,32 +377,107 @@ def test_model_with_tools(tool_style: str):
         create_react_agent(model.bind_tools([tool1]), [tool2])
 
 
+def test__infer_handled_types() -> None:
+    def handle(e):  # type: ignore
+        return ""
+
+    def handle2(e: Exception) -> str:
+        return ""
+
+    def handle3(e: Union[ValueError, ToolException]) -> str:
+        return ""
+
+    class Handler:
+        def handle(self, e: ValueError) -> str:
+            return ""
+
+    handle4 = Handler().handle
+
+    def handle5(e: Union[Union[TypeError, ValueError], ToolException]):
+        return ""
+
+    expected: tuple = (Exception,)
+    actual = _infer_handled_types(handle)
+    assert expected == actual
+
+    expected = (Exception,)
+    actual = _infer_handled_types(handle2)
+    assert expected == actual
+
+    expected = (ValueError, ToolException)
+    actual = _infer_handled_types(handle3)
+    assert expected == actual
+
+    expected = (ValueError,)
+    actual = _infer_handled_types(handle4)
+    assert expected == actual
+
+    expected = (TypeError, ValueError, ToolException)
+    actual = _infer_handled_types(handle5)
+    assert expected == actual
+
+    with pytest.raises(ValueError):
+
+        def handler(e: str):
+            return ""
+
+        _infer_handled_types(handler)
+
+    with pytest.raises(ValueError):
+
+        def handler(e: list[Exception]):
+            return ""
+
+        _infer_handled_types(handler)
+
+    with pytest.raises(ValueError):
+
+        def handler(e: Union[str, int]):
+            return ""
+
+        _infer_handled_types(handler)
+
+
+# tools for testing Too
+def tool1(some_val: int, some_other_val: str) -> str:
+    """Tool 1 docstring."""
+    if some_val == 0:
+        raise ValueError("Test error")
+    return f"{some_val} - {some_other_val}"
+
+
+async def tool2(some_val: int, some_other_val: str) -> str:
+    """Tool 2 docstring."""
+    if some_val == 0:
+        raise ToolException("Test error")
+    return f"tool2: {some_val} - {some_other_val}"
+
+
+async def tool3(some_val: int, some_other_val: str) -> str:
+    """Tool 3 docstring."""
+    return [
+        {"key_1": some_val, "key_2": "foo"},
+        {"key_1": some_other_val, "key_2": "baz"},
+    ]
+
+
+async def tool4(some_val: int, some_other_val: str) -> str:
+    """Tool 4 docstring."""
+    return [
+        {"type": "image_url", "image_url": {"url": "abdc"}},
+    ]
+
+
+@dec_tool
+def tool5(some_val: int):
+    """Tool 5 docstring."""
+    raise ToolException("Test error")
+
+
+tool5.handle_tool_error = "foo"
+
+
 async def test_tool_node():
-    def tool1(some_val: int, some_other_val: str) -> str:
-        """Tool 1 docstring."""
-        if some_val == 0:
-            raise ValueError("Test error")
-        return f"{some_val} - {some_other_val}"
-
-    async def tool2(some_val: int, some_other_val: str) -> str:
-        """Tool 2 docstring."""
-        if some_val == 0:
-            raise ValueError("Test error")
-        return f"tool2: {some_val} - {some_other_val}"
-
-    async def tool3(some_val: int, some_other_val: str) -> str:
-        """Tool 3 docstring."""
-        return [
-            {"key_1": some_val, "key_2": "foo"},
-            {"key_1": some_other_val, "key_2": "baz"},
-        ]
-
-    async def tool4(some_val: int, some_other_val: str) -> str:
-        """Tool 4 docstring."""
-        return [
-            {"type": "image_url", "image_url": {"url": "abdc"}},
-        ]
-
     result = ToolNode([tool1]).invoke(
         {
             "messages": [
@@ -410,31 +500,6 @@ async def test_tool_node():
     assert tool_message.content == "1 - foo"
     assert tool_message.tool_call_id == "some 0"
 
-    result_error = ToolNode([tool1]).invoke(
-        {
-            "messages": [
-                AIMessage(
-                    "hi?",
-                    tool_calls=[
-                        {
-                            "name": "tool1",
-                            "args": {"some_val": 0, "some_other_val": "foo"},
-                            "id": "some 0",
-                        }
-                    ],
-                )
-            ]
-        }
-    )
-
-    tool_message: ToolMessage = result_error["messages"][-1]
-    assert tool_message.type == "tool"
-    assert (
-        tool_message.content
-        == f"Error: {repr(ValueError('Test error'))}\n Please fix your mistakes."
-    )
-    assert tool_message.tool_call_id == "some 0"
-
     result2 = await ToolNode([tool2]).ainvoke(
         {
             "messages": [
@@ -451,11 +516,232 @@ async def test_tool_node():
             ]
         }
     )
+
     tool_message: ToolMessage = result2["messages"][-1]
     assert tool_message.type == "tool"
     assert tool_message.content == "tool2: 2 - bar"
 
-    with pytest.raises(ValueError):
+    # list of dicts tool content
+    result3 = await ToolNode([tool3]).ainvoke(
+        {
+            "messages": [
+                AIMessage(
+                    "hi?",
+                    tool_calls=[
+                        {
+                            "name": "tool3",
+                            "args": {"some_val": 2, "some_other_val": "bar"},
+                            "id": "some 2",
+                        }
+                    ],
+                )
+            ]
+        }
+    )
+    tool_message: ToolMessage = result3["messages"][-1]
+    assert tool_message.type == "tool"
+    assert (
+        tool_message.content
+        == '[{"key_1": 2, "key_2": "foo"}, {"key_1": "bar", "key_2": "baz"}]'
+    )
+    assert tool_message.tool_call_id == "some 2"
+
+    # list of content blocks tool content
+    result4 = await ToolNode([tool4]).ainvoke(
+        {
+            "messages": [
+                AIMessage(
+                    "hi?",
+                    tool_calls=[
+                        {
+                            "name": "tool4",
+                            "args": {"some_val": 2, "some_other_val": "bar"},
+                            "id": "some 3",
+                        }
+                    ],
+                )
+            ]
+        }
+    )
+    tool_message: ToolMessage = result4["messages"][-1]
+    assert tool_message.type == "tool"
+    assert tool_message.content == [{"type": "image_url", "image_url": {"url": "abdc"}}]
+    assert tool_message.tool_call_id == "some 3"
+
+
+async def test_tool_node_error_handling():
+    def handle_all(e: Union[ValueError, ToolException, ValidationError]):
+        return TOOL_CALL_ERROR_TEMPLATE.format(error=repr(e))
+
+    # test catching all exceptions, via:
+    # - handle_tool_errors = True
+    # - passing a tuple of all exceptions
+    # - passing a callable with all exceptions in the signature
+    for handle_tool_errors in (
+        True,
+        (ValueError, ToolException, ValidationError),
+        handle_all,
+    ):
+        result_error = await ToolNode(
+            [tool1, tool2, tool3], handle_tool_errors=handle_tool_errors
+        ).ainvoke(
+            {
+                "messages": [
+                    AIMessage(
+                        "hi?",
+                        tool_calls=[
+                            {
+                                "name": "tool1",
+                                "args": {"some_val": 0, "some_other_val": "foo"},
+                                "id": "some id",
+                            },
+                            {
+                                "name": "tool2",
+                                "args": {"some_val": 0, "some_other_val": "bar"},
+                                "id": "some other id",
+                            },
+                            {
+                                "name": "tool3",
+                                "args": {"some_val": 0},
+                                "id": "another id",
+                            },
+                        ],
+                    )
+                ]
+            }
+        )
+
+        assert all(m.type == "tool" for m in result_error["messages"])
+        assert all(m.status == "error" for m in result_error["messages"])
+        assert (
+            result_error["messages"][0].content
+            == f"Error: {repr(ValueError('Test error'))}\n Please fix your mistakes."
+        )
+        assert (
+            result_error["messages"][1].content
+            == f"Error: {repr(ToolException('Test error'))}\n Please fix your mistakes."
+        )
+        assert (
+            "ValidationError" in result_error["messages"][2].content
+            or "validation error" in result_error["messages"][2].content
+        )
+
+        assert result_error["messages"][0].tool_call_id == "some id"
+        assert result_error["messages"][1].tool_call_id == "some other id"
+        assert result_error["messages"][2].tool_call_id == "another id"
+
+
+async def test_tool_node_error_handling_callable():
+    def handle_value_error(e: ValueError):
+        return "Value error"
+
+    def handle_tool_exception(e: ToolException):
+        return "Tool exception"
+
+    for handle_tool_errors in ("Value error", handle_value_error):
+        result_error = await ToolNode(
+            [tool1], handle_tool_errors=handle_tool_errors
+        ).ainvoke(
+            {
+                "messages": [
+                    AIMessage(
+                        "hi?",
+                        tool_calls=[
+                            {
+                                "name": "tool1",
+                                "args": {"some_val": 0, "some_other_val": "foo"},
+                                "id": "some id",
+                            },
+                        ],
+                    )
+                ]
+            }
+        )
+        tool_message: ToolMessage = result_error["messages"][-1]
+        assert tool_message.type == "tool"
+        assert tool_message.status == "error"
+        assert tool_message.content == "Value error"
+
+    # test raising for an unhandled exception, via:
+    # - passing a tuple of all exceptions
+    # - passing a callable with all exceptions in the signature
+    for handle_tool_errors in ((ValueError,), handle_value_error):
+        with pytest.raises(ToolException) as exc_info:
+            await ToolNode(
+                [tool1, tool2], handle_tool_errors=handle_tool_errors
+            ).ainvoke(
+                {
+                    "messages": [
+                        AIMessage(
+                            "hi?",
+                            tool_calls=[
+                                {
+                                    "name": "tool1",
+                                    "args": {"some_val": 0, "some_other_val": "foo"},
+                                    "id": "some id",
+                                },
+                                {
+                                    "name": "tool2",
+                                    "args": {"some_val": 0, "some_other_val": "bar"},
+                                    "id": "some other id",
+                                },
+                            ],
+                        )
+                    ]
+                }
+            )
+        assert str(exc_info.value) == "Test error"
+
+    for handle_tool_errors in ((ToolException,), handle_tool_exception):
+        with pytest.raises(ValueError) as exc_info:
+            await ToolNode(
+                [tool1, tool2], handle_tool_errors=handle_tool_errors
+            ).ainvoke(
+                {
+                    "messages": [
+                        AIMessage(
+                            "hi?",
+                            tool_calls=[
+                                {
+                                    "name": "tool1",
+                                    "args": {"some_val": 0, "some_other_val": "foo"},
+                                    "id": "some id",
+                                },
+                                {
+                                    "name": "tool2",
+                                    "args": {"some_val": 0, "some_other_val": "bar"},
+                                    "id": "some other id",
+                                },
+                            ],
+                        )
+                    ]
+                }
+            )
+        assert str(exc_info.value) == "Test error"
+
+
+async def test_tool_node_handle_tool_errors_false():
+    with pytest.raises(ValueError) as exc_info:
+        ToolNode([tool1], handle_tool_errors=False).invoke(
+            {
+                "messages": [
+                    AIMessage(
+                        "hi?",
+                        tool_calls=[
+                            {
+                                "name": "tool1",
+                                "args": {"some_val": 0, "some_other_val": "foo"},
+                                "id": "some id",
+                            }
+                        ],
+                    )
+                ]
+            }
+        )
+
+    assert str(exc_info.value) == "Test error"
+
+    with pytest.raises(ToolException):
         await ToolNode([tool2], handle_tool_errors=False).ainvoke(
             {
                 "messages": [
@@ -465,7 +751,7 @@ async def test_tool_node():
                             {
                                 "name": "tool2",
                                 "args": {"some_val": 0, "some_other_val": "bar"},
-                                "id": "some 1",
+                                "id": "some id",
                             }
                         ],
                     )
@@ -473,7 +759,57 @@ async def test_tool_node():
             }
         )
 
-    # incorrect tool name
+    assert str(exc_info.value) == "Test error"
+
+    # test validation errors get raised if handle_tool_errors is False
+    with pytest.raises((ValidationError, ValidationErrorV1)):
+        ToolNode([tool1], handle_tool_errors=False).invoke(
+            {
+                "messages": [
+                    AIMessage(
+                        "hi?",
+                        tool_calls=[
+                            {
+                                "name": "tool1",
+                                "args": {"some_val": 0},
+                                "id": "some id",
+                            }
+                        ],
+                    )
+                ]
+            }
+        )
+
+
+def test_tool_node_individual_tool_error_handling():
+    # test error handling on individual tools (and that it overrides overall error handling!)
+    result_individual_tool_error_handler = ToolNode(
+        [tool5], handle_tool_errors="bar"
+    ).invoke(
+        {
+            "messages": [
+                AIMessage(
+                    "hi?",
+                    tool_calls=[
+                        {
+                            "name": "tool5",
+                            "args": {"some_val": 0},
+                            "id": "some 0",
+                        }
+                    ],
+                )
+            ]
+        }
+    )
+
+    tool_message: ToolMessage = result_individual_tool_error_handler["messages"][-1]
+    assert tool_message.type == "tool"
+    assert tool_message.status == "error"
+    assert tool_message.content == "foo"
+    assert tool_message.tool_call_id == "some 0"
+
+
+def test_tool_node_incorrect_tool_name():
     result_incorrect_name = ToolNode([tool1, tool2]).invoke(
         {
             "messages": [
@@ -490,60 +826,94 @@ async def test_tool_node():
             ]
         }
     )
+
     tool_message: ToolMessage = result_incorrect_name["messages"][-1]
     assert tool_message.type == "tool"
+    assert tool_message.status == "error"
     assert (
         tool_message.content
         == "Error: tool3 is not a valid tool, try one of [tool1, tool2]."
     )
     assert tool_message.tool_call_id == "some 0"
 
-    # list of dicts tool content
-    result3 = await ToolNode([tool3]).ainvoke(
-        {
-            "messages": [
-                AIMessage(
-                    "hi?",
-                    tool_calls=[
-                        {
-                            "name": "tool3",
-                            "args": {"some_val": 2, "some_other_val": "bar"},
-                            "id": "some 0",
-                        }
-                    ],
-                )
-            ]
-        }
-    )
-    tool_message: ToolMessage = result3["messages"][-1]
-    assert tool_message.type == "tool"
-    assert (
-        tool_message.content
-        == '[{"key_1": 2, "key_2": "foo"}, {"key_1": "bar", "key_2": "baz"}]'
-    )
-    assert tool_message.tool_call_id == "some 0"
 
-    # list of content blocks tool content
-    result4 = await ToolNode([tool4]).ainvoke(
-        {
-            "messages": [
-                AIMessage(
-                    "hi?",
-                    tool_calls=[
-                        {
-                            "name": "tool4",
-                            "args": {"some_val": 2, "some_other_val": "bar"},
-                            "id": "some 0",
-                        }
-                    ],
-                )
-            ]
-        }
+def test_tool_node_node_interrupt():
+    def tool_normal(some_val: int) -> str:
+        """Tool docstring."""
+        return "normal"
+
+    def tool_interrupt(some_val: int) -> str:
+        """Tool docstring."""
+        raise NodeInterrupt("foo")
+
+    def handle(e: NodeInterrupt):
+        return "handled"
+
+    for handle_tool_errors in (True, (NodeInterrupt,), "handled", handle, False):
+        node = ToolNode([tool_interrupt], handle_tool_errors=handle_tool_errors)
+        with pytest.raises(NodeInterrupt) as exc_info:
+            node.invoke(
+                {
+                    "messages": [
+                        AIMessage(
+                            "hi?",
+                            tool_calls=[
+                                {
+                                    "name": "tool_interrupt",
+                                    "args": {"some_val": 0},
+                                    "id": "some 0",
+                                }
+                            ],
+                        )
+                    ]
+                }
+            )
+            assert exc_info.value == "foo"
+
+    # test inside react agent
+    model = FakeToolCallingModel(
+        tool_calls=[
+            [
+                ToolCall(name="tool_interrupt", args={"some_val": 0}, id="1"),
+                ToolCall(name="tool_normal", args={"some_val": 1}, id="2"),
+            ],
+            [],
+        ]
     )
-    tool_message: ToolMessage = result4["messages"][-1]
-    assert tool_message.type == "tool"
-    assert tool_message.content == [{"type": "image_url", "image_url": {"url": "abdc"}}]
-    assert tool_message.tool_call_id == "some 0"
+    checkpointer = MemorySaver()
+    config = {"configurable": {"thread_id": "1"}}
+    agent = create_react_agent(
+        model, [tool_interrupt, tool_normal], checkpointer=checkpointer
+    )
+    result = agent.invoke({"messages": [HumanMessage("hi?")]}, config)
+    assert result["messages"] == [
+        _AnyIdHumanMessage(
+            content="hi?",
+        ),
+        AIMessage(
+            content="hi?",
+            id="0",
+            tool_calls=[
+                {
+                    "name": "tool_interrupt",
+                    "args": {"some_val": 0},
+                    "id": "1",
+                    "type": "tool_call",
+                },
+                {
+                    "name": "tool_normal",
+                    "args": {"some_val": 1},
+                    "id": "2",
+                    "type": "tool_call",
+                },
+            ],
+        ),
+    ]
+    state = agent.get_state(config)
+    assert state.next == ("tools",)
+    task = state.tasks[0]
+    assert task.name == "tools"
+    assert task.interrupts == (Interrupt(value="foo", when="during"),)
 
 
 def my_function(some_val: int, some_other_val: str) -> str:
@@ -826,6 +1196,47 @@ def test_tool_node_ensure_utf8() -> None:
     assert outputs[0].content == json.dumps(data, ensure_ascii=False)
 
 
+def test_tool_node_messages_key() -> None:
+    @dec_tool
+    def add(a: int, b: int):
+        """Adds a and b."""
+        return a + b
+
+    model = FakeToolCallingModel(
+        tool_calls=[[ToolCall(name=add.name, args={"a": 1, "b": 2}, id="test_id")]]
+    )
+
+    class State(TypedDict):
+        subgraph_messages: Annotated[list[AnyMessage], add_messages]
+
+    def call_model(state: State):
+        response = model.invoke(state["subgraph_messages"])
+        model.tool_calls = []
+        return {"subgraph_messages": response}
+
+    builder = StateGraph(State)
+    builder.add_node("agent", call_model)
+    builder.add_node("tools", ToolNode([add], messages_key="subgraph_messages"))
+    builder.add_conditional_edges(
+        "agent", partial(tools_condition, messages_key="subgraph_messages")
+    )
+    builder.add_edge(START, "agent")
+    builder.add_edge("tools", "agent")
+
+    graph = builder.compile()
+    result = graph.invoke({"subgraph_messages": [HumanMessage(content="hi")]})
+    assert result["subgraph_messages"] == [
+        _AnyIdHumanMessage(content="hi"),
+        AIMessage(
+            content="hi",
+            id="0",
+            tool_calls=[ToolCall(name=add.name, args={"a": 1, "b": 2}, id="test_id")],
+        ),
+        _AnyIdToolMessage(content="3", name=add.name, tool_call_id="test_id"),
+        AIMessage(content="hi-hi-3", id="1"),
+    ]
+
+
 async def test_return_direct() -> None:
     @dec_tool(return_direct=True)
     def tool_return_direct(input: str) -> str:

libs/langgraph/tests/test_pregel_async.py

@@ -1922,6 +1922,53 @@ async def test_cond_edge_after_send() -> None:
     assert await graph.ainvoke(["0"]) == ["0", "1", "2", "2", "3"]
 
 
+async def test_max_concurrency() -> None:
+    class Node:
+        def __init__(self, name: str):
+            self.name = name
+            setattr(self, "__name__", name)
+            self.currently = 0
+            self.max_currently = 0
+
+        async def __call__(self, state):
+            self.currently += 1
+            if self.currently > self.max_currently:
+                self.max_currently = self.currently
+            await asyncio.sleep(0.1)
+            self.currently -= 1
+            return [self.name]
+
+    async def send_to_many(state):
+        return [Send("2", state)] * 100
+
+    async def route_to_three(state) -> Literal["3"]:
+        return "3"
+
+    node2 = Node("2")
+    builder = StateGraph(Annotated[list, operator.add])
+    builder.add_node(Node("1"))
+    builder.add_node(node2)
+    builder.add_node(Node("3"))
+    builder.add_edge(START, "1")
+    builder.add_conditional_edges("1", send_to_many)
+    builder.add_conditional_edges("2", route_to_three)
+    graph = builder.compile()
+
+    assert await graph.ainvoke(["0"]) == ["0", "1", *(["2"] * 100), "3"]
+    assert node2.max_currently == 100
+    assert node2.currently == 0
+    node2.max_currently = 0
+
+    assert await graph.ainvoke(["0"], {"max_concurrency": 10}) == [
+        "0",
+        "1",
+        *(["2"] * 100),
+        "3",
+    ]
+    assert node2.max_currently == 10
+    assert node2.currently == 0
+
+
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
 async def test_invoke_checkpoint_three(
     mocker: MockerFixture, checkpointer_name: str

libs/langgraph/tests/test_remote_graph.py

@@ -7,7 +7,9 @@ from langchain_core.runnables.graph import (
 from langchain_core.runnables.graph import (
     Node as DrawableNode,
 )
+from langgraph_sdk.schema import StreamPart
 
+from langgraph.errors import GraphInterrupt
 from langgraph.pregel.remote import RemoteGraph
 from langgraph.pregel.types import StateSnapshot
 
@@ -15,7 +17,7 @@ from langgraph.pregel.types import StateSnapshot
 def test_with_config():
     # set up test
     remote_pregel = RemoteGraph(
-        graph_id="test_graph_id",
+        "test_graph_id",
         config={
             "configurable": {
                 "foo": "bar",
@@ -62,7 +64,7 @@ def test_get_graph():
         ],
     }
 
-    remote_pregel = RemoteGraph(sync_client=mock_sync_client, graph_id="test_graph_id")
+    remote_pregel = RemoteGraph("test_graph_id", sync_client=mock_sync_client)
 
     # call method / assertions
     drawable_graph = remote_pregel.get_graph()
@@ -109,7 +111,7 @@ async def test_aget_graph():
         ],
     }
 
-    remote_pregel = RemoteGraph(client=mock_async_client, graph_id="test_graph_id")
+    remote_pregel = RemoteGraph("test_graph_id", client=mock_async_client)
 
     # call method / assertions
     drawable_graph = await remote_pregel.aget_graph()
@@ -133,92 +135,6 @@ async def test_aget_graph():
     ]
 
 
-def test_get_subgraphs():
-    # set up test
-    mock_sync_client = MagicMock()
-    mock_sync_client.assistants.get_subgraphs.return_value = {
-        "namespace_1": {
-            "graph_id": "test_graph_id_2",
-            "input_schema": {},
-            "output_schema": {},
-            "state_schema": {},
-            "config_schema": {},
-        },
-        "namespace_2": {
-            "graph_id": "test_graph_id_3",
-            "input_schema": {},
-            "output_schema": {},
-            "state_schema": {},
-            "config_schema": {},
-        },
-    }
-
-    remote_pregel = RemoteGraph(
-        sync_client=mock_sync_client, graph_id="test_graph_id_1"
-    )
-
-    # call method / assertions
-    subgraphs = list(remote_pregel.get_subgraphs())
-    assert len(subgraphs) == 2
-
-    subgraph_1 = subgraphs[0]
-    ns_1 = subgraph_1[0]
-    remote_pregel_1: RemoteGraph = subgraph_1[1]
-    assert ns_1 == "namespace_1"
-    assert remote_pregel_1.graph_id == "test_graph_id_2"
-
-    subgraph_2 = subgraphs[1]
-    ns_2 = subgraph_2[0]
-    remote_pregel_2: RemoteGraph = subgraph_2[1]
-    assert ns_2 == "namespace_2"
-    assert remote_pregel_2.graph_id == "test_graph_id_3"
-
-
-@pytest.mark.anyio
-async def test_aget_subgraphs():
-    # set up test
-    mock_async_client = AsyncMock()
-    mock_async_client.assistants.get_subgraphs.return_value = {
-        "namespace_1": {
-            "graph_id": "test_graph_id_2",
-            "input_schema": {},
-            "output_schema": {},
-            "state_schema": {},
-            "config_schema": {},
-        },
-        "namespace_2": {
-            "graph_id": "test_graph_id_3",
-            "input_schema": {},
-            "output_schema": {},
-            "state_schema": {},
-            "config_schema": {},
-        },
-    }
-
-    remote_pregel = RemoteGraph(
-        client=mock_async_client,
-        graph_id="test_graph_id_1",
-    )
-
-    # call method / assertions
-    subgraphs = []
-    async for subgraph in remote_pregel.aget_subgraphs():
-        subgraphs.append(subgraph)
-    assert len(subgraphs) == 2
-
-    subgraph_1 = subgraphs[0]
-    ns_1 = subgraph_1[0]
-    remote_pregel_1: RemoteGraph = subgraph_1[1]
-    assert ns_1 == "namespace_1"
-    assert remote_pregel_1.graph_id == "test_graph_id_2"
-
-    subgraph_2 = subgraphs[1]
-    ns_2 = subgraph_2[0]
-    remote_pregel_2: RemoteGraph = subgraph_2[1]
-    assert ns_2 == "namespace_2"
-    assert remote_pregel_2.graph_id == "test_graph_id_3"
-
-
 def test_get_state():
     # set up test
     mock_sync_client = MagicMock()
@@ -238,7 +154,10 @@ def test_get_state():
     }
 
     # call method / assertions
-    remote_pregel = RemoteGraph(sync_client=mock_sync_client, graph_id="test_graph_id")
+    remote_pregel = RemoteGraph(
+        "test_graph_id",
+        sync_client=mock_sync_client,
+    )
 
     config = {"configurable": {"thread_id": "thread1"}}
     state_snapshot = remote_pregel.get_state(config)
@@ -285,7 +204,10 @@ async def test_aget_state():
     }
 
     # call method / assertions
-    remote_pregel = RemoteGraph(client=mock_async_client, graph_id="test_graph_id")
+    remote_pregel = RemoteGraph(
+        "test_graph_id",
+        client=mock_async_client,
+    )
 
     config = {"configurable": {"thread_id": "thread1"}}
     state_snapshot = await remote_pregel.aget_state(config)
@@ -336,7 +258,10 @@ def test_get_state_history():
     ]
 
     # call method / assertions
-    remote_pregel = RemoteGraph(sync_client=mock_sync_client, graph_id="test_graph_id")
+    remote_pregel = RemoteGraph(
+        "test_graph_id",
+        sync_client=mock_sync_client,
+    )
 
     config = {"configurable": {"thread_id": "thread1"}}
     state_history_snapshot = list(
@@ -384,7 +309,10 @@ async def test_aget_state_history():
     ]
 
     # call method / assertions
-    remote_pregel = RemoteGraph(client=mock_async_client, graph_id="test_graph_id")
+    remote_pregel = RemoteGraph(
+        "test_graph_id",
+        client=mock_async_client,
+    )
 
     config = {"configurable": {"thread_id": "thread1"}}
     state_history_snapshot = []
@@ -425,7 +353,10 @@ def test_update_state():
     }
 
     # call method / assertions
-    remote_pregel = RemoteGraph(sync_client=mock_sync_client, graph_id="test_graph_id")
+    remote_pregel = RemoteGraph(
+        "test_graph_id",
+        sync_client=mock_sync_client,
+    )
 
     config = {"configurable": {"thread_id": "thread1"}}
     response = remote_pregel.update_state(config, {"key": "value"})
@@ -454,7 +385,10 @@ async def test_aupdate_state():
     }
 
     # call method / assertions
-    remote_pregel = RemoteGraph(client=mock_async_client, graph_id="test_graph_id")
+    remote_pregel = RemoteGraph(
+        "test_graph_id",
+        client=mock_async_client,
+    )
 
     config = {"configurable": {"thread_id": "thread1"}}
     response = await remote_pregel.aupdate_state(config, {"key": "value"})
@@ -473,17 +407,100 @@ def test_stream():
     # set up test
     mock_sync_client = MagicMock()
     mock_sync_client.runs.stream.return_value = [
+        StreamPart(event="values", data={"chunk": "data1"}),
+        StreamPart(event="values", data={"chunk": "data2"}),
+        StreamPart(event="values", data={"chunk": "data3"}),
+        StreamPart(event="updates", data={"chunk": "data4"}),
+        StreamPart(event="updates", data={"__interrupt__": ()}),
+    ]
+
+    # call method / assertions
+    remote_pregel = RemoteGraph(
+        "test_graph_id",
+        sync_client=mock_sync_client,
+    )
+
+    # stream modes doesn't include 'updates'
+    stream_parts = []
+    with pytest.raises(GraphInterrupt):
+        for stream_part in remote_pregel.stream(
+            {"input": "data"},
+            config={"configurable": {"thread_id": "thread_1"}},
+            stream_mode="values",
+        ):
+            stream_parts.append(stream_part)
+
+    assert stream_parts == [
         {"chunk": "data1"},
         {"chunk": "data2"},
         {"chunk": "data3"},
     ]
 
-    # call method / assertions
-    remote_pregel = RemoteGraph(sync_client=mock_sync_client, graph_id="test_graph_id")
+    mock_sync_client.runs.stream.return_value = [
+        StreamPart(event="updates", data={"chunk": "data3"}),
+        StreamPart(event="updates", data={"chunk": "data4"}),
+        StreamPart(event="updates", data={"__interrupt__": ()}),
+    ]
 
-    config = {"configurable": {"thread_id": "thread_1"}}
-    result = list(remote_pregel.stream({"input": "data"}, config))
-    assert result == [{"chunk": "data1"}, {"chunk": "data2"}, {"chunk": "data3"}]
+    # default stream_mode is updates
+    stream_parts = []
+    with pytest.raises(GraphInterrupt):
+        for stream_part in remote_pregel.stream(
+            {"input": "data"},
+            config={"configurable": {"thread_id": "thread_1"}},
+        ):
+            stream_parts.append(stream_part)
+
+    assert stream_parts == [
+        {"chunk": "data3"},
+        {"chunk": "data4"},
+    ]
+
+    # list stream_mode includes mode names
+    stream_parts = []
+    with pytest.raises(GraphInterrupt):
+        for stream_part in remote_pregel.stream(
+            {"input": "data"},
+            config={"configurable": {"thread_id": "thread_1"}},
+            stream_mode=["updates"],
+        ):
+            stream_parts.append(stream_part)
+
+    assert stream_parts == [
+        ("updates", {"chunk": "data3"}),
+        ("updates", {"chunk": "data4"}),
+    ]
+
+    # subgraphs + list modes
+    stream_parts = []
+    with pytest.raises(GraphInterrupt):
+        for stream_part in remote_pregel.stream(
+            {"input": "data"},
+            config={"configurable": {"thread_id": "thread_1"}},
+            stream_mode=["updates"],
+            subgraphs=True,
+        ):
+            stream_parts.append(stream_part)
+
+    assert stream_parts == [
+        ((), "updates", {"chunk": "data3"}),
+        ((), "updates", {"chunk": "data4"}),
+    ]
+
+    # subgraphs + single mode
+    stream_parts = []
+    with pytest.raises(GraphInterrupt):
+        for stream_part in remote_pregel.stream(
+            {"input": "data"},
+            config={"configurable": {"thread_id": "thread_1"}},
+            subgraphs=True,
+        ):
+            stream_parts.append(stream_part)
+
+    assert stream_parts == [
+        ((), {"chunk": "data3"}),
+        ((), {"chunk": "data4"}),
+    ]
 
 
 @pytest.mark.anyio
@@ -492,20 +509,142 @@ async def test_astream():
     mock_async_client = MagicMock()
     async_iter = MagicMock()
     async_iter.__aiter__.return_value = [
-        {"chunk": "data1"},
-        {"chunk": "data2"},
-        {"chunk": "data3"},
+        StreamPart(event="values", data={"chunk": "data1"}),
+        StreamPart(event="values", data={"chunk": "data2"}),
+        StreamPart(event="values", data={"chunk": "data3"}),
+        StreamPart(event="updates", data={"chunk": "data4"}),
+        StreamPart(event="updates", data={"__interrupt__": ()}),
     ]
     mock_async_client.runs.stream.return_value = async_iter
 
     # call method / assertions
-    remote_pregel = RemoteGraph(client=mock_async_client, graph_id="test_graph_id")
+    remote_pregel = RemoteGraph(
+        "test_graph_id",
+        client=mock_async_client,
+    )
 
-    config = {"configurable": {"thread_id": "thread_1"}}
-    chunks = []
-    async for chunk in remote_pregel.astream({"input": "data"}, config):
-        chunks.append(chunk)
-    assert chunks == [{"chunk": "data1"}, {"chunk": "data2"}, {"chunk": "data3"}]
+    # stream modes doesn't include 'updates'
+    stream_parts = []
+    with pytest.raises(GraphInterrupt):
+        async for stream_part in remote_pregel.astream(
+            {"input": "data"},
+            config={"configurable": {"thread_id": "thread_1"}},
+            stream_mode="values",
+        ):
+            stream_parts.append(stream_part)
+
+    assert stream_parts == [
+        {"chunk": "data1"},
+        {"chunk": "data2"},
+        {"chunk": "data3"},
+    ]
+
+    async_iter = MagicMock()
+    async_iter.__aiter__.return_value = [
+        StreamPart(event="updates", data={"chunk": "data3"}),
+        StreamPart(event="updates", data={"chunk": "data4"}),
+        StreamPart(event="updates", data={"__interrupt__": ()}),
+    ]
+    mock_async_client.runs.stream.return_value = async_iter
+
+    # default stream_mode is updates
+    stream_parts = []
+    with pytest.raises(GraphInterrupt):
+        async for stream_part in remote_pregel.astream(
+            {"input": "data"},
+            config={"configurable": {"thread_id": "thread_1"}},
+        ):
+            stream_parts.append(stream_part)
+
+    assert stream_parts == [
+        {"chunk": "data3"},
+        {"chunk": "data4"},
+    ]
+
+    # list stream_mode includes mode names
+    stream_parts = []
+    with pytest.raises(GraphInterrupt):
+        async for stream_part in remote_pregel.astream(
+            {"input": "data"},
+            config={"configurable": {"thread_id": "thread_1"}},
+            stream_mode=["updates"],
+        ):
+            stream_parts.append(stream_part)
+
+    assert stream_parts == [
+        ("updates", {"chunk": "data3"}),
+        ("updates", {"chunk": "data4"}),
+    ]
+
+    # subgraphs + list modes
+    stream_parts = []
+    with pytest.raises(GraphInterrupt):
+        async for stream_part in remote_pregel.astream(
+            {"input": "data"},
+            config={"configurable": {"thread_id": "thread_1"}},
+            stream_mode=["updates"],
+            subgraphs=True,
+        ):
+            stream_parts.append(stream_part)
+
+    assert stream_parts == [
+        ((), "updates", {"chunk": "data3"}),
+        ((), "updates", {"chunk": "data4"}),
+    ]
+
+    # subgraphs + single mode
+    stream_parts = []
+    with pytest.raises(GraphInterrupt):
+        async for stream_part in remote_pregel.astream(
+            {"input": "data"},
+            config={"configurable": {"thread_id": "thread_1"}},
+            subgraphs=True,
+        ):
+            stream_parts.append(stream_part)
+
+    assert stream_parts == [
+        ((), {"chunk": "data3"}),
+        ((), {"chunk": "data4"}),
+    ]
+
+    async_iter = MagicMock()
+    async_iter.__aiter__.return_value = [
+        StreamPart(event="updates|my|subgraph", data={"chunk": "data3"}),
+        StreamPart(event="updates|hello|subgraph", data={"chunk": "data4"}),
+        StreamPart(event="updates|bye|subgraph", data={"__interrupt__": ()}),
+    ]
+    mock_async_client.runs.stream.return_value = async_iter
+
+    # subgraphs + list modes
+    stream_parts = []
+    with pytest.raises(GraphInterrupt):
+        async for stream_part in remote_pregel.astream(
+            {"input": "data"},
+            config={"configurable": {"thread_id": "thread_1"}},
+            stream_mode=["updates"],
+            subgraphs=True,
+        ):
+            stream_parts.append(stream_part)
+
+    assert stream_parts == [
+        (("my", "subgraph"), "updates", {"chunk": "data3"}),
+        (("hello", "subgraph"), "updates", {"chunk": "data4"}),
+    ]
+
+    # subgraphs + single mode
+    stream_parts = []
+    with pytest.raises(GraphInterrupt):
+        async for stream_part in remote_pregel.astream(
+            {"input": "data"},
+            config={"configurable": {"thread_id": "thread_1"}},
+            subgraphs=True,
+        ):
+            stream_parts.append(stream_part)
+
+    assert stream_parts == [
+        (("my", "subgraph"), {"chunk": "data3"}),
+        (("hello", "subgraph"), {"chunk": "data4"}),
+    ]
 
 
 def test_invoke():
@@ -516,7 +655,10 @@ def test_invoke():
     }
 
     # call method / assertions
-    remote_pregel = RemoteGraph(sync_client=mock_sync_client, graph_id="test_graph_id")
+    remote_pregel = RemoteGraph(
+        "test_graph_id",
+        sync_client=mock_sync_client,
+    )
 
     config = {"configurable": {"thread_id": "thread_1"}}
     result = remote_pregel.invoke(
@@ -535,7 +677,10 @@ async def test_ainvoke():
     }
 
     # call method / assertions
-    remote_pregel = RemoteGraph(client=mock_async_client, graph_id="test_graph_id")
+    remote_pregel = RemoteGraph(
+        "test_graph_id",
+        client=mock_async_client,
+    )
 
     config = {"configurable": {"thread_id": "thread_1"}}
     result = await remote_pregel.ainvoke(
@@ -557,7 +702,9 @@ async def test_langgraph_cloud_integration():
     client = get_client()
     sync_client = get_sync_client()
     remote_pregel = RemoteGraph(
-        client=client, sync_client=sync_client, graph_id="agent"
+        "agent",
+        client=client,
+        sync_client=sync_client,
     )
 
     # define graph
@@ -572,7 +719,7 @@ async def test_langgraph_cloud_integration():
         "messages": [
             {
                 "role": "human",
-                "content": "Hello world!",
+                "content": "What's the weather in SF?",
             }
         ]
     }
@@ -580,7 +727,8 @@ async def test_langgraph_cloud_integration():
     # test invoke
     response = app.invoke(
         input,
-        config={"configurable": {"thread_id": "2dc3e3e7-39ac-4597-aa57-4404b944e82a"}},
+        config={"configurable": {"thread_id": "39a6104a-34e7-4f83-929c-d9eb163003c9"}},
+        interrupt_before=["agent"],
     )
     print("response:", response["messages"][-1].content)
 
@@ -634,9 +782,3 @@ async def test_langgraph_cloud_integration():
     remote_pregel.graph_id = "fe096781-5601-53d2-b2f6-0d3403f7e9ca"  # must be UUID
     graph = await remote_pregel.aget_graph(xray=True)
     print("graph:", graph)
-
-    # test get subgraphs
-    remote_pregel.graph_id = "fe096781-5601-53d2-b2f6-0d3403f7e9ca"  # must be UUID
-    async for name, pregel in remote_pregel.aget_subgraphs():
-        print("name:", name)
-        print("pregel:", pregel)

libs/langgraph/tests/test_state.py

@@ -61,6 +61,9 @@ def test_state_schema_with_type_hint():
     class OutputState(TypedDict):
         input_state: InputState
 
+    class FooState(InputState):
+        foo: str
+
     def complete_hint(state: InputState) -> OutputState:
         return {"input_state": state}
 
@@ -73,24 +76,46 @@ def test_state_schema_with_type_hint():
     def miss_all_hint(state, config):
         return {"input_state": state}
 
+    def pre_foo(_) -> FooState:
+        return {"foo": "bar"}
+
+    class Foo:
+        def __call__(self, state: FooState) -> OutputState:
+            assert state.pop("foo") == "bar"
+            return {"input_state": state}
+
     graph = StateGraph(InputState, output=OutputState)
-    actions = [complete_hint, miss_first_hint, only_return_hint, miss_all_hint]
+    actions = [
+        complete_hint,
+        miss_first_hint,
+        only_return_hint,
+        miss_all_hint,
+        pre_foo,
+        Foo(),
+    ]
 
     for action in actions:
         graph.add_node(action)
 
-    graph.set_entry_point(actions[0].__name__)
+    def get_name(action) -> str:
+        return getattr(action, "__name__", action.__class__.__name__)
+
+    graph.set_entry_point(get_name(actions[0]))
     for i in range(len(actions) - 1):
-        graph.add_edge(actions[i].__name__, actions[i + 1].__name__)
-    graph.set_finish_point(actions[-1].__name__)
+        graph.add_edge(get_name(actions[i]), get_name(actions[i + 1]))
+    graph.set_finish_point(get_name(actions[-1]))
 
     graph = graph.compile()
 
     input_state = InputState(question="Hello World!")
     output_state = OutputState(input_state=input_state)
+    foo_state = FooState(foo="bar")
     for i, c in enumerate(graph.stream(input_state, stream_mode="updates")):
-        node_name = actions[i].__name__
-        assert c[node_name] == output_state
+        node_name = get_name(actions[i])
+        if node_name == get_name(pre_foo):
+            assert c[node_name] == foo_state
+        else:
+            assert c[node_name] == output_state
 
 
 @pytest.mark.parametrize("total_", [True, False])

libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py

@@ -191,7 +191,7 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager):
                 step=saved.metadata["step"] + 1,
                 stop=saved.metadata["step"] + 2,
             ),
-        ) as (channels, managed), AsyncBackgroundExecutor() as submit:
+        ) as (channels, managed), AsyncBackgroundExecutor(msg["config"]) as submit:
             if task := await asyncio.to_thread(
                 prepare_single_task,
                 msg["task"]["path"],

libs/sdk-js/package.json

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

libs/sdk-js/src/schema.ts

@@ -10,7 +10,7 @@ type RunStatus =
   | "timeout"
   | "interrupted";
 
-type ThreadStatus = "idle" | "busy" | "interrupted";
+type ThreadStatus = "idle" | "busy" | "interrupted" | "error";
 
 type MultitaskStrategy = "reject" | "interrupt" | "rollback" | "enqueue";
 

libs/sdk-py/langgraph_sdk/client.py

@@ -39,6 +39,7 @@ from langgraph_sdk.schema import (
     Cron,
     DisconnectMode,
     GraphSchema,
+    IfNotExists,
     Item,
     Json,
     ListNamespaceResponse,
@@ -1169,18 +1170,19 @@ class RunsClient:
         assistant_id: str,
         *,
         input: Optional[dict] = None,
-        stream_mode: Union[StreamMode, list[StreamMode]] = "values",
+        stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
         stream_subgraphs: bool = False,
         metadata: Optional[dict] = None,
         config: Optional[Config] = None,
         checkpoint: Optional[Checkpoint] = None,
         checkpoint_id: Optional[str] = None,
-        interrupt_before: Optional[Union[All, list[str]]] = None,
-        interrupt_after: Optional[Union[All, list[str]]] = None,
-        feedback_keys: Optional[list[str]] = None,
+        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
+        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
+        feedback_keys: Optional[Sequence[str]] = None,
         on_disconnect: Optional[DisconnectMode] = None,
         webhook: Optional[str] = None,
         multitask_strategy: Optional[MultitaskStrategy] = None,
+        if_not_exists: Optional[IfNotExists] = None,
         after_seconds: Optional[int] = None,
     ) -> AsyncIterator[StreamPart]: ...
 
@@ -1191,15 +1193,16 @@ class RunsClient:
         assistant_id: str,
         *,
         input: Optional[dict] = None,
-        stream_mode: Union[StreamMode, list[StreamMode]] = "values",
+        stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
         stream_subgraphs: bool = False,
         metadata: Optional[dict] = None,
         config: Optional[Config] = None,
-        interrupt_before: Optional[Union[All, list[str]]] = None,
-        interrupt_after: Optional[Union[All, list[str]]] = None,
-        feedback_keys: Optional[list[str]] = None,
+        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
+        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
+        feedback_keys: Optional[Sequence[str]] = None,
         on_disconnect: Optional[DisconnectMode] = None,
         on_completion: Optional[OnCompletionBehavior] = None,
+        if_not_exists: Optional[IfNotExists] = None,
         webhook: Optional[str] = None,
         after_seconds: Optional[int] = None,
     ) -> AsyncIterator[StreamPart]: ...
@@ -1210,19 +1213,20 @@ class RunsClient:
         assistant_id: str,
         *,
         input: Optional[dict] = None,
-        stream_mode: Union[StreamMode, list[StreamMode]] = "values",
+        stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
         stream_subgraphs: bool = False,
         metadata: Optional[dict] = None,
         config: Optional[Config] = None,
         checkpoint: Optional[Checkpoint] = None,
         checkpoint_id: Optional[str] = None,
-        interrupt_before: Optional[Union[All, list[str]]] = None,
-        interrupt_after: Optional[Union[All, list[str]]] = None,
-        feedback_keys: Optional[list[str]] = None,
+        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
+        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
+        feedback_keys: Optional[Sequence[str]] = None,
         on_disconnect: Optional[DisconnectMode] = None,
         on_completion: Optional[OnCompletionBehavior] = None,
         webhook: Optional[str] = None,
         multitask_strategy: Optional[MultitaskStrategy] = None,
+        if_not_exists: Optional[IfNotExists] = None,
         after_seconds: Optional[int] = None,
     ) -> AsyncIterator[StreamPart]:
         """Create a run and stream the results.
@@ -1248,6 +1252,8 @@ class RunsClient:
             webhook: Webhook to call after LangGraph API call is done.
             multitask_strategy: Multitask strategy to use.
                 Must be one of 'reject', 'interrupt', 'rollback', or 'enqueue'.
+            if_not_exists: How to handle missing thread. Defaults to 'reject'.
+                Must be either 'reject' (raise error if missing), or 'create' (create new thread).
             after_seconds: The number of seconds to wait before starting the run.
                 Use to schedule future runs.
 
@@ -1293,6 +1299,7 @@ class RunsClient:
             "checkpoint": checkpoint,
             "checkpoint_id": checkpoint_id,
             "multitask_strategy": multitask_strategy,
+            "if_not_exists": if_not_exists,
             "on_disconnect": on_disconnect,
             "on_completion": on_completion,
             "after_seconds": after_seconds,
@@ -1313,14 +1320,15 @@ class RunsClient:
         assistant_id: str,
         *,
         input: Optional[dict] = None,
-        stream_mode: Union[StreamMode, list[StreamMode]] = "values",
+        stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
         stream_subgraphs: bool = False,
         metadata: Optional[dict] = None,
         config: Optional[Config] = None,
-        interrupt_before: Optional[Union[All, list[str]]] = None,
-        interrupt_after: Optional[Union[All, list[str]]] = None,
+        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
+        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
         webhook: Optional[str] = None,
         on_completion: Optional[OnCompletionBehavior] = None,
+        if_not_exists: Optional[IfNotExists] = None,
         after_seconds: Optional[int] = None,
     ) -> Run: ...
 
@@ -1331,16 +1339,17 @@ class RunsClient:
         assistant_id: str,
         *,
         input: Optional[dict] = None,
-        stream_mode: Union[StreamMode, list[StreamMode]] = "values",
+        stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
         stream_subgraphs: bool = False,
         metadata: Optional[dict] = None,
         config: Optional[Config] = None,
         checkpoint: Optional[Checkpoint] = None,
         checkpoint_id: Optional[str] = None,
-        interrupt_before: Optional[Union[All, list[str]]] = None,
-        interrupt_after: Optional[Union[All, list[str]]] = None,
+        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
+        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
         webhook: Optional[str] = None,
         multitask_strategy: Optional[MultitaskStrategy] = None,
+        if_not_exists: Optional[IfNotExists] = None,
         after_seconds: Optional[int] = None,
     ) -> Run: ...
 
@@ -1350,16 +1359,17 @@ class RunsClient:
         assistant_id: str,
         *,
         input: Optional[dict] = None,
-        stream_mode: Union[StreamMode, list[StreamMode]] = "values",
+        stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
         stream_subgraphs: bool = False,
         metadata: Optional[dict] = None,
         config: Optional[Config] = None,
         checkpoint: Optional[Checkpoint] = None,
         checkpoint_id: Optional[str] = None,
-        interrupt_before: Optional[Union[All, list[str]]] = None,
-        interrupt_after: Optional[Union[All, list[str]]] = None,
+        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
+        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
         webhook: Optional[str] = None,
         multitask_strategy: Optional[MultitaskStrategy] = None,
+        if_not_exists: Optional[IfNotExists] = None,
         on_completion: Optional[OnCompletionBehavior] = None,
         after_seconds: Optional[int] = None,
     ) -> Run:
@@ -1383,6 +1393,8 @@ class RunsClient:
                 Must be one of 'reject', 'interrupt', 'rollback', or 'enqueue'.
             on_completion: Whether to delete or keep the thread created for a stateless run.
                 Must be one of 'delete' or 'keep'.
+            if_not_exists: How to handle missing thread. Defaults to 'reject'.
+                Must be either 'reject' (raise error if missing), or 'create' (create new thread).
             after_seconds: The number of seconds to wait before starting the run.
                 Use to schedule future runs.
 
@@ -1466,6 +1478,7 @@ class RunsClient:
             "checkpoint": checkpoint,
             "checkpoint_id": checkpoint_id,
             "multitask_strategy": multitask_strategy,
+            "if_not_exists": if_not_exists,
             "on_completion": on_completion,
             "after_seconds": after_seconds,
         }
@@ -1495,11 +1508,12 @@ class RunsClient:
         config: Optional[Config] = None,
         checkpoint: Optional[Checkpoint] = None,
         checkpoint_id: Optional[str] = None,
-        interrupt_before: Optional[Union[All, list[str]]] = None,
-        interrupt_after: Optional[Union[All, list[str]]] = None,
+        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
+        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
         webhook: Optional[str] = None,
         on_disconnect: Optional[DisconnectMode] = None,
         multitask_strategy: Optional[MultitaskStrategy] = None,
+        if_not_exists: Optional[IfNotExists] = None,
         after_seconds: Optional[int] = None,
     ) -> Union[list[dict], dict[str, Any]]: ...
 
@@ -1512,11 +1526,12 @@ class RunsClient:
         input: Optional[dict] = None,
         metadata: Optional[dict] = None,
         config: Optional[Config] = None,
-        interrupt_before: Optional[Union[All, list[str]]] = None,
-        interrupt_after: Optional[Union[All, list[str]]] = None,
+        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
+        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
         webhook: Optional[str] = None,
         on_disconnect: Optional[DisconnectMode] = None,
         on_completion: Optional[OnCompletionBehavior] = None,
+        if_not_exists: Optional[IfNotExists] = None,
         after_seconds: Optional[int] = None,
     ) -> Union[list[dict], dict[str, Any]]: ...
 
@@ -1530,12 +1545,13 @@ class RunsClient:
         config: Optional[Config] = None,
         checkpoint: Optional[Checkpoint] = None,
         checkpoint_id: Optional[str] = None,
-        interrupt_before: Optional[Union[All, list[str]]] = None,
-        interrupt_after: Optional[Union[All, list[str]]] = None,
+        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
+        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
         webhook: Optional[str] = None,
         on_disconnect: Optional[DisconnectMode] = None,
         on_completion: Optional[OnCompletionBehavior] = None,
         multitask_strategy: Optional[MultitaskStrategy] = None,
+        if_not_exists: Optional[IfNotExists] = None,
         after_seconds: Optional[int] = None,
     ) -> Union[list[dict], dict[str, Any]]:
         """Create a run, wait until it finishes and return the final state.
@@ -1558,6 +1574,8 @@ class RunsClient:
                 Must be one of 'delete' or 'keep'.
             multitask_strategy: Multitask strategy to use.
                 Must be one of 'reject', 'interrupt', 'rollback', or 'enqueue'.
+            if_not_exists: How to handle missing thread. Defaults to 'reject'.
+                Must be either 'reject' (raise error if missing), or 'create' (create new thread).
             after_seconds: The number of seconds to wait before starting the run.
                 Use to schedule future runs.
 
@@ -1619,6 +1637,7 @@ class RunsClient:
             "checkpoint": checkpoint,
             "checkpoint_id": checkpoint_id,
             "multitask_strategy": multitask_strategy,
+            "if_not_exists": if_not_exists,
             "on_disconnect": on_disconnect,
             "on_completion": on_completion,
             "after_seconds": after_seconds,
@@ -3232,18 +3251,19 @@ class SyncRunsClient:
         assistant_id: str,
         *,
         input: Optional[dict] = None,
-        stream_mode: Union[StreamMode, list[StreamMode]] = "values",
+        stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
         stream_subgraphs: bool = False,
         metadata: Optional[dict] = None,
         config: Optional[Config] = None,
         checkpoint: Optional[Checkpoint] = None,
         checkpoint_id: Optional[str] = None,
-        interrupt_before: Optional[Union[All, list[str]]] = None,
-        interrupt_after: Optional[Union[All, list[str]]] = None,
-        feedback_keys: Optional[list[str]] = None,
+        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
+        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
+        feedback_keys: Optional[Sequence[str]] = None,
         on_disconnect: Optional[DisconnectMode] = None,
         webhook: Optional[str] = None,
         multitask_strategy: Optional[MultitaskStrategy] = None,
+        if_not_exists: Optional[IfNotExists] = None,
         after_seconds: Optional[int] = None,
     ) -> Iterator[StreamPart]: ...
 
@@ -3254,15 +3274,16 @@ class SyncRunsClient:
         assistant_id: str,
         *,
         input: Optional[dict] = None,
-        stream_mode: Union[StreamMode, list[StreamMode]] = "values",
+        stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
         stream_subgraphs: bool = False,
         metadata: Optional[dict] = None,
         config: Optional[Config] = None,
-        interrupt_before: Optional[Union[All, list[str]]] = None,
-        interrupt_after: Optional[Union[All, list[str]]] = None,
-        feedback_keys: Optional[list[str]] = None,
+        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
+        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
+        feedback_keys: Optional[Sequence[str]] = None,
         on_disconnect: Optional[DisconnectMode] = None,
         on_completion: Optional[OnCompletionBehavior] = None,
+        if_not_exists: Optional[IfNotExists] = None,
         webhook: Optional[str] = None,
         after_seconds: Optional[int] = None,
     ) -> Iterator[StreamPart]: ...
@@ -3273,19 +3294,20 @@ class SyncRunsClient:
         assistant_id: str,
         *,
         input: Optional[dict] = None,
-        stream_mode: Union[StreamMode, list[StreamMode]] = "values",
+        stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
         stream_subgraphs: bool = False,
         metadata: Optional[dict] = None,
         config: Optional[Config] = None,
         checkpoint: Optional[Checkpoint] = None,
         checkpoint_id: Optional[str] = None,
-        interrupt_before: Optional[Union[All, list[str]]] = None,
-        interrupt_after: Optional[Union[All, list[str]]] = None,
-        feedback_keys: Optional[list[str]] = None,
+        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
+        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
+        feedback_keys: Optional[Sequence[str]] = None,
         on_disconnect: Optional[DisconnectMode] = None,
         on_completion: Optional[OnCompletionBehavior] = None,
         webhook: Optional[str] = None,
         multitask_strategy: Optional[MultitaskStrategy] = None,
+        if_not_exists: Optional[IfNotExists] = None,
         after_seconds: Optional[int] = None,
     ) -> Iterator[StreamPart]:
         """Create a run and stream the results.
@@ -3311,6 +3333,8 @@ class SyncRunsClient:
             webhook: Webhook to call after LangGraph API call is done.
             multitask_strategy: Multitask strategy to use.
                 Must be one of 'reject', 'interrupt', 'rollback', or 'enqueue'.
+            if_not_exists: How to handle missing thread. Defaults to 'reject'.
+                Must be either 'reject' (raise error if missing), or 'create' (create new thread).
             after_seconds: The number of seconds to wait before starting the run.
                 Use to schedule future runs.
 
@@ -3356,6 +3380,7 @@ class SyncRunsClient:
             "checkpoint": checkpoint,
             "checkpoint_id": checkpoint_id,
             "multitask_strategy": multitask_strategy,
+            "if_not_exists": if_not_exists,
             "on_disconnect": on_disconnect,
             "on_completion": on_completion,
             "after_seconds": after_seconds,
@@ -3376,14 +3401,15 @@ class SyncRunsClient:
         assistant_id: str,
         *,
         input: Optional[dict] = None,
-        stream_mode: Union[StreamMode, list[StreamMode]] = "values",
+        stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
         stream_subgraphs: bool = False,
         metadata: Optional[dict] = None,
         config: Optional[Config] = None,
-        interrupt_before: Optional[Union[All, list[str]]] = None,
-        interrupt_after: Optional[Union[All, list[str]]] = None,
+        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
+        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
         webhook: Optional[str] = None,
         on_completion: Optional[OnCompletionBehavior] = None,
+        if_not_exists: Optional[IfNotExists] = None,
         after_seconds: Optional[int] = None,
     ) -> Run: ...
 
@@ -3394,16 +3420,17 @@ class SyncRunsClient:
         assistant_id: str,
         *,
         input: Optional[dict] = None,
-        stream_mode: Union[StreamMode, list[StreamMode]] = "values",
+        stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
         stream_subgraphs: bool = False,
         metadata: Optional[dict] = None,
         config: Optional[Config] = None,
         checkpoint: Optional[Checkpoint] = None,
         checkpoint_id: Optional[str] = None,
-        interrupt_before: Optional[Union[All, list[str]]] = None,
-        interrupt_after: Optional[Union[All, list[str]]] = None,
+        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
+        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
         webhook: Optional[str] = None,
         multitask_strategy: Optional[MultitaskStrategy] = None,
+        if_not_exists: Optional[IfNotExists] = None,
         after_seconds: Optional[int] = None,
     ) -> Run: ...
 
@@ -3413,17 +3440,18 @@ class SyncRunsClient:
         assistant_id: str,
         *,
         input: Optional[dict] = None,
-        stream_mode: Union[StreamMode, list[StreamMode]] = "values",
+        stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
         stream_subgraphs: bool = False,
         metadata: Optional[dict] = None,
         config: Optional[Config] = None,
         checkpoint: Optional[Checkpoint] = None,
         checkpoint_id: Optional[str] = None,
-        interrupt_before: Optional[Union[All, list[str]]] = None,
-        interrupt_after: Optional[Union[All, list[str]]] = None,
+        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
+        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
         webhook: Optional[str] = None,
         multitask_strategy: Optional[MultitaskStrategy] = None,
         on_completion: Optional[OnCompletionBehavior] = None,
+        if_not_exists: Optional[IfNotExists] = None,
         after_seconds: Optional[int] = None,
     ) -> Run:
         """Create a background run.
@@ -3446,6 +3474,8 @@ class SyncRunsClient:
                 Must be one of 'reject', 'interrupt', 'rollback', or 'enqueue'.
             on_completion: Whether to delete or keep the thread created for a stateless run.
                 Must be one of 'delete' or 'keep'.
+            if_not_exists: How to handle missing thread. Defaults to 'reject'.
+                Must be either 'reject' (raise error if missing), or 'create' (create new thread).
             after_seconds: The number of seconds to wait before starting the run.
                 Use to schedule future runs.
 
@@ -3529,6 +3559,7 @@ class SyncRunsClient:
             "checkpoint": checkpoint,
             "checkpoint_id": checkpoint_id,
             "multitask_strategy": multitask_strategy,
+            "if_not_exists": if_not_exists,
             "on_completion": on_completion,
             "after_seconds": after_seconds,
         }
@@ -3558,11 +3589,12 @@ class SyncRunsClient:
         config: Optional[Config] = None,
         checkpoint: Optional[Checkpoint] = None,
         checkpoint_id: Optional[str] = None,
-        interrupt_before: Optional[Union[All, list[str]]] = None,
-        interrupt_after: Optional[Union[All, list[str]]] = None,
+        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
+        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
         webhook: Optional[str] = None,
         on_disconnect: Optional[DisconnectMode] = None,
         multitask_strategy: Optional[MultitaskStrategy] = None,
+        if_not_exists: Optional[IfNotExists] = None,
         after_seconds: Optional[int] = None,
     ) -> Union[list[dict], dict[str, Any]]: ...
 
@@ -3575,11 +3607,12 @@ class SyncRunsClient:
         input: Optional[dict] = None,
         metadata: Optional[dict] = None,
         config: Optional[Config] = None,
-        interrupt_before: Optional[Union[All, list[str]]] = None,
-        interrupt_after: Optional[Union[All, list[str]]] = None,
+        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
+        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
         webhook: Optional[str] = None,
         on_disconnect: Optional[DisconnectMode] = None,
         on_completion: Optional[OnCompletionBehavior] = None,
+        if_not_exists: Optional[IfNotExists] = None,
         after_seconds: Optional[int] = None,
     ) -> Union[list[dict], dict[str, Any]]: ...
 
@@ -3593,12 +3626,13 @@ class SyncRunsClient:
         config: Optional[Config] = None,
         checkpoint: Optional[Checkpoint] = None,
         checkpoint_id: Optional[str] = None,
-        interrupt_before: Optional[Union[All, list[str]]] = None,
-        interrupt_after: Optional[Union[All, list[str]]] = None,
+        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
+        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
         webhook: Optional[str] = None,
         on_disconnect: Optional[DisconnectMode] = None,
         on_completion: Optional[OnCompletionBehavior] = None,
         multitask_strategy: Optional[MultitaskStrategy] = None,
+        if_not_exists: Optional[IfNotExists] = None,
         after_seconds: Optional[int] = None,
     ) -> Union[list[dict], dict[str, Any]]:
         """Create a run, wait until it finishes and return the final state.
@@ -3621,6 +3655,8 @@ class SyncRunsClient:
                 Must be one of 'delete' or 'keep'.
             multitask_strategy: Multitask strategy to use.
                 Must be one of 'reject', 'interrupt', 'rollback', or 'enqueue'.
+            if_not_exists: How to handle missing thread. Defaults to 'reject'.
+                Must be either 'reject' (raise error if missing), or 'create' (create new thread).
             after_seconds: The number of seconds to wait before starting the run.
                 Use to schedule future runs.
 
@@ -3682,6 +3718,7 @@ class SyncRunsClient:
             "checkpoint": checkpoint,
             "checkpoint_id": checkpoint_id,
             "multitask_strategy": multitask_strategy,
+            "if_not_exists": if_not_exists,
             "on_disconnect": on_disconnect,
             "on_completion": on_completion,
             "after_seconds": after_seconds,

libs/sdk-py/langgraph_sdk/schema.py

@@ -17,12 +17,13 @@ Represents the status of a run:
 - "interrupted": The run was manually stopped or interrupted.
 """
 
-ThreadStatus = Literal["idle", "busy", "interrupted"]
+ThreadStatus = Literal["idle", "busy", "interrupted", "error"]
 """
 Represents the status of a thread:
 - "idle": The thread is not currently processing any task.
 - "busy": The thread is actively processing a task.
 - "interrupted": The thread's execution was interrupted.
+- "error": An exception occurred during task processing.
 """
 
 StreamMode = Literal["values", "messages", "updates", "events", "debug", "custom"]
@@ -69,6 +70,13 @@ Defines action after completion:
 All = Literal["*"]
 """Represents a wildcard or 'all' selector."""
 
+IfNotExists = Literal["create", "reject"]
+"""
+Specifies behavior if the thread doesn't exist:
+- "create": Create a new thread if it doesn't exist.
+- "reject": Reject the operation if the thread doesn't exist.
+"""
+
 
 class Config(TypedDict, total=False):
     """Configuration options for a call."""

libs/sdk-py/pyproject.toml

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